pieter/eveAI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Correcting the retrieval of relevant documents

Browse Source
This commit is contained in:
Josako
2024-06-12 16:15:48 +02:00
parent be311c440b
commit fd510c8fcd
8 changed files with 62 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
common/langchain/EveAIRetriever.py
View File

@@ -4,19 +4,21 @@ 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 common.models.document import Document, DocumentVersion, Embedding
from common.models.document import Document, DocumentVersion
from common.utils.datetime_utils import get_date_in_timezone
class EveAIRetriever(BaseRetriever):
model_variables: Dict[str, Any] = Field(...)
tenant_info: Dict[str, Any] = Field(...)
def __init__(self, model_variables: Dict[str, Any]):
def __init__(self, model_variables: Dict[str, Any], tenant_info: Dict[str, Any]):
super().__init__()
current_app.logger.debug('Initializing EveAIRetriever')
self.model_variables = model_variables
self.tenant_info = tenant_info
current_app.logger.debug('EveAIRetriever initialized')
def _get_relevant_documents(self, query: str):
@@ -27,7 +29,7 @@ class EveAIRetriever(BaseRetriever):
k = self.model_variables['k']
try:
current_date = date.today()
current_date = get_date_in_timezone(self.tenant_info['timezone'])
# Subquery to find the latest version of each document
subquery = (
db.session.query(
@@ -53,18 +55,17 @@ class EveAIRetriever(BaseRetriever):
.limit(k)
)
# 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'---------------------------------------')
if self.tenant_info['rag_tuning']:
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')
if self.tenant_info['rag_tuning']:
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: