kcode-k1-chat-core 0.2.1

use std::{
    future::Future,
    pin::Pin,
    sync::{
        Arc,
        atomic::{AtomicU8, AtomicU64, Ordering},
    },
    time::Duration,
};

use rust_decimal::Decimal;

pub type LlmFuture<'a> = Pin<Box<dyn Future<Output = Result<Inference, LlmError>> + Send + 'a>>;
pub type ToolFuture = Pin<Box<dyn Future<Output = ToolOutput> + Send + 'static>>;
pub type CompactFuture = Pin<Box<dyn Future<Output = Result<String, String>> + Send + 'static>>;

pub trait Llm: Send + Sync {
    fn start(&self) -> Box<dyn LlmThread>;
}

pub trait LlmThread: Send {
    fn infer<'a>(&'a mut self, delta: &'a str) -> LlmFuture<'a>;
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LlmError {
    Transient(String),
    Permanent(String),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Inference {
    pub text: String,
    pub calls: Vec<Call>,
    pub continue_inference: bool,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Call {
    Tool(ToolRequest),
    Worker(WorkerRequest),
    Compact(CompactRequest),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolRequest {
    pub name: String,
    pub input: String,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct WorkerRequest {
    pub llm: String,
    pub prompt: String,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CompactRequest {
    pub instruction: String,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolOutput {
    pub text: String,
    pub cost_cents: Decimal,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ToolMode {
    Fast,
    Queued,
}

pub struct ToolStart {
    pub mode: ToolMode,
    pub queued: String,
    pub future: ToolFuture,
}

pub struct WorkerStart {
    pub llm: Arc<dyn Llm>,
    pub queued: String,
}

pub trait Runtime: Send + Sync {
    fn start_tool(&self, request: ToolRequest, updates: Updates) -> Result<ToolStart, String>;
    fn start_worker(&self, request: &WorkerRequest) -> Result<WorkerStart, String>;
    fn compact(&self, request: CompactRequest, primary: String) -> CompactFuture;
}

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ActionId {
    session: [u8; 12],
    sequence: u64,
}

impl ActionId {
    pub const fn new(session: [u8; 12], sequence: u64) -> Self {
        Self { session, sequence }
    }

    pub const fn session(self) -> [u8; 12] {
        self.session
    }

    pub const fn sequence(self) -> u64 {
        self.sequence
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SubmittedUpdate {
    Activity(String),
    Append(String),
}

pub trait UpdateSink: Send + Sync + 'static {
    fn submit(
        &self,
        action: ActionId,
        identity: u64,
        update: SubmittedUpdate,
    ) -> Result<(), ChatError>;
}

#[derive(Clone)]
pub struct Updates {
    action: ActionId,
    next: Arc<AtomicU64>,
    sink: Arc<dyn UpdateSink>,
}

impl Updates {
    pub fn bind(action: ActionId, sink: Arc<dyn UpdateSink>) -> Self {
        Self {
            action,
            next: Arc::new(AtomicU64::new(1)),
            sink,
        }
    }

    pub const fn action_id(&self) -> ActionId {
        self.action
    }

    pub fn activity(&self, text: String) -> PreparedUpdate {
        self.prepare(SubmittedUpdate::Activity(text))
    }

    pub fn append(&self, text: String) -> PreparedUpdate {
        self.prepare(SubmittedUpdate::Append(text))
    }

    fn prepare(&self, update: SubmittedUpdate) -> PreparedUpdate {
        PreparedUpdate {
            action: self.action,
            identity: self.next.fetch_add(1, Ordering::Relaxed),
            update,
            sink: self.sink.clone(),
        }
    }
}

#[derive(Clone)]
pub struct PreparedUpdate {
    action: ActionId,
    identity: u64,
    update: SubmittedUpdate,
    sink: Arc<dyn UpdateSink>,
}

impl PreparedUpdate {
    pub fn send(&self) -> Result<(), ChatError> {
        self.sink
            .submit(self.action, self.identity, self.update.clone())
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ChatView {
    pub primary: String,
    pub pending: String,
    pub history: Vec<String>,
    pub actions: Vec<PendingAction>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PendingAction {
    Inference { attempt: u8 },
    Tool { name: String },
    Worker { llm: String },
    Compaction,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ChatEvent {
    Text(String),
    Activity(String),
    Stalled(String),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ChatError {
    Empty,
    NotStalled,
    Busy,
    Closed,
}

pub async fn infer_with_retry(
    mut thread: Box<dyn LlmThread>,
    delta: String,
    attempt: Arc<AtomicU8>,
) -> (Box<dyn LlmThread>, Result<Inference, String>) {
    let waits = [10, 20, 40, 80];
    for number in 1..=5 {
        attempt.store(number, Ordering::Relaxed);
        match thread.infer(&delta).await {
            Ok(inference) => return (thread, Ok(inference)),
            Err(LlmError::Permanent(error)) => return (thread, Err(error)),
            Err(LlmError::Transient(error)) if number == 5 => return (thread, Err(error)),
            Err(LlmError::Transient(_)) => {
                tokio::time::sleep(Duration::from_secs(waits[number as usize - 1])).await;
            }
        }
    }
    unreachable!("five inference attempts exhaust every branch")
}