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

# Конфигурация локальной модели LM Studio
llm = ChatOpenAI(
    model="gpt-4o-mini",
    base_url="http://localhost:1234/v1",
    api_key=SecretStr("fake"),
    temperature=0.7,
)

@tool
def get_price(product: str, city: str) -> str:
    """Получить примерную цену продукта в указанном городе."""
    # Создаём субагента, который генерирует таблицу цен
    sub_agent = create_agent(
        model=llm,
        tools=[],
        system_prompt=(
            f"Ты генератор цены для продукта {product} в городе {city}. "
            "Сгенерируй таблицу в формате:\n"
            "| Продукт | Цена (руб.) | Магазин |\n"
            "|---------|-------------|---------|\n"
            f"Например:\n"
            f"| {product} | 100 | Магнит |\n"
            "Ответ должен содержать только таблицу без лишних слов."
        ),
    )
    # Запускаем субагента
    result = sub_agent.invoke(
        {"messages": [{"role": "user", "content": f"Сгенерируй цену для {product} в {city}"}]}
    )
    # Извлекаем текст ответа
    content = result["messages"][-1]["content"]
    return content

# Главный агент
main_agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты помощник по планированию покупок.",
)

def format_message(message: dict) -> str:
    """Форматируем сообщение для вывода."""
    if "content" in message and message["content"]:
        return message["content"]
    if "tool_calls" in message and message["tool_calls"]:
        calls = []
        for call in message["tool_calls"]:
            name = call["name"]
            args_str = call.get("arguments", "{}")
            try:
                args = json.loads(args_str)
            except json.JSONDecodeError:
                args = args_str
            calls.append(f"{name}({args})")
        return " | ".join(calls)
    return ""

def main():
    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    result = main_agent.invoke(
        {"messages": [{"role": "human", "content": user_query}]}
    )
    for msg in result["messages"]:
        print(format_message(msg))
        print("---")

if __name__ == "__main__":
    main()