Correcting the retrieval of relevant documents

common/langchain/EveAIRetriever.py

@@ -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:

common/models/user.py

@@ -67,6 +67,7 @@ class Tenant(db.Model):
             'website': self.website,
             'default_language': self.default_language,
             'allowed_languages': self.allowed_languages,
+            'timezone': self.timezone,
             'embedding_model': self.embedding_model,
             'llm_model': self.llm_model,
             'license_start_date': self.license_start_date,

common/utils/datetime_utils.py

@@ -0,0 +1,19 @@
+from datetime import datetime
+import pytz
+
+
+def get_date_in_timezone(timezone_str):
+    try:
+        # Get the timezone object from the string
+        timezone = pytz.timezone(timezone_str)
+
+        # Get the current time in the specified timezone
+        current_time = datetime.now(timezone)
+
+        # Extract the date part
+        current_date = current_time.date()
+
+        return current_date
+    except Exception as e:
+        print(f"Error getting date in timezone {timezone_str}: {e}")
+        return None