diff --git a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/cli.py b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/cli.py index 7bb6f98..fa76f8f 100644 --- a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/cli.py +++ b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/cli.py @@ -112,12 +112,12 @@ async def handle_stdio(config: Config, log_level: str): # Initialize components search_engine = SearchEngine() - query_processor = QueryProcessor(config.openai) + query_processor = QueryProcessor(config.openai_query) mcp_handler = MCPHandler(search_engine, query_processor) # Initialize search engine try: - await search_engine.initialize(config.qdrant, config.openai) + await search_engine.initialize(config.qdrant, config.openai_embedding) if not disable_console_logging: logger.info("Search engine initialized successfully") except Exception as e: @@ -303,7 +303,6 @@ def cli(log_level: str = "INFO", config: Path | None = None) -> None: # Create and set the event loop loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) - # Set up signal handlers for sig in (signal.SIGTERM, signal.SIGINT): loop.add_signal_handler(sig, lambda: asyncio.create_task(shutdown(loop))) diff --git a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/config.py b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/config.py index a993536..89f1d26 100644 --- a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/config.py +++ b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/config.py @@ -4,6 +4,7 @@ from dotenv import load_dotenv from pydantic import BaseModel +from abc import ABC # Load environment variables from .env file load_dotenv() @@ -23,6 +24,7 @@ class QdrantConfig(BaseModel): url: str = "http://localhost:6333" api_key: str | None = None collection_name: str = "documents" + vector_size: int = 4096 def __init__(self, **data): """Initialize with environment variables if not provided.""" @@ -32,15 +34,43 @@ def __init__(self, **data): data["api_key"] = os.getenv("QDRANT_API_KEY") if "collection_name" not in data: data["collection_name"] = os.getenv("QDRANT_COLLECTION_NAME", "documents") + if "vector_size" not in data: + data["vector_size"] = int(os.getenv("QDRANT_VECTOR_SIZE", "4096")) super().__init__(**data) -class OpenAIConfig(BaseModel): - """OpenAI configuration settings.""" +class OpenAIConfig(BaseModel, ABC): + """Abstract base class for OpenAI configuration settings.""" - api_key: str - model: str = "text-embedding-3-small" - chat_model: str = "gpt-3.5-turbo" + api_key: str = "" + model: str = "" + base_url: str = "" # if None, the OpenAI package defaults to "https://api.openai.com/v1" + +class OpenAIEmbeddingConfig(OpenAIConfig): + """OpenAI configuration for embedding.""" + + def __init__(self, **data): + """Initialize with environment variables if not provided""" + if "api_key" not in data: + data["api_key"] = os.getenv("OPENAI_EMBEDDING_API_KEY", "") + if "model" not in data: + data["model"] = os.getenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small") + if "base_url" not in data: + data["base_url"] = os.getenv("OPENAI_EMBEDDING_BASE_URL", "http://localhost:1234/v1") + super().__init__(**data) + +class OpenAIQueryConfig(OpenAIConfig): + """OpenAI configuration for queries.""" + + def __init__(self, **data): + """Initialize with environment variables if not provided""" + if "api_key" not in data: + data["api_key"] = os.getenv("OPENAI_QUERY_API_KEY", "") + if "model" not in data: + data["model"] = os.getenv("OPENAI_QUERY_MODEL", "gpt-3.5-turbo") + if "base_url" not in data: + data["base_url"] = os.getenv("OPENAI_QUERY_BASE_URL", "http://localhost:1234/v1") + super().__init__(**data) class Config(BaseModel): @@ -48,7 +78,8 @@ class Config(BaseModel): server: ServerConfig qdrant: QdrantConfig - openai: OpenAIConfig + openai_query: OpenAIQueryConfig + openai_embedding: OpenAIEmbeddingConfig def __init__(self, **data): """Initialize configuration with environment variables.""" @@ -57,6 +88,8 @@ def __init__(self, **data): data["server"] = ServerConfig() if "qdrant" not in data: data["qdrant"] = QdrantConfig() - if "openai" not in data: - data["openai"] = {"api_key": os.getenv("OPENAI_API_KEY")} + if "openai_query" not in data: + data["openai_query"] = OpenAIQueryConfig() + if "openai_embedding" not in data: + data["openai_embedding"] = OpenAIEmbeddingConfig() super().__init__(**data) diff --git a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/patch/__init__.py b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/patch/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/patch/openai.py b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/patch/openai.py new file mode 100644 index 0000000..ab36d41 --- /dev/null +++ b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/patch/openai.py @@ -0,0 +1,18 @@ +from openai import AsyncOpenAI +from qdrant_loader_mcp_server.config import OpenAIConfig + +class PatchedAsyncOpenAI(AsyncOpenAI): + + model: str + + def __init__(self, *args, model: str = "", **kwargs): + super().__init__(*args, **kwargs) + self.model = model + + @classmethod + def from_config(cls, config: OpenAIConfig): + return cls( + api_key=config.api_key, + base_url=config.base_url, + model = config.model + ) \ No newline at end of file diff --git a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/engine.py b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/engine.py index a638535..91da0e6 100644 --- a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/engine.py +++ b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/engine.py @@ -1,11 +1,11 @@ """Search engine implementation for the MCP server.""" -from openai import AsyncOpenAI +from ..patch.openai import PatchedAsyncOpenAI from qdrant_client import QdrantClient from qdrant_client.http import models -from ..config import OpenAIConfig, QdrantConfig +from ..config import OpenAIEmbeddingConfig, QdrantConfig from ..utils.logging import LoggingConfig from .hybrid_search import HybridSearchEngine from .models import SearchResult @@ -20,18 +20,18 @@ def __init__(self): """Initialize the search engine.""" self.client: QdrantClient | None = None self.config: QdrantConfig | None = None - self.openai_client: AsyncOpenAI | None = None + self.openai_client: PatchedAsyncOpenAI | None = None self.hybrid_search: HybridSearchEngine | None = None self.logger = LoggingConfig.get_logger(__name__) async def initialize( - self, config: QdrantConfig, openai_config: OpenAIConfig + self, config: QdrantConfig, openai_config: OpenAIEmbeddingConfig ) -> None: """Initialize the search engine with configuration.""" self.config = config try: self.client = QdrantClient(url=config.url, api_key=config.api_key) - self.openai_client = AsyncOpenAI(api_key=openai_config.api_key) + self.openai_client = PatchedAsyncOpenAI.from_config(openai_config) # Ensure collection exists if self.client is None: @@ -42,7 +42,7 @@ async def initialize( self.client.create_collection( collection_name=config.collection_name, vectors_config=models.VectorParams( - size=1536, # Default size for OpenAI embeddings + size=config.vector_size, distance=models.Distance.COSINE, ), ) diff --git a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/hybrid_search.py b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/hybrid_search.py index ed8e5ce..c8dc682 100644 --- a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/hybrid_search.py +++ b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/hybrid_search.py @@ -1,11 +1,11 @@ """Hybrid search implementation combining vector and keyword search.""" import re -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import numpy as np -from openai import AsyncOpenAI +from ..patch.openai import PatchedAsyncOpenAI from qdrant_client import QdrantClient from qdrant_client.http import models from rank_bm25 import BM25Okapi @@ -74,17 +74,17 @@ class HybridSearchResult: paragraph_count: int | None = None # 🔥 NEW: Semantic analysis (NLP results) - entities: list[dict | str] = None - topics: list[dict | str] = None - key_phrases: list[dict | str] = None - pos_tags: list[dict] = None + entities: list[dict | str] = field(default_factory=list) + topics: list[dict | str] = field(default_factory=list) + key_phrases: list[dict | str] = field(default_factory=list) + pos_tags: list[dict] = field(default_factory=list) # 🔥 NEW: Navigation context previous_section: str | None = None next_section: str | None = None - sibling_sections: list[str] = None - subsections: list[str] = None - document_hierarchy: list[str] = None + sibling_sections: list[str] = field(default_factory=list) + subsections: list[str] = field(default_factory=list) + document_hierarchy: list[str] = field(default_factory=list) # 🔥 NEW: Chunking context chunk_index: int | None = None @@ -98,7 +98,7 @@ class HybridSearchResult: is_converted: bool = False # 🔥 NEW: Cross-references and enhanced context - cross_references: list[dict] = None + cross_references: list[dict] = field(default_factory=list) topic_analysis: dict | None = None content_type_context: str | None = None # Human-readable content description @@ -128,7 +128,7 @@ class HybridSearchEngine: def __init__( self, qdrant_client: QdrantClient, - openai_client: AsyncOpenAI, + openai_client: PatchedAsyncOpenAI, collection_name: str, vector_weight: float = 0.6, keyword_weight: float = 0.3, @@ -208,7 +208,7 @@ async def _get_embedding(self, text: str) -> list[float]: """Get embedding for text using OpenAI.""" try: response = await self.openai_client.embeddings.create( - model="text-embedding-3-small", + model=self.openai_client.model, input=text, ) return response.data[0].embedding diff --git a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/processor.py b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/processor.py index 6663113..38f4d2a 100644 --- a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/processor.py +++ b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/search/processor.py @@ -3,20 +3,18 @@ import re from typing import Any -from openai import AsyncOpenAI +from ..patch.openai import PatchedAsyncOpenAI -from ..config import OpenAIConfig +from ..config import OpenAIQueryConfig from ..utils.logging import LoggingConfig class QueryProcessor: """Query processor for handling search queries.""" - def __init__(self, openai_config: OpenAIConfig): + def __init__(self, openai_config: OpenAIQueryConfig): """Initialize the query processor.""" - self.openai_client: AsyncOpenAI | None = AsyncOpenAI( - api_key=openai_config.api_key - ) + self.openai_client: PatchedAsyncOpenAI | None = PatchedAsyncOpenAI.from_config(openai_config) self.logger = LoggingConfig.get_logger(__name__) async def process_query(self, query: str) -> dict[str, Any]: @@ -90,7 +88,7 @@ async def _infer_intent(self, query: str) -> tuple[str, bool]: raise RuntimeError("OpenAI client not initialized") response = await self.openai_client.chat.completions.create( - model="gpt-3.5-turbo", + model=self.openai_client.model, messages=[ { "role": "system", diff --git a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/utils/logging.py b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/utils/logging.py index c75e93d..6ac9fa5 100644 --- a/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/utils/logging.py +++ b/packages/qdrant-loader-mcp-server/src/qdrant_loader_mcp_server/utils/logging.py @@ -29,7 +29,7 @@ def filter(self, record): or record.name == "__main__" # Allow logs from main module or record.name == "asyncio" # Allow logs from asyncio or record.name == "main" # Allow logs when started as a script - or record.name == "qdrant_loader_mcp_server" # Allow logs from the package + or record.name.startswith("qdrant_loader_mcp_server") # Allow logs from the package ) diff --git a/packages/qdrant-loader-mcp-server/tests/integration/test_mcp_integration.py b/packages/qdrant-loader-mcp-server/tests/integration/test_mcp_integration.py index 8e95ad7..ba3039f 100644 --- a/packages/qdrant-loader-mcp-server/tests/integration/test_mcp_integration.py +++ b/packages/qdrant-loader-mcp-server/tests/integration/test_mcp_integration.py @@ -48,9 +48,9 @@ async def integration_handler(): search_engine = SearchEngine() # Create OpenAI config for query processor - from qdrant_loader_mcp_server.config import OpenAIConfig + from qdrant_loader_mcp_server.config import OpenAIQueryConfig - openai_config = OpenAIConfig(api_key="test_key") + openai_config = OpenAIQueryConfig(api_key="test_key") query_processor = QueryProcessor(openai_config) # Patch external dependencies @@ -60,20 +60,23 @@ async def integration_handler(): return_value=mock_qdrant_client, ), patch( - "qdrant_loader_mcp_server.search.engine.AsyncOpenAI", + "qdrant_loader_mcp_server.search.engine.PatchedAsyncOpenAI", return_value=mock_openai_client, - ), + ) as PatchedOpenAIMock, patch( - "qdrant_loader_mcp_server.search.processor.AsyncOpenAI", + "qdrant_loader_mcp_server.search.processor.PatchedAsyncOpenAI", return_value=mock_openai_client, ), ): + # Override named constructor + PatchedOpenAIMock.from_config.return_value = mock_openai_client + # Initialize components - from qdrant_loader_mcp_server.config import OpenAIConfig, QdrantConfig + from qdrant_loader_mcp_server.config import OpenAIEmbeddingConfig, QdrantConfig qdrant_config = QdrantConfig(api_key="test_key") - openai_config = OpenAIConfig(api_key="test_key") + openai_config = OpenAIEmbeddingConfig(api_key="test_key") await search_engine.initialize(qdrant_config, openai_config) diff --git a/packages/qdrant-loader-mcp-server/tests/unit/search/test_hybrid_search.py b/packages/qdrant-loader-mcp-server/tests/unit/search/test_hybrid_search.py index 6b20cb7..922afe6 100644 --- a/packages/qdrant-loader-mcp-server/tests/unit/search/test_hybrid_search.py +++ b/packages/qdrant-loader-mcp-server/tests/unit/search/test_hybrid_search.py @@ -4,7 +4,7 @@ import numpy as np import pytest -from openai import AsyncOpenAI +from qdrant_loader_mcp_server.patch.openai import PatchedAsyncOpenAI from qdrant_loader_mcp_server.search.hybrid_search import ( HybridSearchEngine, ) @@ -81,7 +81,8 @@ def mock_qdrant_client(): @pytest.fixture def mock_openai_client(): """Create a mock OpenAI client.""" - client = AsyncMock(spec=AsyncOpenAI) + client = AsyncMock(spec=PatchedAsyncOpenAI) + client.model = "test-model" # Mock embeddings response embedding_response = MagicMock() @@ -565,5 +566,5 @@ async def test_get_embedding_success(hybrid_search, mock_openai_client): # Verify OpenAI API was called correctly mock_openai_client.embeddings.create.assert_called_once_with( - model="text-embedding-3-small", input="test text" + model=mock_openai_client.model, input="test text" ) diff --git a/packages/qdrant-loader-mcp-server/tests/unit/search/test_project_search.py b/packages/qdrant-loader-mcp-server/tests/unit/search/test_project_search.py index b60d0e8..be8149e 100644 --- a/packages/qdrant-loader-mcp-server/tests/unit/search/test_project_search.py +++ b/packages/qdrant-loader-mcp-server/tests/unit/search/test_project_search.py @@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from openai import AsyncOpenAI +from qdrant_loader_mcp_server.patch.openai import PatchedAsyncOpenAI from qdrant_client import QdrantClient from qdrant_loader_mcp_server.search.engine import SearchEngine from qdrant_loader_mcp_server.search.hybrid_search import ( @@ -56,7 +56,7 @@ def mock_qdrant_client(): @pytest.fixture def mock_openai_client(): """Create a mock OpenAI client.""" - client = AsyncMock(spec=AsyncOpenAI) + client = AsyncMock(spec=PatchedAsyncOpenAI) # Mock embedding response mock_response = MagicMock() diff --git a/packages/qdrant-loader-mcp-server/tests/unit/search/test_search_engine.py b/packages/qdrant-loader-mcp-server/tests/unit/search/test_search_engine.py index 18adcbb..24af066 100644 --- a/packages/qdrant-loader-mcp-server/tests/unit/search/test_search_engine.py +++ b/packages/qdrant-loader-mcp-server/tests/unit/search/test_search_engine.py @@ -41,7 +41,7 @@ async def test_search_engine_initialization( return_value=mock_qdrant_client, ), patch( - "qdrant_loader_mcp_server.search.engine.AsyncOpenAI", + "qdrant_loader_mcp_server.search.engine.PatchedAsyncOpenAI", return_value=mock_openai_client, ), patch( @@ -94,7 +94,7 @@ async def test_search_engine_search( return_value=mock_qdrant_client, ), patch( - "qdrant_loader_mcp_server.search.engine.AsyncOpenAI", + "qdrant_loader_mcp_server.search.engine.PatchedAsyncOpenAI", return_value=mock_openai_client, ), patch( @@ -132,7 +132,7 @@ async def test_search_engine_cleanup( return_value=mock_qdrant_client, ), patch( - "qdrant_loader_mcp_server.search.engine.AsyncOpenAI", + "qdrant_loader_mcp_server.search.engine.PatchedAsyncOpenAI", return_value=mock_openai_client, ), patch("qdrant_loader_mcp_server.search.engine.HybridSearchEngine"), @@ -161,7 +161,7 @@ async def test_search_engine_collection_creation( return_value=mock_qdrant_client, ), patch( - "qdrant_loader_mcp_server.search.engine.AsyncOpenAI", + "qdrant_loader_mcp_server.search.engine.PatchedAsyncOpenAI", return_value=mock_openai_client, ), patch("qdrant_loader_mcp_server.search.engine.HybridSearchEngine"), @@ -190,7 +190,7 @@ async def test_search_engine_collection_exists( return_value=mock_qdrant_client, ), patch( - "qdrant_loader_mcp_server.search.engine.AsyncOpenAI", + "qdrant_loader_mcp_server.search.engine.PatchedAsyncOpenAI", return_value=mock_openai_client, ), patch("qdrant_loader_mcp_server.search.engine.HybridSearchEngine"), diff --git a/packages/qdrant-loader-mcp-server/tests/unit/test_config.py b/packages/qdrant-loader-mcp-server/tests/unit/test_config.py index c5b0371..bedff59 100644 --- a/packages/qdrant-loader-mcp-server/tests/unit/test_config.py +++ b/packages/qdrant-loader-mcp-server/tests/unit/test_config.py @@ -3,7 +3,7 @@ import os from unittest.mock import patch -from qdrant_loader_mcp_server.config import Config, OpenAIConfig, QdrantConfig +from qdrant_loader_mcp_server.config import Config, OpenAIQueryConfig, OpenAIEmbeddingConfig, QdrantConfig def test_config_creation(): @@ -11,7 +11,8 @@ def test_config_creation(): config = Config() assert config is not None assert hasattr(config, "qdrant") - assert hasattr(config, "openai") + assert hasattr(config, "openai_query") + assert hasattr(config, "openai_embedding") def test_qdrant_config_defaults(monkeypatch): @@ -20,11 +21,13 @@ def test_qdrant_config_defaults(monkeypatch): monkeypatch.delenv("QDRANT_URL", raising=False) monkeypatch.delenv("QDRANT_API_KEY", raising=False) monkeypatch.delenv("QDRANT_COLLECTION_NAME", raising=False) + monkeypatch.delenv("QDRANT_VECTOR_SIZE", raising=False) config = QdrantConfig() assert config.url == "http://localhost:6333" assert config.collection_name == "documents" assert config.api_key is None + assert config.vector_size == 4096 def test_qdrant_config_from_env(monkeypatch): @@ -33,36 +36,44 @@ def test_qdrant_config_from_env(monkeypatch): monkeypatch.setenv("QDRANT_URL", "http://test:6333") monkeypatch.setenv("QDRANT_API_KEY", "test_key") monkeypatch.setenv("QDRANT_COLLECTION_NAME", "test_collection") + monkeypatch.setenv("QDRANT_VECTOR_SIZE", "12345") config = QdrantConfig() assert config.url == "http://test:6333" assert config.api_key == "test_key" assert config.collection_name == "test_collection" + assert config.vector_size == 12345 -def test_openai_config_defaults(): +def test_openai_query_config_defaults(monkeypatch): """Test OpenAI configuration defaults.""" - config = OpenAIConfig(api_key="test_key") - assert config.model == "text-embedding-3-small" - assert config.api_key == "test_key" + # Clear all related env variables + monkeypatch.delenv("OPENAI_QUERY_API_KEY", raising=False) + monkeypatch.delenv("OPENAI_QUERY_MODEL", raising=False) + monkeypatch.delenv("OPENAI_QUERY_BASE_URL", raising=False) + + config = OpenAIQueryConfig() + assert config.model == "gpt-3.5-turbo" + assert config.api_key == "" + assert config.base_url == "http://localhost:1234/v1" -def test_openai_config_from_env(): +def test_openai_query_config_from_env(): """Test OpenAI configuration from environment variables.""" with patch.dict( os.environ, - {"OPENAI_API_KEY": "test_key", "OPENAI_MODEL": "text-embedding-ada-002"}, + {"OPENAI_QUERY_API_KEY": "test_key", "OPENAI_QUERY_MODEL": "text-embedding-ada-002", "OPENAI_QUERY_BASE_URL": "http://localhost:111234/v1"}, ): - config = OpenAIConfig(api_key="test_key") + config = OpenAIQueryConfig() assert config.api_key == "test_key" - assert ( - config.model == "text-embedding-3-small" - ) # Model is not read from env in this config + assert config.model == "text-embedding-ada-002" + assert config.base_url == "http://localhost:111234/v1" +# TODO: Add unit tests for OpenAI embedding config -def test_openai_config_with_api_key(): +def test_openai__query_config_with_api_key(): """Test OpenAI configuration with explicit API key.""" - config = OpenAIConfig(api_key="explicit_key") + config = OpenAIQueryConfig(api_key="explicit_key") assert config.api_key == "explicit_key" @@ -74,9 +85,6 @@ def test_config_validation(): ) assert qdrant_config.url == "http://localhost:6333" - openai_config = OpenAIConfig(api_key="test_key") - assert openai_config.api_key == "test_key" - def test_config_integration(monkeypatch): """Test full configuration integration.""" @@ -84,9 +92,9 @@ def test_config_integration(monkeypatch): monkeypatch.setenv("QDRANT_URL", "http://localhost:6333") monkeypatch.setenv("QDRANT_COLLECTION_NAME", "test_collection") monkeypatch.delenv("QDRANT_API_KEY", raising=False) - monkeypatch.setenv("OPENAI_API_KEY", "test_key") + monkeypatch.setenv("OPENAI_QUERY_API_KEY", "test_keyaaa") config = Config() assert config.qdrant.url == "http://localhost:6333" assert config.qdrant.collection_name == "test_collection" - assert config.openai.api_key == "test_key" + assert config.openai_query.api_key == "test_keyaaa"