Extract and resolve entities

This guide shows how to extract entities and relations from chunks, and how to resolve duplicate mentions into groups.

Extract with the default schema

import asyncio
from agrag.common.data_models.graph_schema import GENERIC
from agrag.ingestion.extract import BAMLExtractor, ExtractionLLMSettings


async def main():
    # Point at your OpenAI-compatible endpoint via environment variables:
    #   LLM_BASE_URL, LLM_MODEL_ID, and optionally LLM_API_KEY
    settings = ExtractionLLMSettings.from_openai_compatible_env()
    extractor = BAMLExtractor(settings=settings)
    result = await extractor.extract(chunk, GENERIC)
    print(result.entities)


asyncio.run(main())

BAMLExtractor needs the llm extra:

pip install "agentic-graphrag[llm]"

Run locally first, escalating to an LLM

import asyncio
from agrag.common.data_models.graph_schema import GENERIC
from agrag.ingestion.extract import (
    BAMLExtractor,
    EscalatingExtractor,
    ExtractionLLMSettings,
    GlinerExtractor,
)


async def main():
    settings = ExtractionLLMSettings.from_openai_compatible_env()
    extractor = EscalatingExtractor(
        primary=GlinerExtractor(),
        escalate_to=BAMLExtractor(settings=settings),
    )
    result = await extractor.extract(chunk, GENERIC)
    print(result.entities)


asyncio.run(main())

GlinerExtractor needs the extract extra:

pip install "agentic-graphrag[extract]"

Configure the LLM provider

Set EXTRACTION_LLM_CLIENTS (or a .env file) to point at any supported provider — Anthropic, OpenAI, Azure OpenAI, AWS Bedrock, Google AI, Vertex AI, or a self-hosted OpenAI-compatible endpoint via openai-generic.

Alternatively, use ExtractionLLMSettings.from_openai_compatible_env() to read from LLM_BASE_URL, LLM_API_KEY, and LLM_MODEL_ID environment variables.

Resolve duplicate mentions

import asyncio
from agrag.ingestion.resolve import (
    ExactMatch,
    FuzzyMatch,
    InBatchCandidateSource,
    LLMVerify,
    Resolver,
)


async def main():
    # entities come from an extraction call, e.g. result.entities
    # chunks_by_id maps each chunk ID to its Chunk for LLM context
    resolver = Resolver(
        comparators=[
            ExactMatch(),
            FuzzyMatch(),
            LLMVerify(
                chunks_by_id=chunks_by_id,
                settings=ExtractionLLMSettings.from_openai_compatible_env(),
            ),
        ],
        candidate_source=InBatchCandidateSource(),
    )
    groups = await resolver.resolve(entities)
    for group in groups:
        print(group.entity_indices)


asyncio.run(main())

Each group lists the indices of entities resolution decided are the same thing. A group of one means resolution found no match for that entity.

Define your own schema

from agrag.common.data_models.graph_schema import (
    EntityType,
    GraphSchema,
    RelationType,
)

schema = GraphSchema(
    name="clinical",
    version="1",
    entities=[EntityType(label="Drug", description="A medication.")],
    relations=[],
)

Next steps

• See the API reference for every Extractor, Comparator, and schema field.