1066 lines
43 KiB
Rust
1066 lines
43 KiB
Rust
use super::*;
|
|
|
|
pub(super) fn final_response(
|
|
stream: &mut TcpStream,
|
|
state: &State,
|
|
request: ResponseOptions,
|
|
active: crate::runtime::ActiveGeneration,
|
|
id: &str,
|
|
) -> Result<(), (u16, String)> {
|
|
let output = match wait_for_output(stream, active) {
|
|
Ok(output) => output,
|
|
Err(error) if error == "client disconnected" => return Ok(()),
|
|
Err(error) => return send_generation_error(stream, request.protocol, request.cors, error),
|
|
};
|
|
let (content, calls) = parse_generated_tools(state, &output.message.content, request.protocol);
|
|
let finish = if calls.is_empty() {
|
|
output.finish_reason
|
|
} else {
|
|
"tool_calls"
|
|
};
|
|
let reasoning = output.message.reasoning.filter(|value| !value.is_empty());
|
|
let usage = usage_json(
|
|
output.prompt_tokens,
|
|
output.cached_tokens,
|
|
output.completion_tokens,
|
|
);
|
|
let body = match request.protocol {
|
|
Protocol::Chat => {
|
|
let mut message = json!({"role": "assistant", "content": content});
|
|
if let Some(reasoning) = reasoning {
|
|
message["reasoning_content"] = Value::String(reasoning);
|
|
}
|
|
if !calls.is_empty() {
|
|
message["tool_calls"] = tool_calls_json(&calls);
|
|
}
|
|
json!({
|
|
"id": id, "object": "chat.completion", "created": unix_time(),
|
|
"model": request.model_id,
|
|
"choices": [{"index": 0, "message": message, "finish_reason": finish}],
|
|
"usage": usage,
|
|
})
|
|
}
|
|
Protocol::Completion => json!({
|
|
"id": id, "object": "text_completion", "created": unix_time(),
|
|
"model": request.model_id,
|
|
"choices": [{"text": content, "index": 0, "finish_reason": finish}],
|
|
"usage": usage,
|
|
}),
|
|
Protocol::Anthropic => {
|
|
let mut blocks = Vec::new();
|
|
if let Some(reasoning) = reasoning {
|
|
blocks.push(json!({"type": "thinking", "thinking": reasoning, "signature": id}));
|
|
}
|
|
if !content.is_empty() {
|
|
blocks.push(json!({"type": "text", "text": content}));
|
|
}
|
|
for call in &calls {
|
|
blocks.push(json!({
|
|
"type": "tool_use", "id": call.id, "name": call.function.name,
|
|
"input": serde_json::from_str::<Value>(&call.function.arguments).unwrap_or_else(|_| json!({}))
|
|
}));
|
|
}
|
|
if blocks.is_empty() || (blocks.iter().all(|block| block["type"] == "thinking")) {
|
|
blocks.push(json!({"type": "text", "text": ""}));
|
|
}
|
|
let cached = output.cached_tokens.min(output.prompt_tokens);
|
|
let written = output.prompt_tokens - cached;
|
|
json!({
|
|
"id": id, "type": "message", "role": "assistant", "model": request.model_id,
|
|
"content": blocks,
|
|
"stop_reason": if finish == "tool_calls" { "tool_use" } else if finish == "length" { "max_tokens" } else { "end_turn" },
|
|
"stop_sequence": Value::Null,
|
|
"usage": {
|
|
"input_tokens": output.prompt_tokens - cached - written,
|
|
"output_tokens": output.completion_tokens,
|
|
"cache_read_input_tokens": cached,
|
|
"cache_creation_input_tokens": written
|
|
}
|
|
})
|
|
}
|
|
Protocol::Responses => {
|
|
let status = if finish == "length" {
|
|
"incomplete"
|
|
} else if finish == "error" {
|
|
"failed"
|
|
} else {
|
|
"completed"
|
|
};
|
|
let mut items = Vec::new();
|
|
if let Some(reasoning) = reasoning {
|
|
items.push(json!({
|
|
"id": random_id("rs_"), "type": "reasoning", "status": status,
|
|
"summary": [{"type": "summary_text", "text": reasoning}]
|
|
}));
|
|
}
|
|
if !content.is_empty() {
|
|
items.push(json!({
|
|
"id": random_id("msg_"), "type": "message", "status": status,
|
|
"role": "assistant", "content": [{"type": "output_text", "text": content, "annotations": []}]
|
|
}));
|
|
}
|
|
for call in &calls {
|
|
items.push(json!({
|
|
"id": random_id("fc_"), "type": "function_call", "status": status,
|
|
"name": call.function.name, "call_id": call.id, "arguments": call.function.arguments
|
|
}));
|
|
}
|
|
json!({
|
|
"id": id, "object": "response", "created_at": unix_time(), "status": status,
|
|
"model": request.model_id, "output": items,
|
|
"usage": {
|
|
"input_tokens": output.prompt_tokens,
|
|
"input_tokens_details": {"cached_tokens": output.cached_tokens.min(output.prompt_tokens), "cache_write_tokens": output.prompt_tokens - output.cached_tokens.min(output.prompt_tokens)},
|
|
"output_tokens": output.completion_tokens,
|
|
"output_tokens_details": {"reasoning_tokens": 0},
|
|
"total_tokens": output.prompt_tokens + output.completion_tokens
|
|
}
|
|
})
|
|
}
|
|
};
|
|
send_json_with_cors(stream, 200, &body, request.cors).map_err(|error| (500, error))
|
|
}
|
|
|
|
pub(super) fn stream_response(
|
|
stream: &mut TcpStream,
|
|
state: &State,
|
|
request: ResponseOptions,
|
|
active: crate::runtime::ActiveGeneration,
|
|
id: &str,
|
|
) -> Result<(), (u16, String)> {
|
|
if request.protocol != Protocol::Chat {
|
|
return structured_stream_response(stream, state, request, active, id);
|
|
}
|
|
stream_response_with_keepalive(
|
|
stream,
|
|
state,
|
|
request,
|
|
active,
|
|
id,
|
|
PREFILL_KEEPALIVE_INTERVAL,
|
|
)
|
|
}
|
|
|
|
fn structured_stream_response(
|
|
stream: &mut TcpStream,
|
|
state: &State,
|
|
request: ResponseOptions,
|
|
active: crate::runtime::ActiveGeneration,
|
|
id: &str,
|
|
) -> Result<(), (u16, String)> {
|
|
match request.protocol {
|
|
Protocol::Completion => completion_stream_response(stream, request, active, id),
|
|
Protocol::Anthropic => anthropic_stream_response(stream, state, request, active, id),
|
|
Protocol::Responses => responses_stream_response(stream, state, request, active, id),
|
|
Protocol::Chat => unreachable!(),
|
|
}
|
|
}
|
|
|
|
fn completion_stream_response(
|
|
stream: &mut TcpStream,
|
|
request: ResponseOptions,
|
|
active: crate::runtime::ActiveGeneration,
|
|
id: &str,
|
|
) -> Result<(), (u16, String)> {
|
|
send_sse_headers_with_cors(stream, request.cors).map_err(|error| (500, error))?;
|
|
while let Ok(event) = active.events.recv() {
|
|
match event {
|
|
GenerationEvent::Compacted(_) => {
|
|
return Err((
|
|
500,
|
|
"The model runtime returned an unexpected compaction event.".into(),
|
|
));
|
|
}
|
|
GenerationEvent::Chunk {
|
|
reasoning: false,
|
|
content,
|
|
} if !content.is_empty() => send_sse(
|
|
stream,
|
|
&json!({
|
|
"id": id, "object": "text_completion", "created": unix_time(),
|
|
"model": request.model_id,
|
|
"choices": [{"text": content, "index": 0, "finish_reason": Value::Null}]
|
|
}),
|
|
)
|
|
.map_err(|error| (500, error))?,
|
|
GenerationEvent::Finished(Ok(output)) => {
|
|
send_sse(
|
|
stream,
|
|
&json!({
|
|
"id": id, "object": "text_completion", "created": unix_time(),
|
|
"model": request.model_id,
|
|
"choices": [{"text": "", "index": 0, "finish_reason": output.finish_reason}]
|
|
}),
|
|
)
|
|
.map_err(|error| (500, error))?;
|
|
if request.include_usage {
|
|
send_sse(
|
|
stream,
|
|
&json!({
|
|
"id": id, "object": "text_completion", "created": unix_time(),
|
|
"model": request.model_id, "choices": [],
|
|
"usage": usage_json(output.prompt_tokens, output.cached_tokens, output.completion_tokens)
|
|
}),
|
|
)
|
|
.map_err(|error| (500, error))?;
|
|
}
|
|
return stream
|
|
.write_all(b"data: [DONE]\n\n")
|
|
.map_err(|error| (500, error.to_string()));
|
|
}
|
|
GenerationEvent::Finished(Err(error)) => {
|
|
let _ = send_sse_error(stream, &error);
|
|
return Ok(());
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
Err((500, "The model runtime stopped unexpectedly.".into()))
|
|
}
|
|
|
|
fn anthropic_stream_response(
|
|
stream: &mut TcpStream,
|
|
state: &State,
|
|
request: ResponseOptions,
|
|
active: crate::runtime::ActiveGeneration,
|
|
id: &str,
|
|
) -> Result<(), (u16, String)> {
|
|
send_sse_headers_with_cors(stream, request.cors).map_err(|error| (500, error))?;
|
|
let mut prompt_tokens = 0;
|
|
let mut started = false;
|
|
let mut block = None::<(usize, bool)>;
|
|
let mut next_index = 0;
|
|
let mut projector = ToolProjector::new();
|
|
let mut tool_indices = Vec::new();
|
|
while let Ok(event) = active.events.recv() {
|
|
match event {
|
|
GenerationEvent::Context {
|
|
used,
|
|
tokens_per_second,
|
|
..
|
|
} => {
|
|
if tokens_per_second.is_none() {
|
|
prompt_tokens = used;
|
|
} else if !started {
|
|
anthropic_stream_start(stream, &request, id, prompt_tokens, 0)?;
|
|
started = true;
|
|
}
|
|
}
|
|
GenerationEvent::Chunk { reasoning, content } => {
|
|
if !started {
|
|
anthropic_stream_start(stream, &request, id, prompt_tokens, 0)?;
|
|
started = true;
|
|
}
|
|
if !reasoning && request.has_tools {
|
|
let events = projector.push(&content, false, "toolu_");
|
|
send_anthropic_projection_events(
|
|
stream,
|
|
events,
|
|
&mut block,
|
|
&mut next_index,
|
|
&mut tool_indices,
|
|
)?;
|
|
continue;
|
|
}
|
|
if block.is_some_and(|(_, current_reasoning)| current_reasoning != reasoning) {
|
|
let (index, _) = block.take().unwrap();
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_stop",
|
|
&json!({"type": "content_block_stop", "index": index}),
|
|
)?;
|
|
}
|
|
let index = if let Some((index, _)) = block {
|
|
index
|
|
} else {
|
|
let index = next_index;
|
|
next_index += 1;
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_start",
|
|
&json!({
|
|
"type": "content_block_start", "index": index,
|
|
"content_block": if reasoning { json!({"type": "thinking", "thinking": "", "signature": ""}) } else { json!({"type": "text", "text": ""}) }
|
|
}),
|
|
)?;
|
|
block = Some((index, reasoning));
|
|
index
|
|
};
|
|
let delta = if reasoning {
|
|
json!({"type": "thinking_delta", "thinking": content})
|
|
} else {
|
|
json!({"type": "text_delta", "text": content})
|
|
};
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_delta",
|
|
&json!({"type": "content_block_delta", "index": index, "delta": delta}),
|
|
)?;
|
|
}
|
|
GenerationEvent::Finished(Ok(output)) => {
|
|
if !started {
|
|
anthropic_stream_start(
|
|
stream,
|
|
&request,
|
|
id,
|
|
output.prompt_tokens,
|
|
output.cached_tokens,
|
|
)?;
|
|
}
|
|
if request.has_tools {
|
|
let events = projector.finish("toolu_");
|
|
send_anthropic_projection_events(
|
|
stream,
|
|
events,
|
|
&mut block,
|
|
&mut next_index,
|
|
&mut tool_indices,
|
|
)?;
|
|
}
|
|
if let Some((index, _)) = block.take() {
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_stop",
|
|
&json!({"type": "content_block_stop", "index": index}),
|
|
)?;
|
|
}
|
|
let (_, calls) = parse_generated_tools_with_ids(
|
|
state,
|
|
&output.message.content,
|
|
Protocol::Anthropic,
|
|
&projector.ids,
|
|
);
|
|
if projector.ids.is_empty() {
|
|
for call in &calls {
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_start",
|
|
&json!({"type": "content_block_start", "index": next_index, "content_block": {"type": "tool_use", "id": call.id, "name": call.function.name, "input": {}}}),
|
|
)?;
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_delta",
|
|
&json!({"type": "content_block_delta", "index": next_index, "delta": {"type": "input_json_delta", "partial_json": call.function.arguments}}),
|
|
)?;
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_stop",
|
|
&json!({"type": "content_block_stop", "index": next_index}),
|
|
)?;
|
|
next_index += 1;
|
|
}
|
|
}
|
|
let finish = if calls.is_empty() {
|
|
output.finish_reason
|
|
} else {
|
|
"tool_calls"
|
|
};
|
|
send_named_sse(
|
|
stream,
|
|
"message_delta",
|
|
&json!({"type": "message_delta", "delta": {"stop_reason": if finish == "tool_calls" { "tool_use" } else if finish == "length" { "max_tokens" } else { "end_turn" }, "stop_sequence": Value::Null}, "usage": {"output_tokens": output.completion_tokens}}),
|
|
)?;
|
|
return send_named_sse(stream, "message_stop", &json!({"type": "message_stop"}));
|
|
}
|
|
GenerationEvent::Finished(Err(error)) => {
|
|
return send_named_sse(
|
|
stream,
|
|
"error",
|
|
&json!({"type": "error", "error": {"type": "api_error", "message": error}}),
|
|
);
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
Err((500, "The model runtime stopped unexpectedly.".into()))
|
|
}
|
|
|
|
fn send_anthropic_projection_events(
|
|
stream: &mut impl Write,
|
|
events: Vec<ToolProjectionEvent>,
|
|
block: &mut Option<(usize, bool)>,
|
|
next_index: &mut usize,
|
|
tool_indices: &mut Vec<usize>,
|
|
) -> Result<(), (u16, String)> {
|
|
for event in events {
|
|
match event {
|
|
ToolProjectionEvent::Text(text) if !text.is_empty() => {
|
|
if block.is_some_and(|(_, reasoning)| reasoning) {
|
|
let (index, _) = block.take().unwrap();
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_stop",
|
|
&json!({"type": "content_block_stop", "index": index}),
|
|
)?;
|
|
}
|
|
let index = if let Some((index, _)) = *block {
|
|
index
|
|
} else {
|
|
let index = *next_index;
|
|
*next_index += 1;
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_start",
|
|
&json!({"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}}),
|
|
)?;
|
|
*block = Some((index, false));
|
|
index
|
|
};
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_delta",
|
|
&json!({"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}}),
|
|
)?;
|
|
}
|
|
ToolProjectionEvent::Start { index, id, name } => {
|
|
if let Some((open_index, _)) = block.take() {
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_stop",
|
|
&json!({"type": "content_block_stop", "index": open_index}),
|
|
)?;
|
|
}
|
|
let content_index = *next_index;
|
|
*next_index += 1;
|
|
if tool_indices.len() == index {
|
|
tool_indices.push(content_index);
|
|
}
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_start",
|
|
&json!({"type": "content_block_start", "index": content_index, "content_block": {"type": "tool_use", "id": id, "name": name, "input": {}}}),
|
|
)?;
|
|
}
|
|
ToolProjectionEvent::Arguments { index, fragment } => {
|
|
if let Some(content_index) = tool_indices.get(index) {
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_delta",
|
|
&json!({"type": "content_block_delta", "index": content_index, "delta": {"type": "input_json_delta", "partial_json": fragment}}),
|
|
)?;
|
|
}
|
|
}
|
|
ToolProjectionEvent::End { index } => {
|
|
if let Some(content_index) = tool_indices.get(index) {
|
|
send_named_sse(
|
|
stream,
|
|
"content_block_stop",
|
|
&json!({"type": "content_block_stop", "index": content_index}),
|
|
)?;
|
|
}
|
|
}
|
|
ToolProjectionEvent::Text(_) => {}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn anthropic_stream_start(
|
|
stream: &mut impl Write,
|
|
request: &ResponseOptions,
|
|
id: &str,
|
|
prompt_tokens: u32,
|
|
cached_tokens: u32,
|
|
) -> Result<(), (u16, String)> {
|
|
let cached = cached_tokens.min(prompt_tokens);
|
|
let written = prompt_tokens - cached;
|
|
send_named_sse(
|
|
stream,
|
|
"message_start",
|
|
&json!({"type": "message_start", "message": {"id": id, "type": "message", "role": "assistant", "model": request.model_id, "content": [], "stop_reason": Value::Null, "stop_sequence": Value::Null, "usage": {"input_tokens": prompt_tokens - cached - written, "output_tokens": 0, "cache_read_input_tokens": cached, "cache_creation_input_tokens": written}}}),
|
|
)
|
|
}
|
|
|
|
fn responses_stream_response(
|
|
stream: &mut TcpStream,
|
|
state: &State,
|
|
request: ResponseOptions,
|
|
active: crate::runtime::ActiveGeneration,
|
|
id: &str,
|
|
) -> Result<(), (u16, String)> {
|
|
send_sse_headers_with_cors(stream, request.cors).map_err(|error| (500, error))?;
|
|
let created = unix_time();
|
|
let message_id = random_id("msg_");
|
|
let reasoning_id = random_id("rs_");
|
|
let mut sequence = 0;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.created", "response": {"id": id, "object": "response", "created_at": created, "status": "in_progress", "model": request.model_id, "output": []}}),
|
|
)?;
|
|
let mut reasoning_open = false;
|
|
let mut message_open = false;
|
|
let mut reasoning = String::new();
|
|
let mut content = String::new();
|
|
while let Ok(event) = active.events.recv() {
|
|
match event {
|
|
GenerationEvent::Chunk {
|
|
reasoning: true,
|
|
content: chunk,
|
|
} => {
|
|
if !request.reasoning_summary {
|
|
continue;
|
|
}
|
|
if !reasoning_open {
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_item.added", "output_index": 0, "item": {"id": reasoning_id, "type": "reasoning", "status": "in_progress", "summary": []}}),
|
|
)?;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.reasoning_summary_part.added", "item_id": reasoning_id, "output_index": 0, "summary_index": 0, "part": {"type": "summary_text", "text": ""}}),
|
|
)?;
|
|
reasoning_open = true;
|
|
}
|
|
reasoning.push_str(&chunk);
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.reasoning_summary_text.delta", "item_id": reasoning_id, "output_index": 0, "summary_index": 0, "delta": chunk}),
|
|
)?;
|
|
}
|
|
GenerationEvent::Chunk {
|
|
reasoning: false,
|
|
content: chunk,
|
|
} => {
|
|
if request.has_tools {
|
|
content.push_str(&chunk);
|
|
continue;
|
|
}
|
|
if !message_open {
|
|
let output_index = usize::from(reasoning_open);
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_item.added", "output_index": output_index, "item": {"id": message_id, "type": "message", "status": "in_progress", "role": "assistant", "content": []}}),
|
|
)?;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.content_part.added", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": []}}),
|
|
)?;
|
|
message_open = true;
|
|
}
|
|
content.push_str(&chunk);
|
|
let output_index = usize::from(reasoning_open);
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_text.delta", "item_id": message_id, "output_index": output_index, "content_index": 0, "delta": chunk}),
|
|
)?;
|
|
}
|
|
GenerationEvent::Finished(Ok(output)) => {
|
|
let (parsed_content, calls) =
|
|
parse_generated_tools(state, &output.message.content, Protocol::Responses);
|
|
if request.has_tools {
|
|
content = parsed_content;
|
|
}
|
|
let finish = if calls.is_empty() {
|
|
output.finish_reason
|
|
} else {
|
|
"tool_calls"
|
|
};
|
|
let status = if finish == "length" {
|
|
"incomplete"
|
|
} else if finish == "error" {
|
|
"failed"
|
|
} else {
|
|
"completed"
|
|
};
|
|
let mut terminal_items = Vec::new();
|
|
let mut output_index = 0;
|
|
if reasoning_open {
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.reasoning_summary_text.done", "item_id": reasoning_id, "output_index": output_index, "summary_index": 0, "text": reasoning}),
|
|
)?;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.reasoning_summary_part.done", "item_id": reasoning_id, "output_index": output_index, "summary_index": 0, "part": {"type": "summary_text", "text": reasoning}}),
|
|
)?;
|
|
let item = json!({"id": reasoning_id, "type": "reasoning", "status": status, "summary": [{"type": "summary_text", "text": reasoning}]});
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_item.done", "output_index": output_index, "item": item}),
|
|
)?;
|
|
terminal_items.push(item);
|
|
output_index += 1;
|
|
}
|
|
if !content.is_empty() {
|
|
if !message_open {
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_item.added", "output_index": output_index, "item": {"id": message_id, "type": "message", "status": "in_progress", "role": "assistant", "content": []}}),
|
|
)?;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.content_part.added", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": []}}),
|
|
)?;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_text.delta", "item_id": message_id, "output_index": output_index, "content_index": 0, "delta": content}),
|
|
)?;
|
|
}
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_text.done", "item_id": message_id, "output_index": output_index, "content_index": 0, "text": content}),
|
|
)?;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.content_part.done", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": content, "annotations": []}}),
|
|
)?;
|
|
let item = json!({"id": message_id, "type": "message", "status": status, "role": "assistant", "content": [{"type": "output_text", "text": content, "annotations": []}]});
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_item.done", "output_index": output_index, "item": item}),
|
|
)?;
|
|
terminal_items.push(item);
|
|
output_index += 1;
|
|
}
|
|
for call in &calls {
|
|
let item_id = random_id("fc_");
|
|
let mut item = json!({"id": item_id, "type": "function_call", "status": status, "name": call.function.name, "call_id": call.id, "arguments": call.function.arguments});
|
|
let mut added = item.clone();
|
|
added["status"] = Value::String("in_progress".into());
|
|
added["arguments"] = Value::String(String::new());
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_item.added", "output_index": output_index, "item": added}),
|
|
)?;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.function_call_arguments.delta", "item_id": item_id, "output_index": output_index, "delta": call.function.arguments}),
|
|
)?;
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.function_call_arguments.done", "item_id": item_id, "output_index": output_index, "name": call.function.name, "arguments": call.function.arguments}),
|
|
)?;
|
|
item["id"] = Value::String(item_id);
|
|
send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": "response.output_item.done", "output_index": output_index, "item": item}),
|
|
)?;
|
|
terminal_items.push(item);
|
|
output_index += 1;
|
|
}
|
|
let event_type = if finish == "length" {
|
|
"response.incomplete"
|
|
} else if finish == "error" {
|
|
"response.failed"
|
|
} else {
|
|
"response.completed"
|
|
};
|
|
let cached = output.cached_tokens.min(output.prompt_tokens);
|
|
return send_responses_sse(
|
|
stream,
|
|
&mut sequence,
|
|
json!({"type": event_type, "response": {"id": id, "object": "response", "created_at": created, "status": status, "model": request.model_id, "output": terminal_items, "usage": {"input_tokens": output.prompt_tokens, "input_tokens_details": {"cached_tokens": cached, "cache_write_tokens": output.prompt_tokens - cached}, "output_tokens": output.completion_tokens, "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": output.prompt_tokens + output.completion_tokens}}}),
|
|
);
|
|
}
|
|
GenerationEvent::Finished(Err(error)) => {
|
|
let _ = send_sse_error(stream, &error);
|
|
return Ok(());
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
Err((500, "The model runtime stopped unexpectedly.".into()))
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn send_named_sse(
|
|
stream: &mut impl Write,
|
|
event: &str,
|
|
value: &Value,
|
|
) -> Result<(), (u16, String)> {
|
|
let body = serde_json::to_vec(value).map_err(|error| (500, error.to_string()))?;
|
|
stream
|
|
.write_all(b"event: ")
|
|
.and_then(|()| stream.write_all(event.as_bytes()))
|
|
.and_then(|()| stream.write_all(b"\ndata: "))
|
|
.and_then(|()| stream.write_all(&body))
|
|
.and_then(|()| stream.write_all(b"\n\n"))
|
|
.map_err(|error| (500, error.to_string()))
|
|
}
|
|
|
|
fn send_responses_sse(
|
|
stream: &mut impl Write,
|
|
sequence: &mut u32,
|
|
value: Value,
|
|
) -> Result<(), (u16, String)> {
|
|
let mut object = value
|
|
.as_object()
|
|
.cloned()
|
|
.ok_or_else(|| (500, "Responses event is not an object".to_owned()))?;
|
|
let event_type = object.shift_remove("type").unwrap_or(Value::Null);
|
|
let mut ordered = Map::new();
|
|
ordered.insert("type".into(), event_type);
|
|
ordered.insert("sequence_number".into(), Value::from(*sequence));
|
|
ordered.extend(object);
|
|
*sequence += 1;
|
|
send_sse(stream, &Value::Object(ordered)).map_err(|error| (500, error))
|
|
}
|
|
|
|
fn stream_response_with_keepalive(
|
|
stream: &mut TcpStream,
|
|
state: &State,
|
|
request: ResponseOptions,
|
|
active: crate::runtime::ActiveGeneration,
|
|
id: &str,
|
|
keepalive_interval: Duration,
|
|
) -> Result<(), (u16, String)> {
|
|
let mut projector = ToolProjector::new();
|
|
let mut output = None;
|
|
let mut prefilling = true;
|
|
let mut headers_sent = false;
|
|
let mut role_sent = false;
|
|
let mut last_keepalive = Instant::now();
|
|
loop {
|
|
let event = match receive_stream_event(
|
|
stream,
|
|
&active,
|
|
prefilling && headers_sent,
|
|
&mut last_keepalive,
|
|
keepalive_interval,
|
|
) {
|
|
Ok(Some(event)) => event,
|
|
Ok(None) => break,
|
|
Err(_) => return Ok(()),
|
|
};
|
|
match event {
|
|
GenerationEvent::Compacted(_) => {
|
|
return Err((
|
|
500,
|
|
"The model runtime returned an unexpected compaction event.".into(),
|
|
));
|
|
}
|
|
GenerationEvent::Activity(_) | GenerationEvent::Measured(_) => {}
|
|
GenerationEvent::Loading => {}
|
|
GenerationEvent::Context {
|
|
tokens_per_second, ..
|
|
} => {
|
|
prefilling = tokens_per_second.is_none();
|
|
if prefilling && !headers_sent {
|
|
send_sse_headers_with_cors(stream, request.cors)
|
|
.map_err(|error| (500, error))?;
|
|
headers_sent = true;
|
|
last_keepalive = Instant::now();
|
|
} else if !prefilling {
|
|
send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?;
|
|
}
|
|
}
|
|
GenerationEvent::Chunk { reasoning, content } => {
|
|
prefilling = false;
|
|
send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?;
|
|
if request.has_tools && !reasoning {
|
|
let events = projector.push(&content, false, "call_");
|
|
if send_chat_projection_events(stream, &request, id, events).is_err() {
|
|
active.cancel.store(true, Ordering::Relaxed);
|
|
return Ok(());
|
|
}
|
|
if TOOL_SYNTAXES
|
|
.iter()
|
|
.any(|syntax| projector.raw.contains(syntax.tool_end))
|
|
{
|
|
active.cancel.store(true, Ordering::Relaxed);
|
|
}
|
|
} else if !content.is_empty() {
|
|
let field = if reasoning {
|
|
"reasoning_content"
|
|
} else {
|
|
"content"
|
|
};
|
|
let chunk = chunk_json(id, &request.model_id, json!({field: content}), None);
|
|
if send_sse(stream, &chunk).is_err() {
|
|
active.cancel.store(true, Ordering::Relaxed);
|
|
return Ok(());
|
|
}
|
|
}
|
|
}
|
|
GenerationEvent::Finished(result) => {
|
|
match result {
|
|
Ok(result) => {
|
|
send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?;
|
|
output = Some(result);
|
|
}
|
|
Err(error) => {
|
|
if headers_sent {
|
|
let _ = send_sse_error(stream, &error);
|
|
return Ok(());
|
|
}
|
|
return send_generation_error(
|
|
stream,
|
|
request.protocol,
|
|
request.cors,
|
|
error,
|
|
);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
let output = match output {
|
|
Some(output) => output,
|
|
None if headers_sent => {
|
|
let _ = send_sse_error(stream, "The model runtime stopped unexpectedly.");
|
|
return Ok(());
|
|
}
|
|
None => return Err((500, "The model runtime stopped unexpectedly.".into())),
|
|
};
|
|
if request.has_tools {
|
|
let events = projector.finish("call_");
|
|
send_chat_projection_events(stream, &request, id, events)?;
|
|
}
|
|
let (_, calls) = parse_generated_tools_with_ids(
|
|
state,
|
|
&output.message.content,
|
|
Protocol::Chat,
|
|
&projector.ids,
|
|
);
|
|
if request.has_tools && projector.ids.is_empty() && !calls.is_empty() {
|
|
send_sse(
|
|
stream,
|
|
&chunk_json(
|
|
id,
|
|
&request.model_id,
|
|
json!({"tool_calls": tool_calls_json(&calls)}),
|
|
None,
|
|
),
|
|
)
|
|
.map_err(|error| (500, error))?;
|
|
}
|
|
let finish = if calls.is_empty() {
|
|
output.finish_reason
|
|
} else {
|
|
"tool_calls"
|
|
};
|
|
send_sse(
|
|
stream,
|
|
&chunk_json(id, &request.model_id, json!({}), Some(finish)),
|
|
)
|
|
.map_err(|error| (500, error))?;
|
|
if request.include_usage {
|
|
let usage = json!({
|
|
"id": id,
|
|
"object": "chat.completion.chunk",
|
|
"created": unix_time(),
|
|
"model": request.model_id,
|
|
"choices": [],
|
|
"usage": usage_json(output.prompt_tokens, output.cached_tokens, output.completion_tokens),
|
|
});
|
|
send_sse(stream, &usage).map_err(|error| (500, error))?;
|
|
}
|
|
stream
|
|
.write_all(b"data: [DONE]\n\n")
|
|
.map_err(|error| (500, error.to_string()))
|
|
}
|
|
|
|
fn send_chat_projection_events(
|
|
stream: &mut impl Write,
|
|
request: &ResponseOptions,
|
|
response_id: &str,
|
|
events: Vec<ToolProjectionEvent>,
|
|
) -> Result<(), (u16, String)> {
|
|
for event in events {
|
|
let delta = match event {
|
|
ToolProjectionEvent::Text(content) if !content.is_empty() => {
|
|
json!({"content": content})
|
|
}
|
|
ToolProjectionEvent::Start { index, id, name } => json!({"tool_calls": [{
|
|
"index": index, "id": id, "type": "function",
|
|
"function": {"name": name, "arguments": ""}
|
|
}]}),
|
|
ToolProjectionEvent::Arguments { index, fragment } => json!({
|
|
"tool_calls": [{"index": index, "function": {"arguments": fragment}}]
|
|
}),
|
|
ToolProjectionEvent::End { .. } | ToolProjectionEvent::Text(_) => continue,
|
|
};
|
|
send_sse(
|
|
stream,
|
|
&chunk_json(response_id, &request.model_id, delta, None),
|
|
)
|
|
.map_err(|error| (500, error))?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn send_stream_start(
|
|
stream: &mut impl Write,
|
|
request: &ResponseOptions,
|
|
id: &str,
|
|
headers_sent: &mut bool,
|
|
role_sent: &mut bool,
|
|
) -> Result<(), (u16, String)> {
|
|
if !*headers_sent {
|
|
send_sse_headers_with_cors(stream, request.cors).map_err(|error| (500, error))?;
|
|
*headers_sent = true;
|
|
}
|
|
if !*role_sent {
|
|
let role = chunk_json(id, &request.model_id, json!({"role": "assistant"}), None);
|
|
send_sse(stream, &role).map_err(|error| (500, error))?;
|
|
*role_sent = true;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn receive_stream_event(
|
|
stream: &mut impl Write,
|
|
active: &crate::runtime::ActiveGeneration,
|
|
prefilling: bool,
|
|
last_keepalive: &mut Instant,
|
|
keepalive_interval: Duration,
|
|
) -> Result<Option<GenerationEvent>, String> {
|
|
if !prefilling {
|
|
return Ok(active.events.recv().ok());
|
|
}
|
|
loop {
|
|
if last_keepalive.elapsed() >= keepalive_interval {
|
|
if let Err(error) = stream.write_all(b": prefill\n\n") {
|
|
active.cancel.store(true, Ordering::Relaxed);
|
|
return Err(error.to_string());
|
|
}
|
|
*last_keepalive = Instant::now();
|
|
}
|
|
let remaining = keepalive_interval.saturating_sub(last_keepalive.elapsed());
|
|
match active.events.recv_timeout(remaining) {
|
|
Ok(event) => return Ok(Some(event)),
|
|
Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
|
|
Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => return Ok(None),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn wait_for_output(
|
|
stream: &mut TcpStream,
|
|
active: crate::runtime::ActiveGeneration,
|
|
) -> Result<crate::engine::GenerationOutput, String> {
|
|
let mut content = String::new();
|
|
stream
|
|
.set_nonblocking(true)
|
|
.map_err(|error| error.to_string())?;
|
|
loop {
|
|
let event = match active.events.recv_timeout(Duration::from_millis(100)) {
|
|
Ok(event) => event,
|
|
Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
|
|
let mut byte = [0];
|
|
match stream.peek(&mut byte) {
|
|
Ok(0) => return Err("client disconnected".into()),
|
|
Ok(_) => {}
|
|
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {}
|
|
Err(error) => return Err(error.to_string()),
|
|
}
|
|
continue;
|
|
}
|
|
Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break,
|
|
};
|
|
match event {
|
|
GenerationEvent::Chunk {
|
|
reasoning: false,
|
|
content: chunk,
|
|
} => {
|
|
content.push_str(&chunk);
|
|
if TOOL_SYNTAXES
|
|
.iter()
|
|
.any(|syntax| content.contains(syntax.tool_end))
|
|
{
|
|
active.cancel.store(true, Ordering::Relaxed);
|
|
}
|
|
}
|
|
GenerationEvent::Finished(result) => {
|
|
stream
|
|
.set_nonblocking(false)
|
|
.map_err(|error| error.to_string())?;
|
|
return result;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
let _ = stream.set_nonblocking(false);
|
|
Err("The model runtime stopped unexpectedly.".into())
|
|
}
|
|
|
|
pub(super) fn send_generation_error(
|
|
stream: &mut TcpStream,
|
|
protocol: Protocol,
|
|
cors: bool,
|
|
error: String,
|
|
) -> Result<(), (u16, String)> {
|
|
let Some((prompt_tokens, context)) = context_error_dimensions(&error) else {
|
|
return Err((500, error));
|
|
};
|
|
let body = if protocol == Protocol::Anthropic {
|
|
json!({
|
|
"type": "error",
|
|
"error": {
|
|
"type": "invalid_request_error",
|
|
"message": error,
|
|
"n_prompt_tokens": prompt_tokens,
|
|
"n_ctx": context
|
|
}
|
|
})
|
|
} else {
|
|
let parameter = match protocol {
|
|
Protocol::Completion => "prompt",
|
|
Protocol::Responses => "input",
|
|
Protocol::Chat => "messages",
|
|
Protocol::Anthropic => unreachable!(),
|
|
};
|
|
json!({"error": {
|
|
"message": error,
|
|
"type": "invalid_request_error",
|
|
"param": parameter,
|
|
"code": "context_length_exceeded",
|
|
"n_prompt_tokens": prompt_tokens,
|
|
"n_ctx": context
|
|
}})
|
|
};
|
|
send_json_with_cors(stream, 400, &body, cors).map_err(|error| (500, error))
|
|
}
|
|
|
|
pub(super) fn context_error_dimensions(error: &str) -> Option<(u32, u32)> {
|
|
let values = error
|
|
.strip_prefix("Prompt has ")?
|
|
.strip_suffix(" tokens")?
|
|
.split_once(" tokens, but the configured context size is ")?;
|
|
Some((values.0.parse().ok()?, values.1.parse().ok()?))
|
|
}
|
|
|
|
pub(super) fn usage_json(prompt: u32, cached: u32, completion: u32) -> Value {
|
|
let cached = cached.min(prompt);
|
|
json!({
|
|
"prompt_tokens": prompt,
|
|
"completion_tokens": completion,
|
|
"total_tokens": prompt + completion,
|
|
"prompt_tokens_details": {
|
|
"cached_tokens": cached,
|
|
"cache_write_tokens": prompt - cached
|
|
}
|
|
})
|
|
}
|
|
|
|
pub(super) fn chunk_json(id: &str, model: &str, delta: Value, finish: Option<&str>) -> Value {
|
|
json!({
|
|
"id": id,
|
|
"object": "chat.completion.chunk",
|
|
"created": unix_time(),
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}]
|
|
})
|
|
}
|