diff --git a/main.py b/main.py index 799d73f..b4ee740 100644 --- a/main.py +++ b/main.py @@ -51,18 +51,21 @@ Android, Apache, BGL, Hadoop, HDFS, HealthApp, HPC, Linux, Mac, OpenSSH, OpenSta return agent -def stream_agent_updates(agent, user_input: str): - """Stream agent updates for a user input.""" +def stream_agent_updates(agent, user_input: str, conversation_history: list): + """Stream agent updates for a user input with conversation history.""" # Create a human message message = HumanMessage(content=user_input) + # Add the new message to conversation history + conversation_history.append(message) + print("\nAgent: ", end="", flush=True) - # Use the agent's stream method to get real-time updates + # Use the agent's stream method to get real-time updates with full conversation final_response = "" tool_calls_made = False - for event in agent.stream({"messages": [message]}, stream_mode="updates"): + for event in agent.stream({"messages": conversation_history}, stream_mode="updates"): for node_name, node_output in event.items(): if node_name == "agent" and "messages" in node_output: last_message = node_output["messages"][-1] @@ -91,6 +94,9 @@ def stream_agent_updates(agent, user_input: str): print(f"\n\n{final_response}") else: print(final_response) + # Add the agent's response to conversation history + from langchain_core.messages import AIMessage + conversation_history.append(AIMessage(content=final_response)) else: print("No response generated.") @@ -138,6 +144,7 @@ def main(): print("Type 'quit', 'exit', or 'q' to exit the chat.") print("Type 'help' or 'h' for help and examples.") print("Type 'graph' to see the agent structure.") + print("Type 'clear' or 'reset' to clear conversation history.") print("⚠️ WARNING: This agent has shell access - use with caution!") print("📊 Available log analysis capabilities:") print(" - Analyze log files in the loghub directory") @@ -160,6 +167,8 @@ def main(): return # Start the chat loop + conversation_history = [] # Initialize conversation history + while True: try: user_input = input("\nUser: ") @@ -178,13 +187,18 @@ def main(): print(" - 'List all available log files'") print(" - 'Find error patterns in Linux logs'") print(" - 'Check disk usage on the system'") + print(" - 'clear': Clear conversation history") continue elif user_input.lower() in ["graph", "structure"]: visualize_agent(agent) continue + elif user_input.lower() in ["clear", "reset"]: + conversation_history = [] + print("🗑️ Conversation history cleared!") + continue if user_input.strip(): - stream_agent_updates(agent, user_input) + stream_agent_updates(agent, user_input, conversation_history) else: print("Please enter a message.")