stone_mason/
anthropic.rs

1use crate::ModelVersion::{V1, V2};
2use crate::{FromModelOutput, ModelVersion};
3use derive_builder::Builder;
4use serde::{Deserialize, Serialize};
5use std::fmt::{Display, Error, Formatter};
6
7pub enum AnthropicModel {
8    Claude(ModelVersion),
9    ClaudeInstant(ModelVersion),
10}
11
12impl Display for AnthropicModel {
13    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
14        match self {
15            AnthropicModel::Claude(v) if *v == V1 => write!(f, "claude-v1"),
16            AnthropicModel::Claude(v) if *v == V2 => write!(f, "claude-v2"),
17            AnthropicModel::ClaudeInstant(v) if *v == V1 => write!(f, "claude-instant-v1"),
18            _ => Err(Error),
19        }
20    }
21}
22
23#[derive(Serialize, Builder)]
24#[builder(setter(strip_option))]
25pub struct AnthropicParams {
26    prompt: String,
27    max_tokens_to_sample: u32,
28
29    #[builder(default = "None")]
30    #[serde(skip_serializing_if = "Option::is_none")]
31    temperature: Option<f32>,
32
33    #[builder(default = "None")]
34    #[serde(skip_serializing_if = "Option::is_none")]
35    stop_sequences: Option<Vec<String>>,
36
37    #[builder(default = "None")]
38    #[serde(skip_serializing_if = "Option::is_none")]
39    top_p: Option<f32>,
40
41    #[builder(default = "None")]
42    #[serde(skip_serializing_if = "Option::is_none")]
43    top_k: Option<u32>,
44}
45
46#[derive(Deserialize, Debug)]
47pub struct AnthropicResponse {
48    pub completion: String,
49    pub stop_reason: String,
50    pub stop: String,
51}
52
53impl<'de> FromModelOutput<'de, AnthropicResponse> for AnthropicResponse {}