Skip to main content

open_agent/client/
query.rs

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