scripts/test_sarvam_provider.py

"""
Manual smoke test for Sarvam AI provider.

Loads .env, sends a tiny prompt to Sarvam, prints model_used and response.
Never prints the API key.

Usage:
    python scripts/test_sarvam_provider.py [--base-url https://api.sarvam.ai/v1]

Exit codes:
    0 = success
    1 = failure
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from urllib.error import URLError
from urllib.request import Request, urlopen

BACKEND_DIR = Path(__file__).resolve().parents[1]

# Load .env manually (no dependency on python-dotenv for scripts)
env_path = BACKEND_DIR / ".env"
if env_path.exists():
    for line in env_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" in line:
            key, _, value = line.partition("=")
            key = key.strip()
            value = value.strip().strip('"').strip("'")
            os.environ.setdefault(key, value)


def main() -> None:
    parser = argparse.ArgumentParser(description="Sarvam AI Smoke Test")
    parser.add_argument(
        "--base-url",
        default=os.environ.get("SARVAM_BASE_URL", "https://api.sarvam.ai/v1"),
    )
    args = parser.parse_args()
    base = args.base_url.rstrip("/")
    api_key = os.environ.get("SARVAM_API_KEY", "")

    print("=" * 50)
    print("  DocDoe AI - Sarvam Provider Smoke Test")
    print(f"  Base URL: {base}")
    print(f"  Key configured: {'yes' if api_key else 'NO'}")
    print("=" * 50)

    if not api_key:
        print("\n  ERROR: SARVAM_API_KEY is not set in .env or environment.")
        print("  Set SARVAM_API_KEY and try again.")
        sys.exit(1)

    model = os.environ.get("SARVAM_MODEL_MAIN", "sarvam-30b")
    url = f"{base}/chat/completions"

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": "You are a helpful tutor. Respond with one short JSON object only."},
            {"role": "user", "content": 'Return JSON with keys "greeting" and "status". Keep it short.'},
        ],
        "temperature": 0.1,
        "max_tokens": 800,
        "reasoning_effort": "low",
    }

    print(f"\n  Calling: POST {url}")
    print(f"  Model: {model}")

    data = json.dumps(payload).encode()
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }
    req = Request(url, data=data, headers=headers, method="POST")

    try:
        with urlopen(req, timeout=30) as resp:
            result = json.loads(resp.read().decode())
    except URLError as exc:
        print(f"\n  FAIL: {exc}")
        sys.exit(1)
    except Exception as exc:
        print(f"\n  FAIL: {type(exc).__name__}: {exc}")
        sys.exit(1)

    # Parse response
    content = ""
    model_used = result.get("model", "unknown")
    choices = result.get("choices", [])
    if choices:
        msg = choices[0].get("message", {})
        content = msg.get("content") or ""
    usage = result.get("usage", {})

    print(f"\n  model_used: {model_used}")
    print(f"  content: {(content or '(empty)')[:300]}")
    print(f"  raw keys: {list(result.keys())}")
    if choices:
        print(f"  choices[0] keys: {list(choices[0].keys())}")
        print(f"  message keys: {list(choices[0].get('message', {}).keys())}")
    if usage:
        print(f"  tokens: prompt={usage.get('prompt_tokens', '?')}, "
              f"completion={usage.get('completion_tokens', '?')}, "
              f"total={usage.get('total_tokens', '?')}")

    if content:
        print("\n  RESULT: PASS - Sarvam API is responding.")
        sys.exit(0)
    else:
        # Print full raw response for debugging
        print(f"\n  Full raw response: {json.dumps(result, indent=2)[:500]}")
        print("\n  RESULT: FAIL - Empty content from Sarvam.")
        sys.exit(1)


if __name__ == "__main__":
    main()