41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
|
import json
|
||
|
|
import urllib.request
|
||
|
|
import urllib.error
|
||
|
|
from config import LLM_BASE_URL, LLM_MODEL, LLM_API_KEY, LLM_TIMEOUT
|
||
|
|
|
||
|
|
class LLMClient:
|
||
|
|
class Message:
|
||
|
|
def __init__(self, msg):
|
||
|
|
self.content = msg.get('content', '')
|
||
|
|
self.tool_calls = msg.get('tool_calls', None)
|
||
|
|
|
||
|
|
def __init__(self, base_url=LLM_BASE_URL, model=LLM_MODEL, api_key=LLM_API_KEY):
|
||
|
|
self.base_url = base_url.rstrip('/')
|
||
|
|
self.model = model
|
||
|
|
self.api_key = api_key
|
||
|
|
|
||
|
|
def chat(self, messages, tools=None):
|
||
|
|
url = f"{self.base_url}/chat/completions"
|
||
|
|
payload = {
|
||
|
|
"model": self.model,
|
||
|
|
"messages": messages
|
||
|
|
}
|
||
|
|
if tools:
|
||
|
|
payload["tools"] = tools
|
||
|
|
payload["tool_choice"] = "auto"
|
||
|
|
|
||
|
|
data = json.dumps(payload).encode('utf-8')
|
||
|
|
req = urllib.request.Request(url, data=data, method='POST')
|
||
|
|
req.add_header('Content-Type', 'application/json')
|
||
|
|
req.add_header('Authorization', f'Bearer {self.api_key}')
|
||
|
|
|
||
|
|
try:
|
||
|
|
with urllib.request.urlopen(req, timeout=LLM_TIMEOUT) as resp:
|
||
|
|
response = json.loads(resp.read().decode('utf-8'))
|
||
|
|
message = response['choices'][0]['message']
|
||
|
|
return self.Message(message)
|
||
|
|
except urllib.error.HTTPError as e:
|
||
|
|
return self.Message({'content': f"HTTP Error: {e.code} {e.reason}", 'tool_calls': None})
|
||
|
|
except Exception as e:
|
||
|
|
return self.Message({'content': f"Error: {str(e)}", 'tool_calls': None})
|