mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 18:15:02 +00:00
feat: prune some deps
This commit is contained in:
1 parent
c5b372df1e
commit
a166796f6c
26 files changed
+1937
-1343
No files matched your search
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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]
|
||||
@@ -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 {}
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
+235
-103
@@ -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()
|
||||
]
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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 = (
|
||||
"<html><head><title>My Page</title></head>"
|
||||
"<body><h1>Heading</h1><p>Hello world.</p></body></html>"
|
||||
)
|
||||
|
||||
|
||||
@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("<html><body><p>Just text.</p></body></html>")
|
||||
assert HTMLParser().get_file_metadata(path) == {}
|
||||
|
||||
|
||||
def test_html_parser_metadata_unreadable_file(tmp_path):
|
||||
assert HTMLParser().get_file_metadata(tmp_path / "missing.html") == {}
|
||||
@@ -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()
|
||||
@@ -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(
|
||||
"<html><head><title>Test Page</title></head>"
|
||||
"<body><p>Test web page content</p></body></html>"
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
@@ -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"]
|
||||
|
||||
+10
-10
@@ -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"}
|
||||
|
||||
|
||||
+76
-98
@@ -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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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"}
|
||||
|
||||
+187
-587
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
+162
-209
@@ -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
|
||||
Reference in new issue
Block a user