main.py - master - kander08/MCPExample

main.py
120 строк4 KB
Kander08
rewrite to chatbot
05 июн 2025, 22:51
05 июн 2025, 22:514df621a
import os
import random
import asyncio
import logging

from dotenv import load_dotenv
import aiohttp.web

from typing import Annotated, TypedDict, Sequence
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage, SystemMessage
from langchain_core.tools import Tool

from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import create_react_agent
from langgraph.types import interrupt

load_dotenv()

# === 1. Состояния ===

class DialogState(TypedDict, total=False):
    messages: list
    completed: bool


class ReactAgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]


# === 2. Инструмент погоды ===

def get_weather(city: str) -> str:
    return f"Погода в {city}: " + random.choice(["солнечная", "облачная", "дождливая"])

weather_tool = Tool(
    name="get_weather",
    func=get_weather,
    description="Получает погоду в указанном городе. Вход: строка с названием города."
)


# === 3. GigaChat LLM для ReAct-агента ===

from langchain_gigachat import GigaChat as LangChainGigaChat

llm = LangChainGigaChat(
    credentials=os.getenv("GIGACHAT_API_CREDENTIALS"),
    model=os.getenv("GIGACHAT_MODEL_NAME"),
    verify_ssl_certs=False,
    temperature=0.3,
    max_tokens=8000,
    timeout=600
)


# === 4. ReAct агент ===

react_agent_node = create_react_agent(
    model=llm,
    tools=[weather_tool],
    prompt="Ты — AI-агент, работающий в рамках Model Context Protocol (MCP)."
)

react_graph = StateGraph(ReactAgentState)
react_graph.add_node("agent", react_agent_node)
react_graph.set_entry_point("agent")
react_graph.set_finish_point("agent")
compiled_react_graph = react_graph.compile()


# === 5. Диалоговый граф ===

@node
async def prepare(_: DialogState) -> DialogState:
    return {
        "messages": [
            SystemMessage(content="Отвечай на вопросы подробно и развернуто."),
        ],
        "completed": False,
    }

@node
async def chatbot_call(state: DialogState, gigachat_client: GigaChatClient = get_gigachat()) -> DialogState:
    request = ChatRequest(messages=state["messages"])
    response = await gigachat_client.chat(request)
    return {"messages": [response.message]}

@node
async def human_call(state: DialogState, stopword: str = "stop") -> DialogState:
    user_question = interrupt({"state": "expect user input"})["question"]
    if user_question.lower() == stopword:
        return {"completed": True}
    return {"messages": [HumanMessage(content=user_question)]}

def dialog_router(state: DialogState) -> str:
    return END if state.get("completed") else "chatbot_call"

dialog_graph = StateGraph(DialogState)
dialog_graph.add_node("prepare", prepare)
dialog_graph.add_node("chatbot_call", chatbot_call)
dialog_graph.add_node("human_call", human_call)
dialog_graph.add_edge(START, "prepare")
dialog_graph.add_edge("prepare", "chatbot_call")
dialog_graph.add_edge("chatbot_call", "human_call")
dialog_graph.add_conditional_edges("human_call", dialog_router)
compiled_dialog_graph = dialog_graph.compile()


# === 6. Веб-приложение aiohttp ===

def setup(app: aiohttp.web.Application) -> None:
    app.add_routes([
        aiohttp.web.post("/predict", GraphRouter(compiled_dialog_graph)),     # для поэтапного диалога
        aiohttp.web.post("/react_agent", GraphRouter(compiled_react_graph)), # для ReAct-агента
    ])

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    asyncio.run(run(setup=setup))