mirror of
https://github.com/open-webui/open-webui
synced 2025-01-04 01:53:31 +00:00
355 lines
12 KiB
Python
355 lines
12 KiB
Python
from typing import Optional, List, Dict, Any
|
|
from sqlalchemy import (
|
|
cast,
|
|
column,
|
|
create_engine,
|
|
Column,
|
|
Integer,
|
|
select,
|
|
text,
|
|
Text,
|
|
values,
|
|
)
|
|
from sqlalchemy.sql import true
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from sqlalchemy.orm import declarative_base, scoped_session, sessionmaker
|
|
from sqlalchemy.dialects.postgresql import JSONB, array
|
|
from pgvector.sqlalchemy import Vector
|
|
from sqlalchemy.ext.mutable import MutableDict
|
|
|
|
from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
|
|
from open_webui.config import PGVECTOR_DB_URL
|
|
|
|
VECTOR_LENGTH = 1536
|
|
Base = declarative_base()
|
|
|
|
|
|
class DocumentChunk(Base):
|
|
__tablename__ = "document_chunk"
|
|
|
|
id = Column(Text, primary_key=True)
|
|
vector = Column(Vector(dim=VECTOR_LENGTH), nullable=True)
|
|
collection_name = Column(Text, nullable=False)
|
|
text = Column(Text, nullable=True)
|
|
vmetadata = Column(MutableDict.as_mutable(JSONB), nullable=True)
|
|
|
|
|
|
class PgvectorClient:
|
|
def __init__(self) -> None:
|
|
|
|
# if no pgvector uri, use the existing database connection
|
|
if not PGVECTOR_DB_URL:
|
|
from open_webui.internal.db import Session
|
|
|
|
self.session = Session
|
|
else:
|
|
engine = create_engine(
|
|
PGVECTOR_DB_URL, pool_pre_ping=True, poolclass=NullPool
|
|
)
|
|
SessionLocal = sessionmaker(
|
|
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
|
|
)
|
|
self.session = scoped_session(SessionLocal)
|
|
|
|
try:
|
|
# Ensure the pgvector extension is available
|
|
self.session.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
|
|
|
|
# Create the tables if they do not exist
|
|
# Base.metadata.create_all requires a bind (engine or connection)
|
|
# Get the connection from the session
|
|
connection = self.session.connection()
|
|
Base.metadata.create_all(bind=connection)
|
|
|
|
# Create an index on the vector column if it doesn't exist
|
|
self.session.execute(
|
|
text(
|
|
"CREATE INDEX IF NOT EXISTS idx_document_chunk_vector "
|
|
"ON document_chunk USING ivfflat (vector vector_cosine_ops) WITH (lists = 100);"
|
|
)
|
|
)
|
|
self.session.execute(
|
|
text(
|
|
"CREATE INDEX IF NOT EXISTS idx_document_chunk_collection_name "
|
|
"ON document_chunk (collection_name);"
|
|
)
|
|
)
|
|
self.session.commit()
|
|
print("Initialization complete.")
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
print(f"Error during initialization: {e}")
|
|
raise
|
|
|
|
def adjust_vector_length(self, vector: List[float]) -> List[float]:
|
|
# Adjust vector to have length VECTOR_LENGTH
|
|
current_length = len(vector)
|
|
if current_length < VECTOR_LENGTH:
|
|
# Pad the vector with zeros
|
|
vector += [0.0] * (VECTOR_LENGTH - current_length)
|
|
elif current_length > VECTOR_LENGTH:
|
|
raise Exception(
|
|
f"Vector length {current_length} not supported. Max length must be <= {VECTOR_LENGTH}"
|
|
)
|
|
return vector
|
|
|
|
def insert(self, collection_name: str, items: List[VectorItem]) -> None:
|
|
try:
|
|
new_items = []
|
|
for item in items:
|
|
vector = self.adjust_vector_length(item["vector"])
|
|
new_chunk = DocumentChunk(
|
|
id=item["id"],
|
|
vector=vector,
|
|
collection_name=collection_name,
|
|
text=item["text"],
|
|
vmetadata=item["metadata"],
|
|
)
|
|
new_items.append(new_chunk)
|
|
self.session.bulk_save_objects(new_items)
|
|
self.session.commit()
|
|
print(
|
|
f"Inserted {len(new_items)} items into collection '{collection_name}'."
|
|
)
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
print(f"Error during insert: {e}")
|
|
raise
|
|
|
|
def upsert(self, collection_name: str, items: List[VectorItem]) -> None:
|
|
try:
|
|
for item in items:
|
|
vector = self.adjust_vector_length(item["vector"])
|
|
existing = (
|
|
self.session.query(DocumentChunk)
|
|
.filter(DocumentChunk.id == item["id"])
|
|
.first()
|
|
)
|
|
if existing:
|
|
existing.vector = vector
|
|
existing.text = item["text"]
|
|
existing.vmetadata = item["metadata"]
|
|
existing.collection_name = (
|
|
collection_name # Update collection_name if necessary
|
|
)
|
|
else:
|
|
new_chunk = DocumentChunk(
|
|
id=item["id"],
|
|
vector=vector,
|
|
collection_name=collection_name,
|
|
text=item["text"],
|
|
vmetadata=item["metadata"],
|
|
)
|
|
self.session.add(new_chunk)
|
|
self.session.commit()
|
|
print(f"Upserted {len(items)} items into collection '{collection_name}'.")
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
print(f"Error during upsert: {e}")
|
|
raise
|
|
|
|
def search(
|
|
self,
|
|
collection_name: str,
|
|
vectors: List[List[float]],
|
|
limit: Optional[int] = None,
|
|
) -> Optional[SearchResult]:
|
|
try:
|
|
if not vectors:
|
|
return None
|
|
|
|
# Adjust query vectors to VECTOR_LENGTH
|
|
vectors = [self.adjust_vector_length(vector) for vector in vectors]
|
|
num_queries = len(vectors)
|
|
|
|
def vector_expr(vector):
|
|
return cast(array(vector), Vector(VECTOR_LENGTH))
|
|
|
|
# Create the values for query vectors
|
|
qid_col = column("qid", Integer)
|
|
q_vector_col = column("q_vector", Vector(VECTOR_LENGTH))
|
|
query_vectors = (
|
|
values(qid_col, q_vector_col)
|
|
.data(
|
|
[(idx, vector_expr(vector)) for idx, vector in enumerate(vectors)]
|
|
)
|
|
.alias("query_vectors")
|
|
)
|
|
|
|
# Build the lateral subquery for each query vector
|
|
subq = (
|
|
select(
|
|
DocumentChunk.id,
|
|
DocumentChunk.text,
|
|
DocumentChunk.vmetadata,
|
|
(
|
|
DocumentChunk.vector.cosine_distance(query_vectors.c.q_vector)
|
|
).label("distance"),
|
|
)
|
|
.where(DocumentChunk.collection_name == collection_name)
|
|
.order_by(
|
|
(DocumentChunk.vector.cosine_distance(query_vectors.c.q_vector))
|
|
)
|
|
)
|
|
if limit is not None:
|
|
subq = subq.limit(limit)
|
|
subq = subq.lateral("result")
|
|
|
|
# Build the main query by joining query_vectors and the lateral subquery
|
|
stmt = (
|
|
select(
|
|
query_vectors.c.qid,
|
|
subq.c.id,
|
|
subq.c.text,
|
|
subq.c.vmetadata,
|
|
subq.c.distance,
|
|
)
|
|
.select_from(query_vectors)
|
|
.join(subq, true())
|
|
.order_by(query_vectors.c.qid, subq.c.distance)
|
|
)
|
|
|
|
result_proxy = self.session.execute(stmt)
|
|
results = result_proxy.all()
|
|
|
|
ids = [[] for _ in range(num_queries)]
|
|
distances = [[] for _ in range(num_queries)]
|
|
documents = [[] for _ in range(num_queries)]
|
|
metadatas = [[] for _ in range(num_queries)]
|
|
|
|
if not results:
|
|
return SearchResult(
|
|
ids=ids,
|
|
distances=distances,
|
|
documents=documents,
|
|
metadatas=metadatas,
|
|
)
|
|
|
|
for row in results:
|
|
qid = int(row.qid)
|
|
ids[qid].append(row.id)
|
|
distances[qid].append(row.distance)
|
|
documents[qid].append(row.text)
|
|
metadatas[qid].append(row.vmetadata)
|
|
|
|
return SearchResult(
|
|
ids=ids, distances=distances, documents=documents, metadatas=metadatas
|
|
)
|
|
except Exception as e:
|
|
print(f"Error during search: {e}")
|
|
return None
|
|
|
|
def query(
|
|
self, collection_name: str, filter: Dict[str, Any], limit: Optional[int] = None
|
|
) -> Optional[GetResult]:
|
|
try:
|
|
query = self.session.query(DocumentChunk).filter(
|
|
DocumentChunk.collection_name == collection_name
|
|
)
|
|
|
|
for key, value in filter.items():
|
|
query = query.filter(DocumentChunk.vmetadata[key].astext == str(value))
|
|
|
|
if limit is not None:
|
|
query = query.limit(limit)
|
|
|
|
results = query.all()
|
|
|
|
if not results:
|
|
return None
|
|
|
|
ids = [[result.id for result in results]]
|
|
documents = [[result.text for result in results]]
|
|
metadatas = [[result.vmetadata for result in results]]
|
|
|
|
return GetResult(
|
|
ids=ids,
|
|
documents=documents,
|
|
metadatas=metadatas,
|
|
)
|
|
except Exception as e:
|
|
print(f"Error during query: {e}")
|
|
return None
|
|
|
|
def get(
|
|
self, collection_name: str, limit: Optional[int] = None
|
|
) -> Optional[GetResult]:
|
|
try:
|
|
query = self.session.query(DocumentChunk).filter(
|
|
DocumentChunk.collection_name == collection_name
|
|
)
|
|
if limit is not None:
|
|
query = query.limit(limit)
|
|
|
|
results = query.all()
|
|
|
|
if not results:
|
|
return None
|
|
|
|
ids = [[result.id for result in results]]
|
|
documents = [[result.text for result in results]]
|
|
metadatas = [[result.vmetadata for result in results]]
|
|
|
|
return GetResult(ids=ids, documents=documents, metadatas=metadatas)
|
|
except Exception as e:
|
|
print(f"Error during get: {e}")
|
|
return None
|
|
|
|
def delete(
|
|
self,
|
|
collection_name: str,
|
|
ids: Optional[List[str]] = None,
|
|
filter: Optional[Dict[str, Any]] = None,
|
|
) -> None:
|
|
try:
|
|
query = self.session.query(DocumentChunk).filter(
|
|
DocumentChunk.collection_name == collection_name
|
|
)
|
|
if ids:
|
|
query = query.filter(DocumentChunk.id.in_(ids))
|
|
if filter:
|
|
for key, value in filter.items():
|
|
query = query.filter(
|
|
DocumentChunk.vmetadata[key].astext == str(value)
|
|
)
|
|
deleted = query.delete(synchronize_session=False)
|
|
self.session.commit()
|
|
print(f"Deleted {deleted} items from collection '{collection_name}'.")
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
print(f"Error during delete: {e}")
|
|
raise
|
|
|
|
def reset(self) -> None:
|
|
try:
|
|
deleted = self.session.query(DocumentChunk).delete()
|
|
self.session.commit()
|
|
print(
|
|
f"Reset complete. Deleted {deleted} items from 'document_chunk' table."
|
|
)
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
print(f"Error during reset: {e}")
|
|
raise
|
|
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
def has_collection(self, collection_name: str) -> bool:
|
|
try:
|
|
exists = (
|
|
self.session.query(DocumentChunk)
|
|
.filter(DocumentChunk.collection_name == collection_name)
|
|
.first()
|
|
is not None
|
|
)
|
|
return exists
|
|
except Exception as e:
|
|
print(f"Error checking collection existence: {e}")
|
|
return False
|
|
|
|
def delete_collection(self, collection_name: str) -> None:
|
|
self.delete(collection_name)
|
|
print(f"Collection '{collection_name}' deleted.")
|