-
Notifications
You must be signed in to change notification settings - Fork 94
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #27 from 0xPlaygrounds/feat/anthropic
feat(providers): Integrate anthropic models
- Loading branch information
Showing
7 changed files
with
456 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| use std::env; | ||
|
|
||
| use rig::{ | ||
| completion::Prompt, | ||
| providers::anthropic::{self, CLAUDE_3_5_SONNET}, | ||
| }; | ||
|
|
||
| #[tokio::main] | ||
| async fn main() -> Result<(), anyhow::Error> { | ||
| // Create OpenAI client | ||
| let client = anthropic::ClientBuilder::new( | ||
| &env::var("ANTHROPIC_API_KEY").expect("ANTHROPIC_API_KEY not set"), | ||
| ) | ||
| .build(); | ||
|
|
||
| // Create agent with a single context prompt | ||
| let agent = client | ||
| .agent(CLAUDE_3_5_SONNET) | ||
| .preamble("Be precise and concise.") | ||
| .temperature(0.5) | ||
| .max_tokens(8192) | ||
| .build(); | ||
|
|
||
| // Prompt the agent and print the response | ||
| let response = agent | ||
| .prompt("When and where and what type is the next solar eclipse?") | ||
| .await?; | ||
| println!("{}", response); | ||
|
|
||
| Ok(()) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,149 @@ | ||
| //! Anthropic client api implementation | ||
| use crate::{agent::AgentBuilder, extractor::ExtractorBuilder}; | ||
|
|
||
| use schemars::JsonSchema; | ||
| use serde::{Deserialize, Serialize}; | ||
|
|
||
| use super::completion::{CompletionModel, ANTHROPIC_VERSION_LATEST}; | ||
|
|
||
| // ================================================================ | ||
| // Main Anthropic Client | ||
| // ================================================================ | ||
| const ANTHROPIC_API_BASE_URL: &str = "https://api.anthropic.com"; | ||
|
|
||
| #[derive(Clone)] | ||
| pub struct ClientBuilder<'a> { | ||
| api_key: &'a str, | ||
| base_url: &'a str, | ||
| anthropic_version: &'a str, | ||
| anthropic_betas: Option<Vec<&'a str>>, | ||
| } | ||
|
|
||
| /// Create a new anthropic client using the builder | ||
| /// | ||
| /// # Example | ||
| /// ``` | ||
| /// use rig::providers::anthropic::{ClientBuilder, self}; | ||
| /// | ||
| /// // Initialize the Anthropic client | ||
| /// let anthropic_client = ClientBuilder::new("your-claude-api-key") | ||
| /// .anthropic_version(ANTHROPIC_VERSION_LATEST) | ||
| /// .anthropic_beta("prompt-caching-2024-07-31") | ||
| /// .build() | ||
| /// ``` | ||
| impl<'a> ClientBuilder<'a> { | ||
| pub fn new(api_key: &'a str) -> Self { | ||
| Self { | ||
| api_key, | ||
| base_url: ANTHROPIC_API_BASE_URL, | ||
| anthropic_version: ANTHROPIC_VERSION_LATEST, | ||
| anthropic_betas: None, | ||
| } | ||
| } | ||
|
|
||
| pub fn base_url(mut self, base_url: &'a str) -> Self { | ||
| self.base_url = base_url; | ||
| self | ||
| } | ||
|
|
||
| pub fn anthropic_version(mut self, anthropic_version: &'a str) -> Self { | ||
| self.anthropic_version = anthropic_version; | ||
| self | ||
| } | ||
|
|
||
| pub fn anthropic_beta(mut self, anthropic_beta: &'a str) -> Self { | ||
| if let Some(mut betas) = self.anthropic_betas { | ||
| betas.push(anthropic_beta); | ||
| self.anthropic_betas = Some(betas); | ||
| } else { | ||
| self.anthropic_betas = Some(vec![anthropic_beta]); | ||
| } | ||
| self | ||
| } | ||
|
|
||
| pub fn build(self) -> Client { | ||
| Client::new( | ||
| self.api_key, | ||
| self.base_url, | ||
| self.anthropic_betas, | ||
| self.anthropic_version, | ||
| ) | ||
| } | ||
| } | ||
|
|
||
| #[derive(Clone)] | ||
| pub struct Client { | ||
| base_url: String, | ||
| http_client: reqwest::Client, | ||
| } | ||
|
|
||
| impl Client { | ||
| /// Create a new Anthropic client with the given API key, base URL, betas, and version. | ||
| /// Note, you proably want to use the `ClientBuilder` instead. | ||
| /// | ||
| /// Panics: | ||
| /// - If the API key or version cannot be parsed as a Json value from a String. | ||
| /// - This should really never happen. | ||
| /// - If the reqwest client cannot be built (if the TLS backend cannot be initialized). | ||
| pub fn new(api_key: &str, base_url: &str, betas: Option<Vec<&str>>, version: &str) -> Self { | ||
| Self { | ||
| base_url: base_url.to_string(), | ||
| http_client: reqwest::Client::builder() | ||
| .default_headers({ | ||
| let mut headers = reqwest::header::HeaderMap::new(); | ||
| headers.insert("x-api-key", api_key.parse().expect("API key should parse")); | ||
| headers.insert( | ||
| "anthropic-version", | ||
| version.parse().expect("Anthropic version should parse"), | ||
| ); | ||
| if let Some(betas) = betas { | ||
| headers.insert( | ||
| "anthropic-beta", | ||
| betas | ||
| .join(",") | ||
| .parse() | ||
| .expect("Anthropic betas should parse"), | ||
| ); | ||
| } | ||
| headers | ||
| }) | ||
| .build() | ||
| .expect("Anthropic reqwest client should build"), | ||
| } | ||
| } | ||
|
|
||
| pub fn post(&self, path: &str) -> reqwest::RequestBuilder { | ||
| let url = format!("{}/{}", self.base_url, path).replace("//", "/"); | ||
| self.http_client.post(url) | ||
| } | ||
|
|
||
| pub fn completion_model(&self, model: &str) -> CompletionModel { | ||
| CompletionModel::new(self.clone(), model) | ||
| } | ||
|
|
||
| /// Create an agent builder with the given completion model. | ||
| /// | ||
| /// # Example | ||
| /// ``` | ||
| /// use rig::providers::anthropic::{ClientBuilder, self}; | ||
| /// | ||
| /// // Initialize the Anthropic client | ||
| /// let anthropic = ClientBuilder::new("your-claude-api-key").build(); | ||
| /// | ||
| /// let agent = anthropic.agent(anthropic::CLAUDE_3_5_SONNET) | ||
| /// .preamble("You are comedian AI with a mission to make people laugh.") | ||
| /// .temperature(0.0) | ||
| /// .build(); | ||
| /// ``` | ||
| pub fn agent(&self, model: &str) -> AgentBuilder<CompletionModel> { | ||
| AgentBuilder::new(self.completion_model(model)) | ||
| } | ||
|
|
||
| pub fn extractor<T: JsonSchema + for<'a> Deserialize<'a> + Serialize + Send + Sync>( | ||
| &self, | ||
| model: &str, | ||
| ) -> ExtractorBuilder<T, CompletionModel> { | ||
| ExtractorBuilder::new(self.completion_model(model)) | ||
| } | ||
| } |
Oops, something went wrong.