use super::CtxBuilder;
use crate::config::provider::ProviderConfig;
use crate::conversation::message::Message;
use crate::conversation::{ContextStats, Conversation};
use crate::tool::ToolResult;
#[derive(Debug, Clone)]
pub struct OllamaCtx {
ctx_window: usize,
model_id: String,
}
impl OllamaCtx {
pub fn new(provider: &ProviderConfig) -> Self {
Self {
ctx_window: provider.context_window.max(4000),
model_id: provider.model.to_lowercase(),
}
}
fn tool_output_cap(&self) -> usize {
(self.ctx_window / 8).min(6_000).max(2_000)
}
}
impl CtxBuilder for OllamaCtx {
fn build_messages(
&self,
conv: &Conversation,
system_prompt: &str,
turn_reminder: &str,
) -> (Vec<Message>, ContextStats) {
let sys = crate::ctx::render::apply_model_directives(system_prompt, &self.model_id);
crate::ctx::render::build_messages(conv, &sys, self.ctx_window, turn_reminder)
}
fn needs_compression(&self, conv: &Conversation, system_tokens: usize) -> bool {
crate::ctx::render::needs_compression(conv, system_tokens, self.ctx_window)
}
fn compression_plan(
&self,
conv: &Conversation,
keep_ceiling: usize,
) -> Option<(String, usize)> {
let (content, n) = crate::ctx::render::build_compression_content(conv, keep_ceiling);
if content.is_empty() || n == 0 {
None
} else {
Some((content, n))
}
}
fn truncate_tool_output(&self, result: &mut ToolResult, tool_name: &str) {
crate::ctx::truncate::truncate_output(result, tool_name, self.ctx_window);
let cap = self.tool_output_cap();
if result.output.len() > cap {
let mut boundary = cap;
while boundary > 0 && !result.output.is_char_boundary(boundary) {
boundary -= 1;
}
result.output.truncate(boundary);
result
.output
.push_str("\n[... truncated by OllamaCtx (small window) ...]");
}
}
fn ctx_window(&self) -> usize {
self.ctx_window
}
fn name(&self) -> &'static str {
"ollama"
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::Conversation;
use crate::tool::ToolResult;
fn ollama_provider(ctx: usize) -> ProviderConfig {
ProviderConfig {
provider_type: "ollama".into(),
api_key: None,
model: "llama3-8b".into(),
base_url: Some("http://localhost:11434".into()),
system_prompt: None,
user_agent: None,
context_window: ctx,
max_tokens: None,
thinking_type: None,
thinking_keep: None,
reasoning_history: None,
thinking_enabled: None,
thinking_budget: None,
skip_tls_verify: false,
ephemeral: false,
}
}
#[test]
fn name_is_ollama() {
let o = OllamaCtx::new(&ollama_provider(8_000));
assert_eq!(o.name(), "ollama");
}
#[test]
fn ctx_window_clamped_to_4k_minimum() {
let o = OllamaCtx::new(&ollama_provider(0));
assert_eq!(o.ctx_window, 4_000);
let o = OllamaCtx::new(&ollama_provider(2_000));
assert_eq!(o.ctx_window, 4_000);
let o = OllamaCtx::new(&ollama_provider(8_000));
assert_eq!(o.ctx_window, 8_000);
let o = OllamaCtx::new(&ollama_provider(32_000));
assert_eq!(o.ctx_window, 32_000);
}
#[test]
fn tool_output_cap_follows_spec() {
assert_eq!(
OllamaCtx::new(&ollama_provider(8_000)).tool_output_cap(),
2_000
);
assert_eq!(
OllamaCtx::new(&ollama_provider(16_000)).tool_output_cap(),
2_000
);
assert_eq!(
OllamaCtx::new(&ollama_provider(32_000)).tool_output_cap(),
4_000
);
assert_eq!(
OllamaCtx::new(&ollama_provider(64_000)).tool_output_cap(),
6_000
);
}
#[test]
fn truncate_result_enforces_small_cap() {
let o = OllamaCtx::new(&ollama_provider(8_000));
let mut r = ToolResult {
call_id: "t1".into(),
output: "x".repeat(50_000),
success: true,
};
o.truncate_tool_output(&mut r, "bash");
assert!(
r.output.len() <= 2_200,
"OllamaCtx truncate 后输出 {} 字节超过 cap 2200",
r.output.len()
);
}
#[test]
fn truncate_result_utf8_safe_on_cjk_boundary() {
let o = OllamaCtx::new(&ollama_provider(8_000));
let mut r = ToolResult {
call_id: "t1".into(),
output: "中".repeat(5_000),
success: true,
};
o.truncate_tool_output(&mut r, "bash");
assert!(std::str::from_utf8(r.output.as_bytes()).is_ok());
assert!(r.output.len() <= 2_200);
}
#[test]
fn needs_compression_triggers_earlier_than_default() {
let o = OllamaCtx::new(&ollama_provider(8_000));
let empty = Conversation::new();
assert!(!o.needs_compression(&empty, 100));
let mut conv = Conversation::new();
for i in 0..8 {
conv.add_user_message(&format!("user turn {} with moderate content", i));
conv.add_assistant_tool_calls(
Some(&format!("some assistant reasoning for turn {}", i)),
vec![],
None,
);
}
assert!(!o.needs_compression(&conv, 50));
for _ in 0..20 {
conv.add_user_message(&"lorem ipsum ".repeat(50).repeat(2));
conv.add_assistant_tool_calls(Some(&"dolor sit amet ".repeat(50)), vec![], None);
}
assert!(
o.needs_compression(&conv, 50),
"大对话下 OllamaCtx 应触发压缩(35% threshold)"
);
}
#[test]
fn compression_plan_none_below_threshold() {
let o = OllamaCtx::new(&ollama_provider(8_000));
let conv = Conversation::new();
assert!(o.compression_plan(&conv, usize::MAX).is_none());
}
#[test]
fn build_messages_returns_nonempty_for_simple_conv() {
let o = OllamaCtx::new(&ollama_provider(8_000));
let mut conv = Conversation::new();
conv.add_user_message("hello");
let (msgs, stats) = o.build_messages(&conv, "SYS", "");
assert!(!msgs.is_empty());
assert!(stats.sent_tokens <= 8_000);
}
}