From b4133282a650fc16d820d207db6c1331e21e2bf2 Mon Sep 17 00:00:00 2001 From: viettranx Date: Thu, 12 Mar 2026 18:54:14 +0700 Subject: [PATCH] refactor(kg): improve extraction prompt with few-shot example and controlled vocabulary Add domain context, coreference rules, controlled relation types (15 predefined), few-shot example, and dynamic entity count (3-15). Increase max input from 6000 to 12000 chars, reduce max output tokens from 8192 to 4096. --- internal/knowledgegraph/extractor.go | 4 +- internal/knowledgegraph/extractor_prompt.go | 67 +++++++++++++++++---- 2 files changed, 56 insertions(+), 15 deletions(-) diff --git a/internal/knowledgegraph/extractor.go b/internal/knowledgegraph/extractor.go index b9fb19af..ffddb7da 100644 --- a/internal/knowledgegraph/extractor.go +++ b/internal/knowledgegraph/extractor.go @@ -35,7 +35,7 @@ func NewExtractor(provider providers.Provider, model string, minConfidence float // Extract calls the LLM to extract entities and relations from text. func (e *Extractor) Extract(ctx context.Context, text string) (*ExtractionResult, error) { // Truncate very long texts to avoid overwhelming the LLM - const maxInputChars = 6000 + const maxInputChars = 12000 if len(text) > maxInputChars { text = text[:maxInputChars] + "\n\n[...truncated]" } @@ -47,7 +47,7 @@ func (e *Extractor) Extract(ctx context.Context, text string) (*ExtractionResult }, Model: e.model, Options: map[string]any{ - "max_tokens": 8192, + "max_tokens": 4096, "temperature": 0.0, }, } diff --git a/internal/knowledgegraph/extractor_prompt.go b/internal/knowledgegraph/extractor_prompt.go index b6f67d6b..0bb768c1 100644 --- a/internal/knowledgegraph/extractor_prompt.go +++ b/internal/knowledgegraph/extractor_prompt.go @@ -1,6 +1,6 @@ package knowledgegraph -const extractionSystemPrompt = `You are an entity and relationship extractor. Given text, extract structured knowledge. +const extractionSystemPrompt = `You are a knowledge graph extractor for an AI assistant's memory system. Given text (usually personal notes, work logs, or conversation summaries), extract the most important entities and their relationships. Output valid JSON with this schema: { @@ -9,26 +9,67 @@ Output valid JSON with this schema: "external_id": "unique-lowercase-id", "name": "Display Name", "entity_type": "person|project|task|event|concept|location|organization", - "description": "Brief description (max 50 chars)", + "description": "Brief description of the entity", "confidence": 0.0-1.0 } ], "relations": [ { - "source_entity_id": "external_id of source entity", - "relation_type": "verb_phrase", - "target_entity_id": "external_id of target entity", + "source_entity_id": "external_id of source", + "relation_type": "RELATION_TYPE", + "target_entity_id": "external_id of target", "confidence": 0.0-1.0 } ] } -Rules: -- Extract only the TOP 15 most important entities and their relations -- Keep descriptions very short (under 50 characters) -- external_id: lowercase, no spaces, use hyphens (e.g. "john-doe", "project-alpha") -- entity_type: person, project, task, event, concept, location, or organization -- relation_type: lowercase verb phrase with underscores (e.g. works_on, reported_to) -- confidence: 1.0 = explicitly stated, 0.5 = inferred +## Entity ID Rules +- Use consistent, canonical lowercase IDs with hyphens +- For people: use full name when known (e.g., "john-doe"), not partial ("john") +- For projects/products: use official name (e.g., "project-alpha", "goclaw") +- Same real-world entity MUST always get the same external_id across extractions +- When a pronoun or partial reference clearly refers to a named entity, use that entity's ID — do NOT create a new entity + +## Entity Types +- person: named individuals +- organization: companies, teams, departments +- project: software projects, initiatives, products +- task: specific work items, tickets, TODOs +- event: meetings, releases, incidents, deadlines +- concept: technologies, methodologies, domains +- location: cities, offices, regions + +## Relation Types (use ONLY these) +- works_on, manages, reports_to, collaborates_with (people↔work) +- belongs_to, part_of, depends_on, blocks (structure) +- created, completed, assigned_to, scheduled_for (actions) +- located_in, based_at (location) +- uses, implements, integrates_with (technology) +- related_to (fallback — use sparingly) + +## Rules +- Extract 3-15 entities depending on text density. Short text = fewer entities +- Confidence: 1.0 = explicitly stated, 0.7 = strongly implied, 0.4 = weakly inferred - Keep names in original language -- Output ONLY the JSON object, no markdown, no explanation, no code blocks` +- Descriptions: 1 sentence max, capture the entity's role or significance +- Skip generic/vague entities ("the system", "the team" without specific name) +- Output ONLY the JSON object, no markdown, no code blocks + +## Example + +Input: "Talked to Minh about the GoClaw migration. He'll handle the database schema changes by Friday. The team uses PostgreSQL with pgvector." + +Output: +{ + "entities": [ + {"external_id": "minh", "name": "Minh", "entity_type": "person", "description": "Handling database schema changes for GoClaw", "confidence": 1.0}, + {"external_id": "goclaw", "name": "GoClaw", "entity_type": "project", "description": "Project undergoing migration", "confidence": 1.0}, + {"external_id": "goclaw-migration", "name": "GoClaw Migration", "entity_type": "task", "description": "Database migration task for GoClaw", "confidence": 1.0}, + {"external_id": "postgresql", "name": "PostgreSQL", "entity_type": "concept", "description": "Database technology used with pgvector", "confidence": 1.0} + ], + "relations": [ + {"source_entity_id": "minh", "relation_type": "works_on", "target_entity_id": "goclaw-migration", "confidence": 1.0}, + {"source_entity_id": "goclaw-migration", "relation_type": "part_of", "target_entity_id": "goclaw", "confidence": 1.0}, + {"source_entity_id": "goclaw", "relation_type": "uses", "target_entity_id": "postgresql", "confidence": 1.0} + ] +}`