Improving chat functionality significantly throughout the application.

This commit is contained in:
Josako
2024-06-12 11:07:18 +02:00
parent 27b6de8734
commit be311c440b
22 changed files with 604 additions and 127 deletions

View File

@@ -1,11 +1,13 @@
from langchain_core.retrievers import BaseRetriever
from sqlalchemy import func, and_, or_
from sqlalchemy.exc import SQLAlchemyError
from pydantic import BaseModel, Field
from typing import Any, Dict
from flask import current_app
from datetime import date
from common.extensions import db
from flask import current_app
from config.logging_config import LOGGING
from common.models.document import Document, DocumentVersion, Embedding
class EveAIRetriever(BaseRetriever):
@@ -23,26 +25,53 @@ class EveAIRetriever(BaseRetriever):
db_class = self.model_variables['embedding_db_model']
similarity_threshold = self.model_variables['similarity_threshold']
k = self.model_variables['k']
try:
res = (
current_date = date.today()
# Subquery to find the latest version of each document
subquery = (
db.session.query(
DocumentVersion.doc_id,
func.max(DocumentVersion.id).label('latest_version_id')
)
.group_by(DocumentVersion.doc_id)
.subquery()
)
# Main query to filter embeddings
query_obj = (
db.session.query(db_class,
db_class.embedding.cosine_distance(query_embedding)
.label('distance'))
.filter(db_class.embedding.cosine_distance(query_embedding) < similarity_threshold)
db_class.embedding.cosine_distance(query_embedding).label('distance'))
.join(DocumentVersion, db_class.doc_vers_id == DocumentVersion.id)
.join(Document, DocumentVersion.doc_id == Document.id)
.join(subquery, DocumentVersion.id == subquery.c.latest_version_id)
.filter(
or_(Document.valid_from.is_(None), Document.valid_from <= current_date),
or_(Document.valid_to.is_(None), Document.valid_to >= current_date),
db_class.embedding.cosine_distance(query_embedding) < similarity_threshold
)
.order_by('distance')
.limit(k)
.all()
)
current_app.rag_tuning_logger.debug(f'Retrieved {len(res)} relevant documents')
current_app.rag_tuning_logger.debug(f'---------------------------------------')
# Print the generated SQL statement for debugging
current_app.logger.debug("SQL Statement:\n")
current_app.logger.debug(query_obj.statement.compile(compile_kwargs={"literal_binds": True}))
res = query_obj.all()
# current_app.rag_tuning_logger.debug(f'Retrieved {len(res)} relevant documents')
# current_app.rag_tuning_logger.debug(f'---------------------------------------')
result = []
for doc in res:
current_app.rag_tuning_logger.debug(f'Document ID: {doc[0].id} - Distance: {doc[1]}\n')
current_app.rag_tuning_logger.debug(f'Chunk: \n {doc[0].chunk}\n\n')
# current_app.rag_tuning_logger.debug(f'Document ID: {doc[0].id} - Distance: {doc[1]}\n')
# current_app.rag_tuning_logger.debug(f'Chunk: \n {doc[0].chunk}\n\n')
result.append(f'SOURCE: {doc[0].id}\n\n{doc[0].chunk}\n\n')
except SQLAlchemyError as e:
current_app.logger.error(f'Error retrieving relevant documents: {e}')
db.session.rollback()
return []
return res
return result
def _get_query_embedding(self, query: str):
embedding_model = self.model_variables['embedding_model']