swiftide_integrations/anthropic/
mod.rs

1use std::{pin::Pin, sync::Arc};
2
3use derive_builder::Builder;
4use swiftide_core::chat_completion::Usage;
5
6pub mod chat_completion;
7pub mod simple_prompt;
8
9#[derive(Builder, Clone)]
10pub struct Anthropic {
11    #[builder(
12        default = Arc::new(async_anthropic::Client::default()),
13        setter(custom)
14    )]
15    client: Arc<async_anthropic::Client>,
16
17    #[builder(default)]
18    default_options: Options,
19
20    #[cfg(feature = "metrics")]
21    #[builder(default)]
22    /// Optional metadata to attach to metrics emitted by this client.
23    metric_metadata: Option<std::collections::HashMap<String, String>>,
24
25    /// A callback function that is called when usage information is available.
26    #[builder(default, setter(custom))]
27    #[allow(clippy::type_complexity)]
28    on_usage: Option<
29        Arc<
30            dyn for<'a> Fn(
31                    &'a Usage,
32                ) -> Pin<
33                    Box<dyn std::future::Future<Output = anyhow::Result<()>> + Send + 'a>,
34                > + Send
35                + Sync,
36        >,
37    >,
38}
39
40impl std::fmt::Debug for Anthropic {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        f.debug_struct("Anthropic")
43            .field("client", &self.client)
44            .field("default_options", &self.default_options)
45            .finish()
46    }
47}
48
49#[derive(Debug, Clone, Builder)]
50#[builder(setter(into, strip_option))]
51pub struct Options {
52    #[builder(default)]
53    pub prompt_model: String,
54}
55
56impl Default for Options {
57    fn default() -> Self {
58        Self {
59            prompt_model: "claude-3-5-sonnet-20241022".to_string(),
60        }
61    }
62}
63
64impl Anthropic {
65    pub fn builder() -> AnthropicBuilder {
66        AnthropicBuilder::default()
67    }
68}
69
70impl AnthropicBuilder {
71    /// Adds a callback function that will be called when usage information is available.
72    pub fn on_usage<F>(&mut self, func: F) -> &mut Self
73    where
74        F: Fn(&Usage) -> anyhow::Result<()> + Send + Sync + 'static,
75    {
76        let func = Arc::new(func);
77        self.on_usage = Some(Some(Arc::new(move |usage: &Usage| {
78            let func = func.clone();
79            Box::pin(async move { func(usage) })
80        })));
81
82        self
83    }
84
85    /// Adds an asynchronous callback function that will be called when usage information is
86    /// available.
87    pub fn on_usage_async<F>(&mut self, func: F) -> &mut Self
88    where
89        F: for<'a> Fn(
90                &'a Usage,
91            )
92                -> Pin<Box<dyn std::future::Future<Output = anyhow::Result<()>> + Send + 'a>>
93            + Send
94            + Sync
95            + 'static,
96    {
97        let func = Arc::new(func);
98        self.on_usage = Some(Some(Arc::new(move |usage: &Usage| {
99            let func = func.clone();
100            Box::pin(async move { func(usage).await })
101        })));
102
103        self
104    }
105
106    /// Sets the client for the `Anthropic` instance.
107    ///
108    /// See the `async_anthropic::Client` documentation for more information.
109    ///
110    /// # Parameters
111    /// - `client`: The `Anthropic` client to set.
112    ///
113    /// # Returns
114    /// A mutable reference to the `AnthropicBuilder`.
115    pub fn client(&mut self, client: async_anthropic::Client) -> &mut Self {
116        self.client = Some(Arc::new(client));
117        self
118    }
119
120    /// Sets the default prompt model for the `Anthropic` instance.
121    ///
122    /// # Parameters
123    /// - `model`: The prompt model to set.
124    ///
125    /// # Returns
126    /// A mutable reference to the `AnthropicBuilder`.
127    pub fn default_prompt_model(&mut self, model: impl Into<String>) -> &mut Self {
128        if let Some(options) = self.default_options.as_mut() {
129            options.prompt_model = model.into();
130        } else {
131            self.default_options = Some(Options {
132                prompt_model: model.into(),
133            });
134        }
135        self
136    }
137}