- Organise retrievers

common/langchain/retrievers/eveai_default_rag_retriever.py

@@ -0,0 +1,145 @@
+from langchain_core.retrievers import BaseRetriever
+from sqlalchemy import func, and_, or_, desc
+from sqlalchemy.exc import SQLAlchemyError
+from pydantic import BaseModel, Field, PrivateAttr
+from typing import Any, Dict
+from flask import current_app
+
+from common.extensions import db
+from common.models.document import Document, DocumentVersion
+from common.utils.datetime_utils import get_date_in_timezone
+from common.utils.model_utils import ModelVariables
+
+
+class EveAIDefaultRagRetriever(BaseRetriever, BaseModel):
+    _catalog_id: int = PrivateAttr()
+    _model_variables: ModelVariables = PrivateAttr()
+    _tenant_info: Dict[str, Any] = PrivateAttr()
+
+    def __init__(self, catalog_id: int, model_variables: ModelVariables, tenant_info: Dict[str, Any]):
+        super().__init__()
+        current_app.logger.debug(f'Model variables type: {type(model_variables)}')
+        self._catalog_id = catalog_id
+        self._model_variables = model_variables
+        self._tenant_info = tenant_info
+
+    @property
+    def catalog_id(self) -> int:
+        return self._catalog_id
+
+    @property
+    def model_variables(self) -> ModelVariables:
+        return self._model_variables
+
+    @property
+    def tenant_info(self) -> Dict[str, Any]:
+        return self._tenant_info
+
+    def _get_relevant_documents(self, query: str):
+        current_app.logger.debug(f'Retrieving relevant documents for query: {query}')
+        query_embedding = self._get_query_embedding(query)
+        current_app.logger.debug(f'Model Variables Private: {type(self._model_variables)}')
+        current_app.logger.debug(f'Model Variables Property: {type(self.model_variables)}')
+        db_class = self.model_variables['embedding_db_model']
+        similarity_threshold = self.model_variables['similarity_threshold']
+        k = self.model_variables['k']
+
+        if self.model_variables['rag_tuning']:
+            try:
+                current_date = get_date_in_timezone(self.tenant_info['timezone'])
+                current_app.rag_tuning_logger.debug(f'Current date: {current_date}\n')
+
+                # Debug query to show similarity for all valid documents (without chunk text)
+                debug_query = (
+                    db.session.query(
+                        Document.id.label('document_id'),
+                        DocumentVersion.id.label('version_id'),
+                        db_class.id.label('embedding_id'),
+                        (1 - db_class.embedding.cosine_distance(query_embedding)).label('similarity')
+                    )
+                    .join(DocumentVersion, db_class.doc_vers_id == DocumentVersion.id)
+                    .join(Document, DocumentVersion.doc_id == Document.id)
+                    .filter(
+                        or_(Document.valid_from.is_(None), func.date(Document.valid_from) <= current_date),
+                        or_(Document.valid_to.is_(None), func.date(Document.valid_to) >= current_date)
+                    )
+                    .order_by(desc('similarity'))
+                )
+
+                debug_results = debug_query.all()
+
+                current_app.logger.debug("Debug: Similarity for all valid documents:")
+                for row in debug_results:
+                    current_app.rag_tuning_logger.debug(f"Doc ID: {row.document_id}, "
+                                                        f"Version ID: {row.version_id}, "
+                                                        f"Embedding ID: {row.embedding_id}, "
+                                                        f"Similarity: {row.similarity}")
+                current_app.rag_tuning_logger.debug(f'---------------------------------------\n')
+            except SQLAlchemyError as e:
+                current_app.logger.error(f'Error generating overview: {e}')
+                db.session.rollback()
+
+        if self.model_variables['rag_tuning']:
+            current_app.rag_tuning_logger.debug(f'Parameters for Retrieval of documents: \n')
+            current_app.rag_tuning_logger.debug(f'Similarity Threshold: {similarity_threshold}\n')
+            current_app.rag_tuning_logger.debug(f'K: {k}\n')
+            current_app.rag_tuning_logger.debug(f'---------------------------------------\n')
+
+        try:
+            current_date = get_date_in_timezone(self.tenant_info['timezone'])
+            # 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,
+                                 (1 - db_class.embedding.cosine_distance(query_embedding)).label('similarity'))
+                .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), func.date(Document.valid_from) <= current_date),
+                    or_(Document.valid_to.is_(None), func.date(Document.valid_to) >= current_date),
+                    (1 - db_class.embedding.cosine_distance(query_embedding)) > similarity_threshold,
+                    Document.catalog_id == self._catalog_id
+                )
+                .order_by(desc('similarity'))
+                .limit(k)
+            )
+
+            if self.model_variables['rag_tuning']:
+                current_app.rag_tuning_logger.debug(f'Query executed for Retrieval of documents: \n')
+                current_app.rag_tuning_logger.debug(f'{query_obj.statement}\n')
+                current_app.rag_tuning_logger.debug(f'---------------------------------------\n')
+
+            res = query_obj.all()
+
+            if self.model_variables['rag_tuning']:
+                current_app.rag_tuning_logger.debug(f'Retrieved {len(res)} relevant documents \n')
+                current_app.rag_tuning_logger.debug(f'Data retrieved: \n')
+                current_app.rag_tuning_logger.debug(f'{res}\n')
+                current_app.rag_tuning_logger.debug(f'---------------------------------------\n')
+
+            result = []
+            for doc in res:
+                if self.model_variables['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:
+            current_app.logger.error(f'Error retrieving relevant documents: {e}')
+            db.session.rollback()
+            return []
+        return result
+
+    def _get_query_embedding(self, query: str):
+        embedding_model = self.model_variables['embedding_model']
+        query_embedding = embedding_model.embed_query(query)
+        return query_embedding

common/langchain/retrievers/eveai_dossier_retriever.py

@@ -0,0 +1,154 @@
+from langchain_core.retrievers import BaseRetriever
+from sqlalchemy import func, and_, or_, desc, cast, JSON
+from sqlalchemy.exc import SQLAlchemyError
+from pydantic import BaseModel, Field, PrivateAttr
+from typing import Any, Dict, List, Optional
+from flask import current_app
+from contextlib import contextmanager
+
+from common.extensions import db
+from common.models.document import Document, DocumentVersion, Catalog
+from common.utils.datetime_utils import get_date_in_timezone
+from common.utils.model_utils import ModelVariables
+
+
+class EveAIDossierRetriever(BaseRetriever, BaseModel):
+    _catalog_id: int = PrivateAttr()
+    _model_variables: ModelVariables = PrivateAttr()
+    _tenant_info: Dict[str, Any] = PrivateAttr()
+    _active_filters: Optional[Dict[str, Any]] = PrivateAttr()
+
+    def __init__(self, catalog_id: int, model_variables: ModelVariables, tenant_info: Dict[str, Any]):
+        super().__init__()
+        self._catalog_id = catalog_id
+        self._model_variables = model_variables
+        self._tenant_info = tenant_info
+        self._active_filters = None
+
+    @contextmanager
+    def filtering(self, metadata_filters: Dict[str, Any]):
+        """Context manager for temporarily setting metadata filters"""
+        previous_filters = self._active_filters
+        self._active_filters = metadata_filters
+        try:
+            yield self
+        finally:
+            self._active_filters = previous_filters
+
+    def _build_metadata_filter_conditions(self, query):
+        """Build SQL conditions for metadata filtering"""
+        if not self._active_filters:
+            return query
+
+        conditions = []
+        for field, value in self._active_filters.items():
+            if value is None:
+                continue
+
+            # Handle both single values and lists of values
+            if isinstance(value, (list, tuple)):
+                # Multiple values - create OR condition
+                or_conditions = []
+                for val in value:
+                    or_conditions.append(
+                        cast(DocumentVersion.user_metadata[field].astext, JSON) == str(val)
+                    )
+                if or_conditions:
+                    conditions.append(or_(*or_conditions))
+            else:
+                # Single value - direct comparison
+                conditions.append(
+                    cast(DocumentVersion.user_metadata[field].astext, JSON) == str(value)
+                )
+
+        if conditions:
+            query = query.filter(and_(*conditions))
+
+        return query
+
+    def _get_relevant_documents(self, query: str):
+        current_app.logger.debug(f'Retrieving relevant documents for dossier query: {query}')
+        if self._active_filters:
+            current_app.logger.debug(f'Using metadata filters: {self._active_filters}')
+
+        query_embedding = self._get_query_embedding(query)
+        db_class = self.model_variables['embedding_db_model']
+        similarity_threshold = self.model_variables['similarity_threshold']
+        k = self.model_variables['k']
+
+        try:
+            current_date = get_date_in_timezone(self.tenant_info['timezone'])
+
+            # 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()
+            )
+
+            # Build base query
+            # Build base query
+            query_obj = (
+                db.session.query(db_class,
+                                 (1 - db_class.embedding.cosine_distance(query_embedding)).label('similarity'))
+                .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), func.date(Document.valid_from) <= current_date),
+                    or_(Document.valid_to.is_(None), func.date(Document.valid_to) >= current_date),
+                    (1 - db_class.embedding.cosine_distance(query_embedding)) > similarity_threshold,
+                    Document.catalog_id == self._catalog_id
+                )
+            )
+
+            # Apply metadata filters
+            query_obj = self._build_metadata_filter_conditions(query_obj)
+
+            # Order and limit results
+            query_obj = query_obj.order_by(desc('similarity')).limit(k)
+
+            # Debug logging for RAG tuning if enabled
+            if self.model_variables['rag_tuning']:
+                self._log_rag_tuning(query_obj, query_embedding)
+
+            res = query_obj.all()
+
+            result = []
+            for doc in res:
+                if self.model_variables['rag_tuning']:
+                    current_app.logger.debug(f'Document ID: {doc[0].id} - Distance: {doc[1]}\n')
+                    current_app.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 result
+
+    def _log_rag_tuning(self, query_obj, query_embedding):
+        """Log debug information for RAG tuning"""
+        current_app.rag_tuning_logger.debug("Debug: Query execution plan:")
+        current_app.rag_tuning_logger.debug(f"{query_obj.statement}")
+        if self._active_filters:
+            current_app.rag_tuning_logger.debug("Debug: Active metadata filters:")
+            current_app.rag_tuning_logger.debug(f"{self._active_filters}")
+
+    def _get_query_embedding(self, query: str):
+        """Get embedding for the query text"""
+        embedding_model = self.model_variables['embedding_model']
+        query_embedding = embedding_model.embed_query(query)
+        return query_embedding
+
+    @property
+    def model_variables(self) -> ModelVariables:
+        return self._model_variables
+
+    @property
+    def tenant_info(self) -> Dict[str, Any]:
+        return self._tenant_info

common/langchain/retrievers/eveai_history_retriever.py

@@ -0,0 +1,52 @@
+from langchain_core.retrievers import BaseRetriever
+from sqlalchemy import asc
+from sqlalchemy.exc import SQLAlchemyError
+from pydantic import Field, BaseModel, PrivateAttr
+from typing import Any, Dict
+from flask import current_app
+
+from common.extensions import db
+from common.models.interaction import ChatSession, Interaction
+from common.utils.model_utils import ModelVariables
+
+
+class EveAIHistoryRetriever(BaseRetriever, BaseModel):
+    _model_variables: ModelVariables = PrivateAttr()
+    _session_id: str = PrivateAttr()
+
+    def __init__(self, model_variables: ModelVariables, session_id: str):
+        super().__init__()
+        self._model_variables = model_variables
+        self._session_id = session_id
+
+    @property
+    def model_variables(self) -> ModelVariables:
+        return self._model_variables
+
+    @property
+    def session_id(self) -> str:
+        return self._session_id
+
+    def _get_relevant_documents(self, query: str):
+        current_app.logger.debug(f'Retrieving history of interactions for query: {query}')
+
+        try:
+            query_obj = (
+                db.session.query(Interaction)
+                .join(ChatSession, Interaction.chat_session_id == ChatSession.id)
+                .filter(ChatSession.session_id == self.session_id)
+                .order_by(asc(Interaction.id))
+            )
+
+            interactions = query_obj.all()
+
+            result = []
+            for interaction in interactions:
+                result.append(f'HUMAN:\n{interaction.detailed_question}\n\nAI: \n{interaction.answer}\n\n')
+
+        except SQLAlchemyError as e:
+            current_app.logger.error(f'Error retrieving history of interactions: {e}')
+            db.session.rollback()
+            return []
+
+        return result

common/langchain/retrievers/eveai_retriever.py

@@ -0,0 +1,40 @@
+from pydantic import BaseModel, PrivateAttr
+from typing import Dict, Any
+
+from common.utils.model_utils import ModelVariables
+
+
+class EveAIRetriever(BaseModel):
+    _catalog_id: int = PrivateAttr()
+    _user_metadata: Dict[str, Any] = PrivateAttr()
+    _system_metadata: Dict[str, Any] = PrivateAttr()
+    _configuration: Dict[str, Any] = PrivateAttr()
+    _tenant_info: Dict[str, Any] = PrivateAttr()
+    _model_variables: ModelVariables = PrivateAttr()
+    _tuning: bool = PrivateAttr()
+
+    def __init__(self, catalog_id: int, user_metadata: Dict[str, Any], system_metadata: Dict[str, Any],
+                 configuration: Dict[str, Any]):
+        super().__init__()
+        self._catalog_id = catalog_id
+        self._user_metadata = user_metadata
+        self._system_metadata = system_metadata
+        self._configuration = configuration
+
+    @property
+    def catalog_id(self):
+        return self._catalog_id
+
+    @property
+    def user_metadata(self):
+        return self._user_metadata
+
+    @property
+    def system_metadata(self):
+        return self._system_metadata
+
+    @property
+    def configuration(self):
+        return self._configuration
+
+    # Any common methods that should be shared among retrievers can go here.