From 2565906f887ad7b609b64fe0def35611f53d694e Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 9 Aug 2026 12:30:44 +0100 Subject: [PATCH 1/4] chore(deps): bump 47 backend dependencies Bumps the Python backend dependency set, holding back the ones that are resolver-blocked or that regress behaviour this repo depends on. Notable upgrades: cryptography 46.0.7 -> 50.0.0 (requires msal 1.37.0, which relaxes its cap) protobuf 6.33.6 -> 7.35.1 (floats opentelemetry-* to 1.44.0) openai 2.32.0 -> 2.53.0, anthropic 0.88.0 -> 0.121.0 google-genai 1.73.1 -> 2.17.0 (the 2.0 break is scoped to the Interactions API, which this repo does not use) fastmcp 3.2.4 -> 3.4.6, gunicorn 25.3.0 -> 26.0.0 starlette 1.0.0 -> 1.6.0, uvicorn 0.42.0 -> 0.52.1 sentence-transformers 5.3.0 -> 5.7.0, numpy 2.4.4 -> 2.5.1 faiss-cpu 1.13.2 -> 1.15.0, pillow -> 12.3.0 (pinned; security release) pypdf 6.9.2 -> 6.15.0, lxml 6.0.2 -> 6.1.1 (CVE-2026-41066) Code changes needed by the bumps: - openai >= 2.53 rejects a falsy api_key at construction. Keyless OpenAI-compatible backends (Ollama, llama.cpp, vLLM) legitimately have none, and pydantic-settings yields "" for a bare `API_KEY=` in .env, so both call sites now fall back to a placeholder. - Flask >= 3.1.2 tears a stream_with_context request down twice, and a ContextVar token may only be reset once. The log-context teardown hook is now idempotent. - sentence-transformers renamed get_sentence_embedding_dimension to get_embedding_dimension in 5.4; use the new name with a fallback. Held back deliberately: torch 2.11.0 - 2.13.0 drags torchvision 0.26 -> 0.28 and the whole docling stack; torch is only probed for cuda.empty_cache() tokenizers 0.22.2 - transformers pins <=0.23.0 and no stable 0.23.0 exists transformers - capped <5.9.0: 5.9+ breaks docling's PDF layout model on Apple Silicon (MPS float64), matching docling-core's own darwin pin docling 2.84.0 - 2.118.1 needs rapidocr >=3.9.1 and conflicts with transformers on macOS; wants its own PR with a golden-corpus diff redis 7.4.0 - 8.x defaults socket_timeout to 5s, silently capping the blocking reads in the device broker and SSE tail websockets 16.0 - google-genai caps <17.0 marshmallow - dataclasses-json hard-caps <4; spec tightened to match langchain block - langchain-community 0.4.2 deletes the Qdrant vectorstore this repo imports; langchain 1.3.x needs websockets <16 --- application/app.py | 4 + application/llm/openai.py | 10 ++- application/requirements.txt | 94 +++++++++++---------- application/sandbox/daytona.py | 2 +- application/stt/openai_stt.py | 8 +- application/vectorstore/embeddings_local.py | 8 +- tests/llm/test_openai.py | 31 +++++++ tests/test_app.py | 23 +++++ tests/vectorstore/test_embeddings_local.py | 18 +++- 9 files changed, 148 insertions(+), 50 deletions(-) diff --git a/application/app.py b/application/app.py index ad00d0b6..bf27bf52 100644 --- a/application/app.py +++ b/application/app.py @@ -183,6 +183,10 @@ def _reset_log_context(_exc): # leak into the stream's view of the context. token = getattr(request, _LOG_CTX_TOKEN_ATTR, None) if token is not None: + # Flask >= 3.1.2 tears a stream_with_context request down twice: once + # when the view returns, once when the generator is finalized. Clear + # the token first — resetting one twice raises RuntimeError. + setattr(request, _LOG_CTX_TOKEN_ATTR, None) log_context.reset(token) diff --git a/application/llm/openai.py b/application/llm/openai.py index 3bd65446..8c2043e3 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -11,6 +11,9 @@ from application.core.settings import settings from application.llm.base import BaseLLM from application.storage.storage_creator import StorageCreator +# Placeholder sent to OpenAI-compatible backends that require no credentials. +NO_API_KEY = "sk-no-key" + def _truncate_base64_for_logging(messages): """ @@ -124,7 +127,12 @@ class OpenAILLM(BaseLLM): ): super().__init__(*args, **kwargs) - self.api_key = api_key or settings.OPENAI_API_KEY or settings.API_KEY + # openai>=2.53 rejects a falsy api_key at construction. Keyless + # OpenAI-compatible backends (Ollama, llama.cpp, vLLM) legitimately have + # none, and pydantic-settings yields "" for a bare `API_KEY=` in .env. + self.api_key = ( + api_key or settings.OPENAI_API_KEY or settings.API_KEY or NO_API_KEY + ) self.user_api_key = user_api_key # Priority: 1) Parameter base_url, 2) Settings OPENAI_BASE_URL, 3) Default diff --git a/application/requirements.txt b/application/requirements.txt index 6773338e..bb052ff8 100644 --- a/application/requirements.txt +++ b/application/requirements.txt @@ -1,16 +1,16 @@ a2wsgi==1.10.10 alembic>=1.13,<2 -anthropic==0.88.0 +anthropic==0.121.0 asgiref>=3.11.1 -boto3==1.42.83 -beautifulsoup4==4.14.3 +boto3==1.43.67 +beautifulsoup4==4.15.0 cel-python==0.5.0 celery==5.6.3 -celery-redbeat==2.3.3 -croniter==6.2.2 -cryptography==46.0.7 +celery-redbeat==2.4.2 +croniter==6.2.4 +cryptography==50.0.0 dataclasses-json==0.6.7 -daytona==0.190.1 +daytona==0.203.0 defusedxml==0.7.1 docling==2.84.0 rapidocr>=1.4.0 @@ -18,19 +18,19 @@ onnxruntime>=1.19.0 docx2txt==0.9 ddgs>=8.0.0 fast-ebook -elevenlabs==2.43.0 -Flask==3.1.1 -faiss-cpu==1.13.2 -fastmcp==3.2.4 +elevenlabs==2.62.0 +Flask==3.1.3 +faiss-cpu==1.15.0 +fastmcp==3.4.6 flask-restx==1.3.2 -google-genai==1.73.1 -google-api-python-client==2.193.0 -google-auth-httplib2==0.3.1 -google-auth-oauthlib==1.3.1 +google-genai==2.17.0 +google-api-python-client==2.198.0 +google-auth-httplib2==0.4.1 +google-auth-oauthlib==1.4.0 gTTS==2.5.4 -gunicorn==25.3.0 +gunicorn==26.0.0 jinja2==3.1.6 -jiter==0.13.0 +jiter==0.16.0 jmespath==1.1.0 joblib==1.5.3 jsonpatch==1.33 @@ -38,21 +38,22 @@ jsonpointer==3.1.1 kombu==5.6.2 langchain==1.2.3 langchain-community==0.4.1 -langchain-core==1.2.29 +langchain-core==1.2.31 langchain-openai==1.1.12 -langchain-text-splitters==1.1.1 +langchain-text-splitters==1.1.2 langsmith==0.7.31 lazy-object-proxy==1.12.0 -lxml==6.0.2 +lxml==6.1.1 markupsafe==3.0.3 -marshmallow>=3.18.0,<5.0.0 +# dataclasses-json (via langchain-community) hard-caps marshmallow <4. +marshmallow>=3.18.0,<4.0.0 mpmath==1.3.0 multidict==6.7.1 -msal==1.35.1 +msal==1.37.0 mypy-extensions==1.1.0 networkx==3.6.1 -numpy==2.4.4 -openai==2.32.0 +numpy==2.5.1 +openai==2.53.0 openapi3-parser==1.1.22 opentelemetry-distro>=0.50b0,<1 opentelemetry-exporter-otlp>=1.29.0,<2 @@ -64,23 +65,23 @@ opentelemetry-instrumentation-redis>=0.50b0,<1 opentelemetry-instrumentation-requests>=0.50b0,<1 opentelemetry-instrumentation-sqlalchemy>=0.50b0,<1 opentelemetry-instrumentation-starlette>=0.50b0,<1 -orjson==3.11.7 -packaging==26.0 -pandas==3.0.2 +orjson==3.11.9 +packaging==26.3 +pandas==3.0.5 openpyxl==3.1.5 -pathable==0.5.0 +pathable==0.6.0 pdf2image>=1.17.0 -pgvector>=0.2,<1 -pillow +pgvector>=0.5,<1 +pillow==12.3.0 portalocker>=2.7.0,<4.0.0 -prompt-toolkit==3.0.52 -protobuf==6.33.6 +prompt-toolkit==3.0.53 +protobuf==7.35.1 psycopg[binary,pool]>=3.1,<4 py==1.11.0 pydantic pydantic-core pydantic-settings -pypdf==6.9.2 +pypdf==6.15.0 python-dateutil==2.9.0.post0 python-dotenv python-jose==3.5.0 @@ -88,28 +89,31 @@ python-pptx==1.0.2 PyYAML redis==7.4.0 referencing>=0.28.0,<0.38.0 -regex==2026.4.4 -requests==2.33.1 +regex==2026.7.19 +requests==2.34.2 retry==0.9.2 -sentence-transformers==5.3.0 +sentence-transformers==5.7.0 sqlalchemy>=2.0,<3 starlette>=1.0,<2 -tiktoken==0.12.0 +tiktoken==0.13.0 tokenizers==0.22.2 torch==2.11.0 tqdm==4.67.3 -transformers==5.4.0 -typing-extensions==4.15.0 +# Capped <5.9.0 to match docling-core's own darwin pin: 5.9+ breaks docling's +# PDF layout model on Apple Silicon ("Cannot convert a MPS Tensor to float64"). +transformers==5.8.1 +typing-extensions==4.16.0 typing-inspect==0.9.0 -tzdata==2026.1 -urllib3==2.6.3 +tzdata==2026.3 +urllib3==2.7.0 uvicorn[standard]>=0.30,<1 uvicorn-worker>=0.4,<1 vine==5.1.0 -wcwidth==0.6.0 +wcwidth==0.8.2 werkzeug>=3.1.0 -yarl==1.23.0 -markdownify==1.2.2 -tldextract==5.3.1 +yarl==1.24.5 +markdownify==1.2.3 +tldextract==5.3.2 +# google-genai caps websockets <17.0; 17.x will not resolve. websockets==16.0 websocket-client==1.9.0 diff --git a/application/sandbox/daytona.py b/application/sandbox/daytona.py index 3ce1605d..37fd7682 100644 --- a/application/sandbox/daytona.py +++ b/application/sandbox/daytona.py @@ -326,7 +326,7 @@ class DaytonaSandbox(CodeSandbox): The tools pass a token DIRECTORY (e.g. ``artifacts/{token}``), so the delete is recursive — ``fs.delete_file`` with ``recursive=True`` removes a folder and its - contents (daytona==0.190.1). + contents (daytona==0.203.0). """ try: handle = self._get_handle(session_id) diff --git a/application/stt/openai_stt.py b/application/stt/openai_stt.py index 96a46a58..44eb3c65 100644 --- a/application/stt/openai_stt.py +++ b/application/stt/openai_stt.py @@ -6,6 +6,9 @@ from openai import OpenAI from application.core.settings import settings from application.stt.base import BaseSTT +# Placeholder sent to OpenAI-compatible backends that require no credentials. +NO_API_KEY = "sk-no-key" + class OpenAISTT(BaseSTT): def __init__( @@ -14,7 +17,10 @@ class OpenAISTT(BaseSTT): base_url: Optional[str] = None, model: Optional[str] = None, ): - self.api_key = api_key or settings.OPENAI_API_KEY or settings.API_KEY + # See application/llm/openai.py: openai>=2.53 rejects a falsy api_key. + self.api_key = ( + api_key or settings.OPENAI_API_KEY or settings.API_KEY or NO_API_KEY + ) self.base_url = base_url or settings.OPENAI_BASE_URL or "https://api.openai.com/v1" self.model = model or settings.OPENAI_STT_MODEL self.client = OpenAI(api_key=self.api_key, base_url=self.base_url) diff --git a/application/vectorstore/embeddings_local.py b/application/vectorstore/embeddings_local.py index fb6e5f4f..a811a3bb 100644 --- a/application/vectorstore/embeddings_local.py +++ b/application/vectorstore/embeddings_local.py @@ -24,7 +24,13 @@ class EmbeddingsWrapper: raise ValueError( f"SentenceTransformer model failed to load properly for: {model_name}" ) - self.dimension = self.model.get_sentence_embedding_dimension() + # Renamed in sentence-transformers 5.4; keep the old name as a fallback. + get_dimension = getattr( + self.model, + "get_embedding_dimension", + getattr(self.model, "get_sentence_embedding_dimension", None), + ) + self.dimension = get_dimension() logging.info(f"Successfully loaded model with dimension: {self.dimension}") except Exception as e: logging.error( diff --git a/tests/llm/test_openai.py b/tests/llm/test_openai.py index a58031ef..65707845 100644 --- a/tests/llm/test_openai.py +++ b/tests/llm/test_openai.py @@ -2049,3 +2049,34 @@ class TestInlineFilePartResolution: # Upload happened because the read errored → treated as miss. assert len(llm.client.files.calls) == 1 assert out[0]["content"][1] == {"type": "file", "file": {"file_id": "file-1"}} + + +@pytest.mark.unit +class TestKeylessConstruction: + """openai>=2.53 raises on a falsy api_key at construction. + + Keyless OpenAI-compatible backends (Ollama, llama.cpp, vLLM) have no + credential, and pydantic-settings yields "" — not None — for a bare + `API_KEY=` line in .env, so the empty string must not reach the SDK. + """ + + @pytest.mark.parametrize("blank", ["", None]) + def test_llm_accepts_blank_key(self, blank): + with patch("application.llm.openai.settings") as mock_settings: + mock_settings.OPENAI_API_KEY = blank + mock_settings.API_KEY = blank + mock_settings.OPENAI_BASE_URL = "http://localhost:11434/v1" + llm = OpenAILLM(api_key=blank) + assert llm.api_key == "sk-no-key" + + @pytest.mark.parametrize("blank", ["", None]) + def test_stt_accepts_blank_key(self, blank): + from application.stt.openai_stt import OpenAISTT + + with patch("application.stt.openai_stt.settings") as mock_settings: + mock_settings.OPENAI_API_KEY = blank + mock_settings.API_KEY = blank + mock_settings.OPENAI_BASE_URL = "http://localhost:11434/v1" + mock_settings.OPENAI_STT_MODEL = "whisper-1" + stt = OpenAISTT(api_key=blank) + assert stt.api_key == "sk-no-key" diff --git a/tests/test_app.py b/tests/test_app.py index a423ca42..ac2bc5ba 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -19,3 +19,26 @@ def test_app_config(): assert app.config["UPLOAD_FOLDER"] == "inputs" assert app.config["CELERY_BROKER_URL"] == settings.CELERY_BROKER_URL assert app.config["CELERY_RESULT_BACKEND"] == settings.CELERY_RESULT_BACKEND + + +@pytest.mark.unit +class TestLogContextTeardown: + """Flask >= 3.1.2 tears a stream_with_context request down twice: once when + the view returns, once when the generator is finalized. A ContextVar token + may only be reset once, so the teardown hook has to be idempotent.""" + + def test_reset_is_idempotent(self): + from flask import request + + from application.app import _LOG_CTX_TOKEN_ATTR, _reset_log_context, app + from application.core import log_context + + with app.test_request_context("/"): + token = log_context.bind(activity_id="abc", endpoint="stream") + setattr(request, _LOG_CTX_TOKEN_ATTR, token) + + _reset_log_context(None) + _reset_log_context(None) # second teardown must not raise + + assert getattr(request, _LOG_CTX_TOKEN_ATTR) is None + assert "activity_id" not in log_context.snapshot() diff --git a/tests/vectorstore/test_embeddings_local.py b/tests/vectorstore/test_embeddings_local.py index 7b85bd35..b96fbed3 100644 --- a/tests/vectorstore/test_embeddings_local.py +++ b/tests/vectorstore/test_embeddings_local.py @@ -9,7 +9,7 @@ class TestEmbeddingsWrapper: def test_init_success(self, mock_st_cls): mock_model = MagicMock() mock_model._first_module.return_value = MagicMock() - mock_model.get_sentence_embedding_dimension.return_value = 768 + mock_model.get_embedding_dimension.return_value = 768 mock_st_cls.return_value = mock_model from application.vectorstore.embeddings_local import EmbeddingsWrapper @@ -19,6 +19,22 @@ class TestEmbeddingsWrapper: mock_st_cls.assert_called_once() assert wrapper.dimension == 768 + @patch("application.vectorstore.embeddings_local.SentenceTransformer") + def test_init_falls_back_to_legacy_dimension_api(self, mock_st_cls): + """sentence-transformers < 5.4 only exposes get_sentence_embedding_dimension.""" + mock_model = MagicMock() + mock_model._first_module.return_value = MagicMock() + del mock_model.get_embedding_dimension + mock_model.get_sentence_embedding_dimension.return_value = 768 + mock_st_cls.return_value = mock_model + + from application.vectorstore.embeddings_local import EmbeddingsWrapper + + wrapper = EmbeddingsWrapper("test-model") + + assert wrapper.dimension == 768 + mock_model.get_sentence_embedding_dimension.assert_called_once() + @patch("application.vectorstore.embeddings_local.SentenceTransformer") def test_init_failure(self, mock_st_cls): mock_st_cls.side_effect = Exception("model not found") From c5b372df1eaf153ea8894adb442c4c3984fe4c3c Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 10 Aug 2026 10:37:43 +0100 Subject: [PATCH 2/4] fix: pin rapidocr and onnxruntime --- application/parser/remote/telegram.py | 11 ----------- application/requirements.txt | 4 ++-- 2 files changed, 2 insertions(+), 13 deletions(-) delete mode 100644 application/parser/remote/telegram.py diff --git a/application/parser/remote/telegram.py b/application/parser/remote/telegram.py deleted file mode 100644 index 0e691be4..00000000 --- a/application/parser/remote/telegram.py +++ /dev/null @@ -1,11 +0,0 @@ -from langchain.document_loader import TelegramChatApiLoader -from application.parser.remote.base import BaseRemote - -class TelegramChatApiRemote(BaseRemote): - def _init_parser(self, *args, **load_kwargs): - self.loader = TelegramChatApiLoader(**load_kwargs) - return {} - - def parse_file(self, *args, **load_kwargs): - - return \ No newline at end of file diff --git a/application/requirements.txt b/application/requirements.txt index bb052ff8..535930ba 100644 --- a/application/requirements.txt +++ b/application/requirements.txt @@ -13,8 +13,8 @@ dataclasses-json==0.6.7 daytona==0.203.0 defusedxml==0.7.1 docling==2.84.0 -rapidocr>=1.4.0 -onnxruntime>=1.19.0 +rapidocr==3.9.2 +onnxruntime==1.28.0 docx2txt==0.9 ddgs>=8.0.0 fast-ebook From a166796f6cdb1fe7fefe019a0a8ab227642532b5 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 10 Aug 2026 12:10:54 +0100 Subject: [PATCH 3/4] feat: prune some deps --- application/api/user/sources/routes.py | 9 +- application/parser/chunking_strategies.py | 2 +- application/parser/file/base.py | 8 +- application/parser/file/html_parser.py | 26 +- application/parser/remote/base.py | 8 +- application/parser/remote/reddit_loader.py | 92 ++- application/parser/schema/base.py | 12 +- application/requirements.txt | 11 +- application/vectorstore/base.py | 3 +- application/vectorstore/embeddings_openai.py | 89 +++ application/vectorstore/faiss.py | 338 +++++--- application/vectorstore/faiss_docstore.py | 230 ++++++ application/vectorstore/milvus.py | 215 +++++- application/vectorstore/qdrant.py | 248 +++--- application/worker.py | 6 +- tests/parser/file/test_html_parser.py | 61 +- tests/parser/remote/test_reddit_loader.py | 148 ++-- tests/parser/remote/test_web_loader.py | 6 +- tests/parser/test_chunking_strategies.py | 2 +- tests/parser/test_schema.py | 20 +- tests/test_coverage_gaps.py | 174 ++--- tests/test_remaining_coverage.py | 8 +- tests/vectorstore/test_faiss.py | 774 +++++-------------- tests/vectorstore/test_faiss_docstore.py | 191 +++++ tests/vectorstore/test_milvus.py | 228 ++++-- tests/vectorstore/test_qdrant.py | 371 ++++----- 26 files changed, 1937 insertions(+), 1343 deletions(-) create mode 100644 application/vectorstore/embeddings_openai.py create mode 100644 application/vectorstore/faiss_docstore.py create mode 100644 tests/vectorstore/test_faiss_docstore.py diff --git a/application/api/user/sources/routes.py b/application/api/user/sources/routes.py index 20893382..adf27e16 100644 --- a/application/api/user/sources/routes.py +++ b/application/api/user/sources/routes.py @@ -257,10 +257,11 @@ class DeleteOldIndexes(Resource): try: if settings.VECTOR_STORE == "faiss": index_path = f"indexes/{resolved_id}" - if storage.file_exists(f"{index_path}/index.faiss"): - storage.delete_file(f"{index_path}/index.faiss") - if storage.file_exists(f"{index_path}/index.pkl"): - storage.delete_file(f"{index_path}/index.pkl") + # index.pkl is the legacy sidecar; index.json the current one. + # Older sources have only the former, so clear whichever exist. + for index_file in ("index.faiss", "index.json", "index.pkl"): + if storage.file_exists(f"{index_path}/{index_file}"): + storage.delete_file(f"{index_path}/{index_file}") else: vectorstore = VectorCreator.create_vectorstore( settings.VECTOR_STORE, source_id=resolved_id diff --git a/application/parser/chunking_strategies.py b/application/parser/chunking_strategies.py index 87b40839..6e4a0dcb 100644 --- a/application/parser/chunking_strategies.py +++ b/application/parser/chunking_strategies.py @@ -177,7 +177,7 @@ class ParentChildChunker(_BaseStrategyChunker): 50) tokens. Each child stashes its parent window text in ``extra_info["parent_text"]`` so retrieval can expand to the parent later. The child text is what gets embedded; ``parent_text`` rides through - ``Document.to_langchain_format`` into vector-store metadata. + ``Document.to_vector_format`` into vector-store metadata. """ def _child_size(self) -> int: diff --git a/application/parser/file/base.py b/application/parser/file/base.py index 8e9b1015..5ee30579 100644 --- a/application/parser/file/base.py +++ b/application/parser/file/base.py @@ -2,8 +2,8 @@ from abc import abstractmethod from typing import Any, List -from langchain_core.documents import Document as LCDocument from application.parser.schema.base import Document +from application.vectorstore.document_class import Document as VectorDocument class BaseReader: @@ -13,7 +13,7 @@ class BaseReader: def load_data(self, *args: Any, **load_kwargs: Any) -> List[Document]: """Load data from the input directory.""" - def load_langchain_documents(self, **load_kwargs: Any) -> List[LCDocument]: - """Load data in LangChain document format.""" + def load_vector_documents(self, **load_kwargs: Any) -> List[VectorDocument]: + """Load data in the vector-store document format.""" docs = self.load_data(**load_kwargs) - return [d.to_langchain_format() for d in docs] + return [d.to_vector_format() for d in docs] diff --git a/application/parser/file/html_parser.py b/application/parser/file/html_parser.py index 4d69c870..dbf1c6bf 100644 --- a/application/parser/file/html_parser.py +++ b/application/parser/file/html_parser.py @@ -17,8 +17,26 @@ class HTMLParser(BaseParser): return {} def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, list[str]]: - from langchain_community.document_loaders import BSHTMLLoader + """Extract the visible text of an HTML file. - loader = BSHTMLLoader(file) - data = loader.load() - return data + Returns the text as a string, matching the other file parsers — the + title is recovered separately by ``get_file_metadata``. + """ + from bs4 import BeautifulSoup + + with open(file, "r", errors=errors) as f: + soup = BeautifulSoup(f, "html.parser") + return soup.get_text("\n") + + def get_file_metadata(self, file: Path) -> Dict: + """Return the document title, when the markup carries one.""" + from bs4 import BeautifulSoup + + try: + with open(file, "r", errors="ignore") as f: + soup = BeautifulSoup(f, "html.parser") + except OSError: + return {} + if soup.title and soup.title.string: + return {"title": str(soup.title.string)} + return {} diff --git a/application/parser/remote/base.py b/application/parser/remote/base.py index 74b6fce7..58a6cc43 100644 --- a/application/parser/remote/base.py +++ b/application/parser/remote/base.py @@ -2,8 +2,8 @@ from abc import abstractmethod from typing import Any, List -from langchain_core.documents import Document as LCDocument from application.parser.schema.base import Document +from application.vectorstore.document_class import Document as VectorDocument class BaseRemote: @@ -13,7 +13,7 @@ class BaseRemote: def load_data(self, *args: Any, **load_kwargs: Any) -> List[Document]: """Load data from the input directory.""" - def load_langchain_documents(self, **load_kwargs: Any) -> List[LCDocument]: - """Load data in LangChain document format.""" + def load_vector_documents(self, **load_kwargs: Any) -> List[VectorDocument]: + """Load data in the vector-store document format.""" docs = self.load_data(**load_kwargs) - return [d.to_langchain_format() for d in docs] + return [d.to_vector_format() for d in docs] diff --git a/application/parser/remote/reddit_loader.py b/application/parser/remote/reddit_loader.py index 22f5b185..7016b5e0 100644 --- a/application/parser/remote/reddit_loader.py +++ b/application/parser/remote/reddit_loader.py @@ -1,35 +1,89 @@ -from application.parser.remote.base import BaseRemote -from langchain_community.document_loaders import RedditPostsLoader import json +from typing import Any, Iterable, List + +from application.parser.remote.base import BaseRemote +from application.parser.schema.base import Document + +REQUIRED_FIELDS = ["client_id", "client_secret", "user_agent", "search_queries"] class RedditPostsLoaderRemote(BaseRemote): - def load_data(self, inputs): + """Load Reddit posts through praw. + + Create an app at https://www.reddit.com/prefs/apps/ to obtain the + ``client_id`` and ``client_secret``. + """ + + def load_data(self, inputs) -> List[Document]: + """Load posts for every configured subreddit or user. + + Args: + inputs: JSON string holding the praw credentials plus + ``search_queries``, and optionally ``mode`` + (``subreddit``/``username``), ``categories`` and + ``number_posts``. + + Returns: + One document per post, carrying the post body as text. + """ try: data = json.loads(inputs) except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON input: {e}") - required_fields = ["client_id", "client_secret", "user_agent", "search_queries"] - missing_fields = [field for field in required_fields if field not in data] + missing_fields = [field for field in REQUIRED_FIELDS if field not in data] if missing_fields: raise ValueError(f"Missing required fields: {', '.join(missing_fields)}") - client_id = data.get("client_id") - client_secret = data.get("client_secret") - user_agent = data.get("user_agent") - categories = data.get("categories", ["new", "hot"]) + mode = data.get("mode", "subreddit") + if mode not in ("subreddit", "username"): + raise ValueError( + "mode not correct, please enter 'username' or 'subreddit' as mode" + ) + categories = data.get("categories", ["new", "hot"]) search_queries = data.get("search_queries") number_posts = data.get("number_posts", 10) - self.loader = RedditPostsLoader( - client_id=client_id, - client_secret=client_secret, - user_agent=user_agent, - categories=categories, - mode=mode, - search_queries=search_queries, - number_posts=number_posts, + + import praw + + reddit = praw.Reddit( + client_id=data.get("client_id"), + client_secret=data.get("client_secret"), + user_agent=data.get("user_agent"), ) - documents = self.loader.load() - print(f"Loaded {len(documents)} documents from Reddit") + + documents: List[Document] = [] + for search_query in search_queries: + for category in categories: + if mode == "subreddit": + listing = reddit.subreddit(search_query) + else: + listing = reddit.redditor(search_query).submissions + documents.extend( + self._posts_to_documents( + getattr(listing, category)(limit=number_posts), category + ) + ) + return documents + + @staticmethod + def _posts_to_documents(posts: Iterable[Any], category: str) -> List[Document]: + """Convert praw submissions into documents.""" + documents = [] + for post in posts: + documents.append( + Document( + text=post.selftext, + doc_id=post.id, + extra_info={ + "post_subreddit": post.subreddit_name_prefixed, + "post_category": category, + "title": post.title, + "post_score": post.score, + "post_id": post.id, + "source": post.url, + "post_author": str(post.author), + }, + ) + ) return documents diff --git a/application/parser/schema/base.py b/application/parser/schema/base.py index a7453dd7..0d7a5c39 100644 --- a/application/parser/schema/base.py +++ b/application/parser/schema/base.py @@ -1,8 +1,8 @@ """Base schema for readers.""" from dataclasses import dataclass -from langchain_core.documents import Document as LCDocument from application.parser.schema.schema import BaseDocument +from application.vectorstore.document_class import Document as VectorDocument @dataclass @@ -23,12 +23,12 @@ class Document(BaseDocument): """Get Document type.""" return "Document" - def to_langchain_format(self) -> LCDocument: - """Convert struct to LangChain document format.""" + def to_vector_format(self) -> VectorDocument: + """Convert struct to the page_content/metadata shape vector stores take.""" metadata = self.extra_info or {} - return LCDocument(page_content=self.text, metadata=metadata) + return VectorDocument(page_content=self.text, metadata=metadata) @classmethod - def from_langchain_format(cls, doc: LCDocument) -> "Document": - """Convert struct from LangChain document format.""" + def from_vector_format(cls, doc: VectorDocument) -> "Document": + """Convert struct from the vector-store document shape.""" return cls(text=doc.page_content, extra_info=doc.metadata) diff --git a/application/requirements.txt b/application/requirements.txt index 535930ba..2a0c740d 100644 --- a/application/requirements.txt +++ b/application/requirements.txt @@ -36,16 +36,9 @@ joblib==1.5.3 jsonpatch==1.33 jsonpointer==3.1.1 kombu==5.6.2 -langchain==1.2.3 -langchain-community==0.4.1 -langchain-core==1.2.31 -langchain-openai==1.1.12 -langchain-text-splitters==1.1.2 -langsmith==0.7.31 lazy-object-proxy==1.12.0 lxml==6.1.1 markupsafe==3.0.3 -# dataclasses-json (via langchain-community) hard-caps marshmallow <4. marshmallow>=3.18.0,<4.0.0 mpmath==1.3.0 multidict==6.7.1 @@ -81,7 +74,11 @@ py==1.11.0 pydantic pydantic-core pydantic-settings +praw==8.0.2 +pymilvus==3.0.1 +milvus-lite==3.2.0; sys_platform != "win32" pypdf==6.15.0 +qdrant-client==1.19.0 python-dateutil==2.9.0.post0 python-dotenv python-jose==3.5.0 diff --git a/application/vectorstore/base.py b/application/vectorstore/base.py index b0f73836..e58fb4d9 100644 --- a/application/vectorstore/base.py +++ b/application/vectorstore/base.py @@ -3,9 +3,9 @@ import os from abc import ABC, abstractmethod import requests -from langchain_openai import OpenAIEmbeddings from application.core.settings import settings +from application.vectorstore.embeddings_openai import OpenAIEmbeddings from application.utils import get_encoding @@ -295,7 +295,6 @@ class BaseVectorStore(ABC): if embeddings_name == "openai_text-embedding-ada-002": if self.is_azure_configured(): - os.environ["OPENAI_API_TYPE"] = "azure" embedding_instance = EmbeddingsSingleton.get_instance( embeddings_name, model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME ) diff --git a/application/vectorstore/embeddings_openai.py b/application/vectorstore/embeddings_openai.py new file mode 100644 index 00000000..0ea31e36 --- /dev/null +++ b/application/vectorstore/embeddings_openai.py @@ -0,0 +1,89 @@ +"""OpenAI (and Azure OpenAI) embeddings built on the official ``openai`` SDK.""" + +from typing import List, Optional + +from application.core.settings import settings + +# openai >= 2.53 rejects a falsy api_key at construction; Azure authenticates +# through its own deployment credentials, so a placeholder keeps the client +# constructible when no key is configured. +NO_API_KEY = "sk-no-key" + +DEFAULT_MODEL = "text-embedding-ada-002" + + +class OpenAIEmbeddings: + """Embeddings client for OpenAI and Azure OpenAI. + + Mirrors the ``embed_query``/``embed_documents`` interface the vector + stores expect, matching :class:`RemoteEmbeddings` and + :class:`EmbeddingsWrapper`. + """ + + def __init__( + self, + openai_api_key: Optional[str] = None, + model: Optional[str] = None, + **kwargs, + ) -> None: + """Build the client, routing to Azure when the Azure settings are set. + + Args: + openai_api_key: API key; falls back to ``EMBEDDINGS_KEY`` then + ``OPENAI_API_KEY``. + model: Model name, or the Azure deployment name when running + against Azure. + """ + api_key = ( + openai_api_key + or settings.EMBEDDINGS_KEY + or settings.OPENAI_API_KEY + or NO_API_KEY + ) + self.model = model or DEFAULT_MODEL + self.dimension = None + + is_azure = bool( + settings.OPENAI_API_BASE + and settings.OPENAI_API_VERSION + and settings.AZURE_DEPLOYMENT_NAME + ) + if is_azure: + from openai import AzureOpenAI + + self.client = AzureOpenAI( + api_key=api_key, + azure_endpoint=settings.OPENAI_API_BASE, + api_version=settings.OPENAI_API_VERSION, + ) + else: + from openai import OpenAI + + base_url = settings.OPENAI_BASE_URL or None + self.client = OpenAI(api_key=api_key, base_url=base_url) + + def _embed(self, inputs: List[str]) -> List[List[float]]: + """Embed a batch, returning vectors in request order.""" + response = self.client.embeddings.create(model=self.model, input=inputs) + ordered = sorted(response.data, key=lambda item: item.index) + vectors = [item.embedding for item in ordered] + if vectors and self.dimension is None: + self.dimension = len(vectors[0]) + return vectors + + def embed_query(self, query: str) -> List[float]: + """Embed a single query string.""" + return self._embed([query])[0] + + def embed_documents(self, documents: List[str]) -> List[List[float]]: + """Embed a list of documents.""" + if not documents: + return [] + return self._embed(list(documents)) + + def __call__(self, text): + if isinstance(text, str): + return self.embed_query(text) + elif isinstance(text, list): + return self.embed_documents(text) + raise ValueError("Input must be a string or a list of strings") diff --git a/application/vectorstore/faiss.py b/application/vectorstore/faiss.py index 795669e3..18eac04d 100644 --- a/application/vectorstore/faiss.py +++ b/application/vectorstore/faiss.py @@ -1,13 +1,42 @@ +import io +import logging import os import tempfile -import io +import uuid +from typing import Any, Dict, List, Optional, Tuple -from langchain_community.vectorstores import FAISS +import numpy as np from application.core.settings import settings -from application.parser.schema.base import Document -from application.vectorstore.base import BaseVectorStore from application.storage.storage_creator import StorageCreator +from application.vectorstore.base import BaseVectorStore +from application.vectorstore.document_class import Document +from application.vectorstore.faiss_docstore import ( + dump_json_sidecar, + dump_pickle_sidecar, + load_json_sidecar, + load_pickle_sidecar, +) + +logger = logging.getLogger(__name__) + +# Sidecar holding chunk text and the row->id mapping. ``index.json`` is what +# this version writes; ``index.pkl`` is langchain's historical format, still +# read forever (uploads arrive in it) and still written for backward compat. +JSON_SIDECAR = "index.json" +PICKLE_SIDECAR = "index.pkl" +FAISS_INDEX = "index.faiss" + + +def _dependable_faiss_import(): + """Import faiss, with a clearer message than the raw ImportError.""" + try: + import faiss + except ImportError as e: + raise ImportError( + "Could not import faiss. Install it with `pip install faiss-cpu`." + ) from e + return faiss def get_vectorstore(path: str) -> str: @@ -41,6 +70,17 @@ def get_vectorstore(path: str) -> str: class FaissStore(BaseVectorStore): + """Vector store backed by a local FAISS index. + + Holds a flat L2 index plus an in-memory docstore mapping chunk ids to + their text and metadata, persisted through :class:`StorageCreator`. + """ + + # Ranks by L2 distance (lower is better), not cosine — so the number here + # is NOT comparable to the ``score_threshold`` the other stores honour, + # and must not be shown as one. + score_kind = "l2_distance" + def __init__(self, source_id: str, embeddings_key: str, docs_init=None): super().__init__() self.source_id = source_id @@ -48,101 +88,208 @@ class FaissStore(BaseVectorStore): self.embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key) self.storage = StorageCreator.get_storage() + self.index = None + self.documents: Dict[str, Dict[str, Any]] = {} + self.index_to_docstore_id: Dict[int, str] = {} + try: if docs_init: - self.docsearch = FAISS.from_documents(docs_init, self.embeddings) + self._build_from_documents(docs_init) else: - with tempfile.TemporaryDirectory() as temp_dir: - faiss_path = f"{self.path}/index.faiss" - pkl_path = f"{self.path}/index.pkl" - - if not self.storage.file_exists( - faiss_path - ) or not self.storage.file_exists(pkl_path): - raise FileNotFoundError( - f"Index files not found in storage at {self.path}" - ) - - faiss_file = self.storage.get_file(faiss_path) - pkl_file = self.storage.get_file(pkl_path) - - local_faiss_path = os.path.join(temp_dir, "index.faiss") - local_pkl_path = os.path.join(temp_dir, "index.pkl") - - with open(local_faiss_path, "wb") as f: - f.write(faiss_file.read()) - - with open(local_pkl_path, "wb") as f: - f.write(pkl_file.read()) - - self.docsearch = FAISS.load_local( - temp_dir, self.embeddings, allow_dangerous_deserialization=True - ) + self._load_from_storage() except Exception as e: raise Exception(f"Error loading FAISS index: {str(e)}") self.assert_embedding_dimensions(self.embeddings) - # LangChain's FAISS wrapper ranks by L2 distance (lower is better), not by - # a cosine similarity — so the number here is NOT comparable to the - # ``score_threshold`` the other stores honour, and must not be shown as one. - score_kind = "l2_distance" + # -- Construction ---------------------------------------------------- - def search(self, *args, **kwargs): + def _build_from_documents(self, docs_init) -> None: + """Create a fresh index seeded with ``docs_init``.""" + texts, metadatas = [], [] + for doc in docs_init: + texts.append(getattr(doc, "page_content", None) or getattr(doc, "text", "") or "") + metadatas.append(getattr(doc, "metadata", None) or getattr(doc, "extra_info", None) or {}) + + faiss = _dependable_faiss_import() + vectors = self.embeddings.embed_documents(texts) + self.index = faiss.IndexFlatL2(len(vectors[0])) + self._append(texts, metadatas, vectors) + + def _load_from_storage(self) -> None: + """Load the index and its sidecar, preferring JSON over the pickle.""" + faiss = _dependable_faiss_import() + faiss_path = f"{self.path}/{FAISS_INDEX}" + json_path = f"{self.path}/{JSON_SIDECAR}" + pickle_path = f"{self.path}/{PICKLE_SIDECAR}" + + if not self.storage.file_exists(faiss_path): + raise FileNotFoundError(f"Index files not found in storage at {self.path}") + + if self.storage.file_exists(json_path): + sidecar, loader = json_path, load_json_sidecar + elif self.storage.file_exists(pickle_path): + sidecar, loader = pickle_path, load_pickle_sidecar + else: + raise FileNotFoundError(f"Index files not found in storage at {self.path}") + + with tempfile.TemporaryDirectory() as temp_dir: + local_faiss = os.path.join(temp_dir, FAISS_INDEX) + with open(local_faiss, "wb") as f: + f.write(self.storage.get_file(faiss_path).read()) + self.index = faiss.read_index(local_faiss) + + self.documents, self.index_to_docstore_id = loader( + self.storage.get_file(sidecar).read() + ) + + # -- Internals ------------------------------------------------------- + + def _append(self, texts, metadatas, vectors, ids=None) -> List[str]: + """Add embedded rows to the index and docstore, returning their ids.""" + ids = list(ids) if ids else [str(uuid.uuid4()) for _ in texts] + self.index.add(np.array(vectors, dtype=np.float32)) + start = len(self.index_to_docstore_id) + for offset, (text, metadata, doc_id) in enumerate(zip(texts, metadatas, ids)): + self.documents[doc_id] = { + "page_content": text, + "metadata": dict(metadata or {}), + } + self.index_to_docstore_id[start + offset] = doc_id + return ids + + def _to_document(self, doc_id: str) -> Optional[Document]: + stored = self.documents.get(doc_id) + if stored is None: + return None + return Document( + page_content=stored.get("page_content", ""), + metadata=stored.get("metadata") or {}, + ) + + # -- Search ---------------------------------------------------------- + + def search(self, question: str, k: int = 4, *args, **kwargs) -> List[Document]: + """Return the ``k`` nearest chunks for ``question``.""" + return [doc for doc, _ in self.search_with_scores(question, k, *args, **kwargs)] + + def search_with_scores( + self, question: str, k: int = 4, *args, **kwargs + ) -> List[Tuple[Document, float]]: + """Same search as :meth:`search`, pairing each hit with its L2 distance.""" # FAISS has no relevance-threshold knob; drop it so the per-source # score_threshold is safely ignored rather than crashing the forward. kwargs.pop("score_threshold", None) - return self.docsearch.similarity_search(*args, **kwargs) + if self.index is None or self.index.ntotal == 0: + return [] - def search_with_scores(self, *args, **kwargs): - """Same search as :meth:`search`, pairing each hit with its L2 distance.""" - kwargs.pop("score_threshold", None) - results = self.docsearch.similarity_search_with_score(*args, **kwargs) - # The Documents come straight from the live in-memory docstore — return - # them untouched (the score rides alongside, never in their metadata) so - # the index can't be polluted and later persisted back to storage. - return [(doc, float(score)) for doc, score in results] + vector = np.array([self.embeddings.embed_query(question)], dtype=np.float32) + distances, rows = self.index.search(vector, min(k, self.index.ntotal)) - def add_texts(self, *args, **kwargs): - return self.docsearch.add_texts(*args, **kwargs) + results = [] + for distance, row in zip(distances[0], rows[0]): + if row == -1: + continue + doc_id = self.index_to_docstore_id.get(int(row)) + document = self._to_document(doc_id) if doc_id else None + if document is not None: + results.append((document, float(distance))) + return results - def _save_to_storage(self): - """ - Save the FAISS index to storage using temporary directory pattern. - Works consistently for both local and S3 storage. - """ - with tempfile.TemporaryDirectory() as temp_dir: - self.docsearch.save_local(temp_dir) + # -- Mutation -------------------------------------------------------- - faiss_path = os.path.join(temp_dir, "index.faiss") - pkl_path = os.path.join(temp_dir, "index.pkl") + def add_texts( + self, + texts: List[str], + metadatas: Optional[List[dict]] = None, + *args, + ids: Optional[List[str]] = None, + **kwargs, + ) -> List[str]: + """Embed and append ``texts`` to the index.""" + texts = list(texts) + if not texts: + return [] + metadatas = list(metadatas or [{} for _ in texts]) + vectors = self.embeddings.embed_documents(texts) + if self.index is None: + faiss = _dependable_faiss_import() + self.index = faiss.IndexFlatL2(len(vectors[0])) + return self._append(texts, metadatas, vectors, ids) - with open(faiss_path, "rb") as f_faiss: - faiss_data = f_faiss.read() + def delete_index(self, ids: Optional[List[str]] = None, *args, **kwargs): + """Delete the given chunk ids, or the whole index when ids are omitted.""" + if ids is None: + faiss = _dependable_faiss_import() + dimension = self.index.d if self.index is not None else None + self.index = faiss.IndexFlatL2(dimension) if dimension else None + self.documents = {} + self.index_to_docstore_id = {} + return True - with open(pkl_path, "rb") as f_pkl: - pkl_data = f_pkl.read() + missing = set(ids) - set(self.documents) + if missing: + raise ValueError(f"Chunk ids not found in index: {sorted(missing)}") - storage_path = get_vectorstore(self.source_id) - self.storage.save_file(io.BytesIO(faiss_data), f"{storage_path}/index.faiss") - self.storage.save_file(io.BytesIO(pkl_data), f"{storage_path}/index.pkl") + rows_by_id = {doc_id: row for row, doc_id in self.index_to_docstore_id.items()} + rows_to_drop = {rows_by_id[doc_id] for doc_id in ids} + self.index.remove_ids(np.array(sorted(rows_to_drop), dtype=np.int64)) + for doc_id in ids: + self.documents.pop(doc_id, None) + # remove_ids compacts the index, so the mapping has to be renumbered. + remaining = [ + doc_id + for row, doc_id in sorted(self.index_to_docstore_id.items()) + if row not in rows_to_drop + ] + self.index_to_docstore_id = dict(enumerate(remaining)) return True - def save_local(self, path=None): - if path: - os.makedirs(path, exist_ok=True) - self.docsearch.save_local(path) - + def add_chunk(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> str: + """Add a new chunk and save to storage.""" + ids = self.add_texts([text], [metadata or {}]) self._save_to_storage() + return ids[0] + def delete_chunk(self, chunk_id: str) -> bool: + """Delete a chunk and save to storage.""" + self.delete_index([chunk_id]) + self._save_to_storage() return True - def delete_index(self, *args, **kwargs): - return self.docsearch.delete(*args, **kwargs) + # -- Persistence ----------------------------------------------------- - def assert_embedding_dimensions(self, embeddings): - """Check that the word embedding dimension of the docsearch index matches the dimension of the word embeddings used.""" + def _write_index_files(self, directory: str) -> None: + """Write index.faiss plus both sidecars into ``directory``.""" + faiss = _dependable_faiss_import() + os.makedirs(directory, exist_ok=True) + faiss.write_index(self.index, os.path.join(directory, FAISS_INDEX)) + with open(os.path.join(directory, JSON_SIDECAR), "wb") as f: + f.write(dump_json_sidecar(self.documents, self.index_to_docstore_id)) + with open(os.path.join(directory, PICKLE_SIDECAR), "wb") as f: + f.write(dump_pickle_sidecar(self.documents, self.index_to_docstore_id)) + + def _save_to_storage(self) -> bool: + """Persist the index through the configured storage backend.""" + with tempfile.TemporaryDirectory() as temp_dir: + self._write_index_files(temp_dir) + storage_path = get_vectorstore(self.source_id) + for name in (FAISS_INDEX, JSON_SIDECAR, PICKLE_SIDECAR): + with open(os.path.join(temp_dir, name), "rb") as f: + self.storage.save_file(io.BytesIO(f.read()), f"{storage_path}/{name}") + return True + + def save_local(self, path: Optional[str] = None) -> bool: + if path: + self._write_index_files(path) + self._save_to_storage() + return True + + # -- Introspection --------------------------------------------------- + + def assert_embedding_dimensions(self, embeddings) -> None: + """Check the index width matches the embedding model's width.""" if ( settings.EMBEDDINGS_NAME == "huggingface_sentence-transformers/all-mpnet-base-v2" @@ -152,37 +299,22 @@ class FaissStore(BaseVectorStore): raise AttributeError( "'dimension' attribute not found in embeddings instance." ) - - docsearch_index_dimension = self.docsearch.index.d - if word_embedding_dimension != docsearch_index_dimension: + if self.index is None: + return + if word_embedding_dimension != self.index.d: raise ValueError( - f"Embedding dimension mismatch: embeddings.dimension ({word_embedding_dimension}) != docsearch index dimension ({docsearch_index_dimension})" + f"Embedding dimension mismatch: embeddings.dimension " + f"({word_embedding_dimension}) != docsearch index dimension " + f"({self.index.d})" ) - def get_chunks(self): - chunks = [] - if self.docsearch: - for doc_id, doc in self.docsearch.docstore._dict.items(): - chunk_data = { - "doc_id": doc_id, - "text": doc.page_content, - "metadata": doc.metadata, - } - chunks.append(chunk_data) - return chunks - - def add_chunk(self, text, metadata=None): - """Add a new chunk and save to storage.""" - metadata = metadata or {} - doc = Document(text=text, extra_info=metadata).to_langchain_format() - doc_id = self.docsearch.add_documents([doc]) - self._save_to_storage() - return doc_id - - - - def delete_chunk(self, chunk_id): - """Delete a chunk and save to storage.""" - self.delete_index([chunk_id]) - self._save_to_storage() - return True + def get_chunks(self) -> List[Dict[str, Any]]: + """Return every chunk held in the index.""" + return [ + { + "doc_id": doc_id, + "text": stored.get("page_content", ""), + "metadata": stored.get("metadata") or {}, + } + for doc_id, stored in self.documents.items() + ] diff --git a/application/vectorstore/faiss_docstore.py b/application/vectorstore/faiss_docstore.py new file mode 100644 index 00000000..3dcc4f34 --- /dev/null +++ b/application/vectorstore/faiss_docstore.py @@ -0,0 +1,230 @@ +"""Persistence for the FAISS docstore, independent of langchain. + +A FAISS index is two files: ``index.faiss`` (written by faiss itself) and a +sidecar holding the chunk text plus the row-to-id mapping. Historically that +sidecar was ``index.pkl`` — a pickled +``(InMemoryDocstore, index_to_docstore_id)`` tuple produced by langchain. + +Those files are still an active ingress: ``POST /api/upload_index`` accepts +them, and every index written before this change is in that format. So the +legacy sidecar is read forever, via a :class:`CompatUnpickler` that resolves +langchain's symbols onto local stand-ins and refuses everything else — the +previous ``pickle.load`` would execute arbitrary code from an uploaded file. + +``index.json`` is the format written going forward. ``index.pkl`` is written +alongside it, byte-compatible with langchain's layout, so an older DocsGPT can +still read an index this version has re-saved. +""" + +import io +import json +import pickle +import sys +import types +from contextlib import contextmanager +from typing import Any, Dict, Tuple + +# Module paths the legacy pickle references, mapped onto the stand-ins below. +_LEGACY_DOCUMENT_PATHS = ( + ("langchain_core.documents.base", "Document"), + ("langchain_core.documents", "Document"), + ("langchain.schema.document", "Document"), + ("langchain.docstore.document", "Document"), + ("langchain.schema", "Document"), +) +_LEGACY_DOCSTORE_PATHS = ( + ("langchain_community.docstore.in_memory", "InMemoryDocstore"), + ("langchain.docstore.in_memory", "InMemoryDocstore"), +) + +# The path langchain itself writes, and therefore the one we emit. +_CANONICAL_DOCUMENT_PATH = _LEGACY_DOCUMENT_PATHS[0] +_CANONICAL_DOCSTORE_PATH = _LEGACY_DOCSTORE_PATHS[0] + + +class LegacyDocument: + """Stand-in for ``langchain_core.documents.base.Document``. + + langchain's Document is a pydantic v2 model, which pickles as + ``__newobj__(cls)`` plus a state dict; reconstruction only needs a class + whose ``__setstate__`` understands that shape. + """ + + __slots__ = ("page_content", "metadata", "id", "type") + + def __init__(self, page_content: str = "", metadata: Dict[str, Any] = None, + id: Any = None) -> None: + self.page_content = page_content + self.metadata = metadata if metadata is not None else {} + self.id = id + self.type = "Document" + + def __setstate__(self, state: Any) -> None: + fields = state.get("__dict__", state) if isinstance(state, dict) else {} + self.page_content = fields.get("page_content", "") + self.metadata = fields.get("metadata") or {} + self.id = fields.get("id") + self.type = fields.get("type", "Document") + + def __getstate__(self) -> Dict[str, Any]: + """Emit pydantic v2's pickle state so langchain can read it back.""" + return { + "__dict__": { + "id": self.id, + "metadata": self.metadata, + "page_content": self.page_content, + "type": self.type, + }, + "__pydantic_extra__": None, + "__pydantic_fields_set__": {"page_content", "metadata"}, + "__pydantic_private__": None, + } + + +class LegacyDocstore: + """Stand-in for ``langchain_community.docstore.in_memory.InMemoryDocstore``.""" + + def __init__(self, _dict: Dict[str, LegacyDocument] = None) -> None: + self._dict = _dict if _dict is not None else {} + + def __setstate__(self, state: Any) -> None: + self._dict = state.get("_dict", {}) if isinstance(state, dict) else {} + + def __getstate__(self) -> Dict[str, Any]: + return {"_dict": self._dict} + + +class CompatUnpickler(pickle.Unpickler): + """Unpickler that accepts only the symbols a FAISS sidecar may contain. + + langchain's document and docstore classes resolve to the local stand-ins; + every other global is refused, so a hostile ``index.pkl`` cannot execute + code through ``__reduce__``. + """ + + _ALLOWED = { + **{path: LegacyDocument for path in _LEGACY_DOCUMENT_PATHS}, + **{path: LegacyDocstore for path in _LEGACY_DOCSTORE_PATHS}, + } + + def find_class(self, module: str, name: str): + mapped = self._ALLOWED.get((module, name)) + if mapped is not None: + return mapped + raise pickle.UnpicklingError( + f"Refusing to load unexpected symbol from FAISS index: {module}.{name}" + ) + + +@contextmanager +def _legacy_pickle_names(): + """Expose the stand-ins under langchain's module paths while pickling. + + ``pickle`` records a class by its module and qualified name, and verifies + the pair resolves before writing it. Registering throwaway modules for the + duration of the dump is what lets the output reference langchain's paths — + and therefore stay readable by a DocsGPT that still has langchain — without + langchain being installed here. Any pre-existing modules are restored. + """ + installed = [] + for cls, (module_path, attr) in ( + (LegacyDocument, _CANONICAL_DOCUMENT_PATH), + (LegacyDocstore, _CANONICAL_DOCSTORE_PATH), + ): + original_module = cls.__module__ + original_name = cls.__qualname__ + cls.__module__, cls.__qualname__ = module_path, attr + + parts = module_path.split(".") + for depth in range(1, len(parts) + 1): + parent = ".".join(parts[:depth]) + if parent not in sys.modules: + sys.modules[parent] = types.ModuleType(parent) + installed.append(parent) + setattr(sys.modules[module_path], attr, cls) + installed.append((cls, original_module, original_name)) + try: + yield + finally: + for entry in reversed(installed): + if isinstance(entry, tuple): + cls, original_module, original_name = entry + cls.__module__, cls.__qualname__ = original_module, original_name + else: + sys.modules.pop(entry, None) + + +def load_pickle_sidecar(data: bytes) -> Tuple[Dict[str, Dict[str, Any]], Dict[int, str]]: + """Read a legacy ``index.pkl`` into plain dictionaries. + + Args: + data: Raw bytes of the pickled ``(docstore, index_to_docstore_id)``. + + Returns: + ``(documents, index_to_docstore_id)`` where ``documents`` maps a + chunk id to ``{"page_content": ..., "metadata": ...}`` and the second + maps a FAISS row number to a chunk id. + """ + docstore, mapping = CompatUnpickler(io.BytesIO(data)).load() + documents = { + doc_id: { + "page_content": getattr(doc, "page_content", "") or "", + "metadata": getattr(doc, "metadata", None) or {}, + } + for doc_id, doc in getattr(docstore, "_dict", {}).items() + } + return documents, {int(row): doc_id for row, doc_id in (mapping or {}).items()} + + +def dump_pickle_sidecar( + documents: Dict[str, Dict[str, Any]], index_to_docstore_id: Dict[int, str] +) -> bytes: + """Serialize to the legacy langchain layout, for backward compatibility.""" + docstore = LegacyDocstore( + { + doc_id: LegacyDocument( + page_content=doc.get("page_content", ""), + metadata=doc.get("metadata") or {}, + ) + for doc_id, doc in documents.items() + } + ) + with _legacy_pickle_names(): + return pickle.dumps((docstore, dict(index_to_docstore_id)), protocol=4) + + +def load_json_sidecar(data: bytes) -> Tuple[Dict[str, Dict[str, Any]], Dict[int, str]]: + """Read an ``index.json`` sidecar.""" + payload = json.loads(data.decode("utf-8")) + documents = { + doc_id: { + "page_content": doc.get("page_content", "") or "", + "metadata": doc.get("metadata") or {}, + } + for doc_id, doc in (payload.get("documents") or {}).items() + } + mapping = { + int(row): doc_id + for row, doc_id in (payload.get("index_to_docstore_id") or {}).items() + } + return documents, mapping + + +def dump_json_sidecar( + documents: Dict[str, Dict[str, Any]], index_to_docstore_id: Dict[int, str] +) -> bytes: + """Serialize to the pickle-free sidecar written going forward.""" + payload = { + "version": 1, + "documents": { + doc_id: { + "page_content": doc.get("page_content", ""), + "metadata": doc.get("metadata") or {}, + } + for doc_id, doc in documents.items() + }, + "index_to_docstore_id": { + str(row): doc_id for row, doc_id in index_to_docstore_id.items() + }, + } + return json.dumps(payload, ensure_ascii=False).encode("utf-8") diff --git a/application/vectorstore/milvus.py b/application/vectorstore/milvus.py index 8d67b061..3050a877 100644 --- a/application/vectorstore/milvus.py +++ b/application/vectorstore/milvus.py @@ -1,41 +1,210 @@ -from typing import List, Optional -from uuid import uuid4 - +import logging +import os +import uuid +from contextlib import contextmanager +from typing import Any, Dict, List, Optional, Tuple from application.core.settings import settings from application.vectorstore.base import BaseVectorStore +from application.vectorstore.document_class import Document + + +@contextmanager +def _without_milvus_uri_env(): + """Hide ``MILVUS_URI`` from pymilvus while it is imported. + + pymilvus reads the ``MILVUS_URI`` environment variable at import time and + rejects anything that is not an ``http[s]://`` URL. DocsGPT's setting of + the same name defaults to a Milvus Lite file path, and ``load_dotenv`` + puts it on the environment — so an unguarded import raises before the + store can pass its own ``uri``. The value is restored immediately after. + """ + sentinel = object() + previous = os.environ.pop("MILVUS_URI", sentinel) + try: + yield + finally: + if previous is not sentinel: + os.environ["MILVUS_URI"] = previous class MilvusStore(BaseVectorStore): + """Vector store backed by Milvus through the native ``pymilvus`` client. + + Defaults to Milvus Lite (a local file at ``MILVUS_URI``); point + ``MILVUS_URI`` at a server URL to use a full deployment. Rows carry an + indexed ``source_id`` so one collection can hold many sources. + """ + + score_kind = "cosine_similarity" + def __init__(self, source_id: str = "", embeddings_key: str = "embeddings"): super().__init__() - from langchain_milvus import Milvus + with _without_milvus_uri_env(): + from pymilvus import DataType, MilvusClient - connection_args = { - "uri": settings.MILVUS_URI, - "token": settings.MILVUS_TOKEN, - } - self._docsearch = Milvus( - embedding_function=self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key), - collection_name=settings.MILVUS_COLLECTION_NAME, - connection_args=connection_args, + self._DataType = DataType + self._source_id = str(source_id).replace("application/indexes/", "").rstrip("/") + self._collection = settings.MILVUS_COLLECTION_NAME + self._embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key) + self._client = MilvusClient( + uri=settings.MILVUS_URI, token=settings.MILVUS_TOKEN or "" ) - self._source_id = source_id + self._ensure_collection() - def search(self, question, k=2, *args, **kwargs): - # Drop the per-source score_threshold (unsupported here) so it is safely - # ignored instead of being forwarded into the langchain call. - kwargs.pop("score_threshold", None) - expr = f"source_id == '{self._source_id}'" - return self._docsearch.similarity_search(query=question, k=k, expr=expr, *args, **kwargs) + @property + def _filter(self) -> str: + """Boolean expression scoping every read to this source.""" + escaped = self._source_id.replace('"', '\\"') + return f'source_id == "{escaped}"' - def add_texts(self, texts: List[str], metadatas: Optional[List[dict]], *args, **kwargs): - ids = [str(uuid4()) for _ in range(len(texts))] + def _dimension(self) -> int: + """Resolve the embedding width, probing the model when unset.""" + dimension = getattr(self._embeddings, "dimension", None) + if not dimension: + dimension = len(self._embeddings.embed_query("dimension probe")) + return dimension - return self._docsearch.add_texts(texts=texts, metadatas=metadatas, ids=ids, *args, **kwargs) + def _ensure_collection(self) -> None: + """Create the collection, its vector index and the source_id index.""" + if self._client.has_collection(self._collection): + return + DataType = self._DataType + schema = self._client.create_schema(auto_id=False, enable_dynamic_field=True) + schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=64) + schema.add_field("vector", DataType.FLOAT_VECTOR, dim=self._dimension()) + schema.add_field("text", DataType.VARCHAR, max_length=65535) + schema.add_field("source_id", DataType.VARCHAR, max_length=512) + schema.add_field("metadata", DataType.JSON, nullable=True) + + index_params = self._client.prepare_index_params() + index_params.add_index( + field_name="vector", index_type="AUTOINDEX", metric_type="COSINE" + ) + try: + self._client.create_collection( + collection_name=self._collection, + schema=schema, + index_params=index_params, + ) + except Exception as e: + # A concurrent worker may have created it first. + if "already exist" not in str(e).lower(): + raise + + @staticmethod + def _to_document(row: Dict[str, Any]) -> Document: + row = row or {} + return Document( + page_content=row.get("text") or "", + metadata=row.get("metadata") or {}, + ) + + def search(self, question: str, k: int = 2, *args, **kwargs) -> List[Document]: + """Return the ``k`` nearest chunks for ``question``.""" + return [doc for doc, _ in self.search_with_scores(question, k, *args, **kwargs)] + + def search_with_scores( + self, + question: str, + k: int = 2, + *args, + score_threshold: Optional[float] = None, + **kwargs, + ) -> List[Tuple[Document, float]]: + """Search, pairing each hit with its cosine similarity.""" + results = self._client.search( + collection_name=self._collection, + data=[self._embeddings.embed_query(question)], + filter=self._filter, + limit=k, + output_fields=["text", "metadata"], + ) + hits = results[0] if results else [] + out = [] + for hit in hits: + score = hit.get("distance") + if score_threshold is not None and score is not None and score < score_threshold: + continue + out.append((self._to_document(hit.get("entity") or {}), score)) + return out + + def add_texts( + self, + texts: List[str], + metadatas: Optional[List[dict]] = None, + *args, + **kwargs, + ) -> List[str]: + """Embed and insert ``texts``, stamping each with the source id.""" + texts = list(texts) + if not texts: + return [] + metadatas = list(metadatas or [{} for _ in texts]) + vectors = self._embeddings.embed_documents(texts) + + rows, ids = [], [] + for text, metadata, vector in zip(texts, metadatas, vectors): + row_id = str(uuid.uuid4()) + row_metadata = dict(metadata or {}) + row_metadata["source_id"] = self._source_id + ids.append(row_id) + rows.append( + { + "id": row_id, + "vector": vector, + "text": text, + "source_id": self._source_id, + "metadata": row_metadata, + } + ) + self._client.insert(collection_name=self._collection, data=rows) + return ids def save_local(self, *args, **kwargs): + """No-op: Milvus persists server-side.""" pass def delete_index(self, *args, **kwargs): - pass + """Delete every row belonging to this source.""" + try: + return self._client.delete( + collection_name=self._collection, filter=self._filter + ) + except Exception as e: + logging.error("Error deleting index: %s", e, exc_info=True) + return None + + def get_chunks(self) -> List[Dict[str, Any]]: + """Return every chunk stored for this source.""" + try: + rows = self._client.query( + collection_name=self._collection, + filter=self._filter, + output_fields=["id", "text", "metadata"], + ) + return [ + { + "doc_id": row.get("id"), + "text": row.get("text"), + "metadata": row.get("metadata") or {}, + } + for row in rows + ] + except Exception as e: + logging.error("Error getting chunks: %s", e, exc_info=True) + return [] + + def add_chunk(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> str: + """Add one chunk and return its id.""" + ids = self.add_texts([text], [metadata or {}]) + return ids[0] + + def delete_chunk(self, chunk_id: str) -> bool: + """Delete a single chunk by id.""" + try: + self._client.delete(collection_name=self._collection, ids=[chunk_id]) + return True + except Exception as e: + logging.error("Error deleting chunk: %s", e, exc_info=True) + return False diff --git a/application/vectorstore/qdrant.py b/application/vectorstore/qdrant.py index f0311283..fed977dc 100644 --- a/application/vectorstore/qdrant.py +++ b/application/vectorstore/qdrant.py @@ -1,17 +1,31 @@ import logging -from application.vectorstore.base import BaseVectorStore +import uuid +from typing import Any, Dict, List, Optional, Tuple + from application.core.settings import settings +from application.vectorstore.base import BaseVectorStore from application.vectorstore.document_class import Document class QdrantStore(BaseVectorStore): - def __init__(self, source_id: str = "", embeddings_key: str = "embeddings"): - from qdrant_client import models - from langchain_community.vectorstores.qdrant import Qdrant + """Vector store backed by Qdrant through the native ``qdrant-client``. - # Store the source_id for use in add_chunk + Points carry a ``page_content`` payload plus a nested ``metadata`` object; + every read is filtered to the store's ``source_id`` so one collection can + hold many sources. + """ + + score_kind = "cosine_similarity" + + def __init__(self, source_id: str = "", embeddings_key: str = "embeddings"): + super().__init__() + from qdrant_client import QdrantClient, models + + self._models = models self._source_id = str(source_id).replace("application/indexes/", "").rstrip("/") - + self._collection = settings.QDRANT_COLLECTION_NAME + self._embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key) + self._filter = models.Filter( must=[ models.FieldCondition( @@ -20,120 +34,178 @@ class QdrantStore(BaseVectorStore): ) ] ) + self._client = QdrantClient(**self._client_kwargs()) + self._ensure_collection() - embedding=self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key) - self._docsearch = Qdrant.construct_instance( - ["TEXT_TO_OBTAIN_EMBEDDINGS_DIMENSION"], - embedding=embedding, - collection_name=settings.QDRANT_COLLECTION_NAME, - location=settings.QDRANT_LOCATION, - url=settings.QDRANT_URL, - port=settings.QDRANT_PORT, - grpc_port=settings.QDRANT_GRPC_PORT, - https=settings.QDRANT_HTTPS, - prefer_grpc=settings.QDRANT_PREFER_GRPC, - api_key=settings.QDRANT_API_KEY, - prefix=settings.QDRANT_PREFIX, - timeout=settings.QDRANT_TIMEOUT, - path=settings.QDRANT_PATH, - distance_func=settings.QDRANT_DISTANCE_FUNC, - ) + @staticmethod + def _client_kwargs() -> Dict[str, Any]: + """Build ``QdrantClient`` kwargs, dropping unset optional settings. + + ``location``, ``url`` and ``path`` are mutually exclusive in + qdrant-client, so only the ones actually configured are passed. + """ + kwargs: Dict[str, Any] = { + "prefer_grpc": settings.QDRANT_PREFER_GRPC, + "grpc_port": settings.QDRANT_GRPC_PORT, + } + optional = { + "location": settings.QDRANT_LOCATION, + "url": settings.QDRANT_URL, + "host": settings.QDRANT_HOST, + "port": settings.QDRANT_PORT, + "https": settings.QDRANT_HTTPS, + "api_key": settings.QDRANT_API_KEY, + "prefix": settings.QDRANT_PREFIX, + "timeout": settings.QDRANT_TIMEOUT, + "path": settings.QDRANT_PATH, + } + kwargs.update({k: v for k, v in optional.items() if v is not None}) + return kwargs + + def _dimension(self) -> int: + """Resolve the embedding width, probing the model when unset.""" + dimension = getattr(self._embeddings, "dimension", None) + if not dimension: + dimension = len(self._embeddings.embed_query("dimension probe")) + return dimension + + def _ensure_collection(self) -> None: + """Create the collection and the source_id payload index if missing.""" + models = self._models try: - collections = self._docsearch.client.get_collections() - collection_exists = settings.QDRANT_COLLECTION_NAME in [ - collection.name for collection in collections.collections - ] - - if not collection_exists: - self._docsearch.client.recreate_collection( - collection_name=settings.QDRANT_COLLECTION_NAME, - vectors_config=models.VectorParams(size=embedding.client[1].word_embedding_dimension, distance=models.Distance.COSINE), + if not self._client.collection_exists(self._collection): + self._client.create_collection( + collection_name=self._collection, + vectors_config=models.VectorParams( + size=self._dimension(), + distance=models.Distance[settings.QDRANT_DISTANCE_FUNC.upper()], + ), ) - - # Ensure the required index exists for metadata.source_id - try: - self._docsearch.client.create_payload_index( - collection_name=settings.QDRANT_COLLECTION_NAME, - field_name="metadata.source_id", - field_schema=models.PayloadSchemaType.KEYWORD, - ) - except Exception as index_error: - # Index might already exist, which is fine - if "already exists" not in str(index_error).lower(): - logging.warning(f"Could not create index for metadata.source_id: {index_error}") - + self._client.create_payload_index( + collection_name=self._collection, + field_name="metadata.source_id", + field_schema=models.PayloadSchemaType.KEYWORD, + ) except Exception as e: - logging.warning(f"Could not check for collection: {e}") + # A concurrent worker may have created either object first. + if "already exists" not in str(e).lower(): + logging.warning("Qdrant collection setup: %s", e) - def search(self, *args, **kwargs): - # Drop the per-source score_threshold (unsupported here) so it is safely - # ignored instead of being forwarded into the langchain call. - kwargs.pop("score_threshold", None) - return self._docsearch.similarity_search(filter=self._filter, *args, **kwargs) + @staticmethod + def _to_document(payload: Dict[str, Any]) -> Document: + payload = payload or {} + return Document( + page_content=payload.get("page_content", ""), + metadata=payload.get("metadata") or {}, + ) - def add_texts(self, *args, **kwargs): - return self._docsearch.add_texts(*args, **kwargs) + def search(self, question: str, k: int = 2, *args, **kwargs) -> List[Document]: + """Return the ``k`` nearest chunks for ``question``.""" + return [doc for doc, _ in self.search_with_scores(question, k, *args, **kwargs)] + + def search_with_scores( + self, + question: str, + k: int = 2, + *args, + score_threshold: Optional[float] = None, + **kwargs, + ) -> List[Tuple[Document, float]]: + """Search, pairing each hit with its cosine similarity.""" + hits = self._client.query_points( + collection_name=self._collection, + query=self._embeddings.embed_query(question), + query_filter=self._filter, + limit=k, + with_payload=True, + score_threshold=score_threshold, + ).points + return [(self._to_document(hit.payload), hit.score) for hit in hits] + + def add_texts( + self, + texts: List[str], + metadatas: Optional[List[dict]] = None, + *args, + **kwargs, + ) -> List[str]: + """Embed and upsert ``texts``, stamping each with the source id.""" + texts = list(texts) + if not texts: + return [] + metadatas = list(metadatas or [{} for _ in texts]) + vectors = self._embeddings.embed_documents(texts) + + points, ids = [], [] + for text, metadata, vector in zip(texts, metadatas, vectors): + point_id = str(uuid.uuid4()) + payload_metadata = dict(metadata or {}) + payload_metadata["source_id"] = self._source_id + ids.append(point_id) + points.append( + self._models.PointStruct( + id=point_id, + vector=vector, + payload={"page_content": text, "metadata": payload_metadata}, + ) + ) + self._client.upsert(collection_name=self._collection, points=points) + return ids def save_local(self, *args, **kwargs): + """No-op: Qdrant persists server-side.""" pass def delete_index(self, *args, **kwargs): - return self._docsearch.client.delete( - collection_name=settings.QDRANT_COLLECTION_NAME, points_selector=self._filter + """Delete every point belonging to this source.""" + return self._client.delete( + collection_name=self._collection, + points_selector=self._models.FilterSelector(filter=self._filter), ) - def get_chunks(self): + def get_chunks(self) -> List[Dict[str, Any]]: + """Return every chunk stored for this source.""" + chunks: List[Dict[str, Any]] = [] + offset = None try: - - chunks = [] - offset = None while True: - records, offset = self._docsearch.client.scroll( - collection_name=settings.QDRANT_COLLECTION_NAME, + records, offset = self._client.scroll( + collection_name=self._collection, scroll_filter=self._filter, - limit=10, + limit=100, with_payload=True, with_vectors=False, offset=offset, ) for record in records: - doc_id = record.id - text = record.payload.get("page_content") - metadata = record.payload.get("metadata") + payload = record.payload or {} chunks.append( - {"doc_id": doc_id, "text": text, "metadata": metadata} + { + "doc_id": str(record.id), + "text": payload.get("page_content"), + "metadata": payload.get("metadata") or {}, + } ) if offset is None: break return chunks except Exception as e: - logging.error(f"Error getting chunks: {e}", exc_info=True) + logging.error("Error getting chunks: %s", e, exc_info=True) return [] - def add_chunk(self, text, metadata=None): - import uuid - metadata = metadata or {} - - # Create a copy to avoid modifying the original metadata - final_metadata = metadata.copy() - - # Ensure the source_id is in the metadata so the chunk can be found by filters - final_metadata["source_id"] = self._source_id - - doc = Document(page_content=text, metadata=final_metadata) - # Generate a unique ID for the document - doc_id = str(uuid.uuid4()) - doc.id = doc_id - doc_ids = self._docsearch.add_documents([doc]) - return doc_ids[0] if doc_ids else doc_id + def add_chunk(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> str: + """Add one chunk and return its id.""" + ids = self.add_texts([text], [metadata or {}]) + return ids[0] - def delete_chunk(self, chunk_id): + def delete_chunk(self, chunk_id: str) -> bool: + """Delete a single chunk by id.""" try: - self._docsearch.client.delete( - collection_name=settings.QDRANT_COLLECTION_NAME, - points_selector=[chunk_id], + self._client.delete( + collection_name=self._collection, + points_selector=self._models.PointIdsList(points=[chunk_id]), ) return True except Exception as e: - logging.error(f"Error deleting chunk: {e}", exc_info=True) + logging.error("Error deleting chunk: %s", e, exc_info=True) return False diff --git a/application/worker.py b/application/worker.py index 1b91c231..19b298c5 100755 --- a/application/worker.py +++ b/application/worker.py @@ -679,7 +679,7 @@ def ingest_worker( ) raw_docs = chunker.chunk(documents=raw_docs) - docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs] + docs = [Document.to_vector_format(raw_doc) for raw_doc in raw_docs] vector_store_path = os.path.join(temp_dir, "vector_store") os.makedirs(vector_store_path, exist_ok=True) @@ -1224,7 +1224,7 @@ def remote_worker( duplicate_headers=cfg.chunking.duplicate_headers, ) raw_docs = chunker.chunk(documents=raw_docs) - docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs] + docs = [Document.to_vector_format(raw_doc) for raw_doc in raw_docs] tokens = count_tokens_docs(docs) logging.info("Total tokens calculated: %d", tokens) @@ -2068,7 +2068,7 @@ def ingest_connector( source, start=temp_dir ) - docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs] + docs = [Document.to_vector_format(raw_doc) for raw_doc in raw_docs] # Validate operation_mode here too (the source_uuid path # at the top of the function only branches on the diff --git a/tests/parser/file/test_html_parser.py b/tests/parser/file/test_html_parser.py index 1bbfd21d..9726ddcb 100644 --- a/tests/parser/file/test_html_parser.py +++ b/tests/parser/file/test_html_parser.py @@ -1,16 +1,19 @@ import pytest -from pathlib import Path -from unittest.mock import patch, MagicMock - -import sys -import types from application.parser.file.html_parser import HTMLParser +HTML = ( + "My Page" + "

Heading

Hello world.

" +) + + @pytest.fixture -def html_parser(): - return HTMLParser() +def html_file(tmp_path): + path = tmp_path / "page.html" + path.write_text(HTML) + return path def test_html_init_parser(): @@ -21,23 +24,33 @@ def test_html_init_parser(): assert parser.parser_config_set -def test_html_parser_parse_file(): - parser = HTMLParser() - mock_doc = MagicMock() - mock_doc.page_content = "Extracted HTML content" - mock_doc.metadata = {"source": "test.html"} +def test_html_parser_extracts_text(html_file): + text = HTMLParser().parse_file(html_file) + assert "Heading" in text + assert "Hello world." in text - fake_lc = types.ModuleType("langchain_community") - fake_dl = types.ModuleType("langchain_community.document_loaders") - bshtml_mock = MagicMock(return_value=MagicMock(load=MagicMock(return_value=[mock_doc]))) - fake_dl.BSHTMLLoader = bshtml_mock - fake_lc.document_loaders = fake_dl +def test_html_parser_returns_plain_text(html_file): + """bulk.py stringifies parser output, so it must not be a repr. - with patch.dict(sys.modules, { - "langchain_community": fake_lc, - "langchain_community.document_loaders": fake_dl, - }): - result = parser.parse_file(Path("test.html")) - assert result == [mock_doc] - bshtml_mock.assert_called_once_with(Path("test.html")) + The previous langchain BSHTMLLoader returned a Document whose ``str()`` is + ``page_content='...' metadata={...}``, which leaked into indexed content. + """ + text = HTMLParser().parse_file(html_file) + assert isinstance(text, str) + assert "page_content" not in str(text) + assert "metadata=" not in str(text) + + +def test_html_parser_reports_title_metadata(html_file): + assert HTMLParser().get_file_metadata(html_file) == {"title": "My Page"} + + +def test_html_parser_metadata_without_title(tmp_path): + path = tmp_path / "no_title.html" + path.write_text("

Just text.

") + assert HTMLParser().get_file_metadata(path) == {} + + +def test_html_parser_metadata_unreadable_file(tmp_path): + assert HTMLParser().get_file_metadata(tmp_path / "missing.html") == {} diff --git a/tests/parser/remote/test_reddit_loader.py b/tests/parser/remote/test_reddit_loader.py index 26199490..ddc6671c 100644 --- a/tests/parser/remote/test_reddit_loader.py +++ b/tests/parser/remote/test_reddit_loader.py @@ -1,83 +1,103 @@ import json -from unittest.mock import patch, MagicMock +import sys +import types +from unittest.mock import MagicMock, patch + import pytest from application.parser.remote.reddit_loader import RedditPostsLoaderRemote +def _post(post_id, title, body): + post = MagicMock() + post.id = post_id + post.title = title + post.selftext = body + post.subreddit_name_prefixed = "r/python" + post.score = 42 + post.url = f"https://reddit.com/{post_id}" + post.author = "someone" + return post + + +def _fake_praw(posts): + """Install a stub ``praw`` module whose listings yield ``posts``.""" + praw = types.ModuleType("praw") + reddit = MagicMock() + listing = MagicMock() + listing.new.return_value = posts + listing.hot.return_value = posts + reddit.subreddit.return_value = listing + reddit.redditor.return_value = MagicMock(submissions=listing) + praw.Reddit = MagicMock(return_value=reddit) + return praw, reddit, listing + + +BASE_PAYLOAD = { + "client_id": "cid", + "client_secret": "csecret", + "user_agent": "ua", + "search_queries": ["python"], +} + + +@pytest.mark.unit class TestRedditPostsLoaderRemote: def test_invalid_json_raises(self): - loader = RedditPostsLoaderRemote() - with pytest.raises(ValueError) as exc: - loader.load_data("not a json") - assert "Invalid JSON input" in str(exc.value) + with pytest.raises(ValueError, match="Invalid JSON input"): + RedditPostsLoaderRemote().load_data("not a json") def test_missing_required_fields_raises(self): - loader = RedditPostsLoaderRemote() - payload = json.dumps({"client_id": "id"}) - with pytest.raises(ValueError) as exc: - loader.load_data(payload) - assert "Missing required fields" in str(exc.value) + with pytest.raises(ValueError, match="Missing required fields") as exc: + RedditPostsLoaderRemote().load_data(json.dumps({"client_id": "id"})) assert "client_secret" in str(exc.value) - @patch("application.parser.remote.reddit_loader.RedditPostsLoader") - def test_constructs_loader_and_loads_with_defaults(self, MockRedditLoader): - loader = RedditPostsLoaderRemote() + def test_invalid_mode_raises(self): + payload = {**BASE_PAYLOAD, "mode": "comments"} + praw, _, _ = _fake_praw([]) + with patch.dict(sys.modules, {"praw": praw}): + with pytest.raises(ValueError, match="username.*subreddit|subreddit"): + RedditPostsLoaderRemote().load_data(json.dumps(payload)) - instance = MagicMock() - docs = [MagicMock(), MagicMock()] - instance.load.return_value = docs - MockRedditLoader.return_value = instance + def test_subreddit_mode_returns_documents(self): + praw, reddit, listing = _fake_praw([_post("p1", "First", "Body one")]) + with patch.dict(sys.modules, {"praw": praw}): + docs = RedditPostsLoaderRemote().load_data(json.dumps(BASE_PAYLOAD)) - payload = { - "client_id": "cid", - "client_secret": "csecret", - "user_agent": "ua", - "search_queries": ["r/langchain"], - } - - result = loader.load_data(json.dumps(payload)) - - MockRedditLoader.assert_called_once_with( - client_id="cid", - client_secret="csecret", - user_agent="ua", - categories=["new", "hot"], - mode="subreddit", - search_queries=["r/langchain"], - number_posts=10, + praw.Reddit.assert_called_once_with( + client_id="cid", client_secret="csecret", user_agent="ua" ) - instance.load.assert_called_once() - assert result == docs + reddit.subreddit.assert_called_with("python") + # Default categories are new + hot, so the single post arrives twice. + assert len(docs) == 2 + assert docs[0].text == "Body one" + assert docs[0].doc_id == "p1" + assert docs[0].extra_info["title"] == "First" + assert docs[0].extra_info["source"] == "https://reddit.com/p1" + assert docs[0].extra_info["post_category"] == "new" - @patch("application.parser.remote.reddit_loader.RedditPostsLoader") - def test_constructs_loader_and_loads_with_overrides(self, MockRedditLoader): - loader = RedditPostsLoaderRemote() + def test_documents_convert_to_vector_format(self): + """The remote base class calls to_vector_format on every result.""" + praw, _, _ = _fake_praw([_post("p1", "First", "Body one")]) + with patch.dict(sys.modules, {"praw": praw}): + docs = RedditPostsLoaderRemote().load_data(json.dumps(BASE_PAYLOAD)) + converted = docs[0].to_vector_format() + assert converted.page_content == "Body one" + assert converted.metadata["title"] == "First" - instance = MagicMock() - instance.load.return_value = [] - MockRedditLoader.return_value = instance + def test_username_mode_uses_redditor_submissions(self): + payload = {**BASE_PAYLOAD, "mode": "username", "categories": ["new"]} + praw, reddit, _ = _fake_praw([_post("p2", "Second", "Body two")]) + with patch.dict(sys.modules, {"praw": praw}): + docs = RedditPostsLoaderRemote().load_data(json.dumps(payload)) - payload = { - "client_id": "cid", - "client_secret": "csecret", - "user_agent": "ua", - "search_queries": ["python"], - "categories": ["hot"], - "mode": "comments", - "number_posts": 3, - } - - loader.load_data(json.dumps(payload)) - - MockRedditLoader.assert_called_once_with( - client_id="cid", - client_secret="csecret", - user_agent="ua", - categories=["hot"], - mode="comments", - search_queries=["python"], - number_posts=3, - ) - instance.load.assert_called_once() + reddit.redditor.assert_called_with("python") + assert len(docs) == 1 and docs[0].text == "Body two" + def test_number_posts_and_categories_are_forwarded(self): + payload = {**BASE_PAYLOAD, "categories": ["hot"], "number_posts": 3} + praw, _, listing = _fake_praw([]) + with patch.dict(sys.modules, {"praw": praw}): + RedditPostsLoaderRemote().load_data(json.dumps(payload)) + listing.hot.assert_called_once_with(limit=3) + listing.new.assert_not_called() diff --git a/tests/parser/remote/test_web_loader.py b/tests/parser/remote/test_web_loader.py index 51fa8841..758778a6 100644 --- a/tests/parser/remote/test_web_loader.py +++ b/tests/parser/remote/test_web_loader.py @@ -5,7 +5,7 @@ from urllib.parse import urlparse from application.core.url_validation import SSRFError from application.parser.remote.web_loader import WebLoader, headers from application.parser.schema.base import Document -from langchain_core.documents import Document as LCDocument +from application.vectorstore.document_class import Document as LCDocument def _mock_validate_url(url): @@ -301,13 +301,13 @@ class TestWebLoaderIntegration: @patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url) @patch("application.parser.remote.web_loader.pinned_request") - def test_load_langchain_documents_method(self, mock_pinned_request, mock_validate, web_loader): + def test_load_vector_documents_method(self, mock_pinned_request, mock_validate, web_loader): mock_pinned_request.return_value = _fake_response( "Test Page" "

Test web page content

" ) - result = web_loader.load_langchain_documents(inputs="https://example.com") + result = web_loader.load_vector_documents(inputs="https://example.com") assert len(result) == 1 assert isinstance(result[0], LCDocument) diff --git a/tests/parser/test_chunking_strategies.py b/tests/parser/test_chunking_strategies.py index 4fd53fe9..62973a1d 100644 --- a/tests/parser/test_chunking_strategies.py +++ b/tests/parser/test_chunking_strategies.py @@ -115,7 +115,7 @@ class TestParentChild: def test_parent_text_reaches_vectorstore_metadata(self): chunker = ParentChildChunker(max_tokens=80, min_tokens=20) out = chunker.chunk([Document(text="beta " * 150, doc_id="d")]) - lc = out[0].to_langchain_format() + lc = out[0].to_vector_format() # parent_text must survive the langchain conversion into metadata. assert "parent_text" in lc.metadata assert lc.metadata["parent_text"] diff --git a/tests/parser/test_schema.py b/tests/parser/test_schema.py index b6fdf7aa..9f492657 100644 --- a/tests/parser/test_schema.py +++ b/tests/parser/test_schema.py @@ -73,30 +73,30 @@ class TestDocumentBase: with pytest.raises(ValueError, match="text field not set"): Document(text=None) - def test_document_to_langchain_format(self): - """Cover line 27: Document.to_langchain_format converts correctly.""" + def test_document_to_vector_format(self): + """Cover line 27: Document.to_vector_format converts correctly.""" from application.parser.schema.base import Document doc = Document(text="hello world", extra_info={"source": "test"}) - lc_doc = doc.to_langchain_format() + lc_doc = doc.to_vector_format() assert lc_doc.page_content == "hello world" assert lc_doc.metadata == {"source": "test"} - def test_document_to_langchain_format_no_extra_info(self): - """Cover: to_langchain_format with no extra_info uses empty dict.""" + def test_document_to_vector_format_no_extra_info(self): + """Cover: to_vector_format with no extra_info uses empty dict.""" from application.parser.schema.base import Document doc = Document(text="hello") - lc_doc = doc.to_langchain_format() + lc_doc = doc.to_vector_format() assert lc_doc.metadata == {} - def test_document_from_langchain_format(self): - """Cover line 34: Document.from_langchain_format creates Document.""" + def test_document_from_vector_format(self): + """Cover line 34: Document.from_vector_format creates Document.""" from application.parser.schema.base import Document - from langchain_core.documents import Document as LCDocument + from application.vectorstore.document_class import Document as LCDocument lc_doc = LCDocument(page_content="test content", metadata={"key": "val"}) - doc = Document.from_langchain_format(lc_doc) + doc = Document.from_vector_format(lc_doc) assert doc.text == "test content" assert doc.extra_info == {"key": "val"} diff --git a/tests/test_coverage_gaps.py b/tests/test_coverage_gaps.py index 31bddd14..d3aa3028 100644 --- a/tests/test_coverage_gaps.py +++ b/tests/test_coverage_gaps.py @@ -2031,70 +2031,54 @@ class TestInternalRoutes: # --------------------------------------------------------------------------- @pytest.mark.unit class TestFaissStore: - def test_faiss_init_load_from_storage(self): - mock_emb = MagicMock() - mock_storage = MagicMock() - mock_storage.file_exists.return_value = True - faiss_data = b"faiss_data" - pkl_data = b"pkl_data" - mock_storage.get_file.side_effect = [io.BytesIO(faiss_data), io.BytesIO(pkl_data)] + """Storage round-trip against a real FAISS index and real local storage.""" - mock_faiss_class = MagicMock() - mock_faiss_class.load_local.return_value = MagicMock() + class _Embeddings: + dimension = 3 + + def embed_documents(self, documents): + return [[0.1, 0.2, 0.3] for _ in documents] + + def embed_query(self, query): + return [0.1, 0.2, 0.3] + + class _Doc: + def __init__(self, page_content, metadata): + self.page_content = page_content + self.metadata = metadata + + def _make(self, storage, source_id="test", docs_init=None): + from application.vectorstore.faiss import FaissStore with patch( "application.vectorstore.base.BaseVectorStore._get_embeddings", - return_value=mock_emb, + return_value=self._Embeddings(), ), patch( "application.vectorstore.faiss.StorageCreator.get_storage", - return_value=mock_storage, - ), patch( - "application.vectorstore.faiss.FAISS", mock_faiss_class, - ), patch( - "application.vectorstore.faiss.settings" - ) as ms: + return_value=storage, + ), patch("application.vectorstore.faiss.settings") as ms: ms.EMBEDDINGS_NAME = "test" - from application.vectorstore.faiss import FaissStore + return FaissStore(source_id, "key", docs_init=docs_init) - store = FaissStore(source_id="test", embeddings_key="key") - assert store.docsearch is not None + def test_faiss_save_to_storage_writes_all_three_files(self, tmp_path): + from application.storage.local import LocalStorage - def test_faiss_save_to_storage(self): - mock_emb = MagicMock() - mock_storage = MagicMock() - mock_docsearch = MagicMock() + storage = LocalStorage(base_dir=str(tmp_path)) + store = self._make(storage, docs_init=[self._Doc("hello", {"source": "a"})]) - with patch( - "application.vectorstore.base.BaseVectorStore._get_embeddings", - return_value=mock_emb, - ), patch( - "application.vectorstore.faiss.StorageCreator.get_storage", - return_value=mock_storage, - ), patch( - "application.vectorstore.faiss.settings" - ) as ms: - ms.EMBEDDINGS_NAME = "test" - from application.vectorstore.faiss import FaissStore + assert store._save_to_storage() is True + for name in ("index.faiss", "index.json", "index.pkl"): + assert storage.file_exists(f"indexes/test/{name}"), name - store = FaissStore.__new__(FaissStore) - store.source_id = "test" - store.path = "indexes/test" - store.embeddings = mock_emb - store.storage = mock_storage - store.docsearch = mock_docsearch + def test_faiss_init_load_from_storage(self, tmp_path): + from application.storage.local import LocalStorage - def fake_save_local(temp_dir): - os.makedirs(temp_dir, exist_ok=True) - with open(os.path.join(temp_dir, "index.faiss"), "wb") as f: - f.write(b"faiss") - with open(os.path.join(temp_dir, "index.pkl"), "wb") as f: - f.write(b"pkl") + storage = LocalStorage(base_dir=str(tmp_path)) + self._make(storage, docs_init=[self._Doc("hello", {"source": "a"})]).save_local() - mock_docsearch.save_local.side_effect = fake_save_local - - result = store._save_to_storage() - assert result is True - assert mock_storage.save_file.call_count == 2 + reloaded = self._make(storage) + assert reloaded.index.ntotal == 1 + assert reloaded.get_chunks()[0]["text"] == "hello" # --------------------------------------------------------------------------- @@ -2102,58 +2086,52 @@ class TestFaissStore: # --------------------------------------------------------------------------- @pytest.mark.unit class TestQdrantStoreIndexCreation: - def test_init_index_already_exists_error(self): - mock_models = MagicMock() - mock_qdrant_langchain = MagicMock() + """A concurrent worker may create the collection or index first.""" - with patch( - "application.vectorstore.base.BaseVectorStore._get_embeddings" - ) as mock_get_emb, patch( - "application.vectorstore.qdrant.settings" - ) as mock_settings, patch.dict( - "sys.modules", - { - "qdrant_client": MagicMock(), - "qdrant_client.models": mock_models, - "langchain_community": MagicMock(), - "langchain_community.vectorstores": MagicMock(), - "langchain_community.vectorstores.qdrant": mock_qdrant_langchain, - }, + @staticmethod + def _settings(mock_settings): + mock_settings.EMBEDDINGS_NAME = "test" + mock_settings.QDRANT_COLLECTION_NAME = "coll" + mock_settings.QDRANT_LOCATION = ":memory:" + mock_settings.QDRANT_DISTANCE_FUNC = "Cosine" + mock_settings.QDRANT_PREFER_GRPC = False + mock_settings.QDRANT_GRPC_PORT = 6334 + for unset in ( + "QDRANT_URL", "QDRANT_HOST", "QDRANT_PORT", "QDRANT_HTTPS", + "QDRANT_API_KEY", "QDRANT_PREFIX", "QDRANT_TIMEOUT", "QDRANT_PATH", ): - mock_emb = Mock() - mock_emb.client = [None, Mock(word_embedding_dimension=768)] - mock_get_emb.return_value = mock_emb + setattr(mock_settings, unset, None) - mock_settings.EMBEDDINGS_NAME = "test" - mock_settings.QDRANT_COLLECTION_NAME = "coll" - mock_settings.QDRANT_LOCATION = ":memory:" - mock_settings.QDRANT_URL = None - mock_settings.QDRANT_PORT = 6333 - mock_settings.QDRANT_GRPC_PORT = 6334 - mock_settings.QDRANT_HTTPS = False - mock_settings.QDRANT_PREFER_GRPC = False - mock_settings.QDRANT_API_KEY = None - mock_settings.QDRANT_PREFIX = None - mock_settings.QDRANT_TIMEOUT = None - mock_settings.QDRANT_PATH = None - mock_settings.QDRANT_DISTANCE_FUNC = "Cosine" - - mock_docsearch = MagicMock() - mock_docsearch.client.get_collections.return_value.collections = [ - MagicMock(name="coll") - ] - # Index creation error with "already exists" - mock_docsearch.client.create_payload_index.side_effect = Exception( - "Index already exists" - ) - mock_qdrant_langchain.Qdrant.construct_instance.return_value = ( - mock_docsearch - ) - - from application.vectorstore.qdrant import QdrantStore + def test_init_swallows_already_exists_error(self): + from application.vectorstore.qdrant import QdrantStore + emb = Mock(dimension=8) + with patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=emb, + ), patch("application.vectorstore.qdrant.settings") as mock_settings, patch( + "qdrant_client.QdrantClient.create_payload_index", + side_effect=Exception("Index already exists"), + ): + self._settings(mock_settings) store = QdrantStore(source_id="test", embeddings_key="key") - assert store._docsearch is mock_docsearch + assert store._source_id == "test" + + def test_init_logs_other_setup_errors(self, caplog): + from application.vectorstore.qdrant import QdrantStore + + emb = Mock(dimension=8) + with patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=emb, + ), patch("application.vectorstore.qdrant.settings") as mock_settings, patch( + "qdrant_client.QdrantClient.create_payload_index", + side_effect=Exception("connection refused"), + ): + self._settings(mock_settings) + store = QdrantStore(source_id="test", embeddings_key="key") + assert store._source_id == "test" + assert "connection refused" in caplog.text # --------------------------------------------------------------------------- diff --git a/tests/test_remaining_coverage.py b/tests/test_remaining_coverage.py index 06e6ffe8..61ed844b 100644 --- a/tests/test_remaining_coverage.py +++ b/tests/test_remaining_coverage.py @@ -129,9 +129,9 @@ class TestToolABC: # application/parser/file/base.py (lines 18-19) # --------------------------------------------------------------------------- @pytest.mark.unit -class TestBaseReaderLoadLangchain: - def test_load_langchain_documents(self): - """Cover lines 18-19: BaseReader.load_langchain_documents.""" +class TestBaseReaderLoadVectorDocuments: + def test_load_vector_documents(self): + """Cover lines 18-19: BaseReader.load_vector_documents.""" from application.parser.file.base import BaseReader from application.parser.schema.base import Document @@ -143,7 +143,7 @@ class TestBaseReaderLoadLangchain: ] reader = ConcreteReader() - lc_docs = reader.load_langchain_documents() + lc_docs = reader.load_vector_documents() assert len(lc_docs) == 2 assert lc_docs[0].page_content == "hello" assert lc_docs[0].metadata == {"k": "v"} diff --git a/tests/vectorstore/test_faiss.py b/tests/vectorstore/test_faiss.py index b8840de9..d4228c47 100644 --- a/tests/vectorstore/test_faiss.py +++ b/tests/vectorstore/test_faiss.py @@ -1,628 +1,228 @@ -import io +"""FaissStore tests, run against a real FAISS index and real local storage. + +The store no longer wraps langchain, so these drive the actual index rather +than asserting that calls were forwarded to a mock. +""" + +import json from unittest.mock import Mock, patch import pytest - -@pytest.fixture -def mock_embeddings(): - emb = Mock() - emb.embed_query = Mock(return_value=[0.1, 0.2, 0.3]) - emb.embed_documents = Mock(return_value=[[0.1, 0.2, 0.3]]) - emb.dimension = 3 - return emb +from application.storage.local import LocalStorage -@pytest.fixture -def mock_storage(): - storage = Mock() - storage.file_exists = Mock(return_value=True) - storage.get_file = Mock(return_value=io.BytesIO(b"fake data")) - storage.save_file = Mock() - return storage +class _FakeEmbeddings: + """Deterministic 3-dim embeddings: distinct texts get distinct directions.""" + dimension = 3 -@pytest.fixture -def mock_docsearch(): - ds = Mock() - ds.similarity_search = Mock(return_value=[]) - ds.add_texts = Mock(return_value=["id1"]) - ds.add_documents = Mock(return_value=["id1"]) - ds.save_local = Mock() - ds.delete = Mock() - ds.index = Mock() - ds.index.d = 3 - ds.docstore = Mock() - ds.docstore._dict = { - "doc1": Mock(page_content="text1", metadata={"source": "a"}), - "doc2": Mock(page_content="text2", metadata={"source": "b"}), + _VECTORS = { + "paris": [1.0, 0.0, 0.0], + "database": [0.0, 1.0, 0.0], + "celery": [0.0, 0.0, 1.0], } - return ds + + def _vector(self, text): + lowered = (text or "").lower() + for keyword, vector in self._VECTORS.items(): + if keyword in lowered: + return vector + return [0.5, 0.5, 0.5] + + def embed_query(self, query): + return self._vector(query) + + def embed_documents(self, documents): + return [self._vector(d) for d in documents] + + +class _SeedDoc: + def __init__(self, page_content, metadata): + self.page_content = page_content + self.metadata = metadata + + +@pytest.fixture +def storage(tmp_path): + return LocalStorage(base_dir=str(tmp_path)) + + +@pytest.fixture +def make_store(storage): + from application.vectorstore.faiss import FaissStore + + def _make(source_id="src", docs_init=None): + with patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=_FakeEmbeddings(), + ), patch( + "application.vectorstore.faiss.StorageCreator.get_storage", + return_value=storage, + ), patch("application.vectorstore.faiss.settings") as mock_settings: + mock_settings.EMBEDDINGS_NAME = "test_model" + return FaissStore(source_id, "key", docs_init=docs_init) + + return _make + + +@pytest.fixture +def populated(make_store): + store = make_store( + docs_init=[ + _SeedDoc("The capital of France is Paris.", {"source": "geo.txt"}), + _SeedDoc("Postgres is a relational database.", {"source": "db.txt"}), + _SeedDoc("Celery runs background tasks.", {"source": "queue.txt"}), + ] + ) + store.save_local() + return store @pytest.mark.unit -class TestFaissStoreInit: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_init_with_docs(self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() +class TestFaissStore: + def test_build_from_documents(self, populated): + assert populated.index.ntotal == 3 + assert len(populated.get_chunks()) == 3 - from application.vectorstore.faiss import FaissStore + def test_search_ranks_by_similarity(self, populated): + hits = populated.search("Tell me about Paris", k=2) + assert "Paris" in str(hits[0]) + assert hits[0].metadata["source"] == "geo.txt" - store = FaissStore(source_id="test", embeddings_key="key", docs_init=[Mock()]) - mock_faiss.from_documents.assert_called_once() - assert store.docsearch is mock_ds + def test_search_with_scores_reports_l2(self, populated): + scored = populated.search_with_scores("Tell me about Paris", k=3) + assert populated.score_kind == "l2_distance" + # L2 is lower-is-better, so the ranking must be ascending. + assert [s for _, s in scored] == sorted(s for _, s in scored) + assert scored[0][1] == pytest.approx(0.0, abs=1e-4) - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_init_missing_index_files( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_storage = Mock() - mock_storage.file_exists.return_value = False - mock_storage_creator.get_storage.return_value = mock_storage + def test_search_drops_score_threshold(self, populated): + """FAISS has no threshold knob; the kwarg must be ignored, not raise.""" + assert populated.search("Paris", k=1, score_threshold=0.9) - from application.vectorstore.faiss import FaissStore + def test_search_on_empty_index_returns_empty(self, make_store): + store = make_store(docs_init=[_SeedDoc("only doc", {})]) + store.delete_index() + assert store.search("anything") == [] + def test_add_texts_appends(self, populated): + ids = populated.add_texts(["Redis caches things."], [{"source": "cache.txt"}]) + assert len(ids) == 1 + assert populated.index.ntotal == 4 + assert len(populated.get_chunks()) == 4 + + def test_add_texts_empty_is_noop(self, populated): + assert populated.add_texts([], []) == [] + assert populated.index.ntotal == 3 + + def test_add_and_delete_chunk_roundtrip(self, populated): + chunk_id = populated.add_chunk("Redis caches things.", {"source": "cache.txt"}) + assert len(populated.get_chunks()) == 4 + populated.delete_chunk(chunk_id) + assert len(populated.get_chunks()) == 3 + assert populated.index.ntotal == 3 + # The index must stay searchable after a removal renumbers the mapping. + assert populated.search("Paris", k=1) + + def test_delete_index_with_unknown_id_raises(self, populated): + with pytest.raises(ValueError, match="not found in index"): + populated.delete_index(["nope"]) + + def test_delete_index_without_ids_clears(self, populated): + populated.delete_index() + assert populated.get_chunks() == [] + assert populated.index.ntotal == 0 + + def test_get_chunks_shape(self, populated): + chunk = populated.get_chunks()[0] + assert set(chunk) == {"doc_id", "text", "metadata"} + + +@pytest.mark.unit +class TestFaissPersistence: + def test_save_writes_both_sidecars(self, populated, storage, tmp_path): + for name in ("index.faiss", "index.json", "index.pkl"): + assert storage.file_exists(f"indexes/src/{name}"), name + + def test_reload_prefers_json_sidecar(self, populated, make_store, storage): + reloaded = make_store() + assert len(reloaded.get_chunks()) == 3 + assert reloaded.index.ntotal == 3 + assert "Paris" in str(reloaded.search("Paris", k=1)[0]) + + def test_reload_falls_back_to_legacy_pickle(self, populated, make_store, storage, tmp_path): + (tmp_path / "indexes" / "src" / "index.json").unlink() + reloaded = make_store() + assert len(reloaded.get_chunks()) == 3 + assert "Paris" in str(reloaded.search("Paris", k=1)[0]) + + def test_missing_index_raises(self, make_store): with pytest.raises(Exception, match="Error loading FAISS index"): - FaissStore(source_id="test", embeddings_key="key") + make_store(source_id="never-written") + def test_save_local_writes_to_path(self, populated, tmp_path): + target = tmp_path / "exported" + populated.save_local(str(target)) + assert {p.name for p in target.iterdir()} == { + "index.faiss", "index.json", "index.pkl" + } -@pytest.mark.unit -class TestFaissStoreSearch: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_search_delegates_to_docsearch( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_ds.similarity_search.return_value = ["doc1"] - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - result = store.search("query", k=5) - mock_ds.similarity_search.assert_called_once_with("query", k=5) - assert result == ["doc1"] - - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_search_ignores_score_threshold( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - # FAISS has no relevance-threshold knob; the per-source score_threshold - # must be safely dropped, not forwarded (which would crash langchain). - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_get_emb.return_value = Mock(dimension=3) - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_ds.similarity_search.return_value = ["doc1"] - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - result = store.search("query", k=5, score_threshold=0.9) - # score_threshold is stripped before the forward. - mock_ds.similarity_search.assert_called_once_with("query", k=5) - assert result == ["doc1"] - - -@pytest.mark.unit -class TestFaissStoreAddTexts: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_add_texts_delegates( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_ds.add_texts.return_value = ["id1", "id2"] - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - result = store.add_texts(["text1", "text2"]) - assert result == ["id1", "id2"] - - -@pytest.mark.unit -class TestFaissStoreGetChunks: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_get_chunks(self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - - doc1 = Mock(page_content="text1", metadata={"source": "a"}) - doc2 = Mock(page_content="text2", metadata={"source": "b"}) - - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_ds.docstore._dict = {"id1": doc1, "id2": doc2} - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - chunks = store.get_chunks() - - assert len(chunks) == 2 - texts = {c["text"] for c in chunks} - assert texts == {"text1", "text2"} - - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_get_chunks_empty(self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_ds.docstore._dict = {} - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - assert store.get_chunks() == [] - - -@pytest.mark.unit -class TestFaissStoreSaveLocal: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_save_local_with_path( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_faiss.from_documents.return_value = mock_ds - mock_storage = Mock() - mock_storage_creator.get_storage.return_value = mock_storage - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - - # Mock _save_to_storage to avoid file I/O - store._save_to_storage = Mock(return_value=True) - - with patch("os.makedirs"): - result = store.save_local(path="/tmp/test_save") - - mock_ds.save_local.assert_called_once_with("/tmp/test_save") - store._save_to_storage.assert_called_once() - assert result is True - - -@pytest.mark.unit -class TestFaissStoreDeleteIndex: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_delete_index_delegates( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - store.delete_index(["id1"]) - mock_ds.delete.assert_called_once_with(["id1"]) + def test_json_sidecar_is_readable_json(self, populated, tmp_path): + payload = json.loads((tmp_path / "indexes" / "src" / "index.json").read_text()) + assert payload["version"] == 1 + assert len(payload["documents"]) == 3 + assert len(payload["index_to_docstore_id"]) == 3 @pytest.mark.unit class TestFaissStoreAssertEmbeddingDimensions: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_dimension_mismatch_raises( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = ( - "huggingface_sentence-transformers/all-mpnet-base-v2" - ) - mock_emb = Mock(dimension=768) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=512) # Mismatched dimension - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() + def test_dimension_mismatch_raises(self, populated): + with patch("application.vectorstore.faiss.settings") as mock_settings: + mock_settings.EMBEDDINGS_NAME = ( + "huggingface_sentence-transformers/all-mpnet-base-v2" + ) + with pytest.raises(ValueError, match="Embedding dimension mismatch"): + populated.assert_embedding_dimensions(Mock(dimension=768)) - from application.vectorstore.faiss import FaissStore + def test_missing_dimension_attr_raises(self, populated): + with patch("application.vectorstore.faiss.settings") as mock_settings: + mock_settings.EMBEDDINGS_NAME = ( + "huggingface_sentence-transformers/all-mpnet-base-v2" + ) + embeddings = Mock() + del embeddings.dimension + with pytest.raises(AttributeError, match="'dimension' attribute not found"): + populated.assert_embedding_dimensions(embeddings) - with pytest.raises(ValueError, match="Embedding dimension mismatch"): - FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) + def test_dimension_match_passes(self, populated): + with patch("application.vectorstore.faiss.settings") as mock_settings: + mock_settings.EMBEDDINGS_NAME = ( + "huggingface_sentence-transformers/all-mpnet-base-v2" + ) + assert populated.assert_embedding_dimensions(Mock(dimension=3)) is None - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_missing_dimension_attr_raises( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = ( - "huggingface_sentence-transformers/all-mpnet-base-v2" - ) - mock_emb = Mock(spec=[]) # No dimension attribute - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=768) - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - with pytest.raises(AttributeError, match="dimension"): - FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - - -@pytest.mark.unit -class TestFaissStoreDeleteChunk: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_delete_chunk(self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_faiss.from_documents.return_value = mock_ds - mock_storage = Mock() - mock_storage_creator.get_storage.return_value = mock_storage - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - store._save_to_storage = Mock(return_value=True) - - result = store.delete_chunk("chunk_id") - mock_ds.delete.assert_called_once_with(["chunk_id"]) - store._save_to_storage.assert_called_once() - assert result is True + def test_non_huggingface_skips_dimension_check(self, populated): + with patch("application.vectorstore.faiss.settings") as mock_settings: + mock_settings.EMBEDDINGS_NAME = "openai_text-embedding-ada-002" + assert populated.assert_embedding_dimensions(Mock(dimension=1536)) is None @pytest.mark.unit class TestGetVectorstore: - def test_with_path(self): - from application.vectorstore.faiss import get_vectorstore - - assert get_vectorstore("abc123") == "indexes/abc123" - - def test_without_path(self): + def test_empty_path_returns_base(self): from application.vectorstore.faiss import get_vectorstore assert get_vectorstore("") == "indexes" - assert get_vectorstore(None) == "indexes" - def test_with_nested_path(self): + def test_normal_path(self): from application.vectorstore.faiss import get_vectorstore - assert get_vectorstore("user/source123") == "indexes/user/source123" + assert get_vectorstore("abc") == "indexes/abc" - @pytest.mark.parametrize( - "malicious_path", - [ - "../outside", - "../../etc/passwd", - "nested/../../../outside", - "/tmp/evil", - "..\\outside", - "valid/../../escape", - ], - ) - def test_rejects_path_traversal(self, malicious_path): + @pytest.mark.parametrize("bad", ["../etc", "..\\etc", "a/../../b"]) + def test_traversal_rejected(self, bad): from application.vectorstore.faiss import get_vectorstore with pytest.raises(ValueError, match="Invalid source_id path"): - get_vectorstore(malicious_path) - - def test_allows_mongodb_style_ids(self): - from application.vectorstore.faiss import get_vectorstore - - assert get_vectorstore("65e8f6a8a7a96b1bdad4154f") == "indexes/65e8f6a8a7a96b1bdad4154f" - - -@pytest.mark.unit -class TestFaissStoreAddChunk: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_add_chunk_with_metadata( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_ds.add_documents.return_value = ["new_id"] - mock_faiss.from_documents.return_value = mock_ds - mock_storage = Mock() - mock_storage_creator.get_storage.return_value = mock_storage - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - store._save_to_storage = Mock(return_value=True) - - doc_id = store.add_chunk("new text", metadata={"source": "test"}) - - assert doc_id == ["new_id"] - mock_ds.add_documents.assert_called_once() - store._save_to_storage.assert_called_once() - - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_add_chunk_default_metadata( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_ds.add_documents.return_value = ["new_id"] - mock_faiss.from_documents.return_value = mock_ds - mock_storage = Mock() - mock_storage_creator.get_storage.return_value = mock_storage - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - store._save_to_storage = Mock(return_value=True) - - doc_id = store.add_chunk("new text") - - assert doc_id == ["new_id"] - - -@pytest.mark.unit -class TestFaissStoreSaveLocalNoPath: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_save_local_without_path( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_emb = Mock(dimension=3) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=3) - mock_faiss.from_documents.return_value = mock_ds - mock_storage = Mock() - mock_storage_creator.get_storage.return_value = mock_storage - - from application.vectorstore.faiss import FaissStore - - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - store._save_to_storage = Mock(return_value=True) - - result = store.save_local() - - # Should NOT call docsearch.save_local with a path - mock_ds.save_local.assert_not_called() - store._save_to_storage.assert_called_once() - assert result is True - - -@pytest.mark.unit -class TestFaissStoreAssertEmbeddingDimensionsMatch: - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_dimension_match_passes( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = ( - "huggingface_sentence-transformers/all-mpnet-base-v2" - ) - mock_emb = Mock(dimension=768) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=768) # Matching dimension - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - # Should not raise - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - assert store is not None - - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_non_huggingface_skips_dimension_check( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - mock_settings.EMBEDDINGS_NAME = "openai_text-embedding-ada-002" - mock_emb = Mock(dimension=1536) - mock_get_emb.return_value = mock_emb - mock_ds = Mock() - mock_ds.index = Mock(d=999) # Mismatched but doesn't matter - mock_faiss.from_documents.return_value = mock_ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - # Should not raise since embedding name is not the huggingface one - store = FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - assert store is not None - - -@pytest.mark.unit -class TestFaissSearchWithScores: - @staticmethod - def _build(mock_faiss, mock_get_emb, mock_settings, mock_storage_creator, ds): - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_get_emb.return_value = Mock(dimension=3) - mock_faiss.from_documents.return_value = ds - mock_storage_creator.get_storage.return_value = Mock() - - from application.vectorstore.faiss import FaissStore - - return FaissStore(source_id="t", embeddings_key="k", docs_init=[Mock()]) - - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_reports_l2_distance_and_drops_threshold( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - doc = Mock(page_content="text1", metadata={"source": "a"}) - ds = Mock(index=Mock(d=3)) - ds.similarity_search_with_score.return_value = [(doc, 0.42)] - store = self._build( - mock_faiss, mock_get_emb, mock_settings, mock_storage_creator, ds - ) - - results = store.search_with_scores("query", k=5, score_threshold=0.9) - - # LangChain's FAISS ranks by L2 distance, NOT cosine similarity — the - # label must say so, and score_threshold must be dropped (FAISS has no - # such knob and would crash on it). - assert store.score_kind == "l2_distance" - ds.similarity_search_with_score.assert_called_once_with("query", k=5) - assert results == [(doc, 0.42)] - - @patch("application.vectorstore.faiss.StorageCreator") - @patch("application.vectorstore.faiss.FAISS") - @patch.object( - __import__("application.vectorstore.base", fromlist=["BaseVectorStore"]).BaseVectorStore, - "_get_embeddings", - ) - @patch("application.vectorstore.faiss.settings") - def test_does_not_mutate_docstore_metadata( - self, mock_settings, mock_get_emb, mock_faiss, mock_storage_creator - ): - # similarity_search_with_score hands back the LIVE docstore Documents; - # writing a score into their metadata would pollute the in-memory index - # and could be persisted back to storage by a later add_chunk. - doc = Mock(page_content="text1", metadata={"source": "a"}) - ds = Mock(index=Mock(d=3)) - ds.similarity_search_with_score.return_value = [(doc, 0.42)] - store = self._build( - mock_faiss, mock_get_emb, mock_settings, mock_storage_creator, ds - ) - - store.search_with_scores("query", k=1) - - assert doc.metadata == {"source": "a"} + get_vectorstore(bad) diff --git a/tests/vectorstore/test_faiss_docstore.py b/tests/vectorstore/test_faiss_docstore.py new file mode 100644 index 00000000..1723764f --- /dev/null +++ b/tests/vectorstore/test_faiss_docstore.py @@ -0,0 +1,191 @@ +"""Tests for the FAISS sidecar formats and the langchain compatibility layer. + +Two properties matter here and neither can regress quietly: + +1. Indexes written by langchain must stay readable forever — ``POST + /api/upload_index`` still accepts them, so this is not a one-off migration. +2. The reader must refuse arbitrary pickled globals. The previous + ``pickle.load`` ran whatever an uploaded ``index.pkl`` told it to. +""" + +import io +import pickle +import pickletools + +import pytest + +from application.vectorstore.faiss_docstore import ( + CompatUnpickler, + LegacyDocstore, + LegacyDocument, + dump_json_sidecar, + dump_pickle_sidecar, + load_json_sidecar, + load_pickle_sidecar, +) + +DOCUMENTS = { + "id-1": {"page_content": "The capital of France is Paris.", "metadata": {"source": "geo.txt"}}, + "id-2": {"page_content": "Postgres is a database.", "metadata": {"source": "db.txt"}}, +} +MAPPING = {0: "id-1", 1: "id-2"} + + +def _emitted_symbols(data: bytes): + """Return the (module, name) pairs a pickle payload references.""" + strings, symbols = [], [] + for op, arg, _ in pickletools.genops(io.BytesIO(data)): + if op.name in ("SHORT_BINUNICODE", "BINUNICODE"): + strings.append(arg) + if op.name == "STACK_GLOBAL": + symbols.append((strings[-2], strings[-1])) + return set(symbols) + + +@pytest.mark.unit +class TestPickleSidecar: + def test_round_trip(self): + documents, mapping = load_pickle_sidecar(dump_pickle_sidecar(DOCUMENTS, MAPPING)) + assert documents == DOCUMENTS + assert mapping == MAPPING + + def test_written_pickle_references_langchain_paths(self): + """An older DocsGPT still has langchain, so our output must name it.""" + assert _emitted_symbols(dump_pickle_sidecar(DOCUMENTS, MAPPING)) == { + ("langchain_core.documents.base", "Document"), + ("langchain_community.docstore.in_memory", "InMemoryDocstore"), + } + + def test_dump_does_not_leak_modules(self): + import sys + + dump_pickle_sidecar(DOCUMENTS, MAPPING) + assert not [m for m in sys.modules if m.startswith("langchain")] + assert LegacyDocument.__module__.startswith("application.") + assert LegacyDocstore.__module__.startswith("application.") + + def test_reads_pydantic_state_shape(self): + """langchain's Document is a pydantic model; that is the state we get.""" + doc = LegacyDocument.__new__(LegacyDocument) + doc.__setstate__( + { + "__dict__": { + "id": None, + "metadata": {"a": 1}, + "page_content": "hello", + "type": "Document", + }, + "__pydantic_extra__": None, + "__pydantic_fields_set__": {"page_content"}, + "__pydantic_private__": None, + } + ) + assert doc.page_content == "hello" + assert doc.metadata == {"a": 1} + + def test_missing_metadata_becomes_empty_dict(self): + doc = LegacyDocument.__new__(LegacyDocument) + doc.__setstate__({"__dict__": {"page_content": "x", "metadata": None}}) + assert doc.metadata == {} + + @pytest.mark.parametrize( + "module,name", + [ + ("langchain_core.documents.base", "Document"), + ("langchain_core.documents", "Document"), + ("langchain.schema.document", "Document"), + ("langchain.docstore.document", "Document"), + ("langchain_community.docstore.in_memory", "InMemoryDocstore"), + ("langchain.docstore.in_memory", "InMemoryDocstore"), + ], + ) + def test_historical_module_paths_resolve(self, module, name): + """Older langchain releases wrote these classes at different paths.""" + assert CompatUnpickler(io.BytesIO(b"")).find_class(module, name) is not None + + +@pytest.mark.unit +class TestUnpicklerRefusesArbitraryCode: + def test_refuses_os_system(self): + class Exploit: + def __reduce__(self): + import os + + return (os.system, ("echo pwned",)) + + payload = pickle.dumps(Exploit()) + with pytest.raises(pickle.UnpicklingError, match="Refusing to load"): + CompatUnpickler(io.BytesIO(payload)).load() + + def test_refuses_subprocess(self): + class Exploit: + def __reduce__(self): + import subprocess + + return (subprocess.Popen, (["true"],)) + + with pytest.raises(pickle.UnpicklingError, match="Refusing to load"): + CompatUnpickler(io.BytesIO(pickle.dumps(Exploit()))).load() + + def test_refuses_unrelated_builtin(self): + with pytest.raises(pickle.UnpicklingError, match="Refusing to load"): + CompatUnpickler(io.BytesIO(b"")).find_class("builtins", "eval") + + def test_plain_unpickler_would_have_resolved_it(self): + """Confirms the guard is load-bearing, not decorative. + + A stock Unpickler hands back ``os.system`` for the same input the + CompatUnpickler rejects — that resolution is what made the previous + ``pickle.load`` on uploaded indexes an arbitrary-code-execution path. + """ + import os + + assert pickle.Unpickler(io.BytesIO(b"")).find_class("os", "system") is os.system + with pytest.raises(pickle.UnpicklingError, match="Refusing to load"): + CompatUnpickler(io.BytesIO(b"")).find_class("os", "system") + + +@pytest.mark.unit +class TestJsonSidecar: + def test_round_trip(self): + documents, mapping = load_json_sidecar(dump_json_sidecar(DOCUMENTS, MAPPING)) + assert documents == DOCUMENTS + assert mapping == MAPPING + + def test_contains_no_pickle(self): + assert b"langchain" not in dump_json_sidecar(DOCUMENTS, MAPPING) + + def test_row_keys_come_back_as_ints(self): + _, mapping = load_json_sidecar(dump_json_sidecar(DOCUMENTS, MAPPING)) + assert all(isinstance(row, int) for row in mapping) + + def test_unicode_survives(self): + documents = {"id": {"page_content": "café — naïve 日本語", "metadata": {}}} + restored, _ = load_json_sidecar(dump_json_sidecar(documents, {0: "id"})) + assert restored["id"]["page_content"] == "café — naïve 日本語" + + def test_empty_index(self): + documents, mapping = load_json_sidecar(dump_json_sidecar({}, {})) + assert documents == {} and mapping == {} + + +@pytest.mark.unit +class TestRealLegacyFixture: + """The index.pkl checked into the repo was written by langchain in 2025.""" + + def test_reads_committed_legacy_index(self): + with open("application/index.pkl", "rb") as f: + documents, mapping = load_pickle_sidecar(f.read()) + assert len(documents) == 3 + assert len(mapping) == 3 + assert all(d["page_content"] for d in documents.values()) + assert all(d["metadata"].get("title") for d in documents.values()) + + def test_legacy_survives_conversion_to_json(self): + with open("application/index.pkl", "rb") as f: + documents, mapping = load_pickle_sidecar(f.read()) + restored, restored_mapping = load_json_sidecar( + dump_json_sidecar(documents, mapping) + ) + assert restored == documents + assert restored_mapping == mapping diff --git a/tests/vectorstore/test_milvus.py b/tests/vectorstore/test_milvus.py index e993ea8b..67b6f4c6 100644 --- a/tests/vectorstore/test_milvus.py +++ b/tests/vectorstore/test_milvus.py @@ -1,95 +1,173 @@ -from unittest.mock import MagicMock, Mock, patch -from uuid import UUID +"""MilvusStore tests, run against a real Milvus Lite database. + +pymilvus is a first-party dependency now that the store no longer goes through +langchain, so these drive a real local Milvus (a file under ``tmp_path``) +rather than asserting against mocks. +""" + +import os +from unittest.mock import patch import pytest +pytest.importorskip("milvus_lite", reason="Milvus Lite is not available on this platform") + + +class _FakeEmbeddings: + """Deterministic 3-dim embeddings: distinct texts get distinct directions.""" + + dimension = 3 + + _VECTORS = { + "paris": [1.0, 0.0, 0.0], + "database": [0.0, 1.0, 0.0], + "celery": [0.0, 0.0, 1.0], + } + + def _vector(self, text): + lowered = (text or "").lower() + for keyword, vector in self._VECTORS.items(): + if keyword in lowered: + return vector + return [0.577, 0.577, 0.577] + + def embed_query(self, query): + return self._vector(query) + + def embed_documents(self, documents): + return [self._vector(d) for d in documents] + + +@pytest.fixture +def store(tmp_path): + from application.vectorstore.milvus import MilvusStore -def _make_milvus_store(source_id="test-source"): - """Helper to create a MilvusStore with mocked deps.""" with patch( - "application.vectorstore.base.BaseVectorStore._get_embeddings" - ) as mock_get_emb, patch( - "application.vectorstore.milvus.settings" - ) as mock_settings, patch.dict( - "sys.modules", - { - "langchain_milvus": MagicMock(), - }, - ): - mock_emb = Mock() - mock_get_emb.return_value = mock_emb + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=_FakeEmbeddings(), + ), patch("application.vectorstore.milvus.settings") as mock_settings: mock_settings.EMBEDDINGS_NAME = "test_model" - mock_settings.MILVUS_URI = "http://localhost:19530" - mock_settings.MILVUS_TOKEN = "token" mock_settings.MILVUS_COLLECTION_NAME = "test_collection" + mock_settings.MILVUS_URI = str(tmp_path / "milvus.db") + mock_settings.MILVUS_TOKEN = "" + yield MilvusStore(source_id="src-A", embeddings_key="k") - from langchain_milvus import Milvus - mock_docsearch = MagicMock() - Milvus.return_value = mock_docsearch +@pytest.fixture +def populated(store): + store.add_texts( + ["The capital of France is Paris.", + "Postgres is a relational database.", + "Celery runs background tasks."], + [{"source": "geo.txt"}, {"source": "db.txt"}, {"source": "queue.txt"}], + ) + return store + +@pytest.mark.unit +class TestMilvusStore: + def test_source_id_stored(self, store): + assert store._source_id == "src-A" + + def test_filter_expression_scopes_to_source(self, store): + assert store._filter == 'source_id == "src-A"' + + def test_add_texts_returns_one_id_per_text(self, store): + ids = store.add_texts(["a", "b"], [{}, {}]) + assert len(ids) == 2 and len(set(ids)) == 2 + + def test_add_texts_empty_is_noop(self, store): + assert store.add_texts([], []) == [] + + def test_search_ranks_by_similarity(self, populated): + hits = populated.search("Tell me about Paris", k=2) + assert "Paris" in str(hits[0]) + assert hits[0].metadata["source"] == "geo.txt" + + def test_search_with_scores_reports_cosine(self, populated): + scored = populated.search_with_scores("Tell me about Paris", k=3) + assert populated.score_kind == "cosine_similarity" + assert scored[0][1] == pytest.approx(1.0, abs=1e-3) + + def test_search_honours_score_threshold(self, populated): + assert populated.search_with_scores("Paris", k=3, score_threshold=0.99) + assert not populated.search_with_scores("Paris", k=3, score_threshold=1.01) + + def test_get_chunks_returns_all(self, populated): + chunks = populated.get_chunks() + assert len(chunks) == 3 + assert {c["metadata"]["source"] for c in chunks} == { + "geo.txt", "db.txt", "queue.txt" + } + + def test_add_texts_stamps_source_id(self, populated): + assert all(c["metadata"]["source_id"] == "src-A" for c in populated.get_chunks()) + + def test_add_and_delete_chunk(self, populated): + chunk_id = populated.add_chunk("Redis caches things.", {"source": "cache.txt"}) + assert len(populated.get_chunks()) == 4 + assert populated.delete_chunk(chunk_id) is True + assert len(populated.get_chunks()) == 3 + + def test_delete_chunks_by_source_path(self, populated): + assert populated.delete_chunks_by_source_path("db.txt") == 1 + assert len(populated.get_chunks()) == 2 + + def test_delete_index_removes_only_this_source(self, populated): from application.vectorstore.milvus import MilvusStore - store = MilvusStore(source_id=source_id, embeddings_key="key") - store._docsearch = mock_docsearch + with patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=_FakeEmbeddings(), + ), patch("application.vectorstore.milvus.settings") as mock_settings: + mock_settings.EMBEDDINGS_NAME = "test_model" + mock_settings.MILVUS_COLLECTION_NAME = "test_collection" + mock_settings.MILVUS_URI = ":memory:" + mock_settings.MILVUS_TOKEN = "" + other = MilvusStore.__new__(MilvusStore) + # Share the open Milvus Lite handle: one file, two logical sources. + other._client = populated._client + other._collection = populated._collection + other._source_id = "src-B" + other._embeddings = _FakeEmbeddings() + other.add_texts(["Another source entirely."], [{"source": "other.txt"}]) - return store, mock_docsearch + populated.delete_index() + assert populated.get_chunks() == [] + assert len(other.get_chunks()) == 1 -@pytest.mark.unit -class TestMilvusStoreInit: - def test_source_id_stored(self): - store, _ = _make_milvus_store(source_id="src1") - assert store._source_id == "src1" + def test_get_chunks_returns_empty_when_client_raises(self, populated): + with patch.object( + populated._client, "query", side_effect=RuntimeError("milvus down") + ): + assert populated.get_chunks() == [] - -@pytest.mark.unit -class TestMilvusStoreSearch: - def test_search(self): - store, mock_ds = _make_milvus_store(source_id="src1") - mock_ds.similarity_search.return_value = ["doc1", "doc2"] - - results = store.search("query", k=3) - - mock_ds.similarity_search.assert_called_once() - call_kwargs = mock_ds.similarity_search.call_args - assert call_kwargs[1]["query"] == "query" - assert call_kwargs[1]["k"] == 3 - assert call_kwargs[1]["expr"] == "source_id == 'src1'" - assert results == ["doc1", "doc2"] - - -@pytest.mark.unit -class TestMilvusStoreAddTexts: - def test_add_texts(self): - store, mock_ds = _make_milvus_store() - mock_ds.add_texts.return_value = ["id1", "id2"] - - result = store.add_texts( - ["text1", "text2"], metadatas=[{"a": 1}, {"b": 2}] - ) - - mock_ds.add_texts.assert_called_once() - call_kwargs = mock_ds.add_texts.call_args - assert call_kwargs[1]["texts"] == ["text1", "text2"] - # ids should be UUIDs - ids = call_kwargs[1]["ids"] - assert len(ids) == 2 - for uid in ids: - UUID(uid) # Validates it's a valid UUID - - assert result == ["id1", "id2"] - - -@pytest.mark.unit -class TestMilvusStoreSaveLocal: - def test_save_local_is_noop(self): - store, _ = _make_milvus_store() + def test_save_local_is_noop(self, store): assert store.save_local() is None @pytest.mark.unit -class TestMilvusStoreDeleteIndex: - def test_delete_index_is_noop(self): - store, _ = _make_milvus_store() - assert store.delete_index() is None +class TestMilvusUriEnvGuard: + """pymilvus reads MILVUS_URI at import and rejects non-http values. + + DocsGPT's setting of the same name defaults to a Milvus Lite file path, so + the import must be shielded from it and the value restored afterwards. + """ + + def test_env_var_hidden_during_import_and_restored(self): + from application.vectorstore.milvus import _without_milvus_uri_env + + with patch.dict(os.environ, {"MILVUS_URI": "./milvus_local.db"}): + with _without_milvus_uri_env(): + assert "MILVUS_URI" not in os.environ + assert os.environ["MILVUS_URI"] == "./milvus_local.db" + + def test_absent_env_var_stays_absent(self): + from application.vectorstore.milvus import _without_milvus_uri_env + + env = {k: v for k, v in os.environ.items() if k != "MILVUS_URI"} + with patch.dict(os.environ, env, clear=True): + with _without_milvus_uri_env(): + assert "MILVUS_URI" not in os.environ + assert "MILVUS_URI" not in os.environ diff --git a/tests/vectorstore/test_qdrant.py b/tests/vectorstore/test_qdrant.py index 2deb6a69..16c40f9e 100644 --- a/tests/vectorstore/test_qdrant.py +++ b/tests/vectorstore/test_qdrant.py @@ -1,230 +1,183 @@ -from unittest.mock import MagicMock, Mock, patch +"""QdrantStore tests, run against a real in-process Qdrant. + +qdrant-client is a first-party dependency now that the store no longer goes +through langchain, so these exercise the real client in ``:memory:`` mode +rather than asserting against mocks. +""" + +from unittest.mock import patch import pytest -def _make_qdrant_store(source_id="test-source"): - """Helper to create a QdrantStore with all external deps mocked.""" - mock_models = MagicMock() - mock_qdrant_langchain = MagicMock() +class _FakeEmbeddings: + """Deterministic 3-dim embeddings: distinct texts get distinct directions.""" + + dimension = 3 + + _VECTORS = { + "paris": [1.0, 0.0, 0.0], + "database": [0.0, 1.0, 0.0], + "celery": [0.0, 0.0, 1.0], + } + + def _vector(self, text): + lowered = (text or "").lower() + for keyword, vector in self._VECTORS.items(): + if keyword in lowered: + return vector + return [0.577, 0.577, 0.577] + + def embed_query(self, query): + return self._vector(query) + + def embed_documents(self, documents): + return [self._vector(d) for d in documents] + + +def _settings(mock_settings, collection="test_collection"): + mock_settings.EMBEDDINGS_NAME = "test_model" + mock_settings.QDRANT_COLLECTION_NAME = collection + mock_settings.QDRANT_LOCATION = ":memory:" + mock_settings.QDRANT_DISTANCE_FUNC = "Cosine" + mock_settings.QDRANT_PREFER_GRPC = False + mock_settings.QDRANT_GRPC_PORT = 6334 + for unset in ( + "QDRANT_URL", "QDRANT_HOST", "QDRANT_PORT", "QDRANT_HTTPS", + "QDRANT_API_KEY", "QDRANT_PREFIX", "QDRANT_TIMEOUT", "QDRANT_PATH", + ): + setattr(mock_settings, unset, None) + + +@pytest.fixture +def store(): + from application.vectorstore.qdrant import QdrantStore with patch( - "application.vectorstore.base.BaseVectorStore._get_embeddings" - ) as mock_get_emb, patch( - "application.vectorstore.qdrant.settings" - ) as mock_settings, patch.dict( - "sys.modules", - { - "qdrant_client": MagicMock(), - "qdrant_client.models": mock_models, - "langchain_community": MagicMock(), - "langchain_community.vectorstores": MagicMock(), - "langchain_community.vectorstores.qdrant": mock_qdrant_langchain, - }, - ): - mock_emb = Mock() - mock_emb.embed_query = Mock(return_value=[0.1, 0.2, 0.3]) - mock_emb.embed_documents = Mock(return_value=[[0.1, 0.2, 0.3]]) - mock_emb.client = [None, Mock(word_embedding_dimension=768)] - mock_get_emb.return_value = mock_emb + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=_FakeEmbeddings(), + ), patch("application.vectorstore.qdrant.settings") as mock_settings: + _settings(mock_settings) + yield QdrantStore(source_id="src-A", embeddings_key="k") - mock_settings.EMBEDDINGS_NAME = "test_model" - mock_settings.QDRANT_COLLECTION_NAME = "test_collection" - mock_settings.QDRANT_LOCATION = ":memory:" - mock_settings.QDRANT_URL = None - mock_settings.QDRANT_PORT = 6333 - mock_settings.QDRANT_GRPC_PORT = 6334 - mock_settings.QDRANT_HTTPS = False - mock_settings.QDRANT_PREFER_GRPC = False - mock_settings.QDRANT_API_KEY = None - mock_settings.QDRANT_PREFIX = None - mock_settings.QDRANT_TIMEOUT = None - mock_settings.QDRANT_PATH = None - mock_settings.QDRANT_DISTANCE_FUNC = "Cosine" - mock_docsearch = MagicMock() - mock_collections = MagicMock() - mock_collections.collections = [MagicMock(name="test_collection")] - mock_docsearch.client.get_collections.return_value = mock_collections - mock_qdrant_langchain.Qdrant.construct_instance.return_value = mock_docsearch +@pytest.fixture +def populated(store): + store.add_texts( + ["The capital of France is Paris.", + "Postgres is a relational database.", + "Celery runs background tasks."], + [{"source": "geo.txt"}, {"source": "db.txt"}, {"source": "queue.txt"}], + ) + return store + +@pytest.mark.unit +class TestQdrantStore: + def test_source_id_strips_index_prefix(self, store): + assert store._source_id == "src-A" + + def test_add_texts_returns_one_id_per_text(self, store): + ids = store.add_texts(["a", "b"], [{}, {}]) + assert len(ids) == 2 and len(set(ids)) == 2 + + def test_add_texts_empty_is_noop(self, store): + assert store.add_texts([], []) == [] + + def test_search_ranks_by_similarity(self, populated): + hits = populated.search("Tell me about Paris", k=2) + assert "Paris" in str(hits[0]) + assert hits[0].metadata["source"] == "geo.txt" + + def test_search_with_scores_reports_cosine(self, populated): + scored = populated.search_with_scores("Tell me about Paris", k=3) + assert populated.score_kind == "cosine_similarity" + assert scored[0][1] == pytest.approx(1.0, abs=1e-3) + # Scores must come back in descending rank order. + assert [s for _, s in scored] == sorted( + (s for _, s in scored), reverse=True + ) + + def test_search_honours_score_threshold(self, populated): + assert populated.search_with_scores("Paris", k=3, score_threshold=0.99) + assert not populated.search_with_scores("Paris", k=3, score_threshold=1.01) + + def test_add_texts_stamps_source_id(self, populated): + assert all(c["metadata"]["source_id"] == "src-A" for c in populated.get_chunks()) + + def test_get_chunks_returns_all(self, populated): + chunks = populated.get_chunks() + assert len(chunks) == 3 + assert {c["metadata"]["source"] for c in chunks} == { + "geo.txt", "db.txt", "queue.txt" + } + + def test_add_and_delete_chunk(self, populated): + chunk_id = populated.add_chunk("Redis caches things.", {"source": "cache.txt"}) + assert len(populated.get_chunks()) == 4 + assert populated.delete_chunk(chunk_id) is True + assert len(populated.get_chunks()) == 3 + + def test_delete_chunk_returns_false_when_client_raises(self, populated): + with patch.object( + populated._client, "delete", side_effect=RuntimeError("qdrant down") + ): + assert populated.delete_chunk("some-id") is False + + def test_get_chunks_returns_empty_when_client_raises(self, populated): + with patch.object( + populated._client, "scroll", side_effect=RuntimeError("qdrant down") + ): + assert populated.get_chunks() == [] + + def test_delete_chunks_by_source_path(self, populated): + assert populated.delete_chunks_by_source_path("db.txt") == 1 + assert len(populated.get_chunks()) == 2 + + def test_delete_index_removes_only_this_source(self, populated): from application.vectorstore.qdrant import QdrantStore - store = QdrantStore(source_id=source_id, embeddings_key="key") + with patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=_FakeEmbeddings(), + ), patch("application.vectorstore.qdrant.settings") as mock_settings: + _settings(mock_settings) + other = QdrantStore(source_id="src-B", embeddings_key="k") + # Share the in-memory backend so both sources live in one collection. + other._client = populated._client + other.add_texts(["Another source entirely."], [{"source": "other.txt"}]) - return store, mock_docsearch, mock_settings + populated.delete_index() + assert populated.get_chunks() == [] + assert len(other.get_chunks()) == 1 -@pytest.mark.unit -class TestQdrantStoreInit: - def test_source_id_cleaned(self): - store, _, _ = _make_qdrant_store(source_id="application/indexes/abc123/") - assert store._source_id == "abc123" - - def test_filter_constructed(self): - store, _, _ = _make_qdrant_store(source_id="src1") - assert store._filter is not None - - -@pytest.mark.unit -class TestQdrantStoreSearch: - def test_search_delegates(self): - store, mock_ds, _ = _make_qdrant_store() - mock_ds.similarity_search.return_value = ["result1"] - - results = store.search("query", k=5) - - mock_ds.similarity_search.assert_called_once() - assert results == ["result1"] - - -@pytest.mark.unit -class TestQdrantStoreAddTexts: - def test_add_texts_delegates(self): - store, mock_ds, _ = _make_qdrant_store() - mock_ds.add_texts.return_value = ["id1"] - - result = store.add_texts(["text1"], metadatas=[{"a": 1}]) - mock_ds.add_texts.assert_called_once_with(["text1"], metadatas=[{"a": 1}]) - assert result == ["id1"] - - -@pytest.mark.unit -class TestQdrantStoreSaveLocal: - def test_save_local_is_noop(self): - store, _, _ = _make_qdrant_store() + def test_save_local_is_noop(self, store): assert store.save_local() is None @pytest.mark.unit -class TestQdrantStoreDeleteIndex: - def test_delete_index(self): - store, mock_ds, _ = _make_qdrant_store() +class TestQdrantClientKwargs: + def test_unset_settings_are_omitted(self): + from application.vectorstore.qdrant import QdrantStore - with patch("application.vectorstore.qdrant.settings") as ms: - ms.QDRANT_COLLECTION_NAME = "test_collection" - store.delete_index() + with patch("application.vectorstore.qdrant.settings") as mock_settings: + _settings(mock_settings) + kwargs = QdrantStore._client_kwargs() + # location/url/path are mutually exclusive in qdrant-client, so only + # the configured one may be forwarded. + assert kwargs["location"] == ":memory:" + assert "url" not in kwargs and "path" not in kwargs and "host" not in kwargs - mock_ds.client.delete.assert_called_once_with( - collection_name="test_collection", - points_selector=store._filter, - ) + def test_configured_settings_are_forwarded(self): + from application.vectorstore.qdrant import QdrantStore - -@pytest.mark.unit -class TestQdrantStoreGetChunks: - def test_get_chunks(self): - store, mock_ds, _ = _make_qdrant_store() - - record1 = MagicMock() - record1.id = "id1" - record1.payload = { - "page_content": "text1", - "metadata": {"source": "test"}, - } - record2 = MagicMock() - record2.id = "id2" - record2.payload = { - "page_content": "text2", - "metadata": {"source": "test2"}, - } - - # First call returns records with offset, second returns empty with None offset - mock_ds.client.scroll.side_effect = [ - ([record1, record2], None), - ] - - chunks = store.get_chunks() - - assert len(chunks) == 2 - assert chunks[0] == { - "doc_id": "id1", - "text": "text1", - "metadata": {"source": "test"}, - } - - def test_get_chunks_pagination(self): - store, mock_ds, _ = _make_qdrant_store() - - record1 = MagicMock() - record1.id = "id1" - record1.payload = {"page_content": "text1", "metadata": {}} - - record2 = MagicMock() - record2.id = "id2" - record2.payload = {"page_content": "text2", "metadata": {}} - - mock_ds.client.scroll.side_effect = [ - ([record1], "offset_token"), - ([record2], None), - ] - - chunks = store.get_chunks() - assert len(chunks) == 2 - assert mock_ds.client.scroll.call_count == 2 - - def test_get_chunks_returns_empty_on_error(self): - store, mock_ds, _ = _make_qdrant_store() - mock_ds.client.scroll.side_effect = Exception("fail") - - assert store.get_chunks() == [] - - -@pytest.mark.unit -class TestQdrantStoreAddChunk: - def test_add_chunk(self): - store, mock_ds, _ = _make_qdrant_store(source_id="src1") - mock_ds.add_documents.return_value = ["new-id"] - - result = store.add_chunk("hello", metadata={"key": "val"}) - - assert result == "new-id" - mock_ds.add_documents.assert_called_once() - doc = mock_ds.add_documents.call_args[0][0][0] - assert doc.page_content == "hello" - assert doc.metadata["source_id"] == "src1" - assert doc.metadata["key"] == "val" - - def test_add_chunk_default_metadata(self): - store, mock_ds, _ = _make_qdrant_store(source_id="src1") - mock_ds.add_documents.return_value = ["id"] - - store.add_chunk("text") - - doc = mock_ds.add_documents.call_args[0][0][0] - assert doc.metadata["source_id"] == "src1" - - def test_add_chunk_fallback_id(self): - store, mock_ds, _ = _make_qdrant_store() - mock_ds.add_documents.return_value = [] - - result = store.add_chunk("text") - # Should return the uuid that was generated - assert result is not None - assert isinstance(result, str) - - -@pytest.mark.unit -class TestQdrantStoreDeleteChunk: - def test_delete_chunk_success(self): - store, mock_ds, _ = _make_qdrant_store() - - with patch("application.vectorstore.qdrant.settings") as ms: - ms.QDRANT_COLLECTION_NAME = "test_collection" - result = store.delete_chunk("chunk-id") - - mock_ds.client.delete.assert_called_once_with( - collection_name="test_collection", - points_selector=["chunk-id"], - ) - assert result is True - - def test_delete_chunk_returns_false_on_error(self): - store, mock_ds, _ = _make_qdrant_store() - - with patch("application.vectorstore.qdrant.settings") as ms: - ms.QDRANT_COLLECTION_NAME = "test_collection" - mock_ds.client.delete.side_effect = Exception("fail") - result = store.delete_chunk("bad-id") - - assert result is False + with patch("application.vectorstore.qdrant.settings") as mock_settings: + _settings(mock_settings) + mock_settings.QDRANT_LOCATION = None + mock_settings.QDRANT_URL = "http://qdrant:6333" + mock_settings.QDRANT_API_KEY = "secret" + kwargs = QdrantStore._client_kwargs() + assert kwargs["url"] == "http://qdrant:6333" + assert kwargs["api_key"] == "secret" + assert "location" not in kwargs From 209c29d75924bca9637f898db2b6ab5f93365f6e Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 10 Aug 2026 12:37:40 +0100 Subject: [PATCH 4/4] fix: test --- tests/test_remaining_coverage.py | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/tests/test_remaining_coverage.py b/tests/test_remaining_coverage.py index 61ed844b..ff5befe6 100644 --- a/tests/test_remaining_coverage.py +++ b/tests/test_remaining_coverage.py @@ -1049,19 +1049,31 @@ class TestEmbeddingPipelineCoverage: embed_and_store_documents([], str(tmp_path / "test"), "src-1", None) def test_embed_and_store_creates_folder(self, tmp_path): - """Cover line 65: folder creation.""" + """Cover line 65: the folder is created before the vector store is built. + + The failure is forced through ``VectorCreator`` rather than left to + whatever the store happens to reject — this previously relied on + langchain's ``from_documents`` reading a ``.id`` the stub lacks, so it + stopped failing the moment the FAISS store stopped going via langchain. + """ + import os + from application.parser.embedding_pipeline import embed_and_store_documents folder = str(tmp_path / "new_dir") - with pytest.raises(Exception): - # Will fail at VectorCreator but folder should be created - embed_and_store_documents( - [type("Doc", (), {"page_content": "text", "metadata": {}})()], - folder, - "src-1", - None, + with patch( + "application.parser.embedding_pipeline.VectorCreator" + ) as mock_vc: + mock_vc.create_vectorstore.side_effect = RuntimeError( + "vector store unavailable" ) - import os + with pytest.raises(RuntimeError, match="vector store unavailable"): + embed_and_store_documents( + [type("Doc", (), {"page_content": "text", "metadata": {}})()], + folder, + "src-1", + None, + ) assert os.path.exists(folder)