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
11pub type LlmFuture<'a> = Pin<Box<dyn Future<Output = Result<Inference, LlmError>> + Send + 'a>>;
12pub type ToolFuture = Pin<Box<dyn Future<Output = String> + Send + 'static>>;
13pub type CompactFuture = Pin<Box<dyn Future<Output = Result<String, String>> + Send + 'static>>;
14
15pub trait Llm: Send + Sync {
16    fn start(&self) -> Box<dyn LlmThread>;
17}
18
19pub trait LlmThread: Send {
20    fn infer<'a>(&'a mut self, delta: &'a str) -> LlmFuture<'a>;
21}
22
23#[derive(Clone, Debug, PartialEq, Eq)]
24pub enum LlmError {
25    Transient(String),
26    Permanent(String),
27}
28
29#[derive(Clone, Debug, PartialEq, Eq)]
30pub struct Inference {
31    pub text: String,
32    pub calls: Vec<Call>,
33    pub continue_inference: bool,
34}
35
36#[derive(Clone, Debug, PartialEq, Eq)]
37pub enum Call {
38    Tool(ToolRequest),
39    Worker(WorkerRequest),
40    Compact(CompactRequest),
41}
42
43#[derive(Clone, Debug, PartialEq, Eq)]
44pub struct ToolRequest {
45    pub name: String,
46    pub input: String,
47}
48
49#[derive(Clone, Debug, PartialEq, Eq)]
50pub struct WorkerRequest {
51    pub llm: String,
52    pub prompt: String,
53}
54
55#[derive(Clone, Debug, PartialEq, Eq)]
56pub struct CompactRequest {
57    pub instruction: String,
58}
59
60#[derive(Copy, Clone, Debug, PartialEq, Eq)]
61pub enum ToolMode {
62    Fast,
63    Queued,
64}
65
66pub struct ToolStart {
67    pub mode: ToolMode,
68    pub queued: String,
69    pub future: ToolFuture,
70}
71
72pub struct WorkerStart {
73    pub llm: Arc<dyn Llm>,
74    pub queued: String,
75}
76
77pub trait Runtime: Send + Sync {
78    fn start_tool(&self, request: ToolRequest, updates: Updates) -> Result<ToolStart, String>;
79    fn start_worker(&self, request: &WorkerRequest) -> Result<WorkerStart, String>;
80    fn compact(&self, request: CompactRequest, primary: String) -> CompactFuture;
81}
82
83#[derive(Clone, Debug, PartialEq, Eq)]
84pub enum SubmittedUpdate {
85    Activity(String),
86    Append(String),
87}
88
89pub trait UpdateSink: Send + Sync + 'static {
90    fn submit(&self, job: u64, identity: u64, update: SubmittedUpdate) -> Result<(), ChatError>;
91}
92
93#[derive(Clone)]
94pub struct Updates {
95    job: u64,
96    next_identity: Arc<AtomicU64>,
97    sink: Arc<dyn UpdateSink>,
98}
99
100impl Updates {
101    pub fn bind(job: u64, sink: Arc<dyn UpdateSink>) -> Self {
102        Self {
103            job,
104            next_identity: Arc::new(AtomicU64::new(1)),
105            sink,
106        }
107    }
108
109    pub fn activity(&self, activity: String) -> PreparedUpdate {
110        self.prepare(SubmittedUpdate::Activity(activity))
111    }
112
113    pub fn append(&self, text: String) -> PreparedUpdate {
114        self.prepare(SubmittedUpdate::Append(text))
115    }
116
117    fn prepare(&self, update: SubmittedUpdate) -> PreparedUpdate {
118        PreparedUpdate {
119            job: self.job,
120            identity: self.next_identity.fetch_add(1, Ordering::Relaxed),
121            update,
122            sink: self.sink.clone(),
123        }
124    }
125}
126
127#[derive(Clone)]
128pub struct PreparedUpdate {
129    job: u64,
130    identity: u64,
131    update: SubmittedUpdate,
132    sink: Arc<dyn UpdateSink>,
133}
134
135impl PreparedUpdate {
136    pub fn send(&self) -> Result<(), ChatError> {
137        self.sink
138            .submit(self.job, self.identity, self.update.clone())
139    }
140}
141
142#[derive(Clone, Debug, PartialEq, Eq)]
143pub struct ChatView {
144    pub primary: String,
145    pub pending: String,
146    pub history: Vec<String>,
147    pub actions: Vec<PendingAction>,
148}
149
150#[derive(Clone, Debug, PartialEq, Eq)]
151pub enum PendingAction {
152    Inference { attempt: u8 },
153    Tool { name: String },
154    Worker { llm: String },
155    Compaction,
156}
157
158#[derive(Clone, Debug, PartialEq, Eq)]
159pub enum ChatEvent {
160    Text(String),
161    Activity(String),
162    Stalled(String),
163}
164
165#[derive(Clone, Debug, PartialEq, Eq)]
166pub enum ChatError {
167    Empty,
168    NotStalled,
169    Busy,
170    Closed,
171}
172
173pub async fn infer_with_retry(
174    mut thread: Box<dyn LlmThread>,
175    delta: String,
176    attempt: Arc<AtomicU8>,
177) -> (Box<dyn LlmThread>, Result<Inference, String>) {
178    let mut attempt_number = 1;
179    loop {
180        attempt.store(attempt_number, Ordering::Relaxed);
181        match thread.infer(&delta).await {
182            Ok(inference) => return (thread, Ok(inference)),
183            Err(LlmError::Permanent(message)) => return (thread, Err(message)),
184            Err(LlmError::Transient(message)) => {
185                if attempt_number == 5 {
186                    return (thread, Err(message));
187                }
188                tokio::time::sleep(Duration::from_secs(10 << (attempt_number - 1))).await;
189                attempt_number += 1;
190            }
191        }
192    }
193}
194
195#[cfg(test)]
196mod tests;
kcode_k1_chat_core/lib.rs

kcode_k1_chat_core/
lib.rs
