diff --git a/main.py b/main.py index 49d8b3c..d252e2e 100644 --- a/main.py +++ b/main.py @@ -1,101 +1,79 @@ +# main.py +# Реализация агента с памятью и подтверждением действий + import os -import json -import asyncio from rich.console import Console -from langchain_openai import ChatOpenAI -from langchain.tools import tool -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend from langgraph.checkpoint.memory import MemorySaver +from langgraph import create_agent +from langchain_ollama import ChatOllama +from langchain.tools import Tool -# --- LLM ------------------------------------------------------------------- -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, -) - -# --- Backend -------------------------------------------------------------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -# --- Tool ------------------------------------------------------------------ -@tool -def get_weather(query: str) -> str: - """Return a fake weather report for a given city and date. - The argument is a JSON string with keys "city" and "date". - """ - try: - data = json.loads(query) - city = data.get("city", "unknown") - date = data.get("date", "unknown") - return f"Weather in {city} on {date}: sunny, 25°C." - except Exception as e: - return f"Error parsing query: {e}" - -# --- Agent ----------------------------------------------------------------- -memory = MemorySaver() - -agent = create_deep_agent( - model=llm, - tools=[get_weather], - backend=backend, - system_prompt="You are a helpful weather assistant. Use the get_weather tool when asked about weather.", - checkpointer=memory, - interrupt_before=["tools"], # pause before each tool call -) - +# Rich console for pretty output console = Console() -# --------------------------------------------------------------------------- -async def ask_and_run(user_input, config): - """Run the agent with optional user input. - If user_input is None, the agent resumes from the paused state. - """ - # Stream the agent's response - async for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]): - state = agent.get_state(config) - chunk_type, chunk_data = chunk +# ---------- Tool definition ---------- +def get_price(city: str, date: str) -> str: + """Return a mock price for a given city and date.""" + return f"Price in {city} on {date} is $100" + +price_tool = Tool( + name="get_price", + description="Retrieve the price for a specified city and date.", + func=get_price, +) + +# ---------- LLM and Agent setup ---------- +llm = ChatOllama(model="llama3") +memory = MemorySaver() +agent = create_agent( + model=llm, + tools=[price_tool], + system_prompt="You are a helpful assistant that can call the get_price tool.", + checkpointer=memory, + interrupt_before=["tools"], # pause before each tool invocation +) + +# ---------- Helper function to handle streaming and confirmation ---------- + +def ask_and_run(user_input, config): + for chunk_type, chunk_data in agent.stream( + user_input if user_input is not None else None, + config=config, + stream_mode=["messages", "updates"], + ): if chunk_type == "messages": - # chunk_data can be a dict with 'content' - if isinstance(chunk_data, dict) and "content" in chunk_data: - console.print(chunk_data["content"], end="") - else: - console.print(chunk_data) - elif chunk_type == "updates": - console.print("\n[bold]Tool call:[/bold]", chunk_data) - - # Detect pause before tool call - if "__interrupt__" in chunk_data and state.next == ("tools",): - # Retrieve the pending tool call + console.print(chunk_data["content"], end="") + console.flush() + continue + if chunk_type == "updates": + console.print(f"\n[bold cyan]Update:[/bold cyan] {chunk_data}") + continue + if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",): + state = agent.get_state(config) last_msg = state.values["messages"][-1] - tool_call = None - if hasattr(last_msg, "tool_calls") and last_msg.tool_calls: - tool_call = last_msg.tool_calls[0] - if tool_call: - tool_name = tool_call["name"] - tool_args = tool_call["args"] - console.print(f"\n[bold]Agent wants to call {tool_name}({tool_args})[/bold]") - answer = input("Разрешить? (Y/n): ") - if answer.lower().strip() == "y": - await ask_and_run(None, config) # resume - else: - console.print("[red]Отменено[/red]") - break + tool_call = last_msg.tool_calls[0] + tool_name = tool_call["name"] + tool_args = tool_call["args"] + console.print(f"\n\n[bold yellow]Agent wants to call tool:[/bold yellow] {tool_name}({tool_args})") + answer = input("Разрешить? (Y/n): ") + if answer.lower().strip() == "y": + ask_and_run(None, config) + else: + console.print("[red]Отменено[/red]") + break -# --------------------------------------------------------------------------- -async def main(): - console.print("\n[bold green]Weather Agent ready! Type 'exit' to quit.[/bold green]") - config = {"configurable": {"thread_id": "session-1"}} +# ---------- Main chat loop ---------- +if __name__ == "__main__": + thread_id = os.getenv("THREAD_ID", "chat-1") + config = {"configurable": {"thread_id": thread_id}} + console.print("[bold green]Добро пожаловать в интерактивный агент![/bold green]") + console.print("Введите 'exit' для выхода.") while True: user_input = input("\nВы: ") - if user_input.lower() == "exit": + if user_input.lower().strip() == "exit": + console.print("[bold magenta]До свидания![/bold magenta]") break - await ask_and_run({"messages": [{"role": "human", "content": user_input}]}, config) - -if __name__ == "__main__": - asyncio.run(main()) + message = {"messages": [{"role": "human", "content": user_input}]} + ask_and_run(message, config) + console.print("\n--- --- ---")