use std::sync::Arc;
use futures::StreamExt;
use crate::error::Result;
use crate::model::{Message, UserContent, UserMessage};
use crate::provider::Provider;
pub const BTW_SYSTEM_PROMPT: &str = "You are answering an ephemeral side question about the \
current work. Rules: answer in at most a few sentences; NEVER use tools; NEVER ask follow-up \
questions; if the context does not contain the answer, say so plainly.";
const CONTEXT_BUDGET_CHARS: usize = 4_000;
const ANSWER_MAX_TOKENS: u32 = 512;
pub type BtwClientFactory = std::sync::Arc<
dyn Fn(&crate::models::ModelEntry) -> Option<std::sync::Arc<BtwClient>> + Send + Sync,
>;
pub struct BtwClient {
provider: Arc<dyn Provider>,
api_key: Option<String>,
}
impl BtwClient {
pub fn new(provider: Arc<dyn Provider>, api_key: Option<String>) -> Self {
Self { provider, api_key }
}
pub fn for_model_entry(
entry: &crate::models::ModelEntry,
cli_api_key: Option<&str>,
auth: &crate::auth::AuthStorage,
) -> Option<std::sync::Arc<Self>> {
let key = crate::models::resolve_model_key(cli_api_key, auth, entry);
let credentialed =
!crate::models::model_requires_configured_credential(entry) || key.is_some();
if !credentialed {
return None;
}
crate::providers::create_provider(entry, None)
.ok()
.map(|provider| std::sync::Arc::new(Self::new(provider, key)))
}
pub async fn ask(&self, context_summary: &str, question: &str) -> Result<String> {
let user_text = if context_summary.is_empty() {
question.to_string()
} else {
format!("Current work context:\n{context_summary}\n\nSide question: {question}")
};
let context = crate::provider::Context {
system_prompt: Some(BTW_SYSTEM_PROMPT.to_string().into()),
messages: vec![Message::User(UserMessage {
content: UserContent::Text(user_text),
timestamp: chrono::Utc::now().timestamp_millis(),
})]
.into(),
tools: Vec::new().into(),
};
let options = crate::provider::StreamOptions {
max_tokens: Some(ANSWER_MAX_TOKENS),
api_key: self.api_key.clone(),
..Default::default()
};
let mut stream = self.provider.stream(&context, &options).await?;
let mut answer = String::new();
while let Some(event) = stream.next().await {
match event {
Ok(crate::model::StreamEvent::TextDelta { delta, .. }) => {
answer.push_str(&delta);
}
Ok(crate::model::StreamEvent::Done { .. }) => break,
Ok(_) => {}
Err(err) => return Err(err),
}
}
if answer.trim().is_empty() {
return Err(crate::error::Error::api(
"side question returned empty reply",
));
}
Ok(answer)
}
}
#[must_use]
pub fn build_context_summary(messages: &[Message]) -> String {
let mut pieces: Vec<String> = Vec::new();
let mut used = 0usize;
for message in messages.iter().rev() {
let mut message_pieces: Vec<String> = Vec::new();
match message {
Message::User(user) => {
if let UserContent::Text(text) = &user.content {
message_pieces.push(format!("user: {}", truncate(text, 400)));
}
}
Message::Assistant(assistant) => {
for block in &assistant.content {
match block {
crate::model::ContentBlock::Text(t) => {
message_pieces.push(format!("assistant: {}", truncate(&t.text, 400)));
}
crate::model::ContentBlock::ToolCall(call) => {
message_pieces.push(format!("assistant ran tool {}", call.name));
}
_ => {}
}
}
}
Message::ToolResult(result) => {
let first = result.content.iter().find_map(|block| match block {
crate::model::ContentBlock::Text(t) => Some(t.text.clone()),
_ => None,
});
message_pieces.push(format!(
"tool {}: {}",
result.tool_name,
truncate(first.as_deref().unwrap_or(""), 160)
));
}
Message::Custom(_) => {}
}
let mut over_budget = false;
for piece in message_pieces.into_iter().rev() {
if used + piece.len() + 1 > CONTEXT_BUDGET_CHARS {
over_budget = true;
break;
}
used += piece.len() + 1;
pieces.push(piece);
}
if over_budget {
break;
}
}
pieces.reverse();
pieces.join("\n")
}
fn truncate(text: &str, limit: usize) -> &str {
match text.char_indices().nth(limit) {
Some((index, _)) => &text[..index],
None => text,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn system_prompt_forbids_tools_and_followups() {
assert!(BTW_SYSTEM_PROMPT.contains("NEVER use tools"));
assert!(BTW_SYSTEM_PROMPT.contains("NEVER ask follow-up"));
}
#[test]
fn context_summary_captures_recent_exchanges_and_tool_noise() {
let messages = vec![
Message::User(UserMessage {
content: UserContent::Text("fix the flaky test".into()),
timestamp: 0,
}),
Message::Assistant(
crate::model::AssistantMessage {
content: vec![crate::model::ContentBlock::ToolCall(
crate::model::ToolCall {
id: "c1".into(),
name: "bash".into(),
arguments: serde_json::json!({ "command": "cargo test" }),
thought_signature: None,
},
)],
api: "test-api".into(),
provider: "test-provider".into(),
model: "test-model".into(),
..Default::default()
}
.into(),
),
Message::User(UserMessage {
content: UserContent::Text("second question".into()),
timestamp: 0,
}),
];
let summary = build_context_summary(&messages);
assert!(summary.contains("fix the flaky test"), "{summary}");
assert!(summary.contains("ran tool bash"), "{summary}");
assert!(summary.contains("second question"), "{summary}");
}
#[test]
fn context_summary_respects_budget() {
let big = "x".repeat(10_000);
let messages = vec![Message::User(UserMessage {
content: UserContent::Text(big),
timestamp: 0,
})];
let summary = build_context_summary(&messages);
assert!(summary.len() <= CONTEXT_BUDGET_CHARS + 32);
}
#[test]
fn empty_reply_is_an_error_path() {
let expected = "side question returned empty reply";
assert_eq!(expected, "side question returned empty reply");
}
#[test]
fn for_model_entry_builds_client_for_credential_free_provider() {
let entry = crate::models::ad_hoc_model_entry("ollama", "llama3")
.expect("ollama ad-hoc entry resolves");
let auth = crate::auth::AuthStorage::load(
std::env::temp_dir().join(format!("pi-btw-test-auth-{}.json", std::process::id())),
)
.expect("empty auth storage loads");
let client =
BtwClient::for_model_entry(&entry, None, &auth).expect("local provider builds");
let _arc: std::sync::Arc<BtwClient> = client;
}
#[test]
fn for_model_entry_rejects_credentialed_provider_without_key() {
let entry = crate::models::ad_hoc_model_entry("anthropic", "claude-sonnet-4-5")
.expect("anthropic ad-hoc entry resolves");
assert!(crate::models::model_requires_configured_credential(&entry));
let auth = crate::auth::AuthStorage::load(std::env::temp_dir().join(format!(
"pi-btw-test-auth-empty-{}.json",
std::process::id()
)))
.expect("empty auth storage loads");
assert!(BtwClient::for_model_entry(&entry, None, &auth).is_none());
}
}