55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
"""
|
|
Agent with memory and confirmation of tool usage.
|
|
"""
|
|
|
|
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
|
from langchain.chat_models import ChatOpenAI
|
|
from langchain.memory import ConversationBufferMemory
|
|
from langchain.tools import BaseTool
|
|
import os
|
|
|
|
# Example tool: get_price (placeholder)
|
|
class GetPriceTool(BaseTool):
|
|
name = "get_price"
|
|
description = "Get price for a city. Input: {city}"
|
|
|
|
def _run(self, city: str) -> str:
|
|
# Dummy implementation
|
|
return f"The price in {city} is $100"
|
|
|
|
# Setup LLM and memory
|
|
llm = ChatOpenAI(temperature=0)
|
|
memory = ConversationBufferMemory()
|
|
|
|
# Create agent with interrupt_before tools
|
|
agent = create_openai_tools_agent(
|
|
llm,
|
|
[GetPriceTool()],
|
|
verbose=True,
|
|
checkpointer=memory,
|
|
interrupt_before=["tools"]
|
|
)
|
|
|
|
executor = AgentExecutor(agent=agent, tools=[GetPriceTool()], verbose=True)
|
|
|
|
# Main loop
|
|
while True:
|
|
user_input = input("You: ")
|
|
if user_input.lower() in {"exit", "quit"}:
|
|
break
|
|
# Run executor and capture interrupt
|
|
try:
|
|
result = executor.run(user_input)
|
|
print(result)
|
|
except Exception as e:
|
|
if "__interrupt__" in str(e):
|
|
# Extract tool call info from exception message (simplified)
|
|
print("Agent wants to use a tool.")
|
|
confirm = input("Allow? (Y/n): ")
|
|
if confirm.lower() in {"", "y", "yes"}:
|
|
executor.run(user_input, run_name=None) # retry
|
|
else:
|
|
print("Action cancelled by user.")
|
|
else:
|
|
raise
|