main.py

"""
AI Research Concierge - Main Application

This module builds and executes the LangGraph workflow for intelligent research.
"""

from langgraph.graph import StateGraph, END
from graph import (
    GraphState,
    analyze_query,
    call_tool_node,
    synthesize_answer,
    handle_error,
    should_continue_after_analyze,
    should_continue_after_fetch,
    increment_retry,
)


def build_graph() -> StateGraph:
    """
    Build the research agent graph.
    
    Graph structure:
        START
          ↓
        analyze_query
          ↓
        [conditional: error check]
          ↓              ↓
        fetch_info    handle_error
          ↓              ↓
        [conditional: retry logic]
          ↓       ↓         ↓
       synthesize retry  handle_error
          ↓       ↓         ↓
         END   analyze    END
    
    Returns:
        Compiled StateGraph ready for execution
    """

    workflow = StateGraph(GraphState)
    
    workflow.add_node("analyze", analyze_query)
    workflow.add_node("fetch", call_tool_node)
    workflow.add_node("synthesize", synthesize_answer)
    workflow.add_node("error", handle_error)
    workflow.add_node("retry", increment_retry)
    
    
    workflow.set_entry_point("analyze")
    
    workflow.add_conditional_edges(
        "analyze",
        should_continue_after_analyze,
        {
            "fetch": "fetch",
            "error": "error",
        }
    )
    
    workflow.add_conditional_edges(
        "fetch",
        should_continue_after_fetch,
        {
            "synthesize": "synthesize",
            "retry": "retry",
            "error": "error",
        }
    )
    
    # Retry loop: increment counter then go back to analyze
    workflow.add_edge("retry", "analyze")
    
    workflow.add_edge("synthesize", END)
    workflow.add_edge("error", END)
    
    return workflow.compile()


def run_research_agent(query: str, verbose: bool = True) -> dict:
    """
    Run the research agent on a query.
    
    Args:
        query: The user's research question
        verbose: Whether to print progress information
        
    Returns:
        Final state dictionary containing the answer
    """
    
    app = build_graph()
    
    initial_state: GraphState = {
        "user_query": query,
        "sub_questions": [],
        "tool_results": {},
        "final_answer": None,
        "error": None,
        "retry_count": 0,
        "current_step": "start",
    }
    
    if verbose:
        print(f"\n{'='*60}")
        print(f"RESEARCH AGENT STARTED")
        print(f"{'='*60}")
        print(f"Query: {query}")
    
    try:
        final_state = app.invoke(initial_state)
        
        if verbose:
            print(f"\n{'='*60}")
            print(f"RESEARCH AGENT COMPLETED")
            print(f"{'='*60}")
        
        return final_state
        
    except Exception as e:
        error_msg = f"Critical error in graph execution: {type(e).__name__}: {str(e)}"
        
        if verbose:
            print(f"\n{'='*60}")
            print(f"CRITICAL ERROR")
            print(f"{'='*60}")
            print(error_msg)
        
        return {
            **initial_state,
            "error": error_msg,
            "final_answer": f"# Critical Error\n\nThe research agent encountered a critical error:\n\n{error_msg}\n\nPlease try again or contact support if the issue persists.",
        }


def main():
    """Main entry point for interactive CLI usage."""
    print("="*60)
    print("AI Research Concierge")
    print("="*60)
    print("\nAsk me any research question and I'll break it down,")
    print("gather information, and provide a comprehensive answer.")
    print("\nType 'quit' or 'exit' to stop.\n")
    
    while True:
        try:
            
            query = input("\n Your question: ").strip()
            
            if not query:
                continue
                
            if query.lower() in ['quit', 'exit', 'q']:
                print("\nGoodbye!")
                break
            
            result = run_research_agent(query, verbose=True)
            
            print("\n" + "="*60)
            print("FINAL ANSWER")
            print("="*60 + "\n")
            print(result.get("final_answer", "No answer generated."))
            
        except KeyboardInterrupt:
            print("\n\nInterrupted by user. Goodbye! ")
            break
            
        except Exception as e:
            print(f"\n Unexpected error: {type(e).__name__}: {str(e)}")
            print("Please try again.\n")


if __name__ == "__main__":
    main()