langbot-app--langbot
RockChinQ
101 行
2.5 KiB
Python
101 行
2.5 KiB
Python
import openai
|
|
from openai.types import completion, completion_choice
|
|
|
|
from . import model
|
|
|
|
|
|
class CompletionRequest(model.RequestBase):
|
|
"""调用Completion接口的请求类。
|
|
|
|
调用方可以一直next completion直到finish_reason为stop。
|
|
"""
|
|
|
|
model: str
|
|
prompt: str
|
|
kwargs: dict
|
|
|
|
stopped: bool = False
|
|
|
|
def __init__(
|
|
self,
|
|
client: openai.Client,
|
|
model: str,
|
|
messages: list[dict[str, str]],
|
|
**kwargs
|
|
):
|
|
self.client = client
|
|
self.model = model
|
|
self.prompt = ""
|
|
|
|
for message in messages:
|
|
self.prompt += message["role"] + ": " + message["content"] + "\n"
|
|
|
|
self.prompt += "assistant: "
|
|
|
|
self.kwargs = kwargs
|
|
|
|
self.req_func = self.client.completions.create
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self) -> dict:
|
|
"""调用Completion接口，返回生成的文本
|
|
|
|
{
|
|
"id": "id",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"type": "text",
|
|
"content": "message"
|
|
},
|
|
"finish_reason": "reason"
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 30
|
|
}
|
|
}
|
|
"""
|
|
|
|
if self.stopped:
|
|
raise StopIteration()
|
|
|
|
resp: completion.Completion = self._req(
|
|
model=self.model,
|
|
prompt=self.prompt,
|
|
**self.kwargs
|
|
)
|
|
|
|
if resp.choices[0].finish_reason == "stop":
|
|
self.stopped = True
|
|
|
|
choice0: completion_choice.CompletionChoice = resp.choices[0]
|
|
|
|
self.prompt += choice0.text
|
|
|
|
return {
|
|
"id": resp.id,
|
|
"choices": [
|
|
{
|
|
"index": choice0.index,
|
|
"message": {
|
|
"role": "assistant",
|
|
"type": "text",
|
|
"content": choice0.text
|
|
},
|
|
"finish_reason": choice0.finish_reason
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": resp.usage.prompt_tokens,
|
|
"completion_tokens": resp.usage.completion_tokens,
|
|
"total_tokens": resp.usage.total_tokens
|
|
}
|
|
}
|