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

# 1. Подключение к модели (используем плейсхолдеры согласно инструкции)
llm = ChatOpenAI(
    model="google/gemma-4-26b-a4b",
    base_url="http://192.168.0.120:1234/v1",
    api_key=SecretStr("lm-studio"),
    temperature=0.7,
)

# 2. Инструмент с субагентом
@tool
def get_price(product: str, city: str) -> str:
    """Узнает примерную цену на продукт в указанном городе, используя исторические данные."""
    # Создаем субагента для генерации цен
    sub_agent = create_agent(
        model=llm,
        tools=[],  # У субагента нет внешних инструментов, он использует свои знания
        system_prompt=(
            f"Ты — эксперт по ценам в городе {city}. "
            "Твоя задача: придумать реалистичную цену для продукта. "
            "Верни ответ строго в формате Markdown таблицы: | Продукты | Цена (руб.) | Магазин |"
        )
    )
    
    # Вызываем субагента
    response = sub_agent.invoke({
        "messages": [{"role": "human", "content": f"Придумай цену для продукта: {product}"}]
    })
    
    # Возвращаем контент последнего сообщения субагента
    return response["messages"][-1].content

# 3. Главный агент
main_agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты помощник по планированию покупок. Твоя задача — составить список покупок, "
                   "используя инструмент get_price для каждого продукта. В конце выведи итоговую стоимость."
)

# 4. Запрос и вывод
def format_message(message) -> str:
    """Форматирует сообщения агента (текст или вызов инструмента)."""
    if hasattr(message, 'tool_calls') and message.tool_calls:
        call = message.tool_calls[0]
        args_str = ", ".join([f"{k}={v!r}" for k, v in call['args'].items()])
        return f"{call['name']}({args_str})"
    if hasattr(message, 'content') and message.content:
        return message.content
    return ""

def run_shopping_planner():
    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    
    print(f"User Query: {user_query}\n")
    print("--- Agent Trace ---")
    
    # Запуск агента
    result = main_agent.invoke({
        "messages": [{"role": "human", "content": user_query}]
    })

    # Вывод всех сообщений из истории (промежуточные шаги + финал)
    for msg in result["messages"]:
        formatted = format_message(msg)
        if formatted:
            print(formatted)

if __name__ == "__main__":
    run_shopping_planner()