File size: 1,662 Bytes
3f3265f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | import os
import requests
import json
import llm_gemini
AVAILABLE_MODELS = [
"claude-3-haiku",
"claude-3-sonnet",
"claude-3-opus",
"gemini-pro",
"gpt-3.5-turbo",
"gpt-4-turbo",
"gpt-4",
]
MODEL = "claude-3-haiku"
URL = "https://api.theb.ai/v1/chat/completions"
# URL = "https://api.baizhi.ai/v1/chat/completions"
API_KEY = os.environ.get("THEB_API_KEY", "[THEB_API_KEY]")
def generate_content(prompt, model=MODEL, api_key=API_KEY):
assert model in AVAILABLE_MODELS, f"Model {model} not available"
if model == "gemini-pro":
return llm_gemini.generate_content(prompt)
payload = json.dumps({
"model": model,
"messages": [
{
"role": "user",
"content": prompt
}
],
"stream": False
})
headers = {
'Authorization': f'Bearer {api_key}',
'Content-Type': 'application/json'
}
response = requests.request("POST", URL, headers=headers, data=payload, timeout=10)
if response.status_code != 200:
retry_count = 0
while retry_count < 5:
response = requests.request("POST", URL, headers=headers, data=payload, timeout=10)
if response.status_code == 200:
break
else:
retry_count += 1
if retry_count == 5:
raise Exception(f"Failed to get completion from LLM. Status code: {response.status_code}. Response: {response.text}")
# print(response)
return response.json()['choices'][0]['message']['content']
if __name__ == "__main__":
print(generate_content("What are you?", "claude-3-haiku"))
|