diff --git a/main.py b/main.py index 4aa4081..fe2d810 100644 --- a/main.py +++ b/main.py @@ -1,60 +1,102 @@ import os from langchain_openai import ChatOpenAI from langgraph.checkpoint.memory import MemorySaver -from langgraph.prebuilt import create_react_agent as create_agent +from langgraph.prebuilt import create_react_agent from rich.console import Console -from rich.markdown import Markdown -from tools import get_price -# LLM configuration using BroJS +# --------------------------------------------------------------- +# LLM +# --------------------------------------------------------------- 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.5, + temperature=0.0, ) -# Memory checkpoint for interrupt handling +# --------------------------------------------------------------- +# Helper create_agent — обёртка, как описано в инструкции +# --------------------------------------------------------------- +def create_agent(model, tools, system_prompt="", checkpointer=None, interrupt_before=None): + """Helper-функция для создания агента с памятью и interrupt.""" + kwargs = dict(model=model, tools=tools) + if system_prompt: + kwargs["prompt"] = system_prompt + if checkpointer is not None: + kwargs["checkpointer"] = checkpointer + if interrupt_before is not None: + kwargs["interrupt_before"] = interrupt_before + return create_react_agent(**kwargs) + +# --------------------------------------------------------------- +# Инструменты +# --------------------------------------------------------------- +from langchain.tools import tool + +@tool +def get_price(city: str, date: str) -> str: + """Возвращает условную цену/погоду для города на указанную дату.""" + return f"Данные для {city} на {date}: 720 руб/кг, 18C, облачно" + +# --------------------------------------------------------------- +# Агент с памятью и паузой перед инструментом +# --------------------------------------------------------------- memory = MemorySaver() -# Create the agent with interrupt before tools agent = create_agent( model=llm, tools=[get_price], + system_prompt="Ты полезный ассистент. Помогай пользователю находить цены и погоду.", checkpointer=memory, interrupt_before=["tools"], ) console = Console() +# --------------------------------------------------------------- +# ask_and_run — точная структура из задания +# --------------------------------------------------------------- +def ask_and_run(user_input, config): + for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]): + state = agent.get_state(config) + chunk_type, chunk_data = chunk -def ask_and_run(user_input: str, config: dict): - """Synchronously run the agent with streaming and handle interrupts.""" - # Stream messages and updates - for chunk in agent.stream( - {"messages": [{"role": "user", "content": user_input}], - "configurable": config}, - stream_mode=["messages", "updates"], - ): - if isinstance(chunk, dict) and "__interrupt__" in chunk: - # Interrupt: ask for confirmation - console.print("[bold red]Agent requested tool execution. Confirm? (y/n): ", end="") - choice = input().strip().lower() - if choice != "y": - # Reject by sending a new message to the agent - config.update({"configurable": {"thread_id": config.get("thread_id", "default")}}) - continue - if isinstance(chunk, dict) and "messages" in chunk: - for msg in chunk["messages"]: - console.print(Markdown(msg["content"])) + if chunk_type == "messages": + msg = chunk_data[0] if isinstance(chunk_data, (list, tuple)) else chunk_data + content = getattr(msg, "content", "") + if content and not getattr(msg, "tool_calls", None): + console.print(content, end="") + if chunk_type == "updates": + console.print("\n --- --- --- \n") + for node_data in chunk_data.values(): + for msg in node_data.get("messages", []): + for tc in getattr(msg, "tool_calls", []): + console.print(f"[bold]{tc['name']}({tc['args']})[/bold]") + + if "__interrupt__" in chunk_data and state.next == ("tools",): + tool_call = state.values["messages"][-1].tool_calls[0] + console.print(f"Агент хочет вызвать утилиту [bold]{tool_call['name']}({tool_call['args']})[/bold]") + answer = input("Разрешить? (Y/n): ") + if answer.lower().strip() == "y": + ask_and_run(None, config) + else: + console.print("Отменено") + break + +# --------------------------------------------------------------- +# Чат-цикл +# --------------------------------------------------------------- if __name__ == "__main__": - thread_id = os.getenv("THREAD_ID", "session-1") - config = {"thread_id": thread_id} - console.print("[bold green]Welcome to the price agent. Type 'exit' to quit.") + config = {"configurable": {"thread_id": "разговор-1"}} + console.print("[bold green]Агент запущен. Введите 'exit' для выхода.[/bold green]") + while True: - user_input = input("Вы: ") - if user_input.lower() in ("exit", "quit"): + user_input = input("\nВы: ") + if user_input == "exit": break - ask_and_run(user_input, config) - console.print("[bold blue]Goodbye!", style="bold") + + ask_and_run( + {"messages": [{"role": "human", "content": user_input}]}, + config, + ) \ No newline at end of file