agent.py

from typing import Any, List

from langchain_core.prompts import PromptTemplate
from langchain.agents import AgentExecutor, create_tool_calling_agent,create_structured_chat_agent
from langchain.tools import BaseTool
from langgraph.prebuilt import create_react_agent

class Agent:
    def __init__(self, llm, tools: List[BaseTool],prompt: PromptTemplate = None, verbose: bool = False):
        self.llm = llm
        self.tools = tools
        self.prompt = prompt

        if not prompt:
            agent = create_react_agent(llm, tools,debug=verbose)
            self.agent_executor = agent
        else:
            agent = create_structured_chat_agent(llm, tools, prompt)
            self.agent_executor = AgentExecutor(agent=agent,tools=tools,verbose=verbose)        

    def exec(self, prompt_args: dict = {}, stream: bool = False):
        # if stream:
        #     for step in self.agent_executor.stream(prompt_args):
        #         yield step  
        return self.agent_executor.invoke(input=prompt_args)
        
    def stream(self, prompt_args: dict = {}):
        for step in self.agent_executor.stream(prompt_args):
            yield step