Newer
Older
cortex-hub / ai-hub / app / core / rag_service.py
import logging
from typing import Literal, List
from sqlalchemy.orm import Session
from app.core.vector_store import FaissVectorStore
from app.core.llm_providers import get_llm_provider
from app.core.retrievers import Retriever
from app.db import models

# Configure logging for the service
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class RAGService:
    """
    A service class to handle all RAG-related business logic.
    This includes adding documents and processing chat requests with context retrieval.
    The retrieval logic is now handled by pluggable Retriever components.
    """
    def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever]):
        """
        Initializes the service.
        
        Args:
            vector_store (FaissVectorStore): The FAISS vector store for document vectors.
            retrievers (List[Retriever]): A list of retriever components to use for
                                          context retrieval.
        """
        self.vector_store = vector_store
        self.retrievers = retrievers

    def add_document(self, db: Session, doc_data: dict) -> int:
        """
        Adds a new document to the database and its vector to the FAISS index.
        """
        try:
            new_document = models.Document(**doc_data)
            db.add(new_document)
            db.commit()
            db.refresh(new_document)

            faiss_id = self.vector_store.add_document(new_document.text)

            vector_meta = models.VectorMetadata(
                document_id=new_document.id,
                faiss_index=faiss_id,
                embedding_model="mock_embedder"
            )
            db.add(vector_meta)
            db.commit()

            logger.info(f"Document '{new_document.title}' added successfully with ID {new_document.id}")
            return new_document.id
        except Exception as e:
            db.rollback()
            logger.error(f"Failed to add document: {e}")
            raise e

    async def chat_with_rag(
        self,
        db: Session,
        prompt: str,
        model: Literal["deepseek", "gemini"]
    ) -> str:
        """
        Handles a chat request by retrieving context from all configured
        retrievers and passing it to the LLM.
        """
        context_docs_text = []
        # The service now iterates through all configured retrievers to gather context
        for retriever in self.retrievers:
            context_docs_text.extend(retriever.retrieve_context(prompt, db))

        combined_context = "\n\n".join(context_docs_text)
        
        if combined_context:
            logger.info(f"Retrieved context for prompt: '{prompt}'")
            rag_prompt = f"""
            You are an AI assistant that answers questions based on the provided context.
            
            Context:
            {combined_context}
            
            Question:
            {prompt}
            
            If the answer is not in the context, say that you cannot answer the question based on the information provided.
            """
        else:
            rag_prompt = prompt
            logger.warning("No context found for the query.")

        provider = get_llm_provider(model)
        response_text = await provider.generate_response(rag_prompt)
        
        return response_text