from typing import Dict, TypedDict
from langgraph.graph import StateGraph

# Определение состояния
class AgentState(TypedDict):
    numbers: list[int]
    operation: str
    result: int

# Узел обработки
def process_node(state: AgentState) -> AgentState:
    operation = state["operation"]
    numbers = state["numbers"]
    if operation == "*":
        rez_value = 1
        for i in numbers:
            rez_value *= i
    elif operation == "+":
        rez_value = sum(numbers)
    else:
        rez_value = 0
    state["result"] = rez_value
    return state

# Создание графа
graph = StateGraph(AgentState)
graph.add_node("process", process_node)
graph.set_entry_point("process")
graph.set_finish_point("process")
app = graph.compile()

# Визуализация графа
with open("graph.png", "wb") as f:
    f.write(app.get_graph().draw_mermaid_png())

# Запуск графа
result = app.invoke({"numbers": [1, 2, 3, 10], "operation": "+"})
print(result["result"])  # 👉 16

# Шаг 1: LLM понимает, что нужно 2 разные задачи
AIMessage: "Я выполню оба ваших запроса - сложу числа и получу информацию о погоде"

# Шаг 2: Автоматически вызывает СРАЗУ 2 инструмента
tool_calls: [
    {'name': 'add_tool', 'args': {'a': 100, 'b': 123432134}},
    {'name': 'get_weather_tool', 'args': {'city': 'Краснодар'}}
]

@tool
def analyze_image(image_path: str) -> str:
    """Анализирует изображение с помощью ИИ"""

@tool  
def send_email(to: str, subject: str, body: str) -> str:
    """Отправляет email"""

@tool
def query_database(sql: str) -> str:
    """Выполняет SQL запрос"""

# Обычный запрос к LLM
response = model.invoke("Расскажи о Франции в JSON")
# Получаем: "Конечно! Вот информация о Франции: {"capital": "Париж"...} Надеюсь, это помогло!"

from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from pydantic import BaseModel, Field

# 1️⃣ Определяем СТРОГУЮ структуру
class CountryInfo(BaseModel):
    capital: str = Field(description="Столица страны")
    population: int = Field(description="Население")
    country_name: str = Field(description="Название страны")

# 2️⃣ Создаем умный парсер
parser = JsonOutputParser(pydantic_object=CountryInfo)

# 3️⃣ Магия format_instructions!
prompt = PromptTemplate(
    template="""Информация о стране в JSON:
    
Страна: {country}
{format_instructions}

ТОЛЬКО JSON, никакого лишнего текста!""",
    input_variables=["country"],
    partial_variables={"format_instructions": parser.get_format_instructions()}
)

chain = prompt | model | parser
result = chain.invoke({"country": "Египет"})

# Получаем чистый dict:
{
    "capital": "Каир",
    "population": 104000000,
    "country_name": "Египет"
}

The output should be formatted as a JSON instance that conforms to the JSON schema below.

{"properties": {"capital": {"title": "Capital", "type": "string"}, ...}}

def debug_chain(input_data):
    step1 = prompt.invoke(input_data)    # Смотрим промпт
    step2 = model.invoke(step1)          # Смотрим ответ модели  
    step3 = parser.invoke(step2)         # Смотрим парсинг
    return step3