Split Rust code into domain modules
This commit is contained in:
Georg Bauer
974
src/server/response.rs
Normal file
974
src/server/response.rs
Normal file
@@ -0,0 +1,974 @@
|
||||
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 = wait_for_output(active)?;
|
||||
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(stream, 200, &body).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(stream).map_err(|error| (500, error))?;
|
||||
while let Ok(event) = active.events.recv() {
|
||||
match event {
|
||||
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(stream).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.push("", true, "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(stream).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::Loading => {}
|
||||
GenerationEvent::Context {
|
||||
tokens_per_second, ..
|
||||
} => {
|
||||
prefilling = tokens_per_second.is_none();
|
||||
if prefilling && !headers_sent {
|
||||
send_sse_headers(stream).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 Err((500, 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.push("", true, "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(stream).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(
|
||||
active: crate::runtime::ActiveGeneration,
|
||||
) -> Result<crate::engine::GenerationOutput, (u16, String)> {
|
||||
let mut content = String::new();
|
||||
while let Ok(event) = active.events.recv() {
|
||||
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) => {
|
||||
return result.map_err(|error| (500, error));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err((500, "The model runtime stopped unexpectedly.".into()))
|
||||
}
|
||||
|
||||
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}]
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user