-
feat: simplify extraction pipeline and add batch entity summarization (#1224)
frostbyte_neo 发布于
2026-02-12 22:10:26 +00:00 - feat(llm): add token usage tracking for LLM calls
Add TokenUsageTracker class to track input/output tokens by prompt type
during LLM calls. This helps analyze token costs across different
operations like extract_nodes, extract_edges, resolve_nodes, etc.Changes:
- Add graphiti_core/llm_client/token_tracker.py with TokenUsageTracker
- Update LLMClient base class to include token_tracker instance
- Update OpenAI base client to capture and record token usage
- Add token_tracker property on Graphiti class for easy access
- Update podcast_runner.py to print token usage summary after ingestion
Usage:
client = Graphiti(...)... run ingestion ...
client.token_tracker.print_summary(sort_by='prompt_name')
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- chore: temporarily disable summary early return optimization
Disable the optimization that skips LLM calls when node summary + edge
facts is under 2000 characters. This forces all summaries to be
generated via LLM for token usage analysis.Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- Revert "chore: temporarily disable summary early return optimization"
This reverts the summary optimization changes.
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- feat: simplify extraction pipeline and add batch entity summarization
-
Remove chunking code for entity-dense episodes (node_operations.py)
- Delete _extract_nodes_chunked, _extract_from_chunk, _merge_extracted_entities
- Always use single LLM call for entity extraction
-
Remove chunking code for edge extraction (edge_operations.py)
- Remove MAX_NODES constant and generate_covering_chunks usage
- Process all nodes in single LLM call instead of covering subsets
-
Add batch entity summarization (node_operations.py, extract_nodes.py)
- New SummarizedEntity and SummarizedEntities Pydantic models
- New extract_summaries_batch prompt for batch processing
- New _extract_entity_summaries_batch function
- Nodes with short summaries get edge facts appended directly (no LLM)
- Only nodes needing LLM summarization are batched together
-
Simplify edge attribute extraction (extract_edges.py, edge_operations.py)
- Remove episode_content from context (attributes from fact only)
- Keep reference_time for temporal resolution
- Add existing_attributes to preserve/update existing values
-
Improve edge deduplication prompt (dedupe_edges.py, edge_operations.py)
- Use continuous indexing across duplicate and invalidation candidates
- Deduplicate invalidation candidates against duplicate candidates
- Allow EXISTING FACTS to be both duplicates AND contradicted
- Consolidate to single contradicted_facts field
-
Remove obsolete chunking tests (test_entity_extraction.py)
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- chore: bump version to 0.27.2pre1
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- Add token tracking for Anthropic/Gemini clients and missing tests
- Implement token tracking in AnthropicClient._generate_response()
and generate_response() using result.usage.input_tokens/output_tokens - Implement token tracking in GeminiClient._generate_response()
and generate_response() using response.usage_metadata - Add comprehensive unit tests for TokenUsageTracker class
- Add tests for _extract_entity_summaries_batch function covering:
- No nodes needing summarization
- Short summaries with edge facts
- Long summaries requiring LLM
- Node filter (should_summarize_node)
- Batch multiple nodes
- Unknown entity handling
- Missing episode and summary
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- Update test_node_operations.py for batch summarization API
- Remove import of extract_attributes_from_node (function was removed)
- Add import of _extract_entity_summaries_batch
- Update tests to use new batch summarization API
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- Add MAX_NODES limit for batch entity summarization
- Add MAX_NODES = 30 constant
- Partition nodes needing summarization into flights of MAX_NODES
- Extract _process_summary_flight helper for processing each flight
- Each flight makes a separate LLM call to avoid context overflow
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- Change default OpenAI models to gpt-5-mini
Update both DEFAULT_MODEL and DEFAULT_SMALL_MODEL to use gpt-5-mini.
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- Update podcast_runner.py to use default OpenAI models
Remove explicit model configuration to use the default gpt-5-mini models
from OpenAIClient.Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- Revert default model changes to gpt-4.1-mini/nano
Restore the original default models instead of gpt-5-mini.
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
- Address PR review comments
- Fix unreachable code in _handle_structured_response (check response.refusal)
- Process node summary flights in parallel using semaphore_gather
- Use case-insensitive name matching for LLM summary responses
- Handle duplicate node names by applying summary to all matching nodes
- Fix edge case when both edge lists are empty in contradiction processing
- Fix potential AttributeError when episode is None in edge attributes
- Add tests for flight partitioning and case-insensitive name matching
Co-Authored-By: Claude Opus 4.5 noreply@anthropic.com
Co-authored-by: Claude Opus 4.5 noreply@anthropic.com
下载附件