mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-25 01:45:23 +08:00
Hyu
5758f20768
A full-text search sends no query vector by design, but search() validated len([]) == 0 as the expected embedding dimension, so the request failed with 'Embedding dimension 0 is not enabled for this deployment' and hid the backend capability error. An empty query vector now means no expected dimension to verify.
516 lines
21 KiB
Python
516 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
|
|
import sqlalchemy
|
|
|
|
from ..api.http.authz import WorkspaceRequiredError
|
|
from ..api.http.context import ExecutionContext
|
|
from ..core import app
|
|
from ..entity.persistence import rag as persistence_rag
|
|
from ..entity.persistence import workspace as persistence_workspace
|
|
from ..workspace.errors import WorkspaceNotFoundError
|
|
from .vdb import VectorDatabase, SearchType
|
|
|
|
|
|
class VectorDBManager:
|
|
ap: app.Application
|
|
vector_db: VectorDatabase = None
|
|
|
|
def __init__(self, ap: app.Application):
|
|
self.ap = ap
|
|
|
|
async def initialize(self):
|
|
kb_config = self.ap.instance_config.data.get('vdb')
|
|
if kb_config:
|
|
vdb_type = kb_config.get('use')
|
|
|
|
if vdb_type == 'chroma':
|
|
from .vdbs.chroma import ChromaVectorDatabase
|
|
|
|
self.vector_db = ChromaVectorDatabase(self.ap)
|
|
self.ap.logger.info('Initialized Chroma vector database backend.')
|
|
|
|
elif vdb_type == 'qdrant':
|
|
from .vdbs.qdrant import QdrantVectorDatabase
|
|
|
|
self.vector_db = QdrantVectorDatabase(self.ap)
|
|
self.ap.logger.info('Initialized Qdrant vector database backend.')
|
|
elif vdb_type == 'seekdb':
|
|
from .vdbs.seekdb import SeekDBVectorDatabase
|
|
|
|
self.vector_db = SeekDBVectorDatabase(self.ap)
|
|
self.ap.logger.info('Initialized SeekDB vector database backend.')
|
|
|
|
elif vdb_type == 'valkey_search':
|
|
from .vdbs.valkey_search import ValkeySearchVectorDatabase
|
|
|
|
self.vector_db = ValkeySearchVectorDatabase(self.ap)
|
|
self.ap.logger.info('Initialized Valkey Search vector database backend.')
|
|
|
|
elif vdb_type == 'milvus':
|
|
from .vdbs.milvus import MilvusVectorDatabase
|
|
|
|
# Get Milvus configuration
|
|
milvus_config = kb_config.get('milvus', {})
|
|
uri = milvus_config.get('uri', './data/milvus.db')
|
|
token = milvus_config.get('token')
|
|
db_name = milvus_config.get('db_name', 'default')
|
|
self.vector_db = MilvusVectorDatabase(self.ap, uri=uri, token=token, db_name=db_name)
|
|
self.ap.logger.info('Initialized Milvus vector database backend.')
|
|
|
|
elif vdb_type == 'pgvector':
|
|
from .vdbs.pgvector_db import PgVectorDatabase
|
|
|
|
# Get pgvector configuration
|
|
pgvector_config = kb_config.get('pgvector', {})
|
|
use_business_database = pgvector_config.get('use_business_database', False)
|
|
allowed_dimensions = pgvector_config.get(
|
|
'allowed_dimensions',
|
|
[384, 512, 768, 1024, 1536, 3072],
|
|
)
|
|
common_options = {
|
|
'use_business_database': use_business_database,
|
|
'allowed_dimensions': allowed_dimensions,
|
|
}
|
|
if use_business_database:
|
|
self.vector_db = PgVectorDatabase(self.ap, **common_options)
|
|
self.ap.logger.info('Initialized pgvector on the shared business PostgreSQL database.')
|
|
return
|
|
connection_string = pgvector_config.get('connection_string')
|
|
if connection_string:
|
|
self.vector_db = PgVectorDatabase(
|
|
self.ap,
|
|
connection_string=connection_string,
|
|
**common_options,
|
|
)
|
|
else:
|
|
# Use individual parameters
|
|
host = pgvector_config.get('host', 'localhost')
|
|
port = pgvector_config.get('port', 5432)
|
|
database = pgvector_config.get('database', 'langbot')
|
|
user = pgvector_config.get('user', 'postgres')
|
|
password = pgvector_config.get('password', 'postgres')
|
|
self.vector_db = PgVectorDatabase(
|
|
self.ap,
|
|
host=host,
|
|
port=port,
|
|
database=database,
|
|
user=user,
|
|
password=password,
|
|
**common_options,
|
|
)
|
|
self.ap.logger.info('Initialized pgvector database backend.')
|
|
|
|
else:
|
|
from .vdbs.chroma import ChromaVectorDatabase
|
|
|
|
self.vector_db = ChromaVectorDatabase(self.ap)
|
|
self.ap.logger.warning('No valid vector database backend configured, defaulting to Chroma.')
|
|
else:
|
|
from .vdbs.chroma import ChromaVectorDatabase
|
|
|
|
self.vector_db = ChromaVectorDatabase(self.ap)
|
|
self.ap.logger.warning('No vector database backend configured, defaulting to Chroma.')
|
|
|
|
async def shutdown(self) -> None:
|
|
"""Release the active vector backend deterministically."""
|
|
|
|
vector_db = self.vector_db
|
|
self.vector_db = None
|
|
if vector_db is not None:
|
|
await vector_db.close()
|
|
|
|
def get_supported_search_types(self) -> list[str]:
|
|
"""Return the search types supported by the current VDB backend."""
|
|
if self.vector_db is None:
|
|
return [SearchType.VECTOR.value]
|
|
return [st.value for st in self.vector_db.supported_search_types()]
|
|
|
|
@staticmethod
|
|
def physical_collection_name(
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
) -> str:
|
|
"""Derive an opaque physical collection from trusted tenant identity.
|
|
|
|
Vector backends have different collection-name constraints, so the
|
|
instance, Workspace and knowledge-base identifiers are encoded through
|
|
UUIDv5 instead of being concatenated into a client-visible handle.
|
|
Placement generation is deliberately not part of the name: generation
|
|
fencing rejects stale work while preserving data across placements.
|
|
"""
|
|
|
|
if not isinstance(execution_context, ExecutionContext):
|
|
raise WorkspaceRequiredError('ExecutionContext is required for vector access')
|
|
instance_uuid = execution_context.instance_uuid.strip()
|
|
workspace_uuid = execution_context.workspace_uuid.strip()
|
|
kb_uuid = knowledge_base_uuid.strip() if isinstance(knowledge_base_uuid, str) else ''
|
|
if not instance_uuid or not workspace_uuid or not kb_uuid:
|
|
raise WorkspaceRequiredError('Instance, Workspace and knowledge-base context are required')
|
|
if execution_context.placement_generation <= 0:
|
|
raise WorkspaceRequiredError('A positive placement generation is required')
|
|
|
|
collection_uuid = uuid.uuid5(
|
|
uuid.NAMESPACE_URL,
|
|
f'langbot:knowledge-vector:{instance_uuid}:{workspace_uuid}:{kb_uuid}',
|
|
)
|
|
return f'lb_{collection_uuid.hex}'
|
|
|
|
async def _validate_execution_context(self, execution_context: ExecutionContext) -> None:
|
|
"""Validate the active placement before touching a vector backend."""
|
|
|
|
# Also performs structural validation before accessing app services.
|
|
self.physical_collection_name(execution_context, 'context-validation')
|
|
workspace_service = getattr(self.ap, 'workspace_service', None)
|
|
if workspace_service is None:
|
|
raise WorkspaceRequiredError('Workspace execution service is unavailable')
|
|
binding = await workspace_service.get_execution_binding(
|
|
execution_context.workspace_uuid,
|
|
expected_generation=execution_context.placement_generation,
|
|
)
|
|
if binding.instance_uuid != execution_context.instance_uuid:
|
|
raise WorkspaceRequiredError('ExecutionContext belongs to another LangBot instance')
|
|
|
|
async def _resolve_physical_collection_name(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
) -> str:
|
|
"""Resolve a scoped collection or an explicitly migrated OSS handle.
|
|
|
|
Legacy handles are server-owned migration state, not caller input.
|
|
They are honored only for the one local Workspace under the OSS
|
|
single-Workspace policy. A projected/cloud Workspace always gets the
|
|
opaque tenant-derived collection, even if its database row was
|
|
incorrectly marked as legacy.
|
|
"""
|
|
|
|
await self._validate_execution_context(execution_context)
|
|
async with self.ap.persistence_mgr.tenant_uow(execution_context.workspace_uuid):
|
|
result = await self.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.select(
|
|
persistence_rag.KnowledgeBase.collection_id,
|
|
persistence_rag.KnowledgeBase.legacy_vector_collection,
|
|
persistence_workspace.Workspace.source,
|
|
)
|
|
.join(
|
|
persistence_workspace.Workspace,
|
|
persistence_workspace.Workspace.uuid == persistence_rag.KnowledgeBase.workspace_uuid,
|
|
)
|
|
.where(
|
|
persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
|
|
persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
|
|
persistence_workspace.Workspace.instance_uuid == execution_context.instance_uuid,
|
|
)
|
|
.limit(1)
|
|
)
|
|
row = result.first()
|
|
if row is None:
|
|
raise WorkspaceNotFoundError('Knowledge base not found')
|
|
|
|
collection_id, legacy_vector_collection, workspace_source = row
|
|
if legacy_vector_collection:
|
|
policy = getattr(self.ap, 'workspace_policy', None)
|
|
is_single_workspace = policy is not None and not getattr(
|
|
policy,
|
|
'multi_workspace_enabled',
|
|
True,
|
|
)
|
|
is_local_workspace = workspace_source == persistence_workspace.WorkspaceSource.LOCAL.value
|
|
if is_single_workspace and is_local_workspace and isinstance(collection_id, str) and collection_id.strip():
|
|
return collection_id
|
|
self.ap.logger.warning(
|
|
'Ignored a legacy vector collection marker outside the local single-Workspace compatibility boundary.'
|
|
)
|
|
|
|
return self.physical_collection_name(execution_context, knowledge_base_uuid)
|
|
|
|
def _pgvector_database(self):
|
|
from .vdbs.pgvector_db import PgVectorDatabase
|
|
|
|
return self.vector_db if isinstance(self.vector_db, PgVectorDatabase) else None
|
|
|
|
async def _resolve_pgvector_scope(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
*,
|
|
expected_dimension: int | None,
|
|
initialize_dimension: bool,
|
|
):
|
|
"""Bind and verify the server-owned knowledge-base vector dimension."""
|
|
|
|
from .vdbs.pgvector_db import PgVectorScope
|
|
|
|
pgvector = self._pgvector_database()
|
|
if pgvector is None: # pragma: no cover - private call invariant
|
|
raise RuntimeError('pgvector scope requested for another vector backend')
|
|
if expected_dimension is not None and expected_dimension not in pgvector.allowed_dimensions:
|
|
raise ValueError(f'Embedding dimension {expected_dimension} is not enabled for this deployment')
|
|
|
|
async with self.ap.persistence_mgr.tenant_uow(execution_context.workspace_uuid):
|
|
query = sqlalchemy.select(persistence_rag.KnowledgeBase.embedding_dimension).where(
|
|
persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
|
|
persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
|
|
)
|
|
current_dimension = (await self.ap.persistence_mgr.execute_async(query)).scalar_one_or_none()
|
|
if current_dimension is None and expected_dimension is not None and initialize_dimension:
|
|
await self.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.update(persistence_rag.KnowledgeBase)
|
|
.where(
|
|
persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
|
|
persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
|
|
persistence_rag.KnowledgeBase.embedding_dimension.is_(None),
|
|
)
|
|
.values(embedding_dimension=expected_dimension)
|
|
)
|
|
current_dimension = (await self.ap.persistence_mgr.execute_async(query)).scalar_one_or_none()
|
|
|
|
if expected_dimension is not None and current_dimension != expected_dimension:
|
|
if current_dimension is None:
|
|
raise ValueError('Knowledge base has no selected pgvector embedding dimension')
|
|
raise ValueError(f'Knowledge base embedding dimension is {current_dimension}, not {expected_dimension}')
|
|
|
|
return PgVectorScope(
|
|
workspace_uuid=execution_context.workspace_uuid,
|
|
knowledge_base_uuid=knowledge_base_uuid,
|
|
embedding_dimension=current_dimension,
|
|
)
|
|
|
|
async def upsert(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
vectors: list[list[float]],
|
|
ids: list[str],
|
|
metadata: list[dict] | None = None,
|
|
documents: list[str] | None = None,
|
|
):
|
|
"""Upsert vectors into a server-derived tenant collection."""
|
|
|
|
collection_name = await self._resolve_physical_collection_name(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
)
|
|
source_metadata = metadata or [{} for _ in vectors]
|
|
scoped_metadata = [
|
|
{
|
|
**item,
|
|
'_langbot_instance_uuid': execution_context.instance_uuid,
|
|
'_langbot_workspace_uuid': execution_context.workspace_uuid,
|
|
'_langbot_knowledge_base_uuid': knowledge_base_uuid,
|
|
}
|
|
for item in source_metadata
|
|
]
|
|
pgvector = self._pgvector_database()
|
|
if pgvector is not None:
|
|
if not vectors:
|
|
return
|
|
scope = await self._resolve_pgvector_scope(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
expected_dimension=len(vectors[0]),
|
|
initialize_dimension=True,
|
|
)
|
|
await pgvector.add_embeddings(
|
|
collection=collection_name,
|
|
ids=ids,
|
|
embeddings_list=vectors,
|
|
metadatas=scoped_metadata,
|
|
documents=documents,
|
|
scope=scope,
|
|
)
|
|
return
|
|
await self.vector_db.add_embeddings(
|
|
collection=collection_name,
|
|
ids=ids,
|
|
embeddings_list=vectors,
|
|
metadatas=scoped_metadata,
|
|
documents=documents,
|
|
)
|
|
|
|
async def search(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
query_vector: list[float],
|
|
limit: int,
|
|
filter: dict | None = None,
|
|
search_type: str = 'vector',
|
|
query_text: str = '',
|
|
vector_weight: float | None = None,
|
|
) -> list[dict]:
|
|
"""Proxy: Search vectors.
|
|
|
|
Returns a list of dicts with keys: 'id', 'distance', 'metadata'.
|
|
The underlying VectorDatabase.search returns Chroma-style format:
|
|
{ 'ids': [['id1']], 'distances': [[0.1]], 'metadatas': [[{}]] }
|
|
"""
|
|
collection_name = await self._resolve_physical_collection_name(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
)
|
|
pgvector = self._pgvector_database()
|
|
if pgvector is not None:
|
|
scope = await self._resolve_pgvector_scope(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
# A text-only search sends no query vector; there is no dimension
|
|
# to bind, and validating an empty vector as dimension 0 would
|
|
# mask the backend capability error.
|
|
expected_dimension=len(query_vector) or None,
|
|
initialize_dimension=False,
|
|
)
|
|
results = await pgvector.search(
|
|
collection=collection_name,
|
|
query_embedding=query_vector,
|
|
k=limit,
|
|
search_type=search_type,
|
|
query_text=query_text,
|
|
filter=filter,
|
|
vector_weight=vector_weight,
|
|
scope=scope,
|
|
)
|
|
else:
|
|
results = await self.vector_db.search(
|
|
collection=collection_name,
|
|
query_embedding=query_vector,
|
|
k=limit,
|
|
search_type=search_type,
|
|
query_text=query_text,
|
|
filter=filter,
|
|
vector_weight=vector_weight,
|
|
)
|
|
|
|
if not results or 'ids' not in results or not results['ids']:
|
|
return []
|
|
|
|
# Flatten nested lists (Chroma returns batch-style: list of lists)
|
|
raw_ids = results['ids']
|
|
raw_dists = results.get('distances', [])
|
|
raw_metas = results.get('metadatas', [])
|
|
|
|
r_ids = raw_ids[0] if raw_ids and isinstance(raw_ids[0], list) else raw_ids
|
|
r_dists = raw_dists[0] if raw_dists and isinstance(raw_dists[0], list) else raw_dists
|
|
r_metas = raw_metas[0] if raw_metas and isinstance(raw_metas[0], list) else raw_metas
|
|
|
|
parsed_results = []
|
|
for i, id_val in enumerate(r_ids):
|
|
parsed_results.append(
|
|
{
|
|
'id': id_val,
|
|
'distance': r_dists[i] if r_dists and i < len(r_dists) else 0.0,
|
|
'metadata': r_metas[i] if r_metas and i < len(r_metas) else {},
|
|
}
|
|
)
|
|
|
|
return parsed_results
|
|
|
|
async def delete_by_file_id(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
file_ids: list[str],
|
|
):
|
|
"""Proxy: Delete vectors by file_id (metadata-level identifier).
|
|
|
|
This delegates to VectorDatabase.delete_by_file_id which removes
|
|
all vectors associated with the given file IDs.
|
|
"""
|
|
collection_name = await self._resolve_physical_collection_name(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
)
|
|
pgvector = self._pgvector_database()
|
|
scope = None
|
|
if pgvector is not None:
|
|
scope = await self._resolve_pgvector_scope(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
expected_dimension=None,
|
|
initialize_dimension=False,
|
|
)
|
|
for file_id in file_ids:
|
|
if pgvector is not None:
|
|
await pgvector.delete_by_file_id(collection_name, file_id, scope=scope)
|
|
else:
|
|
await self.vector_db.delete_by_file_id(collection_name, file_id)
|
|
|
|
async def delete_collection(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
):
|
|
"""Delete one server-derived tenant collection."""
|
|
|
|
collection_name = await self._resolve_physical_collection_name(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
)
|
|
pgvector = self._pgvector_database()
|
|
if pgvector is not None:
|
|
scope = await self._resolve_pgvector_scope(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
expected_dimension=None,
|
|
initialize_dimension=False,
|
|
)
|
|
await pgvector.delete_collection(collection_name, scope=scope)
|
|
else:
|
|
await self.vector_db.delete_collection(collection_name)
|
|
|
|
async def delete_by_filter(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
filter: dict,
|
|
) -> int:
|
|
"""Proxy: Delete vectors by metadata filter.
|
|
|
|
Returns:
|
|
Number of deleted vectors (best-effort; some backends return 0).
|
|
"""
|
|
collection_name = await self._resolve_physical_collection_name(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
)
|
|
pgvector = self._pgvector_database()
|
|
if pgvector is not None:
|
|
scope = await self._resolve_pgvector_scope(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
expected_dimension=None,
|
|
initialize_dimension=False,
|
|
)
|
|
return await pgvector.delete_by_filter(collection_name, filter, scope=scope)
|
|
return await self.vector_db.delete_by_filter(collection_name, filter)
|
|
|
|
async def list_by_filter(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_uuid: str,
|
|
filter: dict | None = None,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> tuple[list[dict], int]:
|
|
"""Proxy: List vectors by metadata filter with pagination.
|
|
|
|
Returns:
|
|
Tuple of (items, total).
|
|
"""
|
|
collection_name = await self._resolve_physical_collection_name(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
)
|
|
pgvector = self._pgvector_database()
|
|
if pgvector is not None:
|
|
scope = await self._resolve_pgvector_scope(
|
|
execution_context,
|
|
knowledge_base_uuid,
|
|
expected_dimension=None,
|
|
initialize_dimension=False,
|
|
)
|
|
return await pgvector.list_by_filter(collection_name, filter, limit, offset, scope=scope)
|
|
return await self.vector_db.list_by_filter(collection_name, filter, limit, offset)
|