mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 12:13:38 +00:00
refactor: rename the application package to docsgpt
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
This commit is contained in:
Alex committed
1 parent
3e7f1f91b4
commit
574f96341e
984 files changed
+9847
-9682
No files matched your search
Whitespace-only changes.
@@ -0,0 +1,455 @@
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.embeddings_openai import OpenAIEmbeddings
|
||||
from docsgpt.vectorstore.model_registry import (
|
||||
dimension_for,
|
||||
max_input_tokens_for,
|
||||
resolve,
|
||||
)
|
||||
|
||||
|
||||
def _embeddings_name_is_explicit() -> bool:
|
||||
"""True when ``EMBEDDINGS_NAME`` names a model somebody actually chose.
|
||||
|
||||
Not ``model_fields_set``: pydantic marks a field as set for any value that
|
||||
reached it, including one read from ``.env``, and every setup script has
|
||||
always written ``EMBEDDINGS_NAME`` unconditionally. An install carrying the
|
||||
legacy name a script wrote years ago would read as a deliberate choice and
|
||||
lend a remote server that model's context window.
|
||||
"""
|
||||
fields = getattr(type(settings), "model_fields", None)
|
||||
if not isinstance(fields, dict):
|
||||
return False
|
||||
field = fields.get("EMBEDDINGS_NAME")
|
||||
if field is None:
|
||||
return False
|
||||
return settings.EMBEDDINGS_NAME != field.default
|
||||
|
||||
|
||||
class RemoteEmbeddings:
|
||||
"""
|
||||
Wrapper for remote embeddings API (OpenAI-compatible).
|
||||
Used when EMBEDDINGS_BASE_URL is configured.
|
||||
Sends requests to {base_url}/v1/embeddings in OpenAI format.
|
||||
"""
|
||||
|
||||
def __init__(self, api_url: str, model_name: str, api_key: str = None):
|
||||
self.api_url = api_url.rstrip("/")
|
||||
self.model_name = model_name
|
||||
self.headers = {"Content-Type": "application/json"}
|
||||
if api_key:
|
||||
self.headers["Authorization"] = f"Bearer {api_key}"
|
||||
# Width comes from the registry. This used to be a hardcoded 768 that
|
||||
# ``embed_query`` claimed to correct on first use -- but the correction
|
||||
# was guarded by ``if self.dimension is None``, which the hardcode made
|
||||
# unreachable, so a remote model of any other width silently produced a
|
||||
# ``vector(768)`` column. ``None`` here means "unknown", and the probe
|
||||
# below now genuinely runs.
|
||||
self.dimension = dimension_for(model_name)
|
||||
|
||||
def _token_counter(self):
|
||||
"""Counter matching the remote model's tokenizer, cached per process."""
|
||||
from docsgpt.parser.tokenization import get_token_counter
|
||||
|
||||
return get_token_counter(self.model_name)
|
||||
|
||||
def _resolve_input_limit(self):
|
||||
"""Token ceiling for a single embed input, or ``None`` for no limit.
|
||||
|
||||
``EMBEDDINGS_MAX_INPUT_TOKENS`` wins when set. Otherwise a registered
|
||||
model contributes its own context window, so a request that the server
|
||||
would reject -- or silently truncate -- is clipped here instead of
|
||||
being sent and paid for.
|
||||
|
||||
That fallback needs the name to mean something. For a remote server it
|
||||
is only a label forwarded as the ``model`` field, so the settings
|
||||
default must not lend the server mpnet's 384-token window: a name
|
||||
nobody chose describes nothing, and clipping on it would silently
|
||||
discard most of every chunk.
|
||||
"""
|
||||
configured = settings.EMBEDDINGS_MAX_INPUT_TOKENS
|
||||
if configured and configured > 0:
|
||||
return configured
|
||||
if not _embeddings_name_is_explicit():
|
||||
return None
|
||||
model_limit = max_input_tokens_for(self.model_name)
|
||||
return model_limit if model_limit and model_limit > 0 else None
|
||||
|
||||
def _truncate_inputs(self, inputs):
|
||||
"""Clip each input to the resolved token limit.
|
||||
|
||||
The remote server (e.g. llama.cpp) hard-rejects any single input
|
||||
larger than its physical batch size with a 500, so oversized inputs are
|
||||
truncated before the request and the overflow is dropped (lossy by
|
||||
design).
|
||||
|
||||
Counting uses the embedding model's own tokenizer where it is known, so
|
||||
the limit and the count are in the same unit. When it is not -- an
|
||||
unregistered model, or no tokenizer available -- this falls back to
|
||||
tiktoken, and the limit should then carry headroom to absorb the skew
|
||||
between the two tokenizers.
|
||||
|
||||
Args:
|
||||
inputs: A single string or a list of strings to embed.
|
||||
|
||||
Returns:
|
||||
The inputs with each string clipped to the token limit, or the
|
||||
inputs unchanged when no limit applies.
|
||||
"""
|
||||
limit = self._resolve_input_limit()
|
||||
if not limit:
|
||||
return inputs
|
||||
|
||||
counter = self._token_counter()
|
||||
|
||||
def clip(text):
|
||||
if not isinstance(text, str):
|
||||
return text
|
||||
count = counter.count(text)
|
||||
if count <= limit:
|
||||
return text
|
||||
logging.warning(
|
||||
"Truncating remote embeddings input from %d to %d tokens (%d dropped)",
|
||||
count,
|
||||
limit,
|
||||
count - limit,
|
||||
)
|
||||
pieces = counter.split(text, limit)
|
||||
return pieces[0] if pieces else text
|
||||
|
||||
if isinstance(inputs, list):
|
||||
return [clip(text) for text in inputs]
|
||||
return clip(inputs)
|
||||
|
||||
def _embed(self, inputs):
|
||||
"""Send embedding request to remote API in OpenAI-compatible format."""
|
||||
inputs = self._truncate_inputs(inputs)
|
||||
payload = {"input": inputs}
|
||||
if self.model_name:
|
||||
payload["model"] = self.model_name
|
||||
|
||||
url = f"{self.api_url}/v1/embeddings"
|
||||
response = requests.post(url, headers=self.headers, json=payload, timeout=180)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Handle OpenAI-compatible response format
|
||||
if isinstance(result, dict):
|
||||
if "error" in result:
|
||||
raise ValueError(f"Remote embeddings API error: {result['error']}")
|
||||
if "data" in result:
|
||||
# Sort by index to ensure correct order
|
||||
data = sorted(result["data"], key=lambda x: x.get("index", 0))
|
||||
return [item["embedding"] for item in data]
|
||||
raise ValueError(
|
||||
f"Unexpected response format from remote embeddings API: {result}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unexpected response format from remote embeddings API: {result}"
|
||||
)
|
||||
|
||||
def embed_query(self, query: str):
|
||||
"""Embed a single query string."""
|
||||
embeddings_list = self._embed(query)
|
||||
if (
|
||||
isinstance(embeddings_list, list)
|
||||
and len(embeddings_list) == 1
|
||||
and isinstance(embeddings_list[0], list)
|
||||
):
|
||||
if self.dimension is None:
|
||||
self.dimension = len(embeddings_list[0])
|
||||
return embeddings_list[0]
|
||||
raise ValueError(
|
||||
f"Unexpected result structure after embedding query: {embeddings_list}"
|
||||
)
|
||||
|
||||
def embed_documents(self, documents: list):
|
||||
"""Embed a list of documents."""
|
||||
if not documents:
|
||||
return []
|
||||
embeddings_list = self._embed(documents)
|
||||
if self.dimension is None and embeddings_list:
|
||||
self.dimension = len(embeddings_list[0])
|
||||
return embeddings_list
|
||||
|
||||
def __call__(self, text):
|
||||
if isinstance(text, str):
|
||||
return self.embed_query(text)
|
||||
elif isinstance(text, list):
|
||||
return self.embed_documents(text)
|
||||
else:
|
||||
raise ValueError("Input must be a string or a list of strings")
|
||||
|
||||
|
||||
def _get_embeddings_wrapper():
|
||||
"""Lazy import of EmbeddingsWrapper, so a remote setup never loads ONNX."""
|
||||
from docsgpt.vectorstore.embeddings_local import EmbeddingsWrapper
|
||||
|
||||
return EmbeddingsWrapper
|
||||
|
||||
|
||||
class EmbeddingsSingleton:
|
||||
_instances = {}
|
||||
|
||||
@staticmethod
|
||||
def _remote_instance(embeddings_name, embeddings_key=None):
|
||||
"""Return a cached ``RemoteEmbeddings`` for the configured remote API.
|
||||
|
||||
Centralizes the ``EMBEDDINGS_BASE_URL`` dispatch so every caller —
|
||||
including code that calls :meth:`get_instance` directly (GraphRAG,
|
||||
semantic chunking) rather than via
|
||||
:meth:`BaseVectorStore._get_embeddings` — routes to the remote
|
||||
embeddings server instead of attempting a local model download.
|
||||
|
||||
Args:
|
||||
embeddings_name: Model name forwarded to the remote API.
|
||||
embeddings_key: Optional API key; falls back to
|
||||
``settings.EMBEDDINGS_KEY`` when not provided.
|
||||
|
||||
Returns:
|
||||
RemoteEmbeddings: Shared instance keyed by base URL and model name.
|
||||
"""
|
||||
api_key = embeddings_key if embeddings_key is not None else settings.EMBEDDINGS_KEY
|
||||
cache_key = f"remote_{settings.EMBEDDINGS_BASE_URL}_{embeddings_name}"
|
||||
if cache_key not in EmbeddingsSingleton._instances:
|
||||
EmbeddingsSingleton._instances[cache_key] = RemoteEmbeddings(
|
||||
api_url=settings.EMBEDDINGS_BASE_URL,
|
||||
model_name=embeddings_name,
|
||||
api_key=api_key,
|
||||
)
|
||||
return EmbeddingsSingleton._instances[cache_key]
|
||||
|
||||
@staticmethod
|
||||
def get_instance(embeddings_name, *args, **kwargs):
|
||||
if settings.EMBEDDINGS_BASE_URL:
|
||||
return EmbeddingsSingleton._remote_instance(embeddings_name)
|
||||
if embeddings_name not in EmbeddingsSingleton._instances:
|
||||
EmbeddingsSingleton._instances[embeddings_name] = (
|
||||
EmbeddingsSingleton._create_instance(embeddings_name, *args, **kwargs)
|
||||
)
|
||||
return EmbeddingsSingleton._instances[embeddings_name]
|
||||
|
||||
@staticmethod
|
||||
def _create_instance(embeddings_name, *args, **kwargs):
|
||||
"""Build the runner for ``embeddings_name``, per the model registry.
|
||||
|
||||
The registry replaced a hand-maintained factory dict whose entries
|
||||
existed only to rewrite a configured name into a repository id. That
|
||||
rewrite is now a registry field, so an unknown name needs no entry
|
||||
here: it is passed through as a Hugging Face repository.
|
||||
"""
|
||||
spec = resolve(embeddings_name)
|
||||
if spec is not None and spec.provider == "openai":
|
||||
return OpenAIEmbeddings(*args, **kwargs)
|
||||
|
||||
EmbeddingsWrapper = _get_embeddings_wrapper()
|
||||
if spec is not None and (args or kwargs):
|
||||
logging.debug(
|
||||
"Dropping %d positional and %d keyword argument(s) for registered "
|
||||
"embeddings model %s: the registry supplies its configuration.",
|
||||
len(args),
|
||||
len(kwargs),
|
||||
embeddings_name,
|
||||
)
|
||||
return EmbeddingsWrapper(embeddings_name)
|
||||
return EmbeddingsWrapper(embeddings_name, *args, **kwargs)
|
||||
|
||||
|
||||
def _azure_configured() -> bool:
|
||||
"""True when the Azure OpenAI deployment settings are all present."""
|
||||
return bool(
|
||||
settings.OPENAI_API_BASE
|
||||
and settings.OPENAI_API_VERSION
|
||||
and settings.AZURE_DEPLOYMENT_NAME
|
||||
)
|
||||
|
||||
|
||||
def _delegation_enabled() -> bool:
|
||||
"""True when this process should embed on the worker rather than locally.
|
||||
|
||||
Compared against ``True`` rather than coerced: tests patch ``settings``
|
||||
with a ``MagicMock``, whose every attribute is a truthy object, and
|
||||
``bool()`` on that would silently route them through the broker.
|
||||
"""
|
||||
return getattr(settings, "EMBEDDINGS_DELEGATE_TO_WORKER", False) is True
|
||||
|
||||
|
||||
def get_embeddings(
|
||||
embeddings_name: Optional[str] = None, embeddings_key: Optional[str] = None
|
||||
):
|
||||
"""Resolve the configured embeddings instance. The single entry point.
|
||||
|
||||
Callers that reach for :meth:`EmbeddingsSingleton.get_instance` directly
|
||||
skip the remote dispatch and the OpenAI/Azure key handling. Route every
|
||||
caller through here.
|
||||
|
||||
With ``EMBEDDINGS_DELEGATE_TO_WORKER`` this returns a client that runs the
|
||||
model on the Celery worker, so an API process never loads one. The client
|
||||
embeds locally when it finds itself inside a worker task, so the worker is
|
||||
unaffected.
|
||||
|
||||
Args:
|
||||
embeddings_name: Model name; defaults to ``settings.EMBEDDINGS_NAME``.
|
||||
embeddings_key: API key; defaults to ``settings.EMBEDDINGS_KEY``.
|
||||
|
||||
Returns:
|
||||
The shared embeddings instance for the resolved model.
|
||||
"""
|
||||
embeddings_name = embeddings_name or settings.EMBEDDINGS_NAME
|
||||
if not settings.EMBEDDINGS_BASE_URL and _delegation_enabled():
|
||||
cache_key = f"delegated_{embeddings_name}"
|
||||
if cache_key not in EmbeddingsSingleton._instances:
|
||||
from docsgpt.vectorstore.embeddings_delegated import DelegatedEmbeddings
|
||||
|
||||
EmbeddingsSingleton._instances[cache_key] = DelegatedEmbeddings(
|
||||
embeddings_name, embeddings_key
|
||||
)
|
||||
return EmbeddingsSingleton._instances[cache_key]
|
||||
return build_local_embeddings(embeddings_name, embeddings_key)
|
||||
|
||||
|
||||
def build_local_embeddings(
|
||||
embeddings_name: Optional[str] = None, embeddings_key: Optional[str] = None
|
||||
):
|
||||
"""Resolve the embeddings instance that runs in *this* process.
|
||||
|
||||
Bypasses worker delegation, so it is what the worker's embed task and the
|
||||
boot hook use. Everything else should call :func:`get_embeddings`.
|
||||
|
||||
Args:
|
||||
embeddings_name: Model name; defaults to ``settings.EMBEDDINGS_NAME``.
|
||||
embeddings_key: API key; defaults to ``settings.EMBEDDINGS_KEY``.
|
||||
|
||||
Returns:
|
||||
The shared in-process embeddings instance for the resolved model.
|
||||
"""
|
||||
embeddings_name = embeddings_name or settings.EMBEDDINGS_NAME
|
||||
embeddings_key = (
|
||||
embeddings_key if embeddings_key is not None else settings.EMBEDDINGS_KEY
|
||||
)
|
||||
|
||||
# Check for remote embeddings first
|
||||
if settings.EMBEDDINGS_BASE_URL:
|
||||
logging.info(
|
||||
f"Using remote embeddings API at: {settings.EMBEDDINGS_BASE_URL}"
|
||||
)
|
||||
return EmbeddingsSingleton._remote_instance(embeddings_name, embeddings_key)
|
||||
|
||||
# Match through the registry, not on the canonical string: the bare
|
||||
# ``text-embedding-ada-002`` alias resolves here too, and skipping this
|
||||
# branch would drop the Azure deployment name and the API key.
|
||||
spec = resolve(embeddings_name)
|
||||
if spec is not None and spec.provider == "openai":
|
||||
if _azure_configured():
|
||||
embedding_instance = EmbeddingsSingleton.get_instance(
|
||||
embeddings_name, model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME
|
||||
)
|
||||
else:
|
||||
embedding_instance = EmbeddingsSingleton.get_instance(
|
||||
embeddings_name, openai_api_key=embeddings_key
|
||||
)
|
||||
else:
|
||||
# No per-model branching: the registry resolves names and FastEmbed
|
||||
# caches artifacts under EMBEDDINGS_CACHE_DIR, which is where the
|
||||
# image warms them at build time.
|
||||
embedding_instance = EmbeddingsSingleton.get_instance(embeddings_name)
|
||||
return embedding_instance
|
||||
|
||||
|
||||
class BaseVectorStore(ABC):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def search(self, *args, **kwargs):
|
||||
"""Search for similar documents/chunks in the vectorstore.
|
||||
|
||||
Implementations accept an optional ``query_vector`` kwarg: the query
|
||||
already embedded by the caller, so a multi-source retrieval embeds once
|
||||
instead of once per store. A store that cannot use it must still swallow
|
||||
the kwarg (every signature here ends in ``**kwargs``) and embed the
|
||||
question itself.
|
||||
"""
|
||||
pass
|
||||
|
||||
def keyword_search(self, question, k=10):
|
||||
"""Keyword/full-text search.
|
||||
|
||||
Default returns no results so hybrid retrieval degrades to vector-only
|
||||
on stores without keyword support. Override in stores that support it.
|
||||
"""
|
||||
return []
|
||||
|
||||
# What ``search_with_scores`` reports, so a caller can label the number.
|
||||
# ``cosine_similarity`` is higher-is-better in [0, 1]; ``l2_distance`` is
|
||||
# lower-is-better and unbounded. None = this store reports no score.
|
||||
score_kind = None
|
||||
|
||||
def search_with_scores(self, question, k=2, *args, **kwargs):
|
||||
"""Search, pairing each hit with its raw relevance score.
|
||||
|
||||
Default pairs every hit from :meth:`search` with ``None`` so stores that
|
||||
surface no score still satisfy the contract. Stores that already compute
|
||||
one override this and set :attr:`score_kind`.
|
||||
|
||||
Returns:
|
||||
A list of ``(Document, score | None)`` in the same rank order
|
||||
:meth:`search` would return.
|
||||
"""
|
||||
return [
|
||||
(doc, None) for doc in self.search(question, k, *args, **kwargs) or []
|
||||
]
|
||||
|
||||
@abstractmethod
|
||||
def add_texts(self, texts, metadatas=None, *args, **kwargs):
|
||||
"""Add texts with their embeddings to the vectorstore"""
|
||||
pass
|
||||
|
||||
def delete_index(self, *args, **kwargs):
|
||||
"""Delete the entire index/collection"""
|
||||
pass
|
||||
|
||||
def save_local(self, *args, **kwargs):
|
||||
"""Save vectorstore to local storage"""
|
||||
pass
|
||||
|
||||
def get_chunks(self, *args, **kwargs):
|
||||
"""Get all chunks from the vectorstore"""
|
||||
pass
|
||||
|
||||
def add_chunk(self, text, metadata=None, *args, **kwargs):
|
||||
"""Add a single chunk to the vectorstore"""
|
||||
pass
|
||||
|
||||
def delete_chunk(self, chunk_id, *args, **kwargs):
|
||||
"""Delete a specific chunk from the vectorstore"""
|
||||
pass
|
||||
|
||||
def delete_chunks_by_source_path(self, path) -> int:
|
||||
"""Delete every chunk whose ``metadata.source`` equals ``path``.
|
||||
|
||||
Default implementation iterates ``get_chunks()`` and deletes the
|
||||
matches via ``delete_chunk()`` — works for any store. Override with a
|
||||
single targeted statement where the store supports it. Returns the
|
||||
number of chunks deleted.
|
||||
"""
|
||||
deleted = 0
|
||||
for chunk in self.get_chunks() or []:
|
||||
if (chunk.get("metadata") or {}).get("source") == path:
|
||||
if self.delete_chunk(chunk.get("doc_id")):
|
||||
deleted += 1
|
||||
return deleted
|
||||
|
||||
def is_azure_configured(self):
|
||||
"""Kept for compatibility; delegates to the module-level check."""
|
||||
return _azure_configured()
|
||||
|
||||
def _get_embeddings(self, embeddings_name, embeddings_key=None):
|
||||
"""Resolve embeddings for this store; see :func:`get_embeddings`."""
|
||||
return get_embeddings(embeddings_name, embeddings_key)
|
||||
@@ -0,0 +1,8 @@
|
||||
class Document(str):
|
||||
"""Class for storing a piece of text and associated metadata."""
|
||||
|
||||
def __new__(cls, page_content: str, metadata: dict):
|
||||
instance = super().__new__(cls, page_content)
|
||||
instance.page_content = page_content
|
||||
instance.metadata = metadata
|
||||
return instance
|
||||
@@ -0,0 +1,225 @@
|
||||
from docsgpt.vectorstore.base import BaseVectorStore
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.document_class import Document
|
||||
|
||||
|
||||
class ElasticsearchStore(BaseVectorStore):
|
||||
_es_connection = None # Class attribute to hold the Elasticsearch connection
|
||||
|
||||
def __init__(self, source_id, embeddings_key, index_name=settings.ELASTIC_INDEX):
|
||||
super().__init__()
|
||||
self.source_id = source_id.replace("docsgpt/indexes/", "").rstrip("/")
|
||||
self.embeddings_key = embeddings_key
|
||||
self.index_name = index_name
|
||||
|
||||
if ElasticsearchStore._es_connection is None:
|
||||
connection_params = {}
|
||||
if settings.ELASTIC_URL:
|
||||
connection_params["hosts"] = [settings.ELASTIC_URL]
|
||||
connection_params["http_auth"] = (settings.ELASTIC_USERNAME, settings.ELASTIC_PASSWORD)
|
||||
elif settings.ELASTIC_CLOUD_ID:
|
||||
connection_params["cloud_id"] = settings.ELASTIC_CLOUD_ID
|
||||
connection_params["basic_auth"] = (settings.ELASTIC_USERNAME, settings.ELASTIC_PASSWORD)
|
||||
else:
|
||||
raise ValueError("Please provide either elasticsearch_url or cloud_id.")
|
||||
|
||||
import elasticsearch
|
||||
ElasticsearchStore._es_connection = elasticsearch.Elasticsearch(**connection_params)
|
||||
|
||||
self.docsearch = ElasticsearchStore._es_connection
|
||||
|
||||
def connect_to_elasticsearch(
|
||||
*,
|
||||
es_url = None,
|
||||
cloud_id = None,
|
||||
api_key = None,
|
||||
username = None,
|
||||
password = None,
|
||||
):
|
||||
try:
|
||||
import elasticsearch
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Could not import elasticsearch python package. "
|
||||
"Please install it with `pip install elasticsearch`."
|
||||
)
|
||||
|
||||
if es_url and cloud_id:
|
||||
raise ValueError(
|
||||
"Both es_url and cloud_id are defined. Please provide only one."
|
||||
)
|
||||
|
||||
connection_params = {}
|
||||
|
||||
if es_url:
|
||||
connection_params["hosts"] = [es_url]
|
||||
elif cloud_id:
|
||||
connection_params["cloud_id"] = cloud_id
|
||||
else:
|
||||
raise ValueError("Please provide either elasticsearch_url or cloud_id.")
|
||||
|
||||
if api_key:
|
||||
connection_params["api_key"] = api_key
|
||||
elif username and password:
|
||||
connection_params["basic_auth"] = (username, password)
|
||||
|
||||
es_client = elasticsearch.Elasticsearch(
|
||||
**connection_params,
|
||||
)
|
||||
try:
|
||||
es_client.info()
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
return es_client
|
||||
|
||||
def search(
|
||||
self,
|
||||
question,
|
||||
k=2,
|
||||
index_name=settings.ELASTIC_INDEX,
|
||||
*args,
|
||||
query_vector=None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Search by kNN + full text, fused with RRF.
|
||||
|
||||
Args:
|
||||
query_vector: Precomputed embedding of ``question``; when given the
|
||||
store skips embedding the query itself.
|
||||
"""
|
||||
vector = query_vector
|
||||
if vector is None:
|
||||
embeddings = self._get_embeddings(
|
||||
settings.EMBEDDINGS_NAME, self.embeddings_key
|
||||
)
|
||||
vector = embeddings.embed_query(question)
|
||||
knn = {
|
||||
"filter": [{"match": {"metadata.source_id.keyword": self.source_id}}],
|
||||
"field": "vector",
|
||||
"k": k,
|
||||
"num_candidates": 100,
|
||||
"query_vector": vector,
|
||||
}
|
||||
full_query = {
|
||||
"knn": knn,
|
||||
"query": {
|
||||
"bool": {
|
||||
"must": [
|
||||
{
|
||||
"match": {
|
||||
"text": {
|
||||
"query": question,
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
"filter": [{"match": {"metadata.source_id.keyword": self.source_id}}],
|
||||
}
|
||||
},
|
||||
"rank": {"rrf": {}},
|
||||
}
|
||||
resp = self.docsearch.search(index=self.index_name, query=full_query['query'], size=k, knn=full_query['knn'])
|
||||
# create Documents objects from the results page_content ['_source']['text'], metadata ['_source']['metadata']
|
||||
doc_list = []
|
||||
for hit in resp['hits']['hits']:
|
||||
|
||||
doc_list.append(Document(page_content = hit['_source']['text'], metadata = hit['_source']['metadata']))
|
||||
return doc_list
|
||||
|
||||
def _create_index_if_not_exists(
|
||||
self, index_name, dims_length
|
||||
):
|
||||
|
||||
if self._es_connection.indices.exists(index=index_name):
|
||||
print(f"Index {index_name} already exists.")
|
||||
|
||||
else:
|
||||
|
||||
indexSettings = self.index(
|
||||
dims_length=dims_length,
|
||||
)
|
||||
self._es_connection.indices.create(index=index_name, **indexSettings)
|
||||
|
||||
def index(
|
||||
self,
|
||||
dims_length,
|
||||
):
|
||||
return {
|
||||
"mappings": {
|
||||
"properties": {
|
||||
"vector": {
|
||||
"type": "dense_vector",
|
||||
"dims": dims_length,
|
||||
"index": True,
|
||||
"similarity": "cosine",
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
def add_texts(
|
||||
self,
|
||||
texts,
|
||||
metadatas = None,
|
||||
ids = None,
|
||||
refresh_indices = True,
|
||||
create_index_if_not_exists = True,
|
||||
bulk_kwargs = None,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
bulk_kwargs = bulk_kwargs or {}
|
||||
import uuid
|
||||
embeddings = []
|
||||
ids = ids or [str(uuid.uuid4()) for _ in texts]
|
||||
requests = []
|
||||
embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key)
|
||||
|
||||
vectors = embeddings.embed_documents(list(texts))
|
||||
|
||||
dims_length = len(vectors[0])
|
||||
|
||||
if create_index_if_not_exists:
|
||||
self._create_index_if_not_exists(
|
||||
index_name=self.index_name, dims_length=dims_length
|
||||
)
|
||||
|
||||
for i, (text, vector) in enumerate(zip(texts, vectors)):
|
||||
metadata = metadatas[i] if metadatas else {}
|
||||
|
||||
requests.append(
|
||||
{
|
||||
"_op_type": "index",
|
||||
"_index": self.index_name,
|
||||
"text": text,
|
||||
"vector": vector,
|
||||
"metadata": metadata,
|
||||
"_id": ids[i],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
if len(requests) > 0:
|
||||
from elasticsearch.helpers import BulkIndexError, bulk
|
||||
try:
|
||||
success, failed = bulk(
|
||||
self._es_connection,
|
||||
requests,
|
||||
stats_only=True,
|
||||
refresh=refresh_indices,
|
||||
**bulk_kwargs,
|
||||
)
|
||||
return ids
|
||||
except BulkIndexError as e:
|
||||
print(f"Error adding texts: {e}")
|
||||
firstError = e.errors[0].get("index", {}).get("error", {})
|
||||
print(f"First error reason: {firstError.get('reason')}")
|
||||
raise e
|
||||
|
||||
else:
|
||||
return []
|
||||
|
||||
def delete_index(self):
|
||||
self._es_connection.delete_by_query(index=self.index_name, query={"match": {
|
||||
"metadata.source_id.keyword": self.source_id}},)
|
||||
@@ -0,0 +1,254 @@
|
||||
"""Query embedding executed in the Celery worker instead of in the API.
|
||||
|
||||
The API embeds every query it serves, so it needs an embedder -- and a local
|
||||
one costs roughly 800 MB of ONNX Runtime per process. That is the whole
|
||||
footprint of an API container that otherwise holds no model.
|
||||
|
||||
This client keeps the interface (``embed_query``/``embed_documents``/
|
||||
``dimension``) and moves only the computation: the text goes to the worker over
|
||||
Celery and the vector comes back. The API pays a broker round trip per query
|
||||
and no resident model.
|
||||
|
||||
Inside a worker there is nothing to delegate to -- dispatching would queue work
|
||||
behind the task already running and wait on itself -- so a call made while a
|
||||
task is executing runs locally, on a model this process loads once and caches.
|
||||
``DOCUMENT_PARSE_QUEUE`` exists for the same reason on the parsing side.
|
||||
|
||||
Production deployments should point ``EMBEDDINGS_BASE_URL`` at a real embedding
|
||||
service instead: that removes the model from *both* processes and costs a
|
||||
network hop rather than a broker round trip.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.model_registry import dimension_for
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: Dispatched by name so the API never imports the task module -- and through
|
||||
#: it ``docsgpt.worker``, which pulls in the whole parsing stack.
|
||||
EMBED_TASK = "docsgpt.vectorstore.embeddings_tasks.embed_texts"
|
||||
|
||||
#: How long after a failed dispatch to fail fast instead of waiting out another
|
||||
#: full ``EMBEDDINGS_DELEGATE_TIMEOUT``. Short enough that a worker restart is
|
||||
#: picked up within one query, long enough to collapse the retries inside a
|
||||
#: single retrieval into one timeout rather than one per source.
|
||||
_FAILURE_COOLDOWN = 30.0
|
||||
|
||||
#: How long a caller waits for the outcome of the dispatch already in flight
|
||||
#: before giving up on its own. Only applies while the worker is unproven --
|
||||
#: once one dispatch has succeeded, every caller goes straight to the broker.
|
||||
#: Comfortably above a healthy round trip (~60 ms on a prefork worker) and far
|
||||
#: below ``EMBEDDINGS_DELEGATE_TIMEOUT``, which is the point.
|
||||
_PROBE_WAIT = 2.0
|
||||
|
||||
_NO_WORKER_HINT = (
|
||||
"Start a worker consuming it, point EMBEDDINGS_BASE_URL at an embedding "
|
||||
"service, or set EMBEDDINGS_DELEGATE_TO_WORKER=false to load the model in "
|
||||
"this process instead."
|
||||
)
|
||||
|
||||
|
||||
def _forget(result) -> None:
|
||||
"""Drop the task's stored vector from the result backend.
|
||||
|
||||
Nothing ever reads it back. The key is ``celery-task-meta-<uuid>``, minted
|
||||
per dispatch rather than derived from the text, so a repeated query is a new
|
||||
task and a new key -- the value is written once, read once by the ``get()``
|
||||
already waiting on it, then dead. Left alone it occupies ~17 KB for
|
||||
``result_expires`` (7 days), in the Redis the broker also runs on.
|
||||
|
||||
Also releases the backend's pub/sub subscription for the task, which
|
||||
``get()`` alone does not.
|
||||
|
||||
Never raises: the vector is already in hand, and a backend that cannot
|
||||
delete must not fail the search. On the timeout path the worker may still
|
||||
store its result afterwards, leaving one orphaned key -- no worse than not
|
||||
forgetting at all, and bounded by the same expiry.
|
||||
"""
|
||||
try:
|
||||
result.forget()
|
||||
except Exception as exc: # noqa: BLE001 — cleanup must never fail a query
|
||||
logger.debug("Could not forget the embed task result: %s", exc)
|
||||
|
||||
|
||||
def _in_worker() -> bool:
|
||||
"""True when a Celery task is executing in this process."""
|
||||
try:
|
||||
from docsgpt.celery_init import celery
|
||||
|
||||
return celery.current_worker_task is not None
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
class DelegatedEmbeddings:
|
||||
"""Embeds by dispatching to the Celery worker, or locally inside one."""
|
||||
|
||||
def __init__(self, embeddings_name: str, embeddings_key: Optional[str] = None) -> None:
|
||||
self.embeddings_name = embeddings_name
|
||||
self.embeddings_key = embeddings_key
|
||||
self._local: Any = None
|
||||
self._dimension: Optional[int] = dimension_for(embeddings_name)
|
||||
self._failed_at: Optional[float] = None
|
||||
# A dispatch has completed successfully, so the worker is known to be
|
||||
# consuming the queue and callers need not take turns proving it.
|
||||
self._verified = False
|
||||
self._probing = False
|
||||
self._state_lock = threading.Lock()
|
||||
self._probe_done = threading.Event()
|
||||
|
||||
def _cooldown_remaining(self) -> float:
|
||||
"""Seconds left of the fail-fast window after a failed dispatch."""
|
||||
# One load: a concurrent success clearing the latch between two reads
|
||||
# would otherwise subtract from None.
|
||||
failed_at = self._failed_at
|
||||
if failed_at is None:
|
||||
return 0.0
|
||||
return max(0.0, _FAILURE_COOLDOWN - (time.monotonic() - failed_at))
|
||||
|
||||
def _local_embeddings(self):
|
||||
"""The in-process model, built once, for use inside a worker task."""
|
||||
if self._local is None:
|
||||
from docsgpt.vectorstore.base import build_local_embeddings
|
||||
|
||||
self._local = build_local_embeddings(self.embeddings_name, self.embeddings_key)
|
||||
return self._local
|
||||
|
||||
def _send(self, texts: List[str], queue: str, timeout: int) -> List[List[float]]:
|
||||
"""Publish the embed task and wait for its vectors."""
|
||||
from docsgpt.celery_init import celery
|
||||
|
||||
result = celery.send_task(EMBED_TASK, args=[texts, self.embeddings_name], queue=queue)
|
||||
try:
|
||||
vectors = result.get(timeout=timeout)
|
||||
except Exception as exc:
|
||||
self._failed_at = time.monotonic()
|
||||
# Drop the proof with the worker that supplied it. ``_verified``
|
||||
# short-circuits ahead of the probe gate, so leaving it set means
|
||||
# the gate only ever covers a worker that was never healthy --
|
||||
# while the case that actually happens is a healthy one being
|
||||
# redeployed or OOM-killed. Every caller would then pay the full
|
||||
# timeout, together, on every wave once the cooldown lapses.
|
||||
self._verified = False
|
||||
raise RuntimeError(
|
||||
f"Embedding request to the Celery worker timed out or failed ({exc}). "
|
||||
f"A worker must be consuming the {queue!r} queue for retrieval to "
|
||||
f"work. {_NO_WORKER_HINT}"
|
||||
) from exc
|
||||
finally:
|
||||
_forget(result)
|
||||
self._failed_at = None
|
||||
self._verified = True
|
||||
return vectors
|
||||
|
||||
def _cooldown_error(self, queue: str, remaining: float) -> RuntimeError:
|
||||
return RuntimeError(
|
||||
f"Skipping the embed dispatch: a previous request to the {queue!r} "
|
||||
f"queue failed and the {_FAILURE_COOLDOWN}s cooldown has "
|
||||
f"{remaining:.0f}s left. {_NO_WORKER_HINT}"
|
||||
)
|
||||
|
||||
def _dispatch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""Run the embed task on the worker and wait for its vectors.
|
||||
|
||||
A missing worker is a property of the deployment, not of this call, so
|
||||
at most one caller waits out ``EMBEDDINGS_DELEGATE_TIMEOUT`` to discover
|
||||
it. Two guards do that:
|
||||
|
||||
The cooldown latch covers requests arriving *after* a failure -- without
|
||||
it a single retrieval pays the timeout twice, once in
|
||||
``fanout.embed_questions`` and again per source when it falls back to
|
||||
letting each store embed its own query.
|
||||
|
||||
The probe covers requests already in flight *alongside* the first one,
|
||||
which the latch cannot: nothing is latched until that first ``get()``
|
||||
returns, so every thread in the opening wave would otherwise block for
|
||||
the full timeout at once -- at the shipped 60s and 96 WSGI threads, an
|
||||
API that serves nothing at all, health checks included.
|
||||
"""
|
||||
queue = getattr(settings, "EMBEDDINGS_QUEUE", "embeddings")
|
||||
timeout = getattr(settings, "EMBEDDINGS_DELEGATE_TIMEOUT", 60)
|
||||
|
||||
remaining = self._cooldown_remaining()
|
||||
if remaining > 0:
|
||||
raise self._cooldown_error(queue, remaining)
|
||||
|
||||
if self._verified:
|
||||
return self._send(texts, queue, timeout)
|
||||
|
||||
with self._state_lock:
|
||||
# "send" -- proven while we waited for the lock, just go.
|
||||
# "wait" -- another caller is already finding out; don't pay a
|
||||
# second full timeout to learn the same thing.
|
||||
# "probe" -- nobody is; this call is the one that finds out.
|
||||
role = "send" if self._verified else "wait" if self._probing else "probe"
|
||||
if role == "probe":
|
||||
self._probing = True
|
||||
self._probe_done.clear()
|
||||
|
||||
if role == "probe":
|
||||
try:
|
||||
return self._send(texts, queue, timeout)
|
||||
finally:
|
||||
with self._state_lock:
|
||||
self._probing = False
|
||||
self._probe_done.set()
|
||||
|
||||
if role == "wait":
|
||||
self._probe_done.wait(_PROBE_WAIT)
|
||||
remaining = self._cooldown_remaining()
|
||||
if remaining > 0:
|
||||
raise self._cooldown_error(queue, remaining)
|
||||
if not self._verified:
|
||||
raise RuntimeError(
|
||||
f"Skipping the embed dispatch: an earlier request to the "
|
||||
f"{queue!r} queue is still unanswered after {_PROBE_WAIT}s, so "
|
||||
f"no worker appears to be consuming it. {_NO_WORKER_HINT}"
|
||||
)
|
||||
|
||||
return self._send(texts, queue, timeout)
|
||||
|
||||
def embed_documents(self, documents: List[str]) -> List[List[float]]:
|
||||
"""Embed a list of texts, preserving order."""
|
||||
if not documents:
|
||||
return []
|
||||
if _in_worker():
|
||||
return self._local_embeddings().embed_documents(documents)
|
||||
vectors = self._dispatch(list(documents))
|
||||
if self._dimension is None and vectors:
|
||||
self._dimension = len(vectors[0])
|
||||
return vectors
|
||||
|
||||
def embed_query(self, query: str) -> List[float]:
|
||||
"""Embed a single query string."""
|
||||
return self.embed_documents([query])[0]
|
||||
|
||||
@property
|
||||
def dimension(self) -> Optional[int]:
|
||||
"""Vector width, from the registry where possible.
|
||||
|
||||
Falls back to one round trip for a model the registry does not
|
||||
describe, and to ``None`` when even that fails -- callers already treat
|
||||
an unknown width as "nothing to compare yet" rather than an error.
|
||||
"""
|
||||
if self._dimension is None:
|
||||
try:
|
||||
self._dimension = len(self.embed_query("dimension probe"))
|
||||
except Exception as exc:
|
||||
logger.warning("Could not determine embedding width: %s", exc)
|
||||
return None
|
||||
return self._dimension
|
||||
|
||||
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")
|
||||
@@ -0,0 +1,356 @@
|
||||
"""Local embeddings via FastEmbed (ONNX Runtime).
|
||||
|
||||
Replaces the previous SentenceTransformer implementation. Both run the same
|
||||
weights; FastEmbed reaches them through ONNX Runtime instead of torch, which
|
||||
removes torch, transformers and sentence-transformers from the dependency set
|
||||
and measurably reduces both resident memory and import cost.
|
||||
|
||||
The swap is numerically transparent for existing indexes: embedding the same
|
||||
text with SentenceTransformer and with FastEmbed's fp32 ONNX graph of mpnet
|
||||
returns vectors at cosine 1.0, so an mpnet index built before this change
|
||||
keeps working unchanged. That result is specific to the fp32 graph -- the
|
||||
granite entries run an int8-quantised one and are not bit-comparable to a
|
||||
fp32 index of the same model.
|
||||
|
||||
Models are described in :mod:`docsgpt.vectorstore.model_registry`. A name
|
||||
the registry does not know is treated as a Hugging Face repository, which is
|
||||
what someone configuring an arbitrary model expects; how to run it is read
|
||||
from the repository itself rather than assumed.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import replace
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.model_registry import EmbeddingModel, resolve, known_names
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ``add_custom_model`` mutates a process-global registry inside FastEmbed, so
|
||||
# repeated registration of the same name is both wasteful and racy under the
|
||||
# thread pool the API serves requests from.
|
||||
_registered: set = set()
|
||||
_register_lock = threading.Lock()
|
||||
|
||||
# Last-resort layout for a repository that declares nothing about itself.
|
||||
_FALLBACK_ONNX_FILE = "onnx/model.onnx"
|
||||
_FALLBACK_POOLING = "mean"
|
||||
|
||||
# Sentence-transformers records how a model turns token vectors into one
|
||||
# vector, and whether it normalises the result, as files in the repository.
|
||||
# Reading them is the difference between running a model and running something
|
||||
# that merely shares its weights: mean-pooling a CLS model returns vectors at
|
||||
# cosine ~0.95 to the correct ones -- close enough to look like it works, far
|
||||
# enough to degrade retrieval, and silent either way.
|
||||
_POOLING_CONFIG = "1_Pooling/config.json"
|
||||
_MODULES_CONFIG = "modules.json"
|
||||
|
||||
|
||||
def _pooling_type(pooling: str):
|
||||
"""Map our ``"cls"``/``"mean"`` spelling onto FastEmbed's enum."""
|
||||
from fastembed.common.model_description import PoolingType
|
||||
|
||||
return {"cls": PoolingType.CLS, "mean": PoolingType.MEAN}[pooling]
|
||||
|
||||
|
||||
def _is_builtin(repo: str) -> bool:
|
||||
"""True when FastEmbed already ships a description for ``repo``."""
|
||||
from fastembed import TextEmbedding
|
||||
|
||||
lowered = repo.lower()
|
||||
return any(
|
||||
str(entry.get("model", "")).lower() == lowered
|
||||
for entry in TextEmbedding.list_supported_models()
|
||||
)
|
||||
|
||||
|
||||
def _register(model: EmbeddingModel) -> None:
|
||||
"""Teach FastEmbed about a model, exactly once per process."""
|
||||
from fastembed import TextEmbedding
|
||||
from fastembed.common.model_description import ModelSource
|
||||
|
||||
with _register_lock:
|
||||
if model.repo in _registered:
|
||||
return
|
||||
if _is_builtin(model.repo):
|
||||
# ``add_custom_model`` refuses a name FastEmbed already ships, and
|
||||
# its own description carries the pooling, width and graph file we
|
||||
# would be supplying, so there is nothing to add. Without this,
|
||||
# configuring any of FastEmbed's ~30 built-in models (bge, e5,
|
||||
# MiniLM, gte, ...) fails every embed call.
|
||||
_registered.add(model.repo)
|
||||
return
|
||||
TextEmbedding.add_custom_model(
|
||||
model=model.repo,
|
||||
pooling=_pooling_type(model.pooling),
|
||||
normalization=model.normalize,
|
||||
sources=ModelSource(hf=model.repo),
|
||||
dim=model.dimension,
|
||||
model_file=model.onnx_file,
|
||||
)
|
||||
_registered.add(model.repo)
|
||||
|
||||
|
||||
def _read_repo_json(repo: str, filename: str) -> Optional[dict]:
|
||||
"""Fetch one small JSON from ``repo``, or ``None`` when it is not there.
|
||||
|
||||
Reads through the Hugging Face hub cache, so a warmed image finds it
|
||||
offline. Every failure -- absent file, no network, malformed JSON -- is the
|
||||
same answer to the caller: this repository does not tell us.
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
try:
|
||||
path = hf_hub_download(repo_id=repo, filename=filename, local_files_only=True)
|
||||
except Exception: # noqa: BLE001 -- not cached: fetch it
|
||||
path = hf_hub_download(repo_id=repo, filename=filename)
|
||||
with open(path, encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
except Exception as exc:
|
||||
logger.debug("No %s for %s (%s)", filename, repo, exc)
|
||||
return None
|
||||
|
||||
|
||||
def _describe_from_repo(repo: str) -> Optional[EmbeddingModel]:
|
||||
"""Build a spec from a repository's sentence-transformers metadata.
|
||||
|
||||
Args:
|
||||
repo: Hugging Face repository id.
|
||||
|
||||
Returns:
|
||||
The described model, or ``None`` when the repository carries no
|
||||
metadata to read -- leaving the caller to fall back to assumptions.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the model has a Dense projection head. FastEmbed runs
|
||||
the transformer and pools it, and nothing else, so the projection
|
||||
would be skipped and the vectors come out both the wrong width and
|
||||
in a different space. There is no correct way to run it here.
|
||||
"""
|
||||
pooling_config = _read_repo_json(repo, _POOLING_CONFIG)
|
||||
if pooling_config is None:
|
||||
return None
|
||||
|
||||
kinds = {
|
||||
str(module.get("type", "")).rsplit(".", 1)[-1]
|
||||
for module in (_read_repo_json(repo, _MODULES_CONFIG) or [])
|
||||
if isinstance(module, dict)
|
||||
}
|
||||
if "Dense" in kinds:
|
||||
raise RuntimeError(
|
||||
f"Embedding model {repo!r} has a Dense projection layer, which FastEmbed "
|
||||
"cannot run: its vectors would be the wrong width and in a different "
|
||||
"space than the model was trained to produce. Choose a model without "
|
||||
"one, or serve this one over EMBEDDINGS_BASE_URL."
|
||||
)
|
||||
|
||||
if pooling_config.get("pooling_mode_cls_token"):
|
||||
pooling = "cls"
|
||||
elif pooling_config.get("pooling_mode_mean_tokens"):
|
||||
pooling = "mean"
|
||||
else:
|
||||
# max, mean_sqrt_len, weighted-mean: FastEmbed offers none of them, so
|
||||
# there is nothing to describe and guessing is what we are avoiding.
|
||||
logger.warning(
|
||||
"Embedding model %s uses a pooling mode FastEmbed cannot reproduce (%s).",
|
||||
repo,
|
||||
", ".join(sorted(k for k, v in pooling_config.items() if v is True)) or "unknown",
|
||||
)
|
||||
return None
|
||||
|
||||
# ``Normalize`` appears in modules.json only when the model L2-normalises.
|
||||
# Its absence is a fact, not missing data: a dot-product model is trained
|
||||
# on unnormalised vectors and normalising re-ranks its results.
|
||||
normalize = "Normalize" in kinds
|
||||
dimension = int(pooling_config.get("word_embedding_dimension") or 0)
|
||||
logger.info(
|
||||
"Embedding model %s declares %s pooling, normalize=%s, dimension=%s.",
|
||||
repo,
|
||||
pooling,
|
||||
normalize,
|
||||
dimension or "unknown",
|
||||
)
|
||||
return EmbeddingModel(
|
||||
name=repo,
|
||||
dimension=dimension,
|
||||
max_input_tokens=512,
|
||||
pooling=pooling,
|
||||
normalize=normalize,
|
||||
repo=repo,
|
||||
onnx_file=_FALLBACK_ONNX_FILE,
|
||||
)
|
||||
|
||||
|
||||
def _apply_overrides(spec: EmbeddingModel) -> EmbeddingModel:
|
||||
"""Let ``EMBEDDINGS_POOLING``/``EMBEDDINGS_NORMALIZE`` win over any source."""
|
||||
pooling = getattr(settings, "EMBEDDINGS_POOLING", None)
|
||||
normalize = getattr(settings, "EMBEDDINGS_NORMALIZE", None)
|
||||
changes = {}
|
||||
if isinstance(pooling, str) and pooling.strip().lower() in ("cls", "mean"):
|
||||
changes["pooling"] = pooling.strip().lower()
|
||||
if isinstance(normalize, bool):
|
||||
changes["normalize"] = normalize
|
||||
if not changes:
|
||||
return spec
|
||||
logger.info("Overriding %s from settings: %s", spec.repo, changes)
|
||||
return replace(spec, **changes)
|
||||
|
||||
|
||||
def _spec_for(model_name: str) -> EmbeddingModel:
|
||||
"""Registry entry for ``model_name``, or a best-effort one for a raw repo.
|
||||
|
||||
Args:
|
||||
model_name: Configured ``EMBEDDINGS_NAME``.
|
||||
|
||||
Returns:
|
||||
The registry entry, else one read from the repository's own
|
||||
sentence-transformers metadata, else a last-resort entry assuming the
|
||||
standard ONNX layout with mean pooling and ``dimension = 0`` so the
|
||||
caller knows to probe for the real width. Settings overrides win over
|
||||
all three.
|
||||
"""
|
||||
spec = resolve(model_name)
|
||||
if spec is not None:
|
||||
return _apply_overrides(spec)
|
||||
|
||||
described = _describe_from_repo(model_name)
|
||||
if described is not None:
|
||||
return _apply_overrides(described)
|
||||
|
||||
logger.warning(
|
||||
"Embedding model %r is not in the registry (known: %s) and its repository "
|
||||
"declares no pooling, so %s with L2 normalisation is assumed. If that is "
|
||||
"wrong the vectors will be quietly poor rather than fail; set "
|
||||
"EMBEDDINGS_POOLING=cls|mean and EMBEDDINGS_NORMALIZE to pin it.",
|
||||
model_name,
|
||||
", ".join(known_names()),
|
||||
_FALLBACK_POOLING,
|
||||
)
|
||||
return _apply_overrides(
|
||||
EmbeddingModel(
|
||||
name=model_name,
|
||||
dimension=0,
|
||||
max_input_tokens=512,
|
||||
pooling=_FALLBACK_POOLING,
|
||||
normalize=True,
|
||||
repo=model_name,
|
||||
onnx_file=_FALLBACK_ONNX_FILE,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _pad_to_longest_in_batch(model: Any) -> None:
|
||||
"""Undo a fixed padding width baked into a model's ``tokenizer.json``."""
|
||||
# FastEmbed enables padding only when the tokenizer declares none, so a
|
||||
# fixed ``length`` survives loading. Shorter inputs are then padded to that
|
||||
# width while longer ones keep their own, the batch is ragged, and the ONNX
|
||||
# tensor build fails. mpnet ships ``length: 128``; granite does not.
|
||||
# Mean pooling masks pad tokens, so the vectors are unaffected.
|
||||
tokenizer = getattr(getattr(model, "model", None), "tokenizer", None)
|
||||
padding = getattr(tokenizer, "padding", None)
|
||||
if not isinstance(padding, dict) or padding.get("length") is None:
|
||||
return
|
||||
tokenizer.enable_padding(
|
||||
direction=padding.get("direction", "right"),
|
||||
pad_id=padding.get("pad_id", 0),
|
||||
pad_type_id=padding.get("pad_type_id", 0),
|
||||
pad_token=padding.get("pad_token", "<pad>"),
|
||||
length=None,
|
||||
pad_to_multiple_of=padding.get("pad_to_multiple_of"),
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingsWrapper:
|
||||
"""Runs an embedding model locally through FastEmbed.
|
||||
|
||||
Exposes the ``embed_query``/``embed_documents``/``dimension`` interface the
|
||||
vector stores rely on, matching ``RemoteEmbeddings`` and ``OpenAIEmbeddings``.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str, *args: Any, **kwargs: Any) -> None:
|
||||
"""Load ``model_name`` locally.
|
||||
|
||||
Args:
|
||||
model_name: Registry name, alias, or a Hugging Face repository id.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the model cannot be loaded, with the configured
|
||||
name and the registry's known names in the message.
|
||||
"""
|
||||
from fastembed import TextEmbedding
|
||||
|
||||
self.spec = _spec_for(model_name)
|
||||
logger.info("Loading embeddings model %s via FastEmbed", self.spec.repo)
|
||||
try:
|
||||
_register(self.spec)
|
||||
init_kwargs = {"model_name": self.spec.repo}
|
||||
threads = getattr(settings, "EMBEDDINGS_THREADS", None)
|
||||
if isinstance(threads, int) and threads > 0:
|
||||
init_kwargs["threads"] = threads
|
||||
cache_dir = getattr(settings, "EMBEDDINGS_CACHE_DIR", None)
|
||||
if cache_dir:
|
||||
init_kwargs["cache_dir"] = cache_dir
|
||||
self.model = TextEmbedding(**init_kwargs)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
f"Could not load embeddings model {model_name!r} via FastEmbed: "
|
||||
f"{exc}. Known models: {', '.join(known_names())}."
|
||||
) from exc
|
||||
|
||||
_pad_to_longest_in_batch(self.model)
|
||||
self.dimension = self.spec.dimension or self._probe_dimension()
|
||||
logger.info("Embeddings model ready (dimension=%d)", self.dimension)
|
||||
|
||||
def _probe_dimension(self) -> int:
|
||||
"""Determine the vector width of a model the registry does not describe."""
|
||||
return len(self.embed_query("dimension probe"))
|
||||
|
||||
def embed_query(self, query: str) -> List[float]:
|
||||
"""Embed a single query string."""
|
||||
return self.embed_documents([query])[0]
|
||||
|
||||
def embed_documents(self, documents: List[str]) -> List[List[float]]:
|
||||
"""Embed a list of documents, preserving input order.
|
||||
|
||||
Batched by ``EMBEDDINGS_MODEL_BATCH_SIZE``, not by the pipeline's
|
||||
``EMBEDDINGS_BATCH_SIZE``: one is documents per forward pass, the other
|
||||
is chunks per store transaction, and sizing the forward pass from the
|
||||
transaction is what made ingest peak at 6.6 GB.
|
||||
|
||||
Inputs are grouped by length first. ONNX needs a rectangular tensor, so
|
||||
every input in a pass is padded up to the longest one in it; with mixed
|
||||
lengths that padding is most of the work. The original order is
|
||||
restored before returning, so callers zipping these against their texts
|
||||
are unaffected.
|
||||
"""
|
||||
if not documents:
|
||||
return []
|
||||
batch_size: Optional[int] = None
|
||||
raw = getattr(settings, "EMBEDDINGS_MODEL_BATCH_SIZE", None)
|
||||
if isinstance(raw, int) and not isinstance(raw, bool) and raw > 0:
|
||||
batch_size = raw
|
||||
|
||||
documents = list(documents)
|
||||
if batch_size is None or len(documents) <= batch_size:
|
||||
# One batch either way: sorting would only add work.
|
||||
return [v.tolist() for v in self.model.embed(documents, batch_size=batch_size)]
|
||||
|
||||
order = sorted(range(len(documents)), key=lambda i: len(documents[i]))
|
||||
grouped = [documents[i] for i in order]
|
||||
vectors = [v.tolist() for v in self.model.embed(grouped, batch_size=batch_size)]
|
||||
|
||||
restored: List[Optional[List[float]]] = [None] * len(documents)
|
||||
for position, original_index in enumerate(order):
|
||||
restored[original_index] = vectors[position]
|
||||
return restored
|
||||
|
||||
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")
|
||||
@@ -0,0 +1,89 @@
|
||||
"""OpenAI (and Azure OpenAI) embeddings built on the official ``openai`` SDK."""
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from docsgpt.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")
|
||||
@@ -0,0 +1,29 @@
|
||||
"""The Celery task behind :mod:`docsgpt.vectorstore.embeddings_delegated`.
|
||||
|
||||
Kept out of ``docsgpt.api.user.tasks`` deliberately: that module imports
|
||||
``docsgpt.worker`` and the whole parsing stack with it, which is the
|
||||
opposite of what delegation is for.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from docsgpt.celery_init import celery
|
||||
from docsgpt.vectorstore.embeddings_delegated import EMBED_TASK
|
||||
|
||||
|
||||
@celery.task(name=EMBED_TASK, acks_late=False, ignore_result=False)
|
||||
def embed_texts(texts: List[str], embeddings_name: Optional[str] = None) -> List[List[float]]:
|
||||
"""Embed ``texts`` with the worker's local model.
|
||||
|
||||
Args:
|
||||
texts: Strings to embed.
|
||||
embeddings_name: Model to use; the configured one when omitted.
|
||||
|
||||
Returns:
|
||||
One vector per input, in input order.
|
||||
"""
|
||||
from docsgpt.vectorstore.base import get_embeddings
|
||||
|
||||
return get_embeddings(embeddings_name).embed_documents(list(texts))
|
||||
@@ -0,0 +1,385 @@
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.storage.storage_creator import StorageCreator
|
||||
from docsgpt.vectorstore.base import BaseVectorStore
|
||||
from docsgpt.vectorstore.document_class import Document
|
||||
from docsgpt.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:
|
||||
"""Build a safe local path for a FAISS index.
|
||||
|
||||
Args:
|
||||
path: Source identifier provided by the caller.
|
||||
|
||||
Returns:
|
||||
The validated vectorstore path rooted under ``indexes``.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``path`` escapes the ``indexes`` directory.
|
||||
"""
|
||||
base_dir = "indexes"
|
||||
if not path:
|
||||
return base_dir
|
||||
|
||||
normalized = str(path).strip()
|
||||
if "\\" in normalized:
|
||||
raise ValueError("Invalid source_id path")
|
||||
|
||||
candidate = os.path.normpath(os.path.join(base_dir, normalized))
|
||||
base_abs = os.path.abspath(base_dir)
|
||||
candidate_abs = os.path.abspath(candidate)
|
||||
|
||||
if not candidate_abs.startswith(base_abs + os.sep) and candidate_abs != base_abs:
|
||||
raise ValueError("Invalid source_id path")
|
||||
|
||||
return candidate
|
||||
|
||||
|
||||
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,
|
||||
ids=None,
|
||||
batch_size=None,
|
||||
skip_dimension_check: bool = False,
|
||||
):
|
||||
"""Open or build one source's FAISS index.
|
||||
|
||||
Args:
|
||||
source_id: Source whose index to open.
|
||||
embeddings_key: API key handed to the embeddings provider.
|
||||
docs_init: Documents to build a fresh index from. Loads the stored
|
||||
index instead when omitted.
|
||||
ids: Chunk ids to keep when building. Generated when omitted.
|
||||
batch_size: Documents per embed call when building.
|
||||
skip_dimension_check: Open an index whose width does not match the
|
||||
configured model. Only for a caller that is about to replace
|
||||
that index, such as the re-embed script, which otherwise cannot
|
||||
read the chunks it needs to rebuild from.
|
||||
"""
|
||||
super().__init__()
|
||||
self.source_id = source_id
|
||||
self.path = get_vectorstore(source_id)
|
||||
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._build_from_documents(docs_init, ids=ids, batch_size=batch_size)
|
||||
else:
|
||||
self._load_from_storage()
|
||||
except Exception as e:
|
||||
raise Exception(f"Error loading FAISS index: {str(e)}")
|
||||
|
||||
if not skip_dimension_check:
|
||||
self.assert_embedding_dimensions(self.embeddings)
|
||||
|
||||
# -- Construction ----------------------------------------------------
|
||||
|
||||
def _build_from_documents(self, docs_init, ids=None, batch_size=None) -> None:
|
||||
"""Create a fresh index seeded with ``docs_init``.
|
||||
|
||||
Args:
|
||||
docs_init: Documents to embed.
|
||||
ids: Chunk ids to keep. Generated when omitted, which renumbers
|
||||
every chunk and orphans anything referencing the old ids.
|
||||
batch_size: Documents per embed call. Without it the whole index
|
||||
goes out in one call, which a remote embeddings server rejects
|
||||
or times out on.
|
||||
"""
|
||||
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()
|
||||
ids = list(ids) if ids else None
|
||||
step = batch_size if batch_size and batch_size > 0 else len(texts)
|
||||
for start in range(0, len(texts), max(1, step)):
|
||||
stop = start + max(1, step)
|
||||
vectors = self.embeddings.embed_documents(texts[start:stop])
|
||||
if self.index is None:
|
||||
self.index = faiss.IndexFlatL2(len(vectors[0]))
|
||||
self._append(
|
||||
texts[start:stop],
|
||||
metadatas[start:stop],
|
||||
vectors,
|
||||
ids[start:stop] if ids else None,
|
||||
)
|
||||
|
||||
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,
|
||||
query_vector: Optional[List[float]] = None,
|
||||
**kwargs,
|
||||
) -> List[Tuple[Document, float]]:
|
||||
"""Same search as :meth:`search`, pairing each hit with its L2 distance.
|
||||
|
||||
Args:
|
||||
query_vector: Precomputed embedding of ``question``; when given the
|
||||
store skips embedding the query itself.
|
||||
"""
|
||||
# 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)
|
||||
if self.index is None or self.index.ntotal == 0:
|
||||
return []
|
||||
|
||||
if query_vector is None:
|
||||
query_vector = self.embeddings.embed_query(question)
|
||||
vector = np.array([query_vector], dtype=np.float32)
|
||||
distances, rows = self.index.search(vector, min(k, self.index.ntotal))
|
||||
|
||||
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
|
||||
|
||||
# -- Mutation --------------------------------------------------------
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
missing = set(ids) - set(self.documents)
|
||||
if missing:
|
||||
raise ValueError(f"Chunk ids not found in index: {sorted(missing)}")
|
||||
|
||||
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 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
|
||||
|
||||
# -- Persistence -----------------------------------------------------
|
||||
|
||||
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.
|
||||
|
||||
Each file is replaced atomically by the backend, so an interrupted save
|
||||
leaves the previous one readable. The three are still written in
|
||||
sequence: a crash between them pairs a new index with an older sidecar,
|
||||
which stays consistent because the ids and row order are preserved.
|
||||
"""
|
||||
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.
|
||||
|
||||
This used to run only when ``EMBEDDINGS_NAME`` was mpnet, so every
|
||||
other model skipped the check entirely -- exactly the models most
|
||||
likely to differ from an index built earlier. It now runs for any
|
||||
model that reports a width.
|
||||
"""
|
||||
word_embedding_dimension = getattr(embeddings, "dimension", None)
|
||||
if word_embedding_dimension is None:
|
||||
# A remote model of unknown width reports None until its first
|
||||
# call; there is nothing to compare yet.
|
||||
return
|
||||
if self.index is None:
|
||||
return
|
||||
if word_embedding_dimension != self.index.d:
|
||||
raise ValueError(
|
||||
f"Embedding dimension mismatch: {settings.EMBEDDINGS_NAME} produces "
|
||||
f"{word_embedding_dimension}-dim vectors but this FAISS index is "
|
||||
f"{self.index.d}-dim. The index was built with a different "
|
||||
f"embedding model; re-embed it with "
|
||||
f"`python -m docsgpt.scripts.reembed` or point "
|
||||
f"EMBEDDINGS_NAME back at the original model."
|
||||
)
|
||||
|
||||
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")
|
||||
@@ -0,0 +1,132 @@
|
||||
from typing import List, Optional
|
||||
import importlib
|
||||
from docsgpt.vectorstore.base import BaseVectorStore
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.model_registry import DEFAULT_EMBEDDING_DIMENSION
|
||||
|
||||
class LanceDBVectorStore(BaseVectorStore):
|
||||
"""Class for LanceDB Vector Store integration."""
|
||||
|
||||
def __init__(self, path: str = settings.LANCEDB_PATH,
|
||||
table_name_prefix: str = settings.LANCEDB_TABLE_NAME,
|
||||
source_id: str = None,
|
||||
embeddings_key: str = "embeddings"):
|
||||
"""Initialize the LanceDB vector store."""
|
||||
super().__init__()
|
||||
self.path = path
|
||||
self.table_name = f"{table_name_prefix}_{source_id}" if source_id else table_name_prefix
|
||||
self.embeddings_key = embeddings_key
|
||||
self._lance_db = None
|
||||
self.docsearch = None
|
||||
self._pa = None # PyArrow (pa) will be lazy loaded
|
||||
|
||||
@property
|
||||
def pa(self):
|
||||
"""Lazy load pyarrow module."""
|
||||
if self._pa is None:
|
||||
self._pa = importlib.import_module("pyarrow")
|
||||
return self._pa
|
||||
|
||||
@property
|
||||
def lancedb(self):
|
||||
"""Lazy load lancedb module."""
|
||||
if not hasattr(self, "_lancedb_module"):
|
||||
self._lancedb_module = importlib.import_module("lancedb")
|
||||
return self._lancedb_module
|
||||
|
||||
@property
|
||||
def lance_db(self):
|
||||
"""Lazy load the LanceDB connection."""
|
||||
if self._lance_db is None:
|
||||
self._lance_db = self.lancedb.connect(self.path)
|
||||
return self._lance_db
|
||||
|
||||
@property
|
||||
def table(self):
|
||||
"""Lazy load the LanceDB table."""
|
||||
if self.docsearch is None:
|
||||
if self.table_name in self.lance_db.table_names():
|
||||
self.docsearch = self.lance_db.open_table(self.table_name)
|
||||
else:
|
||||
self.docsearch = None
|
||||
return self.docsearch
|
||||
|
||||
def ensure_table_exists(self):
|
||||
"""Ensure the table exists before performing operations."""
|
||||
if self.table is None:
|
||||
embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key)
|
||||
# A model outside the registry reports no width until it has run;
|
||||
# ``list_size=None`` is a TypeError, not a permissive schema.
|
||||
dimension = getattr(embeddings, "dimension", None) or DEFAULT_EMBEDDING_DIMENSION
|
||||
schema = self.pa.schema([
|
||||
self.pa.field("vector", self.pa.list_(self.pa.float32(), list_size=dimension)),
|
||||
self.pa.field("text", self.pa.string()),
|
||||
self.pa.field("metadata", self.pa.struct([
|
||||
self.pa.field("key", self.pa.string()),
|
||||
self.pa.field("value", self.pa.string())
|
||||
]))
|
||||
])
|
||||
self.docsearch = self.lance_db.create_table(self.table_name, schema=schema)
|
||||
|
||||
def add_texts(self, texts: List[str], metadatas: Optional[List[dict]] = None, source_id: str = None):
|
||||
"""Add texts with metadata and their embeddings to the LanceDB table."""
|
||||
embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key).embed_documents(texts)
|
||||
vectors = []
|
||||
for embedding, text, metadata in zip(embeddings, texts, metadatas or [{}] * len(texts)):
|
||||
if source_id:
|
||||
metadata["source_id"] = source_id
|
||||
metadata_struct = [{"key": k, "value": str(v)} for k, v in metadata.items()]
|
||||
vectors.append({
|
||||
"vector": embedding,
|
||||
"text": text,
|
||||
"metadata": metadata_struct
|
||||
})
|
||||
self.ensure_table_exists()
|
||||
self.docsearch.add(vectors)
|
||||
|
||||
def search(self, query: str, k: int = 2, *args, query_vector=None, **kwargs):
|
||||
"""Search LanceDB for the top k most similar vectors.
|
||||
|
||||
Args:
|
||||
query_vector: Precomputed embedding of ``query``; when given the
|
||||
store skips embedding the query itself.
|
||||
"""
|
||||
self.ensure_table_exists()
|
||||
query_embedding = query_vector
|
||||
if query_embedding is None:
|
||||
query_embedding = self._get_embeddings(
|
||||
settings.EMBEDDINGS_NAME, self.embeddings_key
|
||||
).embed_query(query)
|
||||
results = self.docsearch.search(query_embedding).limit(k).to_list()
|
||||
return [(result["_distance"], result["text"], result["metadata"]) for result in results]
|
||||
|
||||
def delete_index(self):
|
||||
"""Delete the entire LanceDB index (table)."""
|
||||
if self.table:
|
||||
self.lance_db.drop_table(self.table_name)
|
||||
|
||||
def assert_embedding_dimensions(self, embeddings):
|
||||
"""Ensure that embedding dimensions match the table index dimensions."""
|
||||
word_embedding_dimension = embeddings.dimension
|
||||
if self.table:
|
||||
table_index_dimension = len(self.docsearch.schema["vector"].type.value_type)
|
||||
if word_embedding_dimension != table_index_dimension:
|
||||
raise ValueError(
|
||||
f"Embedding dimension mismatch: embeddings.dimension ({word_embedding_dimension}) "
|
||||
f"!= table index dimension ({table_index_dimension})"
|
||||
)
|
||||
|
||||
def filter_documents(self, filter_condition: dict) -> List[dict]:
|
||||
"""Filter documents based on certain conditions."""
|
||||
self.ensure_table_exists()
|
||||
|
||||
# Ensure source_id exists in the filter condition
|
||||
if 'source_id' not in filter_condition:
|
||||
raise ValueError("filter_condition must contain 'source_id'")
|
||||
|
||||
source_id = filter_condition["source_id"]
|
||||
|
||||
# Use LanceDB's native filtering if supported, otherwise filter manually
|
||||
filtered_data = self.docsearch.filter(lambda x: x.metadata and x.metadata.get("source_id") == source_id).to_list()
|
||||
|
||||
return filtered_data
|
||||
@@ -0,0 +1,220 @@
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from docsgpt.core.optional_deps import require
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.base import BaseVectorStore
|
||||
from docsgpt.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__()
|
||||
with _without_milvus_uri_env():
|
||||
pymilvus = require("pymilvus", "VECTOR_STORE=milvus")
|
||||
DataType, MilvusClient = pymilvus.DataType, pymilvus.MilvusClient
|
||||
|
||||
self._DataType = DataType
|
||||
self._source_id = str(source_id).replace("docsgpt/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._ensure_collection()
|
||||
|
||||
@property
|
||||
def _filter(self) -> str:
|
||||
"""Boolean expression scoping every read to this source."""
|
||||
escaped = self._source_id.replace('"', '\\"')
|
||||
return f'source_id == "{escaped}"'
|
||||
|
||||
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, 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,
|
||||
query_vector: Optional[List[float]] = None,
|
||||
**kwargs,
|
||||
) -> List[Tuple[Document, float]]:
|
||||
"""Search, pairing each hit with its cosine similarity.
|
||||
|
||||
Args:
|
||||
query_vector: Precomputed embedding of ``question``; when given the
|
||||
store skips embedding the query itself.
|
||||
"""
|
||||
if query_vector is None:
|
||||
query_vector = self._embeddings.embed_query(question)
|
||||
results = self._client.search(
|
||||
collection_name=self._collection,
|
||||
data=[query_vector],
|
||||
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):
|
||||
"""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
|
||||
@@ -0,0 +1,177 @@
|
||||
"""Canonical description of every embedding model DocsGPT knows how to run.
|
||||
|
||||
``EMBEDDINGS_NAME`` used to be a free-form string interpreted in half a dozen
|
||||
places: a factory dict here, a bundled-model path probe there, a dimension
|
||||
assertion gated on one model's name, and a hardcoded ``dimension = 768`` on the
|
||||
remote client. Each of those encoded a different subset of the same facts, and
|
||||
they drifted.
|
||||
|
||||
This module is the single place those facts live. A model is described once and
|
||||
every consumer -- the local runner, the remote client, the schema bootstrap, the
|
||||
chunker -- reads the same entry.
|
||||
|
||||
Unknown names are not an error: :func:`resolve` returns ``None`` and callers
|
||||
fall back to treating the name as a Hugging Face repository, which is what a
|
||||
user configuring an arbitrary model expects.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EmbeddingModel:
|
||||
"""Everything the application needs to know about one embedding model.
|
||||
|
||||
Attributes:
|
||||
name: Canonical ``EMBEDDINGS_NAME`` value.
|
||||
dimension: Width of the vectors it produces.
|
||||
max_input_tokens: The model's own context window, counted in *its*
|
||||
tokenizer. Used to bound what we send it, never to silently
|
||||
reshape chunks.
|
||||
pooling: ``"cls"`` or ``"mean"`` -- how token vectors become one vector.
|
||||
normalize: Whether outputs are L2-normalised to unit length.
|
||||
provider: Which runner handles it (``"fastembed"`` or ``"openai"``).
|
||||
repo: Hugging Face repository holding weights and tokenizer.
|
||||
onnx_file: Path within ``repo`` to the ONNX graph to run.
|
||||
aliases: Other spellings of ``name`` accepted from configuration.
|
||||
"""
|
||||
|
||||
name: str
|
||||
dimension: int
|
||||
max_input_tokens: int
|
||||
pooling: str = "mean"
|
||||
normalize: bool = True
|
||||
provider: str = "fastembed"
|
||||
repo: Optional[str] = None
|
||||
onnx_file: Optional[str] = None
|
||||
aliases: Tuple[str, ...] = field(default_factory=tuple)
|
||||
|
||||
|
||||
#: The model DocsGPT installed before the granite migration. Kept as the
|
||||
#: default so an existing deployment that upgrades keeps its index working;
|
||||
#: new installs are pointed at granite by the setup script and env template.
|
||||
MPNET = EmbeddingModel(
|
||||
name="huggingface_sentence-transformers/all-mpnet-base-v2",
|
||||
dimension=768,
|
||||
max_input_tokens=384,
|
||||
pooling="mean",
|
||||
normalize=True,
|
||||
repo="sentence-transformers/all-mpnet-base-v2",
|
||||
onnx_file="onnx/model.onnx",
|
||||
aliases=(
|
||||
"huggingface_sentence-transformers-all-mpnet-base-v2",
|
||||
"sentence-transformers/all-mpnet-base-v2",
|
||||
"all-mpnet-base-v2",
|
||||
),
|
||||
)
|
||||
|
||||
#: Default for new installs: same 768 dimensions as mpnet, an 8x wider
|
||||
#: effective input, and multilingual retrieval.
|
||||
#:
|
||||
#: Runs the int8-quantised graph, which is ~4x smaller than the fp32 one
|
||||
#: (313 MB against 1247 MB) and keeps the image shippable now that both
|
||||
#: defaults are baked in. It costs some numeric fidelity against fp32 --
|
||||
#: measured at cosine 0.96 on short texts, with retrieval rank order
|
||||
#: unchanged -- so it is not bit-comparable to a fp32 granite index.
|
||||
GRANITE_311M = EmbeddingModel(
|
||||
name="ibm-granite/granite-embedding-311m-multilingual-r2",
|
||||
dimension=768,
|
||||
max_input_tokens=32768,
|
||||
pooling="cls",
|
||||
normalize=True,
|
||||
repo="ibm-granite/granite-embedding-311m-multilingual-r2",
|
||||
onnx_file="onnx/model_quint8_avx2.onnx",
|
||||
aliases=("granite-embedding-311m-multilingual-r2", "granite-311m"),
|
||||
)
|
||||
|
||||
#: Smaller granite. Half the vector width, roughly three times the speed.
|
||||
#: Int8-quantised on the same terms as GRANITE_311M above.
|
||||
GRANITE_97M = EmbeddingModel(
|
||||
name="ibm-granite/granite-embedding-97m-multilingual-r2",
|
||||
dimension=384,
|
||||
max_input_tokens=32768,
|
||||
pooling="cls",
|
||||
normalize=True,
|
||||
repo="ibm-granite/granite-embedding-97m-multilingual-r2",
|
||||
onnx_file="onnx/model_quint8_avx2.onnx",
|
||||
aliases=("granite-embedding-97m-multilingual-r2", "granite-97m"),
|
||||
)
|
||||
|
||||
OPENAI_ADA_002 = EmbeddingModel(
|
||||
name="openai_text-embedding-ada-002",
|
||||
dimension=1536,
|
||||
max_input_tokens=8191,
|
||||
pooling="mean",
|
||||
normalize=True,
|
||||
provider="openai",
|
||||
repo=None,
|
||||
onnx_file=None,
|
||||
aliases=("text-embedding-ada-002",),
|
||||
)
|
||||
|
||||
MODELS: Tuple[EmbeddingModel, ...] = (
|
||||
MPNET,
|
||||
GRANITE_311M,
|
||||
GRANITE_97M,
|
||||
OPENAI_ADA_002,
|
||||
)
|
||||
|
||||
#: Fallback vector width when the configured model is unknown -- the width
|
||||
#: every DocsGPT install has used to date, so an unrecognised model does not
|
||||
#: silently reshape an existing table.
|
||||
DEFAULT_EMBEDDING_DIMENSION = 768
|
||||
|
||||
#: Name that a fresh install should be configured with.
|
||||
DEFAULT_NEW_INSTALL = GRANITE_311M.name
|
||||
|
||||
#: Name that ``settings.EMBEDDINGS_NAME`` defaults to, i.e. what an existing
|
||||
#: deployment falls back to when it never pinned one.
|
||||
DEFAULT_LEGACY = MPNET.name
|
||||
|
||||
|
||||
def _index() -> Dict[str, EmbeddingModel]:
|
||||
"""Build the lookup table of every accepted spelling."""
|
||||
table: Dict[str, EmbeddingModel] = {}
|
||||
for model in MODELS:
|
||||
for key in (model.name, *model.aliases):
|
||||
table[key.lower()] = model
|
||||
return table
|
||||
|
||||
|
||||
_LOOKUP = _index()
|
||||
|
||||
|
||||
def resolve(name: Optional[str]) -> Optional[EmbeddingModel]:
|
||||
"""Return the registry entry for ``name``, or ``None`` when unknown.
|
||||
|
||||
Args:
|
||||
name: A configured ``EMBEDDINGS_NAME`` value, in any accepted spelling.
|
||||
|
||||
Returns:
|
||||
The matching :class:`EmbeddingModel`, or ``None`` for a name the
|
||||
registry does not describe -- which callers treat as a Hugging Face
|
||||
repository rather than an error.
|
||||
"""
|
||||
if not name:
|
||||
return None
|
||||
return _LOOKUP.get(name.strip().lower())
|
||||
|
||||
|
||||
def dimension_for(name: Optional[str]) -> Optional[int]:
|
||||
"""Vector width for ``name``, or ``None`` when unknown."""
|
||||
model = resolve(name)
|
||||
return model.dimension if model else None
|
||||
|
||||
|
||||
def max_input_tokens_for(name: Optional[str]) -> Optional[int]:
|
||||
"""Context window for ``name``, or ``None`` when unknown."""
|
||||
model = resolve(name)
|
||||
return model.max_input_tokens if model else None
|
||||
|
||||
|
||||
def known_names() -> Tuple[str, ...]:
|
||||
"""Canonical names of every registered model, for error messages."""
|
||||
return tuple(model.name for model in MODELS)
|
||||
@@ -0,0 +1,236 @@
|
||||
import logging
|
||||
from functools import cached_property
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.base import BaseVectorStore
|
||||
from docsgpt.vectorstore.document_class import Document
|
||||
|
||||
|
||||
def _lazy_import_pymongo():
|
||||
"""Lazy import of pymongo so installations that don't use the MongoDB vectorstore don't need it."""
|
||||
try:
|
||||
import pymongo
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Could not import pymongo python package. "
|
||||
"Please install it with `pip install pymongo`."
|
||||
) from exc
|
||||
return pymongo
|
||||
|
||||
|
||||
class MongoDBVectorStore(BaseVectorStore):
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str = "",
|
||||
embeddings_key: str = "embeddings",
|
||||
collection: str = "documents",
|
||||
index_name: str = "vector_search_index",
|
||||
text_key: str = "text",
|
||||
embedding_key: str = "embedding",
|
||||
database: str = "docsgpt",
|
||||
):
|
||||
self._index_name = index_name
|
||||
self._text_key = text_key
|
||||
self._embedding_key = embedding_key
|
||||
self._embeddings_key = embeddings_key
|
||||
self._mongo_uri = settings.MONGO_URI
|
||||
self._database_name = database
|
||||
self._collection_name = collection
|
||||
self._source_id = source_id.replace("docsgpt/indexes/", "").rstrip("/")
|
||||
self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key)
|
||||
|
||||
@cached_property
|
||||
def _client(self):
|
||||
pymongo = _lazy_import_pymongo()
|
||||
return pymongo.MongoClient(self._mongo_uri)
|
||||
|
||||
@cached_property
|
||||
def _database(self):
|
||||
return self._client[self._database_name]
|
||||
|
||||
@cached_property
|
||||
def _collection(self):
|
||||
return self._database[self._collection_name]
|
||||
|
||||
score_kind = "cosine_similarity"
|
||||
|
||||
def search(
|
||||
self, question, k=2, *args, score_threshold=None, query_vector=None, **kwargs
|
||||
):
|
||||
"""Search via Atlas ``$vectorSearch``.
|
||||
|
||||
Args:
|
||||
question: The query string.
|
||||
k: Maximum number of results.
|
||||
score_threshold: Optional ``vectorSearchScore`` floor in ``[0, 1]``;
|
||||
results scoring below it are dropped.
|
||||
query_vector: Precomputed embedding of ``question``, so a caller
|
||||
searching several sources embeds the query only once.
|
||||
"""
|
||||
return [
|
||||
doc
|
||||
for doc, _ in self.search_with_scores(
|
||||
question,
|
||||
k,
|
||||
*args,
|
||||
score_threshold=score_threshold,
|
||||
query_vector=query_vector,
|
||||
**kwargs,
|
||||
)
|
||||
]
|
||||
|
||||
def search_with_scores(
|
||||
self, question, k=2, *args, score_threshold=None, query_vector=None, **kwargs
|
||||
):
|
||||
"""Same search as :meth:`search`, pairing each hit with its score.
|
||||
|
||||
The score is Atlas' ``vectorSearchScore`` — the same quantity
|
||||
``score_threshold`` is compared against. ``query_vector`` skips the
|
||||
per-store query embedding when the caller already has one.
|
||||
"""
|
||||
if query_vector is None:
|
||||
query_vector = self._embedding.embed_query(question)
|
||||
|
||||
pipeline = [
|
||||
{
|
||||
"$vectorSearch": {
|
||||
"queryVector": query_vector,
|
||||
"path": self._embedding_key,
|
||||
"limit": k,
|
||||
"numCandidates": k * 10,
|
||||
"index": self._index_name,
|
||||
"filter": {"source_id": {"$eq": self._source_id}},
|
||||
}
|
||||
},
|
||||
{"$addFields": {"_score": {"$meta": "vectorSearchScore"}}},
|
||||
]
|
||||
if score_threshold is not None:
|
||||
pipeline.append({"$match": {"_score": {"$gte": float(score_threshold)}}})
|
||||
|
||||
cursor = self._collection.aggregate(pipeline)
|
||||
|
||||
results = []
|
||||
for doc in cursor:
|
||||
text = doc[self._text_key]
|
||||
doc.pop("_id")
|
||||
doc.pop(self._text_key)
|
||||
doc.pop(self._embedding_key)
|
||||
score = doc.pop("_score", None)
|
||||
metadata = doc
|
||||
results.append(
|
||||
(Document(text, metadata), None if score is None else float(score))
|
||||
)
|
||||
return results
|
||||
|
||||
def _insert_texts(self, texts, metadatas):
|
||||
if not texts:
|
||||
return []
|
||||
embeddings = self._embedding.embed_documents(texts)
|
||||
|
||||
to_insert = [
|
||||
{self._text_key: t, self._embedding_key: embedding, **m}
|
||||
for t, m, embedding in zip(texts, metadatas, embeddings)
|
||||
]
|
||||
|
||||
insert_result = self._collection.insert_many(to_insert)
|
||||
return insert_result.inserted_ids
|
||||
|
||||
def add_texts(
|
||||
self,
|
||||
texts,
|
||||
metadatas=None,
|
||||
ids=None,
|
||||
refresh_indices=True,
|
||||
create_index_if_not_exists=True,
|
||||
bulk_kwargs=None,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
# dims = self._embedding.client[1].word_embedding_dimension
|
||||
# # check if index exists
|
||||
# if create_index_if_not_exists:
|
||||
# # check if index exists
|
||||
# info = self._collection.index_information()
|
||||
# if self._index_name not in info:
|
||||
# index_mongo = {
|
||||
# "fields": [{
|
||||
# "type": "vector",
|
||||
# "path": self._embedding_key,
|
||||
# "numDimensions": dims,
|
||||
# "similarity": "cosine",
|
||||
# },
|
||||
# {
|
||||
# "type": "filter",
|
||||
# "path": "store"
|
||||
# }]
|
||||
# }
|
||||
# self._collection.create_index(self._index_name, index_mongo)
|
||||
|
||||
batch_size = 100
|
||||
_metadatas = metadatas or ({} for _ in texts)
|
||||
texts_batch = []
|
||||
metadatas_batch = []
|
||||
result_ids = []
|
||||
for i, (text, metadata) in enumerate(zip(texts, _metadatas)):
|
||||
texts_batch.append(text)
|
||||
metadatas_batch.append(metadata)
|
||||
if (i + 1) % batch_size == 0:
|
||||
result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))
|
||||
texts_batch = []
|
||||
metadatas_batch = []
|
||||
if texts_batch:
|
||||
result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))
|
||||
return result_ids
|
||||
|
||||
def delete_index(self, *args, **kwargs):
|
||||
self._collection.delete_many({"source_id": self._source_id})
|
||||
|
||||
def get_chunks(self):
|
||||
try:
|
||||
chunks = []
|
||||
cursor = self._collection.find({"source_id": self._source_id})
|
||||
for doc in cursor:
|
||||
doc_id = str(doc.get("_id"))
|
||||
text = doc.get(self._text_key)
|
||||
metadata = {
|
||||
k: v
|
||||
for k, v in doc.items()
|
||||
if k
|
||||
not in ["_id", self._text_key, self._embedding_key, "source_id"]
|
||||
}
|
||||
|
||||
if text:
|
||||
chunks.append(
|
||||
{"doc_id": doc_id, "text": text, "metadata": metadata}
|
||||
)
|
||||
|
||||
return chunks
|
||||
except Exception as e:
|
||||
logging.error(f"Error getting chunks: {e}", exc_info=True)
|
||||
return []
|
||||
|
||||
def add_chunk(self, text, metadata=None):
|
||||
metadata = metadata or {}
|
||||
embeddings = self._embedding.embed_documents([text])
|
||||
if not embeddings:
|
||||
raise ValueError("Could not generate embedding for chunk")
|
||||
|
||||
chunk_data = {
|
||||
self._text_key: text,
|
||||
self._embedding_key: embeddings[0],
|
||||
"source_id": self._source_id,
|
||||
**metadata,
|
||||
}
|
||||
result = self._collection.insert_one(chunk_data)
|
||||
return str(result.inserted_id)
|
||||
|
||||
def delete_chunk(self, chunk_id):
|
||||
try:
|
||||
from bson.objectid import ObjectId
|
||||
|
||||
object_id = ObjectId(chunk_id)
|
||||
result = self._collection.delete_one({"_id": object_id})
|
||||
return result.deleted_count > 0
|
||||
except Exception as e:
|
||||
logging.error(f"Error deleting chunk: {e}", exc_info=True)
|
||||
return False
|
||||
@@ -0,0 +1,135 @@
|
||||
"""Process-wide psycopg connection pools keyed by DSN.
|
||||
|
||||
Every component that talks to the pgvector database — the vector store and the
|
||||
GraphRAG graph store — checks out of the *same* pool for a given DSN, so a
|
||||
request that touches N sources costs pooled checkouts instead of N fresh TCP +
|
||||
TLS + auth handshakes. Keeping the registry in its own module (rather than in
|
||||
``pgvector``) is what lets ``docsgpt.graphrag.store`` join it without
|
||||
importing the vector store.
|
||||
|
||||
The boot-time schema hook deliberately does not use this: pools are built
|
||||
lazily on first checkout, i.e. after the fork in a Celery/gunicorn worker, so a
|
||||
preloading parent never hands a live socket to a child.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Dict
|
||||
|
||||
# Seconds a caller waits for a free pooled connection before giving up.
|
||||
POOL_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
# Physical connections one process may hold per DSN. Lives here, not in either
|
||||
# store: ``pool_for`` keys its registry by DSN and applies ``max_size`` only on
|
||||
# creation, so two stores disagreeing about the default would make the
|
||||
# effective pool size depend on which of them touched the DSN first.
|
||||
DEFAULT_POOL_MAX_SIZE = 8
|
||||
|
||||
# dsn -> psycopg_pool.ConnectionPool, one per process.
|
||||
_POOLS: Dict[str, Any] = {}
|
||||
_POOLS_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def configure_pooled_connection(conn) -> None:
|
||||
"""Register pgvector's type adapters once per physical pooled connection.
|
||||
|
||||
Tolerates a database where ``CREATE EXTENSION vector`` has not run yet:
|
||||
``register_vector`` looks the type up in the catalog and raises when it is
|
||||
absent, which would leave the pool unable to open the very connection the
|
||||
schema bootstrap needs. The write path re-registers once the type exists.
|
||||
"""
|
||||
from pgvector.psycopg import register_vector
|
||||
|
||||
try:
|
||||
register_vector(conn)
|
||||
except Exception as e:
|
||||
logging.debug("pgvector types not registered yet: %s", e)
|
||||
|
||||
|
||||
def resolve_pool_max_size() -> int:
|
||||
"""Pool size from settings, defensively — 0 means one direct connection.
|
||||
|
||||
Returns:
|
||||
int: ``PGVECTOR_POOL_MAX_SIZE`` when it is a non-negative, non-bool
|
||||
int, else :data:`DEFAULT_POOL_MAX_SIZE`. The unit suite replaces
|
||||
``settings`` with a MagicMock, whose attributes must never become
|
||||
a pool size.
|
||||
"""
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
value = getattr(settings, "PGVECTOR_POOL_MAX_SIZE", DEFAULT_POOL_MAX_SIZE)
|
||||
if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
|
||||
return value
|
||||
return DEFAULT_POOL_MAX_SIZE
|
||||
|
||||
|
||||
def pool_for(dsn: str, max_size: int):
|
||||
"""Return this process's connection pool for ``dsn``, creating it lazily.
|
||||
|
||||
The pool is keyed by DSN alone and is built on first checkout — i.e. after
|
||||
the fork in a Celery/gunicorn worker, so a preloading parent never hands a
|
||||
live socket to a child. The boot-time schema hook must therefore use a
|
||||
direct connection rather than this pool.
|
||||
|
||||
Args:
|
||||
dsn: libpq connection string; the pool's only identity.
|
||||
max_size: Upper bound on physical connections, used on creation only.
|
||||
|
||||
Returns:
|
||||
psycopg_pool.ConnectionPool: The shared pool for ``dsn``.
|
||||
"""
|
||||
pool = _POOLS.get(dsn)
|
||||
if pool is not None:
|
||||
return pool
|
||||
with _POOLS_LOCK:
|
||||
pool = _POOLS.get(dsn)
|
||||
if pool is None:
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
pool = ConnectionPool(
|
||||
conninfo=dsn,
|
||||
min_size=1,
|
||||
max_size=max_size,
|
||||
open=True,
|
||||
timeout=POOL_TIMEOUT_SECONDS,
|
||||
name="docsgpt-pgvector",
|
||||
configure=configure_pooled_connection,
|
||||
check=ConnectionPool.check_connection,
|
||||
)
|
||||
_POOLS[dsn] = pool
|
||||
return pool
|
||||
|
||||
|
||||
def release(dsn: str, conn, pooled: bool) -> None:
|
||||
"""Hand ``conn`` back: to ``dsn``'s pool when pooled, else close it.
|
||||
|
||||
A pooled connection is rolled back first when it is still in a transaction,
|
||||
so the next borrower gets a clean session. Never raises — releasing is
|
||||
always best-effort cleanup on a path that is often a destructor.
|
||||
|
||||
Args:
|
||||
dsn: Connection string the connection was opened for.
|
||||
conn: The psycopg connection being released; ``None`` is a no-op.
|
||||
pooled: Whether the connection came from :func:`pool_for`.
|
||||
"""
|
||||
if conn is None:
|
||||
return
|
||||
try:
|
||||
if pooled:
|
||||
try:
|
||||
if conn.info.transaction_status.name != "IDLE":
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
# Connection already broken; the pool discards it.
|
||||
pass
|
||||
pool = _POOLS.get(dsn)
|
||||
if pool is not None:
|
||||
pool.putconn(conn)
|
||||
elif not conn.closed:
|
||||
conn.close()
|
||||
elif not conn.closed:
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
logging.debug("Error releasing pooled connection: %s", e)
|
||||
@@ -0,0 +1,773 @@
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
from typing import List, Optional, Any, Dict
|
||||
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore import pgconn
|
||||
from docsgpt.vectorstore.base import BaseVectorStore
|
||||
from docsgpt.vectorstore.document_class import Document
|
||||
|
||||
# table name -> IVFFlat ``lists`` (None when the table has no such index)
|
||||
_IVFFLAT_LISTS_CACHE: Dict[str, Optional[int]] = {}
|
||||
|
||||
DEFAULT_EMBEDDING_DIM = 768
|
||||
# Advisory-lock key shared with the boot hook so concurrent workers serialize DDL.
|
||||
SCHEMA_LOCK_KEY = "docsgpt:vectors:ddl"
|
||||
|
||||
# The connection pools moved to ``docsgpt.vectorstore.pgconn`` so the graph
|
||||
# store can share them without importing this module. Only the names something
|
||||
# actually reaches through *this* module stay bound here — same objects, not
|
||||
# copies: callers and tests read ``pgvector._POOLS``, patch
|
||||
# ``pgvector._pool_for``, and assert on ``pgvector.DEFAULT_POOL_MAX_SIZE``.
|
||||
# Anything else belongs to ``pgconn`` alone; re-exporting it here would just be
|
||||
# a second name to keep in sync.
|
||||
DEFAULT_POOL_MAX_SIZE = pgconn.DEFAULT_POOL_MAX_SIZE
|
||||
_POOLS = pgconn._POOLS
|
||||
_pool_for = pgconn.pool_for
|
||||
|
||||
|
||||
class PGVectorStore(BaseVectorStore):
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str = "",
|
||||
embeddings_key: str = "embeddings",
|
||||
table_name: str = "documents",
|
||||
decoded_token: Optional[str] = None,
|
||||
vector_column: str = "embedding",
|
||||
text_column: str = "text",
|
||||
metadata_column: str = "metadata",
|
||||
connection_string: str = None,
|
||||
):
|
||||
super().__init__()
|
||||
# Store the source_id for use in add_chunk
|
||||
self._source_id = str(source_id).replace("docsgpt/indexes/", "").rstrip("/")
|
||||
self._embeddings_key = embeddings_key
|
||||
self._table_name = table_name
|
||||
self._vector_column = vector_column
|
||||
self._text_column = text_column
|
||||
self._metadata_column = metadata_column
|
||||
self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key)
|
||||
|
||||
# Use provided connection string or fall back to settings.
|
||||
# If PGVECTOR_CONNECTION_STRING is not set but POSTGRES_URI is,
|
||||
# reuse the same cluster — normalize from SQLAlchemy dialect to libpq form.
|
||||
self._connection_string = connection_string or getattr(settings, 'PGVECTOR_CONNECTION_STRING', None)
|
||||
|
||||
if not self._connection_string and getattr(settings, 'POSTGRES_URI', None):
|
||||
from docsgpt.core.db_uri import normalize_pgvector_connection_string
|
||||
self._connection_string = normalize_pgvector_connection_string(settings.POSTGRES_URI)
|
||||
|
||||
if not self._connection_string:
|
||||
raise ValueError(
|
||||
"PostgreSQL connection string is required. "
|
||||
"Set PGVECTOR_CONNECTION_STRING or POSTGRES_URI in settings, "
|
||||
"or pass connection_string parameter."
|
||||
)
|
||||
|
||||
try:
|
||||
import psycopg
|
||||
from pgvector.psycopg import register_vector
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Could not import required packages. "
|
||||
"Please install with `pip install 'psycopg[binary,pool]' pgvector`."
|
||||
)
|
||||
|
||||
self._psycopg = psycopg
|
||||
self._register_vector = register_vector
|
||||
self._connection = None
|
||||
self._pooled = False
|
||||
self._pool_max_size = self._resolve_pool_max_size()
|
||||
# No DDL here. The retriever builds one store per source per request, so
|
||||
# construction must stay free: schema is owned at boot by
|
||||
# ``ensure_vector_schema``, and the write path below re-checks it once as
|
||||
# a safety net for processes that never ran the boot hook.
|
||||
self._schema_ensured = False
|
||||
|
||||
# Shared with ``GraphStore`` via pgconn so the two can never disagree.
|
||||
_resolve_pool_max_size = staticmethod(pgconn.resolve_pool_max_size)
|
||||
|
||||
def _get_connection(self):
|
||||
"""Get or create this store's connection, pooled unless pooling is off."""
|
||||
if self._connection is not None and self._connection.closed:
|
||||
# Hand the dead connection back before replacing it: psycopg_pool
|
||||
# never reclaims a checkout that is not returned, so dropping it
|
||||
# costs the pool a slot for the life of the process.
|
||||
self.close()
|
||||
if self._connection is None:
|
||||
if self._pool_max_size > 0:
|
||||
self._connection = _pool_for(
|
||||
self._connection_string, self._pool_max_size
|
||||
).getconn()
|
||||
self._pooled = True
|
||||
else:
|
||||
self._connection = self._psycopg.connect(self._connection_string)
|
||||
# Register pgvector types
|
||||
self._register_pgvector_types(self._connection)
|
||||
self._pooled = False
|
||||
self._apply_probes_once(self._connection)
|
||||
return self._connection
|
||||
|
||||
def _register_pgvector_types(self, conn) -> None:
|
||||
"""Register pgvector's adapters, tolerating a not-yet-created extension.
|
||||
|
||||
Kept as a method rather than calling :func:`pgconn.configure_pooled_connection`
|
||||
directly because ``self._register_vector`` is the seam the tests patch.
|
||||
"""
|
||||
try:
|
||||
self._register_vector(conn)
|
||||
except Exception as e:
|
||||
logging.debug("pgvector types not registered yet: %s", e)
|
||||
|
||||
def _apply_probes_once(self, conn) -> None:
|
||||
"""Set ``ivfflat.probes`` once per physical connection and table.
|
||||
|
||||
A pooled connection outlives the store that checked it out, and the
|
||||
``SET`` is session-level, so the probe lookup must not repeat on every
|
||||
checkout. The marker lives on the connection object itself because that
|
||||
is what the pool recycles.
|
||||
"""
|
||||
if getattr(conn, "_docsgpt_probes_table", None) == self._table_name:
|
||||
return
|
||||
self._apply_ivfflat_probes(conn)
|
||||
try:
|
||||
# ``SET`` (not SET LOCAL) is undone by a rollback, so make it stick
|
||||
# before any query runs on this connection.
|
||||
conn.commit()
|
||||
conn._docsgpt_probes_table = self._table_name
|
||||
except Exception as e: # never let session tuning break a query
|
||||
logging.debug("Could not persist ivfflat.probes: %s", e)
|
||||
|
||||
def _ivfflat_lists(self, conn) -> Optional[int]:
|
||||
"""Return the ``lists`` value of this table's IVFFlat index, if any.
|
||||
|
||||
Cached per table because it only changes when the index is rebuilt.
|
||||
"""
|
||||
if self._table_name in _IVFFLAT_LISTS_CACHE:
|
||||
return _IVFFLAT_LISTS_CACHE[self._table_name]
|
||||
lists = None
|
||||
try:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"SELECT indexdef FROM pg_indexes "
|
||||
"WHERE tablename = %s AND indexdef ILIKE %s",
|
||||
(self._table_name, "%ivfflat%"),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if row:
|
||||
match = re.search(r"lists\s*=\s*'?(\d+)", row[0])
|
||||
if match:
|
||||
lists = int(match.group(1))
|
||||
except Exception as e: # index introspection must never break search
|
||||
logging.debug("Could not read IVFFlat lists for %s: %s", self._table_name, e)
|
||||
# Only cache a hit: an index may be created after this process booted,
|
||||
# and caching None would keep probes unset for the process's lifetime.
|
||||
if lists:
|
||||
_IVFFLAT_LISTS_CACHE[self._table_name] = lists
|
||||
return lists
|
||||
|
||||
def _apply_ivfflat_probes(self, conn) -> None:
|
||||
"""Raise ``ivfflat.probes`` so a filtered search cannot come back empty.
|
||||
|
||||
An IVFFlat index partitions vectors into ``lists`` clusters and the
|
||||
default ``probes = 1`` scans exactly one of them. Our searches filter by
|
||||
``source_id`` *after* the index picks candidates, so with one probe the
|
||||
candidates frequently all belong to other sources and the query returns
|
||||
nothing — retrieval reports zero documents and the model answers with no
|
||||
source material, silently. ``sqrt(lists)`` is pgvector's own recall
|
||||
guidance and costs a proportional amount of scan.
|
||||
"""
|
||||
probes = settings.PGVECTOR_IVFFLAT_PROBES
|
||||
if probes is None:
|
||||
lists = self._ivfflat_lists(conn)
|
||||
if not lists:
|
||||
return
|
||||
probes = max(1, math.isqrt(lists))
|
||||
try:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(f"SET ivfflat.probes = {int(probes)};")
|
||||
except Exception as e: # older pgvector / no index — search still works
|
||||
logging.debug("Could not set ivfflat.probes: %s", e)
|
||||
|
||||
@staticmethod
|
||||
def create_schema(
|
||||
conn,
|
||||
*,
|
||||
table_name: str = "documents",
|
||||
vector_column: str = "embedding",
|
||||
text_column: str = "text",
|
||||
metadata_column: str = "metadata",
|
||||
dimension: int = DEFAULT_EMBEDDING_DIM,
|
||||
) -> None:
|
||||
"""Create the extension, table and indexes on ``conn`` without committing.
|
||||
|
||||
Shared by the boot hook (``ensure_vector_schema``) and the store's own
|
||||
write-path safety net; the caller owns the transaction.
|
||||
|
||||
Args:
|
||||
conn: Open psycopg connection.
|
||||
table_name: Documents table to create.
|
||||
vector_column: Embedding column name.
|
||||
text_column: Chunk-text column name.
|
||||
metadata_column: JSONB metadata column name.
|
||||
dimension: Width of the embedding vectors.
|
||||
"""
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
# Enable pgvector extension
|
||||
cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
||||
|
||||
# Create table with vector column
|
||||
create_table_query = f"""
|
||||
CREATE TABLE IF NOT EXISTS {table_name} (
|
||||
id SERIAL PRIMARY KEY,
|
||||
{text_column} TEXT NOT NULL,
|
||||
{vector_column} vector({dimension}),
|
||||
{metadata_column} JSONB,
|
||||
source_id TEXT NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
"""
|
||||
cursor.execute(create_table_query)
|
||||
|
||||
# NO vector index is created here, deliberately.
|
||||
#
|
||||
# This runs when the table is first created, i.e. on an EMPTY
|
||||
# table. IVFFlat computes its cluster centroids at build time, so
|
||||
# an index built on no rows gets random centroids and never
|
||||
# recovers — measured recall 0.06 once 5k rows are added. Combined
|
||||
# with our ``WHERE source_id = ...`` post-filter, that returned
|
||||
# ZERO rows for sources with hundreds of chunks: retrieval reported
|
||||
# no documents and the model answered from nothing, silently.
|
||||
# pgvector only warns when sampled_rows < lists, so the common bad
|
||||
# case is silent.
|
||||
#
|
||||
# Exact search is correct and fast well past the sizes most
|
||||
# deployments ever reach. Add an index deliberately, sized to real
|
||||
# data, once a corpus is large enough to need one.
|
||||
# Create index for source_id filtering
|
||||
source_index_query = f"""
|
||||
CREATE INDEX IF NOT EXISTS {table_name}_source_id_idx
|
||||
ON {table_name} (source_id);
|
||||
"""
|
||||
cursor.execute(source_index_query)
|
||||
|
||||
# Functional GIN index backing keyword_search full-text queries.
|
||||
fts_index_query = f"""
|
||||
CREATE INDEX IF NOT EXISTS {table_name}_text_fts_idx
|
||||
ON {table_name} USING gin(to_tsvector('english', {text_column}));
|
||||
"""
|
||||
cursor.execute(fts_index_query)
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
@staticmethod
|
||||
def table_dimension(
|
||||
conn, table_name: str = "documents", vector_column: str = "embedding"
|
||||
) -> Optional[int]:
|
||||
"""Width declared by the table's vector column, or ``None`` if unknown.
|
||||
|
||||
``None`` means the table is absent or the column is not a ``vector``.
|
||||
"""
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute("SELECT to_regclass(%s);", (table_name,))
|
||||
row = cursor.fetchone()
|
||||
if not row or row[0] is None:
|
||||
return None
|
||||
cursor.execute(
|
||||
"SELECT format_type(a.atttypid, a.atttypmod) FROM pg_attribute a "
|
||||
"WHERE a.attrelid = %s::regclass AND a.attname = %s "
|
||||
"AND NOT a.attisdropped",
|
||||
(table_name, vector_column),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if not row or not row[0]:
|
||||
return None
|
||||
match = re.search(r"vector\((\d+)\)", str(row[0]))
|
||||
return int(match.group(1)) if match else None
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def _ensure_table_exists(self) -> None:
|
||||
"""Create this store's schema under an advisory lock, then commit."""
|
||||
conn = self._get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"SELECT pg_advisory_xact_lock(hashtext(%s));", (SCHEMA_LOCK_KEY,)
|
||||
)
|
||||
finally:
|
||||
cursor.close()
|
||||
self.create_schema(
|
||||
conn,
|
||||
table_name=self._table_name,
|
||||
vector_column=self._vector_column,
|
||||
text_column=self._text_column,
|
||||
metadata_column=self._metadata_column,
|
||||
dimension=getattr(self._embedding, "dimension", DEFAULT_EMBEDDING_DIM),
|
||||
)
|
||||
conn.commit()
|
||||
# The extension may have just been created; pick up its adapters so
|
||||
# the insert that follows can bind a vector.
|
||||
self._register_pgvector_types(conn)
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logging.error(f"Error creating table: {e}")
|
||||
raise
|
||||
|
||||
def _ensure_schema_once(self) -> None:
|
||||
"""Create the schema on this instance's first write, at most once.
|
||||
|
||||
Readers never create schema — boot owns it. This is the safety net for a
|
||||
process that never ran the boot hook (scripts, tests, the first ingest on
|
||||
a brand-new deployment); a read against a missing table surfaces
|
||||
psycopg's own error, like every other store.
|
||||
"""
|
||||
if self._schema_ensured:
|
||||
return
|
||||
self._ensure_table_exists()
|
||||
self._schema_ensured = True
|
||||
|
||||
score_kind = "cosine_similarity"
|
||||
|
||||
def search(
|
||||
self,
|
||||
question: str,
|
||||
k: int = 2,
|
||||
*args,
|
||||
score_threshold: float = None,
|
||||
query_vector: Optional[List[float]] = None,
|
||||
**kwargs,
|
||||
) -> List[Document]:
|
||||
"""Search for similar documents using vector similarity.
|
||||
|
||||
Args:
|
||||
question: The query string.
|
||||
k: Maximum number of results.
|
||||
score_threshold: Optional cosine-similarity floor in ``[0, 1]``.
|
||||
Cosine distance = ``1 - similarity``; rows with similarity below
|
||||
the threshold (distance above ``1 - threshold``) are dropped.
|
||||
query_vector: Precomputed embedding of ``question``. Supplied by a
|
||||
caller searching several sources with one query, so the query is
|
||||
embedded once instead of once per store.
|
||||
"""
|
||||
return [
|
||||
doc
|
||||
for doc, _ in self.search_with_scores(
|
||||
question,
|
||||
k,
|
||||
*args,
|
||||
score_threshold=score_threshold,
|
||||
query_vector=query_vector,
|
||||
**kwargs,
|
||||
)
|
||||
]
|
||||
|
||||
def _nearest_sql(self) -> str:
|
||||
"""Build the nearest-neighbour SELECT for this store's table.
|
||||
|
||||
Identifiers (table/column names) come from this instance's
|
||||
construction, never from a request, so they cannot be interpolated by a
|
||||
caller; the query *values* are always bound parameters.
|
||||
"""
|
||||
return (
|
||||
f"SELECT {self._text_column}, {self._metadata_column}, "
|
||||
f"({self._vector_column} <=> %s::vector) AS distance "
|
||||
f"FROM {self._table_name} "
|
||||
"WHERE source_id = %s "
|
||||
f"ORDER BY {self._vector_column} <=> %s::vector "
|
||||
"LIMIT %s;"
|
||||
)
|
||||
|
||||
def _exact_search(self, cursor, query_vector, k: int, ann_results: list) -> list:
|
||||
"""Redo a short indexed search exactly, when the source has more rows.
|
||||
|
||||
Args:
|
||||
cursor: Open cursor on the search connection.
|
||||
query_vector: The embedded query.
|
||||
k: Requested top-k.
|
||||
ann_results: What the indexed search returned.
|
||||
|
||||
Returns:
|
||||
list: Exact rows when the indexed search under-returned, otherwise
|
||||
``ann_results`` unchanged.
|
||||
"""
|
||||
try:
|
||||
cursor.execute(
|
||||
f"SELECT count(*) FROM {self._table_name} WHERE source_id = %s",
|
||||
(self._source_id,),
|
||||
)
|
||||
available = cursor.fetchone()[0]
|
||||
if len(ann_results) >= min(k, available):
|
||||
return ann_results
|
||||
|
||||
cursor.execute("SET LOCAL enable_indexscan = off;")
|
||||
cursor.execute("SET LOCAL enable_bitmapscan = off;")
|
||||
cursor.execute(
|
||||
self._nearest_sql(),
|
||||
(query_vector, self._source_id, query_vector, k),
|
||||
)
|
||||
exact = cursor.fetchall()
|
||||
if len(exact) > len(ann_results):
|
||||
logging.info(
|
||||
"Vector index under-returned for source %s (%d of %d); "
|
||||
"used exact search instead.",
|
||||
self._source_id, len(ann_results), min(k, available),
|
||||
)
|
||||
return exact
|
||||
return ann_results
|
||||
except Exception as e:
|
||||
# Never let the safety net take down the search it is protecting —
|
||||
# but roll back, or the aborted transaction poisons the connection
|
||||
# and every later search on this store returns nothing.
|
||||
logging.warning("Exact-search fallback failed: %s", e)
|
||||
try:
|
||||
cursor.connection.rollback()
|
||||
except Exception:
|
||||
# Connection already gone; nothing left to roll back.
|
||||
pass
|
||||
return ann_results
|
||||
finally:
|
||||
try:
|
||||
# RESET, not "= on": a deployment may disable these globally.
|
||||
cursor.execute("RESET enable_indexscan;")
|
||||
cursor.execute("RESET enable_bitmapscan;")
|
||||
except Exception:
|
||||
# Cursor/transaction already unusable; the settings die with it.
|
||||
pass
|
||||
|
||||
def search_with_scores(
|
||||
self,
|
||||
question: str,
|
||||
k: int = 2,
|
||||
*args,
|
||||
score_threshold: float = None,
|
||||
query_vector: Optional[List[float]] = None,
|
||||
**kwargs,
|
||||
) -> List[tuple]:
|
||||
"""Same search as :meth:`search`, pairing each hit with its similarity.
|
||||
|
||||
The score is the cosine similarity (``1 - cosine_distance``) — the exact
|
||||
quantity ``score_threshold`` is compared against, so a caller can read a
|
||||
result's score and pick a threshold from it directly.
|
||||
|
||||
Args:
|
||||
query_vector: Precomputed embedding of ``question``; when given the
|
||||
store skips embedding the query itself.
|
||||
"""
|
||||
if query_vector is None:
|
||||
query_vector = self._embedding.embed_query(question)
|
||||
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
# Use cosine distance for similarity search with proper vector formatting
|
||||
search_query = self._nearest_sql()
|
||||
|
||||
cursor.execute(search_query, (query_vector, self._source_id, query_vector, k))
|
||||
results = cursor.fetchall()
|
||||
|
||||
# An ANN index filters ``source_id`` *after* choosing candidates, so
|
||||
# a source holding a small share of the table can come back short —
|
||||
# or empty — no matter how the index is tuned. Raising probes /
|
||||
# ef_search only moves that threshold. When the result looks short,
|
||||
# redo the query exactly: correctness is worth one extra scan, and a
|
||||
# silent empty result reaches the model as "no documents exist".
|
||||
if len(results) < k:
|
||||
results = self._exact_search(cursor, query_vector, k, results)
|
||||
|
||||
max_distance = None
|
||||
if score_threshold is not None:
|
||||
max_distance = 1.0 - float(score_threshold)
|
||||
|
||||
documents = []
|
||||
for text, metadata, distance in results:
|
||||
if max_distance is not None and distance is not None and distance > max_distance:
|
||||
continue
|
||||
metadata = metadata or {}
|
||||
score = None if distance is None else 1.0 - float(distance)
|
||||
documents.append(
|
||||
(Document(page_content=text, metadata=metadata), score)
|
||||
)
|
||||
|
||||
# End the read transaction. On a persistent (pooled) connection an
|
||||
# uncommitted SELECT leaves the backend "idle in transaction",
|
||||
# pinning a snapshot and blocking VACUUM for as long as the store
|
||||
# lives.
|
||||
conn.commit()
|
||||
return documents
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error searching documents: {e}", exc_info=True)
|
||||
try:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
# Connection already gone; nothing left to roll back.
|
||||
pass
|
||||
return []
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def keyword_search(self, question: str, k: int = 10) -> List[Document]:
|
||||
"""Full-text keyword search using Postgres ``websearch_to_tsquery``.
|
||||
|
||||
Returns the same ``Document`` shape as :meth:`search`. The question is
|
||||
bound as a query parameter (never interpolated) to prevent injection.
|
||||
"""
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
keyword_query = f"""
|
||||
SELECT {self._text_column}, {self._metadata_column},
|
||||
ts_rank(
|
||||
to_tsvector('english', {self._text_column}),
|
||||
websearch_to_tsquery('english', %s)
|
||||
) AS rank
|
||||
FROM {self._table_name}
|
||||
WHERE source_id = %s
|
||||
AND to_tsvector('english', {self._text_column})
|
||||
@@ websearch_to_tsquery('english', %s)
|
||||
ORDER BY rank DESC
|
||||
LIMIT %s;
|
||||
"""
|
||||
|
||||
cursor.execute(
|
||||
keyword_query, (question, self._source_id, question, k)
|
||||
)
|
||||
results = cursor.fetchall()
|
||||
|
||||
documents = []
|
||||
for text, metadata, _rank in results:
|
||||
metadata = metadata or {}
|
||||
documents.append(Document(page_content=text, metadata=metadata))
|
||||
|
||||
conn.commit()
|
||||
return documents
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error in keyword search: {e}", exc_info=True)
|
||||
try:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
# Connection already gone; nothing left to roll back.
|
||||
pass
|
||||
return []
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def add_texts(
|
||||
self,
|
||||
texts: List[str],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> List[str]:
|
||||
"""Add texts with their embeddings to the vector store"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
embeddings = self._embedding.embed_documents(texts)
|
||||
metadatas = metadatas or [{}] * len(texts)
|
||||
|
||||
self._ensure_schema_once()
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
insert_query = f"""
|
||||
INSERT INTO {self._table_name} ({self._text_column}, {self._vector_column}, {self._metadata_column}, source_id)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
RETURNING id;
|
||||
"""
|
||||
|
||||
inserted_ids = []
|
||||
for text, embedding, metadata in zip(texts, embeddings, metadatas):
|
||||
cursor.execute(
|
||||
insert_query,
|
||||
(text, embedding, Jsonb(metadata), self._source_id)
|
||||
)
|
||||
inserted_id = cursor.fetchone()[0]
|
||||
inserted_ids.append(str(inserted_id))
|
||||
|
||||
conn.commit()
|
||||
return inserted_ids
|
||||
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logging.error(f"Error adding texts: {e}")
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def delete_index(self, *args, **kwargs):
|
||||
"""Delete all documents for this source_id"""
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
delete_query = f"DELETE FROM {self._table_name} WHERE source_id = %s;"
|
||||
cursor.execute(delete_query, (self._source_id,))
|
||||
conn.commit()
|
||||
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logging.error(f"Error deleting index: {e}")
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def save_local(self, *args, **kwargs):
|
||||
"""No-op for PostgreSQL - data is already persisted"""
|
||||
pass
|
||||
|
||||
def get_chunks(self) -> List[Dict[str, Any]]:
|
||||
"""Get all chunks for this source_id"""
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
select_query = f"""
|
||||
SELECT id, {self._text_column}, {self._metadata_column}
|
||||
FROM {self._table_name}
|
||||
WHERE source_id = %s;
|
||||
"""
|
||||
cursor.execute(select_query, (self._source_id,))
|
||||
results = cursor.fetchall()
|
||||
|
||||
chunks = []
|
||||
for doc_id, text, metadata in results:
|
||||
chunks.append({
|
||||
"doc_id": str(doc_id),
|
||||
"text": text,
|
||||
"metadata": metadata or {}
|
||||
})
|
||||
|
||||
conn.commit()
|
||||
return chunks
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error getting chunks: {e}")
|
||||
try:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
# Connection already gone; nothing left to roll back.
|
||||
pass
|
||||
return []
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def add_chunk(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> str:
|
||||
"""Add a single chunk to the vector store"""
|
||||
metadata = metadata or {}
|
||||
|
||||
final_metadata = metadata.copy()
|
||||
|
||||
final_metadata["source_id"] = self._source_id
|
||||
|
||||
embeddings = self._embedding.embed_documents([text])
|
||||
|
||||
if not embeddings:
|
||||
raise ValueError("Could not generate embedding for chunk")
|
||||
|
||||
self._ensure_schema_once()
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
insert_query = f"""
|
||||
INSERT INTO {self._table_name} ({self._text_column}, {self._vector_column}, {self._metadata_column}, source_id)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
RETURNING id;
|
||||
"""
|
||||
|
||||
cursor.execute(
|
||||
insert_query,
|
||||
(text, embeddings[0], Jsonb(final_metadata), self._source_id)
|
||||
)
|
||||
inserted_id = cursor.fetchone()[0]
|
||||
conn.commit()
|
||||
|
||||
return str(inserted_id)
|
||||
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logging.error(f"Error adding chunk: {e}")
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def delete_chunk(self, chunk_id: str) -> bool:
|
||||
"""Delete a specific chunk by its ID"""
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
delete_query = f"DELETE FROM {self._table_name} WHERE id = %s AND source_id = %s;"
|
||||
cursor.execute(delete_query, (int(chunk_id), self._source_id))
|
||||
deleted_count = cursor.rowcount
|
||||
conn.commit()
|
||||
|
||||
return deleted_count > 0
|
||||
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logging.error(f"Error deleting chunk: {e}")
|
||||
return False
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def delete_chunks_by_source_path(self, path: str) -> int:
|
||||
"""Delete this source's chunks whose ``metadata.source`` equals ``path``.
|
||||
|
||||
One targeted statement instead of the base loop+scan. The path is bound
|
||||
as a query parameter (never interpolated); only the internal table name
|
||||
is f-string interpolated. Returns the number of rows deleted.
|
||||
"""
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
delete_query = (
|
||||
f"DELETE FROM {self._table_name} "
|
||||
f"WHERE source_id = %s AND {self._metadata_column}->>'source' = %s;"
|
||||
)
|
||||
cursor.execute(delete_query, (self._source_id, path))
|
||||
deleted_count = cursor.rowcount
|
||||
conn.commit()
|
||||
|
||||
return deleted_count
|
||||
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logging.error(f"Error deleting chunks by source path: {e}")
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Release this store's connection: back to the pool, or closed outright.
|
||||
|
||||
A pooled connection is rolled back first when it is still in a
|
||||
transaction, so the next borrower gets a clean session.
|
||||
"""
|
||||
conn = getattr(self, "_connection", None)
|
||||
if conn is None:
|
||||
return
|
||||
self._connection = None
|
||||
pgconn.release(self._connection_string, conn, getattr(self, "_pooled", False))
|
||||
|
||||
def __del__(self):
|
||||
"""Release the connection when the object is destroyed. Never raises."""
|
||||
try:
|
||||
self.close()
|
||||
except Exception:
|
||||
# Interpreter teardown can null out module globals; never raise here.
|
||||
pass
|
||||
@@ -0,0 +1,219 @@
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.base import BaseVectorStore
|
||||
from docsgpt.vectorstore.document_class import Document
|
||||
|
||||
|
||||
class QdrantStore(BaseVectorStore):
|
||||
"""Vector store backed by Qdrant through the native ``qdrant-client``.
|
||||
|
||||
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("docsgpt/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(
|
||||
key="metadata.source_id",
|
||||
match=models.MatchValue(value=self._source_id),
|
||||
)
|
||||
]
|
||||
)
|
||||
self._client = QdrantClient(**self._client_kwargs())
|
||||
self._ensure_collection()
|
||||
|
||||
@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:
|
||||
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()],
|
||||
),
|
||||
)
|
||||
self._client.create_payload_index(
|
||||
collection_name=self._collection,
|
||||
field_name="metadata.source_id",
|
||||
field_schema=models.PayloadSchemaType.KEYWORD,
|
||||
)
|
||||
except Exception as 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)
|
||||
|
||||
@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 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,
|
||||
query_vector: Optional[List[float]] = None,
|
||||
**kwargs,
|
||||
) -> List[Tuple[Document, float]]:
|
||||
"""Search, pairing each hit with its cosine similarity.
|
||||
|
||||
Args:
|
||||
query_vector: Precomputed embedding of ``question``; when given the
|
||||
store skips embedding the query itself.
|
||||
"""
|
||||
if query_vector is None:
|
||||
query_vector = self._embeddings.embed_query(question)
|
||||
hits = self._client.query_points(
|
||||
collection_name=self._collection,
|
||||
query=query_vector,
|
||||
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):
|
||||
"""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) -> List[Dict[str, Any]]:
|
||||
"""Return every chunk stored for this source."""
|
||||
chunks: List[Dict[str, Any]] = []
|
||||
offset = None
|
||||
try:
|
||||
while True:
|
||||
records, offset = self._client.scroll(
|
||||
collection_name=self._collection,
|
||||
scroll_filter=self._filter,
|
||||
limit=100,
|
||||
with_payload=True,
|
||||
with_vectors=False,
|
||||
offset=offset,
|
||||
)
|
||||
for record in records:
|
||||
payload = record.payload or {}
|
||||
chunks.append(
|
||||
{
|
||||
"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("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,
|
||||
points_selector=self._models.PointIdsList(points=[chunk_id]),
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logging.error("Error deleting chunk: %s", e, exc_info=True)
|
||||
return False
|
||||
@@ -0,0 +1,24 @@
|
||||
from docsgpt.vectorstore.faiss import FaissStore
|
||||
from docsgpt.vectorstore.elasticsearch import ElasticsearchStore
|
||||
from docsgpt.vectorstore.milvus import MilvusStore
|
||||
from docsgpt.vectorstore.mongodb import MongoDBVectorStore
|
||||
from docsgpt.vectorstore.qdrant import QdrantStore
|
||||
from docsgpt.vectorstore.pgvector import PGVectorStore
|
||||
|
||||
|
||||
class VectorCreator:
|
||||
vectorstores = {
|
||||
"faiss": FaissStore,
|
||||
"elasticsearch": ElasticsearchStore,
|
||||
"mongodb": MongoDBVectorStore,
|
||||
"qdrant": QdrantStore,
|
||||
"milvus": MilvusStore,
|
||||
"pgvector": PGVectorStore
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create_vectorstore(cls, type, *args, **kwargs):
|
||||
vectorstore_class = cls.vectorstores.get(type.lower())
|
||||
if not vectorstore_class:
|
||||
raise ValueError(f"No vectorstore class found for type {type}")
|
||||
return vectorstore_class(*args, **kwargs)
|
||||
Reference in new issue
Block a user