import requests
import json

API_KEY = "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"

headers = {
    "Authorization": f"Bearer {API_KEY}",
    "Accept": "application/json",
    "Content-Type": "application/json"
}

system_prompt = """You are an SHL assessment recommender. Respond ONLY with a valid JSON object in this exact format:
{"reply": "your conversational response here", "recommendations": [], "end_of_conversation": false}
IMPORTANT: Respond ONLY with the JSON object. No markdown, no code fences, no extra text outside the JSON."""

payload = {
    "model": "google/diffusiongemma-26b-a4b-it",
    "messages": [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": "hi, I need help finding an assessment"}
    ],
    "max_tokens": 500,
    "temperature": 0.3,
    "stream": False,
}

r = requests.post("https://integrate.api.nvidia.com/v1/chat/completions", headers=headers, json=payload)
print("Status:", r.status_code)
content = r.json()["choices"][0]["message"]["content"]
print("Raw content:")
print(repr(content))
print()
print("Content:")
print(content)

# Try parsing
try:
    parsed = json.loads(content)
    print("\nParsed successfully as JSON!")
    print(json.dumps(parsed, indent=2))
except:
    print("\nFailed to parse as JSON directly")
    # Try extracting JSON
    first = content.find("{")
    last = content.rfind("}")
    if first != -1 and last != -1:
        try:
            parsed = json.loads(content[first:last+1])
            print("Extracted JSON:", json.dumps(parsed, indent=2))
        except:
            print("Failed to extract JSON either")