tools/app.py

from pathlib import Path

import click

from second_brain_online.application.agents import get_agent
from second_brain_online.application.ui.custom_gradio_ui import CustomGradioUI


@click.command()
@click.option(
    "--retriever-config-path",
    type=click.Path(exists=True),
    required=True,
    help="Path to the retriever config file",
)
@click.option(
    "--ui",
    is_flag=True,
    default=False,
    help="Launch with Gradio UI instead of CLI mode",
)
@click.option(
    "--query",
    "-q",
    type=str,
    default="What is the feature/training/inference (FTI) pipelines architecture?",
    help="Query to run in CLI mode",
)
def main(retriever_config_path: Path, ui: bool, query: str) -> None:
    """Run the agent either in Gradio UI or CLI mode.

    Args:
        ui: If True, launches Gradio UI. If False, runs in CLI mode
        query: Query string to run in CLI mode
    """
    agent = get_agent(retriever_config_path=Path(retriever_config_path))
    if ui:
        # Get the actual agent from the wrapper
        actual_agent = agent._AgentWrapper__agent
        CustomGradioUI(actual_agent).launch()
    else:
        assert query, "Query is required in CLI mode"

        result = agent.run(query)

        # DEBUG: Print raw result
        print("\n" + "="*80)
        print("DEBUG: RAW AGENT RESULT")
        print("="*80)
        print(f"Type: {type(result)}")
        print(f"Full Content:\n{result}")
        print("="*80)
        
        # DEBUG: Check agent object attributes
        print("\n" + "="*80)
        print("DEBUG: AGENT OBJECT ATTRIBUTES")
        print("="*80)
        print(f"Agent type: {type(agent)}")
        print(f"Agent attributes: {dir(agent)}")
        if hasattr(agent, '_AgentWrapper__agent'):
            actual_agent = agent._AgentWrapper__agent
            print(f"Actual agent type: {type(actual_agent)}")
            print(f"Actual agent attributes: {dir(actual_agent)}")
            if hasattr(actual_agent, 'conversation_history'):
                print(f"Conversation history: {actual_agent.conversation_history}")
            if hasattr(actual_agent, 'messages'):
                print(f"Messages: {actual_agent.messages}")
            if hasattr(actual_agent, 'logs'):
                print(f"Logs: {actual_agent.logs}")
            if hasattr(actual_agent, 'state'):
                print(f"State: {actual_agent.state}")
        print("="*80)
        
        # Parse the result using the same logic as the UI
        ui_instance = CustomGradioUI(None)  # We don't need the agent for parsing
        
        # Get agent logs if available
        agent_logs = []
        if hasattr(agent, '_AgentWrapper__agent'):
            actual_agent = agent._AgentWrapper__agent
            if hasattr(actual_agent, 'logs'):
                agent_logs = actual_agent.logs
        
        answer, sources, tools_used = ui_instance.parse_agent_response(result, agent_logs)
        
        print("\n" + "="*80)
        print("DEBUG: PARSED RESULTS")
        print("="*80)
        print(f"Answer: {answer}")
        print(f"Sources ({len(sources)}): {sources}")
        print(f"Tools Used: {tools_used}")
        print("="*80)
        
        print("\n" + "="*80)
        print("FINAL OUTPUT")
        print("="*80)
        
        # Format the answer for better display
        formatted_answer = ui_instance.format_answer(answer)
        print(formatted_answer)


if __name__ == "__main__":
    main()