|
@@ -1,4 +1,5 @@
|
|
from typing import Optional, List, Dict, Any
|
|
from typing import Optional, List, Dict, Any
|
|
|
|
+import logging
|
|
from sqlalchemy import (
|
|
from sqlalchemy import (
|
|
cast,
|
|
cast,
|
|
column,
|
|
column,
|
|
@@ -24,9 +25,14 @@ 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
|
|
|
|
|
|
|
|
+from open_webui.env import SRC_LOG_LEVELS
|
|
|
|
+
|
|
VECTOR_LENGTH = PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH
|
|
VECTOR_LENGTH = PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH
|
|
Base = declarative_base()
|
|
Base = declarative_base()
|
|
|
|
|
|
|
|
+log = logging.getLogger(__name__)
|
|
|
|
+log.setLevel(SRC_LOG_LEVELS["RAG"])
|
|
|
|
+
|
|
|
|
|
|
class DocumentChunk(Base):
|
|
class DocumentChunk(Base):
|
|
__tablename__ = "document_chunk"
|
|
__tablename__ = "document_chunk"
|
|
@@ -82,10 +88,10 @@ class PgvectorClient:
|
|
)
|
|
)
|
|
)
|
|
)
|
|
self.session.commit()
|
|
self.session.commit()
|
|
- print("Initialization complete.")
|
|
|
|
|
|
+ log.info("Initialization complete.")
|
|
except Exception as e:
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
self.session.rollback()
|
|
- print(f"Error during initialization: {e}")
|
|
|
|
|
|
+ log.exception(f"Error during initialization: {e}")
|
|
raise
|
|
raise
|
|
|
|
|
|
def check_vector_length(self) -> None:
|
|
def check_vector_length(self) -> None:
|
|
@@ -150,12 +156,12 @@ class PgvectorClient:
|
|
new_items.append(new_chunk)
|
|
new_items.append(new_chunk)
|
|
self.session.bulk_save_objects(new_items)
|
|
self.session.bulk_save_objects(new_items)
|
|
self.session.commit()
|
|
self.session.commit()
|
|
- print(
|
|
|
|
|
|
+ log.info(
|
|
f"Inserted {len(new_items)} items into collection '{collection_name}'."
|
|
f"Inserted {len(new_items)} items into collection '{collection_name}'."
|
|
)
|
|
)
|
|
except Exception as e:
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
self.session.rollback()
|
|
- print(f"Error during insert: {e}")
|
|
|
|
|
|
+ log.exception(f"Error during insert: {e}")
|
|
raise
|
|
raise
|
|
|
|
|
|
def upsert(self, collection_name: str, items: List[VectorItem]) -> None:
|
|
def upsert(self, collection_name: str, items: List[VectorItem]) -> None:
|
|
@@ -184,10 +190,10 @@ class PgvectorClient:
|
|
)
|
|
)
|
|
self.session.add(new_chunk)
|
|
self.session.add(new_chunk)
|
|
self.session.commit()
|
|
self.session.commit()
|
|
- print(f"Upserted {len(items)} items into collection '{collection_name}'.")
|
|
|
|
|
|
+ log.info(f"Upserted {len(items)} items into collection '{collection_name}'.")
|
|
except Exception as e:
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
self.session.rollback()
|
|
- print(f"Error during upsert: {e}")
|
|
|
|
|
|
+ log.exception(f"Error during upsert: {e}")
|
|
raise
|
|
raise
|
|
|
|
|
|
def search(
|
|
def search(
|
|
@@ -278,7 +284,7 @@ class PgvectorClient:
|
|
ids=ids, distances=distances, documents=documents, metadatas=metadatas
|
|
ids=ids, distances=distances, documents=documents, metadatas=metadatas
|
|
)
|
|
)
|
|
except Exception as e:
|
|
except Exception as e:
|
|
- print(f"Error during search: {e}")
|
|
|
|
|
|
+ log.exception(f"Error during search: {e}")
|
|
return None
|
|
return None
|
|
|
|
|
|
def query(
|
|
def query(
|
|
@@ -310,7 +316,7 @@ class PgvectorClient:
|
|
metadatas=metadatas,
|
|
metadatas=metadatas,
|
|
)
|
|
)
|
|
except Exception as e:
|
|
except Exception as e:
|
|
- print(f"Error during query: {e}")
|
|
|
|
|
|
+ log.exception(f"Error during query: {e}")
|
|
return None
|
|
return None
|
|
|
|
|
|
def get(
|
|
def get(
|
|
@@ -334,7 +340,7 @@ class PgvectorClient:
|
|
|
|
|
|
return GetResult(ids=ids, documents=documents, metadatas=metadatas)
|
|
return GetResult(ids=ids, documents=documents, metadatas=metadatas)
|
|
except Exception as e:
|
|
except Exception as e:
|
|
- print(f"Error during get: {e}")
|
|
|
|
|
|
+ log.exception(f"Error during get: {e}")
|
|
return None
|
|
return None
|
|
|
|
|
|
def delete(
|
|
def delete(
|
|
@@ -356,22 +362,22 @@ class PgvectorClient:
|
|
)
|
|
)
|
|
deleted = query.delete(synchronize_session=False)
|
|
deleted = query.delete(synchronize_session=False)
|
|
self.session.commit()
|
|
self.session.commit()
|
|
- print(f"Deleted {deleted} items from collection '{collection_name}'.")
|
|
|
|
|
|
+ log.info(f"Deleted {deleted} items from collection '{collection_name}'.")
|
|
except Exception as e:
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
self.session.rollback()
|
|
- print(f"Error during delete: {e}")
|
|
|
|
|
|
+ log.exception(f"Error during delete: {e}")
|
|
raise
|
|
raise
|
|
|
|
|
|
def reset(self) -> None:
|
|
def reset(self) -> None:
|
|
try:
|
|
try:
|
|
deleted = self.session.query(DocumentChunk).delete()
|
|
deleted = self.session.query(DocumentChunk).delete()
|
|
self.session.commit()
|
|
self.session.commit()
|
|
- print(
|
|
|
|
|
|
+ log.info(
|
|
f"Reset complete. Deleted {deleted} items from 'document_chunk' table."
|
|
f"Reset complete. Deleted {deleted} items from 'document_chunk' table."
|
|
)
|
|
)
|
|
except Exception as e:
|
|
except Exception as e:
|
|
self.session.rollback()
|
|
self.session.rollback()
|
|
- print(f"Error during reset: {e}")
|
|
|
|
|
|
+ log.exception(f"Error during reset: {e}")
|
|
raise
|
|
raise
|
|
|
|
|
|
def close(self) -> None:
|
|
def close(self) -> None:
|
|
@@ -387,9 +393,9 @@ class PgvectorClient:
|
|
)
|
|
)
|
|
return exists
|
|
return exists
|
|
except Exception as e:
|
|
except Exception as e:
|
|
- print(f"Error checking collection existence: {e}")
|
|
|
|
|
|
+ log.exception(f"Error checking collection existence: {e}")
|
|
return False
|
|
return False
|
|
|
|
|
|
def delete_collection(self, collection_name: str) -> None:
|
|
def delete_collection(self, collection_name: str) -> None:
|
|
self.delete(collection_name)
|
|
self.delete(collection_name)
|
|
- print(f"Collection '{collection_name}' deleted.")
|
|
|
|
|
|
+ log.info(f"Collection '{collection_name}' deleted.")
|