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:
unclecode
2026-02-18 06:44:17 +00:00
parent 8576331d4e
commit c9cb0160cf
5 changed files with 726 additions and 3 deletions

View File

@@ -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

View File

@@ -204,7 +204,9 @@ llm_strategy.show_usage()
# e.g. “Total usage: 1241 tokens across 2 chunk calls”
```
If your model provider doesnt 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.
---

View File

@@ -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`)**