39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
import httpx
|
|
from pretor.protocol_plugin.model_protocol.modelbase import ModelBase
|
|
|
|
class OpenAIAdapter(ModelBase):
|
|
def __init__(self, base_url: str, adapter_title: str, api_key: str = "archon-local"):
|
|
self.adapter_title: str = adapter_title
|
|
self.base_url = base_url.rstrip('/')
|
|
if not self.base_url.endswith('/v1'):
|
|
self.base_url += '/v1'
|
|
self.api_key = api_key
|
|
self.model_list = []
|
|
|
|
async def get_model(self):
|
|
url = "{}/models".format(self.base_url)
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}"
|
|
}
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.get(url, headers=headers)
|
|
response.raise_for_status()
|
|
response = response.json()
|
|
self.model_list = [m.get("id", "") for m in response.get("data", [])]
|
|
return self.model_list
|
|
|
|
async def post_message(self,model: str, messages: list, stream: bool = False, **kwargs):
|
|
url = f"{self.base_url}/chat/completions"
|
|
headers = {"Authorization": f"Bearer {self.api_key}"}
|
|
payload = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"stream": stream,
|
|
**kwargs
|
|
}
|
|
async with httpx.AsyncClient(timeout=None) as client:
|
|
if not stream:
|
|
response = await client.post(url, headers=headers, json=payload)
|
|
return response.json()
|
|
else:
|
|
return client.stream("POST", url, headers=headers, json=payload) |