from langchain_openai import ChatOpenAI
from langchain.tools import tool, BaseTool
from langchain.agents import create_agent
from pydantic import SecretStr
import json

# Connect to local LLM via OpenAI-compatible API
llm = ChatOpenAI(
    model="gpt-4o-mini",  # replace with your LM Studio model name
    base_url="http://localhost:1234/v1",
    api_key=SecretStr("fake"),
    temperature=0.7,
)

# Sub-agent that generates a price table for a product in a city
@tool
def get_price(product: str, city: str) -> str:
    """Return a realistic price table for the given product and city."""
    # Create a sub‑agent with a simple prompt to generate a table
    sub_agent = create_agent(
        model=llm,
        tools=[],
        system_prompt=f"You are a market analyst. Provide a realistic price for {product} in {city}. Return the result as a markdown table with columns: Продукт, Цена (руб.), Магазин.",
    )
    response = sub_agent.invoke({"messages": [{"role": "human", "content": f"Generate price for {product} in {city}"}]})
    # The agent returns a dict with messages; the last message contains the table
    return response["messages"][-1]["content"]

# Main agent that uses get_price to build shopping list
main_agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цены каждого продукта.",
)

# Example query
query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
result = main_agent.invoke({"messages": [{"role": "human", "content": query}]})

# Pretty‑print all messages
for msg in result["messages"]:
    if "content" in msg and msg["content"]:
        print(msg["content"])
    elif "tool_calls" in msg and msg["tool_calls"]:
        for call in msg["tool_calls"]:
            name = call["name"]
            args = json.dumps(call["args"], ensure_ascii=False)
            print(f"{name}({args})")
    else:
        print(msg)

print("\n--- End of conversation ---")