80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
# main.py
|
|
# Реализация агента с памятью и подтверждением действий
|
|
|
|
import os
|
|
from rich.console import Console
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langgraph import create_agent
|
|
from langchain_ollama import ChatOllama
|
|
from langchain.tools import Tool
|
|
|
|
# Rich console for pretty output
|
|
console = Console()
|
|
|
|
# ---------- 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":
|
|
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 = 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
|
|
|
|
# ---------- 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().strip() == "exit":
|
|
console.print("[bold magenta]До свидания![/bold magenta]")
|
|
break
|
|
message = {"messages": [{"role": "human", "content": user_input}]}
|
|
ask_and_run(message, config)
|
|
console.print("\n--- --- ---")
|