fix(ollama): metrics handling #1514
Conversation
|
The latest updates on your projects. Learn more about Vercel for Git ↗︎
|
|
Side note: I wasn't able to run the local tests for Ollama after uncommenting them out due to missing images and other odd behaviour with async calls. It might be related to these issues, or it might just be confusion on my part. Perhaps some additional instructions with explicit dependencies at the top of those tests would help making local Ollama testing easier? Another option is to see if we can get ollama deployed via CLI for testing, but that seems pretty ambitious ;) |
litellm/llms/ollama_chat.py
Outdated
| total_tokens=None, | ||
| ) | ||
|
|
||
| prompt_tokens = self.response_json["prompt_eval_count"] |
There was a problem hiding this comment.
prompt_eval_count cannot be found, but there is eval_count in the JSON response.
| @@ -120,6 +120,27 @@ def get_config(cls): | |||
| and v is not None | |||
| } | |||
|
|
|||
| # Usage metrics are only populated when the ollama response indicates `"done": true` | |||
| # https://github.com/jmorganca/ollama/blob/main/docs/api.md#generate-a-completion | |||
| class OllamaUsage: | |||
There was a problem hiding this comment.
this won't actually solve the inconsistent return of prompt eval count being returned by ollama. we already have an existing fix in place for prompt count, i believe all we're missing is the same thing for the completion tokens
There was a problem hiding this comment.
👍 Alright I can mirror the prior fix for eval_count.
I was hoping that we we'd be able to remove some of the fallback behaviour but it looks like the bugfix on Ollama's side (to reliably return values for these keys) might not be a straightforward solution so building for robustness will be a good approach.
There was a problem hiding this comment.
Applied fallbacks for both prompt_eval_count and eval_count below.
They're both encapsulated within the OllamaUsage class to provide a clear internal interface (and easier testing) for that logic.
|
@puffo let me know when the PR is ready for review again |
There was a problem hiding this comment.
Ready to review again @krrishdholakia !
I added some more type annotations and extra test cases to try increase our confidence in these fallbacks but let me know if there's more to do and I'd be happy to help :)
| @@ -120,6 +120,27 @@ def get_config(cls): | |||
| and v is not None | |||
| } | |||
|
|
|||
| # Usage metrics are only populated when the ollama response indicates `"done": true` | |||
| # https://github.com/jmorganca/ollama/blob/main/docs/api.md#generate-a-completion | |||
| class OllamaUsage: | |||
There was a problem hiding this comment.
Applied fallbacks for both prompt_eval_count and eval_count below.
They're both encapsulated within the OllamaUsage class to provide a clear internal interface (and easier testing) for that logic.
| @@ -121,6 +124,61 @@ def get_config(cls): | |||
| } | |||
|
|
|||
|
|
|||
| class ResponseJSON(TypedDict, total=False): | |||
There was a problem hiding this comment.
I found myself frequently having to revisit the documentation to check the response shape so I added these TypedDict specs to make the Ollama API specification a bit clearer.
|
|
||
| def _prompt_tokens_fallback(self): | ||
| print_verbose(f"Warning: `prompt_eval_count` missing from response, estimating by using OpenAI token counting with cl100k_base encoding.") | ||
| return openai_token_counter(messages=self.messages) |
There was a problem hiding this comment.
Token counting is a lot trickier for the /chat/completions endpoint. I figured that we should just reuse the openai_token_counter for the fallback and provide a warning message?
|
Any news on this? I also suffer from this issue :) |
It has been fixed in this commit and released in v1.19.2 . Happy to close this if you'd rather avoid the extra overhead @krrishdholakia |
|
It still persists for me even in 1.19.2 - I tested it yesterday. I'd like to help track it down if it is an issue with LiteLLM. |
The previous fix was flawed as the /chat API is different from the /generate API.
While fixing the regression, I noticed that there is inconsistent behaviour in Ollama where the
prompt_eval_countdisappears after subsequent requests, yetprompt_eval_durationis persisted. I believe this inconsistent behaviour is why the workaround is needed in the first place and adds some unnecessary complexity in litellm (which might otherwise not be required).I suspect this is a bug on Ollama's side and I've opened an issue to confirm my assumptions with the community there. ollama/ollama#2068
I'm pushing my changes here in the meantime for review and it can be merged up once we get clarity from the Ollama team.