Add token usage tracking to generate_schema / agenerate_schema
generate_schema can make up to 5 internal LLM calls (field inference, schema generation, validation retries) with no way to track token consumption. Add an optional `usage: TokenUsage = None` parameter that accumulates prompt/completion/total tokens across all calls in-place. - _infer_target_json: accept and populate usage accumulator - agenerate_schema: track usage after every aperform_completion call in the retry loop, forward usage to _infer_target_json - generate_schema (sync): forward usage to agenerate_schema Fully backward-compatible — omitting usage changes nothing.
This commit is contained in:
@@ -4537,6 +4537,20 @@ xpath_schema = JsonXPathExtractionStrategy.generate_schema(
|
||||
# Use the generated schema for fast, repeated extractions
|
||||
strategy = JsonCssExtractionStrategy(css_schema)
|
||||
```
|
||||
### Token Usage Tracking
|
||||
`generate_schema` may make multiple LLM calls internally (field inference, generation, validation retries). Track total token consumption by passing a `TokenUsage` accumulator:
|
||||
```python
|
||||
from crawl4ai.models import TokenUsage
|
||||
|
||||
usage = TokenUsage()
|
||||
schema = JsonCssExtractionStrategy.generate_schema(
|
||||
url="https://example.com/products",
|
||||
query="extract product name and price",
|
||||
usage=usage,
|
||||
)
|
||||
print(f"Total tokens: {usage.total_tokens}")
|
||||
```
|
||||
The `usage` parameter is optional and fully backward-compatible. Both `generate_schema` (sync) and `agenerate_schema` (async) support it.
|
||||
### LLM Provider Options
|
||||
1. **OpenAI GPT-4 (`openai/gpt4o`)**
|
||||
- Default provider
|
||||
|
||||
@@ -204,7 +204,9 @@ llm_strategy.show_usage()
|
||||
# e.g. “Total usage: 1241 tokens across 2 chunk calls”
|
||||
```
|
||||
|
||||
If your model provider doesn’t return usage info, these fields might be partial or empty.
|
||||
If your model provider doesn't return usage info, these fields might be partial or empty.
|
||||
|
||||
> **Tip:** `JsonCssExtractionStrategy.generate_schema()` also supports token usage tracking via an optional `usage` parameter. See [Token Usage Tracking in Schema Generation](./no-llm-strategies.md#token-usage-tracking) for details.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -761,6 +761,38 @@ schema = JsonCssExtractionStrategy.generate_schema(
|
||||
|
||||
The generator also understands sibling layouts — for sites like Hacker News where data is split across sibling elements, it will automatically use the [`source` field](#sibling-data) to reach sibling data.
|
||||
|
||||
### Token Usage Tracking
|
||||
|
||||
`generate_schema` may make multiple LLM calls internally (field inference, schema generation, validation retries). To track the total token consumption across all of these calls, pass a `TokenUsage` accumulator:
|
||||
|
||||
```python
|
||||
from crawl4ai import JsonCssExtractionStrategy
|
||||
from crawl4ai.models import TokenUsage
|
||||
|
||||
usage = TokenUsage()
|
||||
|
||||
schema = JsonCssExtractionStrategy.generate_schema(
|
||||
url="https://news.ycombinator.com",
|
||||
query="Extract each story: title, url, score, author",
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
print(f"Prompt tokens: {usage.prompt_tokens}")
|
||||
print(f"Completion tokens: {usage.completion_tokens}")
|
||||
print(f"Total tokens: {usage.total_tokens}")
|
||||
```
|
||||
|
||||
The `usage` parameter is optional — omitting it changes nothing (fully backward-compatible). You can also reuse the same accumulator across multiple calls to get a grand total:
|
||||
|
||||
```python
|
||||
usage = TokenUsage()
|
||||
schema1 = JsonCssExtractionStrategy.generate_schema(url=url1, query=q1, usage=usage)
|
||||
schema2 = JsonCssExtractionStrategy.generate_schema(url=url2, query=q2, usage=usage)
|
||||
print(f"Grand total: {usage.total_tokens} tokens")
|
||||
```
|
||||
|
||||
Both `generate_schema` (sync) and `agenerate_schema` (async) support the `usage` parameter.
|
||||
|
||||
### LLM Provider Options
|
||||
|
||||
1. **OpenAI GPT-4 (`openai/gpt4o`)**
|
||||
|
||||
Reference in New Issue
Block a user