UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-03-15 10:09:31 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Merge pull request #8385 from jk-f5/refac/pgvectorlength

Browse Source 
fix: Pgvector vector column size check was failing on initialization …
This commit is contained in:
Timothy Jaeryang Baek 2025-01-07 14:35:47 -08:00 committed by GitHub
commit 81d03540bd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 23 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
backend/open_webui/retrieval/vector/dbs/pgvector.py
View File

@ -9,6 +9,7 @@ from sqlalchemy import (
select, select,
text, text,
Text, Text,
Table,
values, values,
) )
from sqlalchemy.sql import true from sqlalchemy.sql import true
@ -18,6 +19,7 @@ from sqlalchemy.orm import declarative_base, scoped_session, sessionmaker
from sqlalchemy.dialects.postgresql import JSONB, array from sqlalchemy.dialects.postgresql import JSONB, array
from pgvector.sqlalchemy import Vector from pgvector.sqlalchemy import Vector
from sqlalchemy.ext.mutable import MutableDict from sqlalchemy.ext.mutable import MutableDict
from sqlalchemy.exc import NoSuchTableError
from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
from open_webui.config import PGVECTOR_DB_URL, PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH from open_webui.config import PGVECTOR_DB_URL, PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH
@ -92,10 +94,16 @@ class PgvectorClient:
Raises an exception if there is a mismatch. Raises an exception if there is a mismatch.
""" """
metadata = MetaData() metadata = MetaData()
metadata.reflect(bind=self.session.bind, only=["document_chunk"]) try:
# Attempt to reflect the 'document_chunk' table
document_chunk_table = Table(
"document_chunk", metadata, autoload_with=self.session.bind
)
except NoSuchTableError:
# Table does not exist; no action needed
return
if "document_chunk" in metadata.tables: # Proceed to check the vector column
document_chunk_table = metadata.tables["document_chunk"]
if "vector" in document_chunk_table.columns: if "vector" in document_chunk_table.columns:
vector_column = document_chunk_table.columns["vector"] vector_column = document_chunk_table.columns["vector"]
vector_type = vector_column.type vector_type = vector_column.type
@ -114,9 +122,6 @@ class PgvectorClient:
raise Exception( raise Exception(
"The 'vector' column does not exist in the 'document_chunk' table." "The 'vector' column does not exist in the 'document_chunk' table."
) )
else:
# Table does not exist yet; no action needed
pass
def adjust_vector_length(self, vector: List[float]) -> List[float]: def adjust_vector_length(self, vector: List[float]) -> List[float]:
# Adjust vector to have length VECTOR_LENGTH # Adjust vector to have length VECTOR_LENGTH