diff --git a/src/mcp_server_qdrant/qdrant.py b/src/mcp_server_qdrant/qdrant.py index b848045..5a9cc23 100644 --- a/src/mcp_server_qdrant/qdrant.py +++ b/src/mcp_server_qdrant/qdrant.py @@ -9,23 +9,25 @@ class QdrantConnector: :param qdrant_api_key: The API key to use for the Qdrant server. :param collection_name: The name of the collection to use. :param fastembed_model_name: The name of the FastEmbed model to use. + :param qdrant_local_path: The path to the storage directory for the Qdrant client, if local mode is used. """ def __init__( self, - qdrant_url: str, + qdrant_url: Optional[str], qdrant_api_key: Optional[str], collection_name: str, fastembed_model_name: str, + qdrant_local_path: Optional[str] = None, ): - self._qdrant_url = qdrant_url.rstrip("/") + self._qdrant_url = qdrant_url.rstrip("/") if qdrant_url else None self._qdrant_api_key = qdrant_api_key self._collection_name = collection_name self._fastembed_model_name = fastembed_model_name # For the time being, FastEmbed models are the only supported ones. # A list of all available models can be found here: # https://qdrant.github.io/fastembed/examples/Supported_Models/ - self._client = AsyncQdrantClient(qdrant_url, api_key=qdrant_api_key) + self._client = AsyncQdrantClient(location=qdrant_url, api_key=qdrant_api_key, path=qdrant_local_path) self._client.set_model(fastembed_model_name) async def store_memory(self, information: str): diff --git a/src/mcp_server_qdrant/server.py b/src/mcp_server_qdrant/server.py index fb8478a..897b77a 100644 --- a/src/mcp_server_qdrant/server.py +++ b/src/mcp_server_qdrant/server.py @@ -12,10 +12,11 @@ from .qdrant import QdrantConnector def serve( - qdrant_url: str, + qdrant_url: Optional[str], qdrant_api_key: Optional[str], collection_name: str, fastembed_model_name: str, + qdrant_local_path: Optional[str] = None, ) -> Server: """ Instantiate the server and configure tools to store and find memories in Qdrant. @@ -23,11 +24,12 @@ def serve( :param qdrant_api_key: The API key to use for the Qdrant server. :param collection_name: The name of the collection to use. :param fastembed_model_name: The name of the FastEmbed model to use. + :param qdrant_local_path: The path to the storage directory for the Qdrant client, if local mode is used. """ server = Server("qdrant") qdrant = QdrantConnector( - qdrant_url, qdrant_api_key, collection_name, fastembed_model_name + qdrant_url, qdrant_api_key, collection_name, fastembed_model_name, qdrant_local_path ) @server.list_tools() @@ -112,7 +114,7 @@ def serve( @click.option( "--qdrant-url", envvar="QDRANT_URL", - required=True, + required=False, help="Qdrant URL", ) @click.option( @@ -134,12 +136,23 @@ def serve( help="FastEmbed model name", default="sentence-transformers/all-MiniLM-L6-v2", ) +@click.option( + "--qdrant-local-path", + envvar="QDRANT_LOCAL_PATH", + required=False, + help="Qdrant local path", +) def main( - qdrant_url: str, + qdrant_url: Optional[str], qdrant_api_key: str, collection_name: Optional[str], fastembed_model_name: str, + qdrant_local_path: Optional[str], ): + # XOR of url and local path, since we accept only one of them + if not (bool(qdrant_url) ^ bool(qdrant_local_path)): + raise ValueError("Exactly one of qdrant-url or qdrant-local-path must be provided") + async def _run(): async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): server = serve( @@ -147,6 +160,7 @@ def main( qdrant_api_key, collection_name, fastembed_model_name, + qdrant_local_path, ) await server.run( read_stream,