82 lines
2.5 KiB
Rust
82 lines
2.5 KiB
Rust
use goose_provider_types::thinking::{FilterOut, ThinkFilter};
|
|
|
|
pub(crate) struct ThinkingOutputFilter {
|
|
enabled: bool,
|
|
saw_structured_reasoning: bool,
|
|
think_filter: ThinkFilter,
|
|
pending_inline_thinking: String,
|
|
accumulated_thinking: String,
|
|
}
|
|
|
|
impl ThinkingOutputFilter {
|
|
pub(crate) fn new(enable_thinking: bool, generation_prompt: &str) -> Self {
|
|
let mut think_filter = ThinkFilter::new();
|
|
if enable_thinking && !generation_prompt.is_empty() {
|
|
let _ = think_filter.push(generation_prompt);
|
|
}
|
|
|
|
Self {
|
|
enabled: enable_thinking,
|
|
saw_structured_reasoning: false,
|
|
think_filter,
|
|
pending_inline_thinking: String::new(),
|
|
accumulated_thinking: String::new(),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn push_structured_reasoning(&mut self, reasoning: &str) -> Option<String> {
|
|
if reasoning.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
self.saw_structured_reasoning = true;
|
|
self.pending_inline_thinking.clear();
|
|
self.think_filter = ThinkFilter::new();
|
|
self.accumulated_thinking.push_str(reasoning);
|
|
Some(reasoning.to_string())
|
|
}
|
|
|
|
pub(crate) fn push_text(&mut self, text: &str) -> FilterOut {
|
|
if !self.enabled {
|
|
return FilterOut {
|
|
content: text.to_string(),
|
|
thinking: String::new(),
|
|
};
|
|
}
|
|
|
|
let mut filtered = self.think_filter.push(text);
|
|
if self.saw_structured_reasoning {
|
|
filtered.thinking.clear();
|
|
} else if !filtered.thinking.is_empty() {
|
|
self.pending_inline_thinking.push_str(&filtered.thinking);
|
|
filtered.thinking.clear();
|
|
}
|
|
filtered
|
|
}
|
|
|
|
pub(crate) fn finish(&mut self) -> FilterOut {
|
|
let mut filtered = if self.enabled && !self.saw_structured_reasoning {
|
|
std::mem::take(&mut self.think_filter).finish()
|
|
} else {
|
|
FilterOut::default()
|
|
};
|
|
|
|
if !self.saw_structured_reasoning {
|
|
let mut thinking = std::mem::take(&mut self.pending_inline_thinking);
|
|
thinking.push_str(&filtered.thinking);
|
|
if !thinking.is_empty() {
|
|
self.accumulated_thinking.push_str(&thinking);
|
|
}
|
|
filtered.thinking = thinking;
|
|
} else {
|
|
filtered.thinking.clear();
|
|
}
|
|
|
|
filtered
|
|
}
|
|
|
|
pub(crate) fn accumulated_thinking(&self) -> &str {
|
|
&self.accumulated_thinking
|
|
}
|
|
}
|