updated files

agents/research_agent.py

@@ -1,52 +1,14 @@
-from typing import Dict, List
 from langchain.schema import Document
 from config.llm_config import llm_config
-import logging
-
-logger = logging.getLogger(__name__)
 
 class ResearchAgent:
     def __init__(self):
-        """
-        Initialize the research agent with LLM from LLMConfig
-        """
-        logger.info("Initializing ResearchAgent...")
-
-        # Get LLM callable for research
-        self.llm_fn = llm_config.create_llm("research")
-
-        logger.info("ResearchAgent initialized successfully.")
-
-    def generate(self, question: str, documents: List[Document]) -> Dict:
-        """
-        Generate an answer based on documents
-        """
-        logger.info(f"ResearchAgent.generate called with question='{question}' and {len(documents)} documents.")
-
-        context = "\n\n".join([doc.page_content for doc in documents])
-
-        prompt = f"""
-You are an AI assistant designed to provide precise and factual answers based on the given context.
-
-Question: {question}
-
-Context:
-{context}
-
-Provide your answer below:
-"""
-
-        try:
-            draft_answer = self.llm_fn(prompt)
-            logger.info(f"Generated answer successfully. Length: {len(draft_answer)} characters.")
-            return {
-                "draft_answer": draft_answer.strip(),
-                "context_used": context
-            }
-
-        except Exception as e:
-            logger.error(f"Error during answer generation: {e}")
-            return {
-                "draft_answer": f"I cannot answer this question based on the provided documents. Error: {str(e)}",
-                "context_used": context
-            }
+        self.llm = llm_config.create_llm("research")
+
+    def generate(self, question: str, documents: list[Document]):
+        context = "\n\n".join([d.page_content for d in documents])
+        prompt = f"Question: {question}\n\nContext: {context}\n\nAnswer:"
+        
+        # Proper LangChain invocation
+        response = self.llm.invoke(prompt)
+        return {"draft_answer": response.content}

agents/workflow.py

@@ -1,11 +1,10 @@
-from typing import TypedDict, List, Dict
+from typing import TypedDict, List, Dict, Annotated
 from langchain.schema import Document
-from langchain.retrievers import EnsembleRetriever
+from langgraph.graph import StateGraph, END
+import logging
 from .research_agent import ResearchAgent
 from .verification_agent import VerificationAgent
 from .relevance_checker import RelevanceChecker
-from langgraph.graph import StateGraph, END
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -15,89 +14,67 @@ class AgentState(TypedDict):
     draft_answer: str
     verification_report: str
     is_relevant: bool
-    retriever: EnsembleRetriever
+    retry_count: int  # Added to prevent infinite loops
 
 class AgentWorkflow:
     def __init__(self):
         self.researcher = ResearchAgent()
         self.verifier = VerificationAgent()
         self.relevance_checker = RelevanceChecker()
-        self.compiled_workflow = self.build_workflow()
+        self.workflow = self.build_workflow()
 
     def build_workflow(self):
-        workflow = StateGraph(AgentState)
-        workflow.add_node("check_relevance", self._check_relevance_step)
-        workflow.add_node("research", self._research_step)
-        workflow.add_node("verify", self._verification_step)
+        builder = StateGraph(AgentState)
+        
+        builder.add_node("check_relevance", self._check_relevance_step)
+        builder.add_node("research", self._research_step)
+        builder.add_node("verify", self._verification_step)
 
-        workflow.set_entry_point("check_relevance")
-        workflow.add_conditional_edges(
+        builder.set_entry_point("check_relevance")
+        
+        builder.add_conditional_edges(
             "check_relevance",
-            self._decide_after_relevance_check,
-            {"relevant": "research", "irrelevant": END}
+            lambda x: "research" if x["is_relevant"] else "end",
+            {"research": "research", "end": END}
         )
-        workflow.add_edge("research", "verify")
-        workflow.add_conditional_edges(
+        
+        builder.add_edge("research", "verify")
+        
+        builder.add_conditional_edges(
             "verify",
             self._decide_next_step,
             {"re_research": "research", "end": END}
         )
-        return workflow.compile()
-
-    def _check_relevance_step(self, state: AgentState) -> Dict:
-        classification = self.relevance_checker.check(
-            question=state["question"],
-            retriever=state["retriever"],
-            k=20
-        )
-
-        if classification in ["CAN_ANSWER", "PARTIAL"]:
-            return {"is_relevant": True}
-        return {
-            "is_relevant": False,
-            "draft_answer": "This question isn't related to the uploaded document(s)."
-        }
-
-    def _decide_after_relevance_check(self, state: AgentState) -> str:
-        return "relevant" if state["is_relevant"] else "irrelevant"
-
-    def full_pipeline(self, question: str, retriever: EnsembleRetriever):
-        try:
-            documents = retriever.get_relevant_documents(question)  # updated method
-            logger.info(f"Retrieved {len(documents)} documents")
+        
+        return builder.compile()
 
-            initial_state = AgentState(
-                question=question,
-                documents=documents,
-                draft_answer="",
-                verification_report="",
-                is_relevant=False,
-                retriever=retriever
-            )
+    def _check_relevance_step(self, state: AgentState):
+        # Logic to call relevance_checker.check
+        res = self.relevance_checker.check(state["question"], state["documents"])
+        return {"is_relevant": res != "NO_MATCH", "retry_count": 0}
 
-            final_state = self.compiled_workflow.invoke(initial_state)
-            return {
-                "draft_answer": final_state["draft_answer"],
-                "verification_report": final_state["verification_report"]
-            }
+    def _research_step(self, state: AgentState):
+        res = self.researcher.generate(state["question"], state["documents"])
+        return {"draft_answer": res["draft_answer"], "retry_count": state.get("retry_count", 0) + 1}
 
-        except Exception as e:
-            logger.error(f"Workflow execution failed: {e}")
-            return {
-                "draft_answer": "",
-                "verification_report": f"Error: {e}"
-            }
+    def _verification_step(self, state: AgentState):
+        res = self.verifier.check(state["draft_answer"], state["documents"])
+        return {"verification_report": res["verification_report"]}
 
-    def _research_step(self, state: AgentState) -> Dict:
-        result = self.researcher.generate(state["question"], state["documents"])
-        return {"draft_answer": result["draft_answer"]}
+    def _decide_next_step(self, state: AgentState):
+        # Break loop after 2 retries or if supported
+        if "Supported: YES" in state["verification_report"] or state["retry_count"] >= 3:
+            return "end"
+        return "re_research"
 
-    def _verification_step(self, state: AgentState) -> Dict:
-        result = self.verifier.check(state["draft_answer"], state["documents"])
-        return {"verification_report": result["verification_report"]}
-
-    def _decide_next_step(self, state: AgentState) -> str:
-        report = state["verification_report"]
-        if "NO" in report:
-            return "re_research"
-        return "end"
+    def full_pipeline(self, question: str, retriever):
+        docs = retriever.invoke(question) # Updated from get_relevant_documents
+        initial_state = {
+            "question": question,
+            "documents": docs,
+            "draft_answer": "",
+            "verification_report": "",
+            "is_relevant": False,
+            "retry_count": 0
+        }
+        return self.workflow.invoke(initial_state)

config/llm_config.py

@@ -1,113 +1,44 @@
-"""
-LLM Configuration Manager
-Centralizes all LLM model configurations for easy switching
-"""
-from typing import Dict, Any
-from enum import Enum
 import os
 import logging
-
-# Modern Google Generative AI package import
 import google.generativeai as genai
-
-# Environment variables
+from typing import List, Dict, Any
+from langchain_core.embeddings import Embeddings
+from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
 from dotenv import load_dotenv
-load_dotenv()
 
+load_dotenv()
 logger = logging.getLogger(__name__)
 
-class ModelProvider(Enum):
-    GOOGLE = "google"
-    OPENAI = "openai"
-
 class LLMConfig:
-    """Configuration manager for LLM models"""
-    
-    MODELS = {
-        ModelProvider.GOOGLE: {
+    def __init__(self):
+        self.google_api_key = os.getenv("GOOGLE_API_KEY")
+        if not self.google_api_key:
+            raise ValueError("GOOGLE_API_KEY not found")
+        genai.configure(api_key=self.google_api_key)
+
+    def create_llm(self, task: str):
+        # Map tasks to models
+        model_map = {
             "research": "gemini-1.5-pro",
             "verification": "gemini-1.5-flash",
-            "relevance": "gemini-1.5-flash",
-            "embedding": "text-embedding-004",
-        },
-        ModelProvider.OPENAI: {
-            "research": "gpt-4-turbo",
-            "verification": "gpt-4-turbo",
-            "relevance": "gpt-4-turbo",
-            "embedding": "text-embedding-3-large",
+            "relevance": "gemini-1.5-flash"
         }
-    }
-    
-    DEFAULT_PARAMS = {
-        "research": {"temperature": 0.3, "max_tokens": 300, "top_p": 0.95},
-        "verification": {"temperature": 0.0, "max_tokens": 200, "top_p": 0.9},
-        "relevance": {"temperature": 0.0, "max_tokens": 10, "top_p": 0.9}
-    }
-    
-    def __init__(self, provider: ModelProvider = ModelProvider.GOOGLE):
-        self.provider = provider
-        self.api_key = self._get_api_key()
-        self._validate_config()
-        genai.api_key = self.api_key
-
-    def _get_api_key(self) -> str:
-        if self.provider == ModelProvider.GOOGLE:
-            key = os.getenv("GOOGLE_API_KEY")
-            if not key:
-                raise ValueError("GOOGLE_API_KEY environment variable is required")
-            return key
-        elif self.provider == ModelProvider.OPENAI:
-            key = os.getenv("OPENAI_API_KEY")
-            if not key:
-                raise ValueError("OPENAI_API_KEY environment variable is required")
-            return key
-        raise ValueError(f"Unsupported provider: {self.provider}")
-    
-    def _validate_config(self):
-        if self.provider not in self.MODELS:
-            raise ValueError(f"Provider {self.provider} not configured")
-    
-    def get_model_name(self, task: str) -> str:
-        return self.MODELS[self.provider][task]
-    
-    def get_model_params(self, task: str) -> Dict[str, Any]:
-        return self.DEFAULT_PARAMS.get(task, {}).copy()
-    
-    def create_llm(self, task: str):
-        """Return a callable that sends prompt to LLM and returns the response text."""
-        model_name = self.get_model_name(task)
-        params = self.get_model_params(task)
+        params = {
+            "research": {"temperature": 0.3},
+            "verification": {"temperature": 0.0},
+            "relevance": {"temperature": 0.0}
+        }
+        
+        return ChatGoogleGenerativeAI(
+            model=model_map[task],
+            google_api_key=self.google_api_key,
+            temperature=params[task]["temperature"]
+        )
 
-        def llm_callable(prompt: str) -> str:
-            try:
-                response = genai.chat.create(
-                    model=model_name,
-                    messages=[{"role": "user", "content": prompt}],
-                    temperature=params.get("temperature", 0.3),
-                    max_output_tokens=params.get("max_tokens", 300),
-                    top_p=params.get("top_p", 0.95),
-                )
-                return response.last.split("\n")[0] if hasattr(response, 'last') else response.choices[0].content
-            except Exception as e:
-                logger.error(f"LLM call failed: {e}")
-                raise
-        return llm_callable
-    
     def create_embedding(self):
-        """Return a callable for embeddings"""
-        if self.provider == ModelProvider.GOOGLE:
-            model_name = self.get_model_name("embedding")
-            def embed_fn(texts: list[str]) -> list[list[float]]:
-                try:
-                    response = genai.embeddings.create(
-                        model=model_name,
-                        input=texts
-                    )
-                    return [item.embedding for item in response.data]
-                except Exception as e:
-                    logger.error(f"Embedding generation failed: {e}")
-                    raise
-            return embed_fn
+        return GoogleGenerativeAIEmbeddings(
+            model="models/text-embedding-004",
+            google_api_key=self.google_api_key
+        )
 
-# Global configuration instance
-llm_config = LLMConfig()
+llm_config = LLMConfig()

requirements.txt

@@ -1,3 +1,12 @@
+gradio
+langchain
+langchain-google-genai
+langgraph
+chromadb
+pydantic-settings
+python-dotenv
+pypdf
+docx2txt
 # Core Python
 python-dotenv
 pydantic

retriever/builder.py

@@ -9,47 +9,22 @@ logger = logging.getLogger(__name__)
 
 class RetrieverBuilder:
     def __init__(self):
-        """Initialize the retriever builder with embeddings."""
-        logger.info("Initializing RetrieverBuilder...")
-        
-        # Get embeddings from configuration
+        # Correctly get the LangChain Embedding Object
         self.embeddings = llm_config.create_embedding()
         
-        logger.info("RetrieverBuilder initialized successfully.")
-        
     def build_hybrid_retriever(self, docs):
-        """Build a hybrid retriever using BM25 and vector-based retrieval."""
-        try:
-            logger.info(f"Building hybrid retriever with {len(docs)} documents")
-            
-            # Create Chroma vector store
-            vector_store = Chroma.from_documents(
-                documents=docs,
-                embedding=self.embeddings,
-                persist_directory=settings.CHROMA_DB_PATH,
-                collection_name=settings.CHROMA_COLLECTION_NAME
-            )
-            logger.info("Vector store created successfully.")
-            
-            # Create BM25 retriever
-            bm25 = BM25Retriever.from_documents(docs)
-            logger.info("BM25 retriever created successfully.")
-            
-            # Create vector-based retriever
-            vector_retriever = vector_store.as_retriever(
-                search_kwargs={"k": settings.VECTOR_SEARCH_K}
-            )
-            logger.info("Vector retriever created successfully.")
-            
-            # Combine retrievers into a hybrid retriever
-            hybrid_retriever = EnsembleRetriever(
-                retrievers=[bm25, vector_retriever],
-                weights=settings.HYBRID_RETRIEVER_WEIGHTS
-            )
-            logger.info("Hybrid retriever created successfully.")
-            
-            return hybrid_retriever
-            
-        except Exception as e:
-            logger.error(f"Failed to build hybrid retriever: {e}")
-            raise
+        # Use the class-based embedding interface
+        vector_store = Chroma.from_documents(
+            documents=docs,
+            embedding=self.embeddings, 
+            persist_directory=settings.CHROMA_DB_PATH,
+            collection_name=settings.CHROMA_COLLECTION_NAME
+        )
+        
+        bm25 = BM25Retriever.from_documents(docs)
+        vector_retriever = vector_store.as_retriever(search_kwargs={"k": 5})
+        
+        return EnsembleRetriever(
+            retrievers=[bm25, vector_retriever],
+            weights=[0.4, 0.6]
+        )