Skip to content

Knowledge Graph Extraction

Learn how to automatically extract entities and relationships from your documents.

What is Extraction?

Extraction is the process of analyzing a document and identifying:

  • Entities — Things like people, organizations, locations, concepts
  • Relationships — Connections between entities (e.g., "works for", "located in")

GraphAide uses LLMs to understand document context and extract this structured information.

Quick Start: Extract from a PDF

Simplest Example

from graphaide import GraphAide, ModelConfig, get_sample_data_path

graphaide = GraphAide(
    model_config=ModelConfig(
        provider="openai",
        model_name="gpt-5.6-luna",
        api_key="..."
    )
)

# Extract entities and relationships
result = graphaide.extract(file_path="document.pdf")

print(f"Found {len(result.nodes)} entities")
print(f"Found {len(result.edges)} relationships")

# View extracted data
for node in result.nodes:
    print(f"- {node.label}: {node.properties}")

With Sample Data

from graphaide import get_sample_data_path

sample_pdf = get_sample_data_path("PNNL_About.pdf")
sample_ontology = get_sample_data_path("ontology_cyber.nt")

result = graphaide.extract(
    file_path=sample_pdf,
    ontology_path=sample_ontology
)

Input Formats

GraphAide can extract from multiple input types:

From a File

result = graphaide.extract(file_path="document.pdf")
# or .txt, .docx, .md, etc.

From Text String

text = """
GraphAide is developed by Sumit Purohit at PNNL.
It extracts knowledge from documents and loads it into Neo4j.
"""

result = graphaide.extract(input_text=text)

From Multiple Files (JSONL)

For batch processing multiple documents:

# Create JSONL file: one JSON object per line
# {"text": "Document content here"}
# {"text": "Another document"}

result = graphaide.extract(
    jsonl_file_path="documents.jsonl",
    jsonl_field_name="text",
    jsonl_lines_per_batch=4
)

Output: Understanding Results

Access Extracted Data

result = graphaide.extract(file_path="document.pdf")

# Nodes (entities)
for node in result.nodes:
    print(f"Type: {node.type}")
    print(f"Label: {node.label}")
    print(f"Properties: {node.properties}")

# Edges (relationships)
for edge in result.edges:
    print(f"Relationship: {edge.type}")
    print(f"From: {edge.source} → To: {edge.target}")
    print(f"Properties: {edge.properties}")

Save Results to JSON

result = graphaide.extract(file_path="document.pdf")

# Results are automatically saved to:
# GraphAide_KG_Extract_YYYYMMDD_HHMMSS.json

Use a Knowledge Base for Disambiguation

GraphAide comes with a default knowledge base snapshot of Wikidata out of the box. It contains entries for common entities such as countries, US states, companies, weekdays, months, and government departments. During extraction, each entity is matched against this knowledge base to get a stable Wikidata QID, so that "Georgia" the country and "Georgia" the state are kept apart. You can extend the knowledge base with your own entities for your domain.

Load the Bundled Knowledge Base

from graphaide import GraphAide, ModelConfig, VectorDBConfig, get_sample_data_path

graphaide = GraphAide(
    model_config=ModelConfig(provider="openai", model_name="gpt-4o", api_key="..."),
    vector_configs=VectorDBConfig(provider="chromadb", store_name="GA_VDB"),
)

kb_file = get_sample_data_path("graphaide_KB.jsonl")
graphaide.ingest_vectors(file_paths=[str(kb_file)], append=False)

result = graphaide.extract(input_text="Nvidia released new GPUs. Anthropic builds Claude.", merge=True)
for node in result.nodes:
    print(node.node_name, node.wikidata_id)

Add Your Own Entities

Write a JSONL file with one object per line. Each object needs a qid, a label, and optionally a description, aliases, and edges:

{"qid": "Q9000001", "label": "Acme Robotics", "description": "American robotics company", "aliases": ["Acme"], "edges": []}
{"qid": "Q9000002", "label": "Project Helix", "description": "internal battery research program", "aliases": ["Helix"], "edges": []}

Use real Wikidata QIDs when they exist, and make up stable IDs (for example X0001) for internal entities. Load your file alongside the bundled one:

graphaide.ingest_vectors(
    file_paths=[
        str(get_sample_data_path("graphaide_KB.jsonl")),
        "my_entities.jsonl",
    ],
    append=False,
)

Files whose name contains graphaide_kb or wikidata and end in .jsonl are loaded as Wikidata entities. Name your file accordingly, or it will be loaded as a regular document.

Extract with Your Knowledge Base

result = graphaide.extract(
    input_text="Acme Robotics and Project Helix announced a partnership with Nvidia.",
    merge=True,
)
for node in result.nodes:
    print(f"{node.node_name:20} {node.wikidata_id}")

Entities that match one of your entries get its QID. Entities with no match stay without wikidata_id. See Disambiguation for how matches are chosen.

Advanced Options

Control Context Window

result = graphaide.extract(
    file_path="document.pdf",
    max_tokens_per_segment=8192,
    context_window_mode="large"
)

Filter by Ontology

Keep only entities and relationships that match your ontology:

result = graphaide.extract(
    file_path="document.pdf",
    ontology_path="ontology.nt",
    filter_by_ontology=True  # Only keep types defined in ontology
)

Merge Results

Automatically deduplicate and consolidate extracted data:

result = graphaide.extract(
    file_path="document.pdf",
    merge_output=True  # Combine duplicate entities
)

Processing Steps

Here's what GraphAide does internally:

1. Load Document
   ↓
2. Split into Segments (chunks)
   ↓
3. For Each Segment:
   - Retrieve Context (from Vector Store/Ontology)
   - Send to LLM for Extraction
   - Parse Entities & Relationships
   ↓
4. Combine Results
   ↓
5. Save to JSON

Common Extraction Patterns

Business Documents

result = graphaide.extract(
    file_path="contract.pdf",
    ontology_path="ontology_business.nt"
)

Scientific Papers

result = graphaide.extract(
    file_path="research_paper.pdf",
    ontology_path="ontology_science.nt",
    filter_by_ontology=True
)

Internal Documentation

result = graphaide.extract(
    file_path="procedure.txt",
    merge_output=True  # Deduplicate for cleaner results
)

Next Steps