diff --git a/rig-core/examples/extractor_with_deepseek.rs b/rig-core/examples/extractor_with_deepseek.rs new file mode 100644 index 00000000..20fdb868 --- /dev/null +++ b/rig-core/examples/extractor_with_deepseek.rs @@ -0,0 +1,36 @@ +use rig::providers::deepseek; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +/// A record representing a person +struct Person { + /// The person's first name, if provided (null otherwise) + pub first_name: Option, + /// The person's last name, if provided (null otherwise) + pub last_name: Option, + /// The person's job, if provided (null otherwise) + pub job: Option, +} + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + // Create DeepSeek client + let deepseek_client = deepseek::Client::from_env(); + + // Create extractor + let data_extractor = deepseek_client + .extractor::(deepseek::DEEPSEEK_CHAT) + .build(); + + let person = data_extractor + .extract("Hello my name is John Doe! I am a software engineer.") + .await?; + + println!( + "DeepSeek: {}", + serde_json::to_string_pretty(&person).unwrap() + ); + + Ok(()) +} diff --git a/rig-core/src/providers/deepseek.rs b/rig-core/src/providers/deepseek.rs index 5ffe9e9e..3a13d1ad 100644 --- a/rig-core/src/providers/deepseek.rs +++ b/rig-core/src/providers/deepseek.rs @@ -10,9 +10,11 @@ //! ``` use crate::{ completion::{CompletionModel, CompletionRequest, CompletionResponse}, + extractor::ExtractorBuilder, json_utils, }; use reqwest::Client as HttpClient; +use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -66,6 +68,14 @@ impl Client { pub fn agent(&self, model_name: &str) -> crate::agent::AgentBuilder { crate::agent::AgentBuilder::new(self.completion_model(model_name)) } + + /// Create an extractor builder with the given completion model. + pub fn extractor Deserialize<'a> + Serialize + Send + Sync>( + &self, + model: &str, + ) -> ExtractorBuilder { + ExtractorBuilder::new(self.completion_model(model)) + } } /// The response shape from the DeepSeek API @@ -274,3 +284,5 @@ impl CompletionModel for DeepSeekCompletionModel { // ================================================================ /// `deepseek-chat` completion model pub const DEEPSEEK_CHAT: &str = "deepseek-chat"; +/// `deepseek-reasoner` completion model +pub const DEEPSEEK_REASONER: &str = "deepseek-reasoner";