diff --git a/main.py b/main.py new file mode 100644 index 0000000..03ee0ef --- /dev/null +++ b/main.py @@ -0,0 +1,177 @@ +""" +Main entry point for the agent with memory and human‑in‑the‑loop confirmation. + +The agent is built on top of LangGraph's `create_agent` API. It uses a +`MemorySaver` checkpoint to keep conversation history across calls, and it +is configured with `interrupt_before=["tools"]` so that the agent pauses just +before invoking any tool. The pause allows us to ask the user for explicit +confirmation. + +The example includes one simple tool – ``get_price`` – which pretends to +query a price service. In a real project this would be replaced with an +actual API call. + +Three usage examples are demonstrated in ``__main__``: +1. Ask the agent for weather information (uses the built‑in ``web_search`` + tool). +2. Ask for a product price – the agent will pause and ask for confirmation. +3. Continue the conversation to show that memory is preserved. + +The console output is rendered with `rich` for better readability. +""" + +from __future__ import annotations + +import os +import json +from typing import Any, Dict, Iterable, Tuple + +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage, SystemMessage +from langgraph.checkpoint.memory import MemorySaver +from langgraph.types import Command +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages +from rich.console import Console + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- +console = Console() + +llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1", + api_key=os.getenv("JOURNAL_MCP_PAT"), + temperature=0.0, +) + +# --------------------------------------------------------------------------- +# Simple tool – in a real scenario replace with an actual API call. +# --------------------------------------------------------------------------- +async def get_price(query: Dict[str, Any]) -> str: + """Pretend to fetch a price for a product. + + Parameters + ---------- + query: dict + Expected keys are ``product`` and optionally ``currency``. + + Returns + ------- + str + A human‑readable string describing the price. + """ + product = query.get("product", "unknown") + currency = query.get("currency", "USD") + # Dummy logic – in real life call an external service. + return f"The price of {product} is 42.00 {currency}." + +# --------------------------------------------------------------------------- +# Agent definition +# --------------------------------------------------------------------------- +memory = MemorySaver() + +agent = StateGraph( + state_schema=dict(messages=list, next=tuple) +) + +# Node that simply forwards the messages to the LLM. +async def llm_node(state: Dict[str, Any]) -> Tuple[Dict[str, Any], str]: + # The LLM expects a list of messages; we pass the current history. + response = await llm.ainvoke(state["messages"]) + return {"messages": state["messages"] + [response]}, "next" + +# Node that handles tool calls – for this example we only have get_price. +async def tool_node(state: Dict[str, Any]) -> Tuple[Dict[str, Any], str]: + # The last message should contain a tool call. + last_msg = state["messages"][-1] + if not hasattr(last_msg, "tool_calls") or not last_msg.tool_calls: + return state, "next" + + tool_call = last_msg.tool_calls[0] + name = tool_call.name + args = json.loads(tool_call.args) + if name == "get_price": + result = await get_price(args) + # Append the tool output as a new message. + state["messages"].append( + HumanMessage(content=f"Tool {name} returned: {result}") + ) + return state, "next" + +agent.add_node("llm", llm_node) +agent.add_node("tool", tool_node) +agent.set_entry_point("llm") +agent.add_edge("llm", "tool") +agent.add_edge("tool", "llm") + +# Compile the graph with a MemorySaver checkpoint. +graph = agent.compile(checkpointer=memory, interrupt_before=["tools"]) + +# --------------------------------------------------------------------------- +# Helper to run the agent with human‑in‑the‑loop confirmation. +# --------------------------------------------------------------------------- +async def ask_and_run(user_input: Dict[str, Any], config: Dict[str, Any]): + """Run the agent and pause before each tool call. + + Parameters + ---------- + user_input: dict + Dictionary with a ``messages`` key containing a list of messages. + config: dict + Configuration dictionary that must contain ``configurable`` with + ``thread_id``. + """ + async for chunk in graph.stream(user_input, config=config, stream_mode=["messages", "updates"]): + # ``chunk`` is a tuple (type, data). + chunk_type, chunk_data = chunk + state = graph.get_state(config) + + if chunk_type == "messages": + # Stream token by token. + console.print(chunk_data.content, end="", style="cyan") + console.file.flush() + elif chunk_type == "updates": + # Tool call preview – show the user what will be executed. + console.print("\n[bold magenta]Agent wants to call a tool:[/]") + console.print(json.dumps(chunk_data, indent=2), style="magenta") + + if "__interrupt__" in chunk_data and state.next == ("tools",): + # Pause – ask for confirmation. + console.print("\n[bold yellow]Confirmation required:[/] Do you allow the tool call? (y/n)") + answer = input().strip().lower() + if answer != "y": + console.print("[red]Action cancelled by user.[/]") + break + # Resume from the same state. + await graph.ainvoke(Command(resume=None), config=config) + +# --------------------------------------------------------------------------- +# Main loop – three examples as requested. +# --------------------------------------------------------------------------- +if __name__ == "__main__": + thread_id = "demo-thread" + config = {"configurable": {"thread_id": thread_id}} + + console.print("[bold green]Welcome to the agent demo![/]") + console.print("Type 'exit' to quit.") + + # Example 1 – simple chat (no tool call). + console.print("\n[underline]Example 1: Simple question[/]") + user_msg = {"messages": [HumanMessage(content="What is the capital of France?")]} + import asyncio + asyncio.run(ask_and_run(user_msg, config)) + + # Example 2 – tool call with confirmation. + console.print("\n[underline]Example 2: Tool call (price query)[/]") + user_msg = {"messages": [HumanMessage(content="Get price of laptop in USD")]} + asyncio.run(ask_and_run(user_msg, config)) + + # Example 3 – continue conversation to show memory. + console.print("\n[underline]Example 3: Continue conversation[/]") + user_msg = {"messages": [HumanMessage(content="What about the price in EUR?")]} + asyncio.run(ask_and_run(user_msg, config)) + + console.print("\n[bold green]Demo finished.[/]" +)