from langchain_openai import ChatOpenAI
from langchain.tools import tool
from langchain.agents import create_agent, AgentExecutor
import json

# Configure LLM
llm = ChatOpenAI(
    model="gpt-4o-mini",  # replace with your local model name
    base_url="http://localhost:1234/v1",
    api_key="fake",
    temperature=0.7,
)

@tool("Get price for a product in a city")
def get_price(product: str, city: str) -> str:
    """
    Returns a table with product, price and store.
    The function internally creates a sub-agent that generates realistic prices.
    """
    # Sub‑agent to generate price
    sub_llm = ChatOpenAI(
        model="gpt-4o-mini",
        base_url="http://localhost:1234/v1",
        api_key="fake",
        temperature=0.5,
    )
    sub_agent = create_agent(
        model=sub_llm,
        tools=[],
        system_prompt=f"You are a price estimator for {city}. Provide a realistic price and store name for {product} in a table format.",
    )
    result = sub_agent.invoke({"messages": [{"role": "human", "content": f"Give me the price of {product} in {city}"}]})
    return json.dumps(result["messages"][-1]["content"], ensure_ascii=False)

# Main agent with get_price tool
main_agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="You are a shopping assistant. Use the get_price tool to help users plan their purchases.",
)

if __name__ == "__main__":
    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    response = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
    # Print all messages
    for msg in response["messages"]:
        if "content" in msg:
            print(msg["content"])
        elif "tool_calls" in msg:
            for call in msg["tool_calls"]:
                print(f"{call['name']}({json.dumps(call['args'])})")
    # Final answer
    final = response["messages"][-1]["content"] if "content" in response["messages"][-1] else ""
    print("\nFinal answer:\n", final)