diff --git a/README.md b/README.md index 2fd782c8..fc3f70b5 100644 --- a/README.md +++ b/README.md @@ -96,20 +96,20 @@ async def main(): asyncio.run(main()) ``` -### 3. Define a schema (optional) +### 3. Define an ontology (optional) ```python -from graphrag_sdk import GraphSchema, EntityType, RelationType +from graphrag_sdk import Ontology, Entity, Relation -schema = GraphSchema( +ontology = Ontology( entities=[ - EntityType(label="Person", description="A human being"), - EntityType(label="Organization", description="A company or institution"), - EntityType(label="Location", description="A geographic location"), + Entity(label="Person", description="A human being"), + Entity(label="Organization", description="A company or institution"), + Entity(label="Location", description="A geographic location"), ], relations=[ - RelationType(label="WORKS_AT", description="Is employed by", patterns=[("Person", "Organization")]), - RelationType(label="LOCATED_IN", description="Is situated in", patterns=[("Organization", "Location")]), + Relation(label="WORKS_AT", description="Is employed by", patterns=[("Person", "Organization")]), + Relation(label="LOCATED_IN", description="Is situated in", patterns=[("Organization", "Location")]), ], ) @@ -117,7 +117,7 @@ async with GraphRAG( connection=ConnectionConfig(host="localhost", graph_name="my_graph"), llm=LiteLLM(model="openai/gpt-5.5"), embedder=LiteLLMEmbedder(model="openai/text-embedding-3-large", dimensions=256), - schema=schema, + ontology=ontology, ) as rag: ... # ingest / completion as above ``` diff --git a/docs/api-reference.mdx b/docs/api-reference.mdx index b63767e4..9dc23bda 100644 --- a/docs/api-reference.mdx +++ b/docs/api-reference.mdx @@ -11,7 +11,7 @@ Complete reference for all public classes and methods exported by `graphrag_sdk` - [Connection](#connection) - [Providers](#providers) - [Data Models](#data-models) -- [Schema](#schema) +- [Ontology](#ontology) - [Ingestion Strategies](#ingestion-strategies) - [Ingestion Pipeline](#ingestion-pipeline) - [Retrieval Strategies](#retrieval-strategies) @@ -37,7 +37,7 @@ GraphRAG( connection: FalkorDBConnection | ConnectionConfig, llm: LLMInterface, embedder: Embedder, - schema: GraphSchema | None = None, + ontology: Ontology | None = None, retrieval_strategy: RetrievalStrategy | None = None, ) ``` @@ -47,10 +47,10 @@ GraphRAG( | `connection` | `FalkorDBConnection \| ConnectionConfig` | required | Database connection or config to create one | | `llm` | `LLMInterface` | required | LLM provider | | `embedder` | `Embedder` | required | Embedding provider | -| `schema` | `GraphSchema \| None` | `None` | Schema constraints for extraction (empty = unconstrained) | +| `ontology` | `Ontology \| None` | `None` | Ontology constraints for extraction (empty = unconstrained) | | `retrieval_strategy` | `RetrievalStrategy \| None` | `None` | Default retrieval strategy (uses `MultiPathRetrieval` if None) | -**Public attributes:** `llm`, `embedder`, `schema`, `graph_store`, `vector_store` +**Public attributes:** `llm`, `embedder`, `ontology`, `graph_store`, `vector_store` ### ingest() @@ -522,46 +522,46 @@ Deterministic entity ID from normalized name and optional type. When `entity_typ --- -## Schema +## Ontology ```python -from graphrag_sdk import GraphSchema, EntityType, RelationType +from graphrag_sdk import Ontology, Entity, Relation, Attribute ``` -### EntityType +### Entity ```python -class EntityType(DataModel): +class Entity(DataModel): label: str # e.g. "Person" description: str | None = None # Helps LLM understand what to extract - properties: list[PropertyType] = [] # Optional property definitions + properties: list[Attribute] = [] # Optional property definitions ``` -### RelationType +### Relation ```python -class RelationType(DataModel): +class Relation(DataModel): label: str # e.g. "WORKS_AT" description: str | None = None patterns: list[tuple[str, str]] = [] # Allowed (source_label, target_label) pairs + properties: list[Attribute] = [] # Optional property definitions ``` -### PropertyType +### Attribute ```python -class PropertyType(DataModel): +class Attribute(DataModel): name: str type: str = "STRING" # STRING, INTEGER, FLOAT, BOOLEAN, DATE, LIST description: str | None = None - required: bool = False ``` -### GraphSchema +### Ontology ```python -class GraphSchema(DataModel): - entities: list[EntityType] = [] - relations: list[RelationType] = [] +class Ontology(DataModel): + entities: list[Entity] = [] + relations: list[Relation] = [] ``` --- @@ -593,7 +593,7 @@ class ChunkingStrategy(ABC): ```python class ExtractionStrategy(ABC): @abstractmethod - async def extract(self, chunks: TextChunks, schema: GraphSchema, ctx: Context) -> GraphData: ... + async def extract(self, chunks: TextChunks, ontology: Ontology, ctx: Context) -> GraphData: ... ``` **Built-in:** @@ -632,7 +632,7 @@ IngestionPipeline( resolver: ResolutionStrategy, graph_store: GraphStore, vector_store: VectorStore, - schema: GraphSchema | None = None, + ontology: Ontology | None = None, embedder: Embedder | None = None, ) ``` diff --git a/docs/configuration.mdx b/docs/configuration.mdx index b9224d83..d73ceb1e 100644 --- a/docs/configuration.mdx +++ b/docs/configuration.mdx @@ -269,85 +269,84 @@ Both `LiteLLMEmbedder` and `OpenRouterEmbedder` implement binary-split error rec --- -## 4. GraphSchema +## 4. Ontology -`GraphSchema` defines the structure of your knowledge graph. It constrains LLM extraction and powers the pruning step that filters non-conforming data. +`Ontology` defines the structure of your knowledge graph. It constrains LLM extraction and powers the pruning step that filters non-conforming data. ### Components -**EntityType** -- defines a node type: +**Entity** -- defines a node type: | Field | Type | Default | Description | |---------------|---------------------|---------------|----------------------------------------| | `label` | `str` | -- | The node label (e.g. `"Person"`). | | `description` | `str \| None` | `None` | Human-readable description. | -| `properties` | `list[PropertyType]` | `[]` | Expected properties on this node type. | +| `properties` | `list[Attribute]` | `[]` | Expected properties on this node type. | -**RelationType** -- defines a relationship type: +**Relation** -- defines a relationship type: | Field | Type | Default | Description | |---------------|---------------------|---------------|------------------------------------------| | `label` | `str` | -- | The relationship type (e.g. `"KNOWS"`). | | `description` | `str \| None` | `None` | Human-readable description. | -| `properties` | `list[PropertyType]` | `[]` | Expected properties on this relationship.| +| `properties` | `list[Attribute]` | `[]` | Expected properties on this relationship.| -**PropertyType** -- defines a property on a node or relationship: +**Attribute** -- defines a property on a node or relationship: | Field | Type | Default | Description | |---------------|-----------------|-------------|--------------------------------------------------------------| | `name` | `str` | -- | Property name. | | `type` | `str` | `"STRING"` | Type hint: `STRING`, `INTEGER`, `FLOAT`, `BOOLEAN`, `DATE`, `LIST`. | | `description` | `str \| None` | `None` | Human-readable description. | -| `required` | `bool` | `False` | Whether the property is required. | -### Example Schema Definition +### Example Ontology Definition ```python from graphrag_sdk.core.models import ( - EntityType, RelationType, PropertyType, GraphSchema, + Entity, Relation, Attribute, Ontology, ) -schema = GraphSchema( +ontology = Ontology( entities=[ - EntityType( + Entity( label="Person", description="A character or real person", properties=[ - PropertyType(name="name", type="STRING", required=True), - PropertyType(name="age", type="INTEGER"), - PropertyType(name="occupation", type="STRING"), + Attribute(name="name", type="STRING", required=True), + Attribute(name="age", type="INTEGER"), + Attribute(name="occupation", type="STRING"), ], ), - EntityType( + Entity( label="Location", description="A geographical place or setting", properties=[ - PropertyType(name="name", type="STRING", required=True), - PropertyType(name="country", type="STRING"), + Attribute(name="name", type="STRING", required=True), + Attribute(name="country", type="STRING"), ], ), - EntityType( + Entity( label="Organization", description="A company, institution, or group", ), ], relations=[ - RelationType( + Relation( label="LIVES_IN", description="Person resides at location", patterns=[("Person", "Location")], ), - RelationType( + Relation( label="WORKS_FOR", description="Person is employed by organization", patterns=[("Person", "Organization")], ), - RelationType( + Relation( label="LOCATED_IN", description="Organization is located at a place", patterns=[("Organization", "Location")], ), - RelationType( + Relation( label="KNOWS", description="Two people know each other", patterns=[("Person", "Person")], @@ -356,12 +355,12 @@ schema = GraphSchema( ) ``` -Each `RelationType.patterns` entry is a `(source_label, target_label)` tuple. +Each `Relation.patterns` entry is a `(source_label, target_label)` tuple. An empty `patterns` list means the relation is allowed between any entity types. -### Open Schema Mode +### Open Ontology Mode -If no entity types or relation types are defined (empty `GraphSchema()`), the extraction operates in open-schema mode and the pruning step is skipped. This lets the LLM extract any entities and relationships it finds. +If no entity types or relation types are defined (empty `Ontology()`), the extraction operates in open-ontology mode and the pruning step is skipped. This lets the LLM extract any entities and relationships it finds. --- @@ -395,7 +394,7 @@ Larger chunks provide more context per extraction call but increase LLM token us | `llm` | `LLMInterface` | required | LLM provider for step 2 (verify + relationship extraction). | | `entity_extractor` | `EntityExtractor \| None` | `None` (`GLiNERExtractor()`) | Pluggable NER backend for step 1. | | `coref_resolver` | `CorefResolver \| None` | `None` | Optional coreference resolution (e.g. `FastCorefResolver()`). | -| `entity_types` | `list[str] \| None` | `None` (11 default types) | Custom entity types. Overridden by `schema.entities` if set. | +| `entity_types` | `list[str] \| None` | `None` (11 default types) | Custom entity types. Overridden by `ontology.entities` if set. | | `max_concurrency` | `int \| None` | `None` (uses LLM default) | Maximum parallel LLM calls during step 2. | **Built-in entity extractors:** @@ -437,18 +436,18 @@ extractor = GraphExtraction( entity_types=["Gene", "Protein", "Disease", "Drug", "Pathway"], ) -# Or define them in the schema (takes priority) -from graphrag_sdk import GraphSchema, EntityType +# Or define them in the ontology (takes priority) +from graphrag_sdk import Ontology, Entity -schema = GraphSchema(entities=[ - EntityType(label="Gene", description="A gene or genetic locus"), - EntityType(label="Protein", description="A protein or enzyme"), - EntityType(label="Disease", description="A disease or condition"), +ontology = Ontology(entities=[ + Entity(label="Gene", description="A gene or genetic locus"), + Entity(label="Protein", description="A protein or enzyme"), + Entity(label="Disease", description="A disease or condition"), ]) -rag = GraphRAG(connection=conn, llm=llm, embedder=embedder, schema=schema) +rag = GraphRAG(connection=conn, llm=llm, embedder=embedder, ontology=ontology) ``` -Priority: `schema.entities` > `entity_types` param > defaults (Person, Organization, Technology, Product, Location, Date, Event, Concept, Law, Dataset, Method). +Priority: `ontology.entities` > `entity_types` param > defaults (Person, Organization, Technology, Product, Location, Date, Event, Concept, Law, Dataset, Method). ### LLM Concurrency diff --git a/docs/extraction.mdx b/docs/extraction.mdx index 687d7487..f450f773 100644 --- a/docs/extraction.mdx +++ b/docs/extraction.mdx @@ -204,15 +204,15 @@ Entities that don't match any type (or fall below the confidence threshold) are There are three ways to define entity types, listed by priority: -**1. GraphSchema entities (highest priority):** +**1. Ontology entities (highest priority):** ```python -from graphrag_sdk import GraphSchema, EntityType +from graphrag_sdk import Ontology, Entity -schema = GraphSchema(entities=[ - EntityType(label="Gene", description="A gene or genetic locus"), - EntityType(label="Disease", description="A disease or condition"), +ontology = Ontology(entities=[ + Entity(label="Gene", description="A gene or genetic locus"), + Entity(label="Disease", description="A disease or condition"), ]) -rag = GraphRAG(connection=conn, llm=llm, embedder=embedder, schema=schema) +rag = GraphRAG(connection=conn, llm=llm, embedder=embedder, ontology=ontology) # Extraction uses: ["Gene", "Disease"] ``` diff --git a/docs/getting-started.mdx b/docs/getting-started.mdx index 5ce7f1a1..e61f7034 100644 --- a/docs/getting-started.mdx +++ b/docs/getting-started.mdx @@ -64,23 +64,23 @@ If you use a `.env` file, load it yourself before importing the SDK (e.g., via ` --- -## 5. Define a Schema +## 5. Define an Ontology -A `GraphSchema` tells the extraction pipeline which entity and relationship types to look for in your documents. +An `Ontology` tells the extraction pipeline which entity and relationship types to look for in your documents. ```python -from graphrag_sdk import GraphSchema, EntityType, RelationType +from graphrag_sdk import Ontology, Entity, Relation -schema = GraphSchema( +ontology = Ontology( entities=[ - EntityType(label="Person", description="A human being"), - EntityType(label="Organization", description="A company or institution"), - EntityType(label="Place", description="A geographic location"), + Entity(label="Person", description="A human being"), + Entity(label="Organization", description="A company or institution"), + Entity(label="Place", description="A geographic location"), ], relations=[ - RelationType(label="WORKS_AT", description="Employment relationship"), - RelationType(label="LOCATED_IN", description="Geographic location"), - RelationType(label="KNOWS", description="Personal acquaintance"), + Relation(label="WORKS_AT", description="Employment relationship"), + Relation(label="LOCATED_IN", description="Geographic location"), + Relation(label="KNOWS", description="Personal acquaintance"), ], ) ``` @@ -91,7 +91,7 @@ You can add as many entity and relationship types as your domain requires. Descr ## 6. Initialize GraphRAG -Create a `GraphRAG` instance by providing a connection, LLM, embedder, and schema: +Create a `GraphRAG` instance by providing a connection, LLM, embedder, and ontology: ```python from graphrag_sdk import GraphRAG, ConnectionConfig, LiteLLM, LiteLLMEmbedder @@ -100,7 +100,7 @@ rag = GraphRAG( connection=ConnectionConfig(host="localhost", graph_name="my_graph"), llm=LiteLLM(model="azure/gpt-4.1"), embedder=LiteLLMEmbedder(model="azure/text-embedding-3-large", dimensions=256), - schema=schema, + ontology=ontology, embedding_dimension=256, # must match your embedding model's output dimension ) ``` diff --git a/docs/graph-schema.mdx b/docs/graph-schema.mdx index 33b14a4a..52724365 100644 --- a/docs/graph-schema.mdx +++ b/docs/graph-schema.mdx @@ -209,20 +209,20 @@ This is the **Zero-Loss Data** principle: every piece of source material is trac ## Defining Your Own Schema -A `GraphSchema` tells the extraction pipeline which entity and relationship types to look for, and the pruning step uses it to filter non-conforming data. +An `Ontology` tells the extraction pipeline which entity and relationship types to look for, and the pruning step uses it to filter non-conforming data. ### Basic Schema ```python -from graphrag_sdk import GraphSchema, EntityType, RelationType +from graphrag_sdk import Ontology, Entity, Relation -schema = GraphSchema( +ontology = Ontology( entities=[ - EntityType(label="Person", description="A human being"), - EntityType(label="Organization", description="A company or institution"), + Entity(label="Person", description="A human being"), + Entity(label="Organization", description="A company or institution"), ], relations=[ - RelationType(label="WORKS_AT", description="Employment relationship"), + Relation(label="WORKS_AT", description="Employment relationship"), ], ) ``` @@ -230,32 +230,32 @@ schema = GraphSchema( ### Schema with Patterns Patterns define which source-target pairs are valid for each relationship type. -They are specified directly on `RelationType`: +They are specified directly on `Relation`: ```python -schema = GraphSchema( +ontology = Ontology( entities=[ - EntityType(label="Person"), - EntityType(label="Organization"), - EntityType(label="Location"), + Entity(label="Person"), + Entity(label="Organization"), + Entity(label="Location"), ], relations=[ - RelationType(label="WORKS_AT", patterns=[("Person", "Organization")]), - RelationType(label="LOCATED_IN", patterns=[("Organization", "Location")]), + Relation(label="WORKS_AT", patterns=[("Person", "Organization")]), + Relation(label="LOCATED_IN", patterns=[("Organization", "Location")]), ], ) ``` A relationship with an empty `patterns` list is allowed between any entity types. -### Open Schema Mode +### Open Ontology Mode -If you create an empty schema (`GraphSchema()`), the pipeline operates in **open schema mode**: +If you create an empty ontology (`Ontology()`), the pipeline operates in **open ontology mode**: - The LLM extracts any entities and relationships it finds - The pruning step is skipped entirely - The 11 default entity types are used for NER -This is good for exploration. For production, a defined schema produces cleaner, more consistent graphs. +This is good for exploration. For production, a defined ontology produces cleaner, more consistent graphs. --- @@ -298,7 +298,7 @@ result = await rag.graph_store.query_raw( | File | What it contains | |------|-----------------| -| [`core/models.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/src/graphrag_sdk/core/models.py) | GraphSchema, EntityType, RelationType, PropertyType | +| [`core/models.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/src/graphrag_sdk/core/models.py) | Ontology, Entity, Relation, Attribute | | [`storage/graph_store.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py) | Node/relationship upserts, label hints, statistics | | [`storage/vector_store.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/src/graphrag_sdk/storage/vector_store.py) | Index creation, vector search, fulltext search | | [`ingestion/pipeline.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/src/graphrag_sdk/ingestion/pipeline.py) | Lexical graph construction, pruning logic | diff --git a/docs/ingestion.mdx b/docs/ingestion.mdx index 005845ca..25b24360 100644 --- a/docs/ingestion.mdx +++ b/docs/ingestion.mdx @@ -28,8 +28,8 @@ This document explains what each step does, why it exists, and how to tune it. 4b. Quality Filter (remove bad nodes/dangling edges) | - 5. Prune Against Schema - (keep only schema-conforming data) + 5. Prune Against Ontology + (keep only ontology-conforming data) | 6. Resolve Duplicates (merge same-entity mentions) @@ -138,17 +138,17 @@ For a detailed explanation of the extraction process, see [Extraction](/extracti --- -### Step 5 — Prune Against Schema +### Step 5 — Prune Against Ontology -**What it does:** Filters extracted data to only keep entities and relationships that match your schema definition. +**What it does:** Filters extracted data to only keep entities and relationships that match your ontology definition. **How it works:** -- If your schema defines entity types (e.g., Person, Organization, Location), only entities with those labels pass through -- If your schema defines relationship types, only those relationship types pass through +- If your ontology defines entity types (e.g., Person, Organization, Location), only entities with those labels pass through +- If your ontology defines relationship types, only those relationship types pass through - Relationships whose endpoints were pruned are also removed - **Special cases:** `"Unknown"` entities (low-confidence NER) and `"RELATES"` edges (the unified relationship type) always pass through -**Open schema mode:** If you define no entity or relationship types (empty `GraphSchema()`), this step is skipped entirely — everything passes through. +**Open ontology mode:** If you define no entity or relationship types (empty `Ontology()`), this step is skipped entirely — everything passes through. **Code:** `IngestionPipeline._prune()` in [`ingestion/pipeline.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/src/graphrag_sdk/ingestion/pipeline.py) diff --git a/docs/strategies.mdx b/docs/strategies.mdx index 18eeb3e0..5ebca192 100644 --- a/docs/strategies.mdx +++ b/docs/strategies.mdx @@ -217,7 +217,7 @@ class ExtractionStrategy(ABC): async def extract( self, chunks: TextChunks, - schema: GraphSchema, + ontology: Ontology, ctx: Context, ) -> GraphData: """Extract graph data from text chunks.""" @@ -264,24 +264,24 @@ extractor = GraphExtraction( ) ``` -**3. Use `GraphSchema` entities** (schema types override both defaults and `entity_types`): +**3. Use `Ontology` entities** (ontology entity types override both defaults and `entity_types`): ```python -from graphrag_sdk import GraphRAG, GraphSchema, EntityType +from graphrag_sdk import GraphRAG, Ontology, Entity -schema = GraphSchema(entities=[ - EntityType(label="Vehicle", description="Cars, trucks, etc."), - EntityType(label="Road", description="Streets, highways, etc."), - EntityType(label="Location", description="Cities, countries, etc."), +ontology = Ontology(entities=[ + Entity(label="Vehicle", description="Cars, trucks, etc."), + Entity(label="Road", description="Streets, highways, etc."), + Entity(label="Location", description="Cities, countries, etc."), ]) -# Schema entity types are automatically used for extraction -rag = GraphRAG(connection=conn, llm=llm, embedder=embedder, schema=schema) +# Ontology entity types are automatically used for extraction +rag = GraphRAG(connection=conn, llm=llm, embedder=embedder, ontology=ontology) await rag.ingest("traffic_report.txt") # Extraction uses: ["Vehicle", "Road", "Location"] ``` -The priority order is: `schema.entities` > `entity_types` parameter > defaults. +The priority order is: `ontology.entities` > `entity_types` parameter > defaults. **Default entity types:** Person, Organization, Technology, Product, Location, Date, Event, Concept, Law, Dataset, Method. @@ -368,7 +368,7 @@ Replace the entire 2-step pipeline by subclassing `ExtractionStrategy`: ```python class MyExtraction(ExtractionStrategy): - async def extract(self, chunks, schema, ctx): + async def extract(self, chunks, ontology, ctx): nodes, rels = [], [] for chunk in chunks.chunks: # Your extraction logic diff --git a/graphrag_sdk/examples/02_pdf_with_schema.py b/graphrag_sdk/examples/02_pdf_with_schema.py index 78098956..537354b1 100644 --- a/graphrag_sdk/examples/02_pdf_with_schema.py +++ b/graphrag_sdk/examples/02_pdf_with_schema.py @@ -18,48 +18,48 @@ from graphrag_sdk import ( ConnectionConfig, - EntityType, + Entity, GraphRAG, - GraphSchema, LiteLLM, LiteLLMEmbedder, - RelationType, + Ontology, + Relation, ) from graphrag_sdk.ingestion.chunking_strategies.fixed_size import FixedSizeChunking -def create_schema() -> GraphSchema: +def create_ontology() -> Ontology: """Define what entities and relationships the LLM should extract.""" - return GraphSchema( + return Ontology( entities=[ - EntityType(label="Person", description="A human being or character"), - EntityType(label="Organization", description="A company, institution, or group"), - EntityType(label="Place", description="A geographic location or setting"), - EntityType(label="Event", description="A significant occurrence or happening"), - EntityType(label="Concept", description="An abstract idea or theme"), + Entity(label="Person", description="A human being or character"), + Entity(label="Organization", description="A company, institution, or group"), + Entity(label="Place", description="A geographic location or setting"), + Entity(label="Event", description="A significant occurrence or happening"), + Entity(label="Concept", description="An abstract idea or theme"), ], relations=[ - RelationType( + Relation( label="WORKS_AT", description="Is employed by an organization", patterns=[("Person", "Organization")], ), - RelationType( + Relation( label="LOCATED_IN", description="Is physically located in a place", patterns=[("Person", "Place"), ("Organization", "Place")], ), - RelationType( + Relation( label="RELATED_TO", description="Has a general relationship with", patterns=[("Person", "Person")], ), - RelationType( + Relation( label="PART_OF", description="Is a member or component of", patterns=[("Person", "Organization")], ), - RelationType( + Relation( label="PARTICIPATED_IN", description="Took part in an event", patterns=[("Person", "Event")], @@ -94,7 +94,7 @@ async def main(): connection=ConnectionConfig(host="localhost", graph_name="pdf_demo"), llm=llm, embedder=embedder, - schema=create_schema(), + ontology=create_ontology(), ) # Ingest PDF with larger chunks for better context diff --git a/graphrag_sdk/examples/03_custom_strategies.py b/graphrag_sdk/examples/03_custom_strategies.py index 274b2261..b35d94e9 100644 --- a/graphrag_sdk/examples/03_custom_strategies.py +++ b/graphrag_sdk/examples/03_custom_strategies.py @@ -22,12 +22,12 @@ from graphrag_sdk import ( ConnectionConfig, - EntityType, + Entity, GraphRAG, - GraphSchema, LiteLLM, LiteLLMEmbedder, - RelationType, + Ontology, + Relation, ) from graphrag_sdk.core.context import Context from graphrag_sdk.ingestion.chunking_strategies.fixed_size import FixedSizeChunking @@ -53,41 +53,41 @@ ), ] -SCHEMA = GraphSchema( +ONTOLOGY = Ontology( entities=[ - EntityType(label="Person", description="A historical figure or scientist"), - EntityType(label="Place", description="A city, country, or geographic location"), - EntityType(label="Organization", description="A university, institution, or award body"), - EntityType(label="Event", description="A significant event, award, or discovery"), - EntityType(label="Concept", description="A scientific field or abstract idea"), + Entity(label="Person", description="A historical figure or scientist"), + Entity(label="Place", description="A city, country, or geographic location"), + Entity(label="Organization", description="A university, institution, or award body"), + Entity(label="Event", description="A significant event, award, or discovery"), + Entity(label="Concept", description="A scientific field or abstract idea"), ], relations=[ - RelationType( + Relation( label="LOCATED_IN", description="Is located in a place", patterns=[("Person", "Place"), ("Organization", "Place")], ), - RelationType( + Relation( label="WORKS_AT", description="Works at an institution", patterns=[("Person", "Organization")], ), - RelationType( + Relation( label="MARRIED_TO", description="Is married to", patterns=[("Person", "Person")], ), - RelationType( + Relation( label="RELATED_TO", description="Has a relationship with", patterns=[("Person", "Person")], ), - RelationType( + Relation( label="AWARDED", description="Received an award or prize", patterns=[("Person", "Event")], ), - RelationType( + Relation( label="RESEARCHED", description="Conducted research on a topic", patterns=[("Person", "Concept")], @@ -115,7 +115,7 @@ async def main(): connection=ConnectionConfig(host="localhost", graph_name="strategies_demo"), llm=llm, embedder=embedder, - schema=SCHEMA, + ontology=ONTOLOGY, ) # Clear previous data diff --git a/graphrag_sdk/examples/06_markdown_document_aware.py b/graphrag_sdk/examples/06_markdown_document_aware.py index 92508532..d2eaebac 100644 --- a/graphrag_sdk/examples/06_markdown_document_aware.py +++ b/graphrag_sdk/examples/06_markdown_document_aware.py @@ -34,44 +34,44 @@ from graphrag_sdk import ( ConnectionConfig, - EntityType, + Entity, GraphRAG, - GraphSchema, LiteLLM, LiteLLMEmbedder, - RelationType, + Ontology, + Relation, ) from graphrag_sdk.ingestion.chunking_strategies.structural_chunking import StructuralChunking from graphrag_sdk.ingestion.loaders.markdown_loader import MarkdownLoader -def create_schema() -> GraphSchema: - """Generic schema suitable for technical documentation.""" - return GraphSchema( +def create_ontology() -> Ontology: + """Generic ontology suitable for technical documentation.""" + return Ontology( entities=[ - EntityType(label="Concept", description="A technical concept, feature, or abstraction"), - EntityType(label="Component", description="A software module, library, or service"), - EntityType(label="Technology", description="A programming language, framework, or tool"), - EntityType(label="Person", description="An author, contributor, or mentioned individual"), - EntityType(label="Organization", description="A company, team, or standards body"), + Entity(label="Concept", description="A technical concept, feature, or abstraction"), + Entity(label="Component", description="A software module, library, or service"), + Entity(label="Technology", description="A programming language, framework, or tool"), + Entity(label="Person", description="An author, contributor, or mentioned individual"), + Entity(label="Organization", description="A company, team, or standards body"), ], relations=[ - RelationType( + Relation( label="DEPENDS_ON", description="Requires another component or technology", patterns=[("Component", "Component"), ("Component", "Technology")], ), - RelationType( + Relation( label="IMPLEMENTS", description="Provides a concrete implementation of a concept", patterns=[("Component", "Concept")], ), - RelationType( + Relation( label="CREATED_BY", description="Was authored or maintained by a person or organization", patterns=[("Component", "Person"), ("Component", "Organization")], ), - RelationType( + Relation( label="RELATED_TO", description="Has a general relationship with", patterns=[ @@ -172,7 +172,7 @@ async def main(): connection=ConnectionConfig(host="localhost", graph_name="markdown_demo"), llm=llm, embedder=embedder, - schema=create_schema(), + ontology=create_ontology(), ) # delete_all() must come BEFORE ingest().