potato_type/
lib.rs

1pub mod error;
2
3pub use crate::error::TypeError;
4use pyo3::prelude::*;
5use schemars::JsonSchema;
6use serde::de::Error;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::any::type_name;
10use std::fmt;
11use std::fmt::Display;
12use std::path::{Path, PathBuf};
13use tracing::error;
14pub mod anthropic;
15pub mod common;
16pub mod google;
17pub mod openai;
18pub mod prompt;
19pub mod spec;
20pub mod tools;
21pub mod traits;
22
23#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
24#[pyclass(from_py_object)]
25pub enum Model {
26    Undefined,
27}
28
29impl Model {
30    pub fn as_str(&self) -> &str {
31        match self {
32            Model::Undefined => "undefined",
33        }
34    }
35
36    pub fn from_string(s: &str) -> Result<Self, TypeError> {
37        match s.to_lowercase().as_str() {
38            "undefined" => Ok(Model::Undefined),
39            _ => Err(TypeError::UnknownModelError(s.to_string())),
40        }
41    }
42}
43
44pub enum Common {
45    Undefined,
46}
47
48impl Common {
49    pub fn as_str(&self) -> &str {
50        match self {
51            Common::Undefined => "undefined",
52        }
53    }
54}
55
56impl Display for Common {
57    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
58        write!(f, "{}", self.as_str())
59    }
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
63#[pyclass(from_py_object, eq, eq_int)]
64pub enum Provider {
65    OpenAI,
66    Gemini,
67    Google,
68    Vertex,
69    Anthropic,
70    GoogleAdk,
71    Undefined, // Added Undefined for better error handling
72}
73
74impl Provider {
75    pub fn from_string(s: &str) -> Result<Self, TypeError> {
76        match s.to_lowercase().as_str() {
77            "openai" => Ok(Provider::OpenAI),
78            "gemini" => Ok(Provider::Gemini),
79            "google" => Ok(Provider::Google),
80            "vertex" => Ok(Provider::Vertex),
81            "anthropic" => Ok(Provider::Anthropic),
82            "google_adk" => Ok(Provider::GoogleAdk),
83            "undefined" => Ok(Provider::Undefined), // Handle undefined case
84            _ => Err(TypeError::UnknownProviderError(s.to_string())),
85        }
86    }
87
88    pub const DEFAULT_ENV_VAR: &'static str = "POTATO_HEAD_DEFAULT_PROVIDER";
89
90    /// Reads `POTATO_HEAD_DEFAULT_PROVIDER` from the environment.
91    ///
92    /// - Unset, empty, or whitespace-only -> `Ok(None)`
93    /// - Set and parseable -> `Ok(Some(provider))`
94    /// - Set but unparseable -> `Err(TypeError::UnknownProviderError(...))`
95    pub fn from_env_default() -> Result<Option<Provider>, TypeError> {
96        match std::env::var(Self::DEFAULT_ENV_VAR) {
97            Ok(s) => {
98                let trimmed = s.trim();
99                if trimmed.is_empty() {
100                    Ok(None)
101                } else {
102                    Provider::from_string(trimmed).map(Some)
103                }
104            }
105            Err(_) => Ok(None),
106        }
107    }
108
109    /// Resolve a provider from an optional explicit string, falling back to the env var,
110    /// then erroring with `MissingProviderError` if neither is available.
111    pub fn resolve(provided: Option<&str>) -> Result<Provider, TypeError> {
112        if let Some(s) = provided.map(str::trim).filter(|s| !s.is_empty()) {
113            return Provider::from_string(s);
114        }
115
116        Self::from_env_default()?.ok_or(TypeError::MissingProviderError(Self::DEFAULT_ENV_VAR))
117    }
118
119    /// PyO3-friendly resolution: accepts an optional `Bound<'_, PyAny>` (Provider enum or string).
120    /// `None` or `py.None()` triggers env-var fallback.
121    pub fn resolve_from_py(provider: Option<&Bound<'_, PyAny>>) -> Result<Provider, TypeError> {
122        match provider {
123            Some(p) if !p.is_none() => Self::extract_provider(p),
124            _ => Self::from_env_default()?
125                .ok_or(TypeError::MissingProviderError(Self::DEFAULT_ENV_VAR)),
126        }
127    }
128
129    /// Extract provider from a PyAny object
130    ///
131    /// # Arguments
132    /// * `provider` - PyAny object
133    ///
134    /// # Returns
135    /// * `Result<Provider, AgentError>` - Result
136    ///
137    /// # Errors
138    /// * `AgentError` - Error
139    pub fn extract_provider(provider: &Bound<'_, PyAny>) -> Result<Provider, TypeError> {
140        match provider.is_instance_of::<Provider>() {
141            true => Ok(provider.extract::<Provider>().inspect_err(|e| {
142                error!("Failed to extract provider: {}", e);
143            })?),
144            false => {
145                let provider = provider.extract::<String>().unwrap();
146                Ok(Provider::from_string(&provider).inspect_err(|e| {
147                    error!("Failed to convert string to provider: {}", e);
148                })?)
149            }
150        }
151    }
152
153    pub fn as_str(&self) -> &str {
154        match self {
155            Provider::OpenAI => "openai",
156            Provider::Gemini => "gemini",
157            Provider::Vertex => "vertex",
158            Provider::Google => "google",
159            Provider::Anthropic => "anthropic",
160            Provider::GoogleAdk => "google_adk",
161            Provider::Undefined => "undefined", // Added Undefined case
162        }
163    }
164}
165
166impl Display for Provider {
167    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
168        write!(f, "{}", self.as_str())
169    }
170}
171
172#[pyclass(from_py_object, eq, eq_int)]
173#[derive(Debug, PartialEq, Clone)]
174pub enum SaveName {
175    Prompt,
176}
177
178#[pymethods]
179impl SaveName {
180    #[staticmethod]
181    pub fn from_string(s: &str) -> Option<Self> {
182        match s {
183            "prompt" => Some(SaveName::Prompt),
184
185            _ => None,
186        }
187    }
188
189    pub fn as_string(&self) -> &str {
190        match self {
191            SaveName::Prompt => "prompt",
192        }
193    }
194
195    pub fn __str__(&self) -> String {
196        self.to_string()
197    }
198}
199
200impl Display for SaveName {
201    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
202        write!(f, "{}", self.as_string())
203    }
204}
205
206impl AsRef<Path> for SaveName {
207    fn as_ref(&self) -> &Path {
208        match self {
209            SaveName::Prompt => Path::new("prompt"),
210        }
211    }
212}
213
214// impl PathBuf: From<SaveName>
215impl From<SaveName> for PathBuf {
216    fn from(save_name: SaveName) -> Self {
217        PathBuf::from(save_name.as_ref())
218    }
219}
220
221/// A trait for structured output types that can be used with potatohead prompts agents and workflows.
222///
223/// # Example
224/// ```rust
225/// use potatohead_macro::StructureOutput;
226/// use serde::{Serialize, Deserialize};
227/// use schemars::JsonSchema;
228///
229/// #[derive(Serialize, Deserialize, JsonSchema)]
230/// struct MyOutput {
231///     message: String,
232///     value: i32,
233/// }
234///
235/// impl StructuredOutput for MyOutput {}
236///
237/// let schema = MyOutput::get_structured_output_schema();
238/// ```
239pub trait StructuredOutput: for<'de> serde::Deserialize<'de> + JsonSchema {
240    fn type_name() -> &'static str {
241        type_name::<Self>().rsplit("::").next().unwrap_or("Unknown")
242    }
243
244    /// Validates and deserializes a JSON value into its struct type.
245    ///
246    /// # Arguments
247    /// * `value` - The JSON value to deserialize
248    ///
249    /// # Returns
250    /// * `Result<Self, serde_json::Error>` - The deserialized value or error
251    fn model_validate_json_value(value: &Value) -> Result<Self, serde_json::Error> {
252        match &value {
253            Value::String(json_str) => Self::model_validate_json_str(json_str),
254            Value::Object(_) => {
255                // Content is already a JSON object
256                serde_json::from_value(value.clone())
257            }
258            _ => {
259                // If the value is not a string or object, we cannot deserialize it
260                Err(Error::custom("Expected a JSON string or object"))
261            }
262        }
263    }
264
265    fn model_validate_json_str(value: &str) -> Result<Self, serde_json::Error> {
266        serde_json::from_str(value)
267    }
268
269    /// Generates an OpenAI-compatible JSON schema.
270    ///
271    /// # Returns
272    /// * `Value` - The JSON schema wrapped in OpenAI's format
273    fn get_structured_output_schema() -> Value {
274        let schema = ::schemars::schema_for!(Self);
275        schema.into()
276    }
277    // add fallback parsing logic
278}
279
280#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
281#[pyclass(from_py_object)]
282pub enum SettingsType {
283    GoogleChat,
284    OpenAIChat,
285    ModelSettings,
286    Anthropic,
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use std::sync::Mutex;
293
294    static ENV_LOCK: Mutex<()> = Mutex::new(());
295
296    fn with_env_var<F: FnOnce()>(value: Option<&str>, f: F) {
297        let _guard = ENV_LOCK.lock().unwrap();
298        let prev = std::env::var(Provider::DEFAULT_ENV_VAR).ok();
299        match value {
300            Some(v) => std::env::set_var(Provider::DEFAULT_ENV_VAR, v),
301            None => std::env::remove_var(Provider::DEFAULT_ENV_VAR),
302        }
303        f();
304        match prev {
305            Some(v) => std::env::set_var(Provider::DEFAULT_ENV_VAR, v),
306            None => std::env::remove_var(Provider::DEFAULT_ENV_VAR),
307        }
308    }
309
310    #[test]
311    fn test_provider_google_adk_round_trip() {
312        let p = Provider::from_string("google_adk").unwrap();
313        assert_eq!(p, Provider::GoogleAdk);
314        assert_eq!(p.as_str(), "google_adk");
315    }
316
317    #[test]
318    fn test_provider_all_variants_round_trip() {
319        for (s, variant) in [
320            ("openai", Provider::OpenAI),
321            ("gemini", Provider::Gemini),
322            ("google", Provider::Google),
323            ("vertex", Provider::Vertex),
324            ("anthropic", Provider::Anthropic),
325            ("google_adk", Provider::GoogleAdk),
326            ("undefined", Provider::Undefined),
327        ] {
328            let parsed = Provider::from_string(s).unwrap();
329            assert_eq!(parsed, variant);
330            assert_eq!(parsed.as_str(), s);
331        }
332    }
333
334    #[test]
335    fn test_provider_unknown_string_errors() {
336        assert!(Provider::from_string("not_a_provider").is_err());
337    }
338
339    #[test]
340    fn resolve_explicit_wins_over_env() {
341        with_env_var(Some("anthropic"), || {
342            let p = Provider::resolve(Some("openai")).unwrap();
343            assert_eq!(p, Provider::OpenAI);
344        });
345    }
346
347    #[test]
348    fn resolve_uses_env_when_explicit_missing() {
349        with_env_var(Some("openai"), || {
350            let p = Provider::resolve(None).unwrap();
351            assert_eq!(p, Provider::OpenAI);
352        });
353    }
354
355    #[test]
356    fn resolve_uses_env_when_explicit_empty_string() {
357        with_env_var(Some("openai"), || {
358            let p = Provider::resolve(Some("   ")).unwrap();
359            assert_eq!(p, Provider::OpenAI);
360        });
361    }
362
363    #[test]
364    fn resolve_empty_env_treated_as_unset() {
365        with_env_var(Some("   "), || {
366            let err = Provider::resolve(None).unwrap_err();
367            assert!(matches!(err, TypeError::MissingProviderError(_)));
368        });
369    }
370
371    #[test]
372    fn resolve_invalid_env_is_hard_error() {
373        with_env_var(Some("not_a_provider"), || {
374            let err = Provider::resolve(None).unwrap_err();
375            assert!(matches!(err, TypeError::UnknownProviderError(_)));
376        });
377    }
378
379    #[test]
380    fn resolve_no_explicit_no_env_returns_missing() {
381        with_env_var(None, || {
382            let err = Provider::resolve(None).unwrap_err();
383            match err {
384                TypeError::MissingProviderError(name) => {
385                    assert_eq!(name, "POTATO_HEAD_DEFAULT_PROVIDER");
386                }
387                other => panic!("expected MissingProviderError, got {:?}", other),
388            }
389        });
390    }
391
392    #[test]
393    fn from_env_default_unset_returns_none() {
394        with_env_var(None, || {
395            assert!(Provider::from_env_default().unwrap().is_none());
396        });
397    }
398
399    #[test]
400    fn from_env_default_set_and_parseable() {
401        with_env_var(Some("gemini"), || {
402            assert_eq!(
403                Provider::from_env_default().unwrap(),
404                Some(Provider::Gemini)
405            );
406        });
407    }
408
409    #[test]
410    fn from_env_default_set_but_invalid_errors() {
411        with_env_var(Some("garbage"), || {
412            assert!(Provider::from_env_default().is_err());
413        });
414    }
415}
potato_type/lib.rs

potato_type/
lib.rs
