dimgigov
55bc3e862a
- BaraDBStore for Python: add_texts, similarity_search, max_marginal_relevance_search, delete - BaraDBStore for JS: addDocuments, addTexts, similaritySearch, maxMarginalRelevanceSearch, delete - Both use hybrid_search() / hybrid_search_filtered() for vector+FTS+RRF - Multi-tenant support via tenant_id session variable + metadata filter - Embedding function is injected by user (OpenAI, sentence-transformers, etc.) - MMR reranking for result diversity
213 lines
8.0 KiB
Python
213 lines
8.0 KiB
Python
"""
|
|
BaraDB LangChain Vector Store Integration
|
|
|
|
Usage:
|
|
from baradb import Client
|
|
from baradb.langchain_store import BaraDBStore
|
|
from langchain.embeddings import OpenAIEmbeddings
|
|
|
|
client = Client("localhost", 9472)
|
|
await client.connect()
|
|
|
|
store = BaraDBStore(
|
|
client=client,
|
|
table="docs",
|
|
embedding_col="embedding",
|
|
text_col="content",
|
|
embedding_function=OpenAIEmbeddings().embed_query,
|
|
tenant_id="company-a" # optional, for RLS
|
|
)
|
|
|
|
await store.add_texts(["hello world", "quick brown fox"])
|
|
results = await store.similarity_search("hello", k=5)
|
|
"""
|
|
|
|
import json
|
|
from typing import Any, Callable, List, Optional, Sequence, Tuple
|
|
|
|
|
|
class BaraDBStore:
|
|
"""LangChain-compatible Vector Store for BaraDB."""
|
|
|
|
def __init__(
|
|
self,
|
|
client: Any,
|
|
table: str = "documents",
|
|
embedding_col: str = "embedding",
|
|
text_col: str = "content",
|
|
metadata_cols: Optional[List[str]] = None,
|
|
embedding_function: Optional[Callable[[str], List[float]]] = None,
|
|
tenant_id: Optional[str] = None,
|
|
vector_dimension: int = 1536,
|
|
):
|
|
self.client = client
|
|
self.table = table
|
|
self.embedding_col = embedding_col
|
|
self.text_col = text_col
|
|
self.metadata_cols = metadata_cols or []
|
|
self.embedding_function = embedding_function
|
|
self.tenant_id = tenant_id
|
|
self.vector_dimension = vector_dimension
|
|
self._table_created = False
|
|
|
|
async def _ensure_table(self) -> None:
|
|
if self._table_created:
|
|
return
|
|
# Create table with vector + text + tenant_id columns
|
|
cols = f"id SERIAL PRIMARY KEY, {self.embedding_col} VECTOR({self.vector_dimension}), {self.text_col} TEXT"
|
|
if self.tenant_id:
|
|
cols += ", tenant_id TEXT"
|
|
for mc in self.metadata_cols:
|
|
cols += f", {mc} TEXT"
|
|
await self.client.query(f"CREATE TABLE IF NOT EXISTS {self.table} ({cols})")
|
|
|
|
# Create indexes if not exist
|
|
idx_vec = f"idx_{self.table}_vec"
|
|
idx_fts = f"idx_{self.table}_fts"
|
|
await self.client.query(f"CREATE INDEX IF NOT EXISTS {idx_vec} ON {self.table}({self.embedding_col}) USING hnsw")
|
|
await self.client.query(f"CREATE INDEX IF NOT EXISTS {idx_fts} ON {self.table}({self.text_col}) USING FTS")
|
|
self._table_created = True
|
|
|
|
async def add_texts(
|
|
self,
|
|
texts: Sequence[str],
|
|
metadatas: Optional[List[dict]] = None,
|
|
ids: Optional[List[str]] = None,
|
|
) -> List[str]:
|
|
await self._ensure_table()
|
|
if not self.embedding_function:
|
|
raise ValueError("embedding_function is required for add_texts")
|
|
|
|
inserted_ids: List[str] = []
|
|
for i, text in enumerate(texts):
|
|
vec = self.embedding_function(text)
|
|
vec_str = "[" + ",".join(str(v) for v in vec) + "]"
|
|
|
|
meta = metadatas[i] if metadatas and i < len(metadatas) else {}
|
|
meta_cols = []
|
|
meta_vals = []
|
|
if self.tenant_id:
|
|
meta_cols.append("tenant_id")
|
|
meta_vals.append(f"'{self.tenant_id}'")
|
|
for mc in self.metadata_cols:
|
|
if mc in meta:
|
|
meta_cols.append(mc)
|
|
meta_vals.append(f"'{meta[mc]}'")
|
|
|
|
col_list = f"{self.embedding_col}, {self.text_col}"
|
|
val_list = f"'{vec_str}', '{text.replace(\"'\", \"''\")}'"
|
|
if meta_cols:
|
|
col_list += ", " + ", ".join(meta_cols)
|
|
val_list += ", " + ", ".join(meta_vals)
|
|
|
|
sql = f"INSERT INTO {self.table} ({col_list}) VALUES ({val_list}) RETURNING id"
|
|
result = await self.client.query(sql)
|
|
if result.rows:
|
|
inserted_ids.append(result.rows[0].get("id", str(i)))
|
|
else:
|
|
inserted_ids.append(str(i))
|
|
return inserted_ids
|
|
|
|
async def similarity_search(
|
|
self, query: str, k: int = 4, filter_col: Optional[str] = None, filter_val: Optional[str] = None
|
|
) -> List[Tuple[Any, float]]:
|
|
await self._ensure_table()
|
|
if not self.embedding_function:
|
|
raise ValueError("embedding_function is required for similarity_search")
|
|
|
|
vec = self.embedding_function(query)
|
|
vec_str = "[" + ",".join(str(v) for v in vec) + "]"
|
|
|
|
# Set tenant session variable if multi-tenant
|
|
if self.tenant_id:
|
|
await self.client.query(f"SET app.tenant_id = '{self.tenant_id}'")
|
|
|
|
if filter_col and filter_val:
|
|
sql = f"SELECT hybrid_search_filtered('{self.table}', '{self.embedding_col}', '{self.text_col}', '{query.replace(\"'\", \"''\")}', '{vec_str}', {k}, '{filter_col}', '{filter_val}') AS res"
|
|
else:
|
|
sql = f"SELECT hybrid_search('{self.table}', '{self.embedding_col}', '{self.text_col}', '{query.replace(\"'\", \"''\")}', '{vec_str}', {k}) AS res"
|
|
|
|
result = await self.client.query(sql)
|
|
if not result.rows:
|
|
return []
|
|
|
|
raw = result.rows[0].get("res", "[]")
|
|
try:
|
|
arr = json.loads(raw)
|
|
except:
|
|
return []
|
|
|
|
docs: List[Tuple[Any, float]] = []
|
|
for item in arr:
|
|
doc_id = item.get("id", "")
|
|
score = float(item.get("score", 0))
|
|
# Fetch full row
|
|
row_result = await self.client.query(f"SELECT * FROM {self.table} WHERE id = {doc_id}")
|
|
if row_result.rows:
|
|
page_content = row_result.rows[0].get(self.text_col, "")
|
|
metadata = dict(row_result.rows[0])
|
|
# Wrap in a simple Document-like object
|
|
doc = _SimpleDocument(page_content=page_content, metadata=metadata)
|
|
docs.append((doc, score))
|
|
return docs
|
|
|
|
async def max_marginal_relevance_search(
|
|
self, query: str, k: int = 4, fetch_k: int = 20, lambda_mult: float = 0.5
|
|
) -> List[Any]:
|
|
"""MMR: diversify results while maintaining relevance."""
|
|
await self._ensure_table()
|
|
# Fetch more candidates
|
|
candidates = await self.similarity_search(query, k=fetch_k)
|
|
if not candidates:
|
|
return []
|
|
|
|
# Simple MMR: greedily select docs that maximize lambda*relevance - (1-lambda)*max_similarity_to_selected
|
|
selected: List[Tuple[Any, float]] = []
|
|
remaining = list(candidates)
|
|
|
|
while len(selected) < k and remaining:
|
|
best_score = -float("inf")
|
|
best_idx = 0
|
|
for i, (doc, rel_score) in enumerate(remaining):
|
|
# Penalize similarity to already selected docs
|
|
penalty = 0.0
|
|
for sel_doc, _ in selected:
|
|
penalty = max(penalty, _doc_similarity(doc, sel_doc))
|
|
mmr_score = lambda_mult * rel_score - (1 - lambda_mult) * penalty
|
|
if mmr_score > best_score:
|
|
best_score = mmr_score
|
|
best_idx = i
|
|
selected.append(remaining.pop(best_idx))
|
|
|
|
return [doc for doc, _ in selected]
|
|
|
|
async def delete(self, ids: Optional[List[str]] = None) -> None:
|
|
await self._ensure_table()
|
|
if ids:
|
|
id_list = ", ".join(str(i) for i in ids)
|
|
await self.client.query(f"DELETE FROM {self.table} WHERE id IN ({id_list})")
|
|
|
|
async def set_tenant(self, tenant_id: str) -> None:
|
|
self.tenant_id = tenant_id
|
|
await self.client.query(f"SET app.tenant_id = '{tenant_id}'")
|
|
|
|
|
|
class _SimpleDocument:
|
|
def __init__(self, page_content: str, metadata: dict):
|
|
self.page_content = page_content
|
|
self.metadata = metadata
|
|
|
|
def __repr__(self):
|
|
return f"Document(content={self.page_content[:50]}..., metadata={self.metadata})"
|
|
|
|
|
|
def _doc_similarity(a: _SimpleDocument, b: _SimpleDocument) -> float:
|
|
"""Simple Jaccard similarity on text tokens."""
|
|
tokens_a = set(a.page_content.lower().split())
|
|
tokens_b = set(b.page_content.lower().split())
|
|
if not tokens_a or not tokens_b:
|
|
return 0.0
|
|
intersection = tokens_a & tokens_b
|
|
union = tokens_a | tokens_b
|
|
return len(intersection) / len(union)
|