app/core/pipeline.py

from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from app.core.retriever import hybrid_retriever
from app.core.reranker import document_reranker
from app.core.generator import llm_generator
from app.core.cache import semantic_cache
from app.core.memory import conversation_memory
from app.utils.prompts import get_rag_template, get_conversation_template, get_system_prompt
from app.utils.logger import logger
from app.config import config
from typing import AsyncIterator


class RAGPipeline:
    def __init__(self):
        self.retriever = hybrid_retriever
        self.reranker = document_reranker
        self.generator = llm_generator
        self.cache = semantic_cache
        self.memory = conversation_memory
        self.use_cache = config["rag"]["cache"]["enabled"]
        self.use_reranking = config["rag"]["retrieval"]["rerank"]
        self.top_k = config["rag"]["retrieval"]["top_k"]
        
        logger.info("RAG Pipeline initialized")
    
    def _format_context(self, documents: list) -> str:
        """Format retrieved documents into context string."""
        context_parts = []
        for i, doc in enumerate(documents, 1):
            context_parts.append(f"[{i}] {doc.page_content}")
        return "\n\n".join(context_parts)
    
    async def _retrieve_and_rerank(self, query: str) -> list:
        """Retrieve and optionally rerank documents."""
        # Retrieve documents
        documents = await self.retriever.ainvoke(query)
        
        if not documents:
            logger.warning("No documents retrieved")
            return []
        
        logger.info(f"Retrieved {len(documents)} documents")
        
        # Rerank if enabled
        if self.use_reranking:
            documents = self.reranker.rerank(query, documents, top_k=self.top_k)
            logger.info(f"Reranked to top {len(documents)} documents")
        
        return documents[:self.top_k]
    
    async def generate(
        self,
        query: str,
        session_id: str = None,
        use_context: bool = True
    ) -> str:
        """Generate response for query with optional RAG context."""
        
        # Check cache first
        if self.use_cache:
            cached_response = await self.cache.get(query, use_context)
            if cached_response:
                logger.info("Cache hit")
                return cached_response
        
        # Get conversation history if session provided
        history = []
        if session_id:
            history = self.memory.get_messages(session_id)
        
        # Retrieve and rerank documents if context needed
        context = ""
        if use_context:
            documents = await self._retrieve_and_rerank(query)
            if documents:
                context = self._format_context(documents)
        
        # Build prompt
        if context:
            template = get_rag_template()
            prompt = template.format(context=context, question=query)
        else:
            template = get_conversation_template()
            prompt = template.format(question=query)
        
        # Generate response
        system_prompt = get_system_prompt()
        response = await self.generator.agenerate(prompt, system_prompt)
        
        # Save to memory if session provided
        if session_id:
            self.memory.add_message(session_id, "user", query)
            self.memory.add_message(session_id, "assistant", response)
        
        # Cache the response
        if self.use_cache:
            await self.cache.set(query, response, use_context)
        
        logger.info("Response generated successfully")
        return response
    
    async def stream(
        self,
        query: str,
        session_id: str = None,
        use_context: bool = True
    ) -> AsyncIterator[str]:
        """Stream response for query with optional RAG context."""
        
        # Check cache first
        if self.use_cache:
            cached_response = await self.cache.get(query, use_context)
            if cached_response:
                logger.info("Cache hit - streaming cached response")
                yield cached_response
                return
        
        # Get conversation history if session provided
        history = []
        if session_id:
            history = self.memory.get_messages(session_id)
        
        # Retrieve and rerank documents if context needed
        context = ""
        if use_context:
            documents = await self._retrieve_and_rerank(query)
            if documents:
                context = self._format_context(documents)
        
        # Build prompt
        if context:
            template = get_rag_template()
            prompt = template.format(context=context, question=query)
        else:
            template = get_conversation_template()
            prompt = template.format(question=query)
        
        # Stream response
        system_prompt = get_system_prompt()
        full_response = ""
        
        async for chunk in self.generator.stream(prompt, system_prompt):
            full_response += chunk
            yield chunk
        
        # Save to memory if session provided
        if session_id:
            self.memory.add_message(session_id, "user", query)
            self.memory.add_message(session_id, "assistant", full_response)
        
        # Cache the full response
        if self.use_cache:
            await self.cache.set(query, full_response, use_context)
        
        logger.info("Response streamed successfully")


rag_pipeline = RAGPipeline()