apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-06-17 15:25:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

When using ollama token count was not coming in

Browse source 
When using ollama token count was not coming in resulting in token count and other metrics to show up as zero. This was not causing tracing to break.
This commit is contained in:
Adil Hafeez 2025-01-21 17:13:23 -08:00
parent bea0dd4a83
commit ca41f92d67
1 changed files with 7 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

11
crates/llm_gateway/src/stream_context.rs
View file

@ -483,11 +483,14 @@ impl HttpContext for StreamContext {
let tokens_str = chat_completions_chunk_response_events.to_string();
//HACK: add support for tokenizing mistral and other models
//filed issue https://github.com/katanemo/arch/issues/222
if model.as_ref().unwrap().starts_with("mistral")
|| model.as_ref().unwrap().starts_with("ministral")
{
model = Some("gpt-4".to_string());
if !model.as_ref().unwrap().starts_with("gpt") {
warn!(
"tiktoken_rs: unsupported model: {}, using gpt-4 to compute token count",
model.as_ref().unwrap()
);
}
model = Some("gpt-4".to_string());
let token_count =
match tokenizer::token_count(model.as_ref().unwrap().as_str(), tokens_str.as_str())
{