# import json
from typing import Any
import requests

# Ref: https://medium.com/@shrinath.suresh/deploying-tgi-in-runpod-a84738263e2c

RUNPOD_POD_ID = "q39au00uendckb"

def predict(query: str) -> Any:
    url = f"https://{RUNPOD_POD_ID}-8080.proxy.runpod.net/v1/chat/completions"
    headers = {
        'Content-Type': 'application/json'
    }
    data = {
        "model": "tgi",
        "messages": [
            {
                "role": "system",
                "content": "You are a helpful medical expert. Answer the user's medical question as accurately as possible."
            },
            {
                "role": "user",
                "content": query
            }
        ],
        "stream": False,
        "max_tokens": 5000,
    }

    response = requests.post(url, headers=headers, json=data)
    result = response.json()['choices'][0]['message']['content']
    return result

response = predict("What are common treatments for sepsis?")
print(response)