Repository navigation
Expand file tree
/
Copy pathextraction.py
More file actions
84 lines (65 loc) · 2.35 KB
/
Copy pathextraction.py
File metadata and controls
84 lines (65 loc) · 2.35 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
"""High-level extraction utilities combining heuristic and LLM strategies."""
from __future__ import annotations
from typing import Any, Optional
from .extractors import HeuristicExtractor, LLMExtractor
from .extractors.base import ExtractedEntity, ExtractedRelationship, ExtractionResult
from .extractors.llm import load_json_payload
DEFAULT_STRATEGY = "llm"
def get_extractor(
strategy: str = DEFAULT_STRATEGY,
*,
llm_callable=None,
llm_model: str = "gpt-4o-mini",
llm_temperature: float = 0.2,
llm_options: Optional[dict] = None,
) -> "BaseExtractor":
strategy = strategy.lower()
if strategy == "heuristic":
return HeuristicExtractor()
if strategy == "llm":
return LLMExtractor(
send_callable=llm_callable,
model=llm_model,
temperature=llm_temperature,
client_options=llm_options or {},
)
raise ValueError(f"Unknown extractor strategy '{strategy}'")
def preview_extraction(
text: str,
*,
extractor: Optional["BaseExtractor"] = None,
strategy: str = DEFAULT_STRATEGY,
) -> ExtractionResult:
extractor = extractor or get_extractor(strategy)
entities = extractor.extract_entities(text)
relationships = extractor.extract_relationships(text, entities)
return ExtractionResult(entities=entities, relationships=relationships)
class BaseExtractor:
"""Protocol-like base class for IDE support."""
def extract_entities(self, text: str) -> list[ExtractedEntity]: # pragma: no cover
raise NotImplementedError
def extract_relationships(
self, text: str, entities: list[ExtractedEntity]
) -> list[ExtractedRelationship]: # pragma: no cover
raise NotImplementedError
class LLMClient:
def __init__(self, extractor: LLMExtractor):
self._extractor = extractor
def complete(self, prompt: str) -> str:
return self._extractor.generate(prompt)
def get_llm_client(
*,
llm_callable=None,
llm_model: str = "gpt-4o-mini",
llm_temperature: float = 0.2,
llm_options: Optional[dict] = None,
) -> LLMClient:
extractor = LLMExtractor(
send_callable=llm_callable,
model=llm_model,
temperature=llm_temperature,
client_options=llm_options or {},
)
return LLMClient(extractor)
def parse_llm_json(raw: str) -> Any:
return load_json_payload(raw)