Add metadata filtering capabilities to memory retrieval

src/mcp_server_qdrant/qdrant.py

@@ -1,5 +1,5 @@
 import uuid
-from typing import Optional
+from typing import Any, Dict, List, Optional
 
 from qdrant_client import AsyncQdrantClient, models
 
@@ -53,16 +53,22 @@ async def _ensure_collection_exists(self):
                 },
             )
 
-    async def store_memory(self, information: str):
+    async def store_memory(self, information: str, metadata: Optional[Dict[str, Any]] = None):
         """
         Store a memory in the Qdrant collection.
         :param information: The information to store.
+        :param metadata: Optional metadata to associate with the memory.
         """
         await self._ensure_collection_exists()
 
         # Embed the document
         embeddings = await self._embedding_provider.embed_documents([information])
 
+        # Prepare payload with document and optional metadata
+        payload = {"document": information}
+        if metadata:
+            payload["metadata"] = metadata
+
         # Add to Qdrant
         vector_name = self._embedding_provider.get_vector_name()
         await self._client.upsert(
@@ -71,16 +77,23 @@ async def store_memory(self, information: str):
                 models.PointStruct(
                     id=uuid.uuid4().hex,
                     vector={vector_name: embeddings[0]},
-                    payload={"document": information},
+                    payload=payload,
                 )
             ],
         )
 
-    async def find_memories(self, query: str) -> list[str]:
+    async def find_memories(
+        self, 
+        query: str, 
+        limit: int = 10, 
+        filter_metadata: Optional[Dict[str, Any]] = None
+    ) -> List[Dict[str, Any]]:
         """
         Find memories in the Qdrant collection. If there are no memories found, an empty list is returned.
         :param query: The query to use for the search.
-        :return: A list of memories found.
+        :param limit: Maximum number of results to return.
+        :param filter_metadata: Optional metadata filter to apply to the search.
+        :return: A list of dictionaries containing document content and metadata.
         """
         collection_exists = await self._client.collection_exists(self._collection_name)
         if not collection_exists:
@@ -90,11 +103,33 @@ async def find_memories(self, query: str) -> list[str]:
         query_vector = await self._embedding_provider.embed_query(query)
         vector_name = self._embedding_provider.get_vector_name()
 
+        # Prepare filter if needed
+        filter_condition = None
+        if filter_metadata:
+            filter_condition = models.Filter(
+                must=[
+                    models.FieldCondition(
+                        key=f"metadata.{key}",
+                        match=models.MatchValue(value=value)
+                    )
+                    for key, value in filte`r_metadata.items()
+                ]
+            )
+
         # Search in Qdrant
         search_results = await self._client.search(
             collection_name=self._collection_name,
             query_vector=models.NamedVector(name=vector_name, vector=query_vector),
-            limit=10,
+            limit=limit,
+            query_filter=filter_condition,
         )
 
-        return [result.payload["document"] for result in search_results]
+        # Return both document and metadata if available
+        return [
+            {
+                "document": result.payload["document"],
+                "metadata": result.payload.get("metadata", {}),
+                "score": result.score
+            } 
+            for result in search_results
+        ]