46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
from openai import OpenAI
|
||
|
||
|
||
class LLMMessage:
|
||
def __init__(self, role, content):
|
||
self.role = role
|
||
self.content = content
|
||
|
||
def to_dict(self):
|
||
return {"role": self.role, "content": self.content}
|
||
|
||
|
||
class LLMClient:
|
||
"""LLM客户端，对应一个API服务商"""
|
||
|
||
def __init__(self, url, api_key):
|
||
self.client = OpenAI(
|
||
base_url=url,
|
||
api_key=api_key,
|
||
)
|
||
|
||
def send_chat_request(self, model, messages):
|
||
"""发送对话请求，等待返回结果"""
|
||
response = self.client.chat.completions.create(model=model, messages=messages, stream=False)
|
||
if hasattr(response.choices[0].message, "reasoning_content"):
|
||
# 有单独的推理内容块
|
||
reasoning_content = response.choices[0].message.reasoning_content
|
||
content = response.choices[0].message.content
|
||
else:
|
||
# 无单独的推理内容块
|
||
response = response.choices[0].message.content.split("<think>")[-1].split("</think>")
|
||
# 如果有推理内容，则分割推理内容和内容
|
||
if len(response) == 2:
|
||
reasoning_content = response[0]
|
||
content = response[1]
|
||
else:
|
||
reasoning_content = None
|
||
content = response[0]
|
||
|
||
return reasoning_content, content
|
||
|
||
def send_embedding_request(self, model, text):
|
||
"""发送嵌入请求，等待返回结果"""
|
||
text = text.replace("\n", " ")
|
||
return self.client.embeddings.create(input=[text], model=model).data[0].embedding
|