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

# Configure the local LLM
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,
)

@tool
def get_price(product: str, city: str) -> str:
    """Get realistic price for a product in a city. Returns a markdown table."""
    # Create a sub-agent that generates a price table
    sub_agent = create_agent(
        model=llm,
        tools=[],
        system_prompt=f"You are a price estimator. Provide a realistic price for {product} in {city}. Return a markdown table with columns: Product, Price (rub.), Store.",
    )
    # Invoke the sub-agent
    result = sub_agent.invoke(
        {
            "messages": [
                {"role": "user", "content": f"Provide price for {product} in {city}."}
            ]
        }
    )
    # Extract the assistant message content
    messages = result.get("messages", [])
    for msg in messages:
        if msg.get("role") == "assistant" and msg.get("content"):
            return msg["content"]
    return "No price data available."

def main():
    # Main agent that uses the get_price tool
    agent = create_agent(
        model=llm,
        tools=[get_price],
        system_prompt="You are a shopping list planner. Use the get_price tool to find prices for items.",
    )
    # Sample user query
    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    # Invoke the agent
    result = agent.invoke(
        {
            "messages": [
                {"role": "human", "content": user_query}
            ]
        }
    )
    # Print all messages, including tool calls and final answer
    for msg in result.get("messages", []):
        role = msg.get("role")
        content = msg.get("content")
        tool_calls = msg.get("tool_calls")
        if content:
            print(f"{role}: {content}")
        elif tool_calls:
            for call in tool_calls:
                print(f"{role} calls {call.get('name')} with args {call.get('arguments')}")
        else:
            print(f"{role}: (no content)")

if __name__ == "__main__":
    main()