File size: 4,760 Bytes
8e874f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | """Indexing module - knowledge graph building components.
This module provides modular components for building knowledge graphs
from various sources:
Components:
- Chunkers: Text splitting strategies (TokenChunker)
- Extractors: Entity/relationship extraction (LLMEntityExtractor)
- Builders: Schema-based graph construction (DatabaseSchemaBuilder, ExcelSchemaBuilder)
Usage:
# Direct imports
from src.indexing import TokenChunker, DatabaseSchemaBuilder
chunker = TokenChunker(max_token_size=1024)
chunks = chunker.chunk("Long text here...")
# Or use registry pattern
from src.indexing import get_chunker, get_indexer
chunker = get_chunker("token", max_token_size=512)
builder = get_indexer("schema", graph_storage=..., entities_vdb=..., relationships_vdb=...)
Adding Custom Components:
from src.indexing import BaseChunker, register_chunker
class MyChunker(BaseChunker):
def chunk(self, content, **kwargs):
# Implementation
pass
register_chunker("my_chunker", MyChunker)
"""
from .base import (
BaseChunker,
BaseExtractor,
BaseIndexer,
ChunkResult,
ExtractionResult,
IndexingResult,
)
from .chunkers import TokenChunker, chunking_by_token_size
from .extractors import (
extract_entities,
_handle_entity_relation_summary,
_handle_single_entity_extraction,
_handle_single_relationship_extraction,
_merge_nodes_then_upsert,
_merge_edges_then_upsert,
)
from .schema_builder import DatabaseSchemaBuilder, detect_schema_type
from .excel_builder import ExcelSchemaBuilder
# =============================================================================
# Registry Pattern
# =============================================================================
CHUNKER_REGISTRY = {
"token": TokenChunker,
}
INDEXER_REGISTRY = {
"schema": DatabaseSchemaBuilder,
"database": DatabaseSchemaBuilder, # Alias
"excel": ExcelSchemaBuilder,
}
def get_chunker(name: str, **kwargs) -> BaseChunker:
"""Factory function to get a chunker by name.
Args:
name: Chunker name (e.g., "token")
**kwargs: Arguments to pass to the chunker constructor
Returns:
Instantiated chunker
Raises:
ValueError: If chunker name is not found
Example:
chunker = get_chunker("token", max_token_size=512)
"""
if name not in CHUNKER_REGISTRY:
available = ", ".join(CHUNKER_REGISTRY.keys())
raise ValueError(f"Unknown chunker: '{name}'. Available: {available}")
return CHUNKER_REGISTRY[name](**kwargs)
def get_indexer(name: str, **kwargs) -> BaseIndexer:
"""Factory function to get an indexer by name.
Args:
name: Indexer name (e.g., "schema", "excel")
**kwargs: Arguments to pass to the indexer constructor
Returns:
Instantiated indexer
Raises:
ValueError: If indexer name is not found
Example:
indexer = get_indexer("schema", graph_storage=gs, entities_vdb=ev, relationships_vdb=rv)
"""
if name not in INDEXER_REGISTRY:
available = ", ".join(INDEXER_REGISTRY.keys())
raise ValueError(f"Unknown indexer: '{name}'. Available: {available}")
return INDEXER_REGISTRY[name](**kwargs)
def register_chunker(name: str, cls: type):
"""Register a custom chunker class.
Args:
name: Name to register the chunker under
cls: Chunker class (should inherit from BaseChunker)
"""
CHUNKER_REGISTRY[name] = cls
def register_indexer(name: str, cls: type):
"""Register a custom indexer class.
Args:
name: Name to register the indexer under
cls: Indexer class (should inherit from BaseIndexer)
"""
INDEXER_REGISTRY[name] = cls
def list_chunkers() -> list:
"""List all available chunker names."""
return list(CHUNKER_REGISTRY.keys())
def list_indexers() -> list:
"""List all available indexer names."""
return list(INDEXER_REGISTRY.keys())
__all__ = [
# Base classes
"BaseChunker",
"BaseExtractor",
"BaseIndexer",
"ChunkResult",
"ExtractionResult",
"IndexingResult",
# Chunkers
"TokenChunker",
"chunking_by_token_size",
# Extractors
"extract_entities",
"_handle_entity_relation_summary",
"_handle_single_entity_extraction",
"_handle_single_relationship_extraction",
"_merge_nodes_then_upsert",
"_merge_edges_then_upsert",
# Builders
"DatabaseSchemaBuilder",
"ExcelSchemaBuilder",
"detect_schema_type",
# Registry functions
"get_chunker",
"get_indexer",
"register_chunker",
"register_indexer",
"list_chunkers",
"list_indexers",
"CHUNKER_REGISTRY",
"INDEXER_REGISTRY",
]
|