Add another parameter to allow using local Qdrant mode
This commit is contained in:
@@ -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):
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user