47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
import requests
|
|
import os
|
|
|
|
class ExternalModelService:
|
|
def generate(self, provider: str, model: str, prompt: str):
|
|
if provider == "openai":
|
|
return self.openai(model, prompt)
|
|
|
|
if provider == "anthropic":
|
|
return self.anthropic(model, prompt)
|
|
|
|
if provider == "mistral":
|
|
return self.mistral(model, prompt)
|
|
|
|
return "Erro: provider desconhecido."
|
|
|
|
def openai(self, model, prompt):
|
|
url = "https://api.openai.com/v1/chat/completions"
|
|
headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}"}
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}]
|
|
}
|
|
r = requests.post(url, json=payload, headers=headers)
|
|
return r.json()["choices"][0]["message"]["content"]
|
|
|
|
def anthropic(self, model, prompt):
|
|
url = "https://api.anthropic.com/v1/messages"
|
|
headers = {"x-api-key": os.getenv("ANTHROPIC_API_KEY")}
|
|
payload = {
|
|
"model": model,
|
|
"max_tokens": 4096,
|
|
"messages": [{"role": "user", "content": prompt}]
|
|
}
|
|
r = requests.post(url, json=payload, headers=headers)
|
|
return r.json()["content"][0]["text"]
|
|
|
|
def mistral(self, model, prompt):
|
|
url = "https://api.mistral.ai/v1/chat/completions"
|
|
headers = {"Authorization": f"Bearer {os.getenv('MISTRAL_API_KEY')}"}
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}]
|
|
}
|
|
r = requests.post(url, json=payload, headers=headers)
|
|
return r.json()["choices"][0]["message"]["content"]
|