-
Notifications
You must be signed in to change notification settings - Fork 35
Expand file tree
/
Copy path__init__.py
More file actions
153 lines (127 loc) · 4.82 KB
/
Copy path__init__.py
File metadata and controls
153 lines (127 loc) · 4.82 KB
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
"""High-level helpers for working with Lance-backed knowledge graphs."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, Mapping, Optional
import pyarrow as pa
from lance_graph import CypherQuery, GraphConfig
try: # Prefer to import for typing without raising at runtime.
from lance_graph import GraphConfigBuilder
except ImportError: # pragma: no cover - builder is available in normal installs.
GraphConfigBuilder = object # type: ignore[assignment]
from lance_graph import DistanceMetric, VectorSearch
from .component import KnowledgeGraphComponent
from .config import KnowledgeGraphConfig, build_graph_config_from_mapping
from .extraction import (
DEFAULT_STRATEGY,
BaseExtractor,
get_extractor,
preview_extraction,
)
from .extractors import HeuristicExtractor, LLMExtractor
from .service import LanceKnowledgeGraph, create_default_service
from .store import LanceGraphStore
from .webservice import create_app
TableMapping = Mapping[str, pa.Table]
def _ensure_table(name: str, table: pa.Table) -> pa.Table:
if not isinstance(table, pa.Table):
raise TypeError(
f"Dataset '{name}' must be a pyarrow.Table (got {type(table)!r})"
)
return table
@dataclass(frozen=True)
class KnowledgeGraph:
"""Wraps a ``GraphConfig`` alongside the Arrow tables backing it."""
config: GraphConfig
_tables: Dict[str, pa.Table]
def __init__(self, config: GraphConfig, datasets: TableMapping) -> None:
object.__setattr__(self, "config", config)
normalized = {
name: _ensure_table(name, table) for name, table in datasets.items()
}
object.__setattr__(self, "_tables", normalized)
def run(
self,
statement: str,
*,
datasets: Optional[TableMapping] = None,
):
"""Execute a Cypher statement, overriding tables when provided."""
query = CypherQuery(statement).with_config(self.config)
sources: Dict[str, pa.Table] = dict(self._tables)
if datasets:
sources.update(
{name: _ensure_table(name, table) for name, table in datasets.items()}
)
return query.execute(sources)
def run_with_vector_rerank(
self,
statement: str,
vector_search: "VectorSearch",
*,
datasets: Optional[TableMapping] = None,
) -> pa.Table:
"""Execute a Cypher statement and rerank results by vector similarity."""
query = CypherQuery(statement).with_config(self.config)
sources: Dict[str, pa.Table] = dict(self._tables)
if datasets:
sources.update(
{name: _ensure_table(name, table) for name, table in datasets.items()}
)
return query.execute_with_vector_rerank(sources, vector_search)
def tables(self) -> Dict[str, pa.Table]:
"""Return a shallow copy of the registered datasets."""
return dict(self._tables)
class KnowledgeGraphBuilder:
"""Collects nodes, relationships, and datasets before building a graph."""
def __init__(self) -> None:
builder = GraphConfig.builder()
self._builder: GraphConfigBuilder = builder # type: ignore[annotation-unchecked]
self._datasets: Dict[str, pa.Table] = {}
def with_node(
self,
label: str,
primary_key: str,
table: pa.Table,
) -> KnowledgeGraphBuilder:
"""Register a node label and Arrow table."""
self._builder = self._builder.with_node_label(label, primary_key)
self._datasets[label] = _ensure_table(label, table)
return self
def with_relationship(
self,
name: str,
source_key: str,
target_key: str,
table: pa.Table,
) -> KnowledgeGraphBuilder:
"""Register a relationship and its underlying table."""
self._builder = self._builder.with_relationship(name, source_key, target_key)
self._datasets[name] = _ensure_table(name, table)
return self
def with_dataset(self, name: str, table: pa.Table) -> KnowledgeGraphBuilder:
"""Attach arbitrary supporting datasets (e.g., reference tables)."""
self._datasets[name] = _ensure_table(name, table)
return self
def build(self) -> KnowledgeGraph:
"""Materialize the ``KnowledgeGraph`` instance."""
config = self._builder.build()
return KnowledgeGraph(config, self._datasets)
__all__ = [
"KnowledgeGraph",
"KnowledgeGraphBuilder",
"KnowledgeGraphConfig",
"build_graph_config_from_mapping",
"LanceGraphStore",
"LanceKnowledgeGraph",
"create_default_service",
"KnowledgeGraphComponent",
"create_app",
"DEFAULT_STRATEGY",
"BaseExtractor",
"get_extractor",
"preview_extraction",
"HeuristicExtractor",
"LLMExtractor",
"VectorSearch",
"DistanceMetric",
]