main.py

"""
Dynamic RAG Bot - Main CLI Interface (Integrated Version)
Complete integration of main bot with document loader.
"""

from bot import ask_bot
from document_reader import load_document


def parse_fetch_action(action: str) -> tuple[str, str]:
    """
    Parse fetch action to extract country and question.
    
    Args:
        action: Action string like "fetch:france:Tell me about wine regions"
    
    Returns:
        Tuple of (country, question)
    """
    if not action.startswith("fetch:"):
        return None, None
    
    # Remove "fetch:" prefix and split on first two colons
    parts = action[6:].split(":", 2)
    
    if len(parts) >= 2:
        return parts[0], parts[1]
    elif len(parts) == 1:
        return parts[0], "general information"
    
    return None, None


def main():
    """Main CLI loop for the integrated dynamic RAG bot."""
    print("Dynamic RAG Bot - Integrated System")
    print("Ask me about any of these 10 countries: France, Spain, Japan, Germany, Italy, Brazil, India, China, USA, Australia")
    print("I can answer basic facts or fetch detailed information from documents as needed.")
    print("Type 'quit' to exit.")
    print("-" * 80)
    
    conversation_history = []
    
    while True:
        # Get user input
        user_input = input("\nYou: ").strip()
        
        # Check for exit commands
        if user_input.lower() in ['quit', 'exit', 'q']:
            print("Goodbye!")
            break
            
        if not user_input:
            continue
        
        try:
            # Get response from main bot
            bot_response = ask_bot(user_input, conversation_history=conversation_history)
            
            # Debug: Show bot's thinking
            print(f"\n[Bot thinking: {bot_response.thoughts}]")
            print(f"[Action: {bot_response.action}]")
            
            # Handle different actions
            final_response_content = None
            
            if bot_response.action == "answer":
                print(f"\nBot: {bot_response.content}")
                final_response_content = bot_response.content
                
            elif bot_response.action == "ask_specifics":
                print(f"\nBot: {bot_response.content}")
                final_response_content = bot_response.content
                
            elif bot_response.action.startswith("fetch:"):
                # Check if it's a multi-fetch (contains pipe)
                if "|" in bot_response.action:
                    # Multiple fetches for comparison
                    fetch_actions = bot_response.action.split("|")
                    document_responses = {}
                    
                    print(f"\n[MULTI-FETCH] Processing {len(fetch_actions)} documents for comparison...")
                    
                    for fetch in fetch_actions:
                        country, question = parse_fetch_action(fetch.strip())
                        if country and question:
                            print(f"\n[FETCHING] {country.title()}: {question}")
                            doc_response = load_document(country, question)
                            document_responses[country] = doc_response
                            
                            # Show what was extracted from the document
                            print(f"[EXTRACTED] {doc_response.content}")
                            if doc_response.sources:
                                print(f"[SOURCES] {len(doc_response.sources)} citation(s)")
                                for i, source in enumerate(doc_response.sources, 1):
                                    print(f"  {i}. {source.file} - {source.section} (lines {source.line_start}-{source.line_end})")
                        else:
                            print(f"\n[ERROR] Could not parse fetch action: {fetch}")
                    
                    if document_responses:
                        # Call bot again with document context for synthesis
                        print(f"\n[SYNTHESIZING] Combining information from {len(document_responses)} documents...")
                        synthesis_response = ask_bot(
                            user_query=user_input,  # Original query
                            document_context=document_responses,
                            conversation_history=conversation_history
                        )
                        print(f"\nBot: {synthesis_response.content}")
                        final_response_content = synthesis_response.content
                    else:
                        print("Bot: I couldn't fetch any documents for the comparison.")
                        final_response_content = "I couldn't fetch any documents for the comparison."
                        
                else:
                    # Single fetch - now consistent with multi-fetch pattern
                    country, question = parse_fetch_action(bot_response.action)
                    
                    if country and question:
                        print(f"\n[FETCHING] Document for '{country}': {question}")
                        
                        # Load document information
                        doc_response = load_document(country, question)
                        print(f"  ✓ {len(doc_response.content)} characters extracted")
                        
                        # Show what was extracted from the document
                        print(f"[EXTRACTED] {doc_response.content}")
                        if doc_response.sources:
                            print(f"[SOURCES] {len(doc_response.sources)} citation(s)")
                            for i, source in enumerate(doc_response.sources, 1):
                                print(f"  {i}. {source.file} - {source.section} (lines {source.line_start}-{source.line_end})")
                        
                        # Pass document context back to bot for consistent response
                        print(f"\n[SYNTHESIZING] Formatting response...")
                        document_context = {country: doc_response}
                        synthesis_response = ask_bot(
                            user_query=user_input,  # Original user query
                            document_context=document_context,
                            conversation_history=conversation_history
                        )
                        
                        print(f"\nBot: {synthesis_response.content}")
                        final_response_content = synthesis_response.content
                    else:
                        print(f"\n[ERROR] Could not parse fetch action: {bot_response.action}")
                        print("Bot: I wanted to fetch information but couldn't parse the action properly.")
                        final_response_content = "I wanted to fetch information but couldn't parse the action properly."
                    
            else:
                # Unknown action
                print(f"\n[UNKNOWN ACTION] {bot_response.action}")
                if bot_response.content:
                    print(f"Bot: {bot_response.content}")
                    final_response_content = bot_response.content
                else:
                    print("Bot: I'm not sure how to handle that request.")
                    final_response_content = "I'm not sure how to handle that request."
            
            # Add conversation to history (if we got a response)
            if final_response_content:
                conversation_history.append({"role": "user", "content": user_input})
                conversation_history.append({"role": "assistant", "content": final_response_content})
                
                # Keep conversation history reasonable length (last 10 exchanges = 20 messages)
                if len(conversation_history) > 20:
                    conversation_history = conversation_history[-20:]
                
        except Exception as e:
            print(f"\nError: {e}")
            print("Please try your question again.")
    

if __name__ == "__main__":
    main()