Chain architecture and patterns

from langchain_groq import ChatGroq
from vectordb.langchain.components import AgenticRouter
from vectordb.langchain.utils import EmbedderHelper, RAGHelper

class AgenticRAGPipeline:
    """Agentic RAG pipeline with iterative retrieval and reflection."""

    def __init__(self, config: dict) -> None:
        self.config = config
        
        # Initialize components
        self.embedder = EmbedderHelper.create_embedder(config)
        self.llm = RAGHelper.create_llm(config)
        
        # Create router with deterministic LLM
        router_llm = ChatGroq(
            model=config.get("agentic_rag", {}).get("model", "llama-3.3-70b-versatile"),
            temperature=0.0,  # Deterministic routing
        )
        self.router = AgenticRouter(router_llm)
        
        self.max_iterations = config.get("agentic_rag", {}).get("max_iterations", 3)

    def run(self, query: str) -> dict:
        """Execute agentic RAG pipeline.
        
        Args:
            query: User's question.
        
        Returns:
            Dictionary with final answer and iteration history.
        """
        documents = []
        current_answer = None
        iteration = 1
        history = []

        while iteration <= self.max_iterations:
            # Route to next action
            decision = self.router.route(
                query=query,
                has_documents=len(documents) > 0,
                current_answer=current_answer,
                iteration=iteration,
                max_iterations=self.max_iterations,
            )

            action = decision["action"]
            reasoning = decision["reasoning"]
            history.append({"iteration": iteration, "action": action, "reasoning": reasoning})

            if action == "search":
                # Retrieve documents from vector store
                new_docs = self.retrieve(query)
                documents.extend(new_docs)
                
                # Generate draft answer
                if self.llm:
                    current_answer = RAGHelper.generate(self.llm, query, documents)

            elif action == "reflect":
                # Evaluate and improve current answer
                if self.llm and current_answer:
                    current_answer = self.reflect_and_improve(
                        query, current_answer, documents
                    )

            elif action == "generate":
                # Finalize and return answer
                break

            iteration += 1

        # Generate final answer if not already done
        if not current_answer and self.llm and documents:
            current_answer = RAGHelper.generate(self.llm, query, documents)

        return {
            "query": query,
            "answer": current_answer or "No answer generated",
            "documents": documents,
            "iterations": iteration,
            "history": history,
        }

    def retrieve(self, query: str) -> list:
        """Retrieve documents from vector store."""
        query_embedding = EmbedderHelper.embed_query(self.embedder, query)
        # Vector store retrieval logic
        return []

    def reflect_and_improve(self, query: str, answer: str, documents: list) -> str:
        """Reflect on answer quality and improve."""
        reflection_prompt = f"""Given the question and current answer, identify gaps or inaccuracies.

Question: {query}
Current Answer: {answer}

What is missing or incorrect?"""
        
        reflection = self.llm.invoke(reflection_prompt).content
        
        improvement_prompt = f"""Improve the answer based on the reflection.

Question: {query}
Current Answer: {answer}
Reflection: {reflection}

Improved Answer:"""
        
        improved = self.llm.invoke(improvement_prompt).content
        return improved