Skip to content

Commit 1cc1df8

Browse files
kosz78kosz78
committed
feat: add model info with context len and usage cost
fix: OpenRouter messages handling, add App title
1 parent 1d2fbee commit 1cc1df8

File tree

15 files changed

+521
-93
lines changed

15 files changed

+521
-93
lines changed

‎.gitignore‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,5 @@ logs/
77
history.json
88
huly-coder-local.yaml
99
/memory.yaml
10-
/.fastembed_cache/
10+
/.fastembed_cache/
11+
openrouter_models.json

‎Cargo.lock‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/agent/event.rs‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -43,13 +43,6 @@ impl Display for AgentState {
4343
}
4444
}
4545

46-
#[derive(Clone, Debug, Default)]
47-
pub struct AgentStatus {
48-
pub current_tokens: u32,
49-
pub max_tokens: u32,
50-
pub state: AgentState,
51-
}
52-
5346
/// Status of a command tool call
5447
#[derive(Clone, Debug, Default)]
5548
pub struct AgentCommandStatus {
@@ -66,7 +59,7 @@ pub enum AgentOutputEvent {
6659
UpdateMessage(Message),
6760
NewTask,
6861
CommandStatus(Vec<AgentCommandStatus>),
69-
AgentStatus(AgentStatus),
62+
AgentStatus(u32, u32, AgentState),
7063
HighlightFile(String, bool),
7164
}
7265

‎src/agent/mod.rs‎

Lines changed: 59 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,6 @@ pub use event::AgentOutputEvent;
5151
use tokio::sync::RwLock;
5252

5353
use self::event::AgentState;
54-
use self::event::AgentStatus;
5554
use self::utils::*;
5655

5756
pub struct Agent {
@@ -66,7 +65,8 @@ pub struct Agent {
6665
memory: Arc<RwLock<MemoryManager>>,
6766
memory_index: Option<InMemoryVectorIndex<rig_fastembed::EmbeddingModel, Entity>>,
6867
process_registry: Arc<RwLock<ProcessRegistry>>,
69-
current_tokens: u32,
68+
current_input_tokens: u32,
69+
current_completion_tokens: u32,
7070
state: AgentState,
7171
}
7272

@@ -93,6 +93,10 @@ impl Display for AgentError {
9393
}
9494
}
9595

96+
fn count_tokens(system_prompt: &str) -> u32 {
97+
system_prompt.len() as u32 / 4
98+
}
99+
96100
impl Agent {
97101
pub fn new(
98102
config: Config,
@@ -108,7 +112,8 @@ impl Agent {
108112
messages,
109113
stream: None,
110114
assistant_content: None,
111-
current_tokens: 0,
115+
current_input_tokens: 0,
116+
current_completion_tokens: 0,
112117
memory: Arc::new(RwLock::new(MemoryManager::new(false))),
113118
process_registry: Arc::new(RwLock::new(ProcessRegistry::default())),
114119
memory_index: None,
@@ -338,7 +343,7 @@ impl Agent {
338343
self.messages[last_idx] = message;
339344
}
340345

341-
async fn process_messages(&mut self) -> Result<(), AgentError> {
346+
async fn process_messages(&mut self, system_prompt_token_count: u32) -> Result<(), AgentError> {
342347
if self.state.is_paused() {
343348
return Ok(());
344349
}
@@ -494,7 +499,49 @@ impl Agent {
494499
if let Some(raw_response) = response.raw_response {
495500
let usage = raw_response.usage;
496501
tracing::info!("Usage: {:?}", usage);
497-
self.current_tokens = usage.total_tokens as u32;
502+
if usage.total_tokens > 0 {
503+
self.current_input_tokens = usage.prompt_tokens as u32;
504+
self.current_completion_tokens =
505+
(usage.total_tokens - usage.prompt_tokens) as u32;
506+
} else {
507+
// try to calculate aproximate tokens
508+
self.current_input_tokens = system_prompt_token_count
509+
+ self
510+
.messages
511+
.iter()
512+
.map(|m| match m {
513+
Message::User { content } => content
514+
.iter()
515+
.map(|c| match c {
516+
UserContent::Text(text) => count_tokens(&text.text),
517+
UserContent::ToolResult(tool_result) => tool_result
518+
.content
519+
.iter()
520+
.map(|t| match t {
521+
ToolResultContent::Text(text) => {
522+
count_tokens(&text.text)
523+
}
524+
_ => 0,
525+
})
526+
.sum::<u32>(),
527+
_ => 0,
528+
})
529+
.sum::<u32>(),
530+
Message::Assistant { content } => content
531+
.iter()
532+
.map(|c| match c {
533+
AssistantContent::Text(text) => {
534+
count_tokens(&text.text)
535+
}
536+
AssistantContent::ToolCall(tool_call) => count_tokens(
537+
&serde_json::to_string(tool_call).unwrap(),
538+
),
539+
})
540+
.sum::<u32>(),
541+
})
542+
.sum::<u32>();
543+
self.current_completion_tokens = 0;
544+
}
498545
}
499546
self.assistant_content = None;
500547
if matches!(self.state, AgentState::Completed(false)) {
@@ -517,6 +564,7 @@ impl Agent {
517564
);
518565
let system_prompt =
519566
prepare_system_prompt(&self.config.workspace, &self.config.user_instructions).await;
567+
let system_prompt_token_count = count_tokens(&system_prompt);
520568
self.agent = Some(
521569
Self::build_agent(BuildAgentContext {
522570
config: &self.config,
@@ -577,7 +625,7 @@ impl Agent {
577625
}
578626
}
579627
}
580-
if let Err(e) = self.process_messages().await {
628+
if let Err(e) = self.process_messages(system_prompt_token_count).await {
581629
tracing::debug!("persist_history");
582630
persist_history(&self.messages);
583631
tracing::error!("Error processing messages: {}", e);
@@ -615,11 +663,11 @@ impl Agent {
615663
self.state = state;
616664
if !self.sender.is_closed() {
617665
self.sender
618-
.send(AgentOutputEvent::AgentStatus(AgentStatus {
619-
current_tokens: self.current_tokens,
620-
max_tokens: 1,
621-
state: self.state.clone(),
622-
}))
666+
.send(AgentOutputEvent::AgentStatus(
667+
self.current_input_tokens,
668+
self.current_completion_tokens,
669+
self.state.clone(),
670+
))
623671
.unwrap();
624672
}
625673
}

‎src/main.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use crossterm::terminal::disable_raw_mode;
99
use crossterm::terminal::enable_raw_mode;
1010
use crossterm::terminal::EnterAlternateScreen;
1111
use crossterm::terminal::LeaveAlternateScreen;
12+
use providers::model_info::model_info;
1213
use ratatui::prelude::CrosstermBackend;
1314
use ratatui::DefaultTerminal;
1415
use ratatui::Terminal;
@@ -107,6 +108,9 @@ async fn main() -> color_eyre::Result<()> {
107108
Vec::new()
108109
};
109110

111+
let model_info = model_info(&config).await?;
112+
tracing::info!("Model info: {:?}", model_info);
113+
110114
let mut agent = agent::Agent::new(
111115
config.clone(),
112116
control_receiver,
@@ -120,7 +124,7 @@ async fn main() -> color_eyre::Result<()> {
120124
});
121125

122126
let terminal = init_tui().unwrap();
123-
let result = tui::App::new(config, control_sender, output_receiver, history)
127+
let result = tui::App::new(config, model_info, control_sender, output_receiver, history)
124128
.run(terminal)
125129
.await;
126130
let _ = agent_handler.await;

‎src/providers/anthropic_models.json‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
[
2+
{
3+
"model_id": "claude-opus-4",
4+
"input_price": 0.000015,
5+
"output_price": 0.000075,
6+
"max_context_tokens": 200000
7+
},
8+
{
9+
"model_id": "claude-sonnet-4",
10+
"input_price": 0.000003,
11+
"output_price": 0.000015,
12+
"max_context_tokens": 200000
13+
},
14+
{
15+
"model_id": "claude-sonnet-3.7",
16+
"input_price": 0.000003,
17+
"output_price": 0.000015,
18+
"max_context_tokens": 200000
19+
},
20+
{
21+
"model_id": "claude-sonnet-3.5",
22+
"input_price": 0.000003,
23+
"output_price": 0.000015,
24+
"max_context_tokens": 200000
25+
},
26+
{
27+
"model_id": "claude-haiku-3.5",
28+
"input_price": 0.0000008,
29+
"output_price": 0.000004,
30+
"max_context_tokens": 200000
31+
},
32+
{
33+
"model_id": "claude-opus-3",
34+
"input_price": 0.000015,
35+
"output_price": 0.000075,
36+
"max_context_tokens": 200000
37+
},
38+
{
39+
"model_id": "claude-haiku-3",
40+
"input_price": 0.00000025,
41+
"output_price": 0.00000125,
42+
"max_context_tokens": 200000
43+
}
44+
]

‎src/providers/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
// Copyright © 2025 Huly Labs. Use of this source code is governed by the MIT license.
2+
23
use async_trait::async_trait;
34
use rig::agent::Agent;
45
use rig::completion::CompletionError;
56
use rig::message::Message;
67
use rig::streaming::{StreamingCompletion, StreamingCompletionResponse};
78
use rig::tool::ToolSet;
89

10+
pub mod model_info;
911
pub mod openrouter;
1012

1113
#[async_trait]

‎src/providers/model_info.rs‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
// Copyright © 2025 Huly Labs. Use of this source code is governed by the MIT license.
2+
3+
use std::{fs, path::Path};
4+
5+
use serde::Deserialize;
6+
7+
use crate::config::Config;
8+
9+
const OPENROUTER_MODELS_FILE: &str = "openrouter_models.json";
10+
const ANTHROPIC_MODELS: &str = include_str!("anthropic_models.json");
11+
const OPENAI_MODELS: &str = include_str!("openai_models.json");
12+
13+
#[derive(Debug, Clone)]
14+
pub struct ModelInfo {
15+
pub input_price: f64,
16+
pub completion_price: f64,
17+
pub max_tokens: u32,
18+
}
19+
20+
#[derive(Deserialize)]
21+
struct LMStudioModelInfo {
22+
pub id: String,
23+
pub loaded_context_length: Option<u32>,
24+
pub max_context_length: u32,
25+
}
26+
27+
#[derive(Deserialize)]
28+
struct OpenRouterPriceInfo {
29+
pub prompt: String,
30+
pub completion: String,
31+
}
32+
33+
#[derive(Deserialize)]
34+
struct OpenRouterModelInfo {
35+
pub id: String,
36+
pub pricing: OpenRouterPriceInfo,
37+
pub context_length: u32,
38+
}
39+
40+
#[derive(Deserialize)]
41+
struct AnthropicModelInfo {
42+
pub model_id: String,
43+
pub input_price: f64,
44+
pub output_price: f64,
45+
pub max_context_tokens: u32,
46+
}
47+
48+
#[derive(Deserialize)]
49+
struct OpenAIModelInfo {
50+
pub model_id: String,
51+
pub input_price: f64,
52+
pub output_price: f64,
53+
pub max_context_tokens: u32,
54+
}
55+
56+
pub async fn model_info(config: &Config) -> color_eyre::Result<ModelInfo> {
57+
match config.provider {
58+
crate::config::ProviderKind::OpenAI => {
59+
let models: Vec<OpenAIModelInfo> = serde_json::from_str(OPENAI_MODELS)?;
60+
models
61+
.iter()
62+
.find(|model| config.model.contains(&model.model_id))
63+
.map(|model| ModelInfo {
64+
input_price: model.input_price,
65+
completion_price: model.output_price,
66+
max_tokens: model.max_context_tokens,
67+
})
68+
.ok_or_else(|| color_eyre::eyre::eyre!("Model not found"))
69+
}
70+
crate::config::ProviderKind::OpenRouter => {
71+
let models: Vec<OpenRouterModelInfo> =
72+
serde_json::from_value(if Path::new(OPENROUTER_MODELS_FILE).exists() {
73+
let data = fs::read_to_string(OPENROUTER_MODELS_FILE)?;
74+
serde_json::from_str(&data)?
75+
} else {
76+
let mut data = reqwest::get("https://openrouter.ai/api/v1/models")
77+
.await?
78+
.json::<serde_json::Value>()
79+
.await?;
80+
let data = data["data"].take();
81+
fs::write(OPENROUTER_MODELS_FILE, data.to_string())?;
82+
data
83+
})?;
84+
models
85+
.iter()
86+
.find(|model| model.id == config.model)
87+
.map(|model| ModelInfo {
88+
input_price: model.pricing.prompt.parse::<f64>().unwrap_or(0.0),
89+
completion_price: model.pricing.completion.parse::<f64>().unwrap_or(0.0),
90+
max_tokens: model.context_length,
91+
})
92+
.ok_or_else(|| color_eyre::eyre::eyre!("Model not found"))
93+
}
94+
crate::config::ProviderKind::LMStudio => {
95+
let url = config
96+
.provider_base_url
97+
.clone()
98+
.unwrap_or("http://127.0.0.1:1234/v1".to_string())
99+
.replace("/v1", "/api/v0/models");
100+
let mut data = reqwest::get(url).await?.json::<serde_json::Value>().await?;
101+
let models: Vec<LMStudioModelInfo> = serde_json::from_value(data["data"].take())?;
102+
models
103+
.iter()
104+
.find(|model| model.id == config.model)
105+
.map(|model| ModelInfo {
106+
input_price: 0.0,
107+
completion_price: 0.0,
108+
max_tokens: model
109+
.loaded_context_length
110+
.unwrap_or(model.max_context_length),
111+
})
112+
.ok_or_else(|| color_eyre::eyre::eyre!("Model not found"))
113+
}
114+
crate::config::ProviderKind::Anthropic => {
115+
let models: Vec<AnthropicModelInfo> = serde_json::from_str(ANTHROPIC_MODELS)?;
116+
models
117+
.iter()
118+
.find(|model| config.model.contains(&model.model_id))
119+
.map(|model| ModelInfo {
120+
input_price: model.input_price,
121+
completion_price: model.output_price,
122+
max_tokens: model.max_context_tokens,
123+
})
124+
.ok_or_else(|| color_eyre::eyre::eyre!("Model not found"))
125+
}
126+
}
127+
}

0 commit comments

Comments
 (0)