sqz_engine/types.rs
1use std::path::PathBuf;
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4
5// Type aliases
6
7/// Unique identifier for a compression session.
8pub type SessionId = String;
9
10/// Unique identifier for an agent (budget partition).
11pub type AgentId = String;
12
13/// Unique identifier for a correction entry.
14pub type CorrectionId = String;
15
16/// Unique identifier for a tool definition.
17pub type ToolId = String;
18
19// --- Enums ---
20
21/// Supported image formats for the image compressor.
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23pub enum ImageFormat {
24 Png,
25 Jpeg,
26 Gif,
27 WebP,
28}
29
30/// Conversation turn role.
31#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
32#[serde(rename_all = "lowercase")]
33pub enum Role {
34 User,
35 Assistant,
36 System,
37}
38
39/// What kind of content we're compressing. The pipeline uses this to pick
40/// the right compression strategy.
41#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
42pub enum ContentType {
43 Json,
44 CliOutput { command: String },
45 Code { language: String, path: PathBuf },
46 PlainText,
47 Image { format: ImageFormat },
48}
49
50/// Which LLM family we're targeting. Affects token counting and cache
51/// boundary placement.
52#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
53pub enum ModelFamily {
54 AnthropicClaude,
55 OpenAiGpt,
56 GoogleGemini,
57 /// A local model identified by name (e.g. "llama-3.1-8b").
58 Local(String),
59}
60
61/// API provider — determines cache discount rates.
62#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
63pub enum Provider {
64 /// 90% cache discount
65 Anthropic,
66 /// 50% cache discount
67 OpenAI,
68 /// No cache boundary
69 Google,
70}
71
72/// Whether a task is simple (can use a smaller/cheaper model) or complex.
73#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
74pub enum TaskClassification {
75 Simple,
76 Complex,
77}
78
79// --- Content types ---
80
81/// Metadata about where content came from.
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct ContentMetadata {
84 /// Source identifier (e.g. "stdin", "tool:read_file").
85 pub source: Option<String>,
86 /// File path, if the content came from a file.
87 pub path: Option<PathBuf>,
88 /// Programming language, if detected.
89 pub language: Option<String>,
90}
91
92/// Raw content flowing through the compression pipeline. Stages mutate
93/// the `raw` field in place.
94#[derive(Debug, Clone, Serialize, Deserialize)]
95pub struct Content {
96 /// The text content being compressed.
97 pub raw: String,
98 /// Detected content type (JSON, code, CLI output, etc.).
99 pub content_type: ContentType,
100 /// Where this content came from.
101 pub metadata: ContentMetadata,
102 /// Token count before any compression.
103 pub tokens_original: u32,
104}
105
106/// Source provenance for a compressed segment — tracks where content came
107/// from so you can trace back to the original.
108#[derive(Debug, Clone, Serialize, Deserialize, Default)]
109pub struct Provenance {
110 /// File path the content originated from, if known.
111 pub file: Option<PathBuf>,
112 /// Line range within the file (start inclusive, end exclusive).
113 pub line_range: Option<std::ops::Range<usize>>,
114 /// SHA-256 hex digest of the original content.
115 pub content_hash: Option<String>,
116 /// Tool call ID that produced this content, if applicable.
117 pub tool_call_id: Option<String>,
118 /// Human-readable source label (e.g. "git diff", "cargo test output").
119 pub label: Option<String>,
120}
121
122/// The result of compressing content through the pipeline. This is what
123/// you get back from [`crate::engine::SqzEngine::compress`].
124///
125/// ```rust
126/// # use sqz_engine::SqzEngine;
127/// # let dir = tempfile::tempdir().unwrap();
128/// # std::env::set_var("SQZ_DB_PATH", dir.path().join("sessions.db"));
129/// let engine = SqzEngine::new().unwrap();
130/// let result = engine.compress("hello world").unwrap();
131///
132/// println!("output: {}", result.data);
133/// println!("ratio: {:.0}%", (1.0 - result.compression_ratio) * 100.0);
134/// println!("stages: {:?}", result.stages_applied);
135/// ```
136#[derive(Debug, Clone, Serialize, Deserialize)]
137pub struct CompressedContent {
138 /// The compressed output text.
139 pub data: String,
140 /// Token count after compression.
141 pub tokens_compressed: u32,
142 /// Token count before compression.
143 pub tokens_original: u32,
144 /// Names of pipeline stages that ran (e.g. "condense", "toon_encode").
145 pub stages_applied: Vec<String>,
146 /// Ratio of compressed to original tokens (0.0 = perfect, 1.0 = no change).
147 pub compression_ratio: f64,
148 /// Where this content came from.
149 #[serde(default)]
150 pub provenance: Provenance,
151 /// Verifier result — `None` if the verify pass was skipped.
152 #[serde(default)]
153 pub verify: Option<VerifyResult>,
154}
155
156// --- Session types ---
157
158/// A single conversation turn (user message, assistant response, etc.).
159#[derive(Debug, Clone, Serialize, Deserialize)]
160pub struct Turn {
161 pub role: Role,
162 pub content: String,
163 pub tokens: u32,
164 pub pinned: bool,
165 pub timestamp: DateTime<Utc>,
166}
167
168/// Alias for backward compatibility.
169pub type ConversationTurn = Turn;
170
171/// A pinned segment — a conversation turn that should never be evicted
172/// from the context window.
173#[derive(Debug, Clone, Serialize, Deserialize)]
174pub struct PinnedSegment {
175 pub turn_index: usize,
176 pub reason: String,
177 pub tokens: u32,
178}
179
180/// Alias for backward compatibility.
181pub type PinEntry = PinnedSegment;
182
183/// A key-value fact extracted from conversation (e.g. "preferred language" → "Rust").
184#[derive(Debug, Clone, Serialize, Deserialize)]
185pub struct KvFact {
186 pub key: String,
187 pub value: String,
188 pub source_turn: usize,
189}
190
191/// Alias for backward compatibility.
192pub type Learning = KvFact;
193
194/// Current token budget state for a session.
195#[derive(Debug, Clone, Serialize, Deserialize)]
196pub struct WindowUsage {
197 /// Total context window size in tokens.
198 pub window_size: u32,
199 /// Tokens consumed so far.
200 pub consumed: u32,
201 /// Tokens reserved by pinned segments.
202 pub pinned: u32,
203 /// Which model family determines token counting.
204 pub model_family: ModelFamily,
205}
206
207/// Alias for backward compatibility.
208pub type BudgetState = WindowUsage;
209
210/// Record of a single tool call with cost tracking.
211#[derive(Debug, Clone, Serialize, Deserialize)]
212pub struct ToolCall {
213 pub tool_name: String,
214 pub tokens_input: u32,
215 pub tokens_output: u32,
216 pub cost_usd: f64,
217 pub timestamp: DateTime<Utc>,
218}
219
220/// Alias for backward compatibility.
221pub type ToolUsageRecord = ToolCall;
222
223/// A single correction (edit) made during a session.
224#[derive(Debug, Clone, Serialize, Deserialize)]
225pub struct EditRecord {
226 pub id: CorrectionId,
227 pub timestamp: DateTime<Utc>,
228 pub original: String,
229 pub correction: String,
230 pub context: String,
231}
232
233/// Alias for backward compatibility.
234pub type CorrectionEntry = EditRecord;
235
236/// History of corrections made during a session.
237#[derive(Debug, Clone, Serialize, Deserialize, Default)]
238pub struct EditHistory {
239 pub entries: Vec<EditRecord>,
240}
241
242/// Alias for backward compatibility.
243pub type CorrectionLog = EditHistory;
244
245/// Active compression session — the main state object that tracks an entire
246/// conversation, its budget, tool usage, pins, and learnings.
247#[derive(Debug, Clone, Serialize, Deserialize)]
248pub struct Session {
249 pub id: SessionId,
250 pub project_dir: PathBuf,
251 pub conversation: Vec<Turn>,
252 pub corrections: EditHistory,
253 pub pins: Vec<PinnedSegment>,
254 pub learnings: Vec<KvFact>,
255 pub compressed_summary: String,
256 pub budget: WindowUsage,
257 pub tool_usage: Vec<ToolCall>,
258 pub created_at: DateTime<Utc>,
259 pub updated_at: DateTime<Utc>,
260}
261
262/// Alias for backward compatibility.
263pub type SessionState = Session;
264
265/// Configuration for a single compression stage.
266#[derive(Debug, Clone, Serialize, Deserialize, Default)]
267pub struct StageConfig {
268 /// Whether this stage is enabled.
269 pub enabled: bool,
270 /// Stage-specific options as a JSON value.
271 #[serde(default)]
272 pub options: serde_json::Value,
273}
274
275/// Result of the two-pass compression verifier. Checks that important
276/// content (error lines, JSON keys, diff hunks) survived compression.
277#[derive(Debug, Clone, Serialize, Deserialize)]
278pub struct VerifyResult {
279 /// Overall pass/fail.
280 pub passed: bool,
281 /// Confidence score 0.0–1.0 (1.0 = fully verified).
282 pub confidence: f64,
283 /// Which checks passed.
284 pub checks_passed: Vec<String>,
285 /// Which checks failed, with reason.
286 pub checks_failed: Vec<(String, String)>,
287 /// Whether the pipeline fell back to a safer preset.
288 pub fallback_triggered: bool,
289}