diff --git a/README.md b/README.md index c9f5796..f7a078d 100644 --- a/README.md +++ b/README.md @@ -25,15 +25,15 @@ It acts as a semantic memory layer on top of the Qdrant database. - Input: - `information` (string): Information to store - `metadata` (JSON): Optional metadata to store - - `collection_name` (string): Name of the collection to store the information in, optional. If not provided, - the default collection name will be used. + - `collection_name` (string): Name of the collection to store the information in. This field is required if there are no default collection name. + If there is a default collection name, this field is not enabled. - Returns: Confirmation message 2. `qdrant-find` - Retrieve relevant information from the Qdrant database - Input: - `query` (string): Query to use for searching - - `collection_name` (string): Name of the collection to store the information in, optional. If not provided, - the default collection name will be used. + - `collection_name` (string): Name of the collection to store the information in. This field is required if there are no default collection name. + If there is a default collection name, this field is not enabled. - Returns: Information stored in the Qdrant database as separate messages ## Environment Variables @@ -44,7 +44,7 @@ The configuration of the server is done using environment variables: |--------------------------|---------------------------------------------------------------------|-------------------------------------------------------------------| | `QDRANT_URL` | URL of the Qdrant server | None | | `QDRANT_API_KEY` | API key for the Qdrant server | None | -| `COLLECTION_NAME` | Name of the default collection to use. | *Required* | +| `COLLECTION_NAME` | Name of the default collection to use. | None | | `QDRANT_LOCAL_PATH` | Path to the local Qdrant database (alternative to `QDRANT_URL`) | None | | `EMBEDDING_PROVIDER` | Embedding provider to use (currently only "fastembed" is supported) | `fastembed` | | `EMBEDDING_MODEL` | Name of the embedding model to use | `sentence-transformers/all-MiniLM-L6-v2` | @@ -245,6 +245,15 @@ Claude Code should be already able to: 1. Use the `qdrant-store` tool to store code snippets with descriptions. 2. Use the `qdrant-find` tool to search for relevant code snippets using natural language. +### Run MCP server in Development Mode + +The MCP server can be run in development mode using the `mcp dev` command. This will start the server and open the MCP +inspector in your browser. + +```shell +COLLECTION_NAME=mcp-dev mcp dev src/mcp_server_qdrant/server.py +``` + ## Contributing If you have suggestions for how mcp-server-qdrant could be improved, or want to report a bug, open an issue! diff --git a/src/mcp_server_qdrant/mcp_server.py b/src/mcp_server_qdrant/mcp_server.py new file mode 100644 index 0000000..2b92342 --- /dev/null +++ b/src/mcp_server_qdrant/mcp_server.py @@ -0,0 +1,162 @@ +import json +import logging +from typing import List + +from mcp.server.fastmcp import Context, FastMCP + +from mcp_server_qdrant.embeddings.factory import create_embedding_provider +from mcp_server_qdrant.qdrant import Entry, Metadata, QdrantConnector +from mcp_server_qdrant.settings import ( + EmbeddingProviderSettings, + QdrantSettings, + ToolSettings, +) + +logger = logging.getLogger(__name__) + + +# FastMCP is an alternative interface for declaring the capabilities +# of the server. Its API is based on FastAPI. +class QdrantMCPServer(FastMCP): + """ + A MCP server for Qdrant. + """ + + def __init__( + self, + tool_settings: ToolSettings, + qdrant_settings: QdrantSettings, + embedding_provider_settings: EmbeddingProviderSettings, + name: str = "mcp-server-qdrant", + ): + self.tool_settings = tool_settings + self.qdrant_settings = qdrant_settings + self.embedding_provider_settings = embedding_provider_settings + + self.embedding_provider = create_embedding_provider(embedding_provider_settings) + self.qdrant_connector = QdrantConnector( + qdrant_settings.location, + qdrant_settings.api_key, + qdrant_settings.collection_name, + self.embedding_provider, + qdrant_settings.local_path, + ) + + super().__init__(name=name) + + self.setup_tools() + + def format_entry(self, entry: Entry) -> str: + """ + Feel free to override this method in your subclass to customize the format of the entry. + """ + entry_metadata = json.dumps(entry.metadata) if entry.metadata else "" + return f"{entry.content}{entry_metadata}" + + def setup_tools(self): + async def store( + ctx: Context, + information: str, + collection_name: str, + # The `metadata` parameter is defined as non-optional, but it can be None. + # If we set it to be optional, some of the MCP clients, like Cursor, cannot + # handle the optional parameter correctly. + metadata: Metadata = None, + ) -> str: + """ + Store some information in Qdrant. + :param ctx: The context for the request. + :param information: The information to store. + :param metadata: JSON metadata to store with the information, optional. + :param collection_name: The name of the collection to store the information in, optional. If not provided, + the default collection is used. + :return: A message indicating that the information was stored. + """ + await ctx.debug(f"Storing information {information} in Qdrant") + + entry = Entry(content=information, metadata=metadata) + + await self.qdrant_connector.store(entry, collection_name=collection_name) + if collection_name: + return f"Remembered: {information} in collection {collection_name}" + return f"Remembered: {information}" + + async def store_with_default_collection( + ctx: Context, + information: str, + metadata: Metadata = None, + ) -> str: + return await store( + ctx, information, metadata, self.qdrant_settings.collection_name + ) + + async def find( + ctx: Context, + query: str, + collection_name: str, + ) -> List[str]: + """ + Find memories in Qdrant. + :param ctx: The context for the request. + :param query: The query to use for the search. + :param collection_name: The name of the collection to search in, optional. If not provided, + the default collection is used. + :param limit: The maximum number of entries to return, optional. Default is 10. + :return: A list of entries found. + """ + await ctx.debug(f"Finding results for query {query}") + if collection_name: + await ctx.debug( + f"Overriding the collection name with {collection_name}" + ) + + entries = await self.qdrant_connector.search( + query, + collection_name=collection_name, + limit=self.qdrant_settings.search_limit, + ) + if not entries: + return [f"No information found for the query '{query}'"] + content = [ + f"Results for the query '{query}'", + ] + for entry in entries: + content.append(self.format_entry(entry)) + return content + + async def find_with_default_collection( + ctx: Context, + query: str, + ) -> List[str]: + return await find(ctx, query, self.qdrant_settings.collection_name) + + # Register the tools depending on the configuration + + if self.qdrant_settings.collection_name: + self.add_tool( + find_with_default_collection, + name="qdrant-find", + description=self.tool_settings.tool_find_description, + ) + else: + self.add_tool( + find, + name="qdrant-find", + description=self.tool_settings.tool_find_description, + ) + + if not self.qdrant_settings.read_only: + # Those methods can modify the database + + if self.qdrant_settings.collection_name: + self.add_tool( + store_with_default_collection, + name="qdrant-store", + description=self.tool_settings.tool_store_description, + ) + else: + self.add_tool( + store, + name="qdrant-store", + description=self.tool_settings.tool_store_description, + ) diff --git a/src/mcp_server_qdrant/qdrant.py b/src/mcp_server_qdrant/qdrant.py index e12459d..b2a7d1b 100644 --- a/src/mcp_server_qdrant/qdrant.py +++ b/src/mcp_server_qdrant/qdrant.py @@ -66,6 +66,8 @@ class QdrantConnector: await self._ensure_collection_exists(collection_name) # Embed the document + # ToDo: instead of embedding text explicitly, use `models.Document`, + # it should unlock usage of server-side inference. embeddings = await self._embedding_provider.embed_documents([entry.content]) # Add to Qdrant @@ -99,13 +101,17 @@ class QdrantConnector: return [] # Embed the query + # ToDo: instead of embedding text explicitly, use `models.Document`, + # it should unlock usage of server-side inference. + query_vector = await self._embedding_provider.embed_query(query) vector_name = self._embedding_provider.get_vector_name() # Search in Qdrant - search_results = await self._client.search( + search_results = await self._client.query_points( collection_name=collection_name, - query_vector=models.NamedVector(name=vector_name, vector=query_vector), + query=query_vector, + using=vector_name, limit=limit, ) @@ -114,7 +120,7 @@ class QdrantConnector: content=result.payload["document"], metadata=result.payload.get("metadata"), ) - for result in search_results + for result in search_results.points ] async def _ensure_collection_exists(self, collection_name: str): diff --git a/src/mcp_server_qdrant/server.py b/src/mcp_server_qdrant/server.py index 4a726c0..3f8f303 100644 --- a/src/mcp_server_qdrant/server.py +++ b/src/mcp_server_qdrant/server.py @@ -1,136 +1,12 @@ -import json -import logging -from contextlib import asynccontextmanager -from typing import AsyncIterator, List, Optional - -from mcp.server import Server -from mcp.server.fastmcp import Context, FastMCP - -from mcp_server_qdrant.embeddings.factory import create_embedding_provider -from mcp_server_qdrant.qdrant import Entry, Metadata, QdrantConnector +from mcp_server_qdrant.mcp_server import QdrantMCPServer from mcp_server_qdrant.settings import ( EmbeddingProviderSettings, QdrantSettings, ToolSettings, ) -logger = logging.getLogger(__name__) - - -@asynccontextmanager -async def server_lifespan(server: Server) -> AsyncIterator[dict]: # noqa - """ - Context manager to handle the lifespan of the server. - This is used to configure the embedding provider and Qdrant connector. - All the configuration is now loaded from the environment variables. - Settings handle that for us. - """ - try: - # Embedding provider is created with a factory function so we can add - # some more providers in the future. Currently, only FastEmbed is supported. - embedding_provider_settings = EmbeddingProviderSettings() - embedding_provider = create_embedding_provider(embedding_provider_settings) - logger.info( - f"Using embedding provider {embedding_provider_settings.provider_type} with " - f"model {embedding_provider_settings.model_name}" - ) - - qdrant_configuration = QdrantSettings() - qdrant_connector = QdrantConnector( - qdrant_configuration.location, - qdrant_configuration.api_key, - qdrant_configuration.collection_name, - embedding_provider, - qdrant_configuration.local_path, - ) - logger.info( - f"Connecting to Qdrant at {qdrant_configuration.get_qdrant_location()}" - ) - - yield { - "embedding_provider": embedding_provider, - "qdrant_connector": qdrant_connector, - } - except Exception as e: - logger.error(e) - raise e - finally: - pass - - -# FastMCP is an alternative interface for declaring the capabilities -# of the server. Its API is based on FastAPI. -mcp = FastMCP("mcp-server-qdrant", lifespan=server_lifespan) - -# Load the tool settings from the env variables, if they are set, -# or use the default values otherwise. -tool_settings = ToolSettings() - - -@mcp.tool(name="qdrant-store", description=tool_settings.tool_store_description) -async def store( - ctx: Context, - information: str, - # The `metadata` parameter is defined as non-optional, but it can be None. - # If we set it to be optional, some of the MCP clients, like Cursor, cannot - # handle the optional parameter correctly. - metadata: Metadata = None, - collection_name: Optional[str] = None, -) -> str: - """ - Store some information in Qdrant. - :param ctx: The context for the request. - :param information: The information to store. - :param metadata: JSON metadata to store with the information, optional. - :param collection_name: The name of the collection to store the information in, optional. If not provided, - the default collection is used. - :return: A message indicating that the information was stored. - """ - await ctx.debug(f"Storing information {information} in Qdrant") - qdrant_connector: QdrantConnector = ctx.request_context.lifespan_context[ - "qdrant_connector" - ] - entry = Entry(content=information, metadata=metadata) - await qdrant_connector.store(entry, collection_name=collection_name) - if collection_name: - return f"Remembered: {information} in collection {collection_name}" - return f"Remembered: {information}" - - -@mcp.tool(name="qdrant-find", description=tool_settings.tool_find_description) -async def find( - ctx: Context, - query: str, - collection_name: Optional[str] = None, - limit: int = 10, -) -> List[str]: - """ - Find memories in Qdrant. - :param ctx: The context for the request. - :param query: The query to use for the search. - :param collection_name: The name of the collection to search in, optional. If not provided, - the default collection is used. - :param limit: The maximum number of entries to return, optional. Default is 10. - :return: A list of entries found. - """ - await ctx.debug(f"Finding results for query {query}") - if collection_name: - await ctx.debug(f"Overriding the collection name with {collection_name}") - qdrant_connector: QdrantConnector = ctx.request_context.lifespan_context[ - "qdrant_connector" - ] - entries = await qdrant_connector.search( - query, collection_name=collection_name, limit=limit - ) - if not entries: - return [f"No information found for the query '{query}'"] - content = [ - f"Results for the query '{query}'", - ] - for entry in entries: - # Format the metadata as a JSON string and produce XML-like output - entry_metadata = json.dumps(entry.metadata) if entry.metadata else "" - content.append( - f"{entry.content}{entry_metadata}" - ) - return content +mcp = QdrantMCPServer( + tool_settings=ToolSettings(), + qdrant_settings=QdrantSettings(), + embedding_provider_settings=EmbeddingProviderSettings(), +) diff --git a/src/mcp_server_qdrant/settings.py b/src/mcp_server_qdrant/settings.py index cad8834..8855b69 100644 --- a/src/mcp_server_qdrant/settings.py +++ b/src/mcp_server_qdrant/settings.py @@ -53,10 +53,16 @@ class QdrantSettings(BaseSettings): location: Optional[str] = Field(default=None, validation_alias="QDRANT_URL") api_key: Optional[str] = Field(default=None, validation_alias="QDRANT_API_KEY") - collection_name: str = Field(validation_alias="COLLECTION_NAME") + collection_name: Optional[str] = Field( + default=None, validation_alias="COLLECTION_NAME" + ) local_path: Optional[str] = Field( default=None, validation_alias="QDRANT_LOCAL_PATH" ) + search_limit: Optional[int] = Field( + default=None, validation_alias="QDRANT_SEARCH_LIMIT" + ) + read_only: bool = Field(default=False, validation_alias="QDRANT_READ_ONLY") def get_qdrant_location(self) -> str: """ diff --git a/tests/test_fastembed_integration.py b/tests/test_fastembed_integration.py index 8381dd7..3a34db5 100644 --- a/tests/test_fastembed_integration.py +++ b/tests/test_fastembed_integration.py @@ -53,7 +53,7 @@ class TestFastEmbedProviderIntegration: # The embeddings should be identical for the same input np.testing.assert_array_almost_equal(np.array(embedding), np.array(embedding2)) - def test_get_vector_name(self): + async def test_get_vector_name(self): """Test that the vector name is generated correctly.""" provider = FastEmbedProvider("sentence-transformers/all-MiniLM-L6-v2") vector_name = provider.get_vector_name() diff --git a/tests/test_settings.py b/tests/test_settings.py index 8390ff8..d5b2d37 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -1,8 +1,6 @@ import os from unittest.mock import patch -import pytest - from mcp_server_qdrant.embeddings.types import EmbeddingProviderType from mcp_server_qdrant.settings import ( DEFAULT_TOOL_FIND_DESCRIPTION, @@ -16,9 +14,9 @@ from mcp_server_qdrant.settings import ( class TestQdrantSettings: def test_default_values(self): """Test that required fields raise errors when not provided.""" - with pytest.raises(ValueError): - # Should raise error because required fields are missing - QdrantSettings() + + # Should not raise error because there are no required fields + QdrantSettings() @patch.dict( os.environ,