diff options
Diffstat (limited to 'crates/atuin-ai/src/stream.rs')
| -rw-r--r-- | crates/atuin-ai/src/stream.rs | 10 |
1 files changed, 9 insertions, 1 deletions
diff --git a/crates/atuin-ai/src/stream.rs b/crates/atuin-ai/src/stream.rs index 9c21fc05..f4f4d704 100644 --- a/crates/atuin-ai/src/stream.rs +++ b/crates/atuin-ai/src/stream.rs @@ -12,6 +12,7 @@ use eye_declare::Handle; use eyre::{Context, Result}; use futures::StreamExt; use reqwest::Url; +use reqwest::header::USER_AGENT; use crate::{ context::{AppContext, ClientContext}, @@ -19,6 +20,8 @@ use crate::{ tui::{Session, events::AiTuiEvent}, }; +static APP_USER_AGENT: &str = concat!("atuin/", env!("CARGO_PKG_VERSION")); + /// Frames that alter the stream lifecycle — terminal or state-changing. #[derive(Debug, Clone)] pub(crate) enum StreamControl { @@ -57,6 +60,7 @@ pub(crate) struct ChatRequest { pub messages: Vec<serde_json::Value>, pub session_id: Option<String>, pub capabilities: Vec<String>, + pub invocation_id: String, } impl ChatRequest { @@ -64,8 +68,9 @@ impl ChatRequest { messages: Vec<serde_json::Value>, session_id: Option<String>, capabilities: &AiCapabilities, + invocation_id: String, ) -> Self { - let mut caps = vec![]; + let mut caps = vec!["client_invocations".to_string()]; if capabilities.enable_history_search.unwrap_or(true) { caps.push("client_v1_atuin_history".to_string()); } @@ -82,6 +87,7 @@ impl ChatRequest { messages, session_id, capabilities: caps, + invocation_id, } } } @@ -112,6 +118,7 @@ fn create_chat_stream( "messages": request.messages, "context": context, "capabilities": request.capabilities, + "invocation_id": request.invocation_id }); if let Some(ref sid) = request.session_id { @@ -123,6 +130,7 @@ fn create_chat_stream( let response = match client .post(endpoint.clone()) .header("Accept", "text/event-stream") + .header(USER_AGENT, APP_USER_AGENT) .bearer_auth(&token) .json(&request_body) .send() |