claude - move api types to their own module, simplify message buffer computations
This commit is contained in:
@@ -0,0 +1,120 @@
|
||||
//! Anthropic Claude API types.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use signal_gateway_assistant::Tool;
|
||||
|
||||
/// Request body for the Claude Messages API.
|
||||
#[derive(Serialize)]
|
||||
pub(crate) struct MessagesRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub max_tokens: u32,
|
||||
pub system: &'a [SystemContent],
|
||||
pub messages: &'a [MessageContent],
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub tools: Vec<Tool>,
|
||||
}
|
||||
|
||||
/// A content block in the system prompt array.
|
||||
#[derive(Clone, Default, Serialize)]
|
||||
pub(crate) struct SystemContent {
|
||||
#[serde(rename = "type")]
|
||||
content_type: &'static str,
|
||||
text: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
cache_control: Option<CacheControl>,
|
||||
}
|
||||
|
||||
impl SystemContent {
|
||||
pub fn text(text: impl Into<String>) -> Self {
|
||||
Self {
|
||||
content_type: "text",
|
||||
text: text.into(),
|
||||
cache_control: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_cached(&mut self) {
|
||||
self.cache_control = Some(CacheControl::ephemeral());
|
||||
}
|
||||
}
|
||||
|
||||
/// Cache control directive for prompt caching.
|
||||
#[derive(Clone, Serialize)]
|
||||
struct CacheControl {
|
||||
#[serde(rename = "type")]
|
||||
cache_type: &'static str,
|
||||
}
|
||||
|
||||
impl CacheControl {
|
||||
fn ephemeral() -> Self {
|
||||
Self {
|
||||
cache_type: "ephemeral",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A message in the conversation (can have multiple content blocks).
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub(crate) struct MessageContent {
|
||||
pub role: Box<str>,
|
||||
pub content: Vec<ContentBlock>,
|
||||
}
|
||||
|
||||
impl MessageContent {
|
||||
pub fn assistant(blocks: Vec<ContentBlock>) -> Self {
|
||||
Self {
|
||||
role: "assistant".into(),
|
||||
content: blocks,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tool_result(tool_use_id: String, content: String, is_error: bool) -> Self {
|
||||
Self {
|
||||
role: "user".into(),
|
||||
content: vec![ContentBlock::ToolResult {
|
||||
tool_use_id: tool_use_id.into(),
|
||||
content: content.into(),
|
||||
is_error: if is_error { Some(true) } else { None },
|
||||
}],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A content block in the request/response.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub(crate) enum ContentBlock {
|
||||
Text {
|
||||
text: Box<str>,
|
||||
},
|
||||
ToolUse {
|
||||
id: Box<str>,
|
||||
name: Box<str>,
|
||||
input: Value,
|
||||
},
|
||||
ToolResult {
|
||||
tool_use_id: Box<str>,
|
||||
content: Box<str>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
is_error: Option<bool>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Response from the Claude Messages API.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct MessagesResponse {
|
||||
pub content: Vec<ContentBlock>,
|
||||
pub stop_reason: Box<str>,
|
||||
}
|
||||
|
||||
/// Error response from the Claude API.
|
||||
#[derive(Deserialize)]
|
||||
pub(crate) struct ErrorResponse {
|
||||
pub error: ApiErrorDetail,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(crate) struct ApiErrorDetail {
|
||||
pub message: Box<str>,
|
||||
}
|
||||
@@ -1,12 +1,14 @@
|
||||
//! Claude API implementation of the Assistant trait.
|
||||
|
||||
mod api;
|
||||
mod message_buffer;
|
||||
|
||||
use api::{
|
||||
ContentBlock, ErrorResponse, MessageContent, MessagesRequest, MessagesResponse, SystemContent,
|
||||
};
|
||||
use conf::Conf;
|
||||
use message_buffer::MessageBuffer;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use signal_gateway_assistant::{Assistant, AssistantResponse, ChatMessage, Tool, ToolExecutor};
|
||||
use signal_gateway_assistant::{Assistant, AssistantResponse, ChatMessage, ToolExecutor};
|
||||
use std::{path::PathBuf, sync::Weak, time::Duration};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{error, info, warn};
|
||||
@@ -485,119 +487,3 @@ fn message_to_content(msg: ChatMessage) -> MessageContent {
|
||||
content: vec![ContentBlock::Text { text: text.into() }],
|
||||
}
|
||||
}
|
||||
// ---- API Types ----
|
||||
|
||||
/// Request body for the Claude Messages API.
|
||||
#[derive(Serialize)]
|
||||
struct MessagesRequest<'a> {
|
||||
model: &'a str,
|
||||
max_tokens: u32,
|
||||
system: &'a [SystemContent],
|
||||
messages: &'a [MessageContent],
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
tools: Vec<Tool>,
|
||||
}
|
||||
|
||||
/// A content block in the system prompt array.
|
||||
#[derive(Clone, Default, Serialize)]
|
||||
struct SystemContent {
|
||||
#[serde(rename = "type")]
|
||||
content_type: &'static str,
|
||||
text: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
cache_control: Option<CacheControl>,
|
||||
}
|
||||
|
||||
impl SystemContent {
|
||||
fn text(text: impl Into<String>) -> Self {
|
||||
Self {
|
||||
content_type: "text",
|
||||
text: text.into(),
|
||||
cache_control: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn set_cached(&mut self) {
|
||||
self.cache_control = Some(CacheControl::ephemeral());
|
||||
}
|
||||
}
|
||||
|
||||
/// Cache control directive for prompt caching.
|
||||
#[derive(Clone, Serialize)]
|
||||
struct CacheControl {
|
||||
#[serde(rename = "type")]
|
||||
cache_type: &'static str,
|
||||
}
|
||||
|
||||
impl CacheControl {
|
||||
fn ephemeral() -> Self {
|
||||
Self {
|
||||
cache_type: "ephemeral",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A message in the conversation (can have multiple content blocks).
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
struct MessageContent {
|
||||
role: Box<str>,
|
||||
content: Vec<ContentBlock>,
|
||||
}
|
||||
|
||||
impl MessageContent {
|
||||
fn assistant(blocks: Vec<ContentBlock>) -> Self {
|
||||
Self {
|
||||
role: "assistant".into(),
|
||||
content: blocks,
|
||||
}
|
||||
}
|
||||
|
||||
fn tool_result(tool_use_id: String, content: String, is_error: bool) -> Self {
|
||||
Self {
|
||||
role: "user".into(),
|
||||
content: vec![ContentBlock::ToolResult {
|
||||
tool_use_id: tool_use_id.into(),
|
||||
content: content.into(),
|
||||
is_error: if is_error { Some(true) } else { None },
|
||||
}],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A content block in the request/response.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
enum ContentBlock {
|
||||
Text {
|
||||
text: Box<str>,
|
||||
},
|
||||
ToolUse {
|
||||
id: Box<str>,
|
||||
name: Box<str>,
|
||||
input: Value,
|
||||
},
|
||||
ToolResult {
|
||||
tool_use_id: Box<str>,
|
||||
content: Box<str>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
is_error: Option<bool>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Response from the Claude Messages API.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct MessagesResponse {
|
||||
content: Vec<ContentBlock>,
|
||||
stop_reason: Box<str>,
|
||||
}
|
||||
|
||||
/// Error response from the Claude API.
|
||||
#[derive(Deserialize)]
|
||||
struct ErrorResponse {
|
||||
error: ApiErrorDetail,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ApiErrorDetail {
|
||||
message: Box<str>,
|
||||
}
|
||||
|
||||
@@ -1,15 +1,13 @@
|
||||
//! Message buffer with cached character count.
|
||||
|
||||
use crate::api::{ContentBlock, MessageContent};
|
||||
use serde_json::Value;
|
||||
use std::collections::VecDeque;
|
||||
use std::fmt;
|
||||
|
||||
use crate::{ContentBlock, MessageContent};
|
||||
|
||||
/// A buffer of messages with cached total character count.
|
||||
///
|
||||
/// Uses `VecDeque` for efficient front removal during compaction.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct MessageBuffer {
|
||||
/// `VecDeque` to easily remove oldest messages if needed
|
||||
messages: VecDeque<MessageContent>,
|
||||
/// Cached total character count of all messages.
|
||||
total_chars: usize,
|
||||
@@ -19,12 +17,12 @@ impl MessageBuffer {
|
||||
/// Create a new empty message buffer.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
messages: VecDeque::new(),
|
||||
messages: VecDeque::with_capacity(256),
|
||||
total_chars: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Push a message to the back of the buffer.
|
||||
/// Push a message to the back.
|
||||
pub fn push(&mut self, msg: MessageContent) {
|
||||
self.total_chars += message_chars(&msg);
|
||||
self.messages.push_back(msg);
|
||||
@@ -58,7 +56,7 @@ impl MessageBuffer {
|
||||
self.total_chars
|
||||
}
|
||||
|
||||
/// Returns a reference to the last message, if any.
|
||||
/// Get the last message, if any.
|
||||
pub fn last(&self) -> Option<&MessageContent> {
|
||||
self.messages.back()
|
||||
}
|
||||
@@ -71,22 +69,6 @@ impl MessageBuffer {
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for MessageBuffer {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for MessageBuffer {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("MessageBuffer")
|
||||
.field("len", &self.messages.len())
|
||||
.field("total_chars", &self.total_chars)
|
||||
.field("messages", &self.messages)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Calculate the character count for a single message.
|
||||
fn message_chars(msg: &MessageContent) -> usize {
|
||||
msg.content
|
||||
@@ -99,23 +81,25 @@ fn message_chars(msg: &MessageContent) -> usize {
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Estimate the serialized JSON size of a value.
|
||||
// Estimate the serialized JSON size of a value.
|
||||
// https://github.com/serde-rs/json/issues/784#issuecomment-877688512
|
||||
fn estimate_json_size(value: &Value) -> usize {
|
||||
match value {
|
||||
Value::Null => 4,
|
||||
Value::Bool(true) => 4,
|
||||
Value::Bool(false) => 5,
|
||||
Value::Number(n) => n.to_string().len(),
|
||||
Value::String(s) => s.len() + 2,
|
||||
Value::Array(arr) => {
|
||||
2 + arr.iter().map(estimate_json_size).sum::<usize>() + arr.len().saturating_sub(1)
|
||||
use serde::Serialize;
|
||||
use std::io::{Result, Write};
|
||||
|
||||
struct ByteCount(usize);
|
||||
|
||||
impl Write for ByteCount {
|
||||
fn write(&mut self, buf: &[u8]) -> Result<usize> {
|
||||
self.0 += buf.len();
|
||||
Ok(buf.len())
|
||||
}
|
||||
Value::Object(obj) => {
|
||||
2 + obj
|
||||
.iter()
|
||||
.map(|(k, v)| k.len() + 3 + estimate_json_size(v))
|
||||
.sum::<usize>()
|
||||
+ obj.len().saturating_sub(1)
|
||||
fn flush(&mut self) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
let mut ser = serde_json::Serializer::new(ByteCount(0));
|
||||
value.serialize(&mut ser).unwrap();
|
||||
ser.into_inner().0
|
||||
}
|
||||
|
||||
@@ -25,10 +25,11 @@ pub enum Input {
|
||||
/// Background worker that processes assistant requests serially.
|
||||
pub struct AssistantWorker {
|
||||
assistant: Box<dyn Assistant>,
|
||||
// Serialize input
|
||||
input_rx: mpsc::Receiver<Input>,
|
||||
/// Used when the user wants to cancel the current requests, but not shutdown the service
|
||||
// Used when the user wants to cancel the current requests, but not shutdown the service
|
||||
stop_rx: mpsc::Receiver<()>,
|
||||
/// Used when the user wants to shut down the service
|
||||
// Used when the user wants to shut down the service
|
||||
cancellation_token: CancellationToken,
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user