Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for local Qdrant mode #5

Merged
merged 2 commits into from
Dec 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -73,10 +73,13 @@ by passing the `--fastembed-model-name` argument to the server.

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
- `COLLECTION_NAME`: Name of the collection 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

Expand Down
8 changes: 5 additions & 3 deletions src/mcp_server_qdrant/qdrant.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
22 changes: 18 additions & 4 deletions src/mcp_server_qdrant/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,22 +12,24 @@


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.
:param qdrant_url: The URL of 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 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()
Expand Down Expand Up @@ -112,7 +114,7 @@ async def handle_tool_call(
@click.option(
"--qdrant-url",
envvar="QDRANT_URL",
required=True,
required=False,
help="Qdrant URL",
)
@click.option(
Expand All @@ -134,19 +136,31 @@ async def handle_tool_call(
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(
qdrant_url,
qdrant_api_key,
collection_name,
fastembed_model_name,
qdrant_local_path,
)
await server.run(
read_stream,
Expand Down