Merge pull request #5 from qdrant/feature/support-local-mode

Add support for local Qdrant mode
This commit is contained in:
Kacper Łukawski
2024-12-13 17:47:01 +01:00
committed by GitHub
3 changed files with 27 additions and 8 deletions

View File

@@ -82,10 +82,13 @@ by passing the `--fastembed-model-name` argument to the server.
The configuration of the server can be also done using environment variables: The configuration of the server can be also done using environment variables:
- `QDRANT_URL`: URL of the Qdrant server - `QDRANT_URL`: URL of the Qdrant server, e.g. `http://localhost:6333`
- `QDRANT_API_KEY`: API key for the Qdrant server - `QDRANT_API_KEY`: API key for the Qdrant server
- `COLLECTION_NAME`: Name of the collection to use - `COLLECTION_NAME`: Name of the collection to use
- `FASTEMBED_MODEL_NAME`: Name of the FastEmbed model to use - `FASTEMBED_MODEL_NAME`: Name of the FastEmbed model to use
- `QDRANT_LOCAL_PATH`: Path to the local Qdrant database
You cannot provide `QDRANT_URL` and `QDRANT_LOCAL_PATH` at the same time.
## License ## License

View File

@@ -9,23 +9,25 @@ class QdrantConnector:
:param qdrant_api_key: The API key to use for the Qdrant server. :param qdrant_api_key: The API key to use for the Qdrant server.
:param collection_name: The name of the collection to use. :param collection_name: The name of the collection to use.
:param fastembed_model_name: The name of the FastEmbed model 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__( def __init__(
self, self,
qdrant_url: str, qdrant_url: Optional[str],
qdrant_api_key: Optional[str], qdrant_api_key: Optional[str],
collection_name: str, collection_name: str,
fastembed_model_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._qdrant_api_key = qdrant_api_key
self._collection_name = collection_name self._collection_name = collection_name
self._fastembed_model_name = fastembed_model_name self._fastembed_model_name = fastembed_model_name
# For the time being, FastEmbed models are the only supported ones. # For the time being, FastEmbed models are the only supported ones.
# A list of all available models can be found here: # A list of all available models can be found here:
# https://qdrant.github.io/fastembed/examples/Supported_Models/ # 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) self._client.set_model(fastembed_model_name)
async def store_memory(self, information: str): async def store_memory(self, information: str):

View File

@@ -12,10 +12,11 @@ from .qdrant import QdrantConnector
def serve( def serve(
qdrant_url: str, qdrant_url: Optional[str],
qdrant_api_key: Optional[str], qdrant_api_key: Optional[str],
collection_name: str, collection_name: str,
fastembed_model_name: str, fastembed_model_name: str,
qdrant_local_path: Optional[str] = None,
) -> Server: ) -> Server:
""" """
Instantiate the server and configure tools to store and find memories in Qdrant. 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 qdrant_api_key: The API key to use for the Qdrant server.
:param collection_name: The name of the collection to use. :param collection_name: The name of the collection to use.
:param fastembed_model_name: The name of the FastEmbed model 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") server = Server("qdrant")
qdrant = QdrantConnector( 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() @server.list_tools()
@@ -112,7 +114,7 @@ def serve(
@click.option( @click.option(
"--qdrant-url", "--qdrant-url",
envvar="QDRANT_URL", envvar="QDRANT_URL",
required=True, required=False,
help="Qdrant URL", help="Qdrant URL",
) )
@click.option( @click.option(
@@ -134,12 +136,23 @@ def serve(
help="FastEmbed model name", help="FastEmbed model name",
default="sentence-transformers/all-MiniLM-L6-v2", default="sentence-transformers/all-MiniLM-L6-v2",
) )
@click.option(
"--qdrant-local-path",
envvar="QDRANT_LOCAL_PATH",
required=False,
help="Qdrant local path",
)
def main( def main(
qdrant_url: str, qdrant_url: Optional[str],
qdrant_api_key: str, qdrant_api_key: str,
collection_name: Optional[str], collection_name: Optional[str],
fastembed_model_name: 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 def _run():
async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
server = serve( server = serve(
@@ -147,6 +160,7 @@ def main(
qdrant_api_key, qdrant_api_key,
collection_name, collection_name,
fastembed_model_name, fastembed_model_name,
qdrant_local_path,
) )
await server.run( await server.run(
read_stream, read_stream,