diff --git a/agent.py b/agent.py index 5122237..929e290 100644 --- a/agent.py +++ b/agent.py @@ -1,59 +1,69 @@ -from langgraph.graph import StateGraph -from typing import TypedDict, List, Any +from langgraph import create_agent +from langgraph.checkpoint.memory import MemorySaver +from langchain_openai import ChatOpenAI from rich.console import Console - -# Define state schema -class AgentState(TypedDict): - conversation_history: List[dict] - last_tool_call_confirmed: bool | None +from langgraph.tools import tool +import asyncio console = Console() -# Memory node: just returns the existing history -async def memory_node(state: AgentState) -> dict: - return {"memory": state.get("conversation_history", [])} +@tool +def get_price(params: dict) -> str: + """Return a mock price for a city and date.""" + city = params.get("city", "unknown") + date = params.get("date", "unknown") + return f"Цена в {city} на {date}: 1000₽" -# Confirm tool call node -async def confirm_tool_node(state: AgentState, tool_call: Any) -> dict: +llm = ChatOpenAI( + base_url="http://localhost:11434/v1", + api_key="ollama", + model="llama3", +) + +memory = MemorySaver() + +agent = create_agent( + model=llm, + tools=[get_price], + system_prompt="You are a helpful assistant.", + checkpointer=memory, + interrupt_before=["tools"], +) + +config = {"configurable": {"thread_id": "conversation-1"}} + +async def ask_and_run(user_input, config): try: - console.print(f"[bold cyan]Tool call:[/bold cyan] {tool_call}") - confirmation = input("Confirm execution? (y/n): ") - confirmed = confirmation.lower().startswith("y") - state["last_tool_call_confirmed"] = confirmed + async for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]): + chunk_type, chunk_data = chunk + if chunk_type == "messages": + console.print(chunk_data, end="") + elif chunk_type == "updates": + if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",): + state = agent.get_state(config) + tool_calls = state.values.get("messages", [])[-1].tool_calls + if tool_calls: + last_call = tool_calls[0] + console.print(f"\n\nTool call detected: {last_call['name']}{last_call['args']}") + answer = input("Разрешить? (Y/n): ") + if answer.lower().strip() == "y": + await ask_and_run(None, config) + else: + console.print("Отменено") + break except Exception as e: - console.print(f"[red]Error during confirmation:[/red] {e}") - state["last_tool_call_confirmed"] = False - return {"confirmed_tool_call": state["last_tool_call_confirmed"]} + console.print(f"[red]Ошибка во время выполнения:[/red] {e}") + raise -# Agent node: placeholder for actual LangChain agent logic -async def agent_node(state: AgentState, memory: Any, confirmed_tool_call: Any) -> dict: - # In a real implementation we would invoke the agent here. - # For demonstration, echo back the conversation history and confirmation flag. - response = { - "memory": memory, - "confirmed": confirmed_tool_call, - "message": "Agent executed with memory and confirmation." - } - return {"response": response} - -# Build graph following execution flow: MemoryNode -> AgentNode -> ConfirmToolNode -> AgentNode -builder = StateGraph(AgentState) -builder.add_node("MemoryNode", memory_node) -builder.add_node("AgentNode", agent_node) -builder.add_node("ConfirmToolNode", confirm_tool_node) - -# Define transitions according to plan execution flow -builder.set_entry_point("MemoryNode") -builder.add_edge("MemoryNode", "AgentNode") -builder.add_edge("AgentNode", "ConfirmToolNode") -builder.add_edge("ConfirmToolNode", "AgentNode") -builder.set_finish_node("AgentNode") - -# Compile graph -graph = builder.compile() +def main(): + console.print("Добро пожаловать! Введите 'exit' для выхода.") + while True: + user_input = input("\nВы: ") + if user_input.lower() == "exit": + break + msg = {"messages": [{"role": "human", "content": user_input}]} + asyncio.run(ask_and_run(msg, config)) + console.print("\n---") if __name__ == "__main__": - # Example initial state - init_state: AgentState = {"conversation_history": [], "last_tool_call_confirmed": None} - result = graph.invoke(init_state) - console.print(result) \ No newline at end of file + main()