Having recently integrated Langfuse with a locally hosted Llama 3.1 instance for a client's internal RAG pipeline, I found the token counting mechanism to be non-trivial. Unlike with OpenAI's API, where token usage is returned in the response headers, locally hosted models via `ollama` or `vLLM` do not provide this data out-of-the-box. Langfuse's built-in OpenAI integration handles it automatically, but for local models, we must instrument the counting ourselves.
The core challenge is that we need to replicate the tokenization process used by the specific Llama model to get accurate counts. My approach involves using the Hugging Face `transformers` library to load the exact tokenizer, ensuring counts match what the model actually processes. Below is the middleware pattern I implemented, which sits between the application and the model, intercepting prompts and completions.
```python
from transformers import AutoTokenizer
import langfuse
from langfuse.decorators import observe, langfuse_context
# Initialize tokenizer for your specific model
model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
def count_tokens(text: str) -> int:
"""Count tokens using the model's native tokenizer."""
return len(tokenizer.encode(text))
@observe()
def generate_with_local_llama(prompt: str):
# Manually count prompt tokens
prompt_tokens = count_tokens(prompt)
# Your existing logic to call the local model (e.g., via ollama, vLLM API)
# For example:
# completion = requests.post('http://localhost:11434/api/generate', json={...}).json()
completion_text = call_local_llama_model(prompt) # Your function here
# Count completion tokens
completion_tokens = count_tokens(completion_text)
# Set token usage on the current Langfuse observation
langfuse_context.set_usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens
)
return completion_text
```
Key considerations for production:
* **Tokenizer Matching:** The tokenizer must be the same version as your deployed model. Mismatches can cause significant count errors.
* **Overhead:** Loading the tokenizer for every request is inefficient. Implement it as a singleton or a cached global variable.
* **vLLM & TGI:** If using vLLM or Text Generation Inference, they may provide token counts in their API responses. Always check their response schema first before applying manual counting.
* **Cost Tracking:** While local models may not have a monetary cost, accurate token usage in Langfuse is crucial for performance analytics, prompt engineering, and comparing model efficiencies.
For teams using `ollama` directly, the `/api/generate` endpoint does return a `prompt_eval_count` and `eval_count` in its streaming responses. You can aggregate these and set the total usage upon completion of the stream.
API first.
IntegrationWizard