mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
527 lines
19 KiB
Rust
527 lines
19 KiB
Rust
|
use std::collections::BTreeMap;
|
||
|
|
|
||
|
|
use serde_json::{json, Map, Value};
|
||
|
|
|
||
|
|
use crate::finalize::common::{build_generated_tool_call_id, canonicalize_tool_arguments};
|
||
|
|
use crate::finalize::sse::{encode_json_sse, map_claude_stop_reason};
|
||
|
|
use crate::finalize::standard::stream_core::common::*;
|
||
|
|
use crate::finalize::PipelineFinalizeError;
|
||
|
|
|
||
|
|
#[derive(Default)]
|
||
|
|
struct ClaudeProviderToolState {
|
||
|
|
call_id: String,
|
||
|
|
name: String,
|
||
|
|
started_emitted: bool,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Default)]
|
||
|
|
pub struct ClaudeProviderState {
|
||
|
|
message_id: Option<String>,
|
||
|
|
model: Option<String>,
|
||
|
|
started: bool,
|
||
|
|
finished: bool,
|
||
|
|
tool_calls: BTreeMap<usize, ClaudeProviderToolState>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl ClaudeProviderState {
|
||
|
|
fn identity(&self, report_context: &Value) -> (String, String) {
|
||
|
|
resolve_identity(
|
||
|
|
self.message_id.as_deref(),
|
||
|
|
self.model.as_deref(),
|
||
|
|
report_context,
|
||
|
|
"msg-local-stream",
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn ensure_started(&mut self, report_context: &Value, out: &mut Vec<CanonicalStreamFrame>) {
|
||
|
|
if self.started {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
let (id, model) = self.identity(report_context);
|
||
|
|
out.push(CanonicalStreamFrame {
|
||
|
|
id,
|
||
|
|
model,
|
||
|
|
event: CanonicalStreamEvent::Start,
|
||
|
|
});
|
||
|
|
self.started = true;
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn push_line(
|
||
|
|
&mut self,
|
||
|
|
report_context: &Value,
|
||
|
|
line: Vec<u8>,
|
||
|
|
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
|
||
|
|
let Some(value) = decode_json_data_line(&line) else {
|
||
|
|
return Ok(Vec::new());
|
||
|
|
};
|
||
|
|
let Some(event_object) = value.as_object() else {
|
||
|
|
return Ok(Vec::new());
|
||
|
|
};
|
||
|
|
let mut out = Vec::new();
|
||
|
|
match event_object
|
||
|
|
.get("type")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.unwrap_or_default()
|
||
|
|
{
|
||
|
|
"message_start" => {
|
||
|
|
if let Some(message) = event_object.get("message").and_then(Value::as_object) {
|
||
|
|
self.message_id = message
|
||
|
|
.get("id")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.map(ToOwned::to_owned);
|
||
|
|
self.model = message
|
||
|
|
.get("model")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.map(ToOwned::to_owned);
|
||
|
|
}
|
||
|
|
self.ensure_started(report_context, &mut out);
|
||
|
|
}
|
||
|
|
"content_block_delta" => {
|
||
|
|
let index = event_object
|
||
|
|
.get("index")
|
||
|
|
.and_then(Value::as_u64)
|
||
|
|
.map(|value| value as usize)
|
||
|
|
.unwrap_or(0);
|
||
|
|
let Some(delta) = event_object.get("delta").and_then(Value::as_object) else {
|
||
|
|
return Ok(out);
|
||
|
|
};
|
||
|
|
match delta
|
||
|
|
.get("type")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.unwrap_or_default()
|
||
|
|
{
|
||
|
|
"text_delta" => {
|
||
|
|
let Some(piece) = delta.get("text").and_then(Value::as_str) else {
|
||
|
|
return Ok(out);
|
||
|
|
};
|
||
|
|
if piece.is_empty() {
|
||
|
|
return Ok(out);
|
||
|
|
}
|
||
|
|
self.ensure_started(report_context, &mut out);
|
||
|
|
let (id, model) = self.identity(report_context);
|
||
|
|
out.push(CanonicalStreamFrame {
|
||
|
|
id,
|
||
|
|
model,
|
||
|
|
event: CanonicalStreamEvent::TextDelta(piece.to_string()),
|
||
|
|
});
|
||
|
|
}
|
||
|
|
"input_json_delta" => {
|
||
|
|
let Some(partial_json) = delta.get("partial_json").and_then(Value::as_str)
|
||
|
|
else {
|
||
|
|
return Ok(out);
|
||
|
|
};
|
||
|
|
if partial_json.is_empty() {
|
||
|
|
return Ok(out);
|
||
|
|
}
|
||
|
|
self.ensure_started(report_context, &mut out);
|
||
|
|
let (id, model) = self.identity(report_context);
|
||
|
|
let tool_state = self.tool_calls.entry(index).or_default();
|
||
|
|
if !tool_state.started_emitted {
|
||
|
|
out.push(CanonicalStreamFrame {
|
||
|
|
id: id.clone(),
|
||
|
|
model: model.clone(),
|
||
|
|
event: CanonicalStreamEvent::ToolCallStart {
|
||
|
|
index,
|
||
|
|
call_id: if tool_state.call_id.is_empty() {
|
||
|
|
build_generated_tool_call_id(index)
|
||
|
|
} else {
|
||
|
|
tool_state.call_id.clone()
|
||
|
|
},
|
||
|
|
name: if tool_state.name.is_empty() {
|
||
|
|
"unknown".to_string()
|
||
|
|
} else {
|
||
|
|
tool_state.name.clone()
|
||
|
|
},
|
||
|
|
},
|
||
|
|
});
|
||
|
|
tool_state.started_emitted = true;
|
||
|
|
}
|
||
|
|
out.push(CanonicalStreamFrame {
|
||
|
|
id,
|
||
|
|
model,
|
||
|
|
event: CanonicalStreamEvent::ToolCallArgumentsDelta {
|
||
|
|
index,
|
||
|
|
arguments: partial_json.to_string(),
|
||
|
|
},
|
||
|
|
});
|
||
|
|
}
|
||
|
|
_ => {}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
"content_block_start" => {
|
||
|
|
let index = event_object
|
||
|
|
.get("index")
|
||
|
|
.and_then(Value::as_u64)
|
||
|
|
.map(|value| value as usize)
|
||
|
|
.unwrap_or(0);
|
||
|
|
let Some(block) = event_object.get("content_block").and_then(Value::as_object)
|
||
|
|
else {
|
||
|
|
return Ok(out);
|
||
|
|
};
|
||
|
|
let block_type = block
|
||
|
|
.get("type")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.unwrap_or_default();
|
||
|
|
if block_type == "text" {
|
||
|
|
let Some(text) = block.get("text").and_then(Value::as_str) else {
|
||
|
|
return Ok(out);
|
||
|
|
};
|
||
|
|
if text.is_empty() {
|
||
|
|
return Ok(out);
|
||
|
|
}
|
||
|
|
self.ensure_started(report_context, &mut out);
|
||
|
|
let (id, model) = self.identity(report_context);
|
||
|
|
out.push(CanonicalStreamFrame {
|
||
|
|
id,
|
||
|
|
model,
|
||
|
|
event: CanonicalStreamEvent::TextDelta(text.to_string()),
|
||
|
|
});
|
||
|
|
return Ok(out);
|
||
|
|
}
|
||
|
|
if block_type != "tool_use" {
|
||
|
|
return Ok(out);
|
||
|
|
}
|
||
|
|
self.ensure_started(report_context, &mut out);
|
||
|
|
let (id, model) = self.identity(report_context);
|
||
|
|
let tool_state = self.tool_calls.entry(index).or_default();
|
||
|
|
tool_state.call_id = block
|
||
|
|
.get("id")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.unwrap_or_else(|| tool_state.call_id.as_str())
|
||
|
|
.to_string();
|
||
|
|
tool_state.name = block
|
||
|
|
.get("name")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.unwrap_or_else(|| tool_state.name.as_str())
|
||
|
|
.to_string();
|
||
|
|
if !tool_state.started_emitted {
|
||
|
|
out.push(CanonicalStreamFrame {
|
||
|
|
id: id.clone(),
|
||
|
|
model: model.clone(),
|
||
|
|
event: CanonicalStreamEvent::ToolCallStart {
|
||
|
|
index,
|
||
|
|
call_id: if tool_state.call_id.is_empty() {
|
||
|
|
build_generated_tool_call_id(index)
|
||
|
|
} else {
|
||
|
|
tool_state.call_id.clone()
|
||
|
|
},
|
||
|
|
name: if tool_state.name.is_empty() {
|
||
|
|
"unknown".to_string()
|
||
|
|
} else {
|
||
|
|
tool_state.name.clone()
|
||
|
|
},
|
||
|
|
},
|
||
|
|
});
|
||
|
|
tool_state.started_emitted = true;
|
||
|
|
}
|
||
|
|
let arguments = canonicalize_tool_arguments(block.get("input").cloned());
|
||
|
|
if !arguments.is_empty() && arguments != "{}" {
|
||
|
|
out.push(CanonicalStreamFrame {
|
||
|
|
id,
|
||
|
|
model,
|
||
|
|
event: CanonicalStreamEvent::ToolCallArgumentsDelta { index, arguments },
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
"message_delta" => {
|
||
|
|
self.ensure_started(report_context, &mut out);
|
||
|
|
let Some(delta) = event_object.get("delta").and_then(Value::as_object) else {
|
||
|
|
return Ok(out);
|
||
|
|
};
|
||
|
|
let finish_reason = map_claude_stop_reason(
|
||
|
|
delta.get("stop_reason").and_then(Value::as_str),
|
||
|
|
delta.get("stop_reason").and_then(Value::as_str) == Some("tool_use"),
|
||
|
|
)
|
||
|
|
.map(ToOwned::to_owned);
|
||
|
|
let (id, model) = self.identity(report_context);
|
||
|
|
out.push(CanonicalStreamFrame {
|
||
|
|
id,
|
||
|
|
model,
|
||
|
|
event: CanonicalStreamEvent::Finish {
|
||
|
|
finish_reason,
|
||
|
|
usage: canonical_usage_from_claude_usage(event_object.get("usage")),
|
||
|
|
},
|
||
|
|
});
|
||
|
|
self.finished = true;
|
||
|
|
}
|
||
|
|
_ => {}
|
||
|
|
}
|
||
|
|
Ok(out)
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn finish(
|
||
|
|
&mut self,
|
||
|
|
report_context: &Value,
|
||
|
|
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
|
||
|
|
if !self.started || self.finished {
|
||
|
|
return Ok(Vec::new());
|
||
|
|
}
|
||
|
|
self.finished = true;
|
||
|
|
let (id, model) = self.identity(report_context);
|
||
|
|
Ok(vec![CanonicalStreamFrame {
|
||
|
|
id,
|
||
|
|
model,
|
||
|
|
event: CanonicalStreamEvent::Finish {
|
||
|
|
finish_reason: None,
|
||
|
|
usage: None,
|
||
|
|
},
|
||
|
|
}])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
enum ClaudeOpenBlock {
|
||
|
|
Text {
|
||
|
|
block_index: usize,
|
||
|
|
},
|
||
|
|
Tool {
|
||
|
|
tool_index: usize,
|
||
|
|
block_index: usize,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Default)]
|
||
|
|
pub struct ClaudeClientEmitter {
|
||
|
|
message_id: Option<String>,
|
||
|
|
model: Option<String>,
|
||
|
|
started: bool,
|
||
|
|
finished: bool,
|
||
|
|
next_block_index: usize,
|
||
|
|
open_block: Option<ClaudeOpenBlock>,
|
||
|
|
tool_block_indices: BTreeMap<usize, usize>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl ClaudeClientEmitter {
|
||
|
|
fn update_identity(&mut self, frame: &CanonicalStreamFrame) {
|
||
|
|
self.message_id = Some(frame.id.clone());
|
||
|
|
self.model = Some(frame.model.clone());
|
||
|
|
}
|
||
|
|
|
||
|
|
fn ensure_started(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
|
||
|
|
if self.started {
|
||
|
|
return Ok(Vec::new());
|
||
|
|
}
|
||
|
|
self.started = true;
|
||
|
|
encode_json_sse(
|
||
|
|
Some("message_start"),
|
||
|
|
&json!({
|
||
|
|
"type": "message_start",
|
||
|
|
"message": {
|
||
|
|
"id": self.message_id.as_deref().unwrap_or("msg-local-stream"),
|
||
|
|
"type": "message",
|
||
|
|
"role": "assistant",
|
||
|
|
"model": self.model.as_deref().unwrap_or("unknown"),
|
||
|
|
"content": [],
|
||
|
|
"stop_reason": Value::Null,
|
||
|
|
"stop_sequence": Value::Null,
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn close_open_block(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
|
||
|
|
let Some(open_block) = self.open_block.take() else {
|
||
|
|
return Ok(Vec::new());
|
||
|
|
};
|
||
|
|
let block_index = match open_block {
|
||
|
|
ClaudeOpenBlock::Text { block_index } => block_index,
|
||
|
|
ClaudeOpenBlock::Tool { block_index, .. } => block_index,
|
||
|
|
};
|
||
|
|
encode_json_sse(
|
||
|
|
Some("content_block_stop"),
|
||
|
|
&json!({
|
||
|
|
"type": "content_block_stop",
|
||
|
|
"index": block_index,
|
||
|
|
}),
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn ensure_text_block(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
|
||
|
|
let mut out = Vec::new();
|
||
|
|
if let Some(ClaudeOpenBlock::Text { .. }) = self.open_block {
|
||
|
|
return Ok(out);
|
||
|
|
}
|
||
|
|
out.extend(self.close_open_block()?);
|
||
|
|
let block_index = self.next_block_index;
|
||
|
|
self.next_block_index += 1;
|
||
|
|
self.open_block = Some(ClaudeOpenBlock::Text { block_index });
|
||
|
|
out.extend(encode_json_sse(
|
||
|
|
Some("content_block_start"),
|
||
|
|
&json!({
|
||
|
|
"type": "content_block_start",
|
||
|
|
"index": block_index,
|
||
|
|
"content_block": {
|
||
|
|
"type": "text",
|
||
|
|
"text": "",
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
)?);
|
||
|
|
Ok(out)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn ensure_tool_block(
|
||
|
|
&mut self,
|
||
|
|
tool_index: usize,
|
||
|
|
call_id: &str,
|
||
|
|
name: &str,
|
||
|
|
) -> Result<Vec<u8>, PipelineFinalizeError> {
|
||
|
|
let mut out = Vec::new();
|
||
|
|
if let Some(ClaudeOpenBlock::Tool {
|
||
|
|
tool_index: current_tool_index,
|
||
|
|
..
|
||
|
|
}) = self.open_block
|
||
|
|
{
|
||
|
|
if current_tool_index == tool_index {
|
||
|
|
return Ok(out);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
out.extend(self.close_open_block()?);
|
||
|
|
let block_index = self
|
||
|
|
.tool_block_indices
|
||
|
|
.get(&tool_index)
|
||
|
|
.copied()
|
||
|
|
.unwrap_or_else(|| {
|
||
|
|
let block_index = self.next_block_index;
|
||
|
|
self.next_block_index += 1;
|
||
|
|
self.tool_block_indices.insert(tool_index, block_index);
|
||
|
|
block_index
|
||
|
|
});
|
||
|
|
self.open_block = Some(ClaudeOpenBlock::Tool {
|
||
|
|
tool_index,
|
||
|
|
block_index,
|
||
|
|
});
|
||
|
|
out.extend(encode_json_sse(
|
||
|
|
Some("content_block_start"),
|
||
|
|
&json!({
|
||
|
|
"type": "content_block_start",
|
||
|
|
"index": block_index,
|
||
|
|
"content_block": {
|
||
|
|
"type": "tool_use",
|
||
|
|
"id": call_id,
|
||
|
|
"name": name,
|
||
|
|
"input": {},
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
)?);
|
||
|
|
Ok(out)
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, PipelineFinalizeError> {
|
||
|
|
self.update_identity(&frame);
|
||
|
|
match frame.event {
|
||
|
|
CanonicalStreamEvent::Start => self.ensure_started(),
|
||
|
|
CanonicalStreamEvent::TextDelta(text) => {
|
||
|
|
let mut out = self.ensure_started()?;
|
||
|
|
out.extend(self.ensure_text_block()?);
|
||
|
|
let block_index = match self.open_block {
|
||
|
|
Some(ClaudeOpenBlock::Text { block_index }) => block_index,
|
||
|
|
_ => return Ok(out),
|
||
|
|
};
|
||
|
|
out.extend(encode_json_sse(
|
||
|
|
Some("content_block_delta"),
|
||
|
|
&json!({
|
||
|
|
"type": "content_block_delta",
|
||
|
|
"index": block_index,
|
||
|
|
"delta": {
|
||
|
|
"type": "text_delta",
|
||
|
|
"text": text,
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
)?);
|
||
|
|
Ok(out)
|
||
|
|
}
|
||
|
|
CanonicalStreamEvent::ToolCallStart {
|
||
|
|
index,
|
||
|
|
call_id,
|
||
|
|
name,
|
||
|
|
} => {
|
||
|
|
let mut out = self.ensure_started()?;
|
||
|
|
out.extend(self.ensure_tool_block(index, &call_id, &name)?);
|
||
|
|
Ok(out)
|
||
|
|
}
|
||
|
|
CanonicalStreamEvent::ToolCallArgumentsDelta { index, arguments } => {
|
||
|
|
let mut out = self.ensure_started()?;
|
||
|
|
let call_id = format!("tool_{index}");
|
||
|
|
out.extend(self.ensure_tool_block(index, &call_id, "unknown")?);
|
||
|
|
let block_index = match self.open_block {
|
||
|
|
Some(ClaudeOpenBlock::Tool { block_index, .. }) => block_index,
|
||
|
|
_ => return Ok(out),
|
||
|
|
};
|
||
|
|
out.extend(encode_json_sse(
|
||
|
|
Some("content_block_delta"),
|
||
|
|
&json!({
|
||
|
|
"type": "content_block_delta",
|
||
|
|
"index": block_index,
|
||
|
|
"delta": {
|
||
|
|
"type": "input_json_delta",
|
||
|
|
"partial_json": arguments,
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
)?);
|
||
|
|
Ok(out)
|
||
|
|
}
|
||
|
|
CanonicalStreamEvent::Finish {
|
||
|
|
finish_reason,
|
||
|
|
usage,
|
||
|
|
} => {
|
||
|
|
if self.finished {
|
||
|
|
return Ok(Vec::new());
|
||
|
|
}
|
||
|
|
let mut out = self.ensure_started()?;
|
||
|
|
out.extend(self.close_open_block()?);
|
||
|
|
let mut payload = Map::new();
|
||
|
|
payload.insert(
|
||
|
|
"type".to_string(),
|
||
|
|
Value::String("message_delta".to_string()),
|
||
|
|
);
|
||
|
|
payload.insert(
|
||
|
|
"delta".to_string(),
|
||
|
|
json!({
|
||
|
|
"stop_reason": map_openai_finish_reason_to_claude(
|
||
|
|
finish_reason.as_deref()
|
||
|
|
),
|
||
|
|
"stop_sequence": Value::Null,
|
||
|
|
}),
|
||
|
|
);
|
||
|
|
if let Some(usage) = usage {
|
||
|
|
payload.insert(
|
||
|
|
"usage".to_string(),
|
||
|
|
json!({
|
||
|
|
"input_tokens": usage.input_tokens,
|
||
|
|
"output_tokens": usage.output_tokens,
|
||
|
|
}),
|
||
|
|
);
|
||
|
|
}
|
||
|
|
out.extend(encode_json_sse(
|
||
|
|
Some("message_delta"),
|
||
|
|
&Value::Object(payload),
|
||
|
|
)?);
|
||
|
|
out.extend(encode_json_sse(
|
||
|
|
Some("message_stop"),
|
||
|
|
&json!({
|
||
|
|
"type": "message_stop",
|
||
|
|
}),
|
||
|
|
)?);
|
||
|
|
self.finished = true;
|
||
|
|
Ok(out)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn finish(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
|
||
|
|
if !self.started || self.finished {
|
||
|
|
return Ok(Vec::new());
|
||
|
|
}
|
||
|
|
self.emit(CanonicalStreamFrame {
|
||
|
|
id: self
|
||
|
|
.message_id
|
||
|
|
.clone()
|
||
|
|
.unwrap_or_else(|| "msg-local-stream".to_string()),
|
||
|
|
model: self.model.clone().unwrap_or_else(|| "unknown".to_string()),
|
||
|
|
event: CanonicalStreamEvent::Finish {
|
||
|
|
finish_reason: None,
|
||
|
|
usage: None,
|
||
|
|
},
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|