pass globally unique conversation identifier as sessionId in databricks api call (#8576)
This commit is contained in:
@@ -0,0 +1,49 @@
|
|||||||
|
use crate::config::paths::Paths;
|
||||||
|
use once_cell::sync::Lazy;
|
||||||
|
use std::fs;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
static INSTANCE_ID: Lazy<String> = Lazy::new(load_or_create);
|
||||||
|
|
||||||
|
fn file_path() -> std::path::PathBuf {
|
||||||
|
Paths::state_dir().join("instance_id")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_or_create() -> String {
|
||||||
|
let path = file_path();
|
||||||
|
|
||||||
|
if let Ok(id) = fs::read_to_string(&path) {
|
||||||
|
let id = id.trim().to_string();
|
||||||
|
if !id.is_empty() {
|
||||||
|
return id;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let id = Uuid::new_v4().to_string();
|
||||||
|
|
||||||
|
if let Some(parent) = path.parent() {
|
||||||
|
let _ = fs::create_dir_all(parent);
|
||||||
|
}
|
||||||
|
let _ = fs::write(&path, &id);
|
||||||
|
|
||||||
|
id
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns a stable, globally unique identifier for this Goose installation.
|
||||||
|
/// The ID is generated once and persisted to disk, surviving restarts.
|
||||||
|
pub fn get_instance_id() -> &'static str {
|
||||||
|
&INSTANCE_ID
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_instance_id_is_stable() {
|
||||||
|
let id1 = get_instance_id();
|
||||||
|
let id2 = get_instance_id();
|
||||||
|
assert_eq!(id1, id2);
|
||||||
|
assert!(!id1.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -18,6 +18,7 @@ pub mod execution;
|
|||||||
pub mod gateway;
|
pub mod gateway;
|
||||||
pub mod goose_apps;
|
pub mod goose_apps;
|
||||||
pub mod hints;
|
pub mod hints;
|
||||||
|
pub mod instance_id;
|
||||||
pub mod logging;
|
pub mod logging;
|
||||||
pub mod mcp_utils;
|
pub mod mcp_utils;
|
||||||
pub mod model;
|
pub mod model;
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ use super::retry::ProviderRetry;
|
|||||||
use super::utils::{ImageFormat, RequestLog};
|
use super::utils::{ImageFormat, RequestLog};
|
||||||
use crate::config::ConfigError;
|
use crate::config::ConfigError;
|
||||||
use crate::conversation::message::Message;
|
use crate::conversation::message::Message;
|
||||||
|
use crate::instance_id::get_instance_id;
|
||||||
use crate::model::ModelConfig;
|
use crate::model::ModelConfig;
|
||||||
use crate::providers::retry::{
|
use crate::providers::retry::{
|
||||||
RetryConfig, DEFAULT_BACKOFF_MULTIPLIER, DEFAULT_INITIAL_RETRY_INTERVAL_MS,
|
RetryConfig, DEFAULT_BACKOFF_MULTIPLIER, DEFAULT_INITIAL_RETRY_INTERVAL_MS,
|
||||||
@@ -132,6 +133,8 @@ pub struct DatabricksProvider {
|
|||||||
name: String,
|
name: String,
|
||||||
#[serde(skip)]
|
#[serde(skip)]
|
||||||
token_cache: Arc<Mutex<Option<String>>>,
|
token_cache: Arc<Mutex<Option<String>>>,
|
||||||
|
#[serde(skip)]
|
||||||
|
instance_id: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl DatabricksProvider {
|
impl DatabricksProvider {
|
||||||
@@ -186,6 +189,7 @@ impl DatabricksProvider {
|
|||||||
fast_retry_config,
|
fast_retry_config,
|
||||||
name: DATABRICKS_PROVIDER_NAME.to_string(),
|
name: DATABRICKS_PROVIDER_NAME.to_string(),
|
||||||
token_cache,
|
token_cache,
|
||||||
|
instance_id: Self::resolve_instance_id(),
|
||||||
};
|
};
|
||||||
provider.model =
|
provider.model =
|
||||||
model.with_fast(DATABRICKS_DEFAULT_FAST_MODEL, DATABRICKS_PROVIDER_NAME)?;
|
model.with_fast(DATABRICKS_DEFAULT_FAST_MODEL, DATABRICKS_PROVIDER_NAME)?;
|
||||||
@@ -249,9 +253,21 @@ impl DatabricksProvider {
|
|||||||
fast_retry_config: RetryConfig::new(0, 0, 1.0, 0),
|
fast_retry_config: RetryConfig::new(0, 0, 1.0, 0),
|
||||||
name: DATABRICKS_PROVIDER_NAME.to_string(),
|
name: DATABRICKS_PROVIDER_NAME.to_string(),
|
||||||
token_cache,
|
token_cache,
|
||||||
|
instance_id: Self::resolve_instance_id(),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn resolve_instance_id() -> Option<String> {
|
||||||
|
let enabled = crate::config::Config::global()
|
||||||
|
.get_param::<bool>("GOOSE_DATABRICKS_CLIENT_REQUEST_ID")
|
||||||
|
.unwrap_or(false);
|
||||||
|
if enabled {
|
||||||
|
Some(get_instance_id().to_string())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn is_responses_model(model_name: &str) -> bool {
|
fn is_responses_model(model_name: &str) -> bool {
|
||||||
let normalized = model_name.to_ascii_lowercase();
|
let normalized = model_name.to_ascii_lowercase();
|
||||||
normalized.contains("codex")
|
normalized.contains("codex")
|
||||||
@@ -267,16 +283,31 @@ impl DatabricksProvider {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn build_client_request_id(&self, session_id: &str) -> Option<String> {
|
||||||
|
self.instance_id.as_ref().map(|instance_id| {
|
||||||
|
json!({
|
||||||
|
"sessionId": format!("{}_{}", instance_id, session_id),
|
||||||
|
})
|
||||||
|
.to_string()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
async fn post(
|
async fn post(
|
||||||
&self,
|
&self,
|
||||||
session_id: Option<&str>,
|
session_id: Option<&str>,
|
||||||
payload: Value,
|
mut payload: Value,
|
||||||
model_name: Option<&str>,
|
model_name: Option<&str>,
|
||||||
) -> Result<Value, ProviderError> {
|
) -> Result<Value, ProviderError> {
|
||||||
let is_embedding = payload.get("input").is_some() && payload.get("messages").is_none();
|
let is_embedding = payload.get("input").is_some() && payload.get("messages").is_none();
|
||||||
let model_to_use = model_name.unwrap_or(&self.model.model_name);
|
let model_to_use = model_name.unwrap_or(&self.model.model_name);
|
||||||
let path = self.get_endpoint_path(model_to_use, is_embedding);
|
let path = self.get_endpoint_path(model_to_use, is_embedding);
|
||||||
|
|
||||||
|
if let Some(session_id) = session_id {
|
||||||
|
if let Some(client_request_id) = self.build_client_request_id(session_id) {
|
||||||
|
payload["client_request_id"] = Value::String(client_request_id);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let response = self
|
let response = self
|
||||||
.api_client
|
.api_client
|
||||||
.response_post(session_id, &path, &payload)
|
.response_post(session_id, &path, &payload)
|
||||||
@@ -341,10 +372,14 @@ impl Provider for DatabricksProvider {
|
|||||||
tools: &[Tool],
|
tools: &[Tool],
|
||||||
) -> Result<MessageStream, ProviderError> {
|
) -> Result<MessageStream, ProviderError> {
|
||||||
let path = self.get_endpoint_path(&model_config.model_name, false);
|
let path = self.get_endpoint_path(&model_config.model_name, false);
|
||||||
|
let client_request_id = self.build_client_request_id(session_id);
|
||||||
|
|
||||||
if Self::is_responses_model(&model_config.model_name) {
|
if Self::is_responses_model(&model_config.model_name) {
|
||||||
let mut payload = create_responses_request(model_config, system, messages, tools)?;
|
let mut payload = create_responses_request(model_config, system, messages, tools)?;
|
||||||
payload["stream"] = Value::Bool(true);
|
payload["stream"] = Value::Bool(true);
|
||||||
|
if let Some(ref client_request_id) = client_request_id {
|
||||||
|
payload["client_request_id"] = Value::String(client_request_id.clone());
|
||||||
|
}
|
||||||
|
|
||||||
let mut log = RequestLog::start(model_config, &payload)?;
|
let mut log = RequestLog::start(model_config, &payload)?;
|
||||||
|
|
||||||
@@ -383,6 +418,9 @@ impl Provider for DatabricksProvider {
|
|||||||
.as_object_mut()
|
.as_object_mut()
|
||||||
.expect("payload should have model key")
|
.expect("payload should have model key")
|
||||||
.remove("model");
|
.remove("model");
|
||||||
|
if let Some(client_request_id) = client_request_id {
|
||||||
|
payload["client_request_id"] = Value::String(client_request_id);
|
||||||
|
}
|
||||||
|
|
||||||
payload
|
payload
|
||||||
.as_object_mut()
|
.as_object_mut()
|
||||||
|
|||||||
Reference in New Issue
Block a user