agent.py

from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from tools import extract_text
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_openai import ChatOpenAI
from typing import TypedDict,Annotated


class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    input_file: Optional[str]  # Contains file path (PDF/PNG)


class BasicAgent():
    def __init__(self, llm):
        chat = ChatHuggingFace(llm=llm, verbose=True)
        
        search_tool = DuckDuckGoSearchRun()
        vision_llm = ChatOpenAI(model="gpt-4o")
        self.tools = [extract_text, search_tool]
        self.chat_with_tools = chat.bind_tools(self.tools)
        self._initialize_graph()
        print("BasicAgent initialized.")

    
    def _initialize_graph(self):
        builder = StateGraph(AgentState)

        # Define nodes
        builder.add_node("assistant", self.assistant)
        builder.add_node("tools", ToolNode(self.tools))

        # Define edges
        builder.add_edge(START, "assistant")
        builder.add_conditional_edges("assistant",tools_condition)
        builder.add_edge("tools", "assistant")

        # Compile the graph
        self.agent = builder.compile()

    
    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        messages=[HumanMessage(content=question)]
        response = self.agent.invoke({"messages":messages})
        answer = response['messages'][-1].content
        print(f"Agent returning answer: {answer}")
        return answer

    
    def assistant(self, state: AgentState):
        image = state["input_file"]
        textual_description_of_tool="""
        extract_text(img_path: str) -> str:
            Extract text from an image file using a multimodal model.
        
            Args:
                img_path: A local image file path (strings).
        
            Returns:
                A single string containing the concatenated text extracted from each image.
        search_tool(query: str) -> str:
            Search the web using the DuckDuckGoSearchRun to perform a search query and return a summarized textual result.

            Args:
                query: A string representing the search query.
        
            Returns:
                A single string containing the search result or summary.
        """        
        sys_msg = SystemMessage(content=f"You are a helpful assistant. You have access to the following tools:\n{textual_description_of_tool} \n You may have access to some optional images. Currently the loaded image is: {image}. \n When providing the final answer, be as concise as possible, provide only the FINAL ANSWER. \n For example: QUESTION: What was the actual enrollment count of the clinical trial on H. pylori in acne vulgaris patients from Jan-May 2018 as listed on the NIH website? \n FINAL ANSWER: 90 \n")

        return {
            "messages": [self.chat_with_tools.invoke([sys_msg] + state["messages"])],            
            "input_file": image
        }