diff --git a/main.py b/main.py new file mode 100644 index 0000000..0fcfde3 --- /dev/null +++ b/main.py @@ -0,0 +1,91 @@ +import os +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage, SystemMessage +from langchain.tools import tool +from langgraph.checkpoint.memory import MemorySaver +from langgraph.types import Command +from langgraph.graph import StateGraph, START, END +from langgraph.graph.message import add_messages +from rich.console import Console + +# LLM setup +llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), + temperature=0.0, +) + +console = Console() + +# Dummy tool – can be replaced by real implementation +@tool +def get_price(query: str) -> str: + """Return a fake price for demonstration.""" + return f"Цена для {query} – 100₽" + +# State definition +class ChatState(dict): + messages: list = add_messages + location: str = "start" + inventory: list = [] + +# Agent graph +builder = StateGraph(ChatState) + +# Node that runs the LLM and returns tool calls +async def llm_node(state: ChatState): + system = SystemMessage(content="You are a helpful assistant. Use the get_price tool if needed.") + response = await llm.ainvoke([system] + state["messages"], config={"configurable": {"thread_id": "chat-1"}}) + return {"messages": [response]} + +# Node that handles tool calls (only get_price) +async def tool_node(state: ChatState): + last_msg = state["messages"][-1] + if hasattr(last_msg, "tool_calls") and last_msg.tool_calls: + tool_call = last_msg.tool_calls[0] + name = tool_call["name"] + args = tool_call["args"] + if name == "get_price": + result = get_price(args) + # Append tool result to messages + state["messages"].append(HumanMessage(content=result)) + return state + +builder.add_node("llm", llm_node) +builder.add_node("tool", tool_node) +builder.add_edge(START, "llm") +builder.add_edge("llm", "tool") +builder.add_edge("tool", END) + +chat = builder.compile(checkpointer=MemorySaver(), interrupt_before=["tool"]) + +# Main loop +while True: + user_input = console.input("\nВы: ") + if user_input.lower() in {"exit", "quit"}: + break + # Add user message + chat.add_state({"messages": [HumanMessage(content=user_input)]}) + # Run graph with interrupt handling + state = chat.get_state({"configurable": {"thread_id": "chat-1"}}) + while True: + result = await chat.ainvoke(state, {"configurable": {"thread_id": "chat-1"}}) + # If interrupted before tool + if "__interrupt__" in result and result["__interrupt__"][0]["name"] == "tool": + console.print("\nАгент хочет вызвать инструмент: get_price") + ans = console.input("Разрешить? (Y/n): ") + if ans.lower() in {"", "y", "yes"}: + # resume + state = chat.get_state({"configurable": {"thread_id": "chat-1"}}) + continue + else: + console.print("Отменено") + break + # Print assistant response + for msg in result["messages"]: + if msg.role == "assistant": + console.print(msg.content) + break + +console.print("\nДо свидания!")