Skip to main content

open_agent/client/
query.rs

1/// A pinned, boxed stream of content blocks from the model.
2///
3/// This type alias represents an asynchronous stream that yields `ContentBlock` items.
4/// Each item is wrapped in a `Result` to handle potential errors during streaming.
5///
6/// The stream is:
7/// - **Pinned** (`Pin<Box<...>>`): Required for safe async operations and self-referential types
8/// - **Boxed**: Allows dynamic dispatch and hides the concrete stream implementation
9/// - **Send**: Can be safely transferred between threads
10///
11/// # Content Blocks
12///
13/// The stream can yield several types of content blocks:
14///
15/// - **TextBlock**: Incremental text responses from the model
16/// - **ToolUseBlock**: Requests to execute a tool with specific parameters
17/// - **ToolResultBlock**: Results from tool execution (in manual mode)
18///
19/// # Error Handling
20///
21/// Errors in the stream indicate issues like:
22/// - Network failures or timeouts
23/// - Malformed SSE events
24/// - JSON parsing errors
25/// - API errors from the model provider
26///
27/// When an error occurs, the stream typically terminates. It's the caller's responsibility
28/// to handle errors appropriately.
29///
30/// # Examples
31///
32/// ```rust,no_run
33/// use open_agent::{query, AgentOptions, ContentBlock};
34/// use futures::StreamExt;
35///
36/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
37/// let options = AgentOptions::builder()
38///     .model("gpt-4")
39///     .api_key("sk-...")
40///     .build()?;
41///
42/// let mut stream = query("Hello!", &options).await?;
43///
44/// while let Some(result) = stream.next().await {
45///     match result {
46///         Ok(ContentBlock::Text(text)) => print!("{}", text.text),
47///         Ok(_) => {}, // Other block types
48///         Err(e) => eprintln!("Stream error: {}", e),
49///     }
50/// }
51/// # Ok(())
52/// # }
53/// ```
54pub type ContentStream = Pin<Box<dyn Stream<Item = Result<ContentBlock>> + Send>>;
55
56/// Simple query function for single-turn interactions without conversation history.
57///
58/// This is a stateless convenience function for simple queries that don't require
59/// multi-turn conversations. It creates a temporary HTTP client, sends a single
60/// prompt, and returns a stream of content blocks.
61///
62/// For multi-turn conversations or more control over the interaction, use [`Client`] instead.
63///
64/// # Parameters
65///
66/// - `prompt`: The user's message to send to the model
67/// - `options`: Configuration including model, API key, tools, etc.
68///
69/// # Returns
70///
71/// Returns a `ContentStream` that yields content blocks as they arrive from the model.
72/// The stream must be polled to completion to receive all blocks.
73///
74/// # Behavior
75///
76/// 1. Creates a temporary HTTP client with configured timeout
77/// 2. Builds message array (system prompt + user prompt)
78/// 3. Converts tools to OpenAI format if provided
79/// 4. Makes HTTP POST request to `/chat/completions`
80/// 5. Parses Server-Sent Events (SSE) response stream
81/// 6. Aggregates chunks into complete content blocks
82/// 7. Returns stream that yields blocks as they complete
83///
84/// # Error Handling
85///
86/// This function can return errors for:
87/// - HTTP client creation failures
88/// - Network errors during the request
89/// - API errors (authentication, invalid model, rate limits, etc.)
90/// - SSE parsing errors
91/// - JSON deserialization errors
92///
93/// # Performance Notes
94///
95/// - Creates a new HTTP client for each call (consider using `Client` for repeated queries)
96/// - Timeout is configurable via `AgentOptions::timeout` (default: 120 seconds)
97/// - Streaming begins immediately; no buffering of the full response
98///
99/// # Examples
100///
101/// ## Basic Usage
102///
103/// ```rust,no_run
104/// use open_agent::{query, AgentOptions};
105/// use futures::StreamExt;
106///
107/// #[tokio::main]
108/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
109///     let options = AgentOptions::builder()
110///         .system_prompt("You are a helpful assistant")
111///         .model("gpt-4")
112///         .api_key("sk-...")
113///         .build()?;
114///
115///     let mut stream = query("What's the capital of France?", &options).await?;
116///
117///     while let Some(block) = stream.next().await {
118///         match block? {
119///             open_agent::ContentBlock::Text(text) => {
120///                 print!("{}", text.text);
121///             }
122///             open_agent::ContentBlock::ToolUse(_)
123///             | open_agent::ContentBlock::ToolResult(_)
124///             | open_agent::ContentBlock::Image(_) => {}
125///         }
126///     }
127///
128///     Ok(())
129/// }
130/// ```
131///
132/// ## With Tools
133///
134/// ```rust,no_run
135/// use open_agent::{query, AgentOptions, Tool, ContentBlock};
136/// use futures::StreamExt;
137/// use serde_json::json;
138///
139/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
140/// let calculator = Tool::new(
141///     "calculator",
142///     "Performs calculations",
143///     json!({"type": "object"}),
144///     |input| Box::pin(async move { Ok(json!({"result": 42})) })
145/// );
146///
147/// let options = AgentOptions::builder()
148///     .model("gpt-4")
149///     .api_key("sk-...")
150///     .tools(vec![calculator])
151///     .build()?;
152///
153/// let mut stream = query("Calculate 2+2", &options).await?;
154///
155/// while let Some(block) = stream.next().await {
156///     match block? {
157///         ContentBlock::ToolUse(tool_use) => {
158///             println!("Model wants to use: {}", tool_use.name());
159///             // Note: You'll need to manually execute tools and continue
160///             // the conversation. For automatic execution, use Client.
161///         }
162///         ContentBlock::Text(text) => print!("{}", text.text),
163///         ContentBlock::ToolResult(_) | ContentBlock::Image(_) => {}
164///     }
165/// }
166/// # Ok(())
167/// # }
168/// ```
169///
170/// ## Error Handling
171///
172/// ```rust,no_run
173/// use open_agent::{query, AgentOptions};
174/// use futures::StreamExt;
175///
176/// # async fn example() {
177/// let options = AgentOptions::builder()
178///     .model("gpt-4")
179///     .api_key("invalid-key")
180///     .build()
181///     .unwrap();
182///
183/// match query("Hello", &options).await {
184///     Ok(mut stream) => {
185///         while let Some(result) = stream.next().await {
186///             match result {
187///                 Ok(block) => println!("Block: {:?}", block),
188///                 Err(e) => {
189///                     eprintln!("Stream error: {}", e);
190///                     break;
191///                 }
192///             }
193///         }
194///     }
195///     Err(e) => eprintln!("Query failed: {}", e),
196/// }
197/// # }
198/// ```
199pub async fn query(prompt: &str, options: &AgentOptions) -> Result<ContentStream> {
200    // Create HTTP client with configured timeout
201    // The timeout applies to the entire request, not individual chunks
202    let client = reqwest::Client::builder()
203        .timeout(Duration::from_secs(options.timeout()))
204        .build()
205        .map_err(Error::Http)?;
206
207    // Build messages array for the API request
208    // OpenAI format expects an array of message objects with role and content
209    let mut messages = Vec::new();
210
211    // Add system prompt if provided
212    // System prompts set the assistant's behavior and context
213    if !options.system_prompt().is_empty() {
214        messages.push(OpenAIMessage {
215            role: "system".to_string(),
216            content: Some(OpenAIContent::Text(options.system_prompt().to_string())),
217            tool_calls: None,
218            tool_call_id: None,
219        });
220    }
221
222    // Add user prompt
223    // This is the actual query from the user
224    messages.push(OpenAIMessage {
225        role: "user".to_string(),
226        content: Some(OpenAIContent::Text(prompt.to_string())),
227        tool_calls: None,
228        tool_call_id: None,
229    });
230
231    // Convert tools to OpenAI format if any are provided
232    // Tools are described using JSON Schema for parameter validation
233    let tools = if !options.tools().is_empty() {
234        Some(
235            options
236                .tools()
237                .iter()
238                .map(|t| t.to_openai_format())
239                .collect(),
240        )
241    } else {
242        None
243    };
244
245    // Build the OpenAI-compatible request payload
246    // stream=true enables Server-Sent Events for incremental responses
247    let request = OpenAIRequest {
248        model: options.model().to_string(),
249        messages,
250        stream: true, // Critical: enables SSE streaming
251        max_tokens: options.max_tokens(),
252        temperature: Some(options.temperature()),
253        tools,
254    };
255
256    stream_request(&client, options, &request).await
257}