libtenx 0.0.4

A library for building AI-assisted coding tools, with session management, patch validation, and multiple model support.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! Model abstractions and implementations for interacting with AI language models.
//!
//! This module provides traits and implementations for different AI model providers,
//! along with usage tracking and response handling.

mod claude;
mod conversation;
mod dummy_model;
mod google;
mod openai;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;

pub use claude::{Claude, ClaudeUsage};
pub use dummy_model::{DummyModel, DummyUsage};
pub use google::{Google, GoogleUsage};
pub use openai::{OpenAi, OpenAiUsage};

use crate::{
    config::Config,
    events::Event,
    session::{ModelResponse, Session},
    Result,
};

use std::collections::HashMap;

/// Represents usage statistics for different model types.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
pub enum Usage {
    Claude(ClaudeUsage),
    OpenAi(OpenAiUsage),
    Dummy(DummyUsage),
    Google(google::GoogleUsage),
}

impl Usage {
    /// Returns a map of usage statistics.
    pub fn values(&self) -> HashMap<String, u64> {
        match self {
            Usage::Claude(usage) => usage.values(),
            Usage::OpenAi(usage) => usage.values(),
            Usage::Dummy(usage) => usage.values(),
            Usage::Google(usage) => usage.values(),
        }
    }

    /// Returns a tuple of (tokens in, tokens out).
    pub fn totals(&self) -> (u64, u64) {
        match self {
            Usage::Claude(usage) => usage.totals(),
            Usage::OpenAi(usage) => usage.totals(),
            Usage::Dummy(usage) => usage.totals(),
            Usage::Google(usage) => usage.totals(),
        }
    }
}

/// Implemented by types that expose a prompt operation.
#[async_trait]
pub trait ModelProvider {
    /// Returns user-facing name of the model.
    fn name(&self) -> String;

    /// Returns underlying name of the model.
    fn api_model(&self) -> String;

    /// Render and send a session to the model.
    async fn send(
        &mut self,
        config: &Config,
        session: &Session,
        sender: Option<mpsc::Sender<Event>>,
    ) -> Result<ModelResponse>;

    /// Render a session as it would be sent to the model. It's a requirement that this step be
    /// able to render a sessio with no steps, that is, with the system prompt only.
    fn render(&self, config: &Config, session: &Session) -> Result<String>;
}

/// Available model implementations that can be used for AI interactions.
#[derive(Debug, Clone)]
pub enum Model {
    Claude(Claude),
    OpenAi(OpenAi),
    Dummy(DummyModel),
    Google(google::Google),
}

#[async_trait]
impl ModelProvider for Model {
    fn name(&self) -> String {
        match self {
            Model::Claude(c) => c.name(),
            Model::OpenAi(o) => o.name(),
            Model::Dummy(d) => d.name(),
            Model::Google(g) => g.name(),
        }
    }

    fn api_model(&self) -> String {
        match self {
            Model::Claude(c) => c.api_model(),
            Model::OpenAi(o) => o.api_model(),
            Model::Dummy(d) => d.api_model(),
            Model::Google(g) => g.api_model(),
        }
    }

    async fn send(
        &mut self,
        config: &Config,
        session: &Session,
        sender: Option<mpsc::Sender<Event>>,
    ) -> Result<ModelResponse> {
        match self {
            Model::Claude(c) => c.send(config, session, sender).await,
            Model::OpenAi(o) => o.send(config, session, sender).await,
            Model::Dummy(d) => d.send(config, session, sender).await,
            Model::Google(g) => g.send(config, session, sender).await,
        }
    }

    fn render(&self, config: &Config, session: &Session) -> Result<String> {
        match self {
            Model::Claude(c) => c.render(config, session),
            Model::OpenAi(o) => o.render(config, session),
            Model::Dummy(d) => d.render(config, session),
            Model::Google(g) => g.render(config, session),
        }
    }
}