| import json |
| import re |
| import time |
| from openai import OpenAI |
| from src.models import Chunk, Entity, Relationship |
| from src.config import OPENCODE_ZEN_API_KEY, LLM_BASE_URL, LLM_MODEL |
| from src.retry import with_retry |
|
|
|
|
| _client: OpenAI | None = None |
|
|
|
|
| def _get_client() -> OpenAI: |
| global _client |
| if _client is None: |
| _client = OpenAI(api_key=OPENCODE_ZEN_API_KEY, base_url=LLM_BASE_URL) |
| return _client |
|
|
|
|
| _EXTRACTION_SYSTEM_PROMPT = """You are a knowledge graph extractor. Given a text chunk, extract all named entities and the relationships between them. |
| |
| Rules: |
| - Entities are real-world objects, concepts, people, places, organizations, technologies. |
| - Relationships describe how entities connect (e.g., "works_for", "located_in", "part_of", "developed_by", "invented"). |
| - Use simple relation types (lowercase, underscores). |
| |
| Output ONLY valid JSON array with no markdown: |
| [ |
| {"entity": "EntityName", "type": "Person|Organization|Technology|Location|Concept|Field", "relationships": [{"target": "OtherEntity", "relation": "relation_type"}]} |
| ]""" |
|
|
|
|
| @with_retry(max_retries=5, base_delay=3.0) |
| def _call_extraction(text: str): |
| client = _get_client() |
| return client.chat.completions.create( |
| model=LLM_MODEL, |
| messages=[ |
| {"role": "system", "content": _EXTRACTION_SYSTEM_PROMPT}, |
| {"role": "user", "content": text}, |
| ], |
| max_tokens=2048, |
| temperature=0.1, |
| ) |
|
|
|
|
| def extract_knowledge(chunks: list[Chunk], status_callback=None) -> tuple[list[Entity], list[Relationship]]: |
| all_entities: list[Entity] = [] |
| all_relationships: list[Relationship] = [] |
|
|
| for i, chunk in enumerate(chunks): |
| if i > 0: |
| time.sleep(3.0) |
|
|
| if status_callback: |
| status_callback(f"extracting: Chunk {i + 1}/{len(chunks)} - {chunk.source}") |
|
|
| response = _call_extraction(chunk.text) |
|
|
| raw = response.choices[0].message.content or "[]" |
| extracted = _parse_extraction(raw) |
|
|
| for item in extracted: |
| entity_name = item.get("entity", "").strip() |
| entity_type = item.get("type", "Concept") |
| if not entity_name: |
| continue |
|
|
| entity_id = f"{chunk.id}_{entity_name.lower().replace(' ', '_')}" |
| entity = Entity( |
| id=entity_id, |
| name=entity_name, |
| type=entity_type, |
| chunk_id=chunk.id, |
| ) |
| all_entities.append(entity) |
|
|
| for rel in item.get("relationships", []): |
| target = rel.get("target", "").strip() |
| relation = rel.get("relation", "related_to").strip().lower() |
| if not target: |
| continue |
|
|
| target_id = f"{chunk.id}_{target.lower().replace(' ', '_')}" |
| rel_id = f"{entity_id}_to_{target_id}" |
| relationship = Relationship( |
| id=rel_id, |
| source_entity=entity_name, |
| target_entity=target, |
| relation_type=relation, |
| chunk_id=chunk.id, |
| ) |
| all_relationships.append(relationship) |
|
|
| return all_entities, all_relationships |
|
|
|
|
| def _parse_extraction(raw: str) -> list[dict]: |
| cleaned = raw.strip() |
| if cleaned.startswith("```"): |
| cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned) |
| cleaned = re.sub(r"\s*```$", "", cleaned) |
|
|
| try: |
| data = json.loads(cleaned) |
| return data if isinstance(data, list) else [] |
| except (json.JSONDecodeError, ValueError): |
| matches = re.findall(r'"entity"\s*:\s*"([^"]+)"', cleaned) |
| return [{"entity": m, "type": "Concept", "relationships": []} for m in matches] |
|
|
