espionox 0.1.42

A library for creating multi-agent workflows in Rust
Documentation
use super::{
    super::{
        error::{CompletionResult, ProviderResponseError},
        inference::{CompletionRequest, CompletionRequestBuilder, CompletionResponse},
        ModelParameters,
    },
    builder::AnthropicCompletionModel,
    streaming::AnthropicStreamResponse,
};
use crate::agents::memory::{MessageRole, MessageStack};
use crate::language_models::completions::error::CompletionError;
use crate::language_models::completions::inference::ProcessResponseReturn;
use crate::language_models::completions::streaming::{
    CompletionStream, ProviderStreamHandler, StreamedCompletionHandler,
};
use futures::TryStreamExt;
use reqwest_streams::JsonStreamResponse;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Duration;

#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
pub struct AnthropicIoRequest {
    pub model: String,
    pub messages: Value,
    pub temperature: f32,
    pub system: String,
    pub max_tokens: u32,
    pub stream: bool,
}

impl AnthropicIoRequest {
    pub fn new(
        stack: &MessageStack,
        params: &ModelParameters,
        typ: AnthropicCompletionModel,
        stream: bool,
    ) -> Self {
        let system_stack: MessageStack = stack.ref_filter_by(&MessageRole::System, true).into();
        let sans_system_stack: MessageStack =
            stack.ref_filter_by(&MessageRole::System, false).into();
        let system = system_stack
            .as_ref()
            .into_iter()
            .map(|m| m.content.as_str())
            .collect::<Vec<&str>>()
            .join(".");
        let temperature = match params.temperature().ok() {
            Some(t) => t,
            None => 0.7,
        };
        Self {
            model: typ.model_str().to_string(),
            messages: typ.serialize_messages(&sans_system_stack),
            temperature,
            max_tokens: params.max_tokens.unwrap_or(1000),
            system,
            stream,
        }
    }
}

impl CompletionRequest for AnthropicIoRequest {
    fn as_json(&self) -> CompletionResult<Value> {
        Ok(serde_json::to_value(self)?)
    }

    fn process_response<'r>(&'r self, response: reqwest::Response) -> ProcessResponseReturn<'r> {
        Box::pin(async move {
            match self.stream {
                false => {
                    let json = response.json().await?;
                    tracing::warn!("got response:  {json:#?}");
                    let response: AnthropicResponse = serde_json::from_value(json)?;
                    match response {
                        AnthropicResponse::Success(mut suc) => {
                            let content = suc.content.remove(0).text;
                            Ok(CompletionResponse::from(content))
                        }
                        AnthropicResponse::Err { error } => Err(error.into_error()),
                    }
                }
                true => {
                    let response_stream: CompletionStream = Box::new(
                        tokio::time::timeout(Duration::from_secs(10), async {
                            response
                                .json_array_stream::<Value>(1024)
                                .map_err(|err| err.into())
                        })
                        .await
                        .map_err(|_| CompletionError::StreamTimeout)?,
                    );
                    let handler: ProviderStreamHandler =
                        StreamedCompletionHandler::<AnthropicStreamResponse>::from(response_stream)
                            .into();
                    Ok(handler.into())
                }
            }
        })
    }
}

#[derive(Debug, Deserialize, Clone)]
#[serde(untagged)]
pub enum AnthropicResponse {
    Success(AnthropicSuccess),
    Err { error: AnthropicError },
}

#[derive(Debug, Deserialize, Clone)]
pub struct AnthropicSuccess {
    content: Vec<AnthropicResponseContent>,
    usage: AnthropicUsage,
}

#[derive(Debug, Deserialize, Clone)]
pub struct AnthropicError {
    message: String,
}
impl ProviderResponseError for AnthropicError {}

#[derive(Debug, Deserialize, Clone)]
pub struct AnthropicResponseContent {
    text: String,
}

#[derive(Debug, Deserialize, Clone)]
pub struct AnthropicUsage {
    input_tokens: i32,
    output_tokens: i32,
}