chore

agent.py

@@ -3,11 +3,12 @@ from typing import TypedDict, Annotated, Optional
 from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, START
-from langgraph.graph.message import add_messages
+from langgraph.graph.message import add_messages, MessagesState
 from langgraph.prebuilt import ToolNode, tools_condition
 from langchain.vectorstores import Chroma
 from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.chains import RetrievalQA
+
 from agent_tools import *
 
 load_dotenv()
@@ -35,32 +36,8 @@ sys_msg = SystemMessage(
     """
 )
 
-tools = [
-    multiply,
-    add,
-    subtract,
-    divide,
-    modulus,
-    duckduck_websearch,
-    arvix_search,
-    wiki_search,
-    visit_webpage,
-    youtube_search,
-    text_splitter,
-    read_file,
-    excel_read,
-    csv_read,
-    image_caption,
-]
-
 print("agent.py loaded")
 
-
-class AgentState(TypedDict):
-    input_file: Optional[str]
-    messages: Annotated[list[AnyMessage], add_messages]
-
-
 class CUSTOM_AGENT:
     """
         A simple deterministic agent that leverages our tools directly and avoids
@@ -70,17 +47,9 @@ class CUSTOM_AGENT:
     def __init__(self):
         self.llm = ChatOpenAI(name="gpt-4o",
                               api_key=os.getenv("OPENAI_API_KEY"))
-
-        self.tools = tools
+        self.tools = TOOLS
         self.llm_with_tools = self.llm.bind_tools(self.tools)
-        initial_state = {
-            "input_file": None,
-            "messages": [],
-        }
-        self.app = self._graph_compile()
-        self.initial_state = initial_state
         self.sys_msg = sys_msg
-        # --- Chroma vectorstore + retriever ---
         embeddings = OpenAIEmbeddings(api_key=os.getenv("OPENAI_API_KEY"))
         persist_directory = "chroma_db"
         self.vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
@@ -92,42 +61,46 @@ class CUSTOM_AGENT:
         )
 
     def _graph_compile(self):
-        builder = StateGraph(AgentState)
-        # Define nodes: these do the work
+        builder = StateGraph(MessagesState)
+        builder.add_node("retriever", self._retriever_node)
         builder.add_node("assistant", self._assistant)
         builder.add_node("tools", ToolNode(self.tools))
-        # Define edges: these determine how the control flow moves
-        builder.add_edge(START, "assistant")
+
+        builder.add_edge(START, "retriever")
+        builder.add_edge("retriever", "assistant")
+
         builder.add_conditional_edges(
             "assistant",
             tools_condition,
         )
         builder.add_edge("tools", "assistant")
-        react_graph = builder.compile()
-        return react_graph
-
-    def _assistant(self, state: AgentState):
-        last_human = next((m for m in reversed(state["messages"]) if isinstance(m, HumanMessage)), None)
-        messages = [self.sys_msg] + state["messages"]
-
-        if last_human:
-            question_text = last_human.content
-            retrieved_output = self.qa_chain.invoke({"query": question_text})
-            retrieved_docs = retrieved_output["result"]
-            docs_text = "\n\n---\n\n".join(
-                [doc.page_content for doc in retrieved_output.get("source_documents", [])]
-            )
-            context_message = SystemMessage(content=f"Context from vectorstore:\n{docs_text}")
-            messages.append(context_message)
-
-        response = self.llm_with_tools.invoke(messages)
-
-        return {
-            "messages": state["messages"] + [response],
-            "input_file": state["input_file"]
-        }
-
-    def extract_after_final_answer(self, text):
+        return builder.compile()
+
+    def _retriever_node(self, state: MessagesState):
+        """Retriever node"""
+        question = state["messages"][-1].content
+        docs = self.retriever.get_relevant_documents(question)
+
+        if docs:
+            content = "\n".join([d.page_content for d in docs])
+        else:
+            content = "No relevant documents found"
+
+        return {"messages": [HumanMessage(content=content)]}
+
+    def _assistant(self, state: MessagesState):
+        """Assistant node"""
+        if not any(isinstance(m, SystemMessage) for m in state["messages"]):
+            messages = [self.sys_msg] + state["messages"]
+        else:
+            messages = state["messages"]
+
+        llm_response = self.llm_with_tools.invoke(messages)
+
+        return {"messages": [llm_response]}
+
+    @staticmethod
+    def extract_after_final_answer(text):
         keyword = "FINAL ANSWER: "
         index = text.find(keyword)
         if index != -1:
@@ -144,17 +117,17 @@ class CUSTOM_AGENT:
         else:
             question_text = f'{question} with TASK-ID: {task_id}'
 
-        state = self.initial_state.copy()
-        state["messages"] = [HumanMessage(content=question_text)]
+        graph = self._graph_compile()
 
         max_retries = 3
         base_sleep = 1
         for attempt in range(max_retries):
             try:
-                response = self.app.invoke(state)
-                final_ans = self.extract_after_final_answer(response['messages'][-1].content)
-                time.sleep(10)  # avoid rate limit
-                return final_ans
+                messages: list[HumanMessage] = [HumanMessage(content=question_text)]
+                result = graph.invoke({"messages": messages})
+
+                final_text = result["messages"][-1].content
+                return self.extract_after_final_answer(final_text)
             except Exception as e:
                 sleep_time = base_sleep * (attempt + 1)
                 if attempt < max_retries - 1:

agent_tools.py

@@ -1,6 +1,6 @@
 import os
 import re
-from typing import List
+from typing import List, Callable, Any
 
 import pandas as pd
 import requests
@@ -284,3 +284,21 @@ def arvix_search(query: str) -> str:
         ])
     return formatted_search_docs
 
+
+TOOLS: List[Callable[..., Any]] = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
+    duckduck_websearch,
+    arvix_search,
+    wiki_search,
+    visit_webpage,
+    youtube_search,
+    text_splitter,
+    read_file,
+    excel_read,
+    csv_read,
+    image_caption,
+]

requirements.txt

@@ -26,4 +26,8 @@ langchain_google_genai
 langchain_openai
 google-genai
 openpyxl
-chromadb
+chromadb
+langchain-google-genai
+pytesseract
+matplotlib
+sentence_transformers

test.py

@@ -107,7 +107,7 @@ questions = [
 # Test
 if __name__ == "__main__":
     agent = CUSTOM_AGENT()
-    q = questions[1]
+    q = questions[0]
     print("Question:", q["question"])
     answer = agent.run(q)
     print("Answer:", answer)