diff --git a/src/mcp_server_qdrant/qdrant.py b/src/mcp_server_qdrant/qdrant.py index 4340fd0..d444813 100644 --- a/src/mcp_server_qdrant/qdrant.py +++ b/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 + ]