58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
import httpx
|
|
|
|
from .config import Settings
|
|
from .models import KnowledgeTriple
|
|
from .parsers import extract_json_object
|
|
from .prompts import KG_EXTRACT_PROMPT
|
|
|
|
|
|
class KnowledgeGraphExtractor:
|
|
"""Extract entity-relation-entity triples from page text using LLM."""
|
|
|
|
def __init__(self, settings: Settings):
|
|
self.settings = settings
|
|
|
|
async def extract(self, text: str, query: str) -> list[KnowledgeTriple]:
|
|
if not text.strip():
|
|
return []
|
|
return await asyncio.to_thread(self._extract_sync, text[:3000], query)
|
|
|
|
def _extract_sync(self, text: str, query: str) -> list[KnowledgeTriple]:
|
|
prompt = KG_EXTRACT_PROMPT.format(query=query, text=text)
|
|
payload = {
|
|
"model": self.settings.llm_model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.0,
|
|
"max_tokens": 500,
|
|
}
|
|
try:
|
|
with httpx.Client(timeout=30) as client:
|
|
resp = client.post(
|
|
self.settings.llm_chat_endpoint,
|
|
headers={
|
|
"Authorization": f"Bearer {self.settings.llm_api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json=payload,
|
|
)
|
|
resp.raise_for_status()
|
|
text_out = resp.json()["choices"][0]["message"]["content"].strip()
|
|
raw = extract_json_object(text_out)
|
|
if raw:
|
|
triples = json.loads(raw).get("triples", [])
|
|
result: list[KnowledgeTriple] = []
|
|
for t in triples[:5]:
|
|
try:
|
|
result.append(KnowledgeTriple(**t))
|
|
except Exception:
|
|
pass
|
|
return result
|
|
except Exception:
|
|
pass
|
|
return []
|