141 lines
5.1 KiB
Python
141 lines
5.1 KiB
Python
import os
|
|
from langchain.chat_models import init_chat_model
|
|
from langchain_community.tools.shell.tool import ShellTool
|
|
from langgraph.prebuilt import create_react_agent
|
|
from langchain_core.messages import HumanMessage
|
|
from custom_tools.poem_tool import print_poem
|
|
|
|
|
|
def create_agent():
|
|
"""Create and return a ReAct agent with shell and poem capabilities."""
|
|
|
|
# Initialize the chat model (using OpenAI GPT-4)
|
|
# Make sure you have set your OPENAI_API_KEY environment variable
|
|
llm = init_chat_model("openai:gpt-4o-mini")
|
|
|
|
# Define the tools available to the agent
|
|
shell_tool = ShellTool()
|
|
tools = [shell_tool, print_poem]
|
|
|
|
|
|
# Create a ReAct agent with simplified system prompt
|
|
system_prompt = """You are a helpful assistant with access to shell commands and poem generation capabilities.
|
|
|
|
You can:
|
|
1. Execute shell commands using the shell tool to interact with the system
|
|
2. Generate and print beautiful poems using the print_poem tool
|
|
|
|
The poem tool can create different types of poems:
|
|
- "nature": Poems about nature and the environment
|
|
- "tech": Poems about technology and programming
|
|
- "motivational": Inspirational and motivational poems
|
|
- "friendship": Poems about friendship and relationships
|
|
- "random": Randomly selects from any available poem type
|
|
|
|
When helping users:
|
|
- Be friendly and helpful
|
|
- Use appropriate tools based on the user's request
|
|
- Always be cautious with shell commands and explain what they do
|
|
- For system monitoring commands like 'top', 'ps', etc., use appropriate flags to avoid hanging
|
|
- When users ask for poems or want to brighten their day, use the poem tool
|
|
"""
|
|
|
|
|
|
# Create the ReAct agent
|
|
agent = create_react_agent(
|
|
llm,
|
|
tools,
|
|
prompt=system_prompt
|
|
)
|
|
|
|
return agent
|
|
|
|
|
|
def run_agent_query(agent, user_input: str, conversation_history: list):
|
|
"""Run a simple agent query and display results cleanly."""
|
|
# Create a human message
|
|
message = HumanMessage(content=user_input)
|
|
|
|
# Add the new message to conversation history
|
|
conversation_history.append(message)
|
|
|
|
# Use the agent's stream method for clean output like the LangChain tutorial
|
|
for step in agent.stream({"messages": conversation_history}, stream_mode="values"):
|
|
step["messages"][-1].pretty_print()
|
|
|
|
# Add the agent's response to conversation history
|
|
if step and "messages" in step:
|
|
conversation_history.append(step["messages"][-1])
|
|
|
|
|
|
|
|
def main():
|
|
# Check if required API keys are set
|
|
if not os.getenv("OPENAI_API_KEY"):
|
|
print("Please set your OPENAI_API_KEY environment variable.")
|
|
print("You can set it by running: export OPENAI_API_KEY='your-api-key-here'")
|
|
return
|
|
|
|
print("🤖 LangGraph Simple Demo Agent")
|
|
print("Type 'quit', 'exit', or 'q' to exit the chat.")
|
|
print("Type 'help' or 'h' for help and examples.")
|
|
print("Type 'clear' or 'reset' to clear conversation history.")
|
|
print("⚠️ WARNING: This agent has shell access - use with caution!")
|
|
print("🎭 Available capabilities:")
|
|
print(" - Generate beautiful poems for any occasion")
|
|
print(" - Execute shell commands for system tasks")
|
|
print(" - Help with general assistance")
|
|
print("-" * 60)
|
|
|
|
# Create the agent
|
|
try:
|
|
agent = create_agent()
|
|
print("✅ Simple Demo Agent initialized successfully!")
|
|
print("💡 Try asking: 'Write me a nature poem'")
|
|
print("💡 Or: 'Show me the current directory'")
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error initializing agent: {e}")
|
|
return
|
|
|
|
# Start the chat loop
|
|
conversation_history = [] # Initialize conversation history
|
|
|
|
while True:
|
|
try:
|
|
user_input = input("\nUser: ")
|
|
if user_input.lower() in ["quit", "exit", "q"]:
|
|
print("👋 Goodbye!")
|
|
break
|
|
elif user_input.lower() in ["help", "h"]:
|
|
print("\n🆘 Help:")
|
|
print("Commands:")
|
|
print(" - quit/exit/q: Exit the agent")
|
|
print(" - help/h: Show this help")
|
|
print(" - clear/reset: Clear conversation history")
|
|
print("\nExample queries:")
|
|
print(" - 'Write me a motivational poem'")
|
|
print(" - 'Generate a tech poem'")
|
|
print(" - 'Show me a random poem'")
|
|
print(" - 'List files in current directory'")
|
|
print(" - 'Check disk usage on the system'")
|
|
continue
|
|
elif user_input.lower() in ["clear", "reset"]:
|
|
conversation_history = []
|
|
print("🗑️ Conversation history cleared!")
|
|
continue
|
|
|
|
if user_input.strip():
|
|
run_agent_query(agent, user_input, conversation_history)
|
|
else:
|
|
print("Please enter a message.")
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n👋 Goodbye!")
|
|
break
|
|
except Exception as e:
|
|
print(f"❌ Error: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |