Skip to main content

soma_provider_adapters/
python.rs

1//! The generic Python (/ LangChain / LlamaIndex) provider kind: introspects
2//! and executes a drop-in `.py` provider through a bounded Python sidecar
3//! running `python_bridge::PYTHON_BRIDGE`. Ported from
4//! `soma-service::providers::python`.
5//!
6//! `load_python_catalog` here applies only *generic* manifest validation
7//! (`soma_provider_core::validate_provider_manifest_value`) — Soma's own CLI
8//! reserved-command / env-prefix policy layer is applied downstream by the
9//! host's provider registry when it builds every provider's catalog
10//! (regardless of kind), so dropping the redundant Soma-specific pre-check
11//! here does not weaken overall enforcement — see the PR10 deviation notes.
12
13use std::{
14    io::{Read, Write},
15    path::{Path, PathBuf},
16    process::{Command as StdCommand, Stdio as StdStdio},
17    sync::Arc,
18    thread,
19    time::{Duration, Instant},
20};
21
22use async_trait::async_trait;
23use serde_json::{json, Value};
24use soma_provider_core::{
25    validate_provider_manifest_value, Provider, ProviderCall, ProviderCatalog, ProviderError,
26    ProviderOutput, ProviderTool,
27};
28use tokio::time::Instant as TokioInstant;
29
30use crate::{
31    error::{redact_public, SidecarError},
32    python_bridge::PYTHON_BRIDGE,
33    sidecar::{
34        collect_provider_env, output_exceeded_message, resolve_sidecar_command,
35        run_bounded_sidecar, sidecar_base_env,
36    },
37};
38
39const DEFAULT_TIMEOUT_MS: u64 = 10_000;
40const DEFAULT_MAX_INPUT_BYTES: usize = 64 * 1024;
41const DEFAULT_MAX_OUTPUT_BYTES: usize = 256 * 1024;
42
43#[derive(Clone)]
44pub struct PythonProvider {
45    path: PathBuf,
46    catalog: ProviderCatalog,
47    env_prefix: String,
48}
49
50impl PythonProvider {
51    pub fn new(path: PathBuf, catalog: ProviderCatalog, env_prefix: impl Into<String>) -> Self {
52        Self {
53            path,
54            catalog,
55            env_prefix: env_prefix.into(),
56        }
57    }
58
59    pub fn arc(
60        path: PathBuf,
61        catalog: ProviderCatalog,
62        env_prefix: impl Into<String>,
63    ) -> Arc<Self> {
64        Arc::new(Self::new(path, catalog, env_prefix))
65    }
66}
67
68#[async_trait]
69impl Provider for PythonProvider {
70    fn catalog(&self) -> ProviderCatalog {
71        self.catalog.clone()
72    }
73
74    async fn call(&self, call: ProviderCall) -> Result<ProviderOutput, ProviderError> {
75        let tool = self.tool(&call)?;
76        let runtime = PythonRuntime::from_tool(&self.catalog, tool, &call, &self.env_prefix)?;
77        let source = self.path.display().to_string();
78        let input = python_execution_payload(&self.path, &call, &runtime.env).map_err(|error| {
79            ProviderError::execution(&self.catalog.provider.name, "", error)
80                .with_provider_kind(self.catalog.provider.kind.as_str())
81                .with_source(source.clone())
82                .with_phase("input-serialization")
83        })?;
84
85        if input.len() > runtime.max_input_bytes {
86            return Err(ProviderError::validation(
87                &self.catalog.provider.name,
88                &call.action,
89                "python_input_too_large",
90                format!(
91                    "Python provider input exceeds {} bytes",
92                    runtime.max_input_bytes
93                ),
94            )
95            .with_provider_kind(self.catalog.provider.kind.as_str())
96            .with_source(source)
97            .with_phase("input-validation"));
98        }
99
100        let started = TokioInstant::now();
101        let sidecar = match run_bounded_sidecar(
102            &runtime.command,
103            &["-c", PYTHON_BRIDGE],
104            runtime.env,
105            &input,
106            runtime.timeout_ms,
107            runtime.max_output_bytes,
108        )
109        .await
110        {
111            Ok(sidecar) => sidecar,
112            Err(SidecarError::Timeout) => {
113                return Err(ProviderError::new(
114                    "python_provider_timeout",
115                    &self.catalog.provider.name,
116                    Some(call.action.clone()),
117                    format!("Python provider exceeded {}ms timeout", runtime.timeout_ms),
118                    "Increase tool.limits.timeout_ms or fix the Python provider handler.",
119                )
120                .with_provider_kind(self.catalog.provider.kind.as_str())
121                .with_source(source)
122                .with_phase("execution"));
123            }
124            Err(error) => {
125                return Err(ProviderError::execution(
126                    &self.catalog.provider.name,
127                    call.action.clone(),
128                    error,
129                )
130                .with_provider_kind(self.catalog.provider.kind.as_str())
131                .with_source(source)
132                .with_phase("execution"));
133            }
134        };
135        let output = sidecar.output;
136
137        tracing::debug!(
138            provider = %self.catalog.provider.name,
139            action = %call.action,
140            elapsed_ms = started.elapsed().as_millis(),
141            "Python provider sidecar completed"
142        );
143
144        if sidecar.stdout_exceeded || sidecar.stderr_exceeded {
145            let stream = if sidecar.stdout_exceeded {
146                "stdout"
147            } else {
148                "stderr"
149            };
150            return Err(ProviderError::validation(
151                &self.catalog.provider.name,
152                &call.action,
153                "python_output_too_large",
154                output_exceeded_message(stream, runtime.max_output_bytes),
155            )
156            .with_provider_kind(self.catalog.provider.kind.as_str())
157            .with_source(source)
158            .with_phase("output-validation"));
159        }
160        if !output.status.success() {
161            let stderr = String::from_utf8_lossy(&output.stderr);
162            let code = if stderr.contains("python_provider_unserializable_output") {
163                "python_provider_unserializable_output"
164            } else {
165                "python_provider_failed"
166            };
167            return Err(ProviderError::new(
168                code,
169                &self.catalog.provider.name,
170                Some(call.action),
171                format!("Python provider failed: {}", redact_public(&stderr)),
172                "Fix the Python provider handler and retry.",
173            )
174            .with_provider_kind(self.catalog.provider.kind.as_str())
175            .with_source(source)
176            .with_phase("execution"));
177        }
178
179        let value = serde_json::from_slice(&output.stdout).map_err(|error| {
180            ProviderError::validation(
181                &self.catalog.provider.name,
182                &call.action,
183                "python_invalid_json_output",
184                error.to_string(),
185            )
186            .with_provider_kind(self.catalog.provider.kind.as_str())
187            .with_source(source)
188            .with_phase("output-validation")
189        })?;
190        Ok(ProviderOutput::json(value))
191    }
192}
193
194fn python_execution_payload(
195    path: &Path,
196    call: &ProviderCall,
197    env: &[(String, String)],
198) -> Result<Vec<u8>, serde_json::Error> {
199    let mut payload = serde_json::to_value(crate::sidecar::ExecutionEnvelope::new(call))?;
200    if let Some(object) = payload.as_object_mut() {
201        let env_keys: Vec<&str> = env.iter().map(|(key, _)| key.as_str()).collect();
202        object.insert("mode".to_owned(), json!("call"));
203        object.insert("path".to_owned(), json!(path.to_path_buf()));
204        object.insert("env_keys".to_owned(), json!(env_keys));
205    }
206    serde_json::to_vec(&payload)
207}
208
209impl PythonProvider {
210    fn tool(&self, call: &ProviderCall) -> Result<&ProviderTool, ProviderError> {
211        self.catalog
212            .tools
213            .iter()
214            .find(|tool| tool.name == call.action)
215            .ok_or_else(|| {
216                ProviderError::validation(
217                    &self.catalog.provider.name,
218                    &call.action,
219                    "unknown_python_action",
220                    format!("Python provider has no action `{}`", call.action),
221                )
222            })
223    }
224}
225
226/// Introspects a `.py` provider file by importing it (in "catalog" mode) in
227/// a bounded sidecar and validating the resulting manifest against the
228/// generic provider-core contract. Callers that layer additional product
229/// policy on top of every provider's catalog (e.g. Soma's reserved
230/// CLI-command / env-prefix checks) apply it to the returned catalog
231/// themselves — see the module docs above.
232pub fn load_python_catalog(path: &Path, env_prefix: &str) -> Result<ProviderCatalog, String> {
233    let runtime = PythonRuntime::for_catalog(env_prefix);
234    let input = serde_json::to_vec(&json!({
235        "mode": "catalog",
236        "path": path,
237    }))
238    .map_err(|error| error.to_string())?;
239    let output = run_catalog_sidecar(&runtime, &input)?;
240    let value: Value = serde_json::from_slice(&output).map_err(|error| error.to_string())?;
241    validate_provider_manifest_value(&value).map_err(|error| error.to_string())
242}
243
244struct PythonRuntime {
245    command: String,
246    env: Vec<(String, String)>,
247    timeout_ms: u64,
248    max_input_bytes: usize,
249    max_output_bytes: usize,
250}
251
252impl PythonRuntime {
253    fn for_catalog(env_prefix: &str) -> Self {
254        let prefix = env_prefix.trim_matches('_').to_ascii_uppercase();
255        let timeout_var = format!("{prefix}_PYTHON_CATALOG_TIMEOUT_MS");
256        let timeout_ms = match std::env::var(&timeout_var) {
257            Ok(value) => value.parse().unwrap_or_else(|error| {
258                tracing::warn!(
259                    variable = %timeout_var,
260                    value,
261                    error = %error,
262                    "invalid provider catalog timeout env var; falling back to the default"
263                );
264                DEFAULT_TIMEOUT_MS
265            }),
266            Err(_) => DEFAULT_TIMEOUT_MS,
267        };
268        Self {
269            command: std::env::var(format!("{prefix}_PYTHON_COMMAND"))
270                .unwrap_or_else(|_| default_python_command().to_owned()),
271            env: Vec::new(),
272            timeout_ms,
273            max_input_bytes: DEFAULT_MAX_INPUT_BYTES,
274            max_output_bytes: DEFAULT_MAX_OUTPUT_BYTES,
275        }
276    }
277
278    fn from_tool(
279        catalog: &ProviderCatalog,
280        tool: &ProviderTool,
281        call: &ProviderCall,
282        env_prefix: &str,
283    ) -> Result<Self, ProviderError> {
284        let provider_meta = catalog.meta.get("python");
285        let tool_meta = tool.meta.get("python");
286        let meta_field = |key: &str| {
287            tool_meta
288                .and_then(|value| value.get(key))
289                .or_else(|| provider_meta.and_then(|value| value.get(key)))
290        };
291        let command = meta_field("command")
292            .and_then(Value::as_str)
293            .map(str::to_owned)
294            .or_else(|| {
295                std::env::var(format!(
296                    "{}_PYTHON_COMMAND",
297                    env_prefix.trim_matches('_').to_ascii_uppercase()
298                ))
299                .ok()
300            })
301            .unwrap_or_else(|| default_python_command().to_owned());
302        let timeout_ms = tool
303            .limits
304            .as_ref()
305            .and_then(|limits| limits.timeout_ms)
306            .or_else(|| meta_field("timeout_ms").and_then(Value::as_u64))
307            .unwrap_or(DEFAULT_TIMEOUT_MS);
308        let max_input_bytes = tool
309            .limits
310            .as_ref()
311            .and_then(|limits| limits.max_input_bytes)
312            .unwrap_or(DEFAULT_MAX_INPUT_BYTES);
313        let max_output_bytes = tool
314            .limits
315            .as_ref()
316            .and_then(|limits| limits.max_response_bytes)
317            .unwrap_or(DEFAULT_MAX_OUTPUT_BYTES);
318        Ok(Self {
319            command,
320            env: collect_provider_env(
321                &catalog.env,
322                &tool.env,
323                env_prefix,
324                &call.provider,
325                &call.action,
326            )?,
327            timeout_ms,
328            max_input_bytes,
329            max_output_bytes,
330        })
331    }
332}
333
334#[cfg(windows)]
335fn default_python_command() -> &'static str {
336    "python"
337}
338
339#[cfg(not(windows))]
340fn default_python_command() -> &'static str {
341    "python3"
342}
343
344fn run_catalog_sidecar(runtime: &PythonRuntime, input: &[u8]) -> Result<Vec<u8>, String> {
345    let mut command = StdCommand::new(resolve_sidecar_command(&runtime.command));
346    command
347        .args(["-c", PYTHON_BRIDGE])
348        .env_clear()
349        .stdin(StdStdio::piped())
350        .stdout(StdStdio::piped())
351        .stderr(StdStdio::piped());
352    for (key, value) in sidecar_base_env() {
353        command.env(key, value);
354    }
355    let mut child = command.spawn().map_err(|error| error.to_string())?;
356    let stdout = child
357        .stdout
358        .take()
359        .ok_or_else(|| "Python provider catalog stdout pipe was not captured".to_owned())?;
360    let stderr = child
361        .stderr
362        .take()
363        .ok_or_else(|| "Python provider catalog stderr pipe was not captured".to_owned())?;
364    let max_output_bytes = runtime.max_output_bytes;
365    let stdout_task = thread::spawn(move || read_bounded_sync(stdout, max_output_bytes));
366    let stderr_task = thread::spawn(move || read_bounded_sync(stderr, max_output_bytes));
367
368    if let Some(mut stdin) = child.stdin.take() {
369        stdin.write_all(input).map_err(|error| error.to_string())?;
370    }
371    let deadline = Instant::now() + Duration::from_millis(runtime.timeout_ms);
372    loop {
373        if let Some(status) = child.try_wait().map_err(|error| error.to_string())? {
374            let (stdout, stdout_exceeded) = stdout_task
375                .join()
376                .map_err(|_| "Python provider catalog stdout reader panicked".to_owned())?
377                .map_err(|error| error.to_string())?;
378            let (stderr, stderr_exceeded) = stderr_task
379                .join()
380                .map_err(|_| "Python provider catalog stderr reader panicked".to_owned())?
381                .map_err(|error| error.to_string())?;
382            if stdout_exceeded || stderr_exceeded {
383                let stream = if stdout_exceeded { "stdout" } else { "stderr" };
384                return Err(format!(
385                    "Python provider catalog {}",
386                    output_exceeded_message(stream, runtime.max_output_bytes)
387                ));
388            }
389            if !status.success() {
390                return Err(format!(
391                    "Python provider catalog failed: {}",
392                    redact_public(&String::from_utf8_lossy(&stderr))
393                ));
394            }
395            return Ok(stdout);
396        }
397        if Instant::now() >= deadline {
398            let _ = child.kill();
399            let _ = child.wait();
400            return Err(format!(
401                "Python provider catalog exceeded {}ms timeout",
402                runtime.timeout_ms
403            ));
404        }
405        std::thread::sleep(Duration::from_millis(10));
406    }
407}
408
409fn read_bounded_sync<R: Read>(
410    mut reader: R,
411    max_output_bytes: usize,
412) -> std::io::Result<(Vec<u8>, bool)> {
413    let mut bytes = Vec::new();
414    let mut exceeded = false;
415    let mut chunk = [0u8; 8192];
416    loop {
417        let read = reader.read(&mut chunk)?;
418        if read == 0 {
419            return Ok((bytes, exceeded));
420        }
421        let remaining = max_output_bytes.saturating_sub(bytes.len());
422        if remaining >= read && !exceeded {
423            bytes.extend_from_slice(&chunk[..read]);
424        } else {
425            exceeded = true;
426            if remaining > 0 {
427                bytes.extend_from_slice(&chunk[..remaining]);
428            }
429        }
430    }
431}
432
433#[cfg(test)]
434#[path = "python_tests.rs"]
435mod tests;