diff --git a/pyproject.toml b/pyproject.toml index a58f948..2978a2e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,13 +5,17 @@ description = "Add your description here" readme = "README.md" requires-python = ">=3.13" dependencies = [ - "fastmcp>=2.12.3", "qdrant-client==1.13", "vault-settings>=0.1.0", ] +[project.optional-dependencies] +mcp = [ + "fastmcp>=2.12.3", +] + [project.scripts] -vector-search-mcp = "vector_search_mcp:run" +vector-search-mcp = "vector_search_mcp.mcp_server:run" [build-system] requires = ["uv_build"] diff --git a/src/vector_search_mcp/__init__.py b/src/vector_search_mcp/__init__.py index 541bbd3..e69de29 100644 --- a/src/vector_search_mcp/__init__.py +++ b/src/vector_search_mcp/__init__.py @@ -1,37 +0,0 @@ -"""Vector Search MCP (Model Context Protocol) Package. - -This package provides a Model Context Protocol server for vector similarity search -operations. It supports multiple vector database backends and provides a unified -interface for semantic search functionality. - -The package includes: -- Abstract engine interface for pluggable vector database backends -- Qdrant vector database implementation -- Pydantic models for search operations and conditions -- MCP server implementation with transport support - -Example: - Run the MCP server: - >>> from vector_search_mcp import run - >>> run("sse") # Start with Server-Sent Events transport - -""" - -from fastmcp.server.server import Transport - -from .main import mcp - - -def run(transport: Transport = "sse"): - """Run the vector search MCP server with the specified transport. - - Args: - transport: The transport protocol to use. Either "sse" for Server-Sent Events - or "stdio" for standard input/output communication. - - Example: - >>> run("sse") # Start with Server-Sent Events - >>> run("stdio") # Start with stdio transport - - """ - mcp.run(transport=transport) diff --git a/src/vector_search_mcp/client.py b/src/vector_search_mcp/client.py new file mode 100644 index 0000000..459f0b6 --- /dev/null +++ b/src/vector_search_mcp/client.py @@ -0,0 +1,28 @@ +from typing import final + +from .engine import Backend, get_engine +from .models import Chunk, Condition + + +@final +class Client: + def __init__(self, backend: Backend, collection: str): + self.engine = get_engine(backend) + self.collection = collection + + async def create_index(self, size: int) -> bool: + return await self.engine.create_index(self.collection, size) + + async def upload_chunk(self, chunk: Chunk) -> bool: + return await self.engine.upload_chunk(self.collection, chunk) + + async def semantic_search( + self, + embedding: list[float], + limit: int = 10, + conditions: list[Condition] | None = None, + threshold: float | None = None, + ): + return await self.engine.semantic_search( + embedding, self.collection, limit, conditions, threshold + ) diff --git a/src/vector_search_mcp/engine/base_engine.py b/src/vector_search_mcp/engine/base_engine.py index b9c4a07..70454b6 100644 --- a/src/vector_search_mcp/engine/base_engine.py +++ b/src/vector_search_mcp/engine/base_engine.py @@ -14,15 +14,16 @@ maintaining a consistent interface for the semantic search workflow. from abc import ABC, abstractmethod from typing import TypeVar -from ..models import Condition, SearchRow +from ..models import Chunk, Condition, SearchRow ResponseType = TypeVar("ResponseType") ConditionType = TypeVar("ConditionType") +ChunkType = TypeVar("ChunkType") __all__ = ["BaseEngine"] -class BaseEngine[ResponseType, ConditionType](ABC): +class BaseEngine[ResponseType, ConditionType, ChunkType](ABC): """Abstract base class for vector search engines. This class defines the interface that all vector search engine implementations @@ -34,6 +35,8 @@ class BaseEngine[ResponseType, ConditionType](ABC): For example, list[ScoredPoint] for Qdrant. ConditionType: The backend-specific filter/condition type. For example, models.Filter for Qdrant. + ChunkType: The backend-specific chunk type. + For example, models.Point for Qdrant. The class implements the Template Method pattern where semantic_search() orchestrates calls to the abstract methods that subclasses must implement. @@ -192,3 +195,16 @@ class BaseEngine[ResponseType, ConditionType](ABC): embedding, collection, limit, transformed_conditions, threshold ) return self.transform_response(response) + + @abstractmethod + async def create_index(self, name: str, size: int) -> bool: ... + + @abstractmethod + def transform_chunk(self, chunk: Chunk) -> ChunkType: ... + + @abstractmethod + async def run_upload_chunk(self, index_name: str, chunk: ChunkType) -> bool: ... + + async def upload_chunk(self, index_name: str, chunk: Chunk) -> bool: + transformed_chunk = self.transform_chunk(chunk) + return await self.run_upload_chunk(index_name, transformed_chunk) diff --git a/src/vector_search_mcp/engine/qdrant_engine.py b/src/vector_search_mcp/engine/qdrant_engine.py index 79ac4c5..912cd37 100644 --- a/src/vector_search_mcp/engine/qdrant_engine.py +++ b/src/vector_search_mcp/engine/qdrant_engine.py @@ -16,14 +16,16 @@ from typing import final, override from qdrant_client import AsyncQdrantClient, models from ..config import Settings -from ..models import Condition, Match, MatchAny, MatchExclude, SearchRow +from ..models import Chunk, Condition, Match, MatchAny, MatchExclude, SearchRow from .base_engine import BaseEngine __all__ = ["QdrantEngine"] @final -class QdrantEngine(BaseEngine[list[models.ScoredPoint], models.Filter]): +class QdrantEngine( + BaseEngine[list[models.ScoredPoint], models.Filter, models.PointStruct] +): """Qdrant vector database engine implementation. This class provides a concrete implementation of the BaseEngine interface @@ -195,3 +197,30 @@ class QdrantEngine(BaseEngine[list[models.ScoredPoint], models.Filter]): with_vectors=False, score_threshold=threshold, ) + + @override + async def create_index(self, name: str, size: int) -> bool: + return await self.client.create_collection( + collection_name=name, + vectors_config=models.VectorParams( + size=size, distance=models.Distance.COSINE + ), + ) + + @override + async def run_upload_chunk( + self, index_name: str, chunk: models.PointStruct + ) -> bool: + result = await self.client.upsert( + collection_name=index_name, + points=[chunk], + ) + return result.status == models.UpdateStatus.ACKNOWLEDGED + + @override + def transform_chunk(self, chunk: Chunk) -> models.PointStruct: + return models.PointStruct( + id=chunk.id, + vector=chunk.vector, + payload=chunk.payload.model_dump(), + ) diff --git a/src/vector_search_mcp/mcp_server/__init__.py b/src/vector_search_mcp/mcp_server/__init__.py new file mode 100644 index 0000000..cfb6950 --- /dev/null +++ b/src/vector_search_mcp/mcp_server/__init__.py @@ -0,0 +1,37 @@ +"""Vector Search MCP (Model Context Protocol) Package. + +This package provides a Model Context Protocol server for vector similarity search +operations. It supports multiple vector database backends and provides a unified +interface for semantic search functionality. + +The package includes: +- Abstract engine interface for pluggable vector database backends +- Qdrant vector database implementation +- Pydantic models for search operations and conditions +- MCP server implementation with transport support + +Example: + Run the MCP server: + >>> from vector_search_mcp import run + >>> run("sse") # Start with Server-Sent Events transport + +""" + +from fastmcp.server.server import Transport + +from .server import mcp + + +def run(transport: Transport = "sse"): + """Run the vector search MCP server with the specified transport. + + Args: + transport: The transport protocol to use. Either "sse" for Server-Sent Events + or "stdio" for standard input/output communication. + + Example: + >>> run("sse") # Start with Server-Sent Events + >>> run("stdio") # Start with stdio transport + + """ + mcp.run(transport=transport) diff --git a/src/vector_search_mcp/main.py b/src/vector_search_mcp/mcp_server/server.py similarity index 94% rename from src/vector_search_mcp/main.py rename to src/vector_search_mcp/mcp_server/server.py index e8fed21..356c8b3 100644 --- a/src/vector_search_mcp/main.py +++ b/src/vector_search_mcp/mcp_server/server.py @@ -18,7 +18,7 @@ Example: from fastmcp import FastMCP -from .engine import Backend, get_engine +from ..engine import Backend, get_engine mcp = FastMCP("Vector Search MCP") diff --git a/src/vector_search_mcp/models.py b/src/vector_search_mcp/models.py index 16e506f..26ce35c 100644 --- a/src/vector_search_mcp/models.py +++ b/src/vector_search_mcp/models.py @@ -119,3 +119,15 @@ class MatchExclude(Condition): key: str exclude: list[str] + + +class ChunkData(BaseModel): + page_content: str + filename: str + page: int + + +class Chunk(BaseModel): + id: str + vector: list[float] + payload: ChunkData diff --git a/tests/conftest.py b/tests/test_mcp/conftest.py similarity index 100% rename from tests/conftest.py rename to tests/test_mcp/conftest.py diff --git a/tests/test_mcp.py b/tests/test_mcp/test_mcp.py similarity index 100% rename from tests/test_mcp.py rename to tests/test_mcp/test_mcp.py diff --git a/uv.lock b/uv.lock index 97e2784..8d3126b 100644 --- a/uv.lock +++ b/uv.lock @@ -1620,11 +1620,15 @@ name = "vector-search-mcp" version = "0.1.0" source = { editable = "." } dependencies = [ - { name = "fastmcp" }, { name = "qdrant-client" }, { name = "vault-settings" }, ] +[package.optional-dependencies] +mcp = [ + { name = "fastmcp" }, +] + [package.dev-dependencies] dev = [ { name = "fastembed" }, @@ -1636,10 +1640,11 @@ dev = [ [package.metadata] requires-dist = [ - { name = "fastmcp", specifier = ">=2.12.3" }, + { name = "fastmcp", marker = "extra == 'mcp'", specifier = ">=2.12.3" }, { name = "qdrant-client", specifier = "==1.13" }, { name = "vault-settings", specifier = ">=0.1.0" }, ] +provides-extras = ["mcp"] [package.metadata.requires-dev] dev = [