examples/basic_usage.py

"""Example usage of MbazaAI"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

from mbaza_ai import LegalAI, CountryDatasetLoader, Config
from mbaza_ai.utils import setup_logging

# Setup logging
setup_logging("INFO")

def main():
    # Initialize the MbazaAI system
    print("🏛️  Initializing MbazaAI...")
    config = Config()
    dataset_loader = CountryDatasetLoader(config)
    legal_ai = LegalAI(config, dataset_loader)
    
    print(f"📍 Current country: {dataset_loader.current_country}")
    print(f"🌍 Available countries: {', '.join(config.get_available_countries())}")
    
    # Example questions
    questions = [
        "What are the requirements for Rwandan citizenship?",
        "What is the penalty for theft in Rwanda?",
        "What are the basic constitutional rights?",
    ]
    
    print("\n" + "="*60)
    print("🤖 MbazaAI Legal Question Answering Demo")
    print("="*60)
    
    # Ask questions to the current dataset (Rwanda)
    for question in questions:
        print(f"\n❓ Question: {question}")
        result = legal_ai.predict(question)
        print(f"✅ Answer: {result['answer']}")
        print(f"🎯 Confidence: {result['confidence']:.2f}")
        print(f"📚 Category: {result.get('category', 'N/A')}")
        print("-" * 40)
    
    # Demonstrate adding a new country
    print("\n🌍 Adding a new country (Kenya)...")
    config.add_country(
        country_code="kenya",
        name="Kenya", 
        dataset_path="datasets/kenya/legal_data.json",
        language="en",
        legal_system="common_law"
    )
    
    print(f"✅ Updated available countries: {', '.join(config.get_available_countries())}")
    
    # Switch to Kenya (will create sample data)
    print("\n🔄 Switching to Kenya legal dataset...")
    if legal_ai.switch_country("kenya"):
        print("✅ Successfully switched to Kenya!")
        
        # Ask a question in the new context
        result = legal_ai.predict("What are the basic rights in Kenya?")
        print(f"\n❓ Question: What are the basic rights in Kenya?")
        print(f"✅ Answer: {result['answer']}")
        print(f"📍 Country: {result['country']}")
    
    # Show model information
    print("\n" + "="*40)
    print("📊 Model Information")
    print("="*40)
    info = legal_ai.get_model_info()
    print(f"Model: {info['model_name']}")
    print(f"Current Country: {info['country']}")
    print(f"Legal Categories: {', '.join(info['legal_categories'])}")
    
    print("\n🎉 Demo completed successfully!")

if __name__ == "__main__":
    main()