demeter

 1# tools/db_tools.py
 2import os
 3from qdrant_client import QdrantClient, models
 4from qdrant_client.models import PointStruct
 5from dotenv import load_dotenv
 6
 7load_dotenv()
 8
 9
10class DBTools:
11    def __init__(self, host=None, collection_name="Farm_Memory"):
12        if host is None:
13            host = os.environ.get("QDRANT_URL", "http://localhost:6333")
14        self.client = QdrantClient(url=host, api_key=os.environ.get("QDRANT_API_KEY"))
15        self.collection_name = collection_name
16        self.vector_size = 516  # As seen in Qdrant/Setup.py
17
18    def setup_database(self):
19        """Creates or resets the memory collection."""
20        self.client.recreate_collection(
21            collection_name=self.collection_name,
22            vectors_config=models.VectorParams(
23                size=self.vector_size, distance=models.Distance.COSINE
24            ),
25        )
26        print(f"[DB] Collection '{self.collection_name}' ready.")
27
28    def store_fmu(self, fmu_data):
29        """
30        Stores a Farm Memory Unit (FMU).
31        Derived from Qdrant/Store.py
32        """
33        point = PointStruct(
34            id=fmu_data.id, vector=fmu_data.vector, payload=fmu_data.metadata
35        )
36        self.client.upsert(collection_name=self.collection_name, points=[point])
37        print(f"[DB] Stored FMU ID: {fmu_data.id}")
38
39    def search_similar(self, vector, limit=5):
40        """
41        Finds similar past states.
42        Derived from backend/server/main.py search endpoint
43        """
44        hits = self.client.search(
45            collection_name=self.collection_name, query_vector=vector, limit=limit
46        )
47        return [{"score": hit.score, "payload": hit.payload} for hit in hits]