import os
from dotenv import load_dotenv

# Load environment variables (HF_TOKEN)
load_dotenv()

import config
from smolagents import InferenceClientModel
from agents.orchestrator import create_orchestrator

def main():
    print("Testing AP Schemes Retrieval via Orchestrator...")
    
    # Initialize the model
    model = InferenceClientModel(
        model_id=config.MODEL_ID,
        token=os.getenv("HF_TOKEN")
    )
    
    # Create the agent orchestrator
    orchestrator = create_orchestrator(model)
    
    # Andhra Pradesh farmer profile
    profile_context = "[User Profile - Age: 45, Gender: Male, State: Andhra Pradesh, Category: OBC, Income: 80000, Occupation: Farmer, Area: Rural]"
    query = "What schemes am I eligible for?"
    
    full_query = f"{profile_context}\n\nUser Query: {query}"
    print(f"Running query: {full_query}")
    
    response = orchestrator.run(full_query)
    print("\n=== Agent Response ===")
    print(response)
    print("======================")

if __name__ == "__main__":
    main()