Files
Father Bot a7984073b1 Remove unused is_content_acceptable helper
The moderation helper had no callers anywhere in the codebase. Drop
the dead function.
2026-06-05 12:30:00 +03:00

330 lines
12 KiB
Python

import base64
from io import BytesIO
import config
import logging
import tiktoken
from openai import AsyncOpenAI, BadRequestError
# setup openai client
openai_client = AsyncOpenAI(
api_key=config.openai_api_key,
base_url=config.openai_api_base,
)
# optional OpenRouter client (OpenAI-compatible) for models declared with
# "provider: openrouter" in config/models.yml (e.g. Claude or other vendors)
openrouter_client = None
if config.openrouter_api_key:
openrouter_client = AsyncOpenAI(
api_key=config.openrouter_api_key,
base_url=config.openrouter_api_base,
)
logger = logging.getLogger(__name__)
def _get_client_for_model(model):
provider = config.models["info"].get(model, {}).get("provider", "openai")
if provider == "openrouter":
if openrouter_client is None:
raise ValueError(
"OpenRouter API key is not configured. "
"Set openrouter_api_key in config/config.yml"
)
return openrouter_client
return openai_client
OPENAI_COMPLETION_OPTIONS = {
"temperature": 0.7,
"max_tokens": 1000,
"top_p": 1,
"frequency_penalty": 0,
"presence_penalty": 0,
"timeout": 60.0,
}
class ChatGPT:
def __init__(self, model="gpt-4o-mini"):
assert model in config.models["info"], f"Unknown model: {model}"
self.model = model
self._client = _get_client_for_model(model)
async def send_message(self, message, dialog_messages=[], chat_mode="assistant"):
if chat_mode not in config.chat_modes.keys():
raise ValueError(f"Chat mode {chat_mode} is not supported")
n_dialog_messages_before = len(dialog_messages)
answer = None
while answer is None:
try:
if config.models["info"][self.model]["type"] == "chat_completion":
messages = self._generate_prompt_messages(message, dialog_messages, chat_mode)
r = await self._client.chat.completions.create(
model=self.model,
messages=messages,
**OPENAI_COMPLETION_OPTIONS
)
answer = r.choices[0].message.content
else:
raise ValueError(f"Unknown model: {self.model}")
answer = self._postprocess_answer(answer)
n_input_tokens, n_output_tokens = r.usage.prompt_tokens, r.usage.completion_tokens
except BadRequestError as e: # too many tokens
if len(dialog_messages) == 0:
raise ValueError("Dialog messages is reduced to zero, but still has too many tokens to make completion") from e
# forget first message in dialog_messages
dialog_messages = dialog_messages[1:]
n_first_dialog_messages_removed = n_dialog_messages_before - len(dialog_messages)
return answer, (n_input_tokens, n_output_tokens), n_first_dialog_messages_removed
async def send_message_stream(self, message, dialog_messages=[], chat_mode="assistant"):
if chat_mode not in config.chat_modes.keys():
raise ValueError(f"Chat mode {chat_mode} is not supported")
n_dialog_messages_before = len(dialog_messages)
answer = None
while answer is None:
try:
if config.models["info"][self.model]["type"] == "chat_completion":
messages = self._generate_prompt_messages(message, dialog_messages, chat_mode)
r_gen = await self._client.chat.completions.create(
model=self.model,
messages=messages,
stream=True,
**OPENAI_COMPLETION_OPTIONS
)
answer = ""
async for r_item in r_gen:
if len(r_item.choices) == 0:
continue
delta = r_item.choices[0].delta
if delta.content:
answer += delta.content
n_input_tokens, n_output_tokens = self._count_tokens_from_messages(messages, answer, model=self.model)
n_first_dialog_messages_removed = 0
yield "not_finished", answer, (n_input_tokens, n_output_tokens), n_first_dialog_messages_removed
answer = self._postprocess_answer(answer)
except BadRequestError as e: # too many tokens
if len(dialog_messages) == 0:
raise e
# forget first message in dialog_messages
dialog_messages = dialog_messages[1:]
yield "finished", answer, (n_input_tokens, n_output_tokens), n_first_dialog_messages_removed # sending final answer
async def send_vision_message(
self,
message,
dialog_messages=[],
chat_mode="assistant",
image_buffer: BytesIO = None,
):
n_dialog_messages_before = len(dialog_messages)
answer = None
while answer is None:
try:
if config.models["info"][self.model].get("vision", False):
messages = self._generate_prompt_messages(
message, dialog_messages, chat_mode, image_buffer
)
r = await self._client.chat.completions.create(
model=self.model,
messages=messages,
**OPENAI_COMPLETION_OPTIONS
)
answer = r.choices[0].message.content
else:
raise ValueError(f"Unsupported model: {self.model}")
answer = self._postprocess_answer(answer)
n_input_tokens, n_output_tokens = (
r.usage.prompt_tokens,
r.usage.completion_tokens,
)
except BadRequestError as e: # too many tokens
if len(dialog_messages) == 0:
raise ValueError(
"Dialog messages is reduced to zero, but still has too many tokens to make completion"
) from e
# forget first message in dialog_messages
dialog_messages = dialog_messages[1:]
n_first_dialog_messages_removed = n_dialog_messages_before - len(
dialog_messages
)
return (
answer,
(n_input_tokens, n_output_tokens),
n_first_dialog_messages_removed,
)
async def send_vision_message_stream(
self,
message,
dialog_messages=[],
chat_mode="assistant",
image_buffer: BytesIO = None,
):
n_dialog_messages_before = len(dialog_messages)
answer = None
while answer is None:
try:
if config.models["info"][self.model].get("vision", False):
messages = self._generate_prompt_messages(
message, dialog_messages, chat_mode, image_buffer
)
r_gen = await self._client.chat.completions.create(
model=self.model,
messages=messages,
stream=True,
**OPENAI_COMPLETION_OPTIONS,
)
answer = ""
async for r_item in r_gen:
if len(r_item.choices) == 0:
continue
delta = r_item.choices[0].delta
if delta.content:
answer += delta.content
(
n_input_tokens,
n_output_tokens,
) = self._count_tokens_from_messages(
messages, answer, model=self.model
)
n_first_dialog_messages_removed = (
n_dialog_messages_before - len(dialog_messages)
)
yield "not_finished", answer, (
n_input_tokens,
n_output_tokens,
), n_first_dialog_messages_removed
answer = self._postprocess_answer(answer)
except BadRequestError as e: # too many tokens
if len(dialog_messages) == 0:
raise e
# forget first message in dialog_messages
dialog_messages = dialog_messages[1:]
yield "finished", answer, (
n_input_tokens,
n_output_tokens,
), n_first_dialog_messages_removed
def _encode_image(self, image_buffer: BytesIO) -> bytes:
return base64.b64encode(image_buffer.read()).decode("utf-8")
def _generate_prompt_messages(self, message, dialog_messages, chat_mode, image_buffer: BytesIO = None):
prompt = config.chat_modes[chat_mode]["prompt_start"]
messages = [{"role": "system", "content": prompt}]
for dialog_message in dialog_messages:
messages.append({"role": "user", "content": dialog_message["user"]})
messages.append({"role": "assistant", "content": dialog_message["bot"]})
if image_buffer is not None:
messages.append(
{
"role": "user",
"content": [
{
"type": "text",
"text": message,
},
{
"type": "image_url",
"image_url" : {
"url": f"data:image/jpeg;base64,{self._encode_image(image_buffer)}",
"detail":"high"
}
}
]
}
)
else:
messages.append({"role": "user", "content": message})
return messages
def _postprocess_answer(self, answer):
answer = answer.strip()
return answer
def _count_tokens_from_messages(self, messages, answer, model="gpt-3.5-turbo"):
try:
encoding = tiktoken.encoding_for_model(model)
except KeyError:
# models not known to tiktoken (e.g. Claude or other models routed
# via OpenRouter) fall back to a modern encoding for an estimate
encoding = tiktoken.get_encoding("o200k_base")
# all currently supported chat models (gpt-4o, gpt-4o-mini, gpt-5.5,
# Claude, ...) use the same per-message overhead
tokens_per_message = 3
tokens_per_name = 1
# input
n_input_tokens = 0
for message in messages:
n_input_tokens += tokens_per_message
if isinstance(message["content"], list):
for sub_message in message["content"]:
if "type" in sub_message:
if sub_message["type"] == "text":
n_input_tokens += len(encoding.encode(sub_message["text"]))
elif sub_message["type"] == "image_url":
pass
else:
if "type" in message:
if message["type"] == "text":
n_input_tokens += len(encoding.encode(message["text"]))
elif message["type"] == "image_url":
pass
n_input_tokens += 2
# output
n_output_tokens = 1 + len(encoding.encode(answer))
return n_input_tokens, n_output_tokens
async def transcribe_audio(audio_file) -> str:
r = await openai_client.audio.transcriptions.create(model="whisper-1", file=audio_file)
return r.text or ""
async def generate_images(prompt, n_images=1, size="1024x1024"):
# gpt-image-1 returns base64-encoded images (no URLs), so decode to bytes
r = await openai_client.images.generate(
model="gpt-image-1", prompt=prompt, n=n_images, size=size
)
images = [base64.b64decode(item.b64_json) for item in r.data]
return images