from langchain.agents import AgentExecutor, Agent, create_openai_tools_agent
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    print("Multiply used")
    return a * b

class HelpAgent(AgentExecutor):
    
    def __init__(self, llm, system_prompt):
        prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "human",
                    system_prompt,
                ),
                ("assistant",  "ok" ),  
                MessagesPlaceholder(variable_name="messages"),
                MessagesPlaceholder(variable_name="agent_scratchpad"),
            ]
        )
        #agent = prompt | llm
        agent = create_openai_tools_agent(llm, [multiply], prompt)
        super().__init__(agent=agent, tools=[multiply])