kcode_k1_chat_core/
lib.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    sync::{
5        Arc,
6        atomic::{AtomicU8, AtomicU64, Ordering},
7    },
8    time::Duration,
9};
10
11use rust_decimal::Decimal;
12
13pub type LlmFuture<'a> = Pin<Box<dyn Future<Output = Result<Inference, LlmError>> + Send + 'a>>;
14pub type ToolFuture = Pin<Box<dyn Future<Output = ToolOutput> + Send + 'static>>;
15pub type CompactFuture = Pin<Box<dyn Future<Output = Result<String, String>> + Send + 'static>>;
16
17pub trait Llm: Send + Sync {
18    fn start(&self) -> Box<dyn LlmThread>;
19}
20
21pub trait LlmThread: Send {
22    fn infer<'a>(&'a mut self, delta: &'a str) -> LlmFuture<'a>;
23}
24
25#[derive(Clone, Debug, Eq, PartialEq)]
26pub enum LlmError {
27    Transient(String),
28    Permanent(String),
29}
30
31#[derive(Clone, Debug, Eq, PartialEq)]
32pub struct Inference {
33    pub text: String,
34    pub calls: Vec<Call>,
35    pub continue_inference: bool,
36}
37
38#[derive(Clone, Debug, Eq, PartialEq)]
39pub enum Call {
40    Tool(ToolRequest),
41    Worker(WorkerRequest),
42    Compact(CompactRequest),
43}
44
45#[derive(Clone, Debug, Eq, PartialEq)]
46pub struct ToolRequest {
47    pub name: String,
48    pub input: String,
49}
50
51#[derive(Clone, Debug, Eq, PartialEq)]
52pub struct WorkerRequest {
53    pub llm: String,
54    pub prompt: String,
55}
56
57#[derive(Clone, Debug, Eq, PartialEq)]
58pub struct CompactRequest {
59    pub instruction: String,
60}
61
62#[derive(Clone, Debug, Eq, PartialEq)]
63pub struct ToolOutput {
64    pub text: String,
65    pub cost_cents: Decimal,
66}
67
68#[derive(Clone, Copy, Debug, Eq, PartialEq)]
69pub enum ToolMode {
70    Fast,
71    Queued,
72}
73
74pub struct ToolStart {
75    pub mode: ToolMode,
76    pub queued: String,
77    pub future: ToolFuture,
78}
79
80pub struct WorkerStart {
81    pub llm: Arc<dyn Llm>,
82    pub queued: String,
83}
84
85pub trait Runtime: Send + Sync {
86    fn start_tool(&self, request: ToolRequest, updates: Updates) -> Result<ToolStart, String>;
87    fn start_worker(&self, request: &WorkerRequest) -> Result<WorkerStart, String>;
88    fn compact(&self, request: CompactRequest, primary: String) -> CompactFuture;
89}
90
91#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
92pub struct ActionId {
93    session: [u8; 12],
94    sequence: u64,
95}
96
97impl ActionId {
98    pub const fn new(session: [u8; 12], sequence: u64) -> Self {
99        Self { session, sequence }
100    }
101
102    pub const fn session(self) -> [u8; 12] {
103        self.session
104    }
105
106    pub const fn sequence(self) -> u64 {
107        self.sequence
108    }
109}
110
111#[derive(Clone, Debug, Eq, PartialEq)]
112pub enum SubmittedUpdate {
113    Activity(String),
114    Append(String),
115}
116
117pub trait UpdateSink: Send + Sync + 'static {
118    fn submit(
119        &self,
120        action: ActionId,
121        identity: u64,
122        update: SubmittedUpdate,
123    ) -> Result<(), ChatError>;
124}
125
126#[derive(Clone)]
127pub struct Updates {
128    action: ActionId,
129    next: Arc<AtomicU64>,
130    sink: Arc<dyn UpdateSink>,
131}
132
133impl Updates {
134    pub fn bind(action: ActionId, sink: Arc<dyn UpdateSink>) -> Self {
135        Self {
136            action,
137            next: Arc::new(AtomicU64::new(1)),
138            sink,
139        }
140    }
141
142    pub const fn action_id(&self) -> ActionId {
143        self.action
144    }
145
146    pub fn activity(&self, text: String) -> PreparedUpdate {
147        self.prepare(SubmittedUpdate::Activity(text))
148    }
149
150    pub fn append(&self, text: String) -> PreparedUpdate {
151        self.prepare(SubmittedUpdate::Append(text))
152    }
153
154    fn prepare(&self, update: SubmittedUpdate) -> PreparedUpdate {
155        PreparedUpdate {
156            action: self.action,
157            identity: self.next.fetch_add(1, Ordering::Relaxed),
158            update,
159            sink: self.sink.clone(),
160        }
161    }
162}
163
164#[derive(Clone)]
165pub struct PreparedUpdate {
166    action: ActionId,
167    identity: u64,
168    update: SubmittedUpdate,
169    sink: Arc<dyn UpdateSink>,
170}
171
172impl PreparedUpdate {
173    pub fn send(&self) -> Result<(), ChatError> {
174        self.sink
175            .submit(self.action, self.identity, self.update.clone())
176    }
177}
178
179#[derive(Clone, Debug, Eq, PartialEq)]
180pub struct ChatView {
181    pub primary: String,
182    pub pending: String,
183    pub history: Vec<String>,
184    pub actions: Vec<PendingAction>,
185}
186
187#[derive(Clone, Debug, Eq, PartialEq)]
188pub enum PendingAction {
189    Inference { attempt: u8 },
190    Tool { name: String },
191    Worker { llm: String },
192    Compaction,
193}
194
195#[derive(Clone, Debug, Eq, PartialEq)]
196pub enum ChatEvent {
197    Text(String),
198    Activity(String),
199    Stalled(String),
200}
201
202#[derive(Clone, Debug, Eq, PartialEq)]
203pub enum ChatError {
204    Empty,
205    NotStalled,
206    Busy,
207    Closed,
208}
209
210pub async fn infer_with_retry(
211    mut thread: Box<dyn LlmThread>,
212    delta: String,
213    attempt: Arc<AtomicU8>,
214) -> (Box<dyn LlmThread>, Result<Inference, String>) {
215    let waits = [10, 20, 40, 80];
216    for number in 1..=5 {
217        attempt.store(number, Ordering::Relaxed);
218        match thread.infer(&delta).await {
219            Ok(inference) => return (thread, Ok(inference)),
220            Err(LlmError::Permanent(error)) => return (thread, Err(error)),
221            Err(LlmError::Transient(error)) if number == 5 => return (thread, Err(error)),
222            Err(LlmError::Transient(_)) => {
223                tokio::time::sleep(Duration::from_secs(waits[number as usize - 1])).await;
224            }
225        }
226    }
227    unreachable!("five inference attempts exhaust every branch")
228}
kcode_k1_chat_core/lib.rs

kcode_k1_chat_core/
lib.rs
