test_gemma.py

import asyncio
import httpx
from backend.config import settings
from backend.llm.auth import get_gcp_bearer_token

async def main():
    endpoint_id = settings.VERTEX_LLM_ENDPOINT_ID
    region = settings.VERTEX_LLM_REGION
    project = settings.GCP_PROJECT
    
    if not endpoint_id:
        print("Please set VERTEX_LLM_ENDPOINT_ID in your .env file")
        return

    print(f"Testing Vertex AI Endpoint: {endpoint_id} in {region}")
    
    token = get_gcp_bearer_token()
    if not token:
        print("Failed to get GCP token. Are you logged in via `gcloud auth application-default login`?")
        return

    url = f"https://{region}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{region}/endpoints/{endpoint_id}/chat/completions"
    
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": "google/gemma-4-26B-A4B-it",
        "messages": [
            {"role": "user", "content": "Hello, please reply with a single word: 'Success!'"}
        ],
        "max_tokens": 10,
        "temperature": 0.0
    }

    print("Sending request...")
    async with httpx.AsyncClient(timeout=60) as client:
        response = await client.post(url, headers=headers, json=payload)
        
        print(f"Status Code: {response.status_code}")
        if response.status_code == 200:
            print("Response:", response.json()["choices"][0]["message"]["content"])
        else:
            print("Error:", response.text)

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