Improving chat functionality significantly throughout the application.
This commit is contained in:
@@ -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']
|
||||
|
||||
Reference in New Issue
Block a user