mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-25 11:37:13 +00:00
RockChinQ
505 lines
19 KiB
Python
505 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from decimal import Decimal
|
|
import re
|
|
from typing import Any, Dict, List
|
|
|
|
|
|
from langbot.pkg.core import app
|
|
from langbot.pkg.utils import bounded_executor
|
|
from langbot.pkg.vector.vdb import (
|
|
VectorDatabase,
|
|
SearchType,
|
|
remember_bounded_mapping,
|
|
runtime_cache_limit,
|
|
)
|
|
|
|
try:
|
|
import pyseekdb
|
|
from pyseekdb import HNSWConfiguration
|
|
|
|
SEEKDB_AVAILABLE = True
|
|
except ImportError:
|
|
SEEKDB_AVAILABLE = False
|
|
|
|
SEEKDB_EMBEDDING_MODEL_UUID = 'seekdb-builtin-embedding'
|
|
SEEKDB_EMBEDDING_REQUESTER = 'seekdb-embedding'
|
|
|
|
|
|
class SeekDBVectorDatabase(VectorDatabase):
|
|
"""SeekDB vector database adapter for LangBot.
|
|
|
|
SeekDB is an AI-native search database by OceanBase that unifies
|
|
relational, vector, text, JSON and GIS in a single engine.
|
|
|
|
Supports embedded mode, remote server mode, and full-text/hybrid search.
|
|
"""
|
|
|
|
@classmethod
|
|
def supported_search_types(cls) -> list[SearchType]:
|
|
return [SearchType.VECTOR, SearchType.FULL_TEXT, SearchType.HYBRID]
|
|
|
|
def __init__(self, ap: app.Application):
|
|
if not SEEKDB_AVAILABLE:
|
|
raise ImportError(
|
|
"SeekDB support is not installed. Install LangBot with the 'seekdb' extra: "
|
|
"uv sync --extra seekdb (source) or uvx --from 'langbot[seekdb]@latest' langbot (PyPI)."
|
|
)
|
|
|
|
self.ap = ap
|
|
config = self.ap.instance_config.data['vdb']['seekdb']
|
|
|
|
# Determine connection mode based on config
|
|
mode = config.get('mode', 'embedded') # 'embedded' or 'server'
|
|
|
|
if mode == 'embedded':
|
|
# Embedded mode: local database
|
|
path = config.get('path', './data/seekdb')
|
|
database = config.get('database', 'langbot')
|
|
|
|
# Use AdminClient for database management operations
|
|
admin_client = pyseekdb.AdminClient(path=path)
|
|
# Check if database exists using public API
|
|
existing_dbs = [db.name for db in admin_client.list_databases()]
|
|
if database not in existing_dbs:
|
|
# Use public API to create database
|
|
admin_client.create_database(database)
|
|
self.ap.logger.info(f"Created SeekDB database '{database}'")
|
|
|
|
self.client = pyseekdb.Client(path=path, database=database)
|
|
self.ap.logger.info(f"Initialized SeekDB in embedded mode at '{path}', database '{database}'")
|
|
elif mode == 'server':
|
|
# Server mode: remote SeekDB or OceanBase server
|
|
host = config.get('host', 'localhost')
|
|
port = config.get('port', 2881)
|
|
database = config.get('database', 'langbot')
|
|
user = config.get('user', 'root')
|
|
password = config.get('password', '')
|
|
tenant = config.get('tenant', None) # Optional, for OceanBase
|
|
|
|
connection_params = {
|
|
'host': host,
|
|
'port': int(port),
|
|
'database': database,
|
|
'user': user,
|
|
'password': password,
|
|
}
|
|
|
|
if tenant:
|
|
connection_params['tenant'] = tenant
|
|
|
|
self.client = pyseekdb.Client(**connection_params)
|
|
self.ap.logger.info(
|
|
f"Initialized SeekDB in server mode: {host}:{port}, database '{database}'"
|
|
+ (f", tenant '{tenant}'" if tenant else '')
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid SeekDB mode: {mode}. Must be 'embedded' or 'server'")
|
|
|
|
self._collections: Dict[str, Any] = {}
|
|
self._collection_configs: Dict[str, HNSWConfiguration] = {}
|
|
self._runtime_cache_limit = runtime_cache_limit(ap)
|
|
|
|
self._escape_table = str.maketrans(
|
|
{
|
|
'\x00': '',
|
|
'\\': '\\\\',
|
|
"'": "''", # Standard SQL escaping (OceanBase NO_BACKSLASH_ESCAPES)
|
|
'"': '\\"',
|
|
'\n': '\\n',
|
|
'\r': '\\r',
|
|
'\t': '\\t',
|
|
}
|
|
)
|
|
|
|
async def close(self) -> None:
|
|
self._collections.clear()
|
|
self._collection_configs.clear()
|
|
close = getattr(self.client, 'close', None)
|
|
if callable(close):
|
|
await bounded_executor.run_blocking_cleanup(close)
|
|
|
|
def _normalize_collection_name(self, collection: str) -> str:
|
|
"""SeekDB only accepts [a-zA-Z0-9_], while LangBot uses UUID-like KB IDs."""
|
|
normalized = re.sub(r'[^A-Za-z0-9_]', '_', collection)
|
|
if normalized != collection:
|
|
self.ap.logger.info(f"Normalized SeekDB collection name: '{collection}' -> '{normalized}'")
|
|
return normalized
|
|
|
|
def _json_safe(self, value: Any) -> Any:
|
|
"""Convert SeekDB result values into JSON-serializable Python primitives."""
|
|
if isinstance(value, Decimal):
|
|
return float(value)
|
|
if isinstance(value, dict):
|
|
return {k: self._json_safe(v) for k, v in value.items()}
|
|
if isinstance(value, list):
|
|
return [self._json_safe(v) for v in value]
|
|
if isinstance(value, tuple):
|
|
return [self._json_safe(v) for v in value]
|
|
return value
|
|
|
|
async def _get_or_create_collection_internal(self, collection: str, vector_size: int = None) -> Any:
|
|
"""Internal method to get or create a collection with proper configuration."""
|
|
collection = self._normalize_collection_name(collection)
|
|
if collection in self._collections:
|
|
return self._collections[collection]
|
|
|
|
# Check if collection exists
|
|
if await asyncio.to_thread(self.client.has_collection, collection):
|
|
# Collection exists, get it
|
|
coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None)
|
|
remember_bounded_mapping(
|
|
self._collections,
|
|
collection,
|
|
coll,
|
|
self._runtime_cache_limit,
|
|
)
|
|
self.ap.logger.info(f"SeekDB collection '{collection}' retrieved.")
|
|
return coll
|
|
|
|
# Collection doesn't exist, create it
|
|
if vector_size is None:
|
|
raise ValueError(
|
|
f"Cannot create SeekDB collection '{collection}' without knowing the vector dimension. "
|
|
'Ensure add_embeddings is called before any standalone get_or_create_collection.'
|
|
)
|
|
|
|
# Create HNSW configuration
|
|
config = HNSWConfiguration(dimension=vector_size, distance='cosine')
|
|
remember_bounded_mapping(
|
|
self._collection_configs,
|
|
collection,
|
|
config,
|
|
self._runtime_cache_limit,
|
|
)
|
|
|
|
# Create collection without embedding function (we manage embeddings externally)
|
|
coll = await asyncio.to_thread(
|
|
self.client.create_collection,
|
|
name=collection,
|
|
configuration=config,
|
|
embedding_function=None, # Disable automatic embedding
|
|
)
|
|
|
|
remember_bounded_mapping(
|
|
self._collections,
|
|
collection,
|
|
coll,
|
|
self._runtime_cache_limit,
|
|
)
|
|
self.ap.logger.info(f"SeekDB collection '{collection}' created with dimension={vector_size}, distance='cosine'")
|
|
return coll
|
|
|
|
def _clean_metadata(self, meta: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""SeekDB metadata doesn't support \\ and ", insert will error 3104"""
|
|
return {
|
|
k: v.translate(self._escape_table)
|
|
if isinstance(v, str)
|
|
else v
|
|
if v is None or isinstance(v, (int, float, bool))
|
|
else str(v)
|
|
for k, v in meta.items()
|
|
if v is not None
|
|
}
|
|
|
|
async def get_or_create_collection(self, collection: str):
|
|
"""Get or create collection (without vector size - will use default)."""
|
|
return await self._get_or_create_collection_internal(collection)
|
|
|
|
async def add_embeddings(
|
|
self,
|
|
collection: str,
|
|
ids: List[str],
|
|
embeddings_list: List[List[float]],
|
|
metadatas: List[Dict[str, Any]],
|
|
documents: List[str] | None = None,
|
|
) -> None:
|
|
"""Add vector embeddings to the specified collection.
|
|
|
|
Args:
|
|
collection: Collection name
|
|
ids: List of document IDs
|
|
embeddings_list: List of embedding vectors
|
|
metadatas: List of metadata dictionaries
|
|
documents: Optional raw text documents for full-text search support
|
|
"""
|
|
if not embeddings_list:
|
|
return
|
|
|
|
collection = self._normalize_collection_name(collection)
|
|
# Ensure collection exists with correct dimension
|
|
vector_size = len(embeddings_list[0])
|
|
coll = await self._get_or_create_collection_internal(collection, vector_size)
|
|
|
|
cleaned_metadatas = [self._clean_metadata(meta) for meta in metadatas]
|
|
|
|
kwargs: Dict[str, Any] = dict(ids=ids, embeddings=embeddings_list, metadatas=cleaned_metadatas)
|
|
if documents is not None:
|
|
kwargs['documents'] = [doc.translate(self._escape_table) for doc in documents]
|
|
await asyncio.to_thread(coll.add, **kwargs)
|
|
|
|
self.ap.logger.info(f"Added {len(ids)} embeddings to SeekDB collection '{collection}'")
|
|
|
|
async def search(
|
|
self,
|
|
collection: str,
|
|
query_embedding: List[float],
|
|
k: int = 5,
|
|
search_type: str = 'vector',
|
|
query_text: str = '',
|
|
filter: Dict[str, Any] | None = None,
|
|
vector_weight: float | None = None,
|
|
) -> Dict[str, Any]:
|
|
"""Search for the most similar vectors in the specified collection.
|
|
|
|
SeekDB supports vector, full-text, and hybrid search modes.
|
|
|
|
Args:
|
|
collection: Collection name
|
|
query_embedding: Query vector (used for vector and hybrid modes)
|
|
k: Number of results to return
|
|
search_type: One of 'vector', 'full_text', 'hybrid'
|
|
query_text: Raw query text (used for full_text and hybrid modes)
|
|
filter: Optional metadata filters (Chroma-style ``where`` syntax).
|
|
|
|
Returns:
|
|
Dictionary with 'ids', 'metadatas', 'distances' keys
|
|
"""
|
|
collection = self._normalize_collection_name(collection)
|
|
# Check if collection exists
|
|
exists = await asyncio.to_thread(self.client.has_collection, collection)
|
|
if not exists:
|
|
return {'ids': [[]], 'metadatas': [[]], 'distances': [[]]}
|
|
|
|
# Get collection
|
|
if collection not in self._collections:
|
|
coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None)
|
|
remember_bounded_mapping(
|
|
self._collections,
|
|
collection,
|
|
coll,
|
|
self._runtime_cache_limit,
|
|
)
|
|
else:
|
|
coll = self._collections[collection]
|
|
|
|
# Route by search type.
|
|
# pyseekdb's query() always requires embeddings, so full-text and
|
|
# hybrid modes use hybrid_search() which supports text-only queries
|
|
# and returns the same nested-list format with distances.
|
|
if search_type == SearchType.FULL_TEXT:
|
|
if not query_text:
|
|
return {'ids': [[]], 'metadatas': [[]], 'distances': [[]]}
|
|
|
|
query_cfg: Dict[str, Any] = {
|
|
'where_document': {'$contains': query_text},
|
|
'n_results': k,
|
|
}
|
|
if filter:
|
|
query_cfg['where'] = filter
|
|
|
|
# TODO: pyseekdb hybrid_search with query-only (no knn) returns None
|
|
# for IDs due to column name mismatch (*/_id vs _id).
|
|
# See: https://github.com/oceanbase/pyseekdb/issues/171
|
|
results = await asyncio.to_thread(
|
|
coll.hybrid_search,
|
|
query=query_cfg,
|
|
knn=None,
|
|
n_results=k,
|
|
include=['documents', 'metadatas'],
|
|
)
|
|
|
|
elif search_type == SearchType.HYBRID:
|
|
if not query_text:
|
|
# Fall back to pure vector search when no text is provided
|
|
query_kwargs: Dict[str, Any] = {
|
|
'n_results': k,
|
|
'query_embeddings': query_embedding,
|
|
}
|
|
if filter:
|
|
query_kwargs['where'] = filter
|
|
results = await asyncio.to_thread(coll.query, **query_kwargs)
|
|
else:
|
|
query_cfg = {
|
|
'where_document': {'$contains': query_text},
|
|
'n_results': k,
|
|
}
|
|
knn_cfg: Dict[str, Any] = {
|
|
'query_embeddings': query_embedding,
|
|
'n_results': k,
|
|
}
|
|
if filter:
|
|
query_cfg['where'] = filter
|
|
knn_cfg['where'] = filter
|
|
|
|
# Apply vector_weight via pyseekdb's native boost parameter
|
|
if vector_weight is not None:
|
|
knn_cfg['boost'] = vector_weight
|
|
query_cfg['boost'] = 1.0 - vector_weight
|
|
self.ap.logger.info(
|
|
f"SeekDB hybrid fusion config in '{collection}': "
|
|
f'vector_weight={vector_weight}, '
|
|
f'knn_boost={knn_cfg.get("boost", 1.0)}, '
|
|
f'query_boost={query_cfg.get("boost", 1.0)}'
|
|
)
|
|
|
|
results = await asyncio.to_thread(
|
|
coll.hybrid_search,
|
|
query=query_cfg,
|
|
knn=knn_cfg,
|
|
rank={'rrf': {}},
|
|
n_results=k,
|
|
include=['documents', 'metadatas'],
|
|
)
|
|
self.ap.logger.info(
|
|
f"SeekDB hybrid search in '{collection}' returned {len(results.get('ids', [[]])[0])} results."
|
|
)
|
|
else:
|
|
# Default: vector search via query()
|
|
query_kwargs = {'n_results': k, 'query_embeddings': query_embedding}
|
|
if filter:
|
|
query_kwargs['where'] = filter
|
|
results = await asyncio.to_thread(coll.query, **query_kwargs)
|
|
|
|
results = self._json_safe(results)
|
|
self.ap.logger.info(
|
|
f"SeekDB {search_type} search in '{collection}' returned {len(results.get('ids', [[]])[0])} results"
|
|
)
|
|
|
|
return results
|
|
|
|
async def delete_by_file_id(self, collection: str, file_id: str) -> None:
|
|
"""Delete vectors from the collection by file_id metadata.
|
|
|
|
Args:
|
|
collection: Collection name
|
|
file_id: File ID to delete
|
|
"""
|
|
collection = self._normalize_collection_name(collection)
|
|
# Check if collection exists
|
|
exists = await asyncio.to_thread(self.client.has_collection, collection)
|
|
if not exists:
|
|
self.ap.logger.warning(f"SeekDB collection '{collection}' not found for deletion")
|
|
return
|
|
|
|
# Get collection
|
|
if collection not in self._collections:
|
|
coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None)
|
|
remember_bounded_mapping(
|
|
self._collections,
|
|
collection,
|
|
coll,
|
|
self._runtime_cache_limit,
|
|
)
|
|
else:
|
|
coll = self._collections[collection]
|
|
|
|
# SeekDB's delete() expects a where clause for filtering
|
|
# Delete all records where metadata['file_id'] == file_id
|
|
await asyncio.to_thread(coll.delete, where={'file_id': file_id})
|
|
|
|
self.ap.logger.info(f"Deleted embeddings from SeekDB collection '{collection}' with file_id: {file_id}")
|
|
|
|
async def delete_by_filter(self, collection: str, filter: Dict[str, Any]) -> int:
|
|
"""Delete vectors from the collection by metadata filter.
|
|
|
|
Args:
|
|
collection: Collection name
|
|
filter: Chroma-style ``where`` filter dict
|
|
"""
|
|
collection = self._normalize_collection_name(collection)
|
|
exists = await asyncio.to_thread(self.client.has_collection, collection)
|
|
if not exists:
|
|
self.ap.logger.warning(f"SeekDB collection '{collection}' not found for deletion")
|
|
return 0
|
|
|
|
if collection not in self._collections:
|
|
coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None)
|
|
remember_bounded_mapping(
|
|
self._collections,
|
|
collection,
|
|
coll,
|
|
self._runtime_cache_limit,
|
|
)
|
|
else:
|
|
coll = self._collections[collection]
|
|
|
|
await asyncio.to_thread(coll.delete, where=filter)
|
|
self.ap.logger.info(f"Deleted embeddings from SeekDB collection '{collection}' by filter")
|
|
return 0 # SeekDB delete does not return a count
|
|
|
|
async def list_by_filter(
|
|
self,
|
|
collection: str,
|
|
filter: Dict[str, Any] | None = None,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> tuple[list[Dict[str, Any]], int]:
|
|
collection = self._normalize_collection_name(collection)
|
|
exists = await asyncio.to_thread(self.client.has_collection, collection)
|
|
if not exists:
|
|
return [], 0
|
|
|
|
if collection not in self._collections:
|
|
coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None)
|
|
remember_bounded_mapping(
|
|
self._collections,
|
|
collection,
|
|
coll,
|
|
self._runtime_cache_limit,
|
|
)
|
|
else:
|
|
coll = self._collections[collection]
|
|
|
|
get_kwargs: Dict[str, Any] = dict(
|
|
include=['metadatas', 'documents'],
|
|
limit=limit,
|
|
offset=offset,
|
|
)
|
|
if filter:
|
|
get_kwargs['where'] = filter
|
|
|
|
results = await asyncio.to_thread(coll.get, **get_kwargs)
|
|
|
|
results = self._json_safe(results)
|
|
ids = results.get('ids', [])
|
|
metadatas = results.get('metadatas', []) or [None] * len(ids)
|
|
documents = results.get('documents', []) or [None] * len(ids)
|
|
|
|
items = []
|
|
for i, vid in enumerate(ids):
|
|
items.append(
|
|
{
|
|
'id': vid,
|
|
'document': documents[i] if i < len(documents) else None,
|
|
'metadata': metadatas[i] if i < len(metadatas) else {},
|
|
}
|
|
)
|
|
|
|
total = await asyncio.to_thread(coll.count) if not filter else -1
|
|
return items, total
|
|
|
|
async def delete_collection(self, collection: str):
|
|
"""Delete the entire collection.
|
|
|
|
Args:
|
|
collection: Collection name
|
|
"""
|
|
collection = self._normalize_collection_name(collection)
|
|
# Remove from cache
|
|
if collection in self._collections:
|
|
del self._collections[collection]
|
|
if collection in self._collection_configs:
|
|
del self._collection_configs[collection]
|
|
|
|
# Check if collection exists
|
|
exists = await asyncio.to_thread(self.client.has_collection, collection)
|
|
if not exists:
|
|
self.ap.logger.warning(f"SeekDB collection '{collection}' not found for deletion")
|
|
return
|
|
|
|
# Delete collection
|
|
await asyncio.to_thread(self.client.delete_collection, collection)
|
|
self.ap.logger.info(f"SeekDB collection '{collection}' deleted")
|