main.py

import sys
from langchain_core.messages import HumanMessage, AIMessage
import config

# Load keys first
print("Loading API keys...")
try:
    config.load_keys()
except Exception as e:
    print(f"Failed to load keys: {e}")
    sys.exit(1)

from graph import graph

def main():
    print("\n=== Multi-Agent Research Assistant (LangGraph + Groq + Tavily) ===")
    print("Type 'exit' or 'quit' to stop.\n")

    # Maintain conversation history here
    # Start with empty list or system prompt if needed
    chat_history = [] 

    while True:
        try:
            user_input = input("User: ").strip()
            if not user_input:
                continue
            
            if user_input.lower() in ["exit", "quit"]:
                print("Goodbye!")
                break

            # Add user message to history
            chat_history.append(HumanMessage(content=user_input))
            
            print("\nProcessing...", end="", flush=True)
            
            # Run the graph
            # We pass the full history. The graph will append its steps to it.
            # Note: Since our graph state is `messages: Annotated[list, add_messages]`,
            # passing the existing list might duplicate if we are not careful with checkpointers.
            # WITHOUT checkpointers, we just pass the input for this turn?
            # actually, if we rely on `add_messages` reducer and we don't have persistence,
            # we should pass the full history as `messages`. The output will contain the full history + new messages.
            # We then update our `chat_history` variable with the result.
            
            inputs = {"messages": chat_history, "attempt_count": 0}
            
            final_state = None
            for event in graph.stream(inputs):
                for key, value in event.items():
                    print(f"\n  [{key}] Active")
                    # Optional: Print partial updates
                    if "clarity_status" in value:
                        print(f"    Status: {value['clarity_status']}")
                    if "confidence_score" in value:
                        print(f"    Confidence: {value['confidence_score']}")
            
                # Helper to get state snapshot (if using persistent graph)
                # But here 'event' is the output of the node.
                # We need the final state to update our history.
                pass
            
            # Since stream yields partial updates, we need the final result.
            # Let's use invoke for the state update, or just trust that the last event + history merge?
            # Easier: Just use invoke to get final state for history, 
            # OR parse the stream carefully.
            # Let's do `graph.invoke` to ensure we get the unified final state.
            # (We streamed for UX above, re-invoking might be wasteful but safe for this simple app, 
            # OR we just rely on the last output of stream if it returns state chunks)
            
            # Optimization: Just use invoke and print summary? 
            # Or formatted printing.
            
            result = graph.invoke(inputs)
            
            # Update history with the final state's messages
            chat_history = result["messages"]
            
            # Get the last message to display
            last_msg = chat_history[-1]
            print(f"\nAgent: {last_msg.content}\n")

        except KeyboardInterrupt:
            print("\nGoodbye!")
            break
        except Exception as e:
            print(f"\nAn error occurred: {e}")
            import traceback
            traceback.print_exc()

if __name__ == "__main__":
    main()