54b8a20a19
Messages were mistakenly being sent as Pydantic objects, but templates expect dictionaries. Properly convert these before render. In addition, initialize all Optional lists as an empty list since this will cause the least problems when interacting with other parts of API code, such as templates. Signed-off-by: kingbri <8082010+kingbri1@users.noreply.github.com>
506 lines
17 KiB
Python
506 lines
17 KiB
Python
"""Chat completion utilities for OAI server."""
|
|
|
|
import asyncio
|
|
import json
|
|
import pathlib
|
|
from asyncio import CancelledError
|
|
from typing import List, Optional
|
|
from fastapi import HTTPException, Request
|
|
from jinja2 import TemplateError
|
|
from loguru import logger
|
|
|
|
from common import model
|
|
from common.multimodal import MultimodalEmbeddingWrapper
|
|
from common.networking import (
|
|
get_generator_error,
|
|
handle_request_disconnect,
|
|
handle_request_error,
|
|
request_disconnect_loop,
|
|
)
|
|
from common.utils import unwrap
|
|
from endpoints.OAI.types.chat_completion import (
|
|
ChatCompletionLogprobs,
|
|
ChatCompletionLogprob,
|
|
ChatCompletionMessage,
|
|
ChatCompletionRequest,
|
|
ChatCompletionRespChoice,
|
|
ChatCompletionStreamChunk,
|
|
ChatCompletionResponse,
|
|
ChatCompletionStreamChoice,
|
|
)
|
|
from endpoints.OAI.types.common import UsageStats
|
|
from endpoints.OAI.utils.completion import _parse_gen_request_id, _stream_collector
|
|
from endpoints.OAI.utils.tools import ToolCallProcessor
|
|
|
|
|
|
def _create_response(
|
|
request_id: str, generations: List[dict], model_name: Optional[str]
|
|
):
|
|
"""Create a chat completion response from the provided text."""
|
|
|
|
prompt_tokens = unwrap(generations[-1].get("prompt_tokens"), 0)
|
|
completion_tokens = unwrap(generations[-1].get("generated_tokens"), 0)
|
|
|
|
choices = []
|
|
for index, generation in enumerate(generations):
|
|
message = ChatCompletionMessage(
|
|
role="assistant", content=unwrap(generation.get("text"), "")
|
|
)
|
|
|
|
tool_calls = generation["tool_calls"]
|
|
if tool_calls:
|
|
message.tool_calls = ToolCallProcessor.from_json(tool_calls)
|
|
|
|
logprob_response = None
|
|
|
|
token_probs = unwrap(generation.get("token_probs"), {})
|
|
if token_probs:
|
|
logprobs = unwrap(generation.get("logprobs"), [])
|
|
|
|
collected_token_probs = []
|
|
for index, token in enumerate(token_probs.keys()):
|
|
top_logprobs = [
|
|
ChatCompletionLogprob(token=token, logprob=logprob)
|
|
for token, logprob in logprobs[index].items()
|
|
]
|
|
|
|
collected_token_probs.append(
|
|
ChatCompletionLogprob(
|
|
token=token,
|
|
logprob=token_probs[token],
|
|
top_logprobs=top_logprobs,
|
|
)
|
|
)
|
|
|
|
logprob_response = ChatCompletionLogprobs(content=collected_token_probs)
|
|
|
|
# Initialize finish_reason with a default value or from generation data
|
|
finish_reason = generation.get("finish_reason", "stop")
|
|
|
|
# If a tool call is present, mark the finish reason as such
|
|
if message.tool_calls:
|
|
finish_reason = "tool_calls"
|
|
|
|
choice = ChatCompletionRespChoice(
|
|
index=index,
|
|
finish_reason=finish_reason,
|
|
stop_str=generation.get("stop_str"),
|
|
message=message,
|
|
logprobs=logprob_response,
|
|
)
|
|
|
|
choices.append(choice)
|
|
|
|
response = ChatCompletionResponse(
|
|
id=f"chatcmpl-{request_id}",
|
|
choices=choices,
|
|
model=unwrap(model_name, ""),
|
|
usage=UsageStats(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
total_tokens=prompt_tokens + completion_tokens,
|
|
),
|
|
)
|
|
|
|
return response
|
|
|
|
|
|
def _create_stream_chunk(
|
|
request_id: str,
|
|
generation: Optional[dict] = None,
|
|
model_name: Optional[str] = None,
|
|
is_usage_chunk: bool = False,
|
|
):
|
|
"""Create a chat completion stream chunk from the provided text."""
|
|
|
|
index = generation.get("index")
|
|
choices = []
|
|
usage_stats = None
|
|
|
|
if is_usage_chunk:
|
|
prompt_tokens = unwrap(generation.get("prompt_tokens"), 0)
|
|
completion_tokens = unwrap(generation.get("generated_tokens"), 0)
|
|
|
|
usage_stats = UsageStats(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
total_tokens=prompt_tokens + completion_tokens,
|
|
)
|
|
elif "finish_reason" in generation:
|
|
# Get the finish reason from the generation
|
|
finish_reason = generation.get("finish_reason")
|
|
choice = ChatCompletionStreamChoice(index=index, finish_reason=finish_reason)
|
|
|
|
# lets check if we have tool calls since we are at the end of the generation
|
|
# Mark finish_reason as tool_calls since this is the last chunk
|
|
if "tool_calls" in generation:
|
|
tool_calls = generation["tool_calls"]
|
|
message = ChatCompletionMessage(
|
|
tool_calls=ToolCallProcessor.from_json(tool_calls)
|
|
)
|
|
choice.delta = message
|
|
choice.finish_reason = "tool_calls"
|
|
|
|
choices.append(choice)
|
|
|
|
else:
|
|
message = ChatCompletionMessage(
|
|
role="assistant", content=unwrap(generation.get("text"), "")
|
|
)
|
|
|
|
logprob_response = None
|
|
|
|
token_probs = unwrap(generation.get("token_probs"), {})
|
|
if token_probs:
|
|
logprobs = unwrap(generation.get("logprobs"), {})
|
|
top_logprobs = [
|
|
ChatCompletionLogprob(token=token, logprob=logprob)
|
|
for token, logprob in logprobs.items()
|
|
]
|
|
|
|
generated_token = next(iter(token_probs))
|
|
token_prob_response = ChatCompletionLogprob(
|
|
token=generated_token,
|
|
logprob=token_probs[generated_token],
|
|
top_logprobs=top_logprobs,
|
|
)
|
|
|
|
logprob_response = ChatCompletionLogprobs(content=[token_prob_response])
|
|
|
|
choice = ChatCompletionStreamChoice(
|
|
index=index,
|
|
delta=message,
|
|
logprobs=logprob_response,
|
|
)
|
|
|
|
choices.append(choice)
|
|
|
|
chunk = ChatCompletionStreamChunk(
|
|
id=f"chatcmpl-{request_id}",
|
|
choices=choices,
|
|
model=unwrap(model_name, ""),
|
|
usage=usage_stats,
|
|
)
|
|
|
|
return chunk
|
|
|
|
|
|
async def _append_template_metadata(data: ChatCompletionRequest, template_vars: dict):
|
|
"""Adding metadata is a one-time process."""
|
|
|
|
template_metadata = await model.container.prompt_template.extract_metadata(
|
|
template_vars
|
|
)
|
|
|
|
# Stop strings
|
|
if isinstance(data.stop, str):
|
|
data.stop = [data.stop] + template_metadata.stop_strings
|
|
else:
|
|
data.stop += template_metadata.stop_strings
|
|
|
|
# Tool call start strings
|
|
if template_metadata.tool_starts:
|
|
if data.tool_call_start is None:
|
|
data.tool_call_start = template_metadata.tool_starts
|
|
|
|
# Append to stop strings to halt for a tool call generation
|
|
data.stop.extend(template_metadata.tool_starts)
|
|
|
|
|
|
async def format_messages_with_template(
|
|
messages: List[ChatCompletionMessage],
|
|
existing_template_vars: Optional[dict] = None,
|
|
):
|
|
"""Barebones function to format chat completion messages into a prompt."""
|
|
|
|
template_vars = unwrap(existing_template_vars, {})
|
|
mm_embeddings = MultimodalEmbeddingWrapper() if model.container.use_vision else None
|
|
|
|
# Convert all messages to a dictionary representation
|
|
message_dicts: List[dict] = []
|
|
for message in messages:
|
|
if isinstance(message.content, list):
|
|
concatenated_content = ""
|
|
for content in message.content:
|
|
if content.type == "text":
|
|
concatenated_content += content.text
|
|
elif content.type == "image_url" and mm_embeddings:
|
|
await mm_embeddings.add(content.image_url.url)
|
|
concatenated_content += mm_embeddings.text_alias[-1]
|
|
|
|
# Convert the message content into a concatenated string
|
|
message.content = concatenated_content
|
|
|
|
if message.tool_calls:
|
|
message.tool_calls_json = ToolCallProcessor.to_json(message.tool_calls)
|
|
|
|
# The tools variable is inspectable in the template, so
|
|
# store the list of dicts rather than the ToolCallProcessor object.
|
|
message.tool_calls = ToolCallProcessor.dump(message.tool_calls)
|
|
|
|
message_dicts.append(message.model_dump())
|
|
|
|
# Get all special tokens
|
|
special_tokens_dict = model.container.get_special_tokens()
|
|
|
|
template_vars.update({"messages": message_dicts, **special_tokens_dict})
|
|
|
|
prompt = await model.container.prompt_template.render(template_vars)
|
|
return prompt, mm_embeddings, template_vars
|
|
|
|
|
|
async def apply_chat_template(
|
|
data: ChatCompletionRequest, tool_precursor: Optional[str] = None
|
|
):
|
|
"""
|
|
Compile the prompt and get any additional stop strings from the template.
|
|
Template stop strings can be overriden by sampler overrides if force is true.
|
|
"""
|
|
|
|
# Locally store tools dict
|
|
tools = data.model_dump()["tools"]
|
|
|
|
try:
|
|
data.template_vars.update(
|
|
{
|
|
"add_generation_prompt": data.add_generation_prompt,
|
|
"tools": tools,
|
|
"tools_json": json.dumps(tools, indent=2),
|
|
"functions": data.functions,
|
|
"functions_json": json.dumps(data.functions, indent=2),
|
|
"tool_precursor": tool_precursor,
|
|
}
|
|
)
|
|
|
|
prompt, mm_embeddings, template_vars = await format_messages_with_template(
|
|
data.messages, data.template_vars
|
|
)
|
|
|
|
# Append response prefix if present
|
|
if data.response_prefix:
|
|
if data.add_generation_prompt:
|
|
prompt += data.response_prefix
|
|
else:
|
|
logger.warning(
|
|
"Could not add response prefix because "
|
|
"add_generation_prompt is False"
|
|
)
|
|
|
|
# Add template metadata
|
|
await _append_template_metadata(data, template_vars)
|
|
|
|
return prompt, mm_embeddings
|
|
|
|
except KeyError as exc:
|
|
error_message = handle_request_error(
|
|
"Could not find a Conversation from prompt template "
|
|
f"'{model.container.prompt_template.name}'. "
|
|
"Check your spelling?",
|
|
).error.message
|
|
|
|
raise HTTPException(400, error_message) from exc
|
|
except TemplateError as exc:
|
|
error_message = handle_request_error(f"TemplateError: {str(exc)}").error.message
|
|
|
|
raise HTTPException(400, error_message) from exc
|
|
|
|
|
|
async def stream_generate_chat_completion(
|
|
prompt: str,
|
|
embeddings: MultimodalEmbeddingWrapper,
|
|
data: ChatCompletionRequest,
|
|
request: Request,
|
|
model_path: pathlib.Path,
|
|
):
|
|
"""Generator for the generation process."""
|
|
abort_event = asyncio.Event()
|
|
gen_queue = asyncio.Queue()
|
|
gen_tasks: List[asyncio.Task] = []
|
|
disconnect_task = asyncio.create_task(request_disconnect_loop(request))
|
|
|
|
try:
|
|
logger.info(f"Received chat completion streaming request {request.state.id}")
|
|
|
|
for idx in range(0, data.n):
|
|
task_gen_params = data.model_copy(deep=True)
|
|
request_id = _parse_gen_request_id(data.n, request.state.id, idx)
|
|
|
|
gen_task = asyncio.create_task(
|
|
_stream_collector(
|
|
idx,
|
|
gen_queue,
|
|
request_id,
|
|
prompt,
|
|
task_gen_params,
|
|
abort_event,
|
|
mm_embeddings=embeddings,
|
|
)
|
|
)
|
|
|
|
gen_tasks.append(gen_task)
|
|
|
|
# We need to keep track of the text generated so we can resume the tool calls
|
|
current_generation_text = ""
|
|
|
|
# Consumer loop
|
|
while True:
|
|
if disconnect_task.done():
|
|
abort_event.set()
|
|
handle_request_disconnect(
|
|
f"Chat completion generation {request.state.id} cancelled by user."
|
|
)
|
|
|
|
generation = await gen_queue.get()
|
|
# lets only append the text if we need it for tool calls later
|
|
if data.tool_call_start and "text" in generation:
|
|
current_generation_text += generation["text"]
|
|
|
|
# check if we are running a tool model, and that we are at stop
|
|
if data.tool_call_start and "stop_str" in generation:
|
|
generations = await generate_tool_calls(
|
|
data,
|
|
[generation],
|
|
request,
|
|
current_generations=current_generation_text,
|
|
)
|
|
generation = generations[0] # We only have one generation in this case
|
|
|
|
# Stream collector will push an exception to the queue if it fails
|
|
if isinstance(generation, Exception):
|
|
raise generation
|
|
|
|
response = _create_stream_chunk(
|
|
request.state.id, generation, model_path.name
|
|
)
|
|
yield response.model_dump_json()
|
|
|
|
# Check if all tasks are completed
|
|
if all(task.done() for task in gen_tasks) and gen_queue.empty():
|
|
# Send a usage chunk
|
|
if data.stream_options and data.stream_options.include_usage:
|
|
usage_chunk = _create_stream_chunk(
|
|
request.state.id,
|
|
generation,
|
|
model_path.name,
|
|
is_usage_chunk=True,
|
|
)
|
|
yield usage_chunk.model_dump_json()
|
|
|
|
logger.info(
|
|
f"Finished chat completion streaming request {request.state.id}"
|
|
)
|
|
|
|
yield "[DONE]"
|
|
break
|
|
except CancelledError:
|
|
# Get out if the request gets disconnected
|
|
|
|
if not disconnect_task.done():
|
|
abort_event.set()
|
|
handle_request_disconnect("Chat completion generation cancelled by user.")
|
|
except Exception:
|
|
yield get_generator_error(
|
|
"Chat completion aborted. Please check the server console."
|
|
)
|
|
|
|
|
|
async def generate_chat_completion(
|
|
prompt: str,
|
|
embeddings: MultimodalEmbeddingWrapper,
|
|
data: ChatCompletionRequest,
|
|
request: Request,
|
|
model_path: pathlib.Path,
|
|
):
|
|
gen_tasks: List[asyncio.Task] = []
|
|
|
|
try:
|
|
logger.info(f"Received chat completion request {request.state.id}")
|
|
|
|
for idx in range(0, data.n):
|
|
request_id = _parse_gen_request_id(data.n, request.state.id, idx)
|
|
|
|
gen_tasks.append(
|
|
asyncio.create_task(
|
|
model.container.generate(
|
|
request_id,
|
|
prompt,
|
|
data,
|
|
mm_embeddings=embeddings,
|
|
)
|
|
)
|
|
)
|
|
|
|
generations = await asyncio.gather(*gen_tasks)
|
|
|
|
# Let's not waste our time if we arn't running a tool model
|
|
if data.tool_call_start:
|
|
generations = await generate_tool_calls(data, generations, request)
|
|
|
|
response = _create_response(request.state.id, generations, model_path.name)
|
|
|
|
logger.info(f"Finished chat completion request {request.state.id}")
|
|
|
|
return response
|
|
except Exception as exc:
|
|
error_message = handle_request_error(
|
|
f"Chat completion {request.state.id} aborted. "
|
|
"Maybe the model was unloaded? "
|
|
"Please check the server console."
|
|
).error.message
|
|
|
|
# Server error if there's a generation exception
|
|
raise HTTPException(503, error_message) from exc
|
|
|
|
|
|
async def generate_tool_calls(
|
|
data: ChatCompletionRequest,
|
|
generations: List[str],
|
|
request: Request,
|
|
current_generations: str = None,
|
|
):
|
|
gen_tasks: List[asyncio.Task] = []
|
|
tool_idx: List[int] = []
|
|
|
|
# Copy to make sure the parent JSON schema doesn't get modified
|
|
# FIXME: May not be necessary depending on how the codebase evolves
|
|
tool_data = data.model_copy(deep=True)
|
|
tool_data.json_schema = tool_data.tool_call_schema
|
|
|
|
for idx, gen in enumerate(generations):
|
|
if gen["stop_str"] in tool_data.tool_call_start:
|
|
logger.info(
|
|
f"Detected tool call in chat completion request {request.state.id}"
|
|
)
|
|
|
|
if "text" in gen:
|
|
# non streaming, all generations will have the text they generated
|
|
pre_tool_prompt, embeddings = await apply_chat_template(
|
|
data, gen["text"]
|
|
)
|
|
elif current_generations is not None:
|
|
# streaming, we wont have text in the generation,
|
|
# we'll have to use the current_generations
|
|
pre_tool_prompt, embeddings = await apply_chat_template(
|
|
data, current_generations
|
|
)
|
|
|
|
request_id = _parse_gen_request_id(data.n, request.state.id, idx)
|
|
|
|
gen_tasks.append(
|
|
asyncio.create_task(
|
|
model.container.generate(
|
|
request_id,
|
|
pre_tool_prompt,
|
|
tool_data,
|
|
mm_embeddings=embeddings,
|
|
)
|
|
)
|
|
)
|
|
tool_idx.append(idx)
|
|
|
|
tool_calls = await asyncio.gather(*gen_tasks)
|
|
for outer_idx in range(0, len(tool_idx)):
|
|
gen_idx = tool_idx[outer_idx]
|
|
generations[gen_idx]["tool_calls"] = tool_calls[outer_idx]["text"]
|
|
|
|
return generations
|