mirror of
https://github.com/instructkr/claude-code.git
synced 2026-04-03 22:38:48 +03:00
initial commit scaffold
This commit is contained in:
141
rust/crates/api/src/client.rs
Normal file
141
rust/crates/api/src/client.rs
Normal file
@@ -0,0 +1,141 @@
|
||||
use crate::error::ApiError;
|
||||
use crate::providers::claw_provider::{self, ClawApiClient, AuthSource};
|
||||
use crate::providers::openai_compat::{self, OpenAiCompatClient, OpenAiCompatConfig};
|
||||
use crate::providers::{self, Provider, ProviderKind};
|
||||
use crate::types::{MessageRequest, MessageResponse, StreamEvent};
|
||||
|
||||
async fn send_via_provider<P: Provider>(
|
||||
provider: &P,
|
||||
request: &MessageRequest,
|
||||
) -> Result<MessageResponse, ApiError> {
|
||||
provider.send_message(request).await
|
||||
}
|
||||
|
||||
async fn stream_via_provider<P: Provider>(
|
||||
provider: &P,
|
||||
request: &MessageRequest,
|
||||
) -> Result<P::Stream, ApiError> {
|
||||
provider.stream_message(request).await
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ProviderClient {
|
||||
ClawApi(ClawApiClient),
|
||||
Xai(OpenAiCompatClient),
|
||||
OpenAi(OpenAiCompatClient),
|
||||
}
|
||||
|
||||
impl ProviderClient {
|
||||
pub fn from_model(model: &str) -> Result<Self, ApiError> {
|
||||
Self::from_model_with_default_auth(model, None)
|
||||
}
|
||||
|
||||
pub fn from_model_with_default_auth(
|
||||
model: &str,
|
||||
default_auth: Option<AuthSource>,
|
||||
) -> Result<Self, ApiError> {
|
||||
let resolved_model = providers::resolve_model_alias(model);
|
||||
match providers::detect_provider_kind(&resolved_model) {
|
||||
ProviderKind::ClawApi => Ok(Self::ClawApi(match default_auth {
|
||||
Some(auth) => ClawApiClient::from_auth(auth),
|
||||
None => ClawApiClient::from_env()?,
|
||||
})),
|
||||
ProviderKind::Xai => Ok(Self::Xai(OpenAiCompatClient::from_env(
|
||||
OpenAiCompatConfig::xai(),
|
||||
)?)),
|
||||
ProviderKind::OpenAi => Ok(Self::OpenAi(OpenAiCompatClient::from_env(
|
||||
OpenAiCompatConfig::openai(),
|
||||
)?)),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn provider_kind(&self) -> ProviderKind {
|
||||
match self {
|
||||
Self::ClawApi(_) => ProviderKind::ClawApi,
|
||||
Self::Xai(_) => ProviderKind::Xai,
|
||||
Self::OpenAi(_) => ProviderKind::OpenAi,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn send_message(
|
||||
&self,
|
||||
request: &MessageRequest,
|
||||
) -> Result<MessageResponse, ApiError> {
|
||||
match self {
|
||||
Self::ClawApi(client) => send_via_provider(client, request).await,
|
||||
Self::Xai(client) | Self::OpenAi(client) => send_via_provider(client, request).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn stream_message(
|
||||
&self,
|
||||
request: &MessageRequest,
|
||||
) -> Result<MessageStream, ApiError> {
|
||||
match self {
|
||||
Self::ClawApi(client) => stream_via_provider(client, request)
|
||||
.await
|
||||
.map(MessageStream::ClawApi),
|
||||
Self::Xai(client) | Self::OpenAi(client) => stream_via_provider(client, request)
|
||||
.await
|
||||
.map(MessageStream::OpenAiCompat),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum MessageStream {
|
||||
ClawApi(claw_provider::MessageStream),
|
||||
OpenAiCompat(openai_compat::MessageStream),
|
||||
}
|
||||
|
||||
impl MessageStream {
|
||||
#[must_use]
|
||||
pub fn request_id(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::ClawApi(stream) => stream.request_id(),
|
||||
Self::OpenAiCompat(stream) => stream.request_id(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn next_event(&mut self) -> Result<Option<StreamEvent>, ApiError> {
|
||||
match self {
|
||||
Self::ClawApi(stream) => stream.next_event().await,
|
||||
Self::OpenAiCompat(stream) => stream.next_event().await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub use claw_provider::{
|
||||
oauth_token_is_expired, resolve_saved_oauth_token, resolve_startup_auth_source, OAuthTokenSet,
|
||||
};
|
||||
#[must_use]
|
||||
pub fn read_base_url() -> String {
|
||||
claw_provider::read_base_url()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn read_xai_base_url() -> String {
|
||||
openai_compat::read_base_url(OpenAiCompatConfig::xai())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::providers::{detect_provider_kind, resolve_model_alias, ProviderKind};
|
||||
|
||||
#[test]
|
||||
fn resolves_existing_and_grok_aliases() {
|
||||
assert_eq!(resolve_model_alias("opus"), "claude-opus-4-6");
|
||||
assert_eq!(resolve_model_alias("grok"), "grok-3");
|
||||
assert_eq!(resolve_model_alias("grok-mini"), "grok-3-mini");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_detection_prefers_model_family() {
|
||||
assert_eq!(detect_provider_kind("grok-3"), ProviderKind::Xai);
|
||||
assert_eq!(
|
||||
detect_provider_kind("claude-sonnet-4-6"),
|
||||
ProviderKind::ClawApi
|
||||
);
|
||||
}
|
||||
}
|
||||
135
rust/crates/api/src/error.rs
Normal file
135
rust/crates/api/src/error.rs
Normal file
@@ -0,0 +1,135 @@
|
||||
use std::env::VarError;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::time::Duration;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum ApiError {
|
||||
MissingCredentials {
|
||||
provider: &'static str,
|
||||
env_vars: &'static [&'static str],
|
||||
},
|
||||
ExpiredOAuthToken,
|
||||
Auth(String),
|
||||
InvalidApiKeyEnv(VarError),
|
||||
Http(reqwest::Error),
|
||||
Io(std::io::Error),
|
||||
Json(serde_json::Error),
|
||||
Api {
|
||||
status: reqwest::StatusCode,
|
||||
error_type: Option<String>,
|
||||
message: Option<String>,
|
||||
body: String,
|
||||
retryable: bool,
|
||||
},
|
||||
RetriesExhausted {
|
||||
attempts: u32,
|
||||
last_error: Box<ApiError>,
|
||||
},
|
||||
InvalidSseFrame(&'static str),
|
||||
BackoffOverflow {
|
||||
attempt: u32,
|
||||
base_delay: Duration,
|
||||
},
|
||||
}
|
||||
|
||||
impl ApiError {
|
||||
#[must_use]
|
||||
pub const fn missing_credentials(
|
||||
provider: &'static str,
|
||||
env_vars: &'static [&'static str],
|
||||
) -> Self {
|
||||
Self::MissingCredentials { provider, env_vars }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_retryable(&self) -> bool {
|
||||
match self {
|
||||
Self::Http(error) => error.is_connect() || error.is_timeout() || error.is_request(),
|
||||
Self::Api { retryable, .. } => *retryable,
|
||||
Self::RetriesExhausted { last_error, .. } => last_error.is_retryable(),
|
||||
Self::MissingCredentials { .. }
|
||||
| Self::ExpiredOAuthToken
|
||||
| Self::Auth(_)
|
||||
| Self::InvalidApiKeyEnv(_)
|
||||
| Self::Io(_)
|
||||
| Self::Json(_)
|
||||
| Self::InvalidSseFrame(_)
|
||||
| Self::BackoffOverflow { .. } => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for ApiError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::MissingCredentials { provider, env_vars } => write!(
|
||||
f,
|
||||
"missing {provider} credentials; export {} before calling the {provider} API",
|
||||
env_vars.join(" or ")
|
||||
),
|
||||
Self::ExpiredOAuthToken => {
|
||||
write!(
|
||||
f,
|
||||
"saved OAuth token is expired and no refresh token is available"
|
||||
)
|
||||
}
|
||||
Self::Auth(message) => write!(f, "auth error: {message}"),
|
||||
Self::InvalidApiKeyEnv(error) => {
|
||||
write!(f, "failed to read credential environment variable: {error}")
|
||||
}
|
||||
Self::Http(error) => write!(f, "http error: {error}"),
|
||||
Self::Io(error) => write!(f, "io error: {error}"),
|
||||
Self::Json(error) => write!(f, "json error: {error}"),
|
||||
Self::Api {
|
||||
status,
|
||||
error_type,
|
||||
message,
|
||||
body,
|
||||
..
|
||||
} => match (error_type, message) {
|
||||
(Some(error_type), Some(message)) => {
|
||||
write!(f, "api returned {status} ({error_type}): {message}")
|
||||
}
|
||||
_ => write!(f, "api returned {status}: {body}"),
|
||||
},
|
||||
Self::RetriesExhausted {
|
||||
attempts,
|
||||
last_error,
|
||||
} => write!(f, "api failed after {attempts} attempts: {last_error}"),
|
||||
Self::InvalidSseFrame(message) => write!(f, "invalid sse frame: {message}"),
|
||||
Self::BackoffOverflow {
|
||||
attempt,
|
||||
base_delay,
|
||||
} => write!(
|
||||
f,
|
||||
"retry backoff overflowed on attempt {attempt} with base delay {base_delay:?}"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ApiError {}
|
||||
|
||||
impl From<reqwest::Error> for ApiError {
|
||||
fn from(value: reqwest::Error) -> Self {
|
||||
Self::Http(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<std::io::Error> for ApiError {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::Io(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for ApiError {
|
||||
fn from(value: serde_json::Error) -> Self {
|
||||
Self::Json(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<VarError> for ApiError {
|
||||
fn from(value: VarError) -> Self {
|
||||
Self::InvalidApiKeyEnv(value)
|
||||
}
|
||||
}
|
||||
23
rust/crates/api/src/lib.rs
Normal file
23
rust/crates/api/src/lib.rs
Normal file
@@ -0,0 +1,23 @@
|
||||
mod client;
|
||||
mod error;
|
||||
mod providers;
|
||||
mod sse;
|
||||
mod types;
|
||||
|
||||
pub use client::{
|
||||
oauth_token_is_expired, read_base_url, read_xai_base_url, resolve_saved_oauth_token,
|
||||
resolve_startup_auth_source, MessageStream, OAuthTokenSet, ProviderClient,
|
||||
};
|
||||
pub use error::ApiError;
|
||||
pub use providers::claw_provider::{ClawApiClient, ClawApiClient as ApiClient, AuthSource};
|
||||
pub use providers::openai_compat::{OpenAiCompatClient, OpenAiCompatConfig};
|
||||
pub use providers::{
|
||||
detect_provider_kind, max_tokens_for_model, resolve_model_alias, ProviderKind,
|
||||
};
|
||||
pub use sse::{parse_frame, SseParser};
|
||||
pub use types::{
|
||||
ContentBlockDelta, ContentBlockDeltaEvent, ContentBlockStartEvent, ContentBlockStopEvent,
|
||||
InputContentBlock, InputMessage, MessageDelta, MessageDeltaEvent, MessageRequest,
|
||||
MessageResponse, MessageStartEvent, MessageStopEvent, OutputContentBlock, StreamEvent,
|
||||
ToolChoice, ToolDefinition, ToolResultContentBlock, Usage,
|
||||
};
|
||||
1046
rust/crates/api/src/providers/claw_provider.rs
Normal file
1046
rust/crates/api/src/providers/claw_provider.rs
Normal file
File diff suppressed because it is too large
Load Diff
239
rust/crates/api/src/providers/mod.rs
Normal file
239
rust/crates/api/src/providers/mod.rs
Normal file
@@ -0,0 +1,239 @@
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use crate::error::ApiError;
|
||||
use crate::types::{MessageRequest, MessageResponse};
|
||||
|
||||
pub mod claw_provider;
|
||||
pub mod openai_compat;
|
||||
|
||||
pub type ProviderFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, ApiError>> + Send + 'a>>;
|
||||
|
||||
pub trait Provider {
|
||||
type Stream;
|
||||
|
||||
fn send_message<'a>(
|
||||
&'a self,
|
||||
request: &'a MessageRequest,
|
||||
) -> ProviderFuture<'a, MessageResponse>;
|
||||
|
||||
fn stream_message<'a>(
|
||||
&'a self,
|
||||
request: &'a MessageRequest,
|
||||
) -> ProviderFuture<'a, Self::Stream>;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ProviderKind {
|
||||
ClawApi,
|
||||
Xai,
|
||||
OpenAi,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ProviderMetadata {
|
||||
pub provider: ProviderKind,
|
||||
pub auth_env: &'static str,
|
||||
pub base_url_env: &'static str,
|
||||
pub default_base_url: &'static str,
|
||||
}
|
||||
|
||||
const MODEL_REGISTRY: &[(&str, ProviderMetadata)] = &[
|
||||
(
|
||||
"opus",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::ClawApi,
|
||||
auth_env: "ANTHROPIC_API_KEY",
|
||||
base_url_env: "ANTHROPIC_BASE_URL",
|
||||
default_base_url: claw_provider::DEFAULT_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"sonnet",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::ClawApi,
|
||||
auth_env: "ANTHROPIC_API_KEY",
|
||||
base_url_env: "ANTHROPIC_BASE_URL",
|
||||
default_base_url: claw_provider::DEFAULT_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"haiku",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::ClawApi,
|
||||
auth_env: "ANTHROPIC_API_KEY",
|
||||
base_url_env: "ANTHROPIC_BASE_URL",
|
||||
default_base_url: claw_provider::DEFAULT_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"claude-opus-4-6",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::ClawApi,
|
||||
auth_env: "ANTHROPIC_API_KEY",
|
||||
base_url_env: "ANTHROPIC_BASE_URL",
|
||||
default_base_url: claw_provider::DEFAULT_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"claude-sonnet-4-6",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::ClawApi,
|
||||
auth_env: "ANTHROPIC_API_KEY",
|
||||
base_url_env: "ANTHROPIC_BASE_URL",
|
||||
default_base_url: claw_provider::DEFAULT_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"claude-haiku-4-5-20251213",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::ClawApi,
|
||||
auth_env: "ANTHROPIC_API_KEY",
|
||||
base_url_env: "ANTHROPIC_BASE_URL",
|
||||
default_base_url: claw_provider::DEFAULT_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"grok",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::Xai,
|
||||
auth_env: "XAI_API_KEY",
|
||||
base_url_env: "XAI_BASE_URL",
|
||||
default_base_url: openai_compat::DEFAULT_XAI_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"grok-3",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::Xai,
|
||||
auth_env: "XAI_API_KEY",
|
||||
base_url_env: "XAI_BASE_URL",
|
||||
default_base_url: openai_compat::DEFAULT_XAI_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"grok-mini",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::Xai,
|
||||
auth_env: "XAI_API_KEY",
|
||||
base_url_env: "XAI_BASE_URL",
|
||||
default_base_url: openai_compat::DEFAULT_XAI_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"grok-3-mini",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::Xai,
|
||||
auth_env: "XAI_API_KEY",
|
||||
base_url_env: "XAI_BASE_URL",
|
||||
default_base_url: openai_compat::DEFAULT_XAI_BASE_URL,
|
||||
},
|
||||
),
|
||||
(
|
||||
"grok-2",
|
||||
ProviderMetadata {
|
||||
provider: ProviderKind::Xai,
|
||||
auth_env: "XAI_API_KEY",
|
||||
base_url_env: "XAI_BASE_URL",
|
||||
default_base_url: openai_compat::DEFAULT_XAI_BASE_URL,
|
||||
},
|
||||
),
|
||||
];
|
||||
|
||||
#[must_use]
|
||||
pub fn resolve_model_alias(model: &str) -> String {
|
||||
let trimmed = model.trim();
|
||||
let lower = trimmed.to_ascii_lowercase();
|
||||
MODEL_REGISTRY
|
||||
.iter()
|
||||
.find_map(|(alias, metadata)| {
|
||||
(*alias == lower).then_some(match metadata.provider {
|
||||
ProviderKind::ClawApi => match *alias {
|
||||
"opus" => "claude-opus-4-6",
|
||||
"sonnet" => "claude-sonnet-4-6",
|
||||
"haiku" => "claude-haiku-4-5-20251213",
|
||||
_ => trimmed,
|
||||
},
|
||||
ProviderKind::Xai => match *alias {
|
||||
"grok" | "grok-3" => "grok-3",
|
||||
"grok-mini" | "grok-3-mini" => "grok-3-mini",
|
||||
"grok-2" => "grok-2",
|
||||
_ => trimmed,
|
||||
},
|
||||
ProviderKind::OpenAi => trimmed,
|
||||
})
|
||||
})
|
||||
.map_or_else(|| trimmed.to_string(), ToOwned::to_owned)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn metadata_for_model(model: &str) -> Option<ProviderMetadata> {
|
||||
let canonical = resolve_model_alias(model);
|
||||
let lower = canonical.to_ascii_lowercase();
|
||||
if let Some((_, metadata)) = MODEL_REGISTRY.iter().find(|(alias, _)| *alias == lower) {
|
||||
return Some(*metadata);
|
||||
}
|
||||
if lower.starts_with("grok") {
|
||||
return Some(ProviderMetadata {
|
||||
provider: ProviderKind::Xai,
|
||||
auth_env: "XAI_API_KEY",
|
||||
base_url_env: "XAI_BASE_URL",
|
||||
default_base_url: openai_compat::DEFAULT_XAI_BASE_URL,
|
||||
});
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn detect_provider_kind(model: &str) -> ProviderKind {
|
||||
if let Some(metadata) = metadata_for_model(model) {
|
||||
return metadata.provider;
|
||||
}
|
||||
if claw_provider::has_auth_from_env_or_saved().unwrap_or(false) {
|
||||
return ProviderKind::ClawApi;
|
||||
}
|
||||
if openai_compat::has_api_key("OPENAI_API_KEY") {
|
||||
return ProviderKind::OpenAi;
|
||||
}
|
||||
if openai_compat::has_api_key("XAI_API_KEY") {
|
||||
return ProviderKind::Xai;
|
||||
}
|
||||
ProviderKind::ClawApi
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn max_tokens_for_model(model: &str) -> u32 {
|
||||
let canonical = resolve_model_alias(model);
|
||||
if canonical.contains("opus") {
|
||||
32_000
|
||||
} else {
|
||||
64_000
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{detect_provider_kind, max_tokens_for_model, resolve_model_alias, ProviderKind};
|
||||
|
||||
#[test]
|
||||
fn resolves_grok_aliases() {
|
||||
assert_eq!(resolve_model_alias("grok"), "grok-3");
|
||||
assert_eq!(resolve_model_alias("grok-mini"), "grok-3-mini");
|
||||
assert_eq!(resolve_model_alias("grok-2"), "grok-2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_provider_from_model_name_first() {
|
||||
assert_eq!(detect_provider_kind("grok"), ProviderKind::Xai);
|
||||
assert_eq!(
|
||||
detect_provider_kind("claude-sonnet-4-6"),
|
||||
ProviderKind::ClawApi
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_existing_max_token_heuristic() {
|
||||
assert_eq!(max_tokens_for_model("opus"), 32_000);
|
||||
assert_eq!(max_tokens_for_model("grok-3"), 64_000);
|
||||
}
|
||||
}
|
||||
1050
rust/crates/api/src/providers/openai_compat.rs
Normal file
1050
rust/crates/api/src/providers/openai_compat.rs
Normal file
File diff suppressed because it is too large
Load Diff
279
rust/crates/api/src/sse.rs
Normal file
279
rust/crates/api/src/sse.rs
Normal file
@@ -0,0 +1,279 @@
|
||||
use crate::error::ApiError;
|
||||
use crate::types::StreamEvent;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct SseParser {
|
||||
buffer: Vec<u8>,
|
||||
}
|
||||
|
||||
impl SseParser {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<StreamEvent>, ApiError> {
|
||||
self.buffer.extend_from_slice(chunk);
|
||||
let mut events = Vec::new();
|
||||
|
||||
while let Some(frame) = self.next_frame() {
|
||||
if let Some(event) = parse_frame(&frame)? {
|
||||
events.push(event);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
pub fn finish(&mut self) -> Result<Vec<StreamEvent>, ApiError> {
|
||||
if self.buffer.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let trailing = std::mem::take(&mut self.buffer);
|
||||
match parse_frame(&String::from_utf8_lossy(&trailing))? {
|
||||
Some(event) => Ok(vec![event]),
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn next_frame(&mut self) -> Option<String> {
|
||||
let separator = self
|
||||
.buffer
|
||||
.windows(2)
|
||||
.position(|window| window == b"\n\n")
|
||||
.map(|position| (position, 2))
|
||||
.or_else(|| {
|
||||
self.buffer
|
||||
.windows(4)
|
||||
.position(|window| window == b"\r\n\r\n")
|
||||
.map(|position| (position, 4))
|
||||
})?;
|
||||
|
||||
let (position, separator_len) = separator;
|
||||
let frame = self
|
||||
.buffer
|
||||
.drain(..position + separator_len)
|
||||
.collect::<Vec<_>>();
|
||||
let frame_len = frame.len().saturating_sub(separator_len);
|
||||
Some(String::from_utf8_lossy(&frame[..frame_len]).into_owned())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_frame(frame: &str) -> Result<Option<StreamEvent>, ApiError> {
|
||||
let trimmed = frame.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut data_lines = Vec::new();
|
||||
let mut event_name: Option<&str> = None;
|
||||
|
||||
for line in trimmed.lines() {
|
||||
if line.starts_with(':') {
|
||||
continue;
|
||||
}
|
||||
if let Some(name) = line.strip_prefix("event:") {
|
||||
event_name = Some(name.trim());
|
||||
continue;
|
||||
}
|
||||
if let Some(data) = line.strip_prefix("data:") {
|
||||
data_lines.push(data.trim_start());
|
||||
}
|
||||
}
|
||||
|
||||
if matches!(event_name, Some("ping")) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if data_lines.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let payload = data_lines.join("\n");
|
||||
if payload == "[DONE]" {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
serde_json::from_str::<StreamEvent>(&payload)
|
||||
.map(Some)
|
||||
.map_err(ApiError::from)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{parse_frame, SseParser};
|
||||
use crate::types::{ContentBlockDelta, MessageDelta, OutputContentBlock, StreamEvent, Usage};
|
||||
|
||||
#[test]
|
||||
fn parses_single_frame() {
|
||||
let frame = concat!(
|
||||
"event: content_block_start\n",
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"Hi\"}}\n\n"
|
||||
);
|
||||
|
||||
let event = parse_frame(frame).expect("frame should parse");
|
||||
assert_eq!(
|
||||
event,
|
||||
Some(StreamEvent::ContentBlockStart(
|
||||
crate::types::ContentBlockStartEvent {
|
||||
index: 0,
|
||||
content_block: OutputContentBlock::Text {
|
||||
text: "Hi".to_string(),
|
||||
},
|
||||
},
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_chunked_stream() {
|
||||
let mut parser = SseParser::new();
|
||||
let first = b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hel";
|
||||
let second = b"lo\"}}\n\n";
|
||||
|
||||
assert!(parser
|
||||
.push(first)
|
||||
.expect("first chunk should buffer")
|
||||
.is_empty());
|
||||
let events = parser.push(second).expect("second chunk should parse");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![StreamEvent::ContentBlockDelta(
|
||||
crate::types::ContentBlockDeltaEvent {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::TextDelta {
|
||||
text: "Hello".to_string(),
|
||||
},
|
||||
}
|
||||
)]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_ping_and_done() {
|
||||
let mut parser = SseParser::new();
|
||||
let payload = concat!(
|
||||
": keepalive\n",
|
||||
"event: ping\n",
|
||||
"data: {\"type\":\"ping\"}\n\n",
|
||||
"event: message_delta\n",
|
||||
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\",\"stop_sequence\":null},\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}\n\n",
|
||||
"event: message_stop\n",
|
||||
"data: {\"type\":\"message_stop\"}\n\n",
|
||||
"data: [DONE]\n\n"
|
||||
);
|
||||
|
||||
let events = parser
|
||||
.push(payload.as_bytes())
|
||||
.expect("parser should succeed");
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
StreamEvent::MessageDelta(crate::types::MessageDeltaEvent {
|
||||
delta: MessageDelta {
|
||||
stop_reason: Some("tool_use".to_string()),
|
||||
stop_sequence: None,
|
||||
},
|
||||
usage: Usage {
|
||||
input_tokens: 1,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
output_tokens: 2,
|
||||
},
|
||||
}),
|
||||
StreamEvent::MessageStop(crate::types::MessageStopEvent {}),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_data_less_event_frames() {
|
||||
let frame = "event: ping\n\n";
|
||||
let event = parse_frame(frame).expect("frame without data should be ignored");
|
||||
assert_eq!(event, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_split_json_across_data_lines() {
|
||||
let frame = concat!(
|
||||
"event: content_block_delta\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\n",
|
||||
"data: \"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n"
|
||||
);
|
||||
|
||||
let event = parse_frame(frame).expect("frame should parse");
|
||||
assert_eq!(
|
||||
event,
|
||||
Some(StreamEvent::ContentBlockDelta(
|
||||
crate::types::ContentBlockDeltaEvent {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::TextDelta {
|
||||
text: "Hello".to_string(),
|
||||
},
|
||||
}
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_thinking_content_block_start() {
|
||||
let frame = concat!(
|
||||
"event: content_block_start\n",
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\",\"signature\":null}}\n\n"
|
||||
);
|
||||
|
||||
let event = parse_frame(frame).expect("frame should parse");
|
||||
assert_eq!(
|
||||
event,
|
||||
Some(StreamEvent::ContentBlockStart(
|
||||
crate::types::ContentBlockStartEvent {
|
||||
index: 0,
|
||||
content_block: OutputContentBlock::Thinking {
|
||||
thinking: String::new(),
|
||||
signature: None,
|
||||
},
|
||||
},
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_thinking_related_deltas() {
|
||||
let thinking = concat!(
|
||||
"event: content_block_delta\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"step 1\"}}\n\n"
|
||||
);
|
||||
let signature = concat!(
|
||||
"event: content_block_delta\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"sig_123\"}}\n\n"
|
||||
);
|
||||
|
||||
let thinking_event = parse_frame(thinking).expect("thinking delta should parse");
|
||||
let signature_event = parse_frame(signature).expect("signature delta should parse");
|
||||
|
||||
assert_eq!(
|
||||
thinking_event,
|
||||
Some(StreamEvent::ContentBlockDelta(
|
||||
crate::types::ContentBlockDeltaEvent {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::ThinkingDelta {
|
||||
thinking: "step 1".to_string(),
|
||||
},
|
||||
}
|
||||
))
|
||||
);
|
||||
assert_eq!(
|
||||
signature_event,
|
||||
Some(StreamEvent::ContentBlockDelta(
|
||||
crate::types::ContentBlockDeltaEvent {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::SignatureDelta {
|
||||
signature: "sig_123".to_string(),
|
||||
},
|
||||
}
|
||||
))
|
||||
);
|
||||
}
|
||||
}
|
||||
223
rust/crates/api/src/types.rs
Normal file
223
rust/crates/api/src/types.rs
Normal file
@@ -0,0 +1,223 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct MessageRequest {
|
||||
pub model: String,
|
||||
pub max_tokens: u32,
|
||||
pub messages: Vec<InputMessage>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub system: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<Vec<ToolDefinition>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_choice: Option<ToolChoice>,
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
pub stream: bool,
|
||||
}
|
||||
|
||||
impl MessageRequest {
|
||||
#[must_use]
|
||||
pub fn with_streaming(mut self) -> Self {
|
||||
self.stream = true;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct InputMessage {
|
||||
pub role: String,
|
||||
pub content: Vec<InputContentBlock>,
|
||||
}
|
||||
|
||||
impl InputMessage {
|
||||
#[must_use]
|
||||
pub fn user_text(text: impl Into<String>) -> Self {
|
||||
Self {
|
||||
role: "user".to_string(),
|
||||
content: vec![InputContentBlock::Text { text: text.into() }],
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn user_tool_result(
|
||||
tool_use_id: impl Into<String>,
|
||||
content: impl Into<String>,
|
||||
is_error: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
role: "user".to_string(),
|
||||
content: vec![InputContentBlock::ToolResult {
|
||||
tool_use_id: tool_use_id.into(),
|
||||
content: vec![ToolResultContentBlock::Text {
|
||||
text: content.into(),
|
||||
}],
|
||||
is_error,
|
||||
}],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum InputContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
input: Value,
|
||||
},
|
||||
ToolResult {
|
||||
tool_use_id: String,
|
||||
content: Vec<ToolResultContentBlock>,
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
is_error: bool,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ToolResultContentBlock {
|
||||
Text { text: String },
|
||||
Json { value: Value },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ToolDefinition {
|
||||
pub name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
pub input_schema: Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ToolChoice {
|
||||
Auto,
|
||||
Any,
|
||||
Tool { name: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct MessageResponse {
|
||||
pub id: String,
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String,
|
||||
pub role: String,
|
||||
pub content: Vec<OutputContentBlock>,
|
||||
pub model: String,
|
||||
#[serde(default)]
|
||||
pub stop_reason: Option<String>,
|
||||
#[serde(default)]
|
||||
pub stop_sequence: Option<String>,
|
||||
pub usage: Usage,
|
||||
#[serde(default)]
|
||||
pub request_id: Option<String>,
|
||||
}
|
||||
|
||||
impl MessageResponse {
|
||||
#[must_use]
|
||||
pub fn total_tokens(&self) -> u32 {
|
||||
self.usage.total_tokens()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum OutputContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
input: Value,
|
||||
},
|
||||
Thinking {
|
||||
#[serde(default)]
|
||||
thinking: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
signature: Option<String>,
|
||||
},
|
||||
RedactedThinking {
|
||||
data: Value,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct Usage {
|
||||
pub input_tokens: u32,
|
||||
#[serde(default)]
|
||||
pub cache_creation_input_tokens: u32,
|
||||
#[serde(default)]
|
||||
pub cache_read_input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
}
|
||||
|
||||
impl Usage {
|
||||
#[must_use]
|
||||
pub const fn total_tokens(&self) -> u32 {
|
||||
self.input_tokens + self.output_tokens
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct MessageStartEvent {
|
||||
pub message: MessageResponse,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct MessageDeltaEvent {
|
||||
pub delta: MessageDelta,
|
||||
pub usage: Usage,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct MessageDelta {
|
||||
#[serde(default)]
|
||||
pub stop_reason: Option<String>,
|
||||
#[serde(default)]
|
||||
pub stop_sequence: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ContentBlockStartEvent {
|
||||
pub index: u32,
|
||||
pub content_block: OutputContentBlock,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ContentBlockDeltaEvent {
|
||||
pub index: u32,
|
||||
pub delta: ContentBlockDelta,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ContentBlockDelta {
|
||||
TextDelta { text: String },
|
||||
InputJsonDelta { partial_json: String },
|
||||
ThinkingDelta { thinking: String },
|
||||
SignatureDelta { signature: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ContentBlockStopEvent {
|
||||
pub index: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct MessageStopEvent {}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum StreamEvent {
|
||||
MessageStart(MessageStartEvent),
|
||||
MessageDelta(MessageDeltaEvent),
|
||||
ContentBlockStart(ContentBlockStartEvent),
|
||||
ContentBlockDelta(ContentBlockDeltaEvent),
|
||||
ContentBlockStop(ContentBlockStopEvent),
|
||||
MessageStop(MessageStopEvent),
|
||||
}
|
||||
Reference in New Issue
Block a user