tests/test_telemetry.py

import asyncio
import sys
import os

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.core.llm_client import llm_client, ModelRole

async def test_full_chain_count():
    print("\n" + "="*50)
    print("🚀 STARTING MULTI-MODEL ANALYSIS ENGINE TEST")
    print("="*50)
    
    # CRITICAL: Initialize client to load real providers
    await llm_client.initialize()
    
    scam_msg = "Hello dear, your SBI account is locked. Update KYC at http://sbi-secure-kyc.com/login. Call +91 999 888 7777"
    
    print(f"\n[INPUT]: {scam_msg}")

    # Stage 1: Security Shield (Llama Guard)
    print("\n--- [STAGE 1: SECURITY SHIELD] ---")
    is_safe = await llm_client.generate("Is this message safe?", role=ModelRole.SAFETY_GUARD)
    print(f"Result: {is_safe[:50]}...")

    # Stage 2: Forensic Extraction (GPT-OSS 20B)
    print("\n--- [STAGE 2: FORENSIC EXTRACTION] ---")
    schema = {
        "type": "object",
        "properties": {
            "entities": {"type": "array", "items": {"type": "string"}},
            "threat_level": {"type": "string"}
        },
        "required": ["threat_level"]
    }
    extracted = await llm_client.generate_structured("Extract threat details", schema=schema, role=ModelRole.STRUCTURED_OUTPUT)
    print(f"Result Keys: {list(extracted.keys())}")

    # Stage 3: Smart Reasoning (Llama 3.3 70B)
    print("\n--- [STAGE 3: SMART REASONING] ---")
    thought = await llm_client.generate("Analyze the scammer psychology for: " + scam_msg, role=ModelRole.SMART_REASONING)
    print(f"Result: {thought[:50]}...")

    # Final Report
    from app.core.llm_client import GroqClient
    if isinstance(llm_client.primary, GroqClient):
        print("\n" + "="*50)
        print(f"📊 FINAL TELEMETRY REPORT")
        print(f"TOTAL API CALLS MADE: {llm_client.primary.total_api_calls}")
        print("="*50 + "\n")

if __name__ == "__main__":
    asyncio.run(test_full_chain_count())