From b6d285a79cefcff52b4d137216630aa9202bd482 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Fri, 27 Feb 2026 18:31:11 -0500 Subject: [PATCH] Remove dead LLM crate code and add --schema to ullm prompt Remove unused items that have no callers outside their own tests: middleware wrap_stream_with_middleware/process_stream_event, common ApiMessage/send_and_read_body, types ProviderEvent variant, lib CancellationToken re-export, tools execute_all_tools, and catalog get_latest_model. Add --schema/-S flag to `ullm prompt` so generate_object() and stream_object() are exercisable end-to-end. Includes unit test for invalid JSON rejection and two ignored integration tests. Co-Authored-By: Claude Opus 4.6 --- crates/llm/src/bin/ullm.rs | 105 +++++++++++++--- crates/llm/src/catalog.rs | 42 ------- crates/llm/src/lib.rs | 2 - crates/llm/src/middleware.rs | 30 +---- crates/llm/src/providers/common.rs | 54 -------- crates/llm/src/tools.rs | 193 ----------------------------- crates/llm/src/types.rs | 3 - 7 files changed, 91 insertions(+), 338 deletions(-) diff --git a/crates/llm/src/bin/ullm.rs b/crates/llm/src/bin/ullm.rs index 7cd0989a4..aa9ffb4a7 100644 --- a/crates/llm/src/bin/ullm.rs +++ b/crates/llm/src/bin/ullm.rs @@ -40,6 +40,10 @@ enum Command { #[arg(short, long)] usage: bool, + /// JSON schema for structured output (inline JSON string) + #[arg(short = 'S', long)] + schema: Option, + /// key=value options (temperature, `max_tokens`, `top_p`) #[arg(short, long, value_parser = parse_option)] option: Vec<(String, String)>, @@ -180,6 +184,7 @@ struct PromptArgs { system: Option, no_stream: bool, usage: bool, + schema: Option, option: Vec<(String, String)>, } @@ -206,23 +211,48 @@ async fn run_prompt(args: PromptArgs) -> Result<()> { } params = apply_options(params, &args.option)?; - if args.no_stream { - let result = generate::generate(params).await?; - print!("{}", result.text()); - if args.usage { - print_usage(result.usage()); - } - } else { - let mut stream_result = generate::stream(params).await?; - while let Some(event) = stream_result.next().await { - if let llm::types::StreamEvent::TextDelta { delta, .. } = event? { - print!("{delta}"); + let schema: Option = match &args.schema { + Some(s) => Some(serde_json::from_str(s).context("--schema must be valid JSON")?), + None => None, + }; + + match (args.no_stream, schema) { + (true, Some(schema)) => { + let result = generate::generate_object(params, schema).await?; + let object = result.output.as_ref().unwrap_or(&serde_json::Value::Null); + println!("{}", serde_json::to_string_pretty(object)?); + if args.usage { + print_usage(result.usage()); } } - println!(); - if args.usage { - if let Some(response) = stream_result.response() { - print_usage(&response.usage); + (true, None) => { + let result = generate::generate(params).await?; + print!("{}", result.text()); + if args.usage { + print_usage(result.usage()); + } + } + (false, Some(schema)) => { + let mut stream_result = generate::stream_object(params, schema).await?; + while let Some(event) = stream_result.next().await { + event?; + } + if let Some(object) = stream_result.object() { + println!("{}", serde_json::to_string_pretty(object)?); + } + } + (false, None) => { + let mut stream_result = generate::stream(params).await?; + while let Some(event) = stream_result.next().await { + if let llm::types::StreamEvent::TextDelta { delta, .. } = event? { + print!("{delta}"); + } + } + println!(); + if args.usage { + if let Some(response) = stream_result.response() { + print_usage(&response.usage); + } } } } @@ -265,6 +295,7 @@ async fn main() -> Result<()> { system, no_stream, usage, + schema, option, } => { run_prompt(PromptArgs { @@ -273,6 +304,7 @@ async fn main() -> Result<()> { system, no_stream, usage, + schema, option, }) .await?; @@ -566,4 +598,47 @@ mod tests { .success() .stderr(predicate::str::contains("Tokens:")); } + + #[test] + fn prompt_schema_rejects_invalid_json() { + ullm() + .args(["--no-dotenv", "prompt", "--no-stream", "-m", "test-model", "--schema", "not json", "hello"]) + .assert() + .failure() + .stderr(predicate::str::contains("--schema must be valid JSON")); + } + + #[test] + #[ignore = "requires API key"] + fn prompt_schema_no_stream_generates_json() { + let assert = ullm() + .args([ + "prompt", "--no-stream", "-m", "claude-sonnet-4-5", + "--schema", r#"{"type":"object","properties":{"greeting":{"type":"string"}},"required":["greeting"]}"#, + "Return a JSON object with a greeting field set to hello", + ]) + .assert() + .success(); + + let stdout = String::from_utf8(assert.get_output().stdout.clone()).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(stdout.trim()).expect("stdout should be valid JSON"); + assert!(parsed.get("greeting").is_some(), "expected 'greeting' key in output"); + } + + #[test] + #[ignore = "requires API key"] + fn prompt_schema_stream_generates_json() { + let assert = ullm() + .args([ + "prompt", "-m", "claude-sonnet-4-5", + "--schema", r#"{"type":"object","properties":{"greeting":{"type":"string"}},"required":["greeting"]}"#, + "Return a JSON object with a greeting field set to hello", + ]) + .assert() + .success(); + + let stdout = String::from_utf8(assert.get_output().stdout.clone()).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(stdout.trim()).expect("stdout should be valid JSON"); + assert!(parsed.get("greeting").is_some(), "expected 'greeting' key in output"); + } } diff --git a/crates/llm/src/catalog.rs b/crates/llm/src/catalog.rs index f1b08a9be..751e1a85a 100644 --- a/crates/llm/src/catalog.rs +++ b/crates/llm/src/catalog.rs @@ -26,19 +26,6 @@ pub fn list_models(provider: Option<&str>) -> Vec { ) } -/// Get the latest/best model for a provider, optionally filtered by capability (Section 2.9). -#[must_use] -pub fn get_latest_model(provider: &str, capability: Option<&str>) -> Option { - let mut models = BUILT_IN_MODELS.iter().filter(|m| m.provider == provider); - - match capability { - Some("reasoning") => models.find(|m| m.supports_reasoning).cloned(), - Some("vision") => models.find(|m| m.supports_vision).cloned(), - Some("tools") => models.find(|m| m.supports_tools).cloned(), - _ => models.next().cloned(), - } -} - #[cfg(test)] mod tests { use super::*; @@ -94,35 +81,6 @@ mod tests { assert!(unknown.is_empty()); } - #[test] - fn get_latest_model_returns_first_for_provider() { - let model = get_latest_model("anthropic", None).unwrap(); - assert_eq!(model.id, "claude-opus-4-6"); - - let model = get_latest_model("openai", None).unwrap(); - assert_eq!(model.id, "gpt-5.2"); - - let model = get_latest_model("gemini", None).unwrap(); - assert_eq!(model.id, "gemini-3.1-pro-preview"); - } - - #[test] - fn get_latest_model_filtered_by_capability() { - let model = get_latest_model("anthropic", Some("reasoning")).unwrap(); - assert!(model.supports_reasoning); - - let model = get_latest_model("openai", Some("vision")).unwrap(); - assert!(model.supports_vision); - - let model = get_latest_model("gemini", Some("tools")).unwrap(); - assert!(model.supports_tools); - } - - #[test] - fn get_latest_model_returns_none_for_unknown_provider() { - assert!(get_latest_model("unknown", None).is_none()); - } - #[test] fn model_info_costs() { let claude = get_model_info("claude-opus-4-6").unwrap(); diff --git a/crates/llm/src/lib.rs b/crates/llm/src/lib.rs index e56d1568a..c7927b25b 100644 --- a/crates/llm/src/lib.rs +++ b/crates/llm/src/lib.rs @@ -9,7 +9,5 @@ pub mod generate; pub mod catalog; pub mod providers; -pub use tokio_util::sync::CancellationToken; - // Re-export module-level default client helpers (Section 2.5). pub use generate::set_default_client; diff --git a/crates/llm/src/middleware.rs b/crates/llm/src/middleware.rs index af0e3c272..11feacda9 100644 --- a/crates/llm/src/middleware.rs +++ b/crates/llm/src/middleware.rs @@ -1,7 +1,6 @@ use crate::error::SdkError; use crate::provider::StreamEventStream; -use crate::types::{Request, Response, StreamEvent}; -use futures::StreamExt; +use crate::types::{Request, Response}; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -21,10 +20,6 @@ pub type NextStreamFn = Arc< >; /// Middleware for intercepting `complete()` and streaming calls (Section 2.3). -/// -/// Implement `handle_complete` for blocking requests and `handle_stream` for -/// streaming requests. Override `process_stream` to observe or transform -/// individual stream events without replacing the entire stream handler. #[async_trait::async_trait] pub trait Middleware: Send + Sync { async fn handle_complete( @@ -38,27 +33,4 @@ pub trait Middleware: Send + Sync { request: Request, next: NextStreamFn, ) -> Result; - - /// Process an individual stream event. Override to observe or transform - /// events as they pass through the middleware. The default implementation - /// passes events through unchanged. - /// - /// # Errors - /// - /// Returns `SdkError` if the incoming event is an error or if processing fails. - fn process_stream_event( - &self, - event: Result, - ) -> Result { - event - } -} - -/// Wrap a `StreamEventStream` so that each event passes through a middleware's -/// `process_stream_event` method. -pub fn wrap_stream_with_middleware( - stream: StreamEventStream, - middleware: Arc, -) -> StreamEventStream { - Box::pin(stream.map(move |event| middleware.process_stream_event(event))) } diff --git a/crates/llm/src/providers/common.rs b/crates/llm/src/providers/common.rs index 662bd6659..548b81c15 100644 --- a/crates/llm/src/providers/common.rs +++ b/crates/llm/src/providers/common.rs @@ -3,12 +3,6 @@ use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; use crate::error::{error_from_status_code, SdkError}; use crate::types::{Message, RateLimitInfo, Role}; -#[derive(serde::Serialize)] -pub struct ApiMessage { - pub role: String, - pub content: String, -} - /// Parse an error response body, extracting the message and error code. /// /// `error_code_field` is the JSON field name for the error code (e.g. "type" or "status"). @@ -36,54 +30,6 @@ pub fn parse_error_body( ) } -/// Send an HTTP request and read the response body. -/// -/// Returns an error on non-success status. -/// -/// # Errors -/// -/// Returns `SdkError::Network` on connection failure or `SdkError::Provider` on non-success status. -pub async fn send_and_read_body( - request: reqwest::RequestBuilder, - provider: &str, - error_code_field: &str, -) -> Result { - let http_resp = request.send().await.map_err(|e| { - if e.is_timeout() { - SdkError::RequestTimeout { - message: format!("{provider}: {e}"), - } - } else { - SdkError::Network { - message: e.to_string(), - } - } - })?; - - let status = http_resp.status(); - let retry_after = parse_retry_after(http_resp.headers()); - let body = http_resp - .text() - .await - .map_err(|e| SdkError::Network { - message: e.to_string(), - })?; - - if !status.is_success() { - let (msg, code, raw) = parse_error_body(&body, error_code_field); - return Err(error_from_status_code( - status.as_u16(), - msg, - provider.to_string(), - code, - raw, - retry_after, - )); - } - - Ok(body) -} - /// Extract system and developer messages from a message list. /// /// Returns the joined system prompt and the remaining messages. diff --git a/crates/llm/src/tools.rs b/crates/llm/src/tools.rs index 9261e60b7..332b3c6f0 100644 --- a/crates/llm/src/tools.rs +++ b/crates/llm/src/tools.rs @@ -115,71 +115,6 @@ pub fn validate_tool_name(name: &str) -> Result<(), String> { Ok(()) } -/// Execute all tool calls concurrently (Section 5.7). -/// Returns results in the same order as the input calls. -/// -/// # Panics -/// -/// Panics if a matched tool has `is_active() == true` but its `execute` handler is `None`. -pub async fn execute_all_tools( - tools: &[&Tool], - tool_calls: &[ToolCall], - messages: &[Message], - abort_signal: Option<&CancellationToken>, -) -> Vec { - use futures::future::join_all; - - let futures: Vec<_> = tool_calls - .iter() - .map(|call| { - let tool = tools.iter().find(|t| t.definition.name == call.name).copied(); - let call_id = call.id.clone(); - let call_name = call.name.clone(); - let args = call.arguments.clone(); - let ctx = ToolContext { - tool_call_id: call_id.clone(), - messages: messages.to_vec(), - abort_signal: abort_signal.cloned(), - }; - - async move { - match tool { - Some(t) if t.execute.is_some() => { - let handler = t.execute.as_ref().unwrap(); - match handler(args, ctx).await { - Ok(result) => ToolResult { - tool_call_id: call_id, - content: result, - is_error: false, - image_data: None, - image_media_type: None, - }, - Err(err_msg) => ToolResult { - tool_call_id: call_id, - content: serde_json::Value::String(err_msg), - is_error: true, - image_data: None, - image_media_type: None, - }, - } - } - _ => ToolResult { - tool_call_id: call_id, - content: serde_json::Value::String(format!( - "Unknown tool: {call_name}" - )), - is_error: true, - image_data: None, - image_media_type: None, - }, - } - } - }) - .collect(); - - join_all(futures).await -} - /// A callback to repair invalid tool call arguments (Section 5.8). /// Receives the tool call and the validation error message, returns repaired arguments /// or an error if repair is not possible. @@ -401,134 +336,6 @@ mod tests { assert!(tool.is_active()); } - #[tokio::test] - async fn execute_all_tools_with_known_tools() { - let tools = [Tool::active( - "greet", - "Greet someone", - serde_json::json!({"type": "object", "properties": {"name": {"type": "string"}}}), - |args, _ctx| async move { - let name = args["name"].as_str().unwrap_or("world"); - Ok(serde_json::json!(format!("Hello, {}!", name))) - }, - )]; - - let calls = vec![ToolCall::new( - "call_1", - "greet", - serde_json::json!({"name": "Alice"}), - )]; - - let tool_refs: Vec<&Tool> = tools.iter().collect(); - let results = execute_all_tools(&tool_refs, &calls, &[], None).await; - assert_eq!(results.len(), 1); - assert_eq!(results[0].tool_call_id, "call_1"); - assert!(!results[0].is_error); - assert_eq!(results[0].content, serde_json::json!("Hello, Alice!")); - } - - #[tokio::test] - async fn execute_all_tools_with_unknown_tool() { - let tools = []; - - let calls = vec![ToolCall::new( - "call_1", - "nonexistent", - serde_json::json!({}), - )]; - - let tool_refs: Vec<&Tool> = tools.iter().collect(); - let results = execute_all_tools(&tool_refs, &calls, &[], None).await; - assert_eq!(results.len(), 1); - assert!(results[0].is_error); - assert!(results[0] - .content - .as_str() - .unwrap() - .contains("Unknown tool")); - } - - #[tokio::test] - async fn execute_all_tools_handler_error() { - let tools = [Tool::active( - "fail", - "Always fails", - serde_json::json!({"type": "object", "properties": {}}), - |_args, _ctx| async { Err("something went wrong".to_string()) }, - )]; - - let calls = vec![ToolCall::new( - "call_1", - "fail", - serde_json::json!({}), - )]; - - let tool_refs: Vec<&Tool> = tools.iter().collect(); - let results = execute_all_tools(&tool_refs, &calls, &[], None).await; - assert_eq!(results.len(), 1); - assert!(results[0].is_error); - assert_eq!( - results[0].content, - serde_json::json!("something went wrong") - ); - } - - #[tokio::test] - async fn execute_all_tools_concurrent_multiple() { - let tools = [Tool::active( - "tool_a", - "Tool A", - serde_json::json!({"type": "object", "properties": {}}), - |_args, _ctx| async { Ok(serde_json::json!("result_a")) }, - ), - Tool::active( - "tool_b", - "Tool B", - serde_json::json!({"type": "object", "properties": {}}), - |_args, _ctx| async { Ok(serde_json::json!("result_b")) }, - )]; - - let calls = vec![ - ToolCall::new("call_1", "tool_a", serde_json::json!({})), - ToolCall::new("call_2", "tool_b", serde_json::json!({})), - ]; - - let tool_refs: Vec<&Tool> = tools.iter().collect(); - let results = execute_all_tools(&tool_refs, &calls, &[], None).await; - assert_eq!(results.len(), 2); - assert_eq!(results[0].tool_call_id, "call_1"); - assert_eq!(results[0].content, serde_json::json!("result_a")); - assert_eq!(results[1].tool_call_id, "call_2"); - assert_eq!(results[1].content, serde_json::json!("result_b")); - } - - #[tokio::test] - async fn execute_all_tools_partial_failure() { - let tools = [Tool::active( - "succeed", - "Succeeds", - serde_json::json!({"type": "object", "properties": {}}), - |_args, _ctx| async { Ok(serde_json::json!("ok")) }, - ), - Tool::active( - "fail", - "Fails", - serde_json::json!({"type": "object", "properties": {}}), - |_args, _ctx| async { Err("boom".to_string()) }, - )]; - - let calls = vec![ - ToolCall::new("call_1", "succeed", serde_json::json!({})), - ToolCall::new("call_2", "fail", serde_json::json!({})), - ]; - - let tool_refs: Vec<&Tool> = tools.iter().collect(); - let results = execute_all_tools(&tool_refs, &calls, &[], None).await; - assert_eq!(results.len(), 2); - assert!(!results[0].is_error); - assert!(results[1].is_error); - } - #[test] #[should_panic(expected = "Invalid tool name")] fn passive_tool_panics_on_invalid_name() { diff --git a/crates/llm/src/types.rs b/crates/llm/src/types.rs index 8b6b28c04..1dbfd0ea1 100644 --- a/crates/llm/src/types.rs +++ b/crates/llm/src/types.rs @@ -556,9 +556,6 @@ pub enum StreamEvent { error: SdkError, raw: Option, }, - ProviderEvent { - raw: Option, - }, } impl StreamEvent {