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

# 1. LLM connection
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,
)

# 2. Sub‑agent that generates a price table
@tool
def get_price(product: str, city: str) -> str:
    """Return a realistic price for the product in the given city.
    The function internally creates a small agent that asks the LLM to produce a markdown table.
    """
    # Sub‑agent prompt – keep it short and deterministic
    sub_prompt = (
        f"You are a market analyst. Provide a realistic price for {product} in {city}. "
        "Return a markdown table with columns: Продукт, Цена (руб.), Магазин."
    )
    # Create the sub‑agent
    sub_agent = create_agent(
        model=llm,
        tools=[],  # no external tools needed for this simple query
        system_prompt=sub_prompt,
    )
    # Ask the sub‑agent and get its response
    result = sub_agent.invoke({"messages": [{"role": "human", "content": "Generate table"}]})
    return result["messages"][-1]["content"]

# 3. Main agent with get_price tool
main_agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты помощник по планированию покупок.",
)

# 4. Run the main agent on a sample query
query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
response = main_agent.invoke({"messages": [{"role": "human", "content": query}]})
# Pretty‑print all messages (including tool calls)
for msg in response["messages"]:
    if msg.get("content"):
        print(msg["content"])
    elif msg.get("tool_calls"):
        for call in msg["tool_calls"]:
            name = call["name"]
            args = json.dumps(call["args"], ensure_ascii=False)
            print(f"{name}({args})")

# Final answer (last message content)
print("\n---\nAnswer:\n", response["messages"][-1]["content"])