diff --git a/docs-site/astro.config.mjs b/docs-site/astro.config.mjs index d21f0bd6..ba9ce685 100644 --- a/docs-site/astro.config.mjs +++ b/docs-site/astro.config.mjs @@ -46,6 +46,7 @@ export default defineConfig({ { label: 'Metadata Filtering', slug: 'lexical-graph/metadata-filtering' }, { label: 'Reader Providers', slug: 'lexical-graph/readers' }, { label: 'External Properties', slug: 'lexical-graph/external-properties' }, + { label: 'Ontology-Guided Extraction', slug: 'lexical-graph/ontology-guided-extraction' }, ], }, { diff --git a/docs-site/src/content/docs/lexical-graph/graph-model.mdx b/docs-site/src/content/docs/lexical-graph/graph-model.mdx index 7f4e8fc5..dbc36952 100644 --- a/docs-site/src/content/docs/lexical-graph/graph-model.mdx +++ b/docs-site/src/content/docs/lexical-graph/graph-model.mdx @@ -9,6 +9,7 @@ title: Graph Model - [Units of context](#units-of-context) - [Lineage tier](#lineage-tier) - [Entity-Relationship tier](#entity-relationship-tier) + - [Typed properties on entity nodes](#typed-properties-on-entity-nodes) - [Summarisation tier](#summarisation-tier) - [Facts](#facts) - [Statements](#statements) @@ -57,6 +58,50 @@ Extraction uses a lightly guided strategy whereby the extraction process is seed Relationship values are currently unguided (though relatively concise). +#### Typed properties on entity nodes + +Every value extracted from a source is a string. An `__Entity__` node therefore carries a +string `value` (plus a `search_str` for exact-match lookup) and a string `class`, and +nothing on it is numerically comparable: + +``` +(:__Entity__ {entityId: '…', value: 'Halcyon Motors', class: 'Company', + search_str: 'halcyon motors'}) +``` + +[Ontology-guided extraction](/graphrag-toolkit/lexical-graph/ontology-guided-extraction/) +can add *typed* properties to this tier. When an extracted attribute's predicate resolves +to a datatype property the ontology declares, the value is coerced to the declared XSD type +and written as a native graph property. The write is **additive** – `value`, `search_str` +and `class` are untouched – and where it lands depends on the configured placement: + + - **`typed_properties='subject'`** adds the value to the *subject* entity, keyed by the + ontology property's local name. This is what makes an attribute range-queryable from + the entity that has it: + + ``` + (:__Entity__ {entityId: '…', value: 'Halcyon Motors', class: 'Company', + search_str: 'halcyon motors', foundedYear: 1971}) + ``` + + - **`typed_properties='complement'`** adds `typed_value` and `datatype` to the + *complement* entity – the local-context node created for the value string itself, which + exists only when `include_local_entities` is enabled: + + ``` + (:__Entity__ {entityId: '…', value: '1971', class: '__Local_Entity__', + typed_value: 1971, + datatype: 'http://www.w3.org/2001/XMLSchema#integer'}) + ``` + + These properties describe *the value string*, not the assertion: they say that this + literal read as an integer is 1971, and not which entity was founded then, nor under + which property. The subject and predicate remain on the `__Fact__`. + +`typed_value` and `datatype` are reserved by the graph model at complement placement, as +`value`, `search_str` and `class` are at subject placement; an ontology declaring a +property with one of those names is refused at configuration time. + ### Summarisation tier This currently comprises `__Topic__`, `__Statement__` and `__Fact__` nodes. Proceeding from the bottom up: @@ -145,7 +190,7 @@ What each node type embeds: In other words, document-level metadata such as `file_name` or `size` is **not** baked into chunk embeddings by default — it is carried as graph/source metadata (see -[External properties](/lexical-graph/external-properties/) for making such metadata +[External properties](/graphrag-toolkit/lexical-graph/external-properties/) for making such metadata queryable). If you want a piece of metadata to influence the embedding itself, include it in the node's content rather than relying on metadata, since metadata under `source`, `chunk`, `statement`, and the internal index keys is excluded from the vector. diff --git a/docs-site/src/content/docs/lexical-graph/indexing.mdx b/docs-site/src/content/docs/lexical-graph/indexing.mdx index f8a0d6d7..6bb5386a 100644 --- a/docs-site/src/content/docs/lexical-graph/indexing.mdx +++ b/docs-site/src/content/docs/lexical-graph/indexing.mdx @@ -30,6 +30,8 @@ The list of `DEFAULT_ENTITY_CLASSIFICATIONS` used to seed the extraction process Relationship values are currently unguided (though relatively concise). +If you want a stronger guarantee than "reduces but doesn't eliminate", supply an OWL/RDFS ontology as the extraction vocabulary. The ontology's classes and properties are rendered into the extraction prompt, and a deterministic filter then resolves each extracted name against the ontology – rewriting it to the declared spelling, or, at the strictest level, dropping what doesn't conform. Declared datatype properties can additionally be coerced to their declared type and written as native, queryable graph properties. See [Ontology-Guided Extraction](/graphrag-toolkit/lexical-graph/ontology-guided-extraction/). + #### Build In the build stage, the LlamaIndex chunk nodes emitted from the extract stage are broken down further into a stream of individual source, chunk, topic, statement and fact LlamaIndex nodes. Graph construction and vector indexing handlers process these nodes to build and index the graph content. Each of these nodes has an `aws::graph::index` metadata item containing data that can be used to index the node in a vector store (though only the chunk and statement nodes are actually indexed in the current implementation). @@ -275,6 +277,7 @@ The `ExtractionConfig` object has the following parameters: | `enable_proposition_extraction` | Perform proposition extraction before extracting topics, statements, facts and entities | `True` | | `preferred_entity_classifications` | Comma-separated list of preferred entity classifications used to seed the entity extraction | `DEFAULT_ENTITY_CLASSIFICATIONS` | | `preferred_topics` | List of preferred topic names (or a callable that returns them) supplied to the LLM to seed topic extraction. Accepts the same type as `preferred_entity_classifications`. | `[]` | +| `ontology` | An OWL/RDFS ontology used as the extraction vocabulary. Accepts an `OntologyConfig`, an `Ontology`, an `rdflib.Graph`, or a path to a Turtle (`.ttl`) file – a bare path or graph is equivalent to `OntologyConfig(source)`, which defaults to `ontology_authority='align'` and `typed_properties='off'`. See [Ontology-Guided Extraction](/graphrag-toolkit/lexical-graph/ontology-guided-extraction/). | `None` | | `infer_entity_classifications` | Determines whether to pre-process documents to identify significant domain entity classifications. Supply either `True` or `False`, or an `InferClassificationsConfig` object. When `True`, an `InferClassifications` step runs as a **pre-processor** before the main extraction loop — one extra LLM round-trip per batch, not per document. | `False` | | `extract_propositions_prompt_template` | Prompt used to extract propositions from chunks. If `None`, the [default extract propositions template](https://github.com/awslabs/graphrag-toolkit/blob/main/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py#L29-L72) is used. See [Custom prompts](#custom-prompts) below. | `None` | | `extract_topics_prompt_template` | Prompt used to extract topics, statements and entities from chunks. If `None`, the [default extract topics template](https://github.com/awslabs/graphrag-toolkit/blob/main/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py#L74-L191) is used. See [Custom prompts](#custom-prompts) below. | `None` | diff --git a/docs-site/src/content/docs/lexical-graph/ontology-guided-extraction.mdx b/docs-site/src/content/docs/lexical-graph/ontology-guided-extraction.mdx new file mode 100644 index 00000000..30d88346 --- /dev/null +++ b/docs-site/src/content/docs/lexical-graph/ontology-guided-extraction.mdx @@ -0,0 +1,577 @@ +--- +title: Ontology-Guided Extraction +--- + +import { Aside } from '@astrojs/starlight/components'; + +### Topics + + - [Overview](#overview) + - [Ontology authority: how much say the ontology has](#ontology-authority-how-much-say-the-ontology-has) + - [Reading `report_violations`](#reading-report_violations) + - [Quick start](#quick-start) + - [Vocabulary format](#vocabulary-format) + - [Writing an ontology](#writing-an-ontology) + - [Authored names](#authored-names) + - [What the ontology does, and what it does not](#what-the-ontology-does-and-what-it-does-not) + - [Naming folds conventions; it does not map meaning](#naming-folds-conventions-it-does-not-map-meaning) + - [Enforcement checks shape, never meaning](#enforcement-checks-shape-never-meaning) + - [Typed attribute values](#typed-attribute-values) + - [Placements](#placements) + - [What lands on a node](#what-lands-on-a-node) + - [Coercion, and what happens when it refuses](#coercion-and-what-happens-when-it-refuses) + - [Who reads typed properties](#who-reads-typed-properties) + - [Confirming your properties are visible](#confirming-your-properties-are-visible) + - [Two consumption paths](#two-consumption-paths) + - [Operational notes](#operational-notes) + - [Annotations are frozen at extraction](#annotations-are-frozen-at-extraction) + - [Enabling an ontology is a re-index boundary](#enabling-an-ontology-is-a-re-index-boundary) + - [Subject-mode values are last-write-win](#subject-mode-values-are-last-write-win) + +### Overview + +Extraction is, by default, *lightly* guided: the LLM is seeded with a list of preferred +entity classifications and asked to prefer them, and relationship values are unguided +altogether (see [Indexing](/graphrag-toolkit/lexical-graph/indexing/#extract)). That keeps +recall high, but it can produce a graph in which the same idea is named several ways — +`Company` and `CORPORATION`, `WORKS_FOR` and `EMPLOYED BY`. + +Ontology-guided extraction lets you supply an OWL/RDFS ontology to guide the vocabulary. An +ontology declares the entity types you want, the relationships permitted between them, and +the attributes they carry — see [Writing an ontology](#writing-an-ontology). It then does two +things: + + 1. its classes and properties, with their comments, are rendered into the extraction + prompt, so the LLM is told the vocabulary rather than left to invent one; and + 2. after extraction, a deterministic filter resolves each extracted name against the + ontology and — depending on how much authority you give it — rewrites it to the + declared spelling, or drops what does not conform. + +Optionally, values of declared *datatype* properties are coerced to their declared XSD +type and written onto entity nodes as native graph properties, which allows for direct querying and comparison +such as adding a `WHERE c.foundedYear < 2000` clause to a query. + +The aim is **consistency, not correctness**: the ontology makes the graph name things one +way. It does not make the extraction true. See +[What the ontology does, and what it does not](#what-the-ontology-does-and-what-it-does-not). + +### Ontology authority: how much say the ontology has + +`ontology_authority` is a single knob: how much say the ontology has over what ends up in +the graph. Two questions decide it — what happens to a concept the ontology declares, and +what happens to one it does not: + +| Level | Concept **matches** something in the ontology | Concept is **not in** the ontology | +|---|---|---| +| `'off'` | left exactly as the model emitted it | left exactly as the model emitted it | +| `'align'` | **stored under the ontology's name** | kept, under the model's own name | +| `'strict'` | stored under the ontology's name | **excluded** | + +So `align` governs *naming*, and `strict` governs *membership*. + +**`align` is the level to start at**, and it is the default. It is the setting whose +downside is bounded: the worst case is that some names are not folded, which is where you +started. `strict` is not "align, but tidier" — on a corpus wider than the ontology it +removes a large share of the extraction, by design. That is the right behaviour when the +ontology *is* the specification of what belongs in the graph, and the wrong one when it is +a vocabulary covering part of a broader corpus. There is no way to predict how much it +would remove without running it: the level changes the prompt as well as the gates, so the +model emits a different population before any gate sees it. + +`'off'` is useful in one specific combination: with an explicit dimension override, or with +`typed_properties` alone, when you want datatype coercion without the ontology saying +anything about naming. + +Underneath, the level resolves to six independent dimensions, each of which you can also +set on its own: + +| Dimension | `'off'` | `'align'` | `'strict'` | What it does | +|---|---|---|---|---| +| `normalize_names` | `False` | `True` | `True` | Rewrite a resolved name to the ontology's authored spelling | +| `drop_type_restatements` | `False` | `True` | `True` | Drop a fact that carries ontology language rather than a statement about the world | +| `enforce_entity_types` | `False` | `False` | `True` | Drop a fact whose entity classification does not resolve | +| `enforce_relationship_types` | `False` | `False` | `True` | Drop a fact whose predicate does not resolve | +| `enforce_domain_range` | `False` | `False` | `True` | Drop a fact violating a declared `rdfs:domain` / `rdfs:range` | +| `enforce_datatypes` | `False` | `False` | `True` | Drop a fact whose literal does not parse as its declared XSD datatype | + +If every dimension resolves to `False` and `typed_properties` is `'off'`, the filter is not +added to the extraction pipeline at all. + +`drop_type_restatements` is the one row that does not follow "`align` names, `strict` +excludes": it drops, and it is on at `align`. What it drops is narrow enough to be worth +stating precisely, because "`align` keeps what the ontology does not declare" is otherwise +an exact promise: + +- a predicate that is ontology *language* — `rdf:type`, `rdfs:subClassOf`, `owl:sameAs`. + The prefix is required, so an attribute named `DOMAIN` or `RANGE` is untouched; +- a type-asserting predicate whose value repeats the subject's own classification — + `Meridian Freight [Company] |TYPE| Company`. This is a duplicate of the entity's class, + which the graph already records, so removing it loses nothing. + +Both conditions are skipped for a predicate your ontology declares: if you declare a +`classification` property, it means what you said it means. And the second condition needs +the value to match — `Riverside Rovers [Sports Team] |CLASSIFICATION| football club` is kept, +because it says something the classification does not. + +This exists because `vocabulary_format='turtle'` (below) makes models restate class +membership as a fact, and `align` has no other gate that would catch it. Set it to `False` +if you want those facts in your graph. + +#### Reading `report_violations` + +`report_violations` logs what the filter did, at INFO: per-dimension counts of names +rewritten and facts dropped. It is a record of the filter's actions, and not a measure of +how far your corpus diverges from your ontology. At `align` the four `enforce_*` counts are +structurally zero, because those dimensions are off. At `strict` they are low for a +different reason: the prompt has already told the model that unlisted concepts are +discarded, so what reaches a gate is only the remainder that ignored the instruction. + +Counts accumulate per node batch, and each line states the number of nodes it covers, so +lines add up and no single line is the total for a run. + + + +### Quick start + +Pass an ontology to `ExtractionConfig`. Any source `Ontology.load` accepts will do, +including a bare path: + +```python +from graphrag_toolkit.lexical_graph import ( + ExtractionConfig, IndexingConfig, LexicalGraphIndex +) + +graph_index = LexicalGraphIndex( + graph_store, + vector_store, + indexing_config=IndexingConfig( + extraction=ExtractionConfig(ontology='ontologies/news.ttl') + ) +) +``` + +That gives you the defaults: `ontology_authority='align'`, `typed_properties='off'`. For anything +beyond the defaults, pass an `OntologyConfig`: + +```python +from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + +ontology = OntologyConfig( + 'ontologies/news.ttl', + ontology_authority='align', # 'off' | 'align' | 'strict' + typed_properties='subject', # 'off' | 'subject' | 'complement' | 'both' + report_violations=True, # log what the filter rewrote and dropped, at INFO +) + +graph_index = LexicalGraphIndex( + graph_store, + vector_store, + indexing_config=IndexingConfig(extraction=ExtractionConfig(ontology=ontology)) +) +``` + +`OntologyConfig` has the following parameters: + +| Parameter | Description | Default Value | +| ------------- | ------------- | ------------- | +| `ontology` | The ontology. Accepts an `Ontology`, an `rdflib.Graph`, or a path (`str` or `Path`) to a `.ttl` file. | *required* | +| `ontology_authority` | How much authority the ontology has: `'off'`, `'align'` or `'strict'`. See [Ontology authority](#ontology-authority-how-much-say-the-ontology-has). | `'align'` | +| `normalize_names` | Override for the level. `None` follows `ontology_authority`. | `None` | +| `drop_type_restatements` | Override for the level. On at `align` and `strict`; set `False` to keep facts that restate an entity's own class. | `None` | +| `enforce_entity_types` | Override for the level. | `None` | +| `enforce_relationship_types` | Override for the level. | `None` | +| `enforce_domain_range` | Override for the level. | `None` | +| `enforce_datatypes` | Override for the level. | `None` | +| `report_violations` | Log per-dimension counts of what this filter call rewrote and dropped, at INFO, per node batch. A record of the filter's actions, not of how far the corpus diverges from the ontology — see [Reading `report_violations`](#reading-report_violations). | `False` | +| `typed_properties` | Where a coerced attribute value is written: `'off'`, `'subject'`, `'complement'` or `'both'`. See [Typed attribute values](#typed-attribute-values). | `'off'` | +| `vocabulary_format` | How the vocabulary is written into the prompt: `'prose'` renders three labelled sections, `'turtle'` shows the ontology's own source. See [Vocabulary format](#vocabulary-format). | `'prose'` | + +An overriding dimension is honoured whatever the level says, so +`OntologyConfig('news.ttl', ontology_authority='off', enforce_datatypes=True)` — say nothing in +the prompt, but still coerce declared datatypes — is expressible and is not treated as a +contradiction. + +The ontology's class names seed the preferred entity classifications, so +`infer_entity_classifications` composes with an ontology: inference adds the domain terms +the ontology does not declare, which is exactly what `align` permits. The one combination +that is refused is `InferClassificationsConfig(replace_default_classifications=True)` +alongside an ontology, which asks for the ontology's own vocabulary to be discarded before +the model sees it; that raises `ValueError` rather than picking a half to honour. + +### Vocabulary format + +`vocabulary_format` chooses how the vocabulary is written down for the model, which is a +separate question from how much authority it has. + +`'prose'`, the default, generates three labelled sections from the ontology — one for entity +types, one for relationships, one for attributes. Each section states the output channel it +belongs to, the class hierarchy is shown as indentation, and XSD ranges are written in plain +language (`integer`, `true/false`, `date`) rather than as IRIs. + +`'turtle'` shows the ontology's own source instead, so `rdfs:subClassOf`, `rdfs:domain` and +`rdfs:range` are stated formally rather than paraphrased. OWL and RDFS are notations the +models have read a great deal of, and an entailment that indentation only implies is written +out. What it gives up is the separation the prose sections make explicit: Turtle interleaves +object and datatype properties in whatever order you authored them, so nothing in the block +says which output channel a term belongs to except a sentence in the header. It also makes +the prompt longer, on every chunk — by how much depends entirely on your ontology. + +One behaviour worth knowing if you try `'turtle'`: showing the model type triples makes it +more likely to emit an entity's own class back as an attribute — `Meridian Freight [Company] +|TYPE| Company`. `drop_type_restatements` is on from `align` upwards and removes those, so the +combination is usable, but it is the reason that dimension exists. + +The setting applies to the extraction call that produces entities, relationships and +attributes, which is the one given the full vocabulary. Proposition extraction, which runs +first when it is enabled (see +[Indexing](/graphrag-toolkit/lexical-graph/indexing/#extract)), is given the class names alone +whichever format you choose. + +### Writing an ontology + +A worked example, commented throughout with the reason for each choice, ships at +[`examples/lexical-graph/ontologies/news.ttl`](https://github.com/awslabs/graphrag-toolkit/blob/main/examples/lexical-graph/ontologies/news.ttl). +Start from it rather than from an empty file. + +Turtle (`.ttl`) is the only file format read from a path. Any other RDF serialization is +supported by parsing it with `rdflib` and passing the resulting graph. + +The toolkit reads four kinds of declaration: + + - **`owl:Class`** — an entity classification. `rdfs:subClassOf` is respected, so a range + declared as `:Company` admits a `:SportsTeam` that is declared a subclass of it. + - **`owl:ObjectProperty`** — a relationship between two entities, with optional + `rdfs:domain` and `rdfs:range`. + - **`owl:DatatypeProperty`** — an attribute of one entity, whose `rdfs:range` is an XSD + datatype. These are what [typed values](#typed-attribute-values) are written from. + - **`skos:altLabel`** — an alternative phrasing that *maps onto* a term. Alt labels are + input only: they are never stored, they are how a model's wording finds its way to the + term you want. + +`rdfs:comment` is not decoration. It is rendered into the extraction prompt, and it is the +only place you can tell the model what a term means. A term with no comment is a bare name +the model has to guess at. + +#### Authored names + +Matched concepts are stored under the ontology's own names, **spelled exactly as you +authored them**: a term's `rdfs:label` if it has one, and its IRI local name otherwise. +Label `:SportsTeam` as `"Sports Team"` and entities of that class are stored with +`class = 'Sports Team'`; drop the label and they are stored as `SportsTeam`. Pick the +spelling you want to read back out of the graph and put it in the label. + +How the ontology is rendered into the extraction prompt is a separate matter, and an +**internal detail** — the section layout, the ordering, and the wording may change between +releases. It does not affect what is stored. What is stable is the contract above: labels +if present, local names otherwise, and comments reach the model. + +One exception matters if you use `typed_properties`: the *property key* written to a node +comes from the IRI local name, never from the label. Labelling `:revenue` as +`"Annual Revenue"` gives a fact predicate reading `Annual Revenue` and a node property still +keyed `revenue`. Label each datatype property to match its local name unless you want that +split deliberately; classes and object properties are unaffected, because nothing keys a +graph property from them. + +### What the ontology does, and what it does not + +The work is split between the LLM and the deterministic filter, and the split is sharper +than it first looks: + + - **The LLM decides which concept a phrase means.** That is the whole of the semantic + work. If a passage says an executive "was brought on board by" a company, only the + model can decide that this is employment. + - **The filter does the rest, and only the rest**: folding naming conventions on terms + that *already resolve*, coercing literals to declared datatypes, and — at `strict` — + checking conformance. + +Set expectations accordingly. The ontology raises the ceiling on **consistency**, not on +**correctness**. Extraction quality still rests on the model and the prompt. + +Canonicalization also has a bounded reach: it rewrites entity classifications, fact +predicates and typed values, and never the topic or statement text, which is stored as the +model phrased it. + +#### Naming folds conventions; it does not map meaning + +`normalize_names` resolves an extracted name by folding it to a convention-free key — +splitting camelCase, turning underscores into spaces, lowercasing, collapsing whitespace — +and looking that key up, `skos:altLabel`s included. It does not reason about meaning: + +``` +LLM emits "WORKS FOR" -> folds to "works for" -> :worksFor -> stored as worksFor ✓ +LLM emits "Works_For" -> folds to "works for" -> :worksFor -> stored as worksFor ✓ +LLM emits "EMPLOYED BY" -> matches skos:altLabel -> :worksFor -> stored as worksFor ✓ +LLM emits "HIRED BY" -> folds to "hired by" -> no match -> left as "HIRED BY" ✗ +``` + +That last line is the important one. **An unresolved predicate is never mapped.** Nothing +downstream rescues it: `normalize_names` will not rename `HIRED BY` to `worksFor`, because +as far as resolution is concerned those are unrelated strings. Under `align` the fact +passes through under its own predicate; under `strict` it is discarded. Either way the +ontology term is not applied. + +This is why the prompt block is the load-bearing part of the feature, and why a spread of +`HIRED BY` / `JOINED` / `worksFor` in your graph is a prompt-wording problem rather than +something a resolver could fix. `skos:altLabel` is the one lever for widening the +deterministic reach, and it is best spent on in-house jargon a model could not reasonably +guess — not on ordinary paraphrase, which is the model's job. + +#### Enforcement checks shape, never meaning + +The converse limitation matters just as much: **the filter cannot tell a wrong mapping from +a right one.** If the model decides "acquired a stake in" is `worksFor`, the filter +resolves it, rewrites it to the canonical name, annotates it, and — if the subject and +object classes happen to satisfy `rdfs:domain` and `rdfs:range` — passes it through +`strict` enforcement as a conforming fact. + +So a confident semantic error survives, and arrives looking well-typed. Do not read a +`strict` build as a validated one; read it as a build in which everything present has the +declared *shape*. + +### Typed attribute values + +By default, an attribute value reaches the graph as a string, like every other extracted +value. `typed_properties` additionally coerces the value of a *declared datatype property* +to its declared XSD type and writes it as a native graph property, so it can be compared, +sorted and range-queried. + +Given: + +```turtle +:foundedYear a owl:DatatypeProperty ; + rdfs:label "foundedYear" ; + rdfs:domain :Company ; rdfs:range xsd:integer ; + rdfs:comment "The four-digit year the company was founded." . +``` + +and `typed_properties='subject'`, a company entity carries an integer `foundedYear`: + +```cypher +MATCH (c:`__Entity__`) +WHERE c.class = 'Company' AND c.foundedYear < 2000 +RETURN c.value, c.foundedYear +``` + + + +#### Placements + +`typed_properties` says *where* the coerced value is written. It is deliberately not folded +into `ontology_authority`, because asking for `align` is asking about naming and is not thereby +asking for new properties on your nodes. + +With no ontology configured at all, none of this reaches your graph: no filter is added to +the extraction pipeline, `typed_properties` resolves to `'off'` and cannot be set by an +environment variable, and both the emitted graph queries and the extraction prompts are +byte-identical to a build from before the feature existed. + +| Placement | Where the value is written | Needs `include_local_entities` | +| ------------- | ------------- | ------------- | +| `'off'` | Nowhere. The graph is byte-identical to a build without an ontology. | — | +| `'subject'` | On the **subject** `__Entity__`, keyed by the property's own local name. | No | +| `'complement'` | On the **complement** entity, as `typed_value` and `datatype`. | **Yes** | +| `'both'` | Each of the above, to its own node. Not a fallback chain. | **Yes** | + +**`'subject'` is the recommended placement**, and the only one that works without +`include_local_entities`. It is also the one that answers the question typed properties +exist for — "which companies were founded before 2000" — because the value is on the entity +that has it. Complement entities are created only when `include_local_entities` is on (it +defaults to `False` — see the `BuildConfig` parameters in +[Indexing](/graphrag-toolkit/lexical-graph/indexing/#configuring-the-extract-and-build-stages)), +and the typed write is made from inside the branch that creates the node: it will not +conjure a node the setting would have suppressed. Asking for `'complement'` or `'both'` with +`include_local_entities` off is refused with a `ValueError` when the build pipeline is +configured, rather than silently writing nothing. + +An ontology whose vocabulary would collide with the graph model is refused at +configuration time rather than at build time: `__Entity__` owns `value`, `search_str` and +`class`, and complement placement additionally owns `typed_value` and `datatype`. A +declared datatype property with one of those names raises `ValueError` from +`OntologyConfig`, scoped to the placement you actually asked for. + +#### What lands on a node + +The typed write is **additive**. The string `value` stays exactly where it was; nothing is +replaced, renamed or removed, and the `__Fact__` node and its `__SUBJECT__` / `__OBJECT__` +edges are created exactly as they would have been. + +At `'subject'`: + +``` +(:__Entity__ {entityId: '…', value: 'Halcyon Motors', class: 'Company', + search_str: 'halcyon motors', foundedYear: 1971}) +``` + +At `'complement'`, on the entity created for the value string itself: + +``` +(:__Entity__ {entityId: '…', value: '1971', class: '__Local_Entity__', + typed_value: 1971, datatype: 'http://www.w3.org/2001/XMLSchema#integer'}) +``` + + + +#### Coercion, and what happens when it refuses + +Coercion parses the literal as its declared range. It handles lexical variation on a value +that has already been stated — digit grouping (`'1,994'` → `1994`), a named month +(`'March 3rd 2020'` → `'2020-03-03'`). + +It does *not* guess at readings of the text. `'$1.2bn'` under `xsd:double` is **refused**, +and so is `'publicly listed'` under `xsd:boolean`. Inferring those is the extractor's job; +a coercer that guessed would be inventing data under the authority of a declared type. + +A refusal is lossless at `align`: no typed property is written, the fact and its string +value are still in the graph, and you can see the literal that would not parse. The usual +casualty is a currency amount — the fix is a `rdfs:comment` on the property telling the +model to state a bare number, as `news.ttl` does for `:revenue`. + +At `'strict'`, `enforce_datatypes` is on and the fact is **dropped** rather than kept +untyped. Refusal is lossless at `align`; it is not at `strict`. + +### Who reads typed properties + + + +The consumers typed properties exist for today are: + + - **BYOKG-RAG's schema-grounded openCypher generation.** `ByoKGQueryEngine` reads the + graph's schema and threads it into every generation prompt, and the Cypher linker's + prompt instructs the model to use only the node types, relationship types and + **properties** in that schema — and to handle filtering, aggregation and sorting. A + typed property that appears in the schema is a property the generated query can filter + on. + - **The context ontology accelerator.** + - **Direct Cypher** from notebooks or analytics tooling, where you write the query + yourself. + +Entity-scoped filtering and typed-property projection inside `lexical-graph`'s retrievers +are intended follow-ons. + +#### Confirming your properties are visible + +On Neptune Analytics, this is the cheapest end-to-end check that the feature has a working +consumer — it is the same call BYOKG-RAG makes to discover the schema: + +```cypher +CALL neptune.graph.pg_schema() +YIELD schema +RETURN schema +``` + +Your declared property names should appear under +`nodeLabelDetails["__Entity__"]["properties"]`. If they do not, nothing was written — check +that `typed_properties` is not `'off'`, that the predicates actually resolved to declared +datatype properties, and (for `'complement'`) that `include_local_entities` is on. + +#### Two consumption paths + +The cost of a mis-mapped predicate depends entirely on who is consuming the graph, and the +two cases are very different. This is a description of **current behaviour**, and each half +of it rests on a named default: + +*Within `lexical-graph`'s own RAG loop, the source text is the guard.* `include_facts` +defaults to `False`, so statements are emitted without their fact strings, and `ClearChunks` +is one of the default formatting processors, so chunk text is emptied. What reaches the +answering model is the statement — proposition text, which the proposition prompt instructs +to preserve the original phrasing. So by default the answering model never sees a predicate +label, an entity classification, or a fact string at all. Relationship labels are a +*retrieval* mechanism here, not an assertion the answer rests on: a wrong `worksFor` cannot +fabricate a claim in the answer, but it does change traversal and scoring, so its cost is +retrieval precision and recall. (With `include_facts=True`, fact strings appear beside the +proposition they came from, and the model can weigh the label against its own grounding.) + +*For the structured consumers, there is no text guard.* BYOKG-RAG generating openCypher, +the context ontology accelerator, and direct Cypher all consume the structure **as truth** +— there is no proposition alongside it to fall back on. For them a wrong predicate is a +wrong answer. + +`include_facts=False` and `ClearChunks` are configuration defaults rather than design +intent. If facts, canonical class names or typed properties do reach the answering context in +a given configuration, this guard weakens and correct labelling matters more. + +### Operational notes + +#### Annotations are frozen at extraction + +The ontology's influence is applied **once**, at the moment text becomes facts. The +resolved terms — the class IRI, the property IRI, the canonical name and the datatype — are +written into the fact metadata during extraction and never recomputed. + +So changing the ontology, changing `ontology_authority`, or adding a `skos:altLabel` has **no +effect on already-extracted data**. It requires **re-extraction**, not a rebuild. This is +the general form of the re-index boundary below. + +One consequence worth knowing before you hit it: a build with `typed_properties` set, run +against extraction output produced *without* an ontology — a checkpoint, or `FileBasedDocs` +from an earlier run — finds no annotations and writes nothing. That case is reported rather +than passing silently, but the fix is to re-extract. + +#### Enabling an ontology is a re-index boundary + +Turning an ontology on — or changing one in a way that changes an authored class name — +does **not** update an existing graph in place. `normalize_names` can change +`entity.class`, and `include_classification_in_entity_id` defaults to `True`, so the +entity id is derived in part from the classification. A company extracted as +`CORPORATION` before and `Company` after has two different entity ids and becomes two +nodes, not one. + +Treat it as you would a chunking change: build into a fresh +[tenant](/graphrag-toolkit/lexical-graph/multi-tenancy/) or a fresh graph. Building an +ontology-guided extraction into a tenant that already holds an unguided one merges two +vocabularies silently, which is worse than either alone. + +#### Subject-mode values are last-write-win + +A subject-placed typed property is written with an unconditional `SET`, once per fact. +Entities recur across chunks, so if two chunks state a different `foundedYear` for the +same company, the property ends up holding whichever was processed last — and build order +is not deterministic. There is no reconciliation, no first-write-wins, and no conflict +signal. + +This is usually fine, because the two chunks usually agree. When they do not, the graph +records one value and the disagreement is only visible in the underlying `__Fact__` nodes, +which retain both. diff --git a/examples/lexical-graph/ontologies/news.ttl b/examples/lexical-graph/ontologies/news.ttl new file mode 100644 index 00000000..66b3bcff --- /dev/null +++ b/examples/lexical-graph/ontologies/news.ttl @@ -0,0 +1,188 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# A worked ontology for ontology-guided extraction over a business-news corpus. +# +# from graphrag_toolkit.lexical_graph import ExtractionConfig, IndexingConfig +# from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig +# +# ontology = OntologyConfig( +# 'examples/lexical-graph/ontologies/news.ttl', +# ontology_authority='align', +# typed_properties='subject', +# ) +# indexing_config = IndexingConfig(extraction=ExtractionConfig(ontology=ontology)) +# +# Turtle is the only file format the toolkit reads. For any other serialization, +# parse it into an `rdflib.Graph` yourself and pass the graph: +# +# import rdflib +# graph = rdflib.Graph().parse('news.owl', format='xml') +# ontology = OntologyConfig(graph, ontology_authority='align') +# +# Everything below is here for a reason, and the comments say which. Copy this +# file as a starting point rather than writing an ontology from scratch: the +# things that go wrong are mostly things this file gets right by accident. + +@prefix : . +@prefix owl: . +@prefix rdfs: . +@prefix skos: . +@prefix xsd: . + + a owl:Ontology ; + rdfs:comment "A small ontology for business news: who works where, who owns what, and the facts about a company you would want to filter on." . + +################################################################################ +# Classes +# +# `rdfs:label` is what gets stored, and it takes precedence over the IRI local +# name. `:SportsTeam` is stored as `Sports Team` because that is its label; drop +# the label and it would be stored as `SportsTeam`. Pick whichever spelling you +# want to read back out of the graph and put it in the label. +# +# One exception, and it matters if you use `typed_properties`: the *typed property +# key* comes from the IRI local name, not from the label. Label `:revenue` as +# "Annual Revenue" and the fact predicate reads `Annual Revenue` while the property +# on the node is still `revenue`. Every datatype property below is labelled to +# match its local name for exactly that reason. Classes and object properties are +# not affected - `:SportsTeam` is freely labelled `Sports Team` - because nothing +# keys a graph property from them. +# +# `skos:altLabel` is the other half: a phrasing the model might produce, mapped +# onto the term you want. `Corporation`, `Ball Club` and `Firm` are not stored - +# they are how a model's wording finds its way to `Company`. +# +# `rdfs:comment` is not decoration. It is rendered into the extraction prompt, so +# it is the only place you can tell the model what a class means. A class with no +# comment is a bare name the model has to guess at. +################################################################################ + +:Agent a owl:Class ; + rdfs:label "Agent" ; + rdfs:comment "Anything that can act - a person or an organization." . + +:Person a owl:Class ; rdfs:subClassOf :Agent ; + rdfs:label "Person" ; + rdfs:comment "An individual human being, named in the source." . + +:Company a owl:Class ; rdfs:subClassOf :Agent ; + rdfs:label "Company" ; + skos:altLabel "Corporation" , "Firm" ; + rdfs:comment "An incorporated commercial organization." . + +:SportsTeam a owl:Class ; rdfs:subClassOf :Company ; + rdfs:label "Sports Team" ; + skos:altLabel "Ball Club" ; + rdfs:comment "A professional team that competes in a league. A sports team is a kind of company." . + +:Exchange a owl:Class ; + rdfs:label "Stock Exchange" ; + rdfs:comment "A market on which company shares are listed and traded." . + +################################################################################ +# Object properties - relationships between two entities +# +# `rdfs:domain` and `rdfs:range` are declarations about which classes may sit at +# each end. At `ontology_authority='strict'` they are enforced and a fact that violates +# them is dropped; below `strict` they are documentation, and they still shape the +# prompt. Subclassing is respected: `:playsFor` accepts a `Sports Team` in its +# range because `:SportsTeam rdfs:subClassOf :Company`, so declaring the range as +# `:Company` does not exclude teams. +# +# Leave domain and range off - as `:acquired` does - and the property applies +# anywhere. That is the right choice when the relationship genuinely is general, +# and the wrong one when you have just not decided; an undeclared domain cannot +# be enforced later without re-extracting. +################################################################################ + +:worksFor a owl:ObjectProperty ; + rdfs:label "worksFor" ; + rdfs:domain :Person ; rdfs:range :Company ; + skos:altLabel "EMPLOYED BY" , "WORKS AT" ; + rdfs:comment "Employment of a person by a company." . + +:playsFor a owl:ObjectProperty ; + rdfs:label "playsFor" ; + rdfs:domain :Person ; rdfs:range :Company ; + skos:altLabel "PLAYS ON" ; + rdfs:comment "Membership of a person in a sports team." . + +:subsidiaryOf a owl:ObjectProperty ; + rdfs:label "subsidiaryOf" ; + rdfs:domain :Company ; rdfs:range :Company ; + rdfs:comment "Majority ownership of one company by another." . + +:listedOn a owl:ObjectProperty ; + rdfs:label "listedOn" ; + rdfs:domain :Company ; rdfs:range :Exchange ; + rdfs:comment "The exchange on which a company's shares trade." . + +:acquired a owl:ObjectProperty ; + rdfs:label "acquired" ; + rdfs:comment "Purchase of one organization by another. Domain and range are deliberately undeclared: acquisitions in this corpus involve divisions and brands as well as whole companies." . + +################################################################################ +# Datatype properties - facts about one entity, with a declared value type +# +# These are what `typed_properties` writes. With `typed_properties='subject'`, +# `:foundedYear` becomes an integer property named `foundedYear` on the company's +# own `__Entity__` node, so `WHERE c.foundedYear < 2000` is a query you can +# actually run. Leave `typed_properties` at its `'off'` default and these are +# still useful - they shape the prompt and they name the predicate - but the value +# stays a string. +# +# The declared `rdfs:range` is what the value is coerced to, so choose it for what +# you want to query rather than for what the source text looks like. Two rules: +# +# * a property name must not be `value`, `search_str`, `class`, `typed_value` or +# `datatype` - the graph model owns those on `__Entity__`, and configuration +# is refused outright rather than overwriting them; and +# * `xsd:date` is stored as an ISO-8601 **string**, not a native temporal, so +# quote the bound in a range query: `WHERE c.incorporatedOn >= '2016-01-01'`. +# It sorts and compares correctly as a string. +# +# A value that will not parse as its declared range writes no typed property at +# all - the string is still on the graph, so nothing is lost, but there is no +# typed value to filter on. `revenue` is the usual casualty: a model that writes +# "$1.2bn" has stated a number the coercer will not guess at. +################################################################################ + +:foundedYear a owl:DatatypeProperty ; + rdfs:label "foundedYear" ; + rdfs:domain :Company ; rdfs:range xsd:integer ; + skos:altLabel "FOUNDED IN" ; + rdfs:comment "The four-digit year the company was founded." . + +:incorporatedOn a owl:DatatypeProperty ; + rdfs:label "incorporatedOn" ; + rdfs:domain :Company ; rdfs:range xsd:date ; + rdfs:comment "The date of incorporation, as stated in the source." . + +:revenue a owl:DatatypeProperty ; + rdfs:label "revenue" ; + rdfs:domain :Company ; rdfs:range xsd:double ; + skos:altLabel "TURNOVER" ; + rdfs:comment "Annual revenue in US dollars, as a bare number without currency symbols or units." . + +:employeeCount a owl:DatatypeProperty ; + rdfs:label "employeeCount" ; + rdfs:domain :Company ; rdfs:range xsd:integer ; + skos:altLabel "HEADCOUNT" ; + rdfs:comment "Total number of employees." . + +:isPubliclyTraded a owl:DatatypeProperty ; + rdfs:label "isPubliclyTraded" ; + rdfs:domain :Company ; rdfs:range xsd:boolean ; + rdfs:comment "Whether the company's shares are listed on a public exchange." . + +:tickerSymbol a owl:DatatypeProperty ; + rdfs:label "tickerSymbol" ; + rdfs:domain :Company ; rdfs:range xsd:string ; + skos:altLabel "Stock Symbol" , "TICKER" ; + rdfs:comment "The company's exchange ticker symbol." . + +:jobTitle a owl:DatatypeProperty ; + rdfs:label "jobTitle" ; + rdfs:domain :Person ; rdfs:range xsd:string ; + rdfs:comment "The person's role at their employer." . diff --git a/lexical-graph/output.log b/lexical-graph/output.log new file mode 100644 index 00000000..e69de29b diff --git a/lexical-graph/pyproject.toml b/lexical-graph/pyproject.toml index ef3a788a..b3aae963 100644 --- a/lexical-graph/pyproject.toml +++ b/lexical-graph/pyproject.toml @@ -38,6 +38,12 @@ python_files = ["test_*.py"] python_classes = ["Test*"] python_functions = ["test_*"] asyncio_mode = "auto" +# pytest's default norecursedirs contains "build", which silently skipped the +# whole of tests/unit/indexing/build/ - the builder tests ran only when that path +# was named explicitly, and never in a full-suite or CI run. Restated here without +# "build". Dropping it is safe because `testpaths` already confines collection to +# `tests/`, so no generated build/ directory is reachable from here. +norecursedirs = ["*.egg", ".*", "_darcs", "CVS", "dist", "node_modules", "venv", "{arch}"] addopts = [ "-v", "--tb=short", diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/config.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/config.py index 2a0a5ced..2347c9e7 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/config.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/config.py @@ -57,6 +57,10 @@ DEFAULT_INCLUDE_DOMAIN_LABELS = False DEFAULT_INCLUDE_LOCAL_ENTITIES = False DEFAULT_INCLUDE_CLASSIFICATION_IN_ENTITY_ID = True +# Restated here rather than imported from `ontology_config` so +# that config.py stays free of the extract package (and of rdflib); the two are +# held together by a test. +DEFAULT_TYPED_PROPERTIES = 'off' DEFAULT_ENABLE_CACHE = False DEFAULT_METADATA_DATETIME_SUFFIXES = ['_date', '_datetime'] DEFAULT_OPENSEARCH_ENGINE = 'nmslib' @@ -319,6 +323,7 @@ class _GraphRAGConfig: _batch_writes_enabled (Optional[bool]): Flag indicating whether batch writes are enabled. _include_domain_labels (Optional[bool]): Whether domain-specific labels are included in processes. _include_local_entities (Optional[bool]): Whether local entities are included in the graph. + _typed_properties (Optional[str]): Where coerced attribute values are stored, if anywhere. _enable_cache (Optional[bool]): Boolean flag to enable or disable caching mechanisms. _metadata_datetime_suffixes (Optional[List[str]]): List of datetime suffixes included in metadata handling. """ @@ -347,6 +352,7 @@ class _GraphRAGConfig: _include_domain_labels: Optional[bool] = None _include_local_entities: Optional[bool] = None _include_classification_in_entity_id: Optional[bool] = None + _typed_properties: Optional[str] = None _enable_cache: Optional[bool] = None _metadata_datetime_suffixes: Optional[List[str]] = None _opensearch_engine: Optional[str] = None @@ -933,7 +939,32 @@ def include_local_entities(self, include_local_entities: bool) -> None: self._include_local_entities = include_local_entities @property - def include_classification_in_entity_id(self) -> bool: + def typed_properties(self) -> str: + """Where coerced attribute values are stored, if anywhere. + + Deliberately **not** environment-readable, which is where this setting + departs from `include_local_entities` and every other flag around it. + No ambient configuration should be able to turn on graph writes, and the + setting should be unreachable as anything other than `'off'` for a user + with no ontology; an env var would + break both at once, since an exported `TYPED_PROPERTIES` would add + properties to the nodes of a build whose author never mentioned an + ontology. So the value comes from `OntologyConfig`, and this field exists + only so the `coalesce` chain from `BuildPipeline` has a floor. + + Returns: + str: `'off'` unless something set it programmatically. + """ + if self._typed_properties is None: + self.typed_properties = DEFAULT_TYPED_PROPERTIES + return self._typed_properties + + @typed_properties.setter + def typed_properties(self, typed_properties: str) -> None: + self._typed_properties = typed_properties + + @property + def include_classification_in_entity_id(self) -> bool: if self._include_classification_in_entity_id is None: self.include_classification_in_entity_id = string_to_bool(os.environ.get('INCLUDE_CLASSIFICATION_IN_ENTITY_ID'), DEFAULT_INCLUDE_CLASSIFICATION_IN_ENTITY_ID) return self._include_classification_in_entity_id diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/build_pipeline.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/build_pipeline.py index efeef6ec..5ab727f2 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/build_pipeline.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/build_pipeline.py @@ -12,6 +12,7 @@ from graphrag_toolkit.lexical_graph.config import GraphRAGConfig from graphrag_toolkit.lexical_graph.metadata import SourceMetadataFormatter, DefaultSourceMetadataFormatter from graphrag_toolkit.lexical_graph.indexing import NodeHandler, IdGenerator +from graphrag_toolkit.lexical_graph.indexing.constants import COMPLEMENT_PLACEMENTS, TYPED_PROPERTY_PLACEMENTS from graphrag_toolkit.lexical_graph.indexing.utils.pipeline_utils import run_pipeline from graphrag_toolkit.lexical_graph.indexing.model import SourceType, SourceDocument, source_documents_from_source_types from graphrag_toolkit.lexical_graph.indexing.build.node_builder import NodeBuilder @@ -55,6 +56,8 @@ class BuildPipeline(): batch_writes_enabled (bool): Flag indicating whether batch writes are enabled. batch_write_size (int): Size of batches for processing writes. Defaults to a configured size. include_domain_labels (bool): Flag indicating whether domain labels should be included. + typed_properties (str): Where coerced attribute values are stored, if anywhere; + `'off'` unless an ontology asked otherwise. node_builders (NodeBuilders): Object that encapsulates the logic for building nodes, applying filters, and formatting metadata. node_filter (NodeFilter): Filter used for excluding or including nodes based on certain conditions. @@ -74,6 +77,7 @@ def create(components: List[TransformComponent], include_domain_labels:Optional[bool]=None, include_local_entities:Optional[bool]=None, include_classification_in_entity_id:Optional[bool]=None, + typed_properties:Optional[str]=None, tenant_id:Optional[TenantId]=None, progress_monitor:Optional[ProgressMonitor]=None, **kwargs:Any @@ -108,6 +112,9 @@ def create(components: List[TransformComponent], incorporated in the output. Defaults to None. include_local_entities (Optional[bool]): Specifies whether local entities are included in the graph. Defaults to None. + typed_properties (Optional[str]): Where coerced attribute values are + stored: `'off'`, `'subject'`, `'complement'` or `'both'`. Defaults to + None, which resolves to `'off'`. tenant_id (Optional[TenantId]): Identifier for tenant-specific operations or segregations. Defaults to None. **kwargs (Any): Additional keyword arguments to customize further configuration @@ -132,6 +139,7 @@ def create(components: List[TransformComponent], include_domain_labels=include_domain_labels, include_local_entities=include_local_entities, include_classification_in_entity_id=include_classification_in_entity_id, + typed_properties=typed_properties, tenant_id=tenant_id, progress_monitor=progress_monitor, **kwargs @@ -152,6 +160,7 @@ def __init__(self, include_domain_labels:Optional[bool]=None, include_local_entities:Optional[bool]=None, include_classification_in_entity_id:Optional[bool]=None, + typed_properties:Optional[str]=None, tenant_id:Optional[TenantId]=None, progress_monitor:Optional[ProgressMonitor]=None, **kwargs:Any @@ -188,9 +197,16 @@ def __init__(self, included in the output during processing. Defaults to a preconfigured value. include_local_entities (Optional[bool]): Specifies whether local entities are included in the graph. Defaults to a preconfigured value. + typed_properties (Optional[str]): Where coerced attribute values are + stored: `'off'`, `'subject'`, `'complement'` or `'both'`. Defaults to + a preconfigured value, which is `'off'`. tenant_id (Optional[TenantId]): An identifier for the tenant, used for scoping data. Defaults to None. **kwargs (Any): Additional keyword arguments to configure the pipeline behavior. + + Raises: + ValueError: If `typed_properties` is not a supported placement, or if + it requests complement placement without `include_local_entities`. """ components = components or [] num_workers = coalesce(num_workers, GraphRAGConfig.build_num_workers) @@ -200,8 +216,31 @@ def __init__(self, include_domain_labels = coalesce(include_domain_labels, GraphRAGConfig.include_domain_labels) include_local_entities = coalesce(include_local_entities, GraphRAGConfig.include_local_entities) include_classification_in_entity_id = coalesce(include_classification_in_entity_id, GraphRAGConfig.include_classification_in_entity_id) + typed_properties = coalesce(typed_properties, GraphRAGConfig.typed_properties) source_metadata_formatter = source_metadata_formatter or DefaultSourceMetadataFormatter() - + + if typed_properties not in TYPED_PROPERTY_PLACEMENTS: + raise ValueError( + f'Unknown typed_properties placement: {typed_properties!r}. ' + f'Expected one of {", ".join(TYPED_PROPERTY_PLACEMENTS)}.' + ) + + # Complement placement writes to the complement node, + # and with `include_local_entities` off that node is never created - so + # the write would target nothing and the user would be left looking for + # properties on a node that does not exist. Checked here because this is + # the first place both settings are resolved: `typed_properties` comes + # from the ontology and `include_local_entities` from the build config. + if typed_properties in COMPLEMENT_PLACEMENTS and not include_local_entities: + raise ValueError( + f'typed_properties={typed_properties!r} writes typed_value and datatype ' + 'to the complement node, but include_local_entities is not enabled, so ' + 'no complement node is created and the write would have no target. ' + 'Set include_local_entities=True, or use ' + "typed_properties='subject'." + ) + + for c in components: if isinstance(c, NodeHandler): c.show_progress = show_progress @@ -228,6 +267,7 @@ def __init__(self, self.batch_write_size = batch_write_size self.include_domain_labels = include_domain_labels self.include_local_entities = include_local_entities + self.typed_properties = typed_properties self.node_builders = NodeBuilders( builders=builders, build_filters=build_filters, @@ -320,6 +360,7 @@ def build(self, inputs: Iterable[SourceType]): batch_write_size=self.batch_write_size, include_domain_labels=self.include_domain_labels, include_local_entities=self.include_local_entities, + typed_properties=self.typed_properties, versioning_timestamp=build_timestamp, **self.pipeline_kwargs ) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py index 840c9494..4e64f4fe 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py @@ -8,13 +8,228 @@ from graphrag_toolkit.lexical_graph.storage.graph import GraphStore from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import search_string_from, label_from, new_query_var, escape_cypher_label from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder -from graphrag_toolkit.lexical_graph.indexing.constants import DEFAULT_CLASSIFICATION, LOCAL_ENTITY_CLASSIFICATION +from graphrag_toolkit.lexical_graph.indexing.constants import ( + COMPLEMENT_ENTITY_PROPERTIES, + COMPLEMENT_PLACEMENTS, + DEFAULT_CLASSIFICATION, + LOCAL_ENTITY_CLASSIFICATION, + RESERVED_ENTITY_PROPERTIES, + SUBJECT_PLACEMENTS, + TYPED_PROPERTIES_OFF, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.datatype_utils import coerce_literal from graphrag_toolkit.lexical_graph.indexing.utils.fact_utils import string_complement_to_entity from llama_index.core.schema import BaseNode logger = logging.getLogger(__name__) +# The no-annotations warning state, and the reason it is module-level. +# +# A build that asks for typed properties and writes none is the worst outcome +# available, because the user's next move is to doubt the graph store rather than +# the pipeline order - the cause is almost always that extraction ran before the +# ontology existed, or from a checkpoint that predates it, so the facts carry no +# annotations for the builder to key a property from. Rebuilding does not fix +# that; re-extracting does. +# +# Detecting it needs state that outlives one fact: a builder sees facts one at a +# time and is never told a build has ended, so there is nowhere else to notice +# "none of them were annotated". Module-level means once per worker process, +# which is the honest scope - the same compromise `OntologyFilter` makes for its +# unvalidated-datatype warning, and for the same reason. +# +# The threshold exists so a single legitimately-unresolved fact cannot cry wolf. +# One fact whose predicate is not in the ontology carries no annotations and is +# entirely normal at `align`; twenty-five in a row with not one annotated is not +# something a working filter produces. +_NO_ANNOTATIONS_WARNING_AFTER = 25 + +_annotation_seen = False +_unannotated_facts = 0 +_no_annotations_warned = False + +def _has_ontology_annotations(fact:Fact) -> bool: + """Return True if any part of this fact was annotated by the filter. + + Read across the whole fact rather than off the one field a placement needs, + because the question being asked is "did the filter run at all", not "can + this particular fact be written". + """ + parts = [fact.subject, fact.object, fact.complement] + if any(part is not None and part.classIri for part in parts): + return True + if fact.complement is not None and fact.complement.datatype: + return True + return bool(fact.predicate.propertyIri or fact.predicate.canonicalName) + +def _warn_if_no_annotations(fact:Fact, typed_properties:str) -> None: + """Warn once per process that typed properties were asked for and cannot be written. + + Silence is not an option here: nothing downstream fails, so + the only signal the user would otherwise get is an absence. + """ + global _annotation_seen, _unannotated_facts, _no_annotations_warned + + if typed_properties == TYPED_PROPERTIES_OFF or _annotation_seen or _no_annotations_warned: + return + + if _has_ontology_annotations(fact): + _annotation_seen = True + return + + _unannotated_facts += 1 + + if _unannotated_facts < _NO_ANNOTATIONS_WARNING_AFTER: + return + + _no_annotations_warned = True + logger.warning( + 'typed_properties=%r was requested, but none of the first %d facts in this ' + 'build carry ontology annotations, so no typed properties can be written. ' + 'This means extraction ran without the ontology filter - most often because ' + 'the extracted facts predate the ontology, or come from a checkpoint that ' + 'does. Re-extract; rebuilding the same facts will not add the annotations. ' + 'Logged once per process.', + typed_properties, _unannotated_facts + ) + +def _reset_no_annotations_warning() -> None: + """Clear the no-annotations warning state. For tests only.""" + global _annotation_seen, _unannotated_facts, _no_annotations_warned + _annotation_seen = False + _unannotated_facts = 0 + _no_annotations_warned = False + +def _reserved_property_names(typed_properties:str) -> frozenset: + """The `__Entity__` property names a typed property must not key. + + The same set `OntologyConfig._validate_no_reserved_property_names` computes, + widened the same way: `typed_value` and `datatype` are only owned once + complement placement is writing them. Recomputed here rather than shared as a + function because the two callers reach it from opposite ends of the package - + validation from the ontology config, this from a builder in a spawn worker - + and neither module should import the other. + """ + reserved = set(RESERVED_ENTITY_PROPERTIES) + if typed_properties in COMPLEMENT_PLACEMENTS: + reserved.update(COMPLEMENT_ENTITY_PROPERTIES) + return frozenset(reserved) + +def _typed_subject_property(fact:Fact, typed_properties:str): + """Resolve the one typed property a fact contributes to its subject, if any. + + Reserved names, unannotated facts and uncoercible literals are all a question + about *this* fact, so they are answered in one place and returned as data, rather than spread across + guard clauses in `build`. + + `complement.datatype` is the discriminator, not `predicate.canonicalName`. + `OntologyFilter._annotate` sets `canonicalName` for every resolved predicate + including object properties, and sets `datatype` only when the predicate + resolved to a declared *datatype* property with a literal complement. Keying + off `canonicalName` alone would write the object entity's display string into + an attribute slot on the subject. + + Args: + fact: The validated fact, after `string_complement_to_entity`. + typed_properties: The resolved placement. + + Returns: + A `(key, value)` pair, or `None` when nothing should be written. `value` + may be `False`, `0` or `0.0`, so callers must test the pair, not the value. + """ + if typed_properties not in SUBJECT_PLACEMENTS: + return None + + key = fact.predicate.canonicalName + complement = fact.complement + + # An unresolved predicate, or one that resolved to an object + # property, has no declared datatype and contributes no typed property. + if not key or not isinstance(complement, Entity) or not complement.datatype: + return None + + # A hostile property name can be rejected or escaped, and a line break is the + # one case where rejecting is clearly better. + # `escape_cypher_label` handles every other character by doubling backticks, and + # a quoted identifier may legally contain a newline - but it would split the + # query across lines the rest of the builder assumes are one statement each, and + # a raw line break cannot occur in a Turtle IRI or prefixed name, so a name + # carrying one did not come from a well-formed ontology. + if '\n' in key or '\r' in key: + logger.warning( + 'Skipping typed property %r on entity %r: the property name contains a ' + 'line break, which no well-formed ontology term does.', + key, fact.subject.entityId + ) + return None + + # Defence in depth. `OntologyConfig` has already refused an + # ontology that could get here, but a fact can also arrive from a checkpoint + # written under a different config, and this builder is the last place before + # the write. + if key in _reserved_property_names(typed_properties): + logger.warning( + 'Skipping typed property %r on entity %r: the graph model already owns ' + 'that property on __Entity__. Rename the property in the ontology.', + key, fact.subject.entityId + ) + return None + + # Refuse rather than store the raw string under a key whose + # name promises a number. + value = coerce_literal(complement.value, complement.datatype) + if value is None: + logger.debug( + f'Skipping typed property [key: {key}, value: {complement.value!r}, ' + f'datatype: {complement.datatype}] - the literal does not coerce' + ) + return None + + return (key, value) + +def _typed_complement_values(entity:Entity, typed_properties:str): + """Resolve the `typed_value` / `datatype` pair for a complement node, if any. + + The mirror of `_typed_subject_property`, and deliberately not + folded into it: the two placements write different property names, to different + nodes, under different conditions, and the only thing they share is + `coerce_literal`. `'both'` runs both, each to its own node - it is not a + fallback chain. + + No reserved-name check, because both names are fixed by the graph model rather + than taken from the ontology. The collision runs the other way, and + `_reserved_property_names` is where it is handled: an ontology declaring + `:typed_value` cannot key a subject write once complement placement is on. + + Args: + entity: The complement entity, after `string_complement_to_entity`. + typed_properties: The resolved placement. + + Returns: + A `(typed_value, datatype)` pair, or `None` when nothing should be written. + `typed_value` may be `False`, `0` or `0.0`, so callers must test the pair. + """ + if typed_properties not in COMPLEMENT_PLACEMENTS: + return None + + if not isinstance(entity, Entity) or not entity.datatype: + return None + + # The reason `datatype` is never written on its own: a + # `datatype` without a `typed_value` would assert a type for a value that is + # not there, which is worse than the absence of both. The string is still on + # the node as `value`, where every existing consumer reads it. + value = coerce_literal(entity.value, entity.datatype) + if value is None: + logger.debug( + f'Skipping typed complement value [value: {entity.value!r}, ' + f'datatype: {entity.datatype}] - the literal does not coerce' + ) + return None + + return (value, entity.datatype) + class EntityGraphBuilder(GraphBuilder): """ Handles the process of building and interacting with a graph database for entity and fact data @@ -61,17 +276,34 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): node (BaseNode): The node from which fact metadata is to be extracted. graph_client (GraphStore): The graph database client to execute queries. **kwargs (Any): Additional options, such as `include_domain_labels`, which - determines whether domain-specific labels are added to the entities. + determines whether domain-specific labels are added to the entities, + and `typed_properties`, which says where a coerced attribute value is + stored and defaults to `'off'` when absent. At `'subject'` or + `'both'`, a fact whose predicate resolved to a declared datatype + property additionally gets that value written onto the subject + `__Entity__` under the property's canonical name. At `'complement'` + or `'both'`, the complement `__Entity__` - where one is created at + all, which needs `include_local_entities` - additionally gets + `typed_value` and `datatype`, keeping its string `value`. `'both'` + does each of those to its own node; it is not a fallback chain. """ fact_metadata = node.metadata.get('fact', {}) include_domain_labels = kwargs['include_domain_labels'] include_local_entities = kwargs['include_local_entities'] + # `.get` rather than a subscript, unlike the two above. + # Those are always supplied by `BuildPipeline`, but a caller who + # constructed a pipeline before this setting existed - or who calls a + # builder directly, as several tests do - would otherwise get a `KeyError` + # from a feature they never asked for. + typed_properties = kwargs.get('typed_properties', TYPED_PROPERTIES_OFF) if fact_metadata: fact = Fact.model_validate(fact_metadata) fact = string_complement_to_entity(fact) + _warn_if_no_annotations(fact, typed_properties) + if fact.subject.classification == LOCAL_ENTITY_CLASSIFICATION: if not include_local_entities: logger.debug(f'Ignoring local entities for fact [fact_id: {fact.factId}]') @@ -100,12 +332,118 @@ def insert_for_entity(entity:Entity): graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7) + def insert_typed_complement_values(entity:Entity): + """Set `typed_value` and `datatype` on a complement `__Entity__`. + + A second query for the same reason the + subject write is one: `insert_for_entity` stays byte-identical at + every placement, so the complement's string `value` cannot be + displaced by an edit here, and the recorded baseline keeps + proving something. + + Called only from the branch that inserts the complement node, so + the write cannot conjure a node that `include_local_entities` + would not have created. Duplicating the + condition instead would be one refactor away from a `MERGE` that + creates local entities behind the setting's back. + + Both keys are fixed by the graph model, so neither needs escaping, + but the assignments still go through `property_assigment_fn` +. Both are the identity on every store today - + neither name matches `metadata_datetime_suffixes` - and routing + through it anyway keeps the store's decision the store's to make. + """ + typed_values = _typed_complement_values(entity, typed_properties) + + if not typed_values: + return + + (typed_value, datatype) = typed_values + + (value_key, datatype_key) = COMPLEMENT_ENTITY_PROPERTIES + c_var = new_query_var() + value_assigment = graph_client.property_assigment_fn(value_key, typed_value)('params.typedValue') + datatype_assigment = graph_client.property_assigment_fn(datatype_key, datatype)('params.datatype') + + statements = [ + '// insert typed complement values', + 'UNWIND $params AS params', + f'MERGE ({c_var}:`__Entity__`{{{graph_client.node_id("entityId")}: params.entityId}})', + f'SET {c_var}.`{value_key}` = {value_assigment}, ' + f'{c_var}.`{datatype_key}` = {datatype_assigment}' + ] + + query = '\n'.join(statements) + + graph_client.execute_query_with_retry( + query, + self._to_params({ + 'entityId': entity.entityId, + 'typedValue': typed_value, + 'datatype': datatype + }), + max_attempts=5, max_wait=7 + ) + insert_for_entity(fact.subject) if fact.object and fact.object.entityId != fact.subject.entityId: insert_for_entity(fact.object) elif include_local_entities and fact.complement and fact.complement.entityId != fact.subject.entityId: insert_for_entity(fact.complement) + insert_typed_complement_values(fact.complement) + + typed_subject_property = _typed_subject_property(fact, typed_properties) + + if typed_subject_property: + + def insert_typed_subject_property(entity:Entity, key:str, value:Any): + """Set one coerced attribute value on the subject `__Entity__`. + + A second query rather than an extra `SET` on the insert above, + which costs a round trip and buys one thing + outright: the entity insert is byte-identical at every + placement, so `value`, `search_str` and `class` cannot be + disturbed by a change to this query, and there is no branch in + the insert that a future edit could get wrong. Same trade + `insert_domain_entity` below already makes; batching both into + the insert is a future optimization, not a correctness fix. + + `key` is a Cypher identifier, so it + is backtick-quoted and escaped; `value` is bound, and the + assignment goes through the store's `property_assigment_fn` so + Neptune's `datetime(...)` wrapper still applies. + """ + e_var = new_query_var() + e_key = escape_cypher_label(key) + + # Bound under a fixed name rather than under the property's own + # name, which is where this deliberately departs from + # `source_graph_builder`. That builder's keys come from document + # metadata and it orders `sourceId` last to win a collision; + # here the key comes from an ontology, so a property canonically + # named `entityId` would otherwise collide with the merge key in + # the params dict. `property_assigment_fn` still receives the + # real key, so the store's name-based decisions are unchanged. + assigment = graph_client.property_assigment_fn(key, value)('params.typedValue') + + statements = [ + '// insert typed property', + 'UNWIND $params AS params', + f'MERGE ({e_var}:`__Entity__`{{{graph_client.node_id("entityId")}: params.entityId}})', + f'SET {e_var}.`{e_key}` = {assigment}' + ] + + query = '\n'.join(statements) + + graph_client.execute_query_with_retry( + query, + self._to_params({'entityId': entity.entityId, 'typedValue': value}), + max_attempts=5, max_wait=7 + ) + + (typed_key, typed_value) = typed_subject_property + insert_typed_subject_property(fact.subject, typed_key, typed_value) if include_domain_labels: diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/local_entity_rewrites_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/local_entity_rewrites_graph_builder.py index b2325354..788233f9 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/local_entity_rewrites_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/local_entity_rewrites_graph_builder.py @@ -8,7 +8,12 @@ from graphrag_toolkit.lexical_graph.storage.graph import GraphStore, Query, QueryTree from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder from graphrag_toolkit.lexical_graph.indexing.utils.fact_utils import string_complement_to_entity -from graphrag_toolkit.lexical_graph.indexing.constants import LOCAL_ENTITY_CLASSIFICATION +from graphrag_toolkit.lexical_graph.indexing.constants import ( + COMPLEMENT_ENTITY_PROPERTIES, + COMPLEMENT_PLACEMENTS, + LOCAL_ENTITY_CLASSIFICATION, + TYPED_PROPERTIES_OFF, +) from llama_index.core.schema import BaseNode @@ -24,6 +29,10 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): fact_metadata = node.metadata.get('fact', {}) include_local_entities = kwargs['include_local_entities'] + # `.get` rather than a subscript. This builder is in + # `default_builders()` unconditionally, so it runs for callers who have never + # heard of typed properties. + typed_properties = kwargs.get('typed_properties', TYPED_PROPERTIES_OFF) if fact_metadata: @@ -35,6 +44,31 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): logger.debug(f'Ignoring local entity rewrites for fact [fact_id: {fact.factId}]') return + # `c` is about to be `DETACH DELETE`d by the + # sibling query, taking `typed_value` and `datatype` with it, so they are + # carried onto the surviving node while `c` is still bound. + # + # `coalesce` and not a plain assignment: two complements can fold into the + # same real entity, and the second must not overwrite the first. First + # writer wins, which is arbitrary but at least stable within a run - the + # alternative is a value that changes with build order. + # + # Appended conditionally, and not emitted unconditionally on the grounds + # that `coalesce(null, null)` is harmless. It is not harmless: `SET x = + # null` deletes the property in some stores, it changes the query text + # every existing user sends, and it adds write work for people who did not + # ask for the feature. `''` at every other placement means the query below + # is byte-identical to the recorded baseline. + carry_typed_values = '' + + if typed_properties in COMPLEMENT_PLACEMENTS: + (typed_value_key, datatype_key) = COMPLEMENT_ENTITY_PROPERTIES + carry_typed_values = ( + f'SET n.`{typed_value_key}` = coalesce(n.`{typed_value_key}`, c.`{typed_value_key}`), ' + f'n.`{datatype_key}` = coalesce(n.`{datatype_key}`, c.`{datatype_key}`)\n' + ' ' + ) + copy_complement_relationships_to_subject = Query( query=f"""// copy complement relationships to subject UNWIND $params AS params @@ -43,7 +77,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): WHERE {graph_client.node_id('n.entityId')} = params.n_id AND {graph_client.node_id('c.entityId')} = params.c_id MERGE (s)-[:`__RELATION__`{{value:r.value}}]->(n) MERGE (n)-[:`__OBJECT__`]->(f) - """ + {carry_typed_values}""" ) delete_complement_relationships = Query( diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/constants.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/constants.py index cd668c15..de1ec017 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/constants.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/constants.py @@ -9,19 +9,51 @@ DEFAULT_TOPIC = 'context' DEFAULT_CLASSIFICATION = 'unknown' DEFAULT_ENTITY_CLASSIFICATIONS = [ - 'Company', - 'Location', - 'Event', - 'Sports Team', - 'Person', - 'Role', - 'Product', - 'Service', - 'Creative Work', - 'Software', + 'Company', + 'Location', + 'Event', + 'Sports Team', + 'Person', + 'Role', + 'Product', + 'Service', + 'Creative Work', + 'Software', 'Financial Instrument' ] +# Properties the graph model already owns on `__Entity__`, and which a typed +# attribute property therefore must not be allowed to key. +# `value` is the entity's identity string, `search_str` is what lookup matches +# on, and `class` is its classification - overwriting any of them would not add +# a queryable attribute, it would corrupt the node. +# +# Declared here, with the other graph-model names, rather than in +# `ontology_config.py` or in `entity_graph_builder.py`: both sides need the same +# list (validation rejects the ontology, the builder skips the write as defence +# in depth) and neither of those modules should import the other. +RESERVED_ENTITY_PROPERTIES = ('value', 'search_str', 'class') + +# Reserved in addition to the above once complement placement is active, since +# that is the placement which writes them. They are not +# reserved otherwise: nothing writes them, so nothing can be clobbered. +COMPLEMENT_ENTITY_PROPERTIES = ('typed_value', 'datatype') + +# Where a coerced attribute value is stored, if anywhere. `OntologyConfig` owns +# the setting and the `TypedProperties` type alias; the accepted values live here +# because the build pipeline and the builders test against them and this module +# is rdflib-free, which the ontology package is not. +TYPED_PROPERTY_PLACEMENTS = ('off', 'subject', 'complement', 'both') + +# The two placements, as membership tests rather than as string comparisons +# repeated at each call site. `'both'` is in each of them, which is the whole +# reason these exist rather than `== 'subject'`. +SUBJECT_PLACEMENTS = ('subject', 'both') +COMPLEMENT_PLACEMENTS = ('complement', 'both') + +# What `typed_properties` resolves to when nothing asked for anything. +TYPED_PROPERTIES_OFF = 'off' + diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/__init__.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/__init__.py index b984d341..30ae8401 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/__init__.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/__init__.py @@ -12,3 +12,22 @@ from .infer_classifications import InferClassifications from .infer_config import InferClassificationsConfig from .preferred_values import PREFERRED_VALUES_PROVIDER_TYPE, PreferredValuesProvider, default_preferred_values + +# Ontology configuration, re-exported so `ExtractionConfig(ontology=...)` and its +# argument come from the same place a user already imports the extractors from. +# +# This pulls rdflib into every process that imports the extract package, workers +# included, which the split between `Ontology` (owns the rdflib graph, parent +# process) and `OntologyIndex` (plain data, crosses the spawn boundary) was set +# up to avoid. Measured before accepting it: `import rdflib` is ~88ms against +# this package's own ~2.9s import, so ~3%, once per worker. Not worth a lazy +# module `__getattr__` and the tooling breakage that comes with one. The split +# still earns its keep for a different reason - it is what keeps `OntologyIndex` +# picklable - and extractors still take a rendered string, not an ontology. +from .ontology import Ontology, OntologyConfig, OntologyLoadError, OntologyType, TypedProperties, to_ontology_config + +# The one ontology component that goes into the pipeline. Exported here because +# `_configure_extraction_pipeline` appends it alongside the extractors above, and +# a user reading `extraction_components` should be able to import the type they +# find there from the same place. +from .ontology import OntologyFilter diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_extractor_base.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_extractor_base.py index 9ca636cc..2844e994 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_extractor_base.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_extractor_base.py @@ -38,6 +38,13 @@ class BatchExtractorBase(BaseExtractor): source_metadata_field:Optional[str] = Field(description='Metadata field from which to extract propositions') batch_inference_dir:str = Field(description='Directory for batch inputs and outputs') description:str = Field(description='Description') + # Declared here rather than on each subclass: both batch extractors compose + # it into self.prompt_template the same way, and the field has to exist on + # the model before either can be constructed with it. + ontology_constraints:str = Field( + default='', + description='Rendered ontology vocabulary block, composed into the prompt template at render time' + ) @classmethod def class_name(cls) -> str: diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_llm_proposition_extractor_sync.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_llm_proposition_extractor_sync.py index bbef8ae5..99fb20ef 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_llm_proposition_extractor_sync.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_llm_proposition_extractor_sync.py @@ -10,7 +10,7 @@ from graphrag_toolkit.lexical_graph.indexing.model import Propositions from graphrag_toolkit.lexical_graph.indexing.extract.batch_extractor_base import BatchExtractorBase from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY -from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_PROPOSITIONS_PROMPT +from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_PROPOSITIONS_PROMPT, with_ontology_constraints from graphrag_toolkit.lexical_graph.indexing.extract.batch_config import BatchConfig from graphrag_toolkit.lexical_graph.indexing.extract.llm_proposition_extractor import LLMPropositionExtractor @@ -34,7 +34,8 @@ def __init__(self, llm:LLMCacheType=None, prompt_template:str = None, source_metadata_field:Optional[str] = None, - batch_inference_dir:str = None): + batch_inference_dir:str = None, + ontology_constraints:str=''): super().__init__( batch_config = batch_config, @@ -45,7 +46,8 @@ def __init__(self, prompt_template=prompt_template or EXTRACT_PROPOSITIONS_PROMPT, source_metadata_field=source_metadata_field, batch_inference_dir=batch_inference_dir or os.path.join(GraphRAGConfig.local_output_dir, 'batch-propositions'), - description='Proposition' + description='Proposition', + ontology_constraints=ontology_constraints ) def _get_json(self, node, llm, inference_parameters): @@ -56,7 +58,11 @@ def _get_json(self, node, llm, inference_parameters): else: source_info = '' - messages = llm._get_messages(PromptTemplate(self.prompt_template), text=text, source_info=source_info) + messages = llm._get_messages( + PromptTemplate(with_ontology_constraints(self.prompt_template, self.ontology_constraints)), + text=text, + source_info=source_info + ) return { 'recordId': node.node_id, 'modelInput': get_request_body(llm, messages, inference_parameters) @@ -68,7 +74,8 @@ def _run_non_batch_extractor(self, nodes): extractor = LLMPropositionExtractor( prompt_template=self.prompt_template, - source_metadata_field=self.source_metadata_field + source_metadata_field=self.source_metadata_field, + ontology_constraints=self.ontology_constraints ) extracted = extractor.extract(all_nodes) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_topic_extractor_sync.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_topic_extractor_sync.py index 55894eaa..66df4f37 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_topic_extractor_sync.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_topic_extractor_sync.py @@ -8,7 +8,7 @@ from graphrag_toolkit.lexical_graph import GraphRAGConfig from graphrag_toolkit.lexical_graph.utils import LLMCache, LLMCacheType from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY -from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_TOPICS_PROMPT +from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_TOPICS_PROMPT, with_ontology_constraints from graphrag_toolkit.lexical_graph.indexing.extract.batch_config import BatchConfig from graphrag_toolkit.lexical_graph.indexing.extract.batch_extractor_base import BatchExtractorBase from graphrag_toolkit.lexical_graph.indexing.extract.topic_extractor import TopicExtractor @@ -40,7 +40,8 @@ def __init__(self, source_metadata_field:Optional[str] = None, batch_inference_dir:str = None, entity_classification_provider:Optional[PreferredValuesProvider]=None, - topic_provider:Optional[PreferredValuesProvider]=None): + topic_provider:Optional[PreferredValuesProvider]=None, + ontology_constraints:str=''): super().__init__( batch_config = batch_config, @@ -53,7 +54,8 @@ def __init__(self, batch_inference_dir=batch_inference_dir or os.path.join(GraphRAGConfig.local_output_dir, 'batch-topics'), description='Topic', entity_classification_provider=entity_classification_provider or default_preferred_values([]), - topic_provider=topic_provider or default_preferred_values([]) + topic_provider=topic_provider or default_preferred_values([]), + ontology_constraints=ontology_constraints ) def _get_json(self, node, llm, inference_parameters): @@ -65,7 +67,7 @@ def _get_json(self, node, llm, inference_parameters): else node.text ) messages = llm._get_messages( - PromptTemplate(self.prompt_template), + PromptTemplate(with_ontology_constraints(self.prompt_template, self.ontology_constraints)), text=text, preferred_entity_classifications=format_list(classifications), preferred_topics=format_list(topics) @@ -83,7 +85,8 @@ def _run_non_batch_extractor(self, nodes): prompt_template=self.prompt_template, source_metadata_field=self.source_metadata_field, entity_classification_provider=self.entity_classification_provider, - topic_provider=self.topic_provider + topic_provider=self.topic_provider, + ontology_constraints=self.ontology_constraints ) extracted = extractor.extract(all_nodes) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/llm_proposition_extractor.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/llm_proposition_extractor.py index 1b17a60d..f743fde2 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/llm_proposition_extractor.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/llm_proposition_extractor.py @@ -9,7 +9,7 @@ from graphrag_toolkit.lexical_graph.config import GraphRAGConfig from graphrag_toolkit.lexical_graph.indexing.model import Propositions from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY -from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_PROPOSITIONS_PROMPT +from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_PROPOSITIONS_PROMPT, with_ontology_constraints from graphrag_toolkit.lexical_graph.indexing.extract.progress import run_jobs_with_progress from graphrag_toolkit.lexical_graph.utils.arg_utils import coalesce @@ -39,6 +39,9 @@ class LLMPropositionExtractor(BaseExtractor): source_metadata_field (Optional[str]): The metadata field in the input nodes from which propositions are extracted. If not specified, the node text is used instead. + ontology_constraints (str): Rendered ontology entity-type block, + composed into the prompt template at render time. Empty by + default, which leaves the prompt exactly as it is today. """ llm: Optional[LLMCache] = Field( description='The LLM to use for extraction' @@ -52,6 +55,11 @@ class LLMPropositionExtractor(BaseExtractor): description='Metadata field from which to extract propositions' ) + ontology_constraints:str = Field( + default='', + description='Rendered ontology vocabulary block, composed into the prompt template at render time' + ) + @classmethod def class_name(cls) -> str: """ @@ -66,7 +74,8 @@ def __init__(self, llm:LLMCacheType=None, prompt_template=None, source_metadata_field=None, - num_workers:Optional[int]=None): + num_workers:Optional[int]=None, + ontology_constraints:str=''): """ Initializes the class with configuration options for processing language model outputs. @@ -80,6 +89,10 @@ def __init__(self, source_metadata_field: Field name key to store or retrieve associated metadata from source data. num_workers: Number of worker threads to use for processing tasks. + ontology_constraints: Rendered ontology entity-type block, composed + into the prompt template at render time. Empty by default, + which leaves the template - and therefore the LLM cache key - + exactly as it is today. """ num_workers = coalesce(num_workers, GraphRAGConfig.extraction_num_threads_per_worker) @@ -91,7 +104,8 @@ def __init__(self, ), prompt_template=prompt_template or EXTRACT_PROPOSITIONS_PROMPT, source_metadata_field=source_metadata_field, - num_workers=num_workers + num_workers=num_workers, + ontology_constraints=ontology_constraints ) logger.debug(f'Prompt template: {self.prompt_template}') @@ -200,7 +214,9 @@ async def _extract_propositions(self, text, source_info): """ def blocking_llm_call(): return self.llm.predict( - PromptTemplate(template=self.prompt_template), + PromptTemplate( + template=with_ontology_constraints(self.prompt_template, self.ontology_constraints) + ), text=text, source_info=source_info, exclude_cache_keys=['source_info'] diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/__init__.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/__init__.py new file mode 100644 index 00000000..d3b6e692 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/__init__.py @@ -0,0 +1,38 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Importing this package imports rdflib, via ontology.py. Modules that run +# inside extraction workers - anything pickled across the spawn boundary - +# should import from .ontology_index directly instead, which is plain data and +# rdflib-free. + +from .datatype_utils import coerce_literal, validate_literal_against_xsd +from .naming import camel_to_upper_snake, resolution_key, title_case_with_spaces +from .ontology import Ontology, OntologyLoadError +from .ontology_config import ( + ONTOLOGY_AUTHORITY_LEVELS, + OntologyConfig, + OntologySource, + OntologyType, + ResolvedDimensions, + OntologyAuthority, + TypedProperties, + VocabularyFormat, + to_ontology_config, +) +from .ontology_filter import FilterCounters, OntologyFilter, authored_name +from .ontology_index import ( + OWL_THING, + XSD_NAMESPACE, + DatatypeProperty, + ObjectProperty, + OntologyClass, + OntologyIndex, +) +from .prompt_constraint import ( + PROMPT_CONSTRAINT_LEVELS, + PROSE_VOCABULARY, + TURTLE_VOCABULARY, + VOCABULARY_FORMATS, + rendered_class_names, +) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/datatype_utils.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/datatype_utils.py new file mode 100644 index 00000000..9055a91d --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/datatype_utils.py @@ -0,0 +1,368 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic coercion of an extracted literal to its declared XSD type. + +The division of labour on attributes mirrors the rest of the feature: the LLM's +job is to emit a resolvable attribute *name*, and every decision about the +*value* is made here, deterministically. Nothing in this module +asks a model anything. + +`coerce_literal` is the single place that answers "is this string a value of that +declared type, and if so what is it natively?". Two callers, both later phases: +`enforce_datatypes` drops a subject-predicate-complement fact when the answer is +`None`, and `typed_properties` writes the native value when it +is not. + +Three properties the callers depend on: + +* **Every non-`None` return is JSON-serializable.** The value lands in + `node.metadata` and then in Cypher parameters, so a `Decimal`, a `date`, or a + float `nan` would fail somewhere far from here. Integers come back as `int`, + decimals as a *finite* `float`, booleans as `bool`, and dates and times as ISO + strings. +* **The type families agree with what the prompt promised.** `prompt_constraint` + renders `xsd:gYear` to the model as `integer` and `xsd:double` as + `decimal number`; this module coerces them the same way. If the two disagreed, + a value reported in exactly the form the prompt asked for could still fail + coercion. +* **No `rdflib` import.** XSD IRIs are plain strings here, so this module is safe + on both sides of the extraction spawn boundary and can be imported by + `ontology_filter.py`. + +**A wrong value is worse than no value, so parsing does not guess.** This is the +one judgement running through every table below. `enforce_datatypes` treats a +successful coercion as conforming, so a literal parsed to the wrong number is +stored and trusted, whereas `None` is visible and handled. Real extracted output +for `xsd:double` in the replay corpus includes `'41,500,000 dollars'` and +`'2,400,000,000 US dollars'`, and it is tempting to strip the trailing words. +That is refused: nothing distinguishes a harmless unit (`'dollars'`) from a +magnitude (`'nine hundred million'`, `'$1.2bn'`) without interpreting it, and +interpreting it is a conversion, not a parse. Getting `900` for nine hundred +million is a six-order-of-magnitude error that no downstream check would catch. +So a literal must be numeric *in its entirety*, and the unit-bearing values +return `None` and are dropped or skipped by the caller. + +The same reasoning excludes all-numeric slash dates: `'01/03/1994'` is the first +of March or the third of January depending on where the text came from, and there +is no way to tell. Non-ISO dates are accepted only where the month is named. +""" + +import re + +from datetime import date, datetime +from functools import partial +from typing import Any, Optional, Tuple + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_index import ( + XSD_NAMESPACE, +) + +# Integer types, with the bounds their declaration promises. The bounds are the +# reason to enumerate rather than to treat every integer type alike: a property +# declared `xsd:nonNegativeInteger` and given `'-5'` has a value outside its +# declared type, and `enforce_datatypes` exists to catch exactly that. `None` +# means unbounded on that side. +# +# `gYear` is here, not with the dates, because that is how the prompt renders it +# (`_TYPE_NAMES` in `prompt_constraint.py` maps it to `integer`) and a value the +# model reported as the prompt asked must coerce. +_INTEGER_BOUNDS = { + 'integer': (None, None), + 'long': (-(2 ** 63), 2 ** 63 - 1), + 'int': (-(2 ** 31), 2 ** 31 - 1), + 'short': (-32768, 32767), + 'byte': (-128, 127), + 'nonNegativeInteger': (0, None), + 'positiveInteger': (1, None), + 'nonPositiveInteger': (None, 0), + 'negativeInteger': (None, -1), + 'unsignedLong': (0, 2 ** 64 - 1), + 'unsignedInt': (0, 2 ** 32 - 1), + 'unsignedShort': (0, 65535), + 'unsignedByte': (0, 255), + 'gYear': (None, None), +} + +_DECIMAL_TYPES = frozenset({'decimal', 'double', 'float'}) + +# XSD's lexical space for boolean is exactly these four, and case folding is the +# only latitude taken. `'yes'` and `'no'` are deliberately absent: the prompt +# tells the model `true/false`, the corpus shows it complying, and inventing +# synonyms here would be guessing at a vocabulary nobody was offered. +_TRUE_LEXICAL = frozenset({'true', '1'}) +_FALSE_LEXICAL = frozenset({'false', '0'}) + +# Shape first, calendar second. A regex alone accepts '2015-02-29'; `strptime` +# is what rejects it. Both are needed - the regex because `strptime('%Y-%m-%d')` +# would otherwise accept '15-2-9' and silently produce year 15. +_ISO_DATE = re.compile(r'^-?\d{4}-\d{2}-\d{2}$') +_ISO_TIME = re.compile(r'^\d{2}:\d{2}(:\d{2}(\.\d+)?)?$') + +# Non-ISO dates, accepted only where the month is named and the reading is +# therefore unambiguous. 'March 3rd 2020' is real output from the replay corpus. +_NAMED_MONTH_FORMATS = ( + '%d %B %Y', # 3 March 2020 + '%d %b %Y', # 3 Mar 2020 + '%B %d %Y', # March 3 2020 + '%b %d %Y', # Mar 3 2020 + '%B %d, %Y', # March 3, 2020 + '%b %d, %Y', # Mar 3, 2020 + '%Y %B %d', # 2020 March 3 +) + +_ORDINAL_SUFFIX = re.compile(r'(?<=\d)(st|nd|rd|th)\b', re.IGNORECASE) + +# A trailing timezone is part of several XSD lexical spaces and carries nothing +# this module reports, so it is removed before parsing rather than rejected. +_TRAILING_TIMEZONE = re.compile(r'(Z|[+-]\d{2}:\d{2})$') + +# Digit grouping, and only in valid grouping positions. Matching the whole +# number is what keeps a European decimal comma out: '1994,5' does not match, so +# it is refused rather than silently read as 19945. +_GROUPED_NUMBER = re.compile(r'^([+-]?)(\d{1,3}(?:,\d{3})+)(\.\d+)?$') + +# A fraction that adds nothing, so '1994.0' is an integer but '1994.5' is not. +_ZERO_FRACTION = re.compile(r'\.0*$') + +_INTEGER_LEXICAL = re.compile(r'^[+-]?\d+$') + +# `xsd:anyURI` is almost unconstrained in the standard, so the check here is +# narrow on purpose: it rejects the failure actually seen from a model, which is +# a sentence where a URI was asked for. Internal whitespace is the signal. +_ANY_URI = re.compile(r'^(?:[A-Za-z][A-Za-z0-9+.\-]*:|[/#?])\S*$|^\S+\.\S+\S*$') + +def _local_name(xsd_iri:Optional[str]) -> Optional[str]: + """The XSD local name, or None when the IRI is not an XSD datatype. + + `Ontology` rejects a datatype property whose `rdfs:range` is outside the XSD + namespace at load time, so a non-XSD IRI here means a caller built an index + some other way. There is no defensible coercion against a type this module + knows nothing about, so it declines rather than guessing. + """ + if not xsd_iri or not xsd_iri.startswith(XSD_NAMESPACE): + return None + return xsd_iri[len(XSD_NAMESPACE):] + +def _degrouped(text:str) -> Optional[str]: + """`'2,400,000'` -> `'2400000'`; `'1994,5'` -> None; `'1994'` unchanged.""" + match = _GROUPED_NUMBER.match(text) + if match: + return f'{match.group(1)}{match.group(2).replace(",", "")}{match.group(3) or ""}' + return None if ',' in text else text + +def _coerce_integer(text:str, bounds:Tuple[Optional[int], Optional[int]]) -> Optional[int]: + """An integer, tolerating digit grouping and a zero fraction.""" + degrouped = _degrouped(text) + if degrouped is None: + return None + + # '1994.0' is the same integer as '1994'; '1994.5' is not an integer at all. + if '.' in degrouped: + if not _ZERO_FRACTION.search(degrouped): + return None + degrouped = degrouped[:degrouped.index('.')] + + if not _INTEGER_LEXICAL.match(degrouped): + return None + + value = int(degrouped) + (low, high) = bounds + if (low is not None and value < low) or (high is not None and value > high): + return None + + return value + +def _coerce_decimal(text:str) -> Optional[float]: + """A finite float. `nan` and `inf` are refused - they are not JSON.""" + degrouped = _degrouped(text) + if degrouped is None: + return None + + try: + value = float(degrouped) + except ValueError: + return None + + # `float('nan')` and `float('inf')` both succeed, and `json.dumps` emits + # `NaN` / `Infinity` for them, which no JSON parser is required to accept. + # Neither is a value any ontology means by `xsd:double`. + if value != value or value in (float('inf'), float('-inf')): + return None + + return value + +def _coerce_boolean(text:str) -> Optional[bool]: + """`True`, `False`, or None. Note the caller must test `is not None`.""" + folded = text.lower() + if folded in _TRUE_LEXICAL: + return True + if folded in _FALSE_LEXICAL: + return False + return None + +def _coerce_date(text:str) -> Optional[str]: + """An ISO date string, from an ISO or a named-month input.""" + stripped = _TRAILING_TIMEZONE.sub('', text).strip() + + if _ISO_DATE.match(stripped): + try: + return date.fromisoformat(stripped.lstrip('-')).isoformat() + except ValueError: + # Shape was right, calendar was not - '2015-02-29'. + return None + + # 'March 3rd 2020' -> 'March 3 2020'. Done before the format loop because no + # `strptime` directive matches an ordinal suffix. + plain = re.sub(r'\s+', ' ', _ORDINAL_SUFFIX.sub('', stripped)).strip() + + for fmt in _NAMED_MONTH_FORMATS: + try: + return datetime.strptime(plain, fmt).date().isoformat() + except ValueError: + continue + + return None + +def _coerce_datetime(text:str) -> Optional[str]: + """An ISO datetime string. Only the ISO lexical form is accepted.""" + stripped = _TRAILING_TIMEZONE.sub('', text).strip() + try: + return datetime.fromisoformat(stripped).isoformat() + except ValueError: + return None + +def _coerce_time(text:str) -> Optional[str]: + """An ISO time-of-day string.""" + stripped = _TRAILING_TIMEZONE.sub('', text).strip() + if not _ISO_TIME.match(stripped): + return None + try: + return datetime.strptime( + stripped, '%H:%M:%S' if stripped.count(':') == 2 else '%H:%M' + ).time().isoformat() + except ValueError: + return None + +def _coerce_any_uri(text:str) -> Optional[str]: + """The URI, or None when the value is prose rather than a reference.""" + return text if _ANY_URI.match(text) else None + +def _coerce_text(text:str) -> str: + """The trimmed text, for the types whose value space *is* text.""" + return text + +# The string family, where returning the text unchanged is the implementation +# rather than a give-up. Kept apart from the fallback below so the two cases are +# distinguishable: `enforce_datatypes` warns about a declaration it could not +# honour, and `xsd:string` is honoured. +_TEXT_TYPES = frozenset({ + 'string', 'normalizedString', 'token', 'language', + 'Name', 'NCName', 'NMTOKEN', 'ID', 'IDREF', 'ENTITY', +}) + +# One dispatch table, so membership and behaviour cannot disagree. +# +# This is deliberately a table rather than a chain of `if local_name == ...`: +# `validates_datatype` below is derived from its keys, and `enforce_datatypes` +# uses that to decide whether a coercion returning text was a check or a silent +# pass. A separately-maintained list of "implemented types" +# would drift the first time a branch was added, and drift silently, since the +# symptom is a missing warning. +_COERCERS = { + **{name: partial(_coerce_integer, bounds=bounds) for (name, bounds) in _INTEGER_BOUNDS.items()}, + **{name: _coerce_decimal for name in _DECIMAL_TYPES}, + **{name: _coerce_text for name in _TEXT_TYPES}, + 'boolean': _coerce_boolean, + 'date': _coerce_date, + 'dateTime': _coerce_datetime, + 'time': _coerce_time, + 'anyURI': _coerce_any_uri, +} + +def validates_datatype(xsd_iri:Optional[str]) -> bool: + """Whether coercion actually checks values of this declared type. + + False means `coerce_literal` will return the trimmed literal for any input, + so a `True` from `validate_literal_against_xsd` carries no information. The + caller that enforces datatypes is expected to say so out loud rather than + report a validation that did not happen. + + Args: + xsd_iri: A declared `rdfs:range`, as a plain string. + + Returns: + True for every XSD type with an implementation, including the string + family. False for an XSD type with no implementation (`xsd:hexBinary`, + `xsd:duration`, `xsd:gMonthDay`), and False for a non-XSD IRI - which + `coerce_literal` refuses outright, so nothing is stored unvalidated and + there is nothing to warn about. + """ + local_name = _local_name(xsd_iri) + return local_name is not None and local_name in _COERCERS + +def coerce_literal(literal:Optional[str], xsd_iri:Optional[str]) -> Any: + """Parse an extracted literal into a native value of its declared XSD type. + + Best-effort on the *lexical* form and strict on meaning: digit grouping, + surrounding whitespace, a zero fraction, an ordinal suffix and a named month + are all tolerated, while anything needing interpretation is refused. See the + module docstring for why refusing beats guessing. + + Args: + literal: The value the model emitted, typically `Fact.complement.value`. + xsd_iri: The declared `rdfs:range` of the resolved datatype property, as + a plain string - `DatatypeProperty.datatype`. + + Returns: + An `int` for the integer types, a finite `float` for the decimal types, a + `bool` for `xsd:boolean`, an ISO `str` for `xsd:date`, `xsd:dateTime` and + `xsd:time`, and the trimmed text for `xsd:string` and any other XSD type, + matching the `text` the prompt renders for those. `None` when the literal + does not parse, when it is empty, or when `xsd_iri` is not an XSD type. + + **Test the result with `is not None`.** `False`, `0` and `0.0` are all + successful coercions and all falsy. + """ + if literal is None or xsd_iri is None: + return None + + text = literal.strip() + if not text: + return None + + local_name = _local_name(xsd_iri) + if local_name is None: + return None + + coercer = _COERCERS.get(local_name) + if coercer is not None: + return coercer(text) + + # An XSD type with no implementation. Treated as text, which is also what + # the prompt rendered for it (`DEFAULT_TYPE_NAME`), and returned trimmed + # because the trimmed form is what would be stored anyway. + # + # Note that this is *not* a validation: any input at all is accepted. Callers + # that need to distinguish this from a real check ask `validates_datatype` + # first - `enforce_datatypes` does, and warns. + return text + +def validate_literal_against_xsd(literal:Optional[str], xsd_iri:Optional[str]) -> bool: + """Whether `literal` is a value of the type `xsd_iri` declares. + + One implementation, not two: validity *is* coercibility, so the two can + never disagree about a value. Kept as a named function because + `enforce_datatypes` reads as a validity check at its call site even though it + also wants the coerced value. + + Args: + literal: The value the model emitted. + xsd_iri: The declared `rdfs:range`, as a plain string. + + Returns: + True when the literal parses. `'false'` against `xsd:boolean` is valid + and returns True, which is the reason this is `is not None` rather than a + truth test. + """ + return coerce_literal(literal, xsd_iri) is not None diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/naming.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/naming.py new file mode 100644 index 00000000..30e8f0ff --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/naming.py @@ -0,0 +1,118 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The naming-convention contract between the ontology and the extraction path. + +Three separate jobs, which must agree with each other and with +`indexing/utils/topic_utils.py`: + +* **Render** - `camel_to_upper_snake` for properties, `title_case_with_spaces` + for classes. Chosen so the name survives the response parser: `format_value` + turns `_` into a space and `format_classification` then applies `.title()`. +* **Resolve** - `resolution_key` folds a name to a convention-free form, so a + name the LLM emitted in any of the conventions in play still finds its + declared term. +* **Canonicalize** - not here. The canonical stored spelling is the authored + name, taken verbatim from `rdfs:label` or the IRI local name with no + transformation at all, so it needs no helper and must not be folded into one + of the above. + +Rendering and canonicalization are allowed to disagree, and for a class +authored `:SportsTeam` they do: it renders as `Sports Team` and stores as +`SportsTeam`. Collapsing them into a single "normalize" helper is what produced +the bug this module exists to prevent - rendering `SportsTeam` directly comes +back from the parser as `Sportsteam`, which resolves to nothing. + +The guarding invariant, asserted over every term in every test ontology: + + resolution_key(parse_transform(render(term))) == resolution_key(local_name) + +`render` and `resolution_key` therefore share one camel-case boundary rule: +split only where an uppercase letter follows a **lowercase** letter. The +narrower rule matters. Splitting after any non-uppercase character would split +`Company2X`, which is what `.title()` makes of `Company2x`, and the invariant +would fail on a name that only differs from a working one by a digit. + +This module is plain string handling with no `rdflib` import, so it is safe on +both sides of the extraction process boundary. +""" + +import re + +# Split where an uppercase letter follows a lowercase letter, and nowhere else. +# 'worksFor' -> 'works For'; 'HTTPServer' -> 'HTTPServer' (unsplit, but folded +# to 'httpserver' consistently at both ends of the round trip). +_CAMEL_BOUNDARY = re.compile(r'(?<=[a-z])(?=[A-Z])') + +_WHITESPACE_RUN = re.compile(r'\s+') + +def _split_words(name:str) -> str: + """Insert a space at every camel-case boundary and every underscore.""" + return _CAMEL_BOUNDARY.sub(' ', name).replace('_', ' ') + +def resolution_key(name:str) -> str: + """Fold a name to its convention-free form, for index lookup. + + Splits camelCase, replaces underscores with spaces, lowercases, and + collapses whitespace: + + 'worksFor' -> 'works for' + 'WORKS FOR' -> 'works for' <- what the parser actually hands us + 'WORKS_FOR' -> 'works for' + 'Sports Team' -> 'sports team' + 'SportsTeam' -> 'sports team' + + Args: + name: A name from anywhere - an IRI local name, an `rdfs:label`, a + `skos:altLabel`, or a predicate the LLM emitted. `None` and the + empty string fold to the empty string, which is never indexed. + + Returns: + The folded key. + """ + if not name: + return '' + return _WHITESPACE_RUN.sub(' ', _split_words(name)).strip().lower() + +def camel_to_upper_snake(name:str) -> str: + """Render a property name for the prompt, as `UPPER_SNAKE`. + + 'worksFor' -> 'WORKS_FOR' + 'WORKS_FOR' -> 'WORKS_FOR' + 'founded year'-> 'FOUNDED_YEAR' + + `EXTRACT_TOPICS_PROMPT` already instructs the model that relationship names + are "all uppercase, with underscores instead of spaces", so rendering + `:worksFor` verbatim would put two conflicting instructions in one prompt. + + Args: + name: A property's local name, label, or alias. + + Returns: + The rendered name, containing no lowercase letter. + """ + if not name: + return '' + words = _WHITESPACE_RUN.sub(' ', _split_words(name)).strip().split(' ') + return '_'.join(words).upper() + +def title_case_with_spaces(name:str) -> str: + """Render a class name for the prompt, as `Title Case With Spaces`. + + 'SportsTeam' -> 'Sports Team' + 'SPORTS_TEAM' -> 'Sports Team' + 'Sports Team' -> 'Sports Team' + + `.title()` is applied here rather than left to the parser, which makes the + rendered form a fixed point of `format_classification` - the name comes + back from the parser exactly as it went out. + + Args: + name: A class's local name, label, or alias. + + Returns: + The rendered name, containing no underscore. + """ + if not name: + return '' + return _WHITESPACE_RUN.sub(' ', _split_words(name)).strip().title() diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology.py new file mode 100644 index 00000000..3c143e9b --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology.py @@ -0,0 +1,585 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load an OWL/Turtle ontology and validate it at load time. + +Turtle is the only accepted file format. Other RDF serializations are supported +by parsing them into an `rdflib.Graph` with rdflib and passing that to +`Ontology.from_graph`. + +`Ontology` lives in the parent process only - it owns the `rdflib.Graph` and is +never pickled. It yields two artefacts, both computed once before any text is +seen: the prompt constraint string, and `OntologyIndex`, the plain-data read +index that does cross into extraction workers. + +Structural problems are raised as `OntologyLoadError` when the ontology is +loaded, not when a document is extracted, so that an authoring mistake does not +cost an extraction run. +""" + +import logging +from pathlib import Path +from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Union + +from rdflib import RDF, RDFS, BNode, Graph, Namespace + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_index import ( + OWL_THING, + XSD_NAMESPACE, + DatatypeProperty, + ObjectProperty, + OntologyClass, + OntologyIndex, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.prompt_constraint import ( + PROMPT_CONSTRAINT_LEVELS, + PROSE_VOCABULARY, + VOCABULARY_FORMATS, + format_as_prompt_constraint, + format_as_proposition_constraint, + format_turtle_vocabulary, + rendered_class_names, +) + +logger = logging.getLogger(__name__) + +OWL = Namespace('http://www.w3.org/2002/07/owl#') +SKOS = Namespace('http://www.w3.org/2004/02/skos/core#') + +TURTLE_SUFFIX = '.ttl' + +OntologyTerms = Tuple[ + Dict[str, OntologyClass], + Dict[str, ObjectProperty], + Dict[str, DatatypeProperty], +] + +class OntologyLoadError(Exception): + """Raised when an ontology cannot be loaded. + + Covers malformed Turtle, an unsupported file format, and the structural + rules checked at load time: `rdfs:subClassOf` cycles, dangling class + references in `rdfs:subClassOf` / `rdfs:domain` / `rdfs:range`, a property + declared as both `owl:ObjectProperty` and `owl:DatatypeProperty`, and a + datatype property whose `rdfs:range` is missing or is not an XSD datatype. + + No partial `Ontology` is returned when this is raised. + """ + +class Ontology: + """A parsed, validated OWL/Turtle ontology. + + Attributes: + graph (Graph): The canonical representation. Public and immutable by + convention - `graph.serialize()` is the supported way to write an + ontology back out, which is why there is no `serialize()` method + here. + namespace (str): The base IRI, detected from `@prefix :`, rdflib's + default namespace, the first `owl:Ontology` subject, or the + caller-supplied `base_iri`. Empty when none of those resolve; this + is not an error, because nothing in the feature requires it. + """ + + def __init__(self, graph:Graph, base_iri:Optional[str]=None): + """Adopt a parsed graph, index it, and validate it. + + Raises: + OntologyLoadError: If `graph` is not an `rdflib.Graph`, or if the + ontology violates one of the load-time structural rules. + """ + if not isinstance(graph, Graph): + raise OntologyLoadError( + f'Expected an rdflib.Graph, got {type(graph).__name__}.' + ) + + self.graph = graph + self.namespace = _extract_base_iri(graph, base_iri) or '' + + # Index first so that validation has plain data to read, then validate + # before __init__ returns - a caller never sees an unvalidated + # Ontology. Cycle detection raises from within the index build, since + # the ancestor closure is where a cycle shows up. + self._index = self._build_index() + self._validate() + + @classmethod + def from_turtle(cls, path:Union[str, Path]) -> 'Ontology': + """Load an ontology from a Turtle file. + + Args: + path: Path to a `.ttl` file. + + Raises: + OntologyLoadError: If the suffix is not `.ttl`, the file does not + exist, the Turtle does not parse, or the ontology is + structurally invalid. + """ + path = Path(path) + + if path.suffix.lower() != TURTLE_SUFFIX: + raise OntologyLoadError( + f'Unsupported ontology file format: {path.name}. Turtle ' + f'({TURTLE_SUFFIX}) is the only supported file format. Parse ' + f'other RDF serializations with rdflib and pass the resulting ' + f'graph to Ontology.from_graph().' + ) + + if not path.is_file(): + raise OntologyLoadError(f'Ontology file not found: {path}') + + graph = Graph() + try: + graph.parse(source=str(path), format='turtle') + except Exception as e: + raise OntologyLoadError(f'Failed to parse Turtle file {path}: {e}') from e + + return cls(graph) + + @classmethod + def from_turtle_string(cls, turtle:str, base_iri:Optional[str]=None) -> 'Ontology': + """Load an ontology from a Turtle document held in memory. + + Args: + turtle: The Turtle document. + base_iri: Optional base IRI, passed to rdflib as `publicID` to + resolve relative IRIs, and used as the lowest-priority fallback + when detecting `namespace`. + + Raises: + OntologyLoadError: If the Turtle does not parse, or the ontology is + structurally invalid. + """ + graph = Graph() + try: + graph.parse(data=turtle, format='turtle', publicID=base_iri) + except Exception as e: + raise OntologyLoadError(f'Failed to parse Turtle string: {e}') from e + + return cls(graph, base_iri=base_iri) + + @classmethod + def from_graph(cls, graph:Graph) -> 'Ontology': + """Adopt an already-parsed `rdflib.Graph` unchanged. + + The supported route for any RDF serialization other than Turtle: parse + it with rdflib, then pass the graph here. + + Raises: + OntologyLoadError: If `graph` is not an `rdflib.Graph`, or the + ontology is structurally invalid. + """ + return cls(graph) + + @classmethod + def load(cls, source:Union['Ontology', Graph, str, Path]) -> 'Ontology': + """Load an ontology from whatever form the caller has one in. + + The dispatching constructor, and the one configuration goes through: + `ExtractionConfig(ontology=...)` accepts a path because most users have + a file, and an `Ontology` because a user who built one already should not + pay to parse it twice. + + Args: + source: An `Ontology`, returned unchanged; an `rdflib.Graph`, adopted + by `from_graph`; or a `str`/`Path` to a `.ttl` file, parsed by + `from_turtle`. + + Raises: + OntologyLoadError: If `source` is a path with a suffix other than + `.ttl`, or is not one of the forms above. A Turtle document held + in a string is loaded by `from_turtle_string`, not here - a bare + `str` is read as a path, because guessing between the two from + its content would be the kind of ambiguity that surfaces as a + confusing parse error. + """ + if isinstance(source, Ontology): + return source + + if isinstance(source, Graph): + return cls.from_graph(source) + + if isinstance(source, (str, Path)): + return cls.from_turtle(source) + + raise OntologyLoadError( + f'Cannot load an ontology from {type(source).__name__}. Pass an ' + f'Ontology, an rdflib.Graph, or a path to a Turtle ' + f'({TURTLE_SUFFIX}) file. A Turtle document held in a string is ' + f'loaded with Ontology.from_turtle_string().' + ) + + def index(self) -> OntologyIndex: + """Return the read index. Built once at construction.""" + return self._index + + def format_as_prompt_constraint(self, level:str, vocabulary_format:str=PROSE_VOCABULARY) -> str: + """Render the ontology as the constraint block for the prompt. + + The block names the vocabulary the model should reach for and how much + authority it has; it is composed into the extraction prompt rather than + replacing any part of it. Rendered once at pipeline-configuration time, + in this process. + + Two formats. `'prose'` renders the three generated sections and is the + default and the measured configuration. `'turtle'` shows this ontology's + own source instead, on the argument that the prose sections paraphrase + away the RDFS entailments - `rdfs:subClassOf` survives only as + indentation - and that a model has read a great deal of real RDFS. + + Serializing is done here rather than in the renderer because the renderer + works from `OntologyIndex` and imports no `rdflib`, which an import-graph + test enforces. `format='turtle'` round-trips the canonical graph, so the + text the model sees is the ontology as loaded, not as authored - prefix + bindings and triple order may differ from the input file. + + Args: + level: `'off'`, `'align'` or `'strict'`. + vocabulary_format: `'prose'` (default) or `'turtle'`. + + Returns: + The block, with no trailing newline, or `''` for `'off'` and for an + ontology that declares no terms. + + Raises: + ValueError: If `level` is not a known authority level, or + `vocabulary_format` is not a known format. + """ + if vocabulary_format not in VOCABULARY_FORMATS: + raise ValueError( + f'Unknown ontology vocabulary format: {vocabulary_format!r}. ' + f'Expected one of {", ".join(VOCABULARY_FORMATS)}.' + ) + + if vocabulary_format == PROSE_VOCABULARY: + return format_as_prompt_constraint(self._index, level) + + # Validate the level on the same terms the prose path does, and before + # paying to serialize, so a typo is reported identically either way. + if level not in PROMPT_CONSTRAINT_LEVELS: + raise ValueError( + f'Unknown ontology authority level: {level!r}. ' + f'Expected one of {", ".join(PROMPT_CONSTRAINT_LEVELS)}.' + ) + + if level == 'off': + return '' + + # An ontology declaring no terms renders nothing in either format. Tested + # against the index rather than the graph because a graph can hold + # prefix bindings and an `owl:Ontology` header and still declare no + # vocabulary, and a block with a header and no terms is worse than none. + if not (self._index.classes or self._index.object_properties + or self._index.datatype_properties): + return '' + + return format_turtle_vocabulary(self.graph.serialize(format='turtle'), level) + + def format_as_proposition_constraint(self, level:str) -> str: + """Render the entity types alone, for the propositions prompt. + + The propositions stage classifies the entities it names and extracts + nothing else, so it is steered with the class names and not with the full + vocabulary. + + Args: + level: `'off'`, `'align'` or `'strict'`. + + Returns: + The hint, with no trailing newline, or `''` for `'off'` and for an + ontology that declares no classes. + + Raises: + ValueError: If `level` is not a known authority level. + """ + return format_as_proposition_constraint(self._index, level) + + def class_names(self) -> List[str]: + """The rendered class names, for seeding `preferred_entity_classifications`. + + Independent of `ontology_authority`, unlike the two `format_as_*` methods: the + level decides how much the prompt says *about* the vocabulary, not what + the vocabulary is. A caller that wants nothing seeded reads the level + itself and does not call this. + + Returns: + The names as the prompt renders them, sorted; empty when the ontology + declares no classes. + """ + return rendered_class_names(self._index) + + def _build_index(self) -> OntologyIndex: + """Read the graph into plain data. + + Raises: + OntologyLoadError: If `rdfs:subClassOf` contains a cycle. + """ + classes, object_properties, datatype_properties = _collect_terms(self.graph) + + return OntologyIndex( + classes=classes, + object_properties=object_properties, + datatype_properties=datatype_properties, + ) + + def _validate(self) -> None: + """Run the load-time structural checks over the built index. + + Raises: + OntologyLoadError: On the first violation found, naming the + offending term. + """ + _validate_terms( + self._index.classes, + self._index.object_properties, + self._index.datatype_properties, + ) + +def _local_name_of(iri:str) -> str: + """Return an IRI's last segment, split on `#` then `/`.""" + if '#' in iri: + return iri.rsplit('#', 1)[-1] + if '/' in iri: + return iri.rsplit('/', 1)[-1] + return iri + +def _values(graph:Graph, subject, predicate) -> List[str]: + """Return the sorted, deduplicated non-blank objects of a predicate. + + Sorted because rdflib does not guarantee an order over a triple pattern, + and every artefact derived from the ontology has to be deterministic. + """ + return sorted( + {str(o) for o in graph.objects(subject, predicate) if not isinstance(o, BNode)} + ) + +def _first_value(graph:Graph, subject, predicate) -> Optional[str]: + """Return the first of `_values`, or None when the predicate is absent.""" + values = _values(graph, subject, predicate) + return values[0] if values else None + +def _class_reference( + graph:Graph, subject, predicate, term_iri:str +) -> Optional[str]: + """Return a single class IRI for an `rdfs:domain` / `rdfs:range` slot. + + `None` means the slot is unconstrained - either undeclared or declared as + `owl:Thing`, which are the same thing as far as the filter is concerned. + Multiple declarations are narrowed to one, deterministically and with a + warning, because a single domain and range per property is what the + prompt rendering and the domain/range check are defined over. + """ + values = [value for value in _values(graph, subject, predicate) if value != OWL_THING] + + if not values: + return None + + if len(values) > 1: + logger.warning( + 'Property %s declares %d %s values (%s); using %s. One domain and ' + 'one range per property is supported.', + term_iri, len(values), _local_name_of(str(predicate)), + ', '.join(values), values[0], + ) + + return values[0] + +def _collect_terms(graph:Graph) -> OntologyTerms: + """Read classes, object properties and datatype properties out of a graph. + + Blank-node subjects are skipped: anonymous class axioms such as an + `owl:Restriction` body are not terms a user can name in a domain or range, + and are not vocabulary the model can be asked to emit. + + Raises: + OntologyLoadError: If `rdfs:subClassOf` contains a cycle. + """ + classes:Dict[str, OntologyClass] = {} + parents_by_iri:Dict[str, List[str]] = {} + + for subject in graph.subjects(RDF.type, OWL.Class): + if isinstance(subject, BNode): + continue + iri = str(subject) + parents_by_iri[iri] = _values(graph, subject, RDFS.subClassOf) + + ancestors_by_iri = _compute_subclass_closure(parents_by_iri) + + for subject in graph.subjects(RDF.type, OWL.Class): + if isinstance(subject, BNode): + continue + iri = str(subject) + classes[iri] = OntologyClass( + iri=iri, + local_name=_local_name_of(iri), + label=_first_value(graph, subject, RDFS.label), + aliases=_values(graph, subject, SKOS.altLabel), + parents=parents_by_iri[iri], + ancestors=ancestors_by_iri[iri], + description=_first_value(graph, subject, RDFS.comment), + ) + + object_properties:Dict[str, ObjectProperty] = {} + + for subject in graph.subjects(RDF.type, OWL.ObjectProperty): + if isinstance(subject, BNode): + continue + iri = str(subject) + object_properties[iri] = ObjectProperty( + iri=iri, + local_name=_local_name_of(iri), + label=_first_value(graph, subject, RDFS.label), + aliases=_values(graph, subject, SKOS.altLabel), + domain=_class_reference(graph, subject, RDFS.domain, iri), + range=_class_reference(graph, subject, RDFS.range, iri), + description=_first_value(graph, subject, RDFS.comment), + ) + + datatype_properties:Dict[str, DatatypeProperty] = {} + + for subject in graph.subjects(RDF.type, OWL.DatatypeProperty): + if isinstance(subject, BNode): + continue + iri = str(subject) + # An absent or non-XSD range is rejected by _validate_terms; store the + # empty string here so the model invariant (one datatype per property) + # holds even before validation runs. + ranges = _values(graph, subject, RDFS.range) + datatype_properties[iri] = DatatypeProperty( + iri=iri, + local_name=_local_name_of(iri), + label=_first_value(graph, subject, RDFS.label), + aliases=_values(graph, subject, SKOS.altLabel), + domain=_class_reference(graph, subject, RDFS.domain, iri), + datatype=ranges[0] if ranges else '', + description=_first_value(graph, subject, RDFS.comment), + ) + + return classes, object_properties, datatype_properties + +def _compute_subclass_closure( + parents_by_iri:Dict[str, List[str]] +) -> Dict[str, FrozenSet[str]]: + """Compute the reflexive transitive `rdfs:subClassOf` closure. + + A memoized depth-first walk over the parent edges. Parents that are not + themselves declared classes are skipped here so that `ancestors` only ever + holds declared class IRIs - `_validate_terms` is the single place that + raises on a dangling reference. + + Raises: + OntologyLoadError: If a cycle is reached, naming the class at which it + closed. + """ + memo:Dict[str, FrozenSet[str]] = {} + path:List[str] = [] + + def compute(iri:str) -> FrozenSet[str]: + if iri in memo: + return memo[iri] + if iri in path: + cycle = path[path.index(iri):] + [iri] + raise OntologyLoadError( + f'rdfs:subClassOf cycle detected at {iri}: ' + f'{" -> ".join(cycle)}.' + ) + path.append(iri) + + ancestors:Set[str] = {iri} + for parent_iri in parents_by_iri.get(iri, []): + if parent_iri in parents_by_iri: + ancestors |= compute(parent_iri) + + path.pop() + memo[iri] = frozenset(ancestors) + return memo[iri] + + return {iri: compute(iri) for iri in sorted(parents_by_iri)} + +def _validate_terms( + classes:Dict[str, OntologyClass], + object_properties:Dict[str, ObjectProperty], + datatype_properties:Dict[str, DatatypeProperty], +) -> None: + """Check the structural rules that the closure walk does not cover. + + Cheapest check first, and each raises on the first offending term so that a + broken ontology fails fast with one actionable message. Iteration is sorted + so the term named in the message is deterministic. + + Raises: + OntologyLoadError: On the first violation found. + """ + both = sorted(set(object_properties) & set(datatype_properties)) + if both: + raise OntologyLoadError( + f'Property {both[0]} is declared as both an owl:ObjectProperty and ' + f'an owl:DatatypeProperty. A property must be one or the other.' + ) + + for iri, datatype_property in sorted(datatype_properties.items()): + if not datatype_property.datatype: + raise OntologyLoadError( + f'DatatypeProperty {iri} has no rdfs:range. Every ' + f'owl:DatatypeProperty must declare one XSD datatype as its ' + f'range.' + ) + if not datatype_property.datatype.startswith(XSD_NAMESPACE): + raise OntologyLoadError( + f'DatatypeProperty {iri} has rdfs:range ' + f'{datatype_property.datatype}, which is not an XSD datatype. ' + f'Only ranges in {XSD_NAMESPACE} are supported.' + ) + + for iri, ontology_class in sorted(classes.items()): + for parent_iri in ontology_class.parents: + if parent_iri not in classes: + raise OntologyLoadError( + f'Class {iri} has a dangling rdfs:subClassOf reference to ' + f'{parent_iri}, which is not declared as an owl:Class in ' + f'this ontology.' + ) + + for iri, object_property in sorted(object_properties.items()): + for slot, class_iri in ( + ('rdfs:domain', object_property.domain), + ('rdfs:range', object_property.range), + ): + if class_iri is not None and class_iri not in classes: + raise OntologyLoadError( + f'ObjectProperty {iri} has a dangling {slot} reference to ' + f'{class_iri}, which is not declared as an owl:Class in ' + f'this ontology.' + ) + + for iri, datatype_property in sorted(datatype_properties.items()): + class_iri = datatype_property.domain + if class_iri is not None and class_iri not in classes: + raise OntologyLoadError( + f'DatatypeProperty {iri} has a dangling rdfs:domain reference ' + f'to {class_iri}, which is not declared as an owl:Class in ' + f'this ontology.' + ) + +def _extract_base_iri(graph:Graph, base_iri:Optional[str]=None) -> Optional[str]: + """Detect the ontology's base IRI. + + In priority order: the Turtle `@prefix : <...>` declaration, which rdflib + surfaces as the empty prefix; rdflib's parser-level default namespace; the + first `owl:Ontology` subject; and finally the caller-supplied `base_iri`, + last so that an in-file declaration always wins. + + Returns None when none of those resolve. + """ + for prefix, namespace in graph.namespaces(): + if prefix == '' and str(namespace): + return str(namespace) + + default_namespace = getattr(graph, 'default_namespace', None) + if default_namespace and str(default_namespace): + return str(default_namespace) + + for subject in sorted(str(s) for s in graph.subjects(RDF.type, OWL.Ontology)): + if subject: + return subject + + return base_iri or None diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_config.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_config.py new file mode 100644 index 00000000..6bf2164b --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_config.py @@ -0,0 +1,385 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""How much authority the ontology has, expressed as one level. + +`ontology_authority` is the single knob. It reads as inclusion authority - how much say +the ontology has over what ends up in the graph - and it resolves to the six +dimensions the deterministic layer acts on: + +| | `off` | `align` | `strict` | +|---|---|---|---| +| `normalize_names` | False | True | True | +| `drop_type_restatements` | False | True | True | +| `enforce_entity_types` | False | False | True | +| `enforce_relationship_types` | False | False | True | +| `enforce_domain_range` | False | False | True | +| `enforce_datatypes` | False | False | True | + +Three levels, and each one owns a distinct mechanism: `off` says nothing, `align` +names what the ontology declares, `strict` also decides what survives. + +`drop_type_restatements` is the one row that does not fit that summary, and the +exception is deliberate. It is a dropping gate that is on at `align`, which reads +as a contradiction of "keep what the ontology does not declare" - so it is worth +being precise about what it removes: facts whose predicate is ontology language +(`rdf:type`, `rdfs:subClassOf`), and facts that restate the subject's own +classification under a type-asserting predicate (`Meridian Freight [Company] +|TYPE| Company`). Neither is content the level promised to keep. The second is a +duplicate of the entity's own label line, and dropping it is information- +preserving in the strict sense that nothing in the graph changes except the +absence of a redundant edge. `align` still keeps every fact that says something +the ontology did not anticipate; see `OntologyFilter._carries_no_domain_fact`. + +There was a fourth level between `align` and `strict`, `guide`, whose only +mechanism was extra prompt wording - flag-identical to `align`, so nothing in +this table could distinguish it. It was measured against a same-configuration +control and removed. + +`strict` was deliberately absent from `OntologyAuthority` until the enforcement +existed, because accepting it before then would have let a user ask for +exclusion and quietly get `align`; it lands here together with `OntologyFilter`'s +four conformance gates, which is what makes its prompt claim true. + +Each dimension is also settable on its own, as `True`, `False`, or `None` +meaning "follow the level". That is what makes "say nothing in the prompt but +still coerce datatypes" - `ontology_authority='off', enforce_datatypes=True` - +expressible, and it is accepted rather than treated as a contradiction: the +level chooses defaults, it does not constrain overrides. + +`typed_properties` is the second knob, and a different kind of one: `ontology_authority` +governs what reaches the builders, `typed_properties` governs what the builders +then *write*. It is deliberately not folded into the level, because a user who +asks for `align` is asking about naming and has not thereby asked for new +properties on their nodes. It defaults to `'off'` at every layer. + +`OntologyConfig` lives in the parent process. It holds an `Ontology`, which owns +the `rdflib.Graph`, so it is configuration rather than something that crosses +the spawn boundary; what crosses is the rendered constraint string and +`OntologyIndex`. +""" + +import logging +from dataclasses import dataclass, fields +from pathlib import Path +from typing import Literal, Optional, Tuple, Union + +from rdflib import Graph + +from graphrag_toolkit.lexical_graph.indexing.constants import ( + COMPLEMENT_ENTITY_PROPERTIES, + COMPLEMENT_PLACEMENTS, + RESERVED_ENTITY_PROPERTIES, + SUBJECT_PLACEMENTS, + TYPED_PROPERTIES_OFF, + TYPED_PROPERTY_PLACEMENTS, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology import Ontology +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.prompt_constraint import ( + PROMPT_CONSTRAINT_LEVELS, + PROSE_VOCABULARY, + VOCABULARY_FORMATS, +) + +logger = logging.getLogger(__name__) + +OntologyAuthority = Literal['off', 'align', 'strict'] + +# Where a coerced attribute value is stored, if anywhere. +# +# `'subject'` keys the value from the property's own name on the subject entity, +# which is what makes `WHERE c.foundedYear < 2000` possible. `'complement'` +# writes `typed_value` / `datatype` on the complement node instead, which keeps +# the value where the string already is and needs `include_local_entities`. +# `'both'` does each of those things to its own node; it is not a fallback chain. +# +# The accepted values are `TYPED_PROPERTY_PLACEMENTS`, in `indexing.constants` +# alongside the placement membership tuples, so that the build pipeline can read +# them without importing rdflib. A test asserts the two agree. +TypedProperties = Literal['off', 'subject', 'complement', 'both'] + +# The same set the renderer knows, deliberately not a second list. A level this +# module accepted but `prompt_constraint.py` did not would pass configuration +# and then raise from inside pipeline setup. +ONTOLOGY_AUTHORITY_LEVELS:Tuple[str, ...] = PROMPT_CONSTRAINT_LEVELS + +# How the vocabulary is written down for the model, orthogonal to how much +# authority it has. Re-exported from the renderer for the same reason +# `ONTOLOGY_AUTHORITY_LEVELS` is. +VocabularyFormat = Literal['prose', 'turtle'] + +# Anything `Ontology.load` accepts. +OntologySource = Union[Ontology, Graph, str, Path] + +# Anything `ExtractionConfig(ontology=...)` accepts. +OntologyType = Union['OntologyConfig', OntologySource] + +@dataclass(frozen=True) +class ResolvedDimensions: + """The six dimensions after the level defaults and any overrides. + + Attributes: + normalize_names (bool): Rewrite a resolved name to the ontology's own + spelling. + drop_type_restatements (bool): Drop a fact that carries ontology + language rather than a domain statement - an `rdf:type`-family + predicate, or a type-asserting predicate whose value repeats the + subject's own classification. + enforce_entity_types (bool): Drop a fact whose entity classification + does not resolve. + enforce_relationship_types (bool): Drop a fact whose predicate does not + resolve. + enforce_domain_range (bool): Drop a fact whose subject or object class + does not satisfy the declared domain or range. + enforce_datatypes (bool): Drop a fact whose literal does not parse as + the declared XSD datatype. + """ + normalize_names:bool = False + drop_type_restatements:bool = False + enforce_entity_types:bool = False + enforce_relationship_types:bool = False + enforce_domain_range:bool = False + enforce_datatypes:bool = False + + def any_enabled(self) -> bool: + """Return True if any dimension is on.""" + return any(getattr(self, dimension) for dimension in DIMENSIONS) + +# Derived from the dataclass rather than restated, so the defaults table below +# and the override loop cannot fall out of step with the fields themselves. +DIMENSIONS:Tuple[str, ...] = tuple(field.name for field in fields(ResolvedDimensions)) + +# `off` and the all-False default are the same thing, which is +# why `off` maps to a bare `ResolvedDimensions()`. +# +# `strict` turns everything on, and it is the only level that does: the four +# `enforce_*` dimensions are what distinguishes authority over *naming* from +# authority over *membership*, and nothing between `align` and `strict` is a +# level. That last clause is now a measured position rather than a preference - +# `guide` was exactly such an in-between level, distinguishable only by prompt +# wording, and the wording moved extraction in opposite directions on different +# models. A user who wants one gate without the rest asks for it by name, which +# is a request the code can honour precisely. +_LEVEL_DEFAULTS = { + 'off': ResolvedDimensions(), + 'align': ResolvedDimensions(normalize_names=True, drop_type_restatements=True), + 'strict': ResolvedDimensions( + normalize_names=True, + drop_type_restatements=True, + enforce_entity_types=True, + enforce_relationship_types=True, + enforce_domain_range=True, + enforce_datatypes=True, + ), +} + +@dataclass +class OntologyConfig: + """An ontology and the authority it has over extraction. + + Attributes: + ontology (Ontology): The loaded ontology. Any source `Ontology.load` + accepts may be passed in and is normalized to an `Ontology` here, + so `OntologyConfig('company.ttl')` is valid. + ontology_authority (OntologyAuthority): `'off'`, `'align'` or `'strict'`. + Defaults to `'align'` - name what the ontology declares, keep what it + does not. `'strict'` additionally discards what does not conform, and + on a corpus wider than the ontology that is a large share of it. + normalize_names (Optional[bool]): Override; `None` follows the level. + drop_type_restatements (Optional[bool]): Override; `None` follows the + level. On at `align` and `strict`. Set False to keep facts like + `Meridian Freight|TYPE|Company` that restate an entity's own + classification, which the `turtle` vocabulary format provokes and + the prose rendering does not. + enforce_entity_types (Optional[bool]): Override; `None` follows the + level. + enforce_relationship_types (Optional[bool]): As above. + enforce_domain_range (Optional[bool]): As above. + enforce_datatypes (Optional[bool]): As above. + report_violations (bool): Log per-dimension counts of what the filter + rewrote and dropped, at INFO, per node batch rather than per fact. + A record of the filter's actions, not a measure of how far the + corpus diverges from the ontology: at `align` the drop counts are + zero because no gate runs, and at `strict` they are low because the + prompt already asked the model to omit non-conforming facts, so what + reaches the gates is the residue that ignored the instruction. To + count what `strict` would exclude, ask for `align`'s prompt with the + `enforce_*` gates overridden on. + vocabulary_format (VocabularyFormat): How the vocabulary is written down + for the model. `'prose'` (the default) renders the three generated + sections; `'turtle'` shows the ontology's own source instead. Affects + only the topics vocabulary block - the propositions block stays a + class-name list either way, because that stage extracts nothing a + property declaration could steer. Orthogonal to `ontology_authority`, which + still chooses how much authority the vocabulary has. `'turtle'` is + experimental: it gives up the object/datatype channel separation of + the prose sections and their plain-language ranges, in + exchange for stating the RDFS entailments the prose sections + paraphrase away. On the corpus it was measured against, it was a small + net loss; see the Ontology-Guided Extraction documentation. + typed_properties (TypedProperties): Where a coerced attribute value is + stored: `'off'` (the default) stores nothing and leaves the graph + byte-identical to a build without an ontology, `'subject'` keys it + from the property's own name on the subject entity, `'complement'` + writes `typed_value` / `datatype` on the complement node, `'both'` + does each. Anything other than `'off'` needs the annotations the + filter adds, which is why it makes `filter_required()` True on its + own. + """ + ontology:OntologySource + ontology_authority:OntologyAuthority = 'align' + normalize_names:Optional[bool] = None + drop_type_restatements:Optional[bool] = None + enforce_entity_types:Optional[bool] = None + enforce_relationship_types:Optional[bool] = None + enforce_domain_range:Optional[bool] = None + enforce_datatypes:Optional[bool] = None + report_violations:bool = False + typed_properties:TypedProperties = 'off' + vocabulary_format:VocabularyFormat = PROSE_VOCABULARY + + def __post_init__(self) -> None: + """Validate the settings and normalize `ontology` to an `Ontology`. + + Both validations run before the ontology is loaded, so a typo in a + setting is reported without first paying to parse a Turtle file - except + the reserved-name check, which by definition needs the loaded ontology. + + Raises: + ValueError: If `ontology_authority` is not a supported level, if + `typed_properties` is not a supported placement, if + `vocabulary_format` is not a supported format, or if a + declared datatype property would key a reserved `__Entity__` + property under the requested placement. + OntologyLoadError: If `ontology` cannot be loaded. + """ + if self.ontology_authority not in ONTOLOGY_AUTHORITY_LEVELS: + raise ValueError( + f'Unknown ontology authority level: {self.ontology_authority!r}. ' + f'Expected one of {", ".join(ONTOLOGY_AUTHORITY_LEVELS)}.' + ) + + if self.typed_properties not in TYPED_PROPERTY_PLACEMENTS: + raise ValueError( + f'Unknown typed_properties placement: {self.typed_properties!r}. ' + f'Expected one of {", ".join(TYPED_PROPERTY_PLACEMENTS)}.' + ) + + if self.vocabulary_format not in VOCABULARY_FORMATS: + raise ValueError( + f'Unknown ontology vocabulary_format: {self.vocabulary_format!r}. ' + f'Expected one of {", ".join(VOCABULARY_FORMATS)}.' + ) + + self.ontology = Ontology.load(self.ontology) + + self._validate_no_reserved_property_names() + + def _validate_no_reserved_property_names(self) -> None: + """Refuse an ontology whose vocabulary would overwrite the graph model. + + `__Entity__` already owns `value`, `search_str` and + `class`, and complement placement additionally owns `typed_value` and + `datatype`. A declared datatype property named `value` would, under + subject placement, be written with the same key as the entity's own + identity string. + + Raised here rather than left to the builder because this is a + configuration mistake and the builder runs in a spawn worker, one fact at + a time, where the only available response is a warning nobody reads. The + builder skips the write as well - defence in depth, not the primary + check. + + The check is scoped to the placement actually requested, so an ontology + that is fine for `'off'` is not rejected for a write it will never + perform. Only subject placement keys a property from the ontology's own + vocabulary, so `'complement'` alone has nothing to collide: it widens the + reserved set only when subject placement is active too, which is `'both'`. + """ + if self.typed_properties not in SUBJECT_PLACEMENTS: + return + + reserved = set(RESERVED_ENTITY_PROPERTIES) + if self.typed_properties in COMPLEMENT_PLACEMENTS: + reserved.update(COMPLEMENT_ENTITY_PROPERTIES) + + collisions = sorted( + declared.local_name + for declared in self.ontology.index().datatype_properties.values() + if declared.local_name in reserved + ) + + if collisions: + raise ValueError( + f'Ontology declares datatype propert{"ies" if len(collisions) > 1 else "y"} ' + f'{", ".join(repr(name) for name in collisions)}, which ' + f'{"collide" if len(collisions) > 1 else "collides"} with a property the ' + f'graph model already owns on __Entity__ ' + f'({", ".join(sorted(reserved))}). Writing ' + f'{"them" if len(collisions) > 1 else "it"} under ' + f'typed_properties={self.typed_properties!r} would overwrite the entity ' + 'itself. Rename the property in the ontology, or set ' + "typed_properties='off'." + ) + + def resolved(self) -> ResolvedDimensions: + """Resolve the six dimensions: level defaults, then explicit overrides. + + Recomputed on each call rather than cached in `__post_init__`, because + this is a plain mutable dataclass and a caller who reassigns + `ontology_authority` should not be reading a stale answer. + """ + defaults = _LEVEL_DEFAULTS[self.ontology_authority] + overrides = {} + for dimension in DIMENSIONS: + override = getattr(self, dimension) + overrides[dimension] = ( + getattr(defaults, dimension) if override is None else bool(override) + ) + return ResolvedDimensions(**overrides) + + def writes_subject_properties(self) -> bool: + """Return True if a coerced value is keyed onto the subject entity.""" + return self.typed_properties in SUBJECT_PLACEMENTS + + def writes_complement_properties(self) -> bool: + """Return True if `typed_value` / `datatype` go on the complement node.""" + return self.typed_properties in COMPLEMENT_PLACEMENTS + + def filter_required(self) -> bool: + """Return True if the ontology filter transform has work to do. + + When every dimension resolves to False there is nothing + for the filter to normalize, enforce or annotate, and it is left out of + the pipeline rather than added as a no-op. + + `typed_properties` is part of this test because the annotations the + builders read - `Relation.canonicalName` and `Entity.datatype` - are + written by the filter and by nothing else. So + `ontology_authority='off', typed_properties='subject'` is a coherent request that + needs the filter present purely to annotate, and dropping it there would + produce a build that silently writes no typed properties at all. + """ + return self.resolved().any_enabled() or self.typed_properties != TYPED_PROPERTIES_OFF + +def to_ontology_config(source:OntologyType) -> OntologyConfig: + """Normalize whatever the user configured into an `OntologyConfig`. + + The single entry point for `ExtractionConfig.ontology`, so that every + accepted form converges on one object before anything reads it. + + Args: + source: An `OntologyConfig`, an `Ontology`, an `rdflib.Graph`, or a path + to a `.ttl` file. + + Returns: + `source` itself when it is already an `OntologyConfig` - settings a user + made are never rebuilt at defaults - otherwise a new `OntologyConfig` + wrapping the loaded ontology at the default authority level. + + Raises: + OntologyLoadError: If `source` is not a form that can be loaded. + """ + if isinstance(source, OntologyConfig): + return source + return OntologyConfig(source) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_filter.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_filter.py new file mode 100644 index 00000000..f17ee59b --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_filter.py @@ -0,0 +1,818 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The one component in this feature that changes extracted data. + +Everything else the ontology produces is static: a constraint string rendered +into the prompt, and `OntologyIndex`. Neither touches a fact. `OntologyFilter` +runs after the topic extractor and rewrites what the model emitted - renaming +resolved terms to the ontology's own spelling, dropping facts that violate a +declared shape, and recording the resolution on the surviving facts so the build +stage needs no ontology knowledge of its own. + +Per fact, in this order: + +1. **Resolve** every name once - subject class, object class, predicate. +2. **Normalize** the resolved names to their authored spelling. +3. **Enforce**, each dimension gated on its own flag. +4. **Annotate** the survivors. + +Resolution happens before normalization rather than after, which is stronger +than enforcing on the rewritten name would be: enforcement decides +on the resolved *term*, so it can never see a pre-rewrite spelling at all. This +is only sound because resolution is idempotent - a term's authored name is +indexed under its own resolution key, so resolving the rewritten name returns +the same term. `TestIdempotence` asserts that rather than trusting it. + +Normalization writes the authored name **verbatim**: `rdfs:label` if the term +declares one, otherwise the IRI local name, with no case, separator or +whitespace transformation applied. `:worksFor` stores `worksFor` and +`:WORKS_FOR` stores `WORKS_FOR`; this module imposes no house convention on +either. That is deliberately *not* what the prompt renders - a class authored +`:SportsTeam` with no label renders as `Sports Team` so it survives the response +parser's `.title()`, and stores as `SportsTeam`. The two must be allowed to +disagree, which is why the `naming.py` rendering helpers are not reused here. + +**What this cannot do.** The filter's entire reach is exact match modulo case +and separator style, plus declared `rdfs:label` and `skos:altLabel` values. +Word boundaries are *not* folded away: `resolution_key` maps `worksFor`, +`WORKS_FOR` and `WORKS FOR` onto one key, but `Sportsteam` and `SportsTeam` are +one word and two, so they do not meet. Recognising that `HIRED BY` means the +ontology's employment concept is +the *model's* job, done at extraction time in response to the rendered prompt. +`HIRED BY` against an ontology declaring `worksFor` resolves to nothing and is +left exactly as the parser produced it. Enforcement checks shape, never meaning: +a semantically wrong but resolvable mapping is renamed, annotated, passed by +domain and range if the classes happen to fit, and is then indistinguishable in +the graph from a correct one. + +**Process boundary.** Extraction runs `ProcessPoolExecutor(mp_context='spawn')`, +so this component is pickled per node batch per worker. It therefore holds only +`OntologyIndex` - plain `str`, `List[str]`, `Dict[str, ...]` and `FrozenSet[str]` +- never an `rdflib.Graph`, declares **no custom `__init__`** because +`BaseComponent.__setstate__` calls `self.__init__(**state['__dict__'])`, and +imports no rdflib so workers do not pay for the import. Every import below is +from a submodule rather than from the `ontology` package, whose `__init__` +imports `ontology.py` and therefore rdflib. +""" + +import logging +import re + +from dataclasses import dataclass +from typing import Any, List, Optional, Sequence, Union + +from llama_index.core.schema import BaseNode, TransformComponent + +from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.datatype_utils import ( + coerce_literal, + validates_datatype, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.naming import ( + resolution_key, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_index import ( + DatatypeProperty, + ObjectProperty, + OntologyClass, + OntologyIndex, +) +from graphrag_toolkit.lexical_graph.indexing.model import ( + Entity, + Fact, + Topic, + TopicCollection, +) + +logger = logging.getLogger(__name__) + +OntologyProperty = Union[ObjectProperty, DatatypeProperty] + +# The drop counters, in the order `_enforce` applies its gates, each paired with +# the setting that produced it. The report names the *setting* rather than the +# dimension, because the action a count implies is turning that setting off - and +# a gate the user left off has a zero count and so cannot appear at all, which +# means the report can never attribute a drop to a gate that did not run. +_DROP_DIMENSIONS = ( + ('facts_dropped_type_restatement', 'drop_type_restatements'), + ('facts_dropped_entity_type', 'enforce_entity_types'), + ('facts_dropped_relationship_type', 'enforce_relationship_types'), + ('facts_dropped_domain_range', 'enforce_domain_range'), + ('facts_dropped_datatype', 'enforce_datatypes'), +) + +_NON_ALPHANUMERIC = re.compile(r'[^a-z0-9]') + +def _compact(name:str) -> str: + """Fold a name to lowercase alphanumerics, dropping every separator. + + Deliberately coarser than `resolution_key`, which preserves word boundaries + because it has to round-trip against the prompt rendering. Here the opposite + is wanted: the names being matched are ones no ontology declared and no + renderer produced, so there is no round trip to preserve, and every spelling + a model might reach for should land on one key. + + 'rdf:type' 'RDF_TYPE' 'rdf type' -> 'rdftype' + 'subClassOf' 'SUBCLASS_OF' -> 'subclassof' + 'isA' 'IS_A' 'is a' -> 'isa' + + Note what this gives up: it cannot tell `rdfs:label` from a predicate + genuinely named `RDFSLABEL`. That is why the prefixed family below is + matched on the *whole* compacted string, prefix included - stripping the + prefix would fold `rdfs:domain` onto `DOMAIN`, and a web domain is a real + attribute. + """ + return _NON_ALPHANUMERIC.sub('', (name or '').lower()) + +# Predicates that are ontology *language* rather than domain vocabulary. A fact +# whose predicate is one of these is never a statement about the world, whatever +# its value, so `drop_type_restatements` drops it unconditionally. +# +# Matched with the prefix attached, which is what makes the set safe: `rdfs:label` +# is meta, `LABEL` is a perfectly good attribute of a parcel, and only the prefix +# distinguishes them. A model that emits the bare local name is caught by the +# restatement rule below instead, if it is caught at all. +# +# Empirically empty on the recorded corpus - no predicate any model produced +# contains a colon. It is here because the failure it guards is the one the +# `turtle` vocabulary format invites, and a model that does write `rdf:type` +# verbatim should not need a code change to be handled. +_ONTOLOGY_LANGUAGE_PREDICATES = frozenset( + _compact(name) for name in ( + 'rdf:type', 'rdf:Property', 'rdf:value', + 'rdfs:subClassOf', 'rdfs:subPropertyOf', 'rdfs:domain', 'rdfs:range', + 'rdfs:label', 'rdfs:comment', 'rdfs:isDefinedBy', 'rdfs:seeAlso', + 'owl:Class', 'owl:ObjectProperty', 'owl:DatatypeProperty', + 'owl:sameAs', 'owl:equivalentClass', 'owl:equivalentProperty', + 'owl:Thing', 'owl:NamedIndividual', + 'skos:prefLabel', 'skos:altLabel', 'skos:broader', 'skos:narrower', + ) +) + +# Predicates whose meaning is "this entity is of type X". Membership here is not +# enough to drop a fact: the value must also restate the subject's own +# classification, because every one of these words is a legitimate domain +# attribute against some ontology. `Ship|CLASS|Destroyer` stays unless the ship +# is already classified `Destroyer`, at which point the fact is a duplicate of +# the entity's own label and carries nothing. +# +# `a` is Turtle's own type keyword, and `describedBy` is on the list for a reason +# that is empirical rather than semantic: it is the worked example the base +# extraction prompt gives for attribute naming, so models reach for it when they +# have a class name and nowhere to put it. Both are safe here only because of the +# restatement requirement. +# +# The list will drift - these are the forms three models produced on one corpus, +# plus the obvious neighbours. A form that is missing is a cleanup not performed, +# never a wrong drop, which is the asymmetry that makes an incomplete list +# acceptable. +_TYPE_ASSERTING_PREDICATES = frozenset( + _compact(name) for name in ( + 'a', 'type', 'typeOf', 'entityType', 'instanceOf', 'isA', 'isAn', + 'class', 'classification', 'classifiedAs', 'category', 'categoryOf', + 'kind', 'kindOf', 'subClassOf', 'subTypeOf', 'describedBy', + ) +) + +# Datatype IRIs already warned about, per process. Module level rather than +# instance level so a fresh `OntologyFilter` per node batch does not re-warn - +# the component is rebuilt on every unpickle. See +# `OntologyFilter._warn_unvalidated_datatype` for why per-process is the honest +# contract. Exposed under this name so tests can clear it. +_warned_unvalidated_datatypes = set() + +def _complement_literal(complement:Optional[Any]) -> Optional[str]: + """The string form of a fact's complement. + + `Fact.complement` is `Optional[EntityType]` where + `EntityType = Union[Entity, str]`, and both forms occur: the parser builds an + `Entity`, while a hand-constructed or older payload may carry a bare `str`. + """ + if complement is None: + return None + if isinstance(complement, Entity): + return complement.value + return str(complement) + +def authored_name(term:Union[OntologyClass, OntologyProperty]) -> str: + """The term's name as the ontology author wrote it. + + `rdfs:label` when declared, otherwise the IRI local name, returned verbatim. + This is the canonical stored spelling and is + deliberately unrelated to the prompt rendering convention - see the module + docstring, and `naming.py`, which documents collapsing the two into one + helper as the bug it exists to prevent. + + Args: + term: Any declared class or property from the index. + + Returns: + The authored name. Never empty: `local_name` is always populated. + """ + return term.label or term.local_name + +@dataclass +class FilterCounters: + """What the filter changed, over one `__call__`. + + Rewrites are counted per dimension because they are not equivalent: + a classification rewrite changes entity *identity*, since + `include_classification_in_entity_id` defaults to `True` and + `create_entity_id` hashes the classification into the node id, whereas a + predicate rewrite only changes an edge label. + + Accumulated per `__call__` as a local, never as component state - the + component is pickled per node batch per worker, so an instance attribute + would count one worker's share of one batch and read as a total. + + `OntologyFilter._report` turns these into the log line `report_violations` + asks for. Counting happens either way; only the logging is gated. + """ + classifications_rewritten:int = 0 + predicates_rewritten:int = 0 + facts_dropped_type_restatement:int = 0 + facts_dropped_entity_type:int = 0 + facts_dropped_relationship_type:int = 0 + facts_dropped_domain_range:int = 0 + facts_dropped_datatype:int = 0 + + def facts_dropped(self) -> int: + """Total facts dropped, across all five dropping dimensions. + + Summed from `_DROP_DIMENSIONS` rather than by naming the fields, so a + sixth gate cannot be added with a counter that the total silently omits. + """ + return sum(getattr(self, field) for (field, _) in _DROP_DIMENSIONS) + + def any_change(self) -> bool: + """Whether the filter altered anything at all. + + Annotation is excluded deliberately: it happens to every surviving fact + at every level, so counting it would make every call a change and the + report would carry no signal. + """ + return bool( + self.classifications_rewritten + or self.predicates_rewritten + or self.facts_dropped() + ) + + def summary(self) -> str: + """The counts as a log fragment. + + The three totals are always present, so a reader can tell zero from + absent. The per-dimension breakdown lists only the dimensions that + dropped something, since a gate that is off contributes an unbreakable + zero and naming it would suggest it ran. + """ + parts = [ + f'classifications rewritten: {self.classifications_rewritten}', + f'predicates rewritten: {self.predicates_rewritten}', + f'facts dropped: {self.facts_dropped()}', + ] + + breakdown = ', '.join( + f'{setting}: {getattr(self, field)}' + for (field, setting) in _DROP_DIMENSIONS + if getattr(self, field) + ) + + if breakdown: + parts.append(f'dropped by: {breakdown}') + + return ', '.join(parts) + +@dataclass +class FactResolution: + """What the index says about one fact, resolved once and reused. + + Attributes: + subject_class (Optional[OntologyClass]): The subject's declared class, + or `None` if its classification resolved to nothing. + object_class (Optional[OntologyClass]): As above for the object. Always + `None` on a fact with no object. + predicate (Optional[OntologyProperty]): The declared property the + predicate resolved to, or `None`. + predicate_is_datatype (bool): True when `predicate` is a + `DatatypeProperty`. Read rather than re-tested with `isinstance` + at each use. + """ + subject_class:Optional[OntologyClass] = None + object_class:Optional[OntologyClass] = None + predicate:Optional[OntologyProperty] = None + predicate_is_datatype:bool = False + +class OntologyFilter(TransformComponent): + """Rewrites, filters and annotates extracted topics against an ontology. + + Reads and writes exactly one metadata key, `TOPICS_KEY`. A node without it + passes through untouched. + + Attributes: + index (OntologyIndex): The plain-data read index. The only state, and + the only thing pickled into a worker. + normalize_names (bool): Rewrite resolved names to their authored + spelling. + drop_type_restatements (bool): Drop facts that carry ontology language + rather than a statement about the world. See + `_carries_no_domain_fact`. The one dropping gate that is not an + ontology-conformance check, and the only one on at `align`. + enforce_entity_types (bool): Drop facts whose subject or object + classification does not resolve to a declared class. + enforce_relationship_types (bool): Drop facts whose predicate does not + resolve to a declared property. + enforce_domain_range (bool): Drop facts violating a declared + `rdfs:domain` or `rdfs:range`, honouring subclass closure. + enforce_datatypes (bool): Drop facts whose literal does not coerce to + the declared XSD datatype. + report_violations (bool): Log what this call rewrote and dropped, per + dimension, at INFO. Counting happens regardless; only the logging is + gated on this. See `_report` for what "once per extraction" can + honestly mean on the far side of a spawned boundary, and for why + these counts are not a divergence measure. + + Note: + Annotation is not flag-gated. A fact that survives is annotated even + with every `enforce_*` off, because that is how typed storage works at + `ontology_authority='off'`. + """ + + index:OntologyIndex + + normalize_names:bool = False + drop_type_restatements:bool = False + enforce_entity_types:bool = False + enforce_relationship_types:bool = False + enforce_domain_range:bool = False + enforce_datatypes:bool = False + report_violations:bool = False + + @classmethod + def class_name(cls) -> str: + """Stable name for llama-index component serialization.""" + return 'OntologyFilter' + + def __call__(self, nodes:Sequence[BaseNode], **kwargs:Any) -> Sequence[BaseNode]: + """Filter each node's `TopicCollection` in place. + + Nodes are independent, so processing order does not affect the result. + + Args: + nodes: The nodes leaving the topic extractor. + **kwargs: Ignored; present because `TransformComponent` passes them. + + Returns: + The same sequence, with `TOPICS_KEY` rewritten on the nodes that + carried one. + """ + counters = FilterCounters() + filtered = 0 + + for node in nodes: + topics_data = node.metadata.get(TOPICS_KEY) + if topics_data is None: + continue + + topics = TopicCollection.model_validate(topics_data) + for topic in topics.topics: + self._filter_topic(topic, counters) + + node.metadata[TOPICS_KEY] = topics.model_dump() + filtered += 1 + + self._report(counters, filtered) + + return nodes + + def _report(self, counters:FilterCounters, node_count:int) -> None: + """Log what this call changed, gated on `report_violations`. + + Per-dimension counts are wanted once per extraction rather + than per fact. One `__call__` is the largest unit actually available: it + is one node batch in one worker, and the component is pickled per batch + per worker, so a run-wide total would need state that crosses the spawn + boundary. The line therefore says how many nodes it covers, so a reader + adds the lines up rather than mistaking one for the total - the same + honesty `_warn_unvalidated_datatype` settles for, and for the same + reason. + + A call that changed nothing logs at DEBUG instead. Every gate is off by + default, so on a normalize-only run most batches drop nothing at all, and + a line of zeros per batch would bury the batches that did something. + + What the line is not: a measure of how far the corpus diverges from the + ontology. The drop counters see only facts the model actually emitted, + and the prompt is not neutral about that. At `align` no gate runs, so the + drops are zero by construction. At `strict` the closing instruction tells + the model that unlisted concepts are discarded anyway and it should leave + those facts out, so what reaches a gate is the residue that ignored the + instruction - non-compliance, not divergence, and much the smaller of the + two. The combination that does count divergence is `align`'s prompt with + the `enforce_*` gates overridden on: the model is not asked to suppress, + so the gates see its full output. + """ + if not self.report_violations: + return + + logger.log( + logging.INFO if counters.any_change() else logging.DEBUG, + 'Ontology filter [nodes: %d, %s]', + node_count, + counters.summary(), + ) + + def _filter_topic(self, topic:Topic, counters:FilterCounters) -> None: + """Rewrite one topic's entities and facts in place. + + Both `topic.entities` and each fact's subject and object are visited. + When the filter runs directly on parser output these are the *same* + objects - `parse_extracted_topics` looks the subject up in the topic's + entity dict - so the second visit is a no-op, which is sound only + because normalization is idempotent. After a `model_dump()` round trip + through `TOPICS_KEY`, which is how the pipeline actually delivers them, + the sharing is gone and both visits do real work. Relying on the sharing + to propagate a rewrite would therefore work in a unit test and silently + fail in the pipeline. + + Topics are never dropped, and `topic.entities` is never pruned +: that list is not written to the graph, and pruning it + would change entity extraction rather than fact conformance. + """ + for entity in topic.entities: + self._apply_entity_class(entity, counters) + + for statement in topic.statements: + statement.facts = [ + fact for fact in statement.facts + if self._filter_fact(fact, counters) + ] + + def _filter_fact(self, fact:Fact, counters:FilterCounters) -> bool: + """Normalize, enforce and annotate one fact. + + Returns: + True to keep the fact, False to drop it. + """ + resolution = self._resolve(fact) + + if self.normalize_names: + self._apply_names(fact, resolution, counters) + + if not self._enforce(fact, resolution, counters): + return False + + self._annotate(fact, resolution) + return True + + # ------------------------------------------------------------------ + # Resolve + # ------------------------------------------------------------------ + + def _resolve(self, fact:Fact) -> FactResolution: + """Look up every name on `fact`, once. + + Predicate resolution follows the fact's shape, which is the same + convention used throughout: a fact with an object is a + relation and resolves against object properties; a fact with only a + complement is an attribute and resolves against datatype properties. + + The attribute case falls back to object properties on a miss, because + the shape is not always what it looks like. `parse_extracted_topics` + produces a complement whenever it could not match the object text to an + entity it had already seen, so a genuine `worksFor` relation arrives + complement-shaped whenever the employer was not listed in the entity + block. Normalization rewrites a predicate that resolves to a declared + *property*, not specifically to an object property, so refusing the + fallback would leave those unnormalized. + """ + resolution = FactResolution() + + resolution.subject_class = self.index.resolve_class(fact.subject.classification or '') + + if fact.object is not None: + resolution.object_class = self.index.resolve_class(fact.object.classification or '') + resolution.predicate = self.index.resolve_object_predicate(fact.predicate.value or '') + return resolution + + datatype_property = self.index.resolve_datatype_predicate(fact.predicate.value or '') + if datatype_property is not None: + resolution.predicate = datatype_property + resolution.predicate_is_datatype = True + return resolution + + resolution.predicate = self.index.resolve_object_predicate(fact.predicate.value or '') + return resolution + + # ------------------------------------------------------------------ + # Normalize + # ------------------------------------------------------------------ + + def _apply_names(self, fact:Fact, resolution:FactResolution, counters:FilterCounters) -> None: + """Rewrite the fact's resolved names to their authored spelling. + + Unresolved names are left exactly as the parser produced them +. + """ + self._rewrite_classification(fact.subject, resolution.subject_class, counters) + if fact.object is not None: + self._rewrite_classification(fact.object, resolution.object_class, counters) + + if resolution.predicate is not None: + name = authored_name(resolution.predicate) + if fact.predicate.value != name: + fact.predicate.value = name + counters.predicates_rewritten += 1 + + def _apply_entity_class(self, entity:Entity, counters:FilterCounters) -> None: + """Normalize and annotate a `Topic.entities` member. + + Topic entities carry no predicate, so this is the classification half of + the fact path. Annotation is unconditional, matching facts: an entity + whose class resolved records the IRI even with every `enforce_*` off. + """ + ontology_class = self.index.resolve_class(entity.classification or '') + if self.normalize_names: + self._rewrite_classification(entity, ontology_class, counters) + if ontology_class is not None: + entity.classIri = ontology_class.iri + + def _rewrite_classification( + self, + entity:Optional[Entity], + ontology_class:Optional[OntologyClass], + counters:FilterCounters, + ) -> None: + """Rewrite one entity's classification, in place. + + In place rather than by replacement: within a topic + straight off the parser, this object is also a member of + `topic.entities`, and rebinding would desynchronize the two views. + """ + if entity is None or ontology_class is None: + return + name = authored_name(ontology_class) + if entity.classification != name: + entity.classification = name + counters.classifications_rewritten += 1 + + # ------------------------------------------------------------------ + # Enforce + # ------------------------------------------------------------------ + + def _enforce(self, fact:Fact, resolution:FactResolution, counters:FilterCounters) -> bool: + """Apply the enforcement gates, each on its own flag. + + The gates are **independent**: each one drops only what its own name + says, and no gate implies another. That is the whole point of the + per-dimension escape hatches - `enforce_domain_range` + alone must not start rejecting unresolvable classifications, or a user + who turned `enforce_entity_types` off would find it still on. + + The consequence is that "unknown" is not "violating". A fact whose + subject classification resolves to nothing cannot be *shown* to breach a + declared `rdfs:domain`, so the domain gate passes it and only + `enforce_entity_types` rejects it. Likewise an unresolved predicate has + no declared domain, range or datatype to violate, so only + `enforce_relationship_types` rejects it. This is a deliberate divergence + from the prior implementation, which bundled all three behind one + `strict` flag and could not tell them apart. + + Order is type restatements, then entity types, then relationship types, + then domain and range, then datatypes - broadest first. A fact breaching + two dimensions is dropped once and counted once, under the first gate to + reject it, because the counters break down *drops* rather than violations + and must sum to the number of facts lost. + + `drop_type_restatements` goes first deliberately. At `strict` most of + what it drops would fail `enforce_relationship_types` anyway - the + predicates it matches are by construction ones the ontology does not + declare - so the ordering does not change what survives, only which + counter reports it. "This was a restatement of the entity's own class" + is the more actionable of the two readings, and the less alarming: it + says the model was redundant, not that the ontology was too narrow. + + Returns: + True to keep the fact, False to drop it. + """ + if self.drop_type_restatements and self._carries_no_domain_fact(fact, resolution): + counters.facts_dropped_type_restatement += 1 + return False + + if self.enforce_entity_types: + if resolution.subject_class is None: + counters.facts_dropped_entity_type += 1 + return False + if fact.object is not None and resolution.object_class is None: + counters.facts_dropped_entity_type += 1 + return False + + if self.enforce_relationship_types and resolution.predicate is None: + counters.facts_dropped_relationship_type += 1 + return False + + if self.enforce_domain_range and not self._domain_range_holds(fact, resolution): + counters.facts_dropped_domain_range += 1 + return False + + if self.enforce_datatypes and not self._datatype_holds(fact, resolution): + counters.facts_dropped_datatype += 1 + return False + + return True + + def _carries_no_domain_fact(self, fact:Fact, resolution:FactResolution) -> bool: + """Whether the fact states ontology language instead of something about the world. + + Two independent sufficient conditions, each with its own justification: + + 1. **The predicate is ontology language.** `rdf:type`, `rdfs:subClassOf`, + `owl:sameAs` and their neighbours describe a vocabulary, never a + company. Dropped whatever the value is. Safe unconditionally only + because the prefix is required - see `_compact`. + + 2. **The value restates the subject's own classification**, under a + predicate whose meaning is type membership. `Meridian Freight + [Company] |TYPE| Company` duplicates the entity's own label line and + adds nothing; dropping it cannot lose information. This is the + condition that fires in practice, and the two halves are both + necessary: the value test alone would drop + `Rovers|COMPETES_WITH|Sports Team`, which is at least arguable, and + the predicate test alone would drop + `Rovers|CLASSIFICATION|football club`, which is a real fact the + classification does not carry. + + A predicate the ontology **declares** is never touched by either + condition. An author who declares a datatype property called + `classification` has said what it means, and this gate does not get to + disagree - the same escape the rest of the module gives declared terms. + + Why this is a gate rather than a prompt fix: both failure modes are + base-prompt behaviour, worst with no ontology at all (`|TYPE|` appears 43 + times in the `off` recordings against 0 at prose `strict`). The `turtle` + vocabulary format re-exposes them because it shows type triples in the + same shape as the requested output. Whether prompt wording could suppress + them is an open question; this gate is what makes the format usable + without waiting on that answer. + + Args: + fact: The fact under consideration, after resolution. + resolution: Its resolved terms. Only `predicate` is read, to spare a + declared property. + + Returns: + True to drop the fact. + """ + if resolution.predicate is not None: + return False + + predicate = _compact(fact.predicate.value or '') + + if predicate in _ONTOLOGY_LANGUAGE_PREDICATES: + return True + + if predicate not in _TYPE_ASSERTING_PREDICATES: + return False + + # Attribute shape only. A type restatement whose value matched a named + # entity arrives relation-shaped, and then `object.value` is an entity in + # its own right with its own classification rather than a bare class + # word - a different thing, and one no model produced in seven recorded + # arms, so it is left alone rather than guessed at. + if fact.object is not None: + return False + + value = _complement_literal(fact.complement) + if not value: + return False + + return resolution_key(value) == resolution_key(fact.subject.classification or '') + + def _domain_range_holds(self, fact:Fact, resolution:FactResolution) -> bool: + """Whether the fact's classes satisfy the predicate's declared domain and range. + + Subclass closure is honoured through `OntologyClass.ancestors`, which is + the reflexive transitive closure of `rdfs:subClassOf` - so a declared + domain of `Agent` is satisfied by a `Company`, and by an `Agent`. + + An undeclared domain or range is `owl:Thing` and constrains nothing. An + unresolved class or predicate is unknown rather than wrong; see + `_enforce`. A datatype property has no class range at all - its range is + the XSD type, which `_datatype_holds` checks - so only its domain is + tested here. + + **This checks types, never meaning**. The corpus's one + reproducible mislabelling, a board member recorded as + `Priya Raman | WORKS FOR | Halcyon Motors`, satisfies + `:worksFor rdfs:domain :Person; rdfs:range :Company` exactly and is kept + by this gate at every authority level. Domain and range bound which + entities a property may relate; they cannot see a wrong predicate whose + endpoints happen to fit. + """ + if resolution.predicate is None: + return True + + if not self._class_satisfies(resolution.subject_class, resolution.predicate.domain): + return False + + if resolution.predicate_is_datatype: + return True + + # Reached only for an object property, since the datatype case returned + # above, so `range` is always present. + return self._class_satisfies(resolution.object_class, resolution.predicate.range) + + def _class_satisfies(self, ontology_class:Optional[OntologyClass], constraint:Optional[str]) -> bool: + """Whether `ontology_class` is `constraint` or one of its subclasses.""" + if constraint is None: + return True + if ontology_class is None: + return True + return self.index.is_subclass_of(ontology_class.iri, constraint) + + def _datatype_holds(self, fact:Fact, resolution:FactResolution) -> bool: + """Whether the complement literal is a value of the declared XSD type. + + Only subject-predicate-complement facts resolving to a datatype property + have a declared datatype, so everything else passes. + + When the declared type is one coercion does not implement, the literal is + kept as text and the fact survives - but one WARN is emitted naming the + type, because a value stored without validation must not look like a + validated one. + """ + if not resolution.predicate_is_datatype: + return True + + datatype = resolution.predicate.datatype + + if not validates_datatype(datatype): + self._warn_unvalidated_datatype(datatype, resolution.predicate.local_name) + return True + + return coerce_literal(_complement_literal(fact.complement), datatype) is not None + + @staticmethod + def _warn_unvalidated_datatype(datatype:str, property_name:str) -> None: + """Warn once per process that a declared datatype was not enforced. + + Unconditional on `report_violations`: this reports a + declaration the system did not honour, not a violation count, and silence + would make an unvalidated value indistinguishable from a validated one. + + Deduped in a module-level set, so the honest contract is once per + *process*. Extraction runs `mp_context='spawn'` and each worker imports + this module afresh, so a run with `num_workers=4` can log the same type + four times. Stated rather than papered over: the alternative is shared + state across the process boundary, which is not worth it for a log line, + and the prior implementation's "once per run" was only ever this too. + + The message does not overstate its reach. `Ontology` rejects a datatype + property whose range is outside the XSD namespace at load time, so only + an unimplemented *XSD* type reaches here - a non-XSD IRI is refused + outright by `coerce_literal`. + """ + if datatype in _warned_unvalidated_datatypes: + return + _warned_unvalidated_datatypes.add(datatype) + logger.warning( + 'enforce_datatypes cannot validate %s, declared as the range of %s. ' + 'Values for this and any other property declaring %s are stored as ' + 'text without being checked, and facts carrying them are not ' + 'dropped. Declare a range that coercion implements, or treat these ' + 'values as unvalidated. Logged once per datatype per process.', + datatype, property_name, datatype, + ) + + # ------------------------------------------------------------------ + # Annotate + # ------------------------------------------------------------------ + + def _annotate(self, fact:Fact, resolution:FactResolution) -> None: + """Record the resolution on a surviving fact. + + Not gated on any flag. Written once here and never + recomputed downstream, so the build stage reads an + answer rather than an ontology, and a checkpoint cannot lose it. + + Every value is a plain `str` straight off the index, which holds no + `rdflib` terms - the metadata is written through `json.dump` and + revalidated under `ConfigDict(strict=True)`. + + `canonicalName` is the property's `local_name`, not its authored name: + it is the key typed-property storage writes under, and `rdfs:label` may + contain spaces. + """ + if resolution.subject_class is not None: + fact.subject.classIri = resolution.subject_class.iri + if fact.object is not None and resolution.object_class is not None: + fact.object.classIri = resolution.object_class.iri + + if resolution.predicate is None: + return + + fact.predicate.propertyIri = resolution.predicate.iri + fact.predicate.canonicalName = resolution.predicate.local_name + + if resolution.predicate_is_datatype and isinstance(fact.complement, Entity): + fact.complement.datatype = resolution.predicate.datatype diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_index.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_index.py new file mode 100644 index 00000000..418a8492 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/ontology_index.py @@ -0,0 +1,254 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Plain-data read index over a parsed ontology. + +`OntologyIndex` is the artefact that crosses the extraction process boundary: +extraction runs `ProcessPoolExecutor(mp_context='spawn')`, so the transform +components in the pipeline are pickled per node batch per worker. Every value +here is therefore a `str`, `List[str]`, `Dict[str, ...]` or `FrozenSet[str]` - +never an `rdflib` term, and this module does not import `rdflib` at all. The +`rdflib.Graph` stays behind in the parent process on `Ontology`. + +`is_subclass_of` is a containment check against a precomputed reflexive +ancestor set rather than a graph traversal, because the ontology filter calls +it per fact per node. The `*_by_key` maps are precomputed for the same reason - +`resolve_*` is called per emitted name per fact. +""" + +from typing import Dict, FrozenSet, List, Optional, Union + +from pydantic import BaseModel, ConfigDict, model_validator + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.naming import resolution_key + +XSD_NAMESPACE = 'http://www.w3.org/2001/XMLSchema#' + +OWL_THING = 'http://www.w3.org/2002/07/owl#Thing' + +class OntologyClass(BaseModel): + """An `owl:Class` declared in the ontology. + + Attributes: + iri (str): The full class IRI. + local_name (str): The IRI's last segment, split on `#` or `/`. + label (Optional[str]): The `rdfs:label` value, if declared. This is the + canonical stored spelling when `normalize_names` is on, in + preference to `local_name`. + aliases (List[str]): `skos:altLabel` values, sorted. + parents (List[str]): Direct `rdfs:subClassOf` class IRIs, sorted. + ancestors (FrozenSet[str]): Reflexive transitive closure of `parents` - + includes `iri` itself. What makes `is_subclass_of` O(1). + description (Optional[str]): The `rdfs:comment` value, if declared. + """ + model_config = ConfigDict(frozen=True, strict=True) + + iri:str + local_name:str + label:Optional[str]=None + aliases:List[str]=[] + parents:List[str]=[] + ancestors:FrozenSet[str]=frozenset() + description:Optional[str]=None + +class ObjectProperty(BaseModel): + """An `owl:ObjectProperty` - an entity-to-entity predicate. + + Attributes: + iri (str): The full property IRI. + local_name (str): The IRI's last segment. + label (Optional[str]): The `rdfs:label` value, if declared. + aliases (List[str]): `skos:altLabel` values, sorted. + domain (Optional[str]): The `rdfs:domain` class IRI. `None` means + `owl:Thing` - the property matches any subject class. + range (Optional[str]): The `rdfs:range` class IRI. `None` means + `owl:Thing`. + description (Optional[str]): The `rdfs:comment` value, if declared. + """ + model_config = ConfigDict(frozen=True, strict=True) + + iri:str + local_name:str + label:Optional[str]=None + aliases:List[str]=[] + domain:Optional[str]=None + range:Optional[str]=None + description:Optional[str]=None + +class DatatypeProperty(BaseModel): + """An `owl:DatatypeProperty` - an entity-to-literal predicate. + + Attributes: + iri (str): The full property IRI. + local_name (str): The IRI's last segment. + label (Optional[str]): The `rdfs:label` value, if declared. + aliases (List[str]): `skos:altLabel` values, sorted. + domain (Optional[str]): The `rdfs:domain` class IRI. `None` means + `owl:Thing`. + datatype (str): The `rdfs:range` XSD IRI, as a plain string. Always + populated - a datatype property without an XSD range is rejected + at load time. + description (Optional[str]): The `rdfs:comment` value, if declared. + """ + model_config = ConfigDict(frozen=True, strict=True) + + iri:str + local_name:str + label:Optional[str]=None + aliases:List[str]=[] + domain:Optional[str]=None + datatype:str + description:Optional[str]=None + +OntologyTerm = Union[OntologyClass, ObjectProperty, DatatypeProperty] + +# Resolution precedence. A term's own local name outranks its rdfs:label, which +# outranks any skos:altLabel, so that a name folding to the same key as another +# term's alias still resolves to the term that owns it outright. Ties within a +# rank break on the IRI, which makes resolution a function of the ontology's +# content and not of dict iteration order. +_LOCAL_NAME_RANK = 0 +_LABEL_RANK = 1 +_ALIAS_RANK = 2 + +def _ranked_keys(term:OntologyTerm) -> List[tuple]: + """Return `(resolution_key, rank)` for a term's local name, label, aliases. + + Empty keys are dropped rather than indexed - a term whose alias is a blank + string should not answer a lookup for the empty name. + """ + ranked = [(resolution_key(term.local_name), _LOCAL_NAME_RANK)] + if term.label: + ranked.append((resolution_key(term.label), _LABEL_RANK)) + for alias in term.aliases: + ranked.append((resolution_key(alias), _ALIAS_RANK)) + return [(key, rank) for (key, rank) in ranked if key] + +def _rank_by_key(terms:Dict[str, OntologyTerm]) -> Dict[str, Dict[str, int]]: + """Group terms by resolution key, keeping each term's best rank.""" + ranked:Dict[str, Dict[str, int]] = {} + for term in terms.values(): + for (key, rank) in _ranked_keys(term): + by_iri = ranked.setdefault(key, {}) + if rank < by_iri.get(term.iri, _ALIAS_RANK + 1): + by_iri[term.iri] = rank + return ranked + +def _single_valued_key_index(terms:Dict[str, OntologyTerm]) -> Dict[str, str]: + """Map each resolution key to the one term that wins it.""" + return { + key: min(by_iri, key=lambda iri: (by_iri[iri], iri)) + for (key, by_iri) in _rank_by_key(terms).items() + } + +def _multi_valued_key_index(terms:Dict[str, OntologyTerm]) -> Dict[str, List[str]]: + """Map each resolution key to every term that claims it, best first.""" + return { + key: sorted(by_iri, key=lambda iri: (by_iri[iri], iri)) + for (key, by_iri) in _rank_by_key(terms).items() + } + +class OntologyIndex(BaseModel): + """The read side of an ontology, keyed by IRI. + + A pure function of the `rdflib.Graph` it was built from. Constructed once + at pipeline-configuration time by `Ontology.index()` and read-only + thereafter. + + Attributes: + classes (Dict[str, OntologyClass]): Declared classes by IRI. + object_properties (Dict[str, ObjectProperty]): Declared object + properties by IRI. + datatype_properties (Dict[str, DatatypeProperty]): Declared datatype + properties by IRI. + class_by_key (Dict[str, str]): Class IRI by `resolution_key`. Derived - + recomputed on construction, so a value passed in is discarded. + obj_property_by_key (Dict[str, List[str]]): Object property IRIs by + `resolution_key`, in resolution precedence order. Derived. + Multi-valued because a key collision between two declared + properties is an authoring choice to report on, not a reason to + lose one of them. + dt_property_by_key (Dict[str, List[str]]): Datatype property IRIs by + `resolution_key`, in resolution precedence order. Derived. + """ + model_config = ConfigDict(frozen=True, strict=True) + + classes:Dict[str, OntologyClass]={} + object_properties:Dict[str, ObjectProperty]={} + datatype_properties:Dict[str, DatatypeProperty]={} + class_by_key:Dict[str, str]={} + obj_property_by_key:Dict[str, List[str]]={} + dt_property_by_key:Dict[str, List[str]]={} + + @model_validator(mode='after') + def _build_key_indexes(self) -> 'OntologyIndex': + """Derive the three key maps from the terms. + + Done here rather than in `Ontology._build_index` so that an index built + by any route - constructed directly in a test, unpickled in a worker, + round-tripped through `model_dump` - cannot carry keys that disagree + with its terms. `object.__setattr__` because the model is frozen. + """ + object.__setattr__(self, 'class_by_key', _single_valued_key_index(self.classes)) + object.__setattr__(self, 'obj_property_by_key', _multi_valued_key_index(self.object_properties)) + object.__setattr__(self, 'dt_property_by_key', _multi_valued_key_index(self.datatype_properties)) + return self + + def resolve_class(self, name:str) -> Optional[OntologyClass]: + """Resolve an emitted classification to a declared class. + + Args: + name: A name in any convention - `'Sports Team'` as the parser + hands it over, `'SPORTS_TEAM'`, `'SportsTeam'`, or a declared + label or alias. + + Returns: + The declared class, or `None` if nothing in the ontology carries + that name. + """ + iri = self.class_by_key.get(resolution_key(name)) + return self.classes.get(iri) if iri else None + + def resolve_object_predicate(self, name:str) -> Optional[ObjectProperty]: + """Resolve an emitted predicate to a declared object property. + + Args: + name: A name in any convention - `'WORKS FOR'` as the parser hands + it over, `'WORKS_FOR'`, `'works for'`, `'worksFor'`, or a + declared label or alias. + + Returns: + The highest-precedence declared object property carrying that name, + or `None`. + """ + return self._first(self.obj_property_by_key, self.object_properties, name) + + def resolve_datatype_predicate(self, name:str) -> Optional[DatatypeProperty]: + """Resolve an emitted attribute name to a declared datatype property. + + Args: + name: A name in any convention, as for `resolve_object_predicate`. + + Returns: + The highest-precedence declared datatype property carrying that + name, or `None`. + """ + return self._first(self.dt_property_by_key, self.datatype_properties, name) + + @staticmethod + def _first(by_key:Dict[str, List[str]], terms:Dict[str, OntologyTerm], name:str) -> Optional[OntologyTerm]: + """Return the first term claiming `name`'s resolution key.""" + iris = by_key.get(resolution_key(name)) + return terms.get(iris[0]) if iris else None + + def is_subclass_of(self, child_iri:str, parent_iri:str) -> bool: + """Return True if `parent_iri` is an ancestor of `child_iri`. + + Reflexive: `is_subclass_of(c, c)` is True for any declared class `c`. + An IRI that is not a declared class returns False rather than raising, + so callers can pass an unresolved classification straight through. + """ + ontology_class = self.classes.get(child_iri) + if ontology_class is None: + return False + return parent_iri in ontology_class.ancestors diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/prompt_constraint.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/prompt_constraint.py new file mode 100644 index 00000000..04c3c7ed --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/ontology/prompt_constraint.py @@ -0,0 +1,541 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Render an ontology as the constraint block that goes into the prompt. + +Three labelled sections, mapped onto the two output channels +`EXTRACT_TOPICS_PROMPT` already defines plus its entity classification: + +| Section | Prompt channel | Fed by | +|----------------|---------------------------------|-----------------------| +| Entity types | `entity\\|label` | `owl:Class` | +| Relationships | `entity\\|RELATIONSHIP\\|entity` | `owl:ObjectProperty` | +| Attributes | `entity\\|ATTRIBUTE_NAME\\|value` | `owl:DatatypeProperty`| + +The two property sections are separate and separately labelled on purpose. +Merged into one list they invite the model to emit an attribute where a +relationship belongs, and the attribute path depends on the distinction +surviving all the way to the response: an emitted attribute name that resolves +to a declared `owl:DatatypeProperty` is what yields a canonical name and a +declared datatype, and nothing downstream of that fires without it. + +Datatype ranges are named in terms a model can act on - `integer`, `decimal +number`, `date`, `true/false`, `text` - never as XSD IRIs. The range is rendered +to help the model pick the right *attribute* and report the value in a sensible +form. It is not an instruction to convert or validate anything: value handling +is deterministic and happens after extraction. + +Rendering reads only `OntologyIndex`, so it is plain-data in and a string out, +with no `rdflib` involved. `Ontology.format_as_prompt_constraint` is the public +entry point. + +Determinism: every iteration is over a `sorted(...)` by +`local_name`, so the output is a pure function of the ontology's content and not +of triple insertion order, and two calls at the same level are byte-identical. +""" + +from typing import Dict, List, Optional + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.naming import ( + camel_to_upper_snake, + resolution_key, + title_case_with_spaces, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_index import ( + XSD_NAMESPACE, + DatatypeProperty, + OntologyClass, + OntologyIndex, +) + +# The levels this renderer knows, weakest first. The order is the ladder of +# prompt pressure, and each closing below adds to the one before it rather than +# replacing its reasoning. +# +# There was a fourth, `guide`, between `align` and `strict`, whose entire +# mechanism was one extra sentence of preference wording. It was measured against +# a same-configuration control on three models and removed: the sentence had a +# real effect with no consistent direction, raising resolution on one model and +# lowering it 11-16 points on another, so no user could pick it in advance. The +# recordings that showed it are in git history. +PROMPT_CONSTRAINT_LEVELS = ('off', 'align', 'strict') + +# How the vocabulary itself is presented, independent of the level. `'prose'` is +# the three generated sections above; `'turtle'` shows the ontology's own source +# instead and keeps only the protocol section. +# +# The point of `'turtle'` is that the prose sections are a *paraphrase*: +# `rdfs:subClassOf` becomes indentation, which carries no formal meaning, so the +# entailment that a subclass may carry its parent's properties is nowhere stated. +# Turtle carries it inherently. Whether that is worth the two things it costs - +# the object/datatype channel separation, and the plain-language ranges - is a +# measurement, not a preference, and it measured slightly worse overall. +VOCABULARY_FORMATS = ('prose', 'turtle') + +PROSE_VOCABULARY, TURTLE_VOCABULARY = VOCABULARY_FORMATS + +# What a domain or range of owl:Thing - or none at all - is called in the +# prompt. 'owl:Thing' is vocabulary for an ontologist, not for a model being +# asked to read text. +ANY_ENTITY = 'anything' + +ANY_SUBJECT_GROUP = 'any entity' + +# XSD range -> the words the model sees. Deliberately plain: the model is being +# helped to pick the right attribute and to report the value in a sensible form, +# not asked to perform a conversion. +_TYPE_NAMES = { + 'integer': 'integer', + 'int': 'integer', + 'long': 'integer', + 'short': 'integer', + 'byte': 'integer', + 'nonNegativeInteger': 'integer', + 'nonPositiveInteger': 'integer', + 'positiveInteger': 'integer', + 'negativeInteger': 'integer', + 'unsignedInt': 'integer', + 'unsignedLong': 'integer', + 'unsignedShort': 'integer', + 'unsignedByte': 'integer', + 'gYear': 'integer', + 'decimal': 'decimal number', + 'double': 'decimal number', + 'float': 'decimal number', + 'boolean': 'true/false', + 'date': 'date', + 'dateTime': 'date and time', + 'time': 'time of day', +} + +DEFAULT_TYPE_NAME = 'text' + +_HEADER = """\ +# Vocabulary for this extraction + +The entity types, relationships and attributes below are the vocabulary for this +extraction. Read the guidance at the end of this section before using them.""" + +# The `'turtle'` counterpart. It has to do two jobs the prose header does not. +# +# It has to say which OWL construct feeds which output channel, because Turtle +# interleaves `owl:ObjectProperty` and `owl:DatatypeProperty` in whatever order +# the author wrote them - the prose rendering separates them into two labelled +# sections precisely so the model cannot emit an attribute where a relationship +# belongs, and that separation is what this format gives up. +# +# And it has to say that the declared ranges are not conversion instructions, +# because `xsd:integer` reads far more like one than the word `integer` does, and +# value handling here is deterministic and post-hoc. +# +# Both sentences are restatements of what the prose *layout* conveys silently. If +# this arm wins anyway, that is the interesting result; if it loses, these are the +# first two things to suspect. +_TURTLE_HEADER = """\ +# Vocabulary for this extraction + +The vocabulary for this extraction is the OWL/RDFS ontology below, given in +Turtle. Read the guidance at the end of this section before using it. + +How to read it for this task: + + - An `owl:Class` is an entity type. `rdfs:subClassOf` means the subject type is + a kind of the object type, with everything that follows from that: a type may + be used wherever any of its ancestors may be used, and it carries every + property declared on any of its ancestors. Use the most specific type the + text actually supports. + - An `owl:ObjectProperty` is a relationship between two entities. Emit these on + the relationship channel, as `entity|RELATIONSHIP|entity`. + - An `owl:DatatypeProperty` is an attribute of one entity holding a literal + value. Emit these on the attribute channel, as `entity|ATTRIBUTE_NAME|value`. + Do not confuse the two: an attribute emitted where a relationship belongs is + discarded. + - `rdfs:domain` and `rdfs:range` say which entity types a property may relate, + honouring `rdfs:subClassOf`. + - An `rdfs:range` naming an XSD datatype describes the kind of value the + attribute holds. It is there to help you pick the right attribute and report + the value in a sensible form. **It is not an instruction to convert, reformat + or validate anything** - report the value as the text states it. + - `rdfs:label` and `skos:altLabel` are the names this vocabulary is known by; + `rdfs:comment` describes what a term means. + +When you write a name from this ontology, write the local name - the part after +the `:` - and not the full IRI or the prefix.""" + +def format_turtle_vocabulary(turtle:str, level:str) -> str: + """Wrap an ontology's Turtle source as a prompt constraint block. + + The Turtle stands in for the three generated sections; the protocol section + is the same text the prose format uses, because how to *use* the vocabulary + is not a function of how the vocabulary was written down, and holding it + constant is what makes the two formats comparable. + + Serialization lives with the caller, in `ontology.py`: this module renders + from `OntologyIndex` with no `rdflib` involved, and an import-graph test + holds it to that. + + Args: + turtle: The ontology serialized as Turtle. + level: `'align'` or `'strict'`. `'off'` never reaches here. + + Returns: + The block, with no trailing newline, or `''` for empty Turtle. + """ + if not turtle.strip(): + return '' + + return '\n\n'.join([ + _TURTLE_HEADER, + f'```turtle\n{turtle.strip()}\n```', + f'{_PROTOCOL_HEADING}\n\n{_MAP_BY_MEANING}\n\n{_THEN_SPELLING}\n\n{_CLOSINGS[level]}', + ]) + +_CLASSES_HEADING = """\ +## Entity types (entity|label) + +Indentation shows specialization: an indented type is a kind of the type above +it. Use the most specific type that the text actually supports.""" + +_OBJECT_PROPERTIES_HEADING = """\ +## Relationships (entity|RELATIONSHIP|entity) + +Each line is a relationship name followed by the entity types it relates, as +subject -> object.""" + +_DATATYPE_PROPERTIES_HEADING = """\ +## Attributes (entity|ATTRIBUTE_NAME|value) + +Grouped by the entity type that carries the attribute. Each line is an +attribute name followed by the kind of value it holds. The kind of value is +there to help you pick the right attribute and report the value in a sensible +form - report what the text says, and do not convert or invent values. + +Write a listed attribute name exactly as it appears here. Do not add a prefix +such as HAS_ to it, and do not reword it.""" + +_PROTOCOL_HEADING = '## Using this vocabulary' + +_MAP_BY_MEANING = """\ +Map by meaning first. When a relationship or attribute in the text means what +one of the entries above means, use that entry - however differently the text +words it. The text will usually not use the listed wording, and that does not +matter; what the entry means is what decides it.""" + +_THEN_SPELLING = """\ +Then use the listed spelling. Once you have chosen an entry, write its name +exactly as it appears above, in the same case and with the same underscores, and +classify its subject and object as the entity types listed for it.""" + +_ALIGN_CLOSING = """\ +When nothing listed carries the meaning, name the concept yourself as you +normally would. Do not stretch a listed entry to cover something it does not +mean - a wrong listed name is worse than an unlisted one.""" + +# Two claims, and the second is the one that changes the +# model's incentive: not just that unlisted names are discarded, but that +# omitting therefore costs nothing while a stretched listed name is *kept*. At +# align a wrong listed name is merely wrong; here it is the only way to get bad +# data past the filter, so it is worth naming as the specific hazard. +# +# The opening preference sentence is the one `guide` was removed over. It is kept +# here for a reason that did not apply there: strict discards what does not match, +# so a closing that says "leave the fact out" without first asking the model to +# look for a genuine match would discard more while helping less. What the guide +# measurement establishes is that the sentence is not free, and its effect here - +# alongside four gates rather than alone - has not been measured on its own. +_STRICT_CLOSING = """\ +Prefer a listed entry whenever the meaning is close, even when the text words it +quite differently. When nothing listed carries the meaning, leave the fact out: +anything not listed above is discarded, so omitting it loses nothing and naming +it gains nothing. Do not stretch a listed entry to cover something it does not +mean - that is worse than omitting, because a wrong listed name is kept.""" + +_CLOSINGS = { + 'align': _ALIGN_CLOSING, + 'strict': _STRICT_CLOSING, +} + +# The propositions prompt gets the entity types and nothing else. Short on +# purpose: that stage decomposes text into atomic statements and classifies the +# entities it names, and a relationship or attribute vocabulary would be +# instructions for work it does not do. +_PROPOSITION_HEADER = '# Entity types for this extraction' + +_PROPOSITION_LEADS = { + 'align': """\ +When a proposition classifies a named entity, use one of these types where the +entity is one of them:""", + 'strict': """\ +When a proposition classifies a named entity, prefer one of these types whenever +the entity is close to one of them, however differently the text words it:""", +} + +# Keyed by level because the align wording is not merely weaker at strict, it is +# wrong there: "classify the entity as you normally would" invites a type that +# `enforce_entity_types` will then drop facts over. The strict variant says what +# actually happens instead. +_PROPOSITION_CLOSINGS = { + 'align': """\ +When none of them fits, classify the entity as you normally would. Do not +stretch a listed type to cover something it does not mean.""", + 'strict': """\ +When none of them fits, classify the entity as you normally would - but note that +facts about an entity classified as anything not listed above are discarded. Do +not stretch a listed type to cover something it does not mean; a wrong listed +type is kept, and is worse than an unlisted one.""", +} + +def format_as_prompt_constraint(index:OntologyIndex, level:str) -> str: + """Render `index` as a prompt constraint block at `level`. + + Args: + index: The read index to render. + level: `'off'`, `'align'` or `'strict'`. + + Returns: + The block, with no trailing newline, or the empty string when `level` is + `'off'` or the ontology declares no terms at all. + + Raises: + ValueError: If `level` is not a level this renderer knows. + """ + if level not in PROMPT_CONSTRAINT_LEVELS: + raise ValueError( + f'Unknown ontology authority level: {level!r}. ' + f'Expected one of {", ".join(PROMPT_CONSTRAINT_LEVELS)}.' + ) + + if level == 'off': + return '' + + sections = [ + _render_classes(index), + _render_object_properties(index), + _render_datatype_properties(index), + ] + sections = [section for section in sections if section] + + if not sections: + return '' + + return '\n\n'.join([ + _HEADER, + *sections, + f'{_PROTOCOL_HEADING}\n\n{_MAP_BY_MEANING}\n\n{_THEN_SPELLING}\n\n{_CLOSINGS[level]}', + ]) + +def format_as_proposition_constraint(index:OntologyIndex, level:str) -> str: + """Render the entity types alone, for the propositions prompt. + + The propositions stage decomposes text before topics see it, and the one + thing it does that an ontology can steer is rule 4 - "add a proposition per + named entity that classifies that entity". So it gets the class names and + nothing else: relationships and attributes are extracted downstream, and + naming them here would spend a large part of the propositions prompt on + vocabulary that stage cannot use. + + Args: + index: The read index to render. + level: `'off'`, `'align'` or `'strict'`. + + Returns: + The hint, with no trailing newline, or the empty string when `level` is + `'off'` or the ontology declares no classes. + + Raises: + ValueError: If `level` is not a level this renderer knows. + """ + if level not in PROMPT_CONSTRAINT_LEVELS: + raise ValueError( + f'Unknown ontology authority level: {level!r}. ' + f'Expected one of {", ".join(PROMPT_CONSTRAINT_LEVELS)}.' + ) + + if level == 'off' or not index.classes: + return '' + + return '\n\n'.join([ + _PROPOSITION_HEADER, + f'{_PROPOSITION_LEADS[level]}\n\n{", ".join(rendered_class_names(index))}.', + _PROPOSITION_CLOSINGS[level], + ]) + +def rendered_class_names(index:OntologyIndex) -> List[str]: + """The class names exactly as the prompt shows them, sorted. + + `preferred_entity_classifications` is seeded from these, so the + vocabulary block and the `{preferred_entity_classifications}` slot cannot + name the same class two different ways. That only holds while both go through + `_render_class_name`, which is why this is the one place either of them gets + a list of names from. + + Flat and alphabetical rather than in the tree order `_render_classes` uses: + that ordering carries the hierarchy, and a preference list has nowhere to put + it. + + Args: + index: The read index to take the class names from. + + Returns: + The rendered names, sorted; empty when the ontology declares no classes. + """ + return sorted( + _render_class_name(ontology_class) for ontology_class in index.classes.values() + ) + +def _render_classes(index:OntologyIndex) -> str: + """Render the class hierarchy as an indented tree. + + Depth is shown by indentation rather than stated, and the tree is walked + parent-first so that a specialization sits directly under the type it + specializes. Alphabetical ordering within each level keeps it deterministic; + a flat alphabetical list would have put `Athlete` above `Company` and lost + the structure the model is being shown. + """ + if not index.classes: + return '' + + children:Dict[str, List[OntologyClass]] = {} + roots:List[OntologyClass] = [] + + for ontology_class in index.classes.values(): + # A class with more than one declared parent is rendered once, under the + # first of them; the others are named on its own line, so nothing is + # lost and no subtree is duplicated. + parents = [parent for parent in ontology_class.parents if parent in index.classes] + if parents: + children.setdefault(parents[0], []).append(ontology_class) + else: + roots.append(ontology_class) + + lines:List[str] = [] + + def render(ontology_class:OntologyClass, depth:int) -> None: + extra_parents = [ + _render_class_name(index.classes[parent]) + for parent in ontology_class.parents[1:] + if parent in index.classes + ] + also = ([f'also a kind of {", ".join(extra_parents)}'] if extra_parents else []) + lines.append(' ' * (depth + 1) + _render_term_line( + _render_class_name(ontology_class), + _render_aliases(ontology_class, title_case_with_spaces), + ontology_class.description, + also + )) + for child in sorted(children.get(ontology_class.iri, []), key=lambda c: c.local_name): + render(child, depth + 1) + + for root in sorted(roots, key=lambda c: c.local_name): + render(root, 0) + + return f'{_CLASSES_HEADING}\n\n' + '\n'.join(lines) + +def _render_object_properties(index:OntologyIndex) -> str: + """Render the object properties as `NAME subject -> object`.""" + if not index.object_properties: + return '' + + properties = sorted(index.object_properties.values(), key=lambda p: p.local_name) + width = max(len(_render_property_name(prop)) for prop in properties) + + lines = [ + ' ' + _render_term_line( + _render_property_name(prop).ljust(width), + _render_aliases(prop, camel_to_upper_snake), + prop.description, + [f'{_class_name(index, prop.domain)} -> {_class_name(index, prop.range)}'] + ) + for prop in properties + ] + + return f'{_OBJECT_PROPERTIES_HEADING}\n\n' + '\n'.join(lines) + +def _render_datatype_properties(index:OntologyIndex) -> str: + """Render the datatype properties grouped by the class that carries them. + + Properties with no declared domain are grouped last, under a heading that + says any entity may carry them, rather than being silently attached to a + class the ontology did not name. + """ + if not index.datatype_properties: + return '' + + properties = sorted(index.datatype_properties.values(), key=lambda p: p.local_name) + width = max(len(_render_property_name(prop)) for prop in properties) + + grouped:Dict[str, List[DatatypeProperty]] = {} + for prop in properties: + grouped.setdefault(prop.domain or '', []).append(prop) + + lines:List[str] = [] + for domain in sorted(grouped, key=lambda iri: _domain_sort_key(index, iri)): + subject = _class_name(index, domain) if domain else ANY_SUBJECT_GROUP + lines.append(f' {subject}:') + for prop in grouped[domain]: + lines.append(' ' + _render_term_line( + _render_property_name(prop).ljust(width), + _render_aliases(prop, camel_to_upper_snake), + prop.description, + [_type_name_of(prop.datatype)] + )) + + return f'{_DATATYPE_PROPERTIES_HEADING}\n\n' + '\n'.join(lines) + +def _render_term_line(name:str, aliases:List[str], description:Optional[str], middle:List[str]) -> str: + """Assemble one vocabulary line: name, then facts, then description. + + `middle` carries whatever the section puts between the name and the + description - a `subject -> object` pair, a value type, an extra parent. + """ + parts = [name, *middle] + if aliases: + parts.append(f'(also known as {", ".join(aliases)})') + if description: + parts.append(f'"{description}"') + return ' '.join(parts) + +def _render_class_name(ontology_class:OntologyClass) -> str: + """Render a class name for the prompt, preferring the declared label.""" + return title_case_with_spaces(ontology_class.label or ontology_class.local_name) + +def _render_property_name(prop) -> str: + """Render a property name for the prompt, preferring the declared label.""" + return camel_to_upper_snake(prop.label or prop.local_name) + +def _render_aliases(term, render) -> List[str]: + """Render a term's `skos:altLabel` values as "also known as" hints. + + Rendered through the same function as the primary name, because an alias is + a name the model may emit and therefore has to survive the response parser + too. An alias that renders to the same name as the term itself is dropped - + it tells the model nothing. + """ + primary = resolution_key(term.label or term.local_name) + rendered:List[str] = [] + for alias in term.aliases: + if resolution_key(alias) == primary: + continue + name = render(alias) + if name and name not in rendered: + rendered.append(name) + return rendered + +def _class_name(index:OntologyIndex, iri:Optional[str]) -> str: + """Render a domain or range class name, or `anything` when unconstrained.""" + if not iri: + return ANY_ENTITY + ontology_class = index.classes.get(iri) + return _render_class_name(ontology_class) if ontology_class else ANY_ENTITY + +def _domain_sort_key(index:OntologyIndex, iri:str) -> tuple: + """Sort domain groups by class local name, with the no-domain group last.""" + if not iri: + return (1, '') + ontology_class = index.classes.get(iri) + return (0, ontology_class.local_name if ontology_class else iri) + +def _type_name_of(datatype:str) -> str: + """Name an XSD range in terms the model can act on.""" + local_name = datatype[len(XSD_NAMESPACE):] if datatype.startswith(XSD_NAMESPACE) else '' + return _TYPE_NAMES.get(local_name, DEFAULT_TYPE_NAME) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py index 37a90622..db3461bb 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py @@ -11,7 +11,7 @@ from graphrag_toolkit.lexical_graph.indexing.extract.preferred_values import PreferredValuesProvider, default_preferred_values from graphrag_toolkit.lexical_graph.indexing.model import TopicCollection from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY -from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_TOPICS_PROMPT +from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_TOPICS_PROMPT, with_ontology_constraints from graphrag_toolkit.lexical_graph.indexing.extract.progress import run_jobs_with_progress from graphrag_toolkit.lexical_graph.utils.arg_utils import coalesce @@ -44,6 +44,11 @@ class TopicExtractor(BaseExtractor): description='Topic provider' ) + ontology_constraints:str = Field( + default='', + description='Rendered ontology vocabulary block, composed into the prompt template at render time' + ) + @classmethod def class_name(cls) -> str: """ @@ -65,7 +70,8 @@ def __init__(self, source_metadata_field=None, num_workers:Optional[int]=None, entity_classification_provider=None, - topic_provider=None + topic_provider=None, + ontology_constraints:str='' ): """ Initializes the instance with the provided or default parameters to facilitate @@ -88,6 +94,10 @@ def __init__(self, topic_provider (FixedScopedValueProvider, optional): Provider for topics. Defaults to a fixed-scoped value provider initialized with an empty list. + ontology_constraints (str, optional): Rendered ontology vocabulary + block, composed into the prompt template at render time. + Defaults to the empty string, which leaves the template - and + therefore the LLM cache key - exactly as it is today. """ num_workers = coalesce(num_workers, GraphRAGConfig.extraction_num_threads_per_worker) @@ -101,7 +111,8 @@ def __init__(self, source_metadata_field=source_metadata_field, num_workers=num_workers, entity_classification_provider=entity_classification_provider or default_preferred_values([]), - topic_provider=topic_provider or default_preferred_values([]) + topic_provider=topic_provider or default_preferred_values([]), + ontology_constraints=ontology_constraints ) logger.debug(f'Prompt template: {self.prompt_template}') @@ -206,7 +217,9 @@ async def _extract_topics(self, text:str, preferred_entity_classifications:List[ """ def blocking_llm_call(): return self.llm.predict( - PromptTemplate(template=self.prompt_template), + PromptTemplate( + template=with_ontology_constraints(self.prompt_template, self.ontology_constraints) + ), text=text, preferred_entity_classifications=format_list(preferred_entity_classifications), preferred_topics=format_list(preferred_topics), diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/model.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/model.py index 4840c1a1..90601eee 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/model.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/model.py @@ -128,15 +128,36 @@ class Entity(BaseModel): attribute. classification (Optional[str]): Optional classification or category of the entity. Defaults to None. + classIri (Optional[str]): IRI of the ontology class this entity's + classification resolved to, when an ontology was configured and the + filter resolved it. Defaults to None. + datatype (Optional[str]): XSD datatype IRI for a complement entity whose + predicate resolved to a datatype property. Defaults to None. """ model_config = ConfigDict(strict=True) - + entityId: Optional[str]=None altEntityId: Optional[str]=None value: str classification: Optional[str]=None + # Ontology annotations. Written once by `OntologyFilter` during extraction and + # never recomputed downstream, so the build stage needs no + # ontology knowledge and a checkpoint cannot lose them. + # + # Plain `str`, never `rdflib.URIRef`: this metadata goes through `json.dump` + # in `file_based_docs.py` and through `mp_context='spawn'` pickling, and + # `ConfigDict(strict=True)` above is in the path too. Because `URIRef` + # subclasses `str`, strict mode *accepts and coerces* one rather than + # rejecting it, which is the behaviour that matters here. + # + # Declared rather than relied on as extras: `strict=True` with pydantic's + # default `extra='ignore'` silently discards undeclared keys at both + # `model_dump()` and `model_validate()`. + classIri: Optional[str]=None + datatype: Optional[str]=None + class Relation(BaseModel): """ Represents a relation with specific configuration attributes. @@ -149,11 +170,22 @@ class Relation(BaseModel): model_config (ConfigDict): Configuration dictionary enforcing a strict model behavior. value (str): The value representing the relation. + propertyIri (Optional[str]): IRI of the ontology property this relation + resolved to. Defaults to None. + canonicalName (Optional[str]): The resolved property's `local_name`, which + is the key typed-property storage writes under. Defaults to None. """ model_config = ConfigDict(strict=True) value: str + # See `Entity` above for why these are declared, optional, and plain `str`. + # `canonicalName` is separate from `propertyIri` because the IRI is the + # identity and the local name is the storage key - a graph property cannot be + # keyed on a full IRI. + propertyIri: Optional[str]=None + canonicalName: Optional[str]=None + EntityType = Union[Entity, str] class Fact(BaseModel): diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py index 4534034b..11758ca9 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py @@ -1,6 +1,8 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import re + EXTRACT_PROPOSITIONS_PROMPT = """ You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph. Your task is to decompose the given text into clear, concise, and context-independent propositions. @@ -230,3 +232,84 @@ Classification3 """ + +# --- Composing an ontology's vocabulary into a prompt ------------------------- +# +# The extraction prompts above carry no {ontology_constraints} placeholder, and +# they are not going to get one. A placeholder cannot be invisible when there is +# nothing to substitute: rendering '' into it still leaves a blank line, which +# changes the prompt text - and therefore the LLMCache key, which is derived +# from prompt.format(**args) - for every user who has never configured an +# ontology. So the vocabulary block is composed into the template instead, and +# only when there is a block to compose, which makes the no-ontology path +# byte-identical by construction rather than by whitespace bookkeeping. + +ONTOLOGY_CONSTRAINTS_PLACEHOLDER = '{ontology_constraints}' + +# Where the block goes in each shipped template, and the string that says so. +# The topics anchor is the final admonition, so the vocabulary lands after the +# instructions it qualifies and before the propositions payload; the +# propositions anchor is the equivalent line in that prompt. Both are documented +# insertion points rather than incidental matches: a custom template that keeps +# either line inherits the same placement. +EXTRACT_TOPICS_ANCHOR = 'Adhere strictly to the provided instructions.' + +EXTRACT_PROPOSITIONS_ANCHOR = ( + 'Do not provide any other explanatory text. Ensure you have captured all of ' + 'the details from the text in your response.' +) + +ONTOLOGY_CONSTRAINT_ANCHORS = (EXTRACT_TOPICS_ANCHOR, EXTRACT_PROPOSITIONS_ANCHOR) + +_FORMAT_FIELD = re.compile(r'\{([A-Za-z_][A-Za-z0-9_.]*)\}') + +def _neutralize_format_fields(text:str) -> str: + """Stop a brace group in inserted text being read as a prompt argument. + + The composed template is rendered by `PromptTemplate.format`, which is not + `str.format`: llama-index substitutes with a regex over `{name}` and leaves + anything it has no argument for exactly as it found it. So a brace arriving + from the ontology - an `rdfs:comment` mentioning JSON, say - is already + harmless, and `{{` is not an escape sequence here; doubling braces would + only put literal `{{` in front of the model. + + The one case that does bite is a comment whose brace group happens to name a + prompt argument, `{text}` above all: that would silently substitute the + chunk into the middle of the vocabulary block, with nothing raised and + nothing logged. Padding the group - `{text}` becomes `{ text }` - takes the + name out of the renderer's reach while still reading as what the ontology + said. Only identifier-shaped groups are touched, so prose and JSON examples + survive unchanged. + """ + return _FORMAT_FIELD.sub(lambda match: f'{{ {match.group(1)} }}', text) + +def with_ontology_constraints(template:str, constraints:str) -> str: + """Insert a rendered ontology constraint block into a prompt template. + + Args: + template: The prompt template, shipped or custom. + constraints: The rendered block, or `''` when no ontology is configured + or its `ontology_authority` is `'off'`. + + Returns: + `template` itself - the same object, not a copy - when `constraints` is + empty. Otherwise a new template with the block inserted at the first of: + an `{ontology_constraints}` placeholder, in which case a custom template + has chosen its own insertion point; the documented anchor for whichever + shipped prompt this is; or the end, so that a custom template which + matches neither still receives the vocabulary rather than silently + dropping it. + """ + if not constraints: + return template + + block = _neutralize_format_fields(constraints) + + if ONTOLOGY_CONSTRAINTS_PLACEHOLDER in template: + return template.replace(ONTOLOGY_CONSTRAINTS_PLACEHOLDER, block) + + for anchor in ONTOLOGY_CONSTRAINT_ANCHORS: + if anchor in template: + return template.replace(anchor, f'{block}\n\n{anchor}', 1) + + return f'{template}\n\n{block}' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/pipeline_utils.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/pipeline_utils.py index 209e3843..3ae0076d 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/pipeline_utils.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/pipeline_utils.py @@ -13,9 +13,10 @@ from llama_index.core.schema import BaseNode, Document from graphrag_toolkit.lexical_graph.config import GraphRAGConfig +from graphrag_toolkit.lexical_graph.logging import apply_logging_config, get_applied_logging_config -def _init_worker(config_snapshot): +def _init_worker(config_snapshot, logging_config=None): """Re-apply the parent's GraphRAGConfig scalars in a spawn-started worker. spawn re-imports config.py in a clean interpreter, so the GraphRAGConfig @@ -24,8 +25,18 @@ def _init_worker(config_snapshot): the ambient role) and mis-placing data (s3_chunk_store -> None falls back to the in-graph chunk store, dropping the intended KMS CMK). Re-applying the snapshot keeps workers consistent with the parent. + + The logging config needs the same treatment for the same reason, and is + passed separately because it is `logging.config.dictConfig` state rather + than a GraphRAGConfig field. Without it a worker's root logger sits at + WARNING with no handler but `lastResort`, so anything an extraction + component logs below WARNING is discarded - and extraction components run + *only* in workers, which makes their INFO logging unreachable rather than + merely quiet. None means the parent never configured logging, in which case + the worker is left at the interpreter default. """ GraphRAGConfig.apply_config_snapshot(config_snapshot) + apply_logging_config(logging_config) def _sink(): @@ -58,11 +69,12 @@ def run_pipeline( # which also drops the GraphRAGConfig singleton's programmatically-set # values - so propagate a picklable snapshot via the worker initializer. config_snapshot = GraphRAGConfig.get_config_snapshot() + logging_config = get_applied_logging_config() with ProcessPoolExecutor( max_workers=num_workers, mp_context=multiprocessing.get_context('spawn'), initializer=_init_worker, - initargs=(config_snapshot,), + initargs=(config_snapshot, logging_config), ) as p: processed_node_batches = p.map(transform, node_batches) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py index ce673150..ba6b8f74 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 import logging -from typing import List, Optional, Union, Any, Dict, overload +from dataclasses import asdict +from typing import List, NamedTuple, Optional, Union, Any, Dict, overload from pipe import Pipe from graphrag_toolkit.lexical_graph import GraphRAGConfig @@ -25,6 +26,8 @@ from graphrag_toolkit.lexical_graph.indexing.extract import TopicExtractor, BatchTopicExtractorSync from graphrag_toolkit.lexical_graph.indexing.extract import ExtractionPipeline from graphrag_toolkit.lexical_graph.indexing.extract import InferClassifications, InferClassificationsConfig +from graphrag_toolkit.lexical_graph.indexing.extract import OntologyType, to_ontology_config +from graphrag_toolkit.lexical_graph.indexing.extract import OntologyFilter from graphrag_toolkit.lexical_graph.indexing.build import BuildPipeline from graphrag_toolkit.lexical_graph.indexing.build import VectorIndexing from graphrag_toolkit.lexical_graph.indexing.build import GraphConstruction @@ -45,6 +48,21 @@ ExtractionLLMType = Union[str, LLM, LLMCache] +class OntologyConstraints(NamedTuple): + """The rendered blocks an ontology contributes to each extraction prompt. + + Named rather than a bare tuple because the two are easy to swap and the + failure would be silent: the propositions prompt would receive the full + vocabulary and the topics prompt only the class names, and both would still + render, extract, and look plausible. + + Attributes: + topics (str): The full vocabulary block, for the topics prompt. + propositions (str): The class names alone, for the propositions prompt. + """ + topics:str + propositions:str + class ExtractionConfig(): """ Configuration for extraction-related operations. @@ -75,7 +93,12 @@ class ExtractionConfig(): be applied during the extraction process. Will be internally converted to a FilterConfig object. extraction_llm (Optional[ExtractionLLMType]): LLM to be used for extracting - propositions and topics. If None, GraphRAGConfig.extract_llm is used. + propositions and topics. If None, GraphRAGConfig.extract_llm is used. + ontology (Optional[OntologyType]): An ontology to guide extraction, as an + OntologyConfig, an Ontology, an rdflib.Graph, or a path to a Turtle file. + Normalized to an OntologyConfig at its default authority level ('align') when + it is not one already. Defaults to None, which leaves extraction exactly + as it was. """ def __init__(self, enable_proposition_extraction: bool = True, @@ -85,8 +108,20 @@ def __init__(self, extract_propositions_prompt_template: Optional[str] = None, extract_topics_prompt_template: Optional[str] = None, extraction_filters: Optional[MetadataFiltersType] = None, - extraction_llm: Optional[ExtractionLLMType] = None): + extraction_llm: Optional[ExtractionLLMType] = None, + ontology: Optional[OntologyType] = None): self.enable_proposition_extraction = enable_proposition_extraction + + # Whether the user named a list, asked for none, or said nothing at all. + # Seeding turns on that distinction, and it is only knowable here: the default is a module-level list, so identity against + # it answers the question exactly, whereas a value check later cannot + # tell a user who passed the defaults from a user who passed nothing. + # `None` counts as set - it is a request for no preferences, and reads + # differently from silence. + self.preferred_entity_classifications_provided = ( + preferred_entity_classifications is not DEFAULT_ENTITY_CLASSIFICATIONS + ) + self.preferred_entity_classifications = preferred_entity_classifications if preferred_entity_classifications is not None else [] self.preferred_topics = preferred_topics if preferred_topics is not None else [] self.infer_entity_classifications = infer_entity_classifications @@ -98,6 +133,40 @@ def __init__(self, else: self.extraction_llm = None + self.ontology = to_ontology_config(ontology) if ontology is not None else None + + if self.ontology is not None: + self._validate_ontology_combination() + + def _validate_ontology_combination(self): + """Reject the one ontology/inference combination that has no coherent reading. + + Inference alongside an ontology is coherent and supported: the ontology's + class names seed `default_classifications` and inference adds the domain + terms the ontology does not declare, which is exactly what 'align' + permits - name what is declared, keep what is not. + + `replace_default_classifications=True` is the exception. It asks for the + seeded names to be discarded, so the user has configured an ontology and + then asked for its vocabulary to be thrown away before the model ever + sees that slot. Silently honouring either half would be a guess, so it + raises instead. + + Raises: + ValueError: If inference is configured to replace the seeded + classifications while an ontology is present. + """ + infer = self.infer_entity_classifications + if isinstance(infer, InferClassificationsConfig) and infer.replace_default_classifications: + raise ValueError( + 'infer_entity_classifications with replace_default_classifications=True ' + 'cannot be combined with an ontology: the ontology seeds the preferred ' + 'entity classifications and replacing them discards the vocabulary the ' + 'ontology was configured to supply. Set ' + 'replace_default_classifications=False to let inference extend the ' + "ontology's classes, or remove the ontology." + ) + class BuildConfig(): """ @@ -272,7 +341,9 @@ class relies on configurable components for batch processing, classification inf extraction_pre_processors (list): List of preprocessing steps used for data processing during the extraction pipeline. extraction_components (list): List of components forming the main data extraction - pipeline, including proposition and topic extractors. + pipeline, including proposition and topic extractors, and - when an + ontology resolves at least one dimension on - an OntologyFilter after + them. allow_batch_inference (bool): Specifies whether batch inference is allowed based on indexing configuration settings. """ @@ -343,38 +414,56 @@ def _configure_extraction_pipeline(self, config: IndexingConfig): for c in config.chunking: components.append(c) + # Rendered once, here in the parent process, and handed to the extractors + # as a plain string. Rendering reads the rdflib graph, which does not + # cross the spawn boundary; a string does. Empty when there is no + # ontology or its authority is 'off', and an empty block leaves every + # prompt byte-identical to what it was. + ontology_constraints = self._render_ontology_constraints(config.extraction.ontology) + if config.extraction.enable_proposition_extraction: if config.batch_config: components.append(BatchLLMPropositionExtractorSync( batch_config=config.batch_config, prompt_template=config.extraction.extract_propositions_prompt_template, - llm=config.extraction.extraction_llm + llm=config.extraction.extraction_llm, + ontology_constraints=ontology_constraints.propositions )) else: components.append(LLMPropositionExtractor( prompt_template=config.extraction.extract_propositions_prompt_template, - llm=config.extraction.extraction_llm + llm=config.extraction.extraction_llm, + ontology_constraints=ontology_constraints.propositions )) entity_classification_provider = None topic_provider = None + # Unchanged, and deliberately ahead of any ontology seeding: with no real + # store there is no scoped-value store to read or write, so both providers + # stay empty. The ontology still reaches the prompt through + # `ontology_constraints` above - what a DummyGraphStore run cannot show is + # the seeded `{preferred_entity_classifications}` slot. if isinstance(self.graph_store, DummyGraphStore): entity_classification_provider = default_preferred_values([]) topic_provider = default_preferred_values([]) else: + # Either the user's list or the ontology's class names, decided once + # so the three branches below cannot disagree about which it is. + preferred_entity_classifications = self._preferred_entity_classifications(config.extraction) + if config.extraction.infer_entity_classifications: if isinstance(config.extraction.infer_entity_classifications, InferClassificationsConfig): - infer_config = config.extraction.infer_entity_classifications + infer_config = config.extraction.infer_entity_classifications else: infer_config = InferClassificationsConfig() default_classifications = [] - if isinstance(config.extraction.preferred_entity_classifications, list): - default_classifications = config.extraction.preferred_entity_classifications + if isinstance(preferred_entity_classifications, list): + default_classifications = preferred_entity_classifications entity_classification_provider = InferClassifications( splitter=SentenceSplitter(chunk_size=256, chunk_overlap=20) if config.chunking else None, @@ -389,10 +478,10 @@ def _configure_extraction_pipeline(self, config: IndexingConfig): pre_processors.append(entity_classification_provider) - elif isinstance(config.extraction.preferred_entity_classifications, list): - entity_classification_provider = default_preferred_values(config.extraction.preferred_entity_classifications) + elif isinstance(preferred_entity_classifications, list): + entity_classification_provider = default_preferred_values(preferred_entity_classifications) else: - entity_classification_provider = config.extraction.preferred_entity_classifications + entity_classification_provider = preferred_entity_classifications if isinstance(config.extraction.preferred_topics, list): topic_provider = default_preferred_values(config.extraction.preferred_topics) @@ -408,7 +497,8 @@ def _configure_extraction_pipeline(self, config: IndexingConfig): entity_classification_provider=entity_classification_provider, topic_provider=topic_provider, prompt_template=config.extraction.extract_topics_prompt_template, - llm=config.extraction.extraction_llm + llm=config.extraction.extraction_llm, + ontology_constraints=ontology_constraints.topics ) else: topic_extractor = TopicExtractor( @@ -416,13 +506,147 @@ def _configure_extraction_pipeline(self, config: IndexingConfig): entity_classification_provider=entity_classification_provider, topic_provider=topic_provider, prompt_template=config.extraction.extract_topics_prompt_template, - llm=config.extraction.extraction_llm + llm=config.extraction.extraction_llm, + ontology_constraints=ontology_constraints.topics ) components.append(topic_extractor) + ontology_filter = self._ontology_filter(config.extraction.ontology) + + if ontology_filter is not None: + components.append(ontology_filter) + return (pre_processors, components) + def _typed_properties(self) -> Optional[str]: + """Where the builders should store coerced attribute values, if anywhere. + + The setting lives on `OntologyConfig` and only there, which is what keeps + it unreachable without one: with no ontology this returns None, `BuildPipeline` + coalesces that to `GraphRAGConfig.typed_properties`, and that is `'off'` + unless something set it programmatically. So a user who never mentioned an + ontology cannot reach a placement that writes, and no environment variable + can reach one on their behalf. + + Returned as None rather than as `'off'` so that the `coalesce` chain in + `BuildPipeline` behaves the same way here as for every other setting - an + unasked-for value defers to the layer below rather than pinning it. + + Returns: + The configured placement, or None when there is no ontology. + """ + ontology_config = self.indexing_config.extraction.ontology + return None if ontology_config is None else ontology_config.typed_properties + + @staticmethod + def _ontology_filter(ontology_config) -> Optional[OntologyFilter]: + """The filter that enforces what the prompt asked for, or None. + + This runs *after* the topic extractor and nowhere else. Everything the + ontology contributes before this point is advisory - a block of text in a + prompt, which a model may ignore - and this is the only component that + makes a level's claim true rather than requested. + + Returns None when there is no ontology, and when every dimension resolves + to False. The second case is not the same as the first: + `ontology_authority='off'` still seeds `{preferred_entity_classifications}` from + the ontology, so the prompt differs from the no-ontology prompt even + though no component here does. What `off` does guarantee is that nothing + rewrites or discards a fact the model produced, and the way it guarantees + it is by this method returning None - not by a filter that runs with every + flag off. A no-op in the pipeline would still round-trip `TOPICS_KEY` + through `model_validate` / `model_dump` and would still annotate, which is + exactly the difference `off` rules out. + + Args: + ontology_config: The normalized `OntologyConfig`, or None. + + Returns: + A configured `OntologyFilter`, or None if it would have nothing to do. + """ + if ontology_config is None or not ontology_config.filter_required(): + return None + + # Spread the resolved dimensions rather than naming them one by one. The + # field names of `ResolvedDimensions` and the flags of `OntologyFilter` + # are deliberately the same six words, and a hand-written argument list + # can omit one - which would leave a gate the user asked for silently not + # running, the one failure in this feature that looks like success. + return OntologyFilter( + index=ontology_config.ontology.index(), + report_violations=ontology_config.report_violations, + **asdict(ontology_config.resolved()), + ) + + @staticmethod + def _render_ontology_constraints(ontology_config) -> OntologyConstraints: + """Render the two prompt blocks an ontology contributes, once. + + Both extraction stages get a block, and they are not the same block: the + topics prompt gets the full vocabulary, the propositions prompt gets the + class names alone, because classifying the entities it names is the only + thing that stage does which an ontology can steer. + + Args: + ontology_config: The normalized `OntologyConfig`, or None. + + Returns: + The two rendered blocks, both empty when there is no ontology. + """ + if ontology_config is None: + return OntologyConstraints('', '') + + ontology = ontology_config.ontology + ontology_authority = ontology_config.ontology_authority + + # `vocabulary_format` reaches the topics block only. The propositions + # block is a class-name list by design - that stage classifies the + # entities it names and extracts nothing else - so there is no property + # vocabulary there for a serialization to present differently. + return OntologyConstraints( + topics=ontology.format_as_prompt_constraint( + ontology_authority, ontology_config.vocabulary_format + ), + propositions=ontology.format_as_proposition_constraint(ontology_authority) + ) + + @staticmethod + def _preferred_entity_classifications(extraction_config: ExtractionConfig) -> PREFERRED_VALUES_PROVIDER_TYPE: + """Decide what fills the `{preferred_entity_classifications}` prompt slot. + + With an ontology and no user list, the slot is seeded + from the ontology's rendered class names, so it cannot name a class + differently from the way the vocabulary block above it does. + + A user who named their own list keeps it, with a + warning. Honouring the ontology instead would discard a setting the user + made deliberately, and merging the two would produce a vocabulary neither + of them asked for. + + Args: + extraction_config: The extraction configuration to read. + + Returns: + The user's value, unchanged, unless an ontology is present and the + user said nothing - in which case the ontology's class names. + """ + if extraction_config.ontology is None: + return extraction_config.preferred_entity_classifications + + if extraction_config.preferred_entity_classifications_provided: + logger.warning( + 'Both an ontology and preferred_entity_classifications were configured. ' + 'Honouring preferred_entity_classifications: %s. The ontology still ' + 'supplies the vocabulary block in the extraction prompt, but its classes ' + 'will not be offered as preferred classifications. Remove ' + 'preferred_entity_classifications to seed that slot from the ontology.', + extraction_config.preferred_entity_classifications + ) + return extraction_config.preferred_entity_classifications + + return extraction_config.ontology.ontology.class_names() + def extract( self, nodes: List[BaseNode] = [], @@ -561,6 +785,7 @@ def build( source_metadata_formatter=build_config.source_metadata_formatter, include_domain_labels=build_config.include_domain_labels, include_local_entities=build_config.include_local_entities, + typed_properties=self._typed_properties(), tenant_id=self.tenant_id, progress_monitor=progress_monitor, **kwargs @@ -631,6 +856,7 @@ def extract_and_build( source_metadata_formatter=build_config.source_metadata_formatter, include_domain_labels=build_config.include_domain_labels, include_local_entities=build_config.include_local_entities, + typed_properties=self._typed_properties(), tenant_id=self.tenant_id, progress_monitor=progress_monitor, **kwargs diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/logging.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/logging.py index d2895f94..d8fcee52 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/logging.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/logging.py @@ -1,6 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import copy import logging import logging.config import warnings @@ -271,7 +272,11 @@ def set_advanced_logging_config( if isinstance(logging_level, int): logging_level = logging.getLevelName(logging_level) - config = BASE_LOGGING_CONFIG.copy() + # deepcopy, not copy: the mutations below reach into nested dicts and into + # `loggers['']['handlers']`, all of which a shallow copy still shares with + # BASE_LOGGING_CONFIG. With a shallow copy, calling this twice accumulates + # handlers and leaks one call's module filters and log filename into the next. + config = copy.deepcopy(BASE_LOGGING_CONFIG) config['loggers']['']['level'] = logging_level.upper() config['filters']['moduleFilter']['included_modules'].update(included_modules or dict()) config['filters']['moduleFilter']['excluded_modules'].update(excluded_modules or dict()) @@ -285,7 +290,45 @@ def set_advanced_logging_config( config['handlers']['file_handler']['filename'] = filename config['loggers']['']['handlers'].append('file_handler') + apply_logging_config(config) + + +# The last config applied in this process, kept so that spawn-started workers can +# be given the same one. `logging.config.dictConfig` is global interpreter state +# rather than a `GraphRAGConfig` field, so it does not travel in the config +# snapshot, and a worker that has not had it applied sits at the root logger's +# default WARNING with no handler but `lastResort`. Anything a component logs at +# INFO from inside extraction - which is where extraction components run - then +# goes nowhere at all. +_applied_logging_config: Optional[Dict] = None + + +def get_applied_logging_config() -> Optional[Dict]: + """The logging config applied in this process, or None if none ever was. + + Returns: + The `dictConfig` dictionary last passed to `apply_logging_config`, or None + if the caller never configured logging - in which case a worker should be + left at the interpreter default rather than given one. + """ + return _applied_logging_config + + +def apply_logging_config(config: Optional[Dict]) -> None: + """Apply a `dictConfig` dictionary and remember it. + + Args: + config: The dictionary to apply. None is a no-op, so that propagating + "the parent never configured logging" needs no special case at the + call site. + """ + global _applied_logging_config + + if config is None: + return + logging.config.dictConfig(config) + _applied_logging_config = config def _is_valid_logging_level(level: Union[str, LoggingLevel]) -> bool: diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt b/lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt index 3f2b430d..3e1de1ec 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt @@ -14,6 +14,7 @@ lru-dict==1.3.0 nltk<3.10.3 pipe==2.2 python-dotenv==1.2.2 +rdflib>=7.0,<8.0 smart_open==7.1.0 spacy==3.8.7 tfidf_matcher==0.3.0 diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_utils.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_utils.py index 4d1ffd8c..460d0971 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_utils.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_utils.py @@ -13,6 +13,13 @@ SEARCH_STRING_PATTERN = re.compile(r'([^\s\w]|_)+') +# Split where an uppercase letter follows a lowercase one, and nowhere else. +# Deliberately the same rule as `_CAMEL_BOUNDARY` in the ontology's naming +# module, and duplicated rather than imported: this is a leaf utility, and +# importing from `indexing.extract.ontology` would pull rdflib and the whole +# extract package into the build and storage paths to save one regex. +CAMEL_BOUNDARY_PATTERN = re.compile(r'(?<=[a-z])(?=[A-Z])') + def new_query_var(): return f'n{uuid.uuid4().hex}' @@ -113,17 +120,31 @@ def relationship_name_from(value:str): """ Generates a formatted relationship name from a given string. - The function transforms the input string by replacing all non-alphanumeric - characters with underscores and converts the resulting string to uppercase. + The function splits camel-case words, replaces all non-alphanumeric + characters with underscores, and converts the result to uppercase. + + The camel-case split exists for ontology-normalized predicates. Without an + ontology a predicate arrives from the response parser as `WORKS FOR`, which + becomes `WORKS_FOR`; with `normalize_names` on it arrives as the authored + `worksFor`, which without the split would become `WORKSFOR` and lose the + word boundary. The only live caller is `GraphSummaryBuilder`, whose value is + read back by `GraphSummary._get_paths` and rendered into a domain-summary + prompt as `(Person)-[WORKS_FOR]->(Company)` - so the boundary is the + difference between a legible path and one word for the LLM to guess at. + + The rule is narrow on purpose: an uppercase letter following a *lowercase* + one, so `HTTPServer` and `Company2X` are left alone. Inputs with no such + boundary - anything already upper case, or spaced - are unaffected. Args: value (str): The input string to be processed. Returns: - str: A formatted string where non-alphanumeric characters are replaced - with underscores and all characters are in uppercase. + str: A formatted string where camel-case boundaries and non-alphanumeric + characters become underscores and all characters are in uppercase. """ - return ''.join([ c if c.isalnum() else '_' for c in value ]).upper() + split = CAMEL_BOUNDARY_PATTERN.sub('_', value) + return ''.join([ c if c.isalnum() else '_' for c in split ]).upper() def node_result(node_ref:str, node_id:Optional[NodeId]=None, diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/neptune_graph_stores.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/neptune_graph_stores.py index 11889ed2..01b05656 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/neptune_graph_stores.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/neptune_graph_stores.py @@ -171,7 +171,16 @@ def create_property_assigment_fn_for_neptune(key:str, value:Any) -> Callable[[st try: format_datetime(value) return lambda x: f'datetime({x})' - except ValueError as e: + except (TypeError, ValueError): + # `TypeError` as well as `ValueError`, because `is_datetime_key` keys on + # the property *name* and says nothing about the value's type, while + # `format_datetime` hands a non-string straight to `dateutil.parse`, + # which raises `TypeError`. Unreachable while every property came from + # document metadata, where values are strings; typed properties + # (`typed_properties`) are the first source of native ints, floats and + # bools, so an ontology property named `founded_date` with a numeric + # range would otherwise abort the build here rather than fall back to + # a plain assignment. return lambda x: x else: return lambda x: x diff --git a/lexical-graph/tests/fixtures/ontologies/company.ttl b/lexical-graph/tests/fixtures/ontologies/company.ttl new file mode 100644 index 00000000..55248f40 --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/company.ttl @@ -0,0 +1,110 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# The test ontology the ontology-guided-extraction tests run on. Small enough to +# read in one screen, and extended to cover every shape the feature has to +# handle: +# +# * a three-level hierarchy :SportsTeam < :Company < :Agent +# * a label that differs from the local name :SportsTeam -> "Sports Team" +# * skos:altLabel on a class and on a property +# * an object property with no domain and no range :acquired +# * rdfs:comment descriptions on some, deliberately not all, terms +# * datatype properties covering xsd:integer, xsd:double, xsd:boolean, +# xsd:date and xsd:string +# * a datatype property with no domain :officialName +# +# The spellings - :worksFor, :SportsTeam, :foundedYear, :revenue - are fixed, +# because the naming-contract tests assert against those exact names. + +@prefix : . +@prefix owl: . +@prefix rdfs: . +@prefix skos: . +@prefix xsd: . + + a owl:Ontology ; + rdfs:comment "Test ontology for ontology-guided extraction." . + +# +# Classes. :Agent is the root; :SportsTeam is three deep. +# + +:Agent a owl:Class ; + rdfs:comment "Anything that can act - a person or an organization." . + +:Person a owl:Class ; rdfs:subClassOf :Agent ; + rdfs:comment "An individual human." . + +:Athlete a owl:Class ; rdfs:subClassOf :Person ; + rdfs:comment "A person who competes in a sport professionally." . + +:Company a owl:Class ; rdfs:subClassOf :Agent ; + rdfs:label "Company" ; + skos:altLabel "Corporation" ; + rdfs:comment "An incorporated commercial organization." . + +# The label differs from the local name on purpose: rendered into the prompt it +# has to survive the response parser's .title(), and stored it has to come back +# as the authored SportsTeam. +:SportsTeam a owl:Class ; rdfs:subClassOf :Company ; + rdfs:label "Sports Team" ; + skos:altLabel "Ball Club" ; + rdfs:comment "A professional team that competes in a league." . + +# +# Object properties. +# + +:worksFor a owl:ObjectProperty ; + rdfs:domain :Person ; rdfs:range :Company ; + skos:altLabel "REQ_TO_HC" ; + rdfs:comment "Employment of a person by a company." . + +# Domain :Athlete, range :SportsTeam - narrower than :worksFor, so a fact +# stated of a :Person and a :Company only satisfies it through the subclass +# closure. +:playsFor a owl:ObjectProperty ; + rdfs:domain :Athlete ; rdfs:range :SportsTeam ; + skos:altLabel "PLAYS_ON" ; + rdfs:comment "Membership of an athlete in a sports team." . + +:subsidiaryOf a owl:ObjectProperty ; + rdfs:domain :Company ; rdfs:range :Company ; + rdfs:comment "Ownership of one company by another." . + +# No domain and no range: matches any subject and any object. +:acquired a owl:ObjectProperty . + +# +# Datatype properties. Between them these cover every XSD range the feature +# claims to coerce. +# + +:foundedYear a owl:DatatypeProperty ; + rdfs:domain :Company ; rdfs:range xsd:integer ; + rdfs:comment "The four-digit year the company was founded." . + +:revenue a owl:DatatypeProperty ; + rdfs:domain :Company ; rdfs:range xsd:double ; + rdfs:comment "Annual revenue in US dollars." . + +:isPubliclyTraded a owl:DatatypeProperty ; + rdfs:domain :Company ; rdfs:range xsd:boolean ; + rdfs:comment "Whether the company's shares are listed on a public exchange." . + +:incorporatedOn a owl:DatatypeProperty ; + rdfs:domain :Company ; rdfs:range xsd:date ; + rdfs:comment "The date of incorporation." . + +:tickerSymbol a owl:DatatypeProperty ; + rdfs:domain :Company ; rdfs:range xsd:string ; + skos:altLabel "Stock Symbol" . + +:jobTitle a owl:DatatypeProperty ; + rdfs:domain :Person ; rdfs:range xsd:string ; + rdfs:comment "The person's role at their employer." . + +# No domain: any subject may carry it. +:officialName a owl:DatatypeProperty ; + rdfs:range xsd:string . diff --git a/lexical-graph/tests/fixtures/ontologies/malformed/dangling_domain_reference.ttl b/lexical-graph/tests/fixtures/ontologies/malformed/dangling_domain_reference.ttl new file mode 100644 index 00000000..6a795969 --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/malformed/dangling_domain_reference.ttl @@ -0,0 +1,14 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# An rdfs:domain naming an undeclared class. Left unchecked +# this makes the property unusable - no emitted fact can ever satisfy a domain +# no entity can be classified into. + +@prefix : . +@prefix owl: . +@prefix rdfs: . + +:Company a owl:Class . + +:worksFor a owl:ObjectProperty ; rdfs:domain :Employee ; rdfs:range :Company . diff --git a/lexical-graph/tests/fixtures/ontologies/malformed/dangling_range_reference.ttl b/lexical-graph/tests/fixtures/ontologies/malformed/dangling_range_reference.ttl new file mode 100644 index 00000000..3d08d0e4 --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/malformed/dangling_range_reference.ttl @@ -0,0 +1,12 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# The same fault on the range side. + +@prefix : . +@prefix owl: . +@prefix rdfs: . + +:Person a owl:Class . + +:worksFor a owl:ObjectProperty ; rdfs:domain :Person ; rdfs:range :Employer . diff --git a/lexical-graph/tests/fixtures/ontologies/malformed/dangling_subclass_reference.ttl b/lexical-graph/tests/fixtures/ontologies/malformed/dangling_subclass_reference.ttl new file mode 100644 index 00000000..0b1f67f0 --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/malformed/dangling_subclass_reference.ttl @@ -0,0 +1,11 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Rdfs:subClassOf points at :Organization, which is never +# declared an owl:Class - almost always a typo or a missing import. + +@prefix : . +@prefix owl: . +@prefix rdfs: . + +:Company a owl:Class ; rdfs:subClassOf :Organization . diff --git a/lexical-graph/tests/fixtures/ontologies/malformed/dual_typed_property.ttl b/lexical-graph/tests/fixtures/ontologies/malformed/dual_typed_property.ttl new file mode 100644 index 00000000..5d5680af --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/malformed/dual_typed_property.ttl @@ -0,0 +1,16 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# One IRI declared both an owl:ObjectProperty and an +# owl:DatatypeProperty. The filter would have to decide whether the object is an +# entity or a literal, and there is no non-arbitrary answer. + +@prefix : . +@prefix owl: . +@prefix rdfs: . +@prefix xsd: . + +:Company a owl:Class . + +:founder a owl:ObjectProperty , owl:DatatypeProperty ; + rdfs:domain :Company ; rdfs:range xsd:string . diff --git a/lexical-graph/tests/fixtures/ontologies/malformed/malformed_syntax.ttl b/lexical-graph/tests/fixtures/ontologies/malformed/malformed_syntax.ttl new file mode 100644 index 00000000..5a8a81dd --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/malformed/malformed_syntax.ttl @@ -0,0 +1,7 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Not parseable Turtle. The prefix is never declared and the +# final statement is unterminated. + +:Company a owl:Class ; diff --git a/lexical-graph/tests/fixtures/ontologies/malformed/missing_datatype_range.ttl b/lexical-graph/tests/fixtures/ontologies/malformed/missing_datatype_range.ttl new file mode 100644 index 00000000..e8a4754d --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/malformed/missing_datatype_range.ttl @@ -0,0 +1,13 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Absence rather than the wrong kind: a datatype property with +# no rdfs:range at all. + +@prefix : . +@prefix owl: . +@prefix rdfs: . + +:Company a owl:Class . + +:foundedYear a owl:DatatypeProperty ; rdfs:domain :Company . diff --git a/lexical-graph/tests/fixtures/ontologies/malformed/non_xsd_range.ttl b/lexical-graph/tests/fixtures/ontologies/malformed/non_xsd_range.ttl new file mode 100644 index 00000000..64a52db8 --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/malformed/non_xsd_range.ttl @@ -0,0 +1,15 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# A datatype property whose rdfs:range is not an XSD datatype. +# There is nothing to coerce a literal to, and nothing to name in the prompt's +# attribute section. + +@prefix : . +@prefix owl: . +@prefix rdfs: . + +:Company a owl:Class . + +:foundedYear a owl:DatatypeProperty ; + rdfs:domain :Company ; rdfs:range . diff --git a/lexical-graph/tests/fixtures/ontologies/malformed/subclass_cycle.ttl b/lexical-graph/tests/fixtures/ontologies/malformed/subclass_cycle.ttl new file mode 100644 index 00000000..61632be4 --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/malformed/subclass_cycle.ttl @@ -0,0 +1,13 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# An rdfs:subClassOf cycle. Three classes rather than two, so +# the error message has to name the ring in order rather than just the pair. + +@prefix : . +@prefix owl: . +@prefix rdfs: . + +:Company a owl:Class ; rdfs:subClassOf :Employer . +:Employer a owl:Class ; rdfs:subClassOf :Business . +:Business a owl:Class ; rdfs:subClassOf :Company . diff --git a/lexical-graph/tests/fixtures/ontologies/upper_snake_names.ttl b/lexical-graph/tests/fixtures/ontologies/upper_snake_names.ttl new file mode 100644 index 00000000..ca714f58 --- /dev/null +++ b/lexical-graph/tests/fixtures/ontologies/upper_snake_names.ttl @@ -0,0 +1,33 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# The same vocabulary as company.ttl's core, authored in the other convention a +# real ontology might use: UPPER_SNAKE property names and an already-spaced +# class label. The naming contract says :WORKS_FOR and :worksFor must resolve +# identically, and whichever the author wrote is what gets stored - which can +# only be tested against an ontology that actually writes it the other way. +# +# Kept as a separate file rather than merged into company.ttl: :worksFor and +# :WORKS_FOR share a resolution key, so declaring both in one ontology would +# make the resolution deliberately ambiguous. + +@prefix : . +@prefix owl: . +@prefix rdfs: . +@prefix skos: . +@prefix xsd: . + + a owl:Ontology . + +:Person a owl:Class . +:Company a owl:Class . +# Authored with the spaces already in the label, and an underscored altLabel. +:SPORTS_TEAM a owl:Class ; rdfs:subClassOf :Company ; + rdfs:label "Sports Team" ; + skos:altLabel "BALL_CLUB" . + +:WORKS_FOR a owl:ObjectProperty ; + rdfs:domain :Person ; rdfs:range :Company . + +:FOUNDED_YEAR a owl:DatatypeProperty ; + rdfs:domain :Company ; rdfs:range xsd:integer . diff --git a/lexical-graph/tests/unit/conftest.py b/lexical-graph/tests/unit/conftest.py index 444ae1ef..a2cf6d90 100644 --- a/lexical-graph/tests/unit/conftest.py +++ b/lexical-graph/tests/unit/conftest.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from pathlib import Path from unittest.mock import Mock, patch, MagicMock from graphrag_toolkit.lexical_graph.tenant_id import TenantId from graphrag_toolkit.lexical_graph.indexing.id_generator import IdGenerator @@ -75,3 +76,75 @@ def custom_id_gen(custom_tenant): Fixture for custom ID generator (backward compatible mode, no delimiter). ''' return IdGenerator(tenant_id=custom_tenant, include_classification_in_entity_id=True, use_chunk_id_delimiter=False) + + +# --- Ontology fixtures ------------------------------------------------------- +# +# The .ttl files and the verification corpus live under tests/fixtures/ and are +# reached through these fixtures rather than through paths built in each test, +# so moving the directory is a one-line change. Declared here rather than in a +# leaf conftest because tests across indexing/extract, indexing/build and the +# model tests all read the same ontology. + +@pytest.fixture +def fixtures_dir(): + ''' + Fixture for the tests/fixtures directory. + ''' + return Path(__file__).parent.parent / 'fixtures' + + +@pytest.fixture +def ontology_fixtures_dir(fixtures_dir): + ''' + Fixture for the directory holding the test ontology .ttl files. + ''' + return fixtures_dir / 'ontologies' + + +@pytest.fixture +def company_ttl_path(ontology_fixtures_dir): + ''' + Fixture for the path to the main test ontology. + ''' + return ontology_fixtures_dir / 'company.ttl' + + +@pytest.fixture +def company_ontology(company_ttl_path): + ''' + Fixture for the loaded main test ontology. Imported lazily so that tests + which do not touch the ontology do not pay for importing rdflib. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import Ontology + return Ontology.from_turtle(company_ttl_path) + + +@pytest.fixture +def upper_snake_ttl_path(ontology_fixtures_dir): + ''' + Fixture for the path to the UPPER_SNAKE-authored test ontology. + ''' + return ontology_fixtures_dir / 'upper_snake_names.ttl' + + +@pytest.fixture +def upper_snake_ontology(upper_snake_ttl_path): + ''' + Fixture for the loaded UPPER_SNAKE-authored test ontology, which declares + :WORKS_FOR and :SPORTS_TEAM where company.ttl declares :worksFor and + :SportsTeam. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import Ontology + return Ontology.from_turtle(upper_snake_ttl_path) + + +@pytest.fixture +def malformed_ttl_path(ontology_fixtures_dir): + ''' + Fixture returning a function that resolves a malformed ontology fixture by + name, e.g. malformed_ttl_path('subclass_cycle'). + ''' + def resolve(name): + return ontology_fixtures_dir / 'malformed' / f'{name}.ttl' + return resolve diff --git a/lexical-graph/tests/unit/indexing/build/test_fact_graph_builder.py b/lexical-graph/tests/unit/indexing/build/test_fact_graph_builder.py index 3357c4fa..e08045ab 100644 --- a/lexical-graph/tests/unit/indexing/build/test_fact_graph_builder.py +++ b/lexical-graph/tests/unit/indexing/build/test_fact_graph_builder.py @@ -1,84 +1,169 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import pytest +"""Unit tests for `FactGraphBuilder`. + +`Fact.subject` and `Fact.predicate` are an `Entity` and a `Relation`, not strings, +and `build` reads `include_local_entities` with a hard subscript. The earlier +version of this file passed strings and omitted the kwarg, so every test in it +raised; nothing noticed, because pytest's default `norecursedirs` contains +`build`, so this whole directory was skipped unless its path was named explicitly. +""" + from unittest.mock import Mock + from graphrag_toolkit.lexical_graph.indexing.build.fact_graph_builder import FactGraphBuilder +from llama_index.core.schema import TextNode + +RELATION_FACT = { + 'factId': 'fact-1', + 'statementId': 'stmt-1', + 'subject': {'entityId': 's-1', 'value': 'GraphRAG', 'classification': 'Framework'}, + 'predicate': {'value': 'combines'}, + 'object': {'entityId': 'o-1', 'value': 'knowledge graphs', 'classification': 'Technology'}, +} + +ATTRIBUTE_FACT = { + 'factId': 'fact-2', + 'statementId': 'stmt-1', + 'subject': {'entityId': 's-1', 'value': 'GraphRAG', 'classification': 'Framework'}, + 'predicate': {'value': 'RELEASE YEAR'}, + 'complement': {'entityId': 'c-1', 'value': '2024', 'classification': '__Local_Entity__'}, +} + +def fact_node(fact:dict, text:str='GraphRAG combines knowledge graphs') -> TextNode: + node = TextNode(text=text, id_=fact.get('factId', 'fact-1')) + node.metadata = {'fact': fact} + return node + +def graph_client() -> Mock: + client = Mock() + client.node_id = Mock(side_effect=lambda field: f'params.{field}') + client.execute_query_with_retry = Mock(return_value=[]) + return client + +def queries(client:Mock) -> list: + return [call.args[0] for call in client.execute_query_with_retry.call_args_list] class TestFactGraphBuilderInitialization: """Tests for FactGraphBuilder initialization.""" - + def test_initialization(self): """Verify FactGraphBuilder initializes correctly.""" - builder = FactGraphBuilder() - assert builder is not None + assert FactGraphBuilder() is not None + def test_index_key(self): + """Verify the builder claims the fact index.""" + assert FactGraphBuilder.index_key() == 'fact' class TestFactGraphBuilding: """Tests for fact graph building functionality.""" - - def test_build_fact_node(self, mock_neptune_store): - """Verify building fact node with metadata.""" - builder = FactGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'fact_001' - mock_node.text = 'GraphRAG is an AI framework' - mock_node.metadata = { - 'fact': { - 'factId': 'fact_001', - 'subject': 'GraphRAG', - 'predicate': 'is', - 'object': 'AI framework', - 'metadata': {'confidence': 0.95} - } + + def test_build_fact_node(self): + """Verify the fact itself is written, with the node's text as its value.""" + client = graph_client() + + FactGraphBuilder().build(fact_node(RELATION_FACT), client, include_local_entities=False) + + (query, params) = client.execute_query_with_retry.call_args_list[0].args + assert '__SUPPORTS__' in query + assert params['params'][0] == { + 'statement_id': 'stmt-1', + 'fact_id': 'fact-1', + 'fact': 'GraphRAG combines knowledge graphs', } - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='factId') - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert mock_graph_client.execute_query_with_retry.called - + + def test_build_links_subject_and_object_entities(self): + """Verify both ends of a relation fact are linked to it.""" + client = graph_client() + + FactGraphBuilder().build(fact_node(RELATION_FACT), client, include_local_entities=False) + + bound = [ + call.args[1]['params'][0] + for call in client.execute_query_with_retry.call_args_list + if call.args[1]['params'] and 'entity_id' in call.args[1]['params'][0] + ] + + assert {b['entity_id'] for b in bound} == {'s-1', 'o-1'} + + def test_complement_is_linked_only_with_local_entities(self): + """Verify the complement is skipped when local entities are excluded. + + The complement node is created by `EntityGraphBuilder` only when + `include_local_entities` is on, so linking to it otherwise would point at + a node that does not exist. + """ + for (include_local_entities, expected) in [(False, {'s-1'}), (True, {'s-1', 'c-1'})]: + client = graph_client() + + FactGraphBuilder().build( + fact_node(ATTRIBUTE_FACT, text='GraphRAG RELEASE YEAR 2024'), + client, + include_local_entities=include_local_entities + ) + + bound = [ + call.args[1]['params'][0] + for call in client.execute_query_with_retry.call_args_list + if call.args[1]['params'] and 'entity_id' in call.args[1]['params'][0] + ] + + assert {b['entity_id'] for b in bound} == expected, include_local_entities + def test_build_multiple_facts(self): """Verify building multiple fact nodes.""" + client = graph_client() builder = FactGraphBuilder() - - facts = [ - Mock(metadata={'fact': {'factId': 'f1', 'subject': 'A', 'predicate': 'is', 'object': 'B'}}), - Mock(metadata={'fact': {'factId': 'f2', 'subject': 'C', 'predicate': 'has', 'object': 'D'}}) - ] - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='factId') - mock_graph_client.execute_query_with_retry = Mock() - - for fact in facts: - fact.node_id = fact.metadata['fact']['factId'] - fact.text = f"{fact.metadata['fact']['subject']} {fact.metadata['fact']['predicate']} {fact.metadata['fact']['object']}" - fact.relationships = {} - builder.build(fact, mock_graph_client) - - assert mock_graph_client.execute_query_with_retry.call_count >= 2 + for fact in (RELATION_FACT, ATTRIBUTE_FACT): + builder.build(fact_node(fact), client, include_local_entities=False) + + assert len([q for q in queries(client) if '__SUPPORTS__' in q]) == 2 + + def test_build_tolerates_typed_properties(self): + """Verify the kwarg this builder does not read is accepted anyway. + + Every builder receives the same kwargs, so a setting only + `EntityGraphBuilder` acts on must not break the others. + """ + client = graph_client() + + FactGraphBuilder().build( + fact_node(RELATION_FACT), + client, + include_local_entities=True, + typed_properties='subject', + ) + + assert client.execute_query_with_retry.called class TestFactGraphBuilderErrorHandling: """Tests for fact graph builder error handling.""" - - def test_build_with_missing_fact_id(self): - """Verify handling of fact with missing ID.""" - builder = FactGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'fact_001' - mock_node.metadata = {'fact': {'subject': 'A', 'predicate': 'is', 'object': 'B'}} - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert not mock_graph_client.execute_query_with_retry.called + + def test_build_with_no_fact_metadata_writes_nothing(self): + """Verify a node carrying no fact is left alone.""" + client = graph_client() + + node = TextNode(text='', id_='fact-1') + node.metadata = {} + + FactGraphBuilder().build(node, client, include_local_entities=False) + + assert not client.execute_query_with_retry.called + + def test_build_requires_include_local_entities(self): + """Verify the kwarg is read with a hard subscript. + + Pinned deliberately: `BuildPipeline` always supplies it, but a caller who + does not gets a `KeyError` rather than a default. The same read in + `EntityGraphBuilder` was made tolerant for `typed_properties`; this one was + left alone, and a change either way should be a visible decision. + """ + import pytest + + client = graph_client() + + with pytest.raises(KeyError): + FactGraphBuilder().build(fact_node(RELATION_FACT), client) diff --git a/lexical-graph/tests/unit/indexing/build/test_local_entity_rewrites_typed_carry.py b/lexical-graph/tests/unit/indexing/build/test_local_entity_rewrites_typed_carry.py new file mode 100644 index 00000000..f3f4592f --- /dev/null +++ b/lexical-graph/tests/unit/indexing/build/test_local_entity_rewrites_typed_carry.py @@ -0,0 +1,146 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Carrying typed values across a local-entity rewrite. + +`LocalEntityRewritesGraphBuilder` folds a complement node into a real entity when +the two turn out to be the same thing, and the fold `DETACH DELETE`s the complement. +At a complement placement that node is where `typed_value` and `datatype` live, so +they have to be copied onto the surviving node while the doomed one is still bound. + +The builder emits Cypher rather than executing it, so these tests read the emitted +query text through a recording stand-in for the graph store. That is the whole of +what can be checked without a live store, and it is the part that regresses: the +clause is appended by string concatenation, and it is appended *conditionally*. +""" + +import logging + +import pytest +from llama_index.core.schema import TextNode + +from graphrag_toolkit.lexical_graph.indexing.build.local_entity_rewrites_graph_builder import ( + LocalEntityRewritesGraphBuilder, +) +from graphrag_toolkit.lexical_graph.indexing.model import Entity, Fact, Relation + +XSD = 'http://www.w3.org/2001/XMLSchema#' +LOCAL = '__Local_Entity__' + +class RecordingGraphStore: + """The two methods this builder asks a store for.""" + + def __init__(self): + self.trees = [] + + def node_id(self, name): + return name + + def execute_query_with_retry(self, query_tree, params, **kwargs): + self.trees.append((query_tree, params)) + +def fact(subject_class='Company', complement_class=LOCAL): + return Fact( + factId='f1', + subject=Entity(entityId='e1', value='Meridian Freight', classification=subject_class), + predicate=Relation(value='foundedYear', canonicalName='foundedYear'), + complement=Entity( + entityId='c1', altEntityId='e1', value='1994', + classification=complement_class, datatype=f'{XSD}integer', + ), + ) + +def build(typed_properties=None, include_local_entities=True, node_fact=None): + store = RecordingGraphStore() + kwargs = {'include_local_entities': include_local_entities} + if typed_properties is not None: + kwargs['typed_properties'] = typed_properties + + metadata = {} if node_fact is False else {'fact': (node_fact or fact()).model_dump()} + LocalEntityRewritesGraphBuilder().build(TextNode(text='x', metadata=metadata), store, **kwargs) + return store + +def queries(store): + """Every query in every tree the builder issued, as text.""" + found = [] + + def walk(query): + found.append(query.query) + for child in getattr(query, 'child_queries', None) or []: + walk(child) + + for (tree, _) in store.trees: + walk(tree.root_query) + return found + +def copy_query(store): + return next(q for q in queries(store) if 'copy complement relationships' in q) + +class TestTheCarryClause: + + @pytest.mark.parametrize('placement', ['complement', 'both']) + def test_a_complement_placement_carries_both_properties(self, placement): + query = copy_query(build(placement)) + + assert 'SET n.`typed_value`' in query + assert 'n.`datatype`' in query + + @pytest.mark.parametrize('placement', ['complement', 'both']) + def test_the_carry_is_first_writer_wins(self, placement): + """Two complements can fold into the same entity, and the second must not + overwrite the first. Arbitrary, but stable within a run - the alternative + is a value that changes with build order.""" + query = copy_query(build(placement)) + + assert 'coalesce(n.`typed_value`, c.`typed_value`)' in query + assert 'coalesce(n.`datatype`, c.`datatype`)' in query + + @pytest.mark.parametrize('placement', [None, 'off', 'subject']) + def test_no_clause_is_emitted_where_nothing_was_asked_for(self, placement): + """`SET x = null` deletes the property in some stores, it changes the query + text every existing user sends, and it is write work nobody asked for. So + the clause is absent rather than emitted as a harmless no-op - which is + also what keeps the query byte-identical for a caller who never set the + kwarg at all.""" + query = copy_query(build(placement)) + + assert 'typed_value' not in query + assert 'coalesce' not in query + + def test_the_query_is_byte_identical_across_the_placements_that_do_not_carry(self): + assert copy_query(build(None)) == copy_query(build('off')) + assert copy_query(build('off')) == copy_query(build('subject')) + + def test_only_the_copy_query_changes(self): + """The delete half of the fold is not a function of the placement.""" + for placement in ('off', 'complement', 'both'): + deletes = [q for q in queries(build(placement)) if 'delete complement relationships' in q] + assert deletes == [q for q in queries(build('off')) if 'delete complement relationships' in q] + +class TestWhatTheBuilderIssues: + + def test_both_directions_of_the_fold_are_attempted(self): + """A subject may turn out to be a local entity, and a complement may turn + out to be a real one, so the builder asks about each.""" + store = build('complement') + + assert len(store.trees) == 2 + assert any('matching subject' in q for q in queries(store)) + assert any('matching complement' in q for q in queries(store)) + + def test_a_local_entity_subject_is_skipped_when_local_entities_are_off(self): + store = build('complement', include_local_entities=False, node_fact=fact(subject_class=LOCAL)) + + assert store.trees == [] + + def test_a_local_entity_subject_is_processed_when_they_are_on(self): + store = build('complement', include_local_entities=True, node_fact=fact(subject_class=LOCAL)) + + assert store.trees + + def test_a_node_carrying_no_fact_warns_and_writes_nothing(self, caplog): + with caplog.at_level(logging.WARNING): + store = build('complement', node_fact=False) + + assert store.trees == [] + assert 'fact_id missing' in caplog.text diff --git a/lexical-graph/tests/unit/indexing/build/test_statement_graph_builder.py b/lexical-graph/tests/unit/indexing/build/test_statement_graph_builder.py index 57c309c8..643b5f8d 100644 --- a/lexical-graph/tests/unit/indexing/build/test_statement_graph_builder.py +++ b/lexical-graph/tests/unit/indexing/build/test_statement_graph_builder.py @@ -1,132 +1,174 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import pytest +"""Unit tests for `StatementGraphBuilder`. + +The `Statement` metadata here is spelled as the model actually defines it - +`value`, `details`, `chunkId` - not `text` and `entities`. The earlier version of +this file used the latter and every test in it raised `ValidationError`; nothing +noticed, because pytest's default `norecursedirs` contains `build`, so this whole +directory was skipped unless its path was named explicitly. +""" + from unittest.mock import Mock + from graphrag_toolkit.lexical_graph.indexing.build.statement_graph_builder import StatementGraphBuilder +from llama_index.core.schema import NodeRelationship, RelatedNodeInfo, TextNode + +def statement_node(statement:dict, node_id:str='stmt-1') -> TextNode: + node = TextNode(text=statement.get('value', ''), id_=node_id) + node.metadata = {'statement': statement} + return node + +def graph_client() -> Mock: + client = Mock() + client.node_id = Mock(side_effect=lambda field: f'params.{field}') + client.execute_query_with_retry = Mock(return_value=[]) + return client class TestStatementGraphBuilderInitialization: """Tests for StatementGraphBuilder initialization.""" - + def test_initialization(self): """Verify StatementGraphBuilder initializes correctly.""" - builder = StatementGraphBuilder() - assert builder is not None + assert StatementGraphBuilder() is not None + def test_index_key(self): + """Verify the builder claims the statement index.""" + assert StatementGraphBuilder.index_key() == 'statement' class TestStatementGraphBuilding: """Tests for statement graph building functionality.""" - - def test_build_statement_node(self, mock_neptune_store): - """Verify building statement node with metadata.""" - builder = StatementGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'stmt_001' - mock_node.text = 'GraphRAG combines knowledge graphs with RAG' - mock_node.metadata = { - 'statement': { - 'statementId': 'stmt_001', - 'text': 'GraphRAG combines knowledge graphs with RAG', - 'metadata': { - 'confidence': 0.92, - 'source_chunk': 'chunk_001' - } - } - } - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='statementId') - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert mock_graph_client.execute_query_with_retry.called - - def test_build_statement_with_entities(self): - """Verify building statement with entity references.""" - builder = StatementGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'stmt_002' - mock_node.text = 'Knowledge graphs store structured information' - mock_node.metadata = { - 'statement': { - 'statementId': 'stmt_002', - 'text': 'Knowledge graphs store structured information', - 'entities': ['entity_001', 'entity_002'], - 'metadata': {} - } + + def test_build_statement_node(self): + """Verify building a statement node writes its value and details.""" + client = graph_client() + + StatementGraphBuilder().build( + statement_node({ + 'statementId': 'stmt-1', + 'value': 'GraphRAG combines knowledge graphs with RAG', + 'details': ['first detail', 'second detail'], + }), + client + ) + + client.execute_query_with_retry.assert_called_once() + (_, params) = client.execute_query_with_retry.call_args.args + assert params['params'][0] == { + 'statement_id': 'stmt-1', + 'value': 'GraphRAG combines knowledge graphs with RAG', + 'details': 'first detail\nsecond detail', } - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='statementId') - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert mock_graph_client.execute_query_with_retry.called - + + def test_build_statement_with_chunk_writes_the_relationship(self): + """Verify a statement carrying a chunk id also gets linked to that chunk.""" + client = graph_client() + + StatementGraphBuilder().build( + statement_node({ + 'statementId': 'stmt-2', + 'value': 'Knowledge graphs store structured information', + 'chunkId': 'chunk-1', + }), + client + ) + + assert client.execute_query_with_retry.call_count == 2 + (query, params) = client.execute_query_with_retry.call_args_list[1].args + assert '__MENTIONED_IN__' in query + assert params['params'][0] == {'statement_id': 'stmt-2', 'chunk_id': 'chunk-1'} + + def test_build_without_chunk_writes_only_the_statement(self): + """Verify no chunk relationship is invented when there is no chunk id.""" + client = graph_client() + + StatementGraphBuilder().build( + statement_node({'statementId': 'stmt-3', 'value': 'A statement'}), + client + ) + + assert client.execute_query_with_retry.call_count == 1 + def test_build_multiple_statements(self): """Verify building multiple statement nodes.""" + client = graph_client() builder = StatementGraphBuilder() - - statements = [ - Mock(metadata={'statement': {'statementId': 's1', 'text': 'Statement 1'}}), - Mock(metadata={'statement': {'statementId': 's2', 'text': 'Statement 2'}}) - ] - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='statementId') - mock_graph_client.execute_query_with_retry = Mock() - - for stmt in statements: - stmt.node_id = stmt.metadata['statement']['statementId'] - stmt.text = stmt.metadata['statement']['text'] - stmt.relationships = {} - builder.build(stmt, mock_graph_client) - - assert mock_graph_client.execute_query_with_retry.call_count >= 2 + for i in range(2): + builder.build( + statement_node({'statementId': f's{i}', 'value': f'Statement {i}'}, node_id=f's{i}'), + client + ) + + assert client.execute_query_with_retry.call_count == 2 + + def test_build_reads_the_previous_statement_relationship(self): + """Verify a PREVIOUS relationship carrying a statement is parsed, not ignored. + + The builder validates the previous node's statement metadata, so a + malformed one would raise here rather than at the write - worth covering, + because the relationship is optional and easy to leave untested. + """ + client = graph_client() + + node = statement_node({'statementId': 'stmt-2', 'value': 'Second'}) + node.relationships[NodeRelationship.PREVIOUS] = RelatedNodeInfo( + node_id='stmt-1', + metadata={'statement': {'statementId': 'stmt-1', 'value': 'First'}}, + ) + + StatementGraphBuilder().build(node, client) + + assert client.execute_query_with_retry.called + + def test_build_tolerates_unknown_kwargs(self): + """Verify the shared build kwargs the pipeline passes are ignored safely.""" + client = graph_client() + + StatementGraphBuilder().build( + statement_node({'statementId': 'stmt-1', 'value': 'A statement'}), + client, + include_domain_labels=False, + include_local_entities=True, + typed_properties='subject', + ) + + assert client.execute_query_with_retry.called class TestStatementGraphBuilderErrorHandling: """Tests for statement graph builder error handling.""" - - def test_build_with_missing_statement_id(self): - """Verify handling of statement with missing ID.""" - builder = StatementGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'stmt_001' - mock_node.metadata = {'statement': {'text': 'Statement without ID'}} - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert not mock_graph_client.execute_query_with_retry.called - - def test_build_with_empty_statement_text(self): - """Verify handling of statement with empty text.""" - builder = StatementGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'stmt_001' - mock_node.text = '' - mock_node.metadata = { - 'statement': { - 'statementId': 'stmt_001', - 'text': '', - 'metadata': {} - } - } - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='statementId') - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert mock_graph_client.execute_query_with_retry.called + + def test_build_with_no_statement_metadata_writes_nothing(self): + """Verify a node carrying no statement is left alone.""" + client = graph_client() + + node = TextNode(text='', id_='stmt-1') + node.metadata = {} + + StatementGraphBuilder().build(node, client) + + assert not client.execute_query_with_retry.called + + def test_build_with_missing_statement_id_still_writes(self): + """Verify the id is not treated as required. + + `Statement.statementId` is optional, so the write goes ahead with a null + id. Asserted because it is surprising, not because it is desirable. + """ + client = graph_client() + + StatementGraphBuilder().build(statement_node({'value': 'Statement without ID'}), client) + + (_, params) = client.execute_query_with_retry.call_args.args + assert params['params'][0]['statement_id'] is None + + def test_build_with_empty_statement_value(self): + """Verify an empty value is written rather than skipped.""" + client = graph_client() + + StatementGraphBuilder().build(statement_node({'statementId': 'stmt-1', 'value': ''}), client) + + (_, params) = client.execute_query_with_retry.call_args.args + assert params['params'][0]['value'] == '' diff --git a/lexical-graph/tests/unit/indexing/build/test_topic_graph_builder.py b/lexical-graph/tests/unit/indexing/build/test_topic_graph_builder.py index feaba402..d33a245f 100644 --- a/lexical-graph/tests/unit/indexing/build/test_topic_graph_builder.py +++ b/lexical-graph/tests/unit/indexing/build/test_topic_graph_builder.py @@ -1,134 +1,142 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import pytest +"""Unit tests for `TopicGraphBuilder`. + +The `Topic` metadata here is spelled as the model actually defines it - `value` +and `chunkIds`, not `name` and `metadata`. The earlier version of this file used +the latter and every test in it raised `ValidationError`; nothing noticed, because +pytest's default `norecursedirs` contains `build`, so this whole directory was +skipped unless its path was named explicitly. +""" + from unittest.mock import Mock + from graphrag_toolkit.lexical_graph.indexing.build.topic_graph_builder import TopicGraphBuilder +from llama_index.core.schema import TextNode + +def topic_node(topic:dict, node_id:str='topic-1') -> TextNode: + node = TextNode(text=topic.get('value', ''), id_=node_id) + node.metadata = {'topic': topic} + return node + +def graph_client() -> Mock: + client = Mock() + client.node_id = Mock(side_effect=lambda field: f'params.{field}') + client.execute_query_with_retry = Mock(return_value=[]) + return client class TestTopicGraphBuilderInitialization: """Tests for TopicGraphBuilder initialization.""" - + def test_initialization(self): """Verify TopicGraphBuilder initializes correctly.""" - builder = TopicGraphBuilder() - assert builder is not None + assert TopicGraphBuilder() is not None + def test_index_key(self): + """Verify the builder claims the topic index.""" + assert TopicGraphBuilder.index_key() == 'topic' class TestTopicGraphBuilding: """Tests for topic graph building functionality.""" - - def test_build_topic_node(self, mock_neptune_store): - """Verify building topic node with metadata.""" - builder = TopicGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'topic_001' - mock_node.text = 'Artificial Intelligence' - mock_node.metadata = { - 'topic': { - 'topicId': 'topic_001', - 'name': 'Artificial Intelligence', - 'metadata': { - 'category': 'Technology', - 'relevance': 0.95 - } - } - } - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='topicId') - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert mock_graph_client.execute_query_with_retry.called - - def test_build_topic_with_subtopics(self): - """Verify building topic with subtopic relationships.""" - builder = TopicGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'topic_002' - mock_node.text = 'Machine Learning' - mock_node.metadata = { - 'topic': { - 'topicId': 'topic_002', - 'name': 'Machine Learning', - 'parent_topic': 'topic_001', - 'subtopics': ['topic_003', 'topic_004'], - 'metadata': {} - } + + def test_build_topic_node(self): + """Verify building a topic node writes it with its value.""" + client = graph_client() + + TopicGraphBuilder().build( + topic_node({'topicId': 'topic-1', 'value': 'Artificial Intelligence'}), + client + ) + + client.execute_query_with_retry.assert_called_once() + (_, params) = client.execute_query_with_retry.call_args.args + assert params['params'][0] == { + 'topic_id': 'topic-1', + 'title': 'Artificial Intelligence', + 'chunk_ids': [], } - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='topicId') - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert mock_graph_client.execute_query_with_retry.called - + + def test_build_topic_with_chunks(self): + """Verify each chunk the topic was mentioned in is bound as a parameter.""" + client = graph_client() + + TopicGraphBuilder().build( + topic_node({ + 'topicId': 'topic-2', + 'value': 'Machine Learning', + 'chunkIds': ['chunk-1', 'chunk-2'], + }), + client + ) + + (query, params) = client.execute_query_with_retry.call_args.args + assert params['params'][0]['chunk_ids'] == [{'chunk_id': 'chunk-1'}, {'chunk_id': 'chunk-2'}] + assert '__MENTIONED_IN__' in query + def test_build_multiple_topics(self): """Verify building multiple topic nodes.""" + client = graph_client() builder = TopicGraphBuilder() - - topics = [ - Mock(metadata={'topic': {'topicId': 't1', 'name': 'Topic 1'}}), - Mock(metadata={'topic': {'topicId': 't2', 'name': 'Topic 2'}}), - Mock(metadata={'topic': {'topicId': 't3', 'name': 'Topic 3'}}) - ] - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='topicId') - mock_graph_client.execute_query_with_retry = Mock() - - for topic in topics: - topic.node_id = topic.metadata['topic']['topicId'] - topic.text = topic.metadata['topic']['name'] - topic.relationships = {} - builder.build(topic, mock_graph_client) - - assert mock_graph_client.execute_query_with_retry.call_count >= 3 + for i in range(3): + builder.build(topic_node({'topicId': f't{i}', 'value': f'Topic {i}'}, node_id=f't{i}'), client) + + assert client.execute_query_with_retry.call_count == 3 + + def test_build_tolerates_unknown_kwargs(self): + """Verify the shared build kwargs the pipeline passes are ignored safely. + + Every builder receives the same kwargs, so a builder that does not read + one must still accept it. + """ + client = graph_client() + + TopicGraphBuilder().build( + topic_node({'topicId': 'topic-1', 'value': 'AI'}), + client, + include_domain_labels=False, + include_local_entities=True, + typed_properties='subject', + ) + + assert client.execute_query_with_retry.called class TestTopicGraphBuilderErrorHandling: """Tests for topic graph builder error handling.""" - - def test_build_with_missing_topic_id(self): - """Verify handling of topic with missing ID.""" - builder = TopicGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'topic_001' - mock_node.metadata = {'topic': {'name': 'Topic without ID'}} - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert not mock_graph_client.execute_query_with_retry.called - - def test_build_with_empty_topic_name(self): - """Verify handling of topic with empty name.""" - builder = TopicGraphBuilder() - - mock_node = Mock() - mock_node.node_id = 'topic_001' - mock_node.text = '' - mock_node.metadata = { - 'topic': { - 'topicId': 'topic_001', - 'name': '', - 'metadata': {} - } - } - mock_node.relationships = {} - - mock_graph_client = Mock() - mock_graph_client.node_id = Mock(return_value='topicId') - mock_graph_client.execute_query_with_retry = Mock() - - builder.build(mock_node, mock_graph_client) - assert mock_graph_client.execute_query_with_retry.called + + def test_build_with_no_topic_metadata_writes_nothing(self): + """Verify a node carrying no topic is left alone.""" + client = graph_client() + + node = TextNode(text='', id_='topic-1') + node.metadata = {} + + TopicGraphBuilder().build(node, client) + + assert not client.execute_query_with_retry.called + + def test_build_with_missing_topic_id_still_writes(self): + """Verify the id is not treated as required. + + `Topic.topicId` is optional, so the write goes ahead with a null id rather + than being skipped. Asserted because it is surprising, not because it is + desirable - a caller relying on the builder to reject an unidentified + topic would be relying on something it does not do. + """ + client = graph_client() + + TopicGraphBuilder().build(topic_node({'value': 'Topic without ID'}), client) + + (_, params) = client.execute_query_with_retry.call_args.args + assert params['params'][0]['topic_id'] is None + + def test_build_with_empty_topic_value(self): + """Verify an empty value is written rather than skipped.""" + client = graph_client() + + TopicGraphBuilder().build(topic_node({'topicId': 'topic-1', 'value': ''}), client) + + (_, params) = client.execute_query_with_retry.call_args.args + assert params['params'][0]['title'] == '' diff --git a/lexical-graph/tests/unit/indexing/build/test_typed_properties_config.py b/lexical-graph/tests/unit/indexing/build/test_typed_properties_config.py new file mode 100644 index 00000000..bacfeca4 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/build/test_typed_properties_config.py @@ -0,0 +1,644 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""How the `typed_properties` setting is plumbed, and what refuses it. + +The setting exists in one place a user can set it, reaches `builder.build()` from +there through `LexicalGraphIndex` and `BuildPipeline`, defaults to `'off'` at every +layer, and the two configurations that cannot work are refused up front rather than +half-performed. + +Much of what these tests assert is an absence, which is the point: the feature is +opt-in, and a user who does not opt in must be unaffected. + +The four things under test, and where each lives: + +* `OntologyConfig.typed_properties` - the only place a user can set it, plus the + reserved-name refusal; +* `GraphRAGConfig.typed_properties` - a `coalesce` floor only, deliberately not + environment-readable; +* `BuildPipeline` - validation of the placement, and the + complement-without-local-entities refusal; +* `EntityGraphBuilder` - reads the kwarg tolerantly, and warns when it was asked + for and cannot be honoured. + +What each placement actually *writes* is in `test_typed_property_values.py` and +`test_local_entity_rewrites_typed_carry.py`. +""" + +import logging +from typing import get_args +from unittest.mock import MagicMock, patch + +import pytest + +from graphrag_toolkit.lexical_graph.config import GraphRAGConfig +from graphrag_toolkit.lexical_graph.indexing.build import entity_graph_builder as entity_graph_builder_module +from graphrag_toolkit.lexical_graph.indexing.build.build_pipeline import BuildPipeline +from graphrag_toolkit.lexical_graph.indexing.build.entity_graph_builder import EntityGraphBuilder +from graphrag_toolkit.lexical_graph.indexing.build.graph_construction import GraphConstruction, default_builders +from graphrag_toolkit.lexical_graph.indexing.constants import ( + COMPLEMENT_ENTITY_PROPERTIES, + COMPLEMENT_PLACEMENTS, + RESERVED_ENTITY_PROPERTIES, + SUBJECT_PLACEMENTS, + TYPED_PROPERTIES_OFF, + TYPED_PROPERTY_PLACEMENTS, +) + +from llama_index.core.schema import TextNode + + +# A minimal ontology declaring a datatype property named `value`, which is a name +# `__Entity__` already owns. Written inline rather than added to +# `tests/fixtures/ontologies/` because it is not a plausible ontology - its only +# purpose is to be rejected. +RESERVED_NAME_TURTLE = """ +@prefix owl: . +@prefix rdfs: . +@prefix xsd: . +@prefix : . + +: a owl:Ontology . + +:Thing a owl:Class ; rdfs:label "Thing" . + +:value a owl:DatatypeProperty ; + rdfs:label "value" ; + rdfs:domain :Thing ; + rdfs:range xsd:string . +""" + +# One well-formed fact, annotated as `OntologyFilter` would leave it. +ANNOTATED_FACT = { + 'factId': 'f-1', + 'subject': { + 'entityId': 's-1', 'value': 'Amazon', 'classification': 'Company', + 'classIri': 'http://example.org/company#Company', + }, + 'predicate': { + 'value': 'FOUNDED YEAR', + 'propertyIri': 'http://example.org/company#foundedYear', + 'canonicalName': 'foundedYear', + }, + 'complement': { + 'entityId': 'c-1', 'value': '1994', 'classification': '__Local_Entity__', + 'datatype': 'http://www.w3.org/2001/XMLSchema#integer', + }, +} + +UNANNOTATED_FACT = { + 'factId': 'f-2', + 'subject': {'entityId': 's-2', 'value': 'Amazon', 'classification': 'Company'}, + 'predicate': {'value': 'acquired'}, + 'object': {'entityId': 'o-2', 'value': 'Whole Foods', 'classification': 'Company'}, +} + +def fact_node(fact:dict, node_id:str='n-1') -> TextNode: + node = TextNode(text='', id_=node_id) + node.metadata = {'fact': fact} + return node + +@pytest.fixture +def reserved_name_ontology(): + ''' + An ontology declaring `:value a owl:DatatypeProperty`. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import Ontology + return Ontology.from_turtle_string(RESERVED_NAME_TURTLE, base_iri='http://example.org/reserved#') + +@pytest.fixture +def reset_annotation_warning(): + ''' + Clear `entity_graph_builder`'s module-level keeping graph writes out of ambient configuration3 state around a + test. It is process-wide by design, so a test that trips it would otherwise + leak into every test that runs after it. + ''' + entity_graph_builder_module._reset_no_annotations_warning() + yield + entity_graph_builder_module._reset_no_annotations_warning() + +@pytest.fixture +def restore_graphrag_typed_properties(): + ''' + Restore `GraphRAGConfig`'s typed_properties around a test that sets it. The + config is a process-wide singleton. + ''' + previous = GraphRAGConfig._typed_properties + yield + GraphRAGConfig._typed_properties = previous + +class TestPlacementVocabulary: + ''' + Tests that the placement names agree everywhere they are spelled out. + ''' + + def test_literal_and_tuple_agree(self): + ''' + `TypedProperties` is what a type checker sees and `TYPED_PROPERTY_PLACEMENTS` + is what validation checks, and they live in different modules - the + `Literal` in `ontology_config`, the tuple in `indexing.constants`, because + `build_pipeline` must read the tuple without importing rdflib. Nothing but + this test keeps them in step. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import TypedProperties + + assert get_args(TypedProperties) == TYPED_PROPERTY_PLACEMENTS + + def test_placement_subsets_cover_every_placement_but_off(self): + ''' + Every placement other than `'off'` writes somewhere, and `'off'` writes + nowhere. A placement in neither subset would be silently inert. + ''' + writing = set(SUBJECT_PLACEMENTS) | set(COMPLEMENT_PLACEMENTS) + + assert writing == set(TYPED_PROPERTY_PLACEMENTS) - {TYPED_PROPERTIES_OFF} + assert TYPED_PROPERTIES_OFF not in writing + + def test_both_is_in_both_subsets(self): + ''' + `'both'` is not a fallback chain: it writes to the subject *and* to the + complement. + ''' + assert 'both' in SUBJECT_PLACEMENTS + assert 'both' in COMPLEMENT_PLACEMENTS + +class TestDefaultsToOff: + ''' + Tests that the default is `'off'` at every layer. + ''' + + def test_graphrag_config_default(self): + ''' + The `coalesce` floor. + ''' + assert GraphRAGConfig.typed_properties == TYPED_PROPERTIES_OFF + + def test_no_environment_variable_can_turn_it_on(self, monkeypatch, restore_graphrag_typed_properties): + ''' + No ambient configuration turns on graph writes. The + property is deliberately not environment-readable, unlike its neighbours, + so that a user with no ontology cannot reach any placement but `'off'` +. Several plausible spellings are tried, because the + failure this guards against is someone adding one of them later. + ''' + for name in ('TYPED_PROPERTIES', 'typed_properties', 'GRAPHRAG_TYPED_PROPERTIES'): + monkeypatch.setenv(name, 'both') + + GraphRAGConfig._typed_properties = None + + assert GraphRAGConfig.typed_properties == TYPED_PROPERTIES_OFF + + def test_ontology_config_default(self, company_ontology): + ''' + An ontology that says nothing about typed properties writes none. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + + config = OntologyConfig(company_ontology, ontology_authority='strict') + + assert config.typed_properties == TYPED_PROPERTIES_OFF + assert config.writes_subject_properties() is False + assert config.writes_complement_properties() is False + + def test_index_with_no_ontology_passes_none(self): + ''' + `LexicalGraphIndex._typed_properties()` returns None with no ontology, + leaving `BuildPipeline` to coalesce to the config floor. None rather than + `'off'` so that a caller's explicit argument still wins. + ''' + from graphrag_toolkit.lexical_graph.lexical_graph_index import LexicalGraphIndex + + index = LexicalGraphIndex.__new__(LexicalGraphIndex) + index.indexing_config = MagicMock() + index.indexing_config.extraction.ontology = None + + assert index._typed_properties() is None + + def test_index_with_ontology_passes_its_placement(self, company_ontology): + ''' + With an ontology, the ontology's own setting is what travels. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + from graphrag_toolkit.lexical_graph.lexical_graph_index import LexicalGraphIndex + + index = LexicalGraphIndex.__new__(LexicalGraphIndex) + index.indexing_config = MagicMock() + index.indexing_config.extraction.ontology = OntologyConfig( + company_ontology, ontology_authority='align', typed_properties='subject' + ) + + assert index._typed_properties() == 'subject' + +class TestOntologyConfigValidation: + ''' + Tests for `OntologyConfig`'s own refusals. + ''' + + @pytest.mark.parametrize('placement', TYPED_PROPERTY_PLACEMENTS) + def test_every_documented_placement_is_accepted(self, company_ontology, placement): + ''' + The four placements the type says exist all construct. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + + assert OntologyConfig(company_ontology, typed_properties=placement).typed_properties == placement + + def test_unknown_placement_raises_listing_the_supported_ones(self, company_ontology): + ''' + A typo should say what was expected, not just that something was wrong. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + + with pytest.raises(ValueError) as excinfo: + OntologyConfig(company_ontology, typed_properties='sub-ject') + + message = str(excinfo.value) + assert "'sub-ject'" in message + for placement in TYPED_PROPERTY_PLACEMENTS: + assert placement in message + + @pytest.mark.parametrize('placement', SUBJECT_PLACEMENTS) + def test_reserved_property_name_raises_under_subject_placement(self, reserved_name_ontology, placement): + ''' + A declared `:value` would be keyed onto `__Entity__` with + the same name as the entity's own identity string. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + + with pytest.raises(ValueError) as excinfo: + OntologyConfig(reserved_name_ontology, typed_properties=placement) + + message = str(excinfo.value) + assert "'value'" in message + assert placement in message + # The message must be actionable: name the offending term, and say what + # to do about it. + assert 'Rename' in message + + def test_reserved_name_message_widens_the_set_under_both(self, reserved_name_ontology): + ''' + `'both'` writes `typed_value` and `datatype` too, so those names are + reserved as well and the message should say so; `'subject'` alone does + not write them and should not claim otherwise. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + + with pytest.raises(ValueError) as both_error: + OntologyConfig(reserved_name_ontology, typed_properties='both') + with pytest.raises(ValueError) as subject_error: + OntologyConfig(reserved_name_ontology, typed_properties='subject') + + # Read the parenthesized reserved set rather than searching the whole + # message: the prose says "declares datatype property", so a bare + # substring test for `datatype` matches text that is not the set. + def reserved_set(error): + message = str(error.value) + listed = message[message.index('__Entity__ (') + len('__Entity__ ('):] + return {name.strip() for name in listed[:listed.index(')')].split(',')} + + assert reserved_set(subject_error) == set(RESERVED_ENTITY_PROPERTIES) + assert reserved_set(both_error) == set(RESERVED_ENTITY_PROPERTIES) | set(COMPLEMENT_ENTITY_PROPERTIES) + + @pytest.mark.parametrize('placement', [TYPED_PROPERTIES_OFF, 'complement']) + def test_reserved_property_name_accepted_when_nothing_keys_from_it(self, reserved_name_ontology, placement): + ''' + The check is scoped to the placement that actually writes. `'off'` writes + nothing, and `'complement'` writes fixed names rather than keying from the + ontology's vocabulary, so neither can collide - and rejecting them would + refuse a configuration that works. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + + assert OntologyConfig(reserved_name_ontology, typed_properties=placement) is not None + +class TestFilterRequired: + ''' + Tests that asking for typed properties keeps the filter in the pipeline. + ''' + + @pytest.mark.parametrize('placement', sorted(set(TYPED_PROPERTY_PLACEMENTS) - {TYPED_PROPERTIES_OFF})) + def test_typed_properties_alone_requires_the_filter(self, company_ontology, placement): + ''' + The annotations the builders read are written by `OntologyFilter` and by + nothing else, so `ontology_authority='off'` plus a placement must still build the + filter. Without this the request would be accepted and silently write + nothing. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + + config = OntologyConfig(company_ontology, ontology_authority='off', typed_properties=placement) + + assert config.resolved().any_enabled() is False + assert config.filter_required() is True + + def test_off_plus_off_does_not_require_the_filter(self, company_ontology): + ''' + The pre-existing behaviour is unchanged: nothing to do means no transform. + ''' + from graphrag_toolkit.lexical_graph.indexing.extract.ontology import OntologyConfig + + config = OntologyConfig(company_ontology, ontology_authority='off', typed_properties='off') + + assert config.filter_required() is False + +class TestBuildPipelineValidation: + ''' + Tests for `BuildPipeline`'s validation, including the local-entities rule. + ''' + + def _create(self, **kwargs): + # `BuildPipeline.create` wraps the pipeline in a `Pipe` and returns that, + # so the object under test is unreachable through it. Constructed directly + # here; `create` is a pass-through and is covered by the `run_pipeline` + # test below. + return BuildPipeline(components=[], **kwargs) + + @pytest.mark.parametrize('placement', [TYPED_PROPERTIES_OFF, 'subject']) + def test_placements_needing_no_complement_are_accepted(self, placement): + ''' + Neither of these writes to the complement node, so neither depends on one + existing. + ''' + pipeline = self._create(typed_properties=placement, include_local_entities=False) + + assert pipeline.typed_properties == placement + + @pytest.mark.parametrize('placement', COMPLEMENT_PLACEMENTS) + def test_complement_placement_without_local_entities_raises(self, placement): + ''' + Complement placement writes to a node that + `include_local_entities=False` never creates, so the write has no target. + Refused at construction: the alternative is a build that completes and + stores nothing, which is the failure mode hardest to diagnose. + ''' + with pytest.raises(ValueError) as excinfo: + self._create(typed_properties=placement, include_local_entities=False) + + message = str(excinfo.value) + assert placement in message + assert 'include_local_entities' in message + # Actionable: both ways out are named. + assert 'include_local_entities=True' in message + assert "typed_properties='subject'" in message + + @pytest.mark.parametrize('placement', COMPLEMENT_PLACEMENTS) + def test_complement_placement_with_local_entities_is_accepted(self, placement): + ''' + With a complement node to write to, the same placement is fine. + ''' + pipeline = self._create(typed_properties=placement, include_local_entities=True) + + assert pipeline.typed_properties == placement + + def test_unknown_placement_raises(self): + ''' + Validated here as well as in `OntologyConfig`, because a caller can + construct a pipeline directly and never touch an `OntologyConfig`. + ''' + with pytest.raises(ValueError) as excinfo: + self._create(typed_properties='on') + + assert "'on'" in str(excinfo.value) + + def test_default_resolves_to_off(self): + ''' + No argument, no ontology, no writes. + ''' + assert self._create().typed_properties == TYPED_PROPERTIES_OFF + + def test_explicit_argument_beats_the_config_floor(self, restore_graphrag_typed_properties): + ''' + `coalesce` order: the caller's argument wins over `GraphRAGConfig`. + ''' + GraphRAGConfig.typed_properties = 'subject' + + assert self._create().typed_properties == 'subject' + assert self._create(typed_properties=TYPED_PROPERTIES_OFF).typed_properties == TYPED_PROPERTIES_OFF + +class TestKwargReachesTheBuilders: + ''' + Tests that the setting travels from the pipeline to `builder.build()`. + ''' + + def test_build_pipeline_passes_it_to_run_pipeline(self): + ''' + `run_pipeline`'s `**kwargs` are what eventually reach `accept`, so this is + the hand-off that matters at the pipeline end. + ''' + pipeline = BuildPipeline( + components=[], + typed_properties='subject', + include_local_entities=True, + ) + + from graphrag_toolkit.lexical_graph.indexing.model import SourceDocument + from llama_index.core.schema import NodeRelationship, RelatedNodeInfo + + chunk = TextNode(text='a', id_='a1') + chunk.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id='src-1') + doc = SourceDocument(nodes=[chunk]) + + with patch( + 'graphrag_toolkit.lexical_graph.indexing.build.build_pipeline.run_pipeline', + return_value=[] + ) as run: + list(pipeline.build([doc])) + + assert run.call_args.kwargs['typed_properties'] == 'subject' + + def test_graph_construction_passes_it_to_build(self, mock_neptune_store): + ''' + And this is the hand-off at the builder end: `accept`'s surviving kwargs + go to every builder for the node's index. + ''' + from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder + from graphrag_toolkit.lexical_graph.storage.constants import INDEX_KEY + + builder = MagicMock(spec=GraphBuilder) + builder.index_key.return_value = 'fact' + + construction = GraphConstruction(graph_client=mock_neptune_store, builders=[builder]) + + node = fact_node(ANNOTATED_FACT) + node.metadata[INDEX_KEY] = {'index': 'fact'} + + list(construction.accept( + [node], + batch_writes_enabled=False, + batch_write_size=10, + include_domain_labels=False, + include_local_entities=True, + typed_properties='subject', + )) + + builder.build.assert_called_once() + assert builder.build.call_args.kwargs['typed_properties'] == 'subject' + + @pytest.mark.parametrize( + 'builder', + default_builders(), + ids=lambda builder: type(builder).__name__ + ) + def test_every_default_builder_tolerates_the_kwarg(self, builder): + ''' + keeping graph writes out of ambient configuration5's other half. Every builder receives the same kwargs, so + adding one must not break the eight that ignore it. Asserted against all + of them rather than just `EntityGraphBuilder`, because the others read + their own kwargs with hard subscripts and a future kwarg added the same + way would fail here first. + ''' + node = fact_node(ANNOTATED_FACT, node_id=ANNOTATED_FACT['factId']) + + graph_client = MagicMock() + graph_client.node_id = lambda field: f'params.{field}' + graph_client.property_assigment_fn = lambda key, value: (lambda x: x) + graph_client.execute_query_with_retry = MagicMock(return_value=[]) + + # Not every builder handles a fact node - most look for their own + # metadata key and do nothing. Doing nothing is a pass; raising is not. + builder.build( + node, + graph_client, + include_domain_labels=False, + include_local_entities=True, + typed_properties='subject', + ) + + @pytest.mark.parametrize( + 'builder', + default_builders(), + ids=lambda builder: type(builder).__name__ + ) + def test_every_default_builder_tolerates_its_absence(self, builder): + ''' + A caller who built a pipeline before this setting + existed, or who calls a builder directly as several tests do, supplies no + `typed_properties` at all - and must not get a `KeyError` from a feature + they never asked for. + ''' + node = fact_node(ANNOTATED_FACT, node_id=ANNOTATED_FACT['factId']) + + graph_client = MagicMock() + graph_client.node_id = lambda field: f'params.{field}' + graph_client.property_assigment_fn = lambda key, value: (lambda x: x) + graph_client.execute_query_with_retry = MagicMock(return_value=[]) + + builder.build( + node, + graph_client, + include_domain_labels=False, + include_local_entities=True, + ) + +class TestNoAnnotationsWarning: + ''' + Typed properties asked for, and none can be written. + ''' + + def _build_unannotated_facts(self, count, typed_properties): + graph_client = MagicMock() + graph_client.node_id = lambda field: f'params.{field}' + graph_client.execute_query_with_retry = MagicMock(return_value=[]) + + builder = EntityGraphBuilder() + for i in range(count): + fact = dict(UNANNOTATED_FACT, factId=f'f-{i}') + builder.build( + fact_node(fact, node_id=fact['factId']), + graph_client, + include_domain_labels=False, + include_local_entities=False, + typed_properties=typed_properties, + ) + + def test_warns_once_after_enough_unannotated_facts(self, caplog, reset_annotation_warning): + ''' + The warning names the cause and the remedy, because the symptom - an empty + property - points at the graph store rather than at the pipeline order. + ''' + threshold = entity_graph_builder_module._NO_ANNOTATIONS_WARNING_AFTER + + with caplog.at_level(logging.WARNING): + self._build_unannotated_facts(threshold * 2, 'subject') + + warnings = [r for r in caplog.records if 'typed_properties' in r.getMessage()] + + assert len(warnings) == 1 + message = warnings[0].getMessage() + assert 'Re-extract' in message + assert 'ontology annotations' in message + + def test_silent_below_the_threshold(self, caplog, reset_annotation_warning): + ''' + One unresolved predicate carries no annotations and is entirely normal at + `align`, so a single unannotated fact must not cry wolf. + ''' + threshold = entity_graph_builder_module._NO_ANNOTATIONS_WARNING_AFTER + + with caplog.at_level(logging.WARNING): + self._build_unannotated_facts(threshold - 1, 'subject') + + assert not [r for r in caplog.records if 'typed_properties' in r.getMessage()] + + def test_silent_at_off(self, caplog, reset_annotation_warning): + ''' + Nobody asked for typed properties, so their absence is not news. + ''' + threshold = entity_graph_builder_module._NO_ANNOTATIONS_WARNING_AFTER + + with caplog.at_level(logging.WARNING): + self._build_unannotated_facts(threshold * 2, TYPED_PROPERTIES_OFF) + + assert not [r for r in caplog.records if 'typed_properties' in r.getMessage()] + + def test_one_annotated_fact_silences_it_permanently(self, caplog, reset_annotation_warning): + ''' + Evidence that the filter ran is evidence enough. Facts the filter could + not resolve are expected at every level below `strict`, and counting them + after that point would warn about normal operation. + ''' + graph_client = MagicMock() + graph_client.node_id = lambda field: f'params.{field}' + graph_client.execute_query_with_retry = MagicMock(return_value=[]) + + builder = EntityGraphBuilder() + + with caplog.at_level(logging.WARNING): + builder.build( + fact_node(ANNOTATED_FACT), + graph_client, + include_domain_labels=False, + include_local_entities=True, + typed_properties='subject', + ) + self._build_unannotated_facts( + entity_graph_builder_module._NO_ANNOTATIONS_WARNING_AFTER * 2, 'subject' + ) + + assert not [r for r in caplog.records if 'typed_properties' in r.getMessage()] + + def test_annotation_detection_reads_the_whole_fact(self): + ''' + The question is "did the filter run", not "can this fact be written", so + any one annotation counts - a fact whose subject resolved but whose + predicate did not is still proof the filter ran. + ''' + from graphrag_toolkit.lexical_graph.indexing.model import Fact + + has = entity_graph_builder_module._has_ontology_annotations + + assert has(Fact.model_validate(ANNOTATED_FACT)) is True + assert has(Fact.model_validate(UNANNOTATED_FACT)) is False + + subject_only = dict( + UNANNOTATED_FACT, + subject=dict(UNANNOTATED_FACT['subject'], classIri='http://example.org/company#Company'), + ) + assert has(Fact.model_validate(subject_only)) is True + + predicate_only = dict( + UNANNOTATED_FACT, + predicate={'value': 'acquired', 'canonicalName': 'acquired'}, + ) + assert has(Fact.model_validate(predicate_only)) is True diff --git a/lexical-graph/tests/unit/indexing/build/test_typed_property_values.py b/lexical-graph/tests/unit/indexing/build/test_typed_property_values.py new file mode 100644 index 00000000..da48c23b --- /dev/null +++ b/lexical-graph/tests/unit/indexing/build/test_typed_property_values.py @@ -0,0 +1,343 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Where a coerced attribute value gets written, if anywhere. + +Most of the decision is in two pure functions that answer "what, +if anything, does *this* fact contribute", so most of this file tests those +directly. The builder is exercised through a recording stand-in for the graph +store, for the one property that is about the queries rather than the values: the +entity insert has to be byte-identical at every placement, so turning typed +properties on cannot disturb `value`, `search_str` or `class`. +""" + +import logging + +import pytest + +from graphrag_toolkit.lexical_graph.indexing.build.entity_graph_builder import ( + EntityGraphBuilder, + _typed_complement_values, + _typed_subject_property, +) +from graphrag_toolkit.lexical_graph.indexing.model import Entity, Fact, Relation + +XSD = 'http://www.w3.org/2001/XMLSchema#' +LOCAL = '__Local_Entity__' + +def attribute_fact(canonical_name='foundedYear', value='1994', datatype=f'{XSD}integer'): + """An annotated attribute fact, as the ontology filter leaves one.""" + return Fact( + factId='f1', + subject=Entity(entityId='e1', value='Meridian Freight', classification='Company'), + predicate=Relation(value='foundedYear', canonicalName=canonical_name), + complement=Entity(entityId='c1', value=value, classification=LOCAL, datatype=datatype), + ) + +def relation_fact(): + """An annotated *relation*: `canonicalName` is set, but there is no literal.""" + return Fact( + factId='f2', + subject=Entity(entityId='e1', value='Priya Raman', classification='Person'), + predicate=Relation(value='worksFor', canonicalName='worksFor'), + object=Entity(entityId='e2', value='Meridian Freight', classification='Company'), + ) + +class RecordingGraphStore: + """The three methods `EntityGraphBuilder` asks a store for.""" + + def __init__(self): + self.calls = [] + + def node_id(self, name): + return name + + def property_assigment_fn(self, key, value): + return lambda placeholder: placeholder + + def execute_query_with_retry(self, query, params, **kwargs): + self.calls.append((query, params)) + +def build(fact, typed_properties, include_local_entities=True): + from llama_index.core.schema import TextNode + + store = RecordingGraphStore() + EntityGraphBuilder().build( + TextNode(text='x', metadata={'fact': fact.model_dump()}), + store, + include_domain_labels=False, + include_local_entities=include_local_entities, + typed_properties=typed_properties, + ) + return store.calls + +def queries(calls): + return [query for (query, _) in calls] + +def row(params): + """The one parameter row `UNWIND $params AS params` is given.""" + return params['params'][0] + +class TestWhichPlacementWrites: + + @pytest.mark.parametrize('placement,writes', [ + ('off', False), ('subject', True), ('complement', False), ('both', True), + ]) + def test_the_subject_property(self, placement, writes): + assert bool(_typed_subject_property(attribute_fact(), placement)) is writes + + @pytest.mark.parametrize('placement,writes', [ + ('off', False), ('subject', False), ('complement', True), ('both', True), + ]) + def test_the_complement_values(self, placement, writes): + complement = attribute_fact().complement + assert bool(_typed_complement_values(complement, placement)) is writes + + def test_both_is_not_a_fallback_chain(self): + """Each placement writes to its own node, and `'both'` does each.""" + fact = attribute_fact() + assert _typed_subject_property(fact, 'both') == ('foundedYear', 1994) + assert _typed_complement_values(fact.complement, 'both') == (1994, f'{XSD}integer') + +class TestTheSubjectProperty: + + def test_the_key_is_the_canonical_name_and_the_value_is_coerced(self): + assert _typed_subject_property(attribute_fact(), 'subject') == ('foundedYear', 1994) + + def test_a_relation_contributes_nothing_even_though_it_is_annotated(self): + """`complement.datatype` is the discriminator, not `canonicalName`. + + The filter sets `canonicalName` for every resolved predicate including + object properties. Keying off it alone would write the object entity's + display string into an attribute slot on the subject. + """ + assert _typed_subject_property(relation_fact(), 'subject') is None + + @pytest.mark.parametrize('kwargs', [ + {'canonical_name': None}, + {'datatype': None}, + ]) + def test_an_unannotated_fact_contributes_nothing(self, kwargs): + """An unresolved predicate has no declared datatype.""" + assert _typed_subject_property(attribute_fact(**kwargs), 'subject') is None + + def test_a_literal_that_does_not_coerce_is_refused_rather_than_stored_raw(self): + """Storing `'nineteen ninety four'` under a key whose + name promises a number is worse than storing nothing.""" + assert _typed_subject_property(attribute_fact(value='nineteen ninety four'), 'subject') is None + + @pytest.mark.parametrize('value,datatype,expected', [ + ('0', f'{XSD}integer', 0), + ('false', f'{XSD}boolean', False), + ('0.0', f'{XSD}double', 0.0), + ]) + def test_a_falsy_value_is_still_written(self, value, datatype, expected): + """Callers must test the pair, not the value.""" + pair = _typed_subject_property(attribute_fact(value=value, datatype=datatype), 'subject') + assert pair == ('foundedYear', expected) + + @pytest.mark.parametrize('name', ['value', 'search_str', 'class']) + def test_a_name_the_graph_model_owns_is_skipped_with_a_warning(self, name, caplog): + """the reserved-name rule, defence in depth: the config already refuses such an + ontology, but a fact can arrive from a checkpoint written under another.""" + with caplog.at_level(logging.WARNING): + assert _typed_subject_property( + attribute_fact(canonical_name=name, datatype=f'{XSD}string'), 'subject', + ) is None + + assert 'already owns' in caplog.text + + @pytest.mark.parametrize('name', ['typed_value', 'datatype']) + def test_the_complement_names_are_owned_only_once_complement_placement_writes(self, name): + fact = attribute_fact(canonical_name=name, datatype=f'{XSD}string', value='x') + + assert _typed_subject_property(fact, 'subject') == (name, 'x') + assert _typed_subject_property(fact, 'both') is None + + @pytest.mark.parametrize('name', ['founded\nYear', 'founded\rYear']) + def test_a_name_carrying_a_line_break_is_rejected(self, name, caplog): + """Every other character is escaped; a line break would + split the query across lines the builder treats as one statement each, and + no well-formed Turtle term contains one.""" + with caplog.at_level(logging.WARNING): + assert _typed_subject_property(attribute_fact(canonical_name=name), 'subject') is None + + assert 'line break' in caplog.text + +class TestTheComplementValues: + + def test_the_pair_is_the_coerced_value_and_the_declared_datatype(self): + complement = attribute_fact().complement + assert _typed_complement_values(complement, 'complement') == (1994, f'{XSD}integer') + + def test_an_unannotated_complement_contributes_nothing(self): + complement = attribute_fact(datatype=None).complement + assert _typed_complement_values(complement, 'complement') is None + + def test_a_datatype_is_never_written_without_a_value(self): + """Asserting a type for a value that is not there is + worse than the absence of both. The string is still on the node as + `value`, where every existing consumer reads it.""" + complement = attribute_fact(value='nineteen ninety four').complement + assert _typed_complement_values(complement, 'complement') is None + + def test_a_falsy_value_is_still_written(self): + complement = attribute_fact(value='false', datatype=f'{XSD}boolean').complement + assert _typed_complement_values(complement, 'complement') == (False, f'{XSD}boolean') + +class TestWhatReachesTheGraph: + + def test_off_issues_no_typed_write(self): + assert not [query for query in queries(build(attribute_fact(), 'off')) if 'typed' in query] + + def test_subject_placement_sets_the_property_under_its_own_name(self): + typed = [(q, p) for (q, p) in build(attribute_fact(), 'subject') if 'insert typed property' in q] + + assert len(typed) == 1 + (query, params) = typed[0] + assert 'SET' in query and '`foundedYear`' in query + assert row(params)['typedValue'] == 1994 + assert row(params)['entityId'] == 'e1' + + def test_the_value_is_bound_and_never_interpolated(self): + """The key is a Cypher identifier and is escaped; the + value is a parameter.""" + typed = [(q, p) for (q, p) in build(attribute_fact(value='1994'), 'subject') + if 'insert typed property' in q] + + (query, params) = typed[0] + assert '1994' not in query + assert row(params)['typedValue'] == 1994 + + def test_complement_placement_sets_typed_value_and_datatype(self): + typed = [(q, p) for (q, p) in build(attribute_fact(), 'complement') + if 'typed_value' in q or 'typedValue' in str(p)] + + assert typed + (_, params) = typed[0] + assert row(params)['typedValue'] == 1994 + assert row(params)['datatype'] == f'{XSD}integer' + + def test_the_entity_insert_is_byte_identical_at_every_placement(self): + """The byte-for-byte guarantee, and the reason each typed write is a separate + query rather than an extra `SET` on the insert. + + Turning typed properties on must not be able to disturb `value`, + `search_str` or `class`, and this is what makes that structural rather than + a promise. + """ + inserts = { + placement: [q for q in queries(build(attribute_fact(), placement)) if 'insert entities' in q] + for placement in ('off', 'subject', 'complement', 'both') + } + + assert inserts['off'] + for placement in ('subject', 'complement', 'both'): + assert inserts[placement] == inserts['off'] + + def test_a_fact_with_no_annotations_warns_once_when_a_placement_was_asked_for(self, caplog): + """A build that silently writes no typed properties is the failure worth a + log line: the setting is on, and nothing arrives.""" + import graphrag_toolkit.lexical_graph.indexing.build.entity_graph_builder as module + + module._reset_no_annotations_warning() + + with caplog.at_level(logging.WARNING): + for _ in range(60): + build(attribute_fact(canonical_name=None, datatype=None), 'subject') + + assert caplog.text.count('typed_properties=') == 1 + + def test_typed_properties_defaults_to_off_when_the_caller_omits_it(self): + """A pipeline built before this setting existed, or a + builder called directly, must not raise a KeyError from a feature it never + asked for.""" + from llama_index.core.schema import TextNode + + store = RecordingGraphStore() + EntityGraphBuilder().build( + TextNode(text='x', metadata={'fact': attribute_fact().model_dump()}), + store, + include_domain_labels=False, + include_local_entities=True, + ) + + assert not [query for query in queries(store.calls) if 'typed' in query] + +class TestTheSurroundingWriteConditions: + """The branches a placement runs inside, which decide whether it runs at all.""" + + def test_a_local_entity_subject_is_skipped_when_local_entities_are_off(self): + """Nothing is written for a fact whose subject is itself a value node, so + the typed write cannot conjure the node the setting suppressed.""" + fact = attribute_fact() + fact.subject.classification = LOCAL + + assert build(fact, 'subject', include_local_entities=False) == [] + assert build(fact, 'subject', include_local_entities=True) + + def test_the_complement_node_is_only_inserted_when_local_entities_are_on(self): + calls = build(attribute_fact(), 'complement', include_local_entities=False) + + assert not [q for (q, _) in calls if 'typed_value' in q] + + def test_domain_labels_are_added_alongside_a_typed_write(self): + """The two features are independent, and both write to the same node.""" + from llama_index.core.schema import TextNode + + store = RecordingGraphStore() + EntityGraphBuilder().build( + TextNode(text='x', metadata={'fact': attribute_fact().model_dump()}), + store, + include_domain_labels=True, + include_local_entities=True, + typed_properties='subject', + ) + + labelled = [q for (q, _) in store.calls if 'awsqid:' in q] + typed = [q for (q, _) in store.calls if 'insert typed property' in q] + + assert labelled and typed + + def test_a_local_entity_never_gets_a_domain_label(self): + """Its classification is a marker, not a type anyone would query by.""" + from llama_index.core.schema import TextNode + + store = RecordingGraphStore() + EntityGraphBuilder().build( + TextNode(text='x', metadata={'fact': attribute_fact().model_dump()}), + store, + include_domain_labels=True, + include_local_entities=True, + typed_properties='complement', + ) + + assert not [q for (q, _) in store.calls if f'`{LOCAL}`' in q] + + def test_a_node_carrying_no_fact_warns_and_writes_nothing(self, caplog): + from llama_index.core.schema import TextNode + + store = RecordingGraphStore() + with caplog.at_level(logging.WARNING): + EntityGraphBuilder().build( + TextNode(text='x', metadata={}), store, + include_domain_labels=False, include_local_entities=True, + typed_properties='subject', + ) + + assert store.calls == [] + assert 'fact_id missing' in caplog.text + + def test_an_annotated_fact_silences_the_no_annotations_warning_permanently(self, caplog): + """One annotated fact proves the filter ran, so the warning can never be + right afterwards however many unannotated facts follow.""" + from graphrag_toolkit.lexical_graph.indexing.build import entity_graph_builder as module + + module._reset_no_annotations_warning() + build(attribute_fact(), 'subject') + + with caplog.at_level(logging.WARNING): + for _ in range(60): + build(attribute_fact(canonical_name=None, datatype=None), 'subject') + + assert 'typed_properties=' not in caplog.text diff --git a/lexical-graph/tests/unit/indexing/build/test_vector_batch_client.py b/lexical-graph/tests/unit/indexing/build/test_vector_batch_client.py index 9ea52880..dfbfee90 100644 --- a/lexical-graph/tests/unit/indexing/build/test_vector_batch_client.py +++ b/lexical-graph/tests/unit/indexing/build/test_vector_batch_client.py @@ -1,101 +1,165 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +"""Unit tests for `VectorBatchClient` and `BatchVectorIndex`. + +Rewritten against the real API. The earlier version of this file called +`batch_add_embeddings`, `batch_update_embeddings` and `batch_delete_embeddings`, +none of which exist - and then assigned each one a `Mock` before calling it, so +every assertion was about the mock's own return value rather than about any code +in this repository. It also depended on a `mock_opensearch_store` fixture that +does not exist, so it errored at setup. Nothing noticed, because pytest's default +`norecursedirs` contains `build`, so this whole directory was skipped unless its +path was named explicitly. +""" + +from unittest.mock import MagicMock + import pytest -from unittest.mock import Mock -from graphrag_toolkit.lexical_graph.indexing.build.vector_batch_client import VectorBatchClient +from graphrag_toolkit.lexical_graph.indexing.build.vector_batch_client import ( + BatchVectorIndex, + VectorBatchClient, +) +from graphrag_toolkit.lexical_graph.storage.constants import ALL_EMBEDDING_INDEXES +from graphrag_toolkit.lexical_graph.storage.vector import DummyVectorIndex + +def vector_index(index_name:str) -> MagicMock: + index = MagicMock() + index.index_name = index_name + return index + +def vector_store(*index_names:str) -> MagicMock: + store = MagicMock() + store.all_indexes.return_value = [vector_index(name) for name in index_names] + return store + +def batch_client(*index_names, batch_writes_enabled=True, batch_write_size=2) -> VectorBatchClient: + return VectorBatchClient( + vector_store=vector_store(*index_names), + batch_writes_enabled=batch_writes_enabled, + batch_write_size=batch_write_size, + ) class TestVectorBatchClientInitialization: """Tests for VectorBatchClient initialization.""" - - def test_initialization(self, mock_opensearch_store): - """Verify VectorBatchClient initializes with vector store.""" - client = VectorBatchClient(vector_store=mock_opensearch_store) - assert client is not None - - -class TestVectorBatchOperations: - """Tests for vector batch operations.""" - - def test_batch_add_embeddings(self, mock_opensearch_store): - """Verify batch addition of embeddings.""" - client = VectorBatchClient(vector_store=mock_opensearch_store) - - embeddings = [ - {'id': 'chunk_001', 'vector': [0.1] * 384, 'text': 'Chunk 1'}, - {'id': 'chunk_002', 'vector': [0.2] * 384, 'text': 'Chunk 2'}, - {'id': 'chunk_003', 'vector': [0.3] * 384, 'text': 'Chunk 3'} - ] - - client.batch_add_embeddings = Mock(return_value={'added': 3}) - result = client.batch_add_embeddings(embeddings) - - assert result is not None - assert result['added'] == 3 - - def test_batch_update_embeddings(self, mock_opensearch_store): - """Verify batch update of embeddings.""" - client = VectorBatchClient(vector_store=mock_opensearch_store) - - updates = [ - {'id': 'chunk_001', 'vector': [0.15] * 384}, - {'id': 'chunk_002', 'vector': [0.25] * 384} - ] - - client.batch_update_embeddings = Mock(return_value={'updated': 2}) - result = client.batch_update_embeddings(updates) - - assert result is not None - assert result['updated'] == 2 - - def test_batch_delete_embeddings(self, mock_opensearch_store): - """Verify batch deletion of embeddings.""" - client = VectorBatchClient(vector_store=mock_opensearch_store) - - ids_to_delete = ['chunk_001', 'chunk_002', 'chunk_003'] - - client.batch_delete_embeddings = Mock(return_value={'deleted': 3}) - result = client.batch_delete_embeddings(ids_to_delete) - - assert result is not None - assert result['deleted'] == 3 - - -class TestVectorBatchClientErrorHandling: - """Tests for vector batch client error handling.""" - - def test_batch_add_with_empty_list(self, mock_opensearch_store): - """Verify handling of empty batch.""" - client = VectorBatchClient(vector_store=mock_opensearch_store) - - client.batch_add_embeddings = Mock(return_value={'added': 0}) - result = client.batch_add_embeddings([]) - - assert result['added'] == 0 - - def test_batch_add_with_invalid_vector_dimension(self, mock_opensearch_store): - """Verify handling of invalid vector dimensions.""" - client = VectorBatchClient(vector_store=mock_opensearch_store) - - invalid_embeddings = [ - {'id': 'chunk_001', 'vector': [0.1] * 100, 'text': 'Wrong dimension'} - ] - - client.batch_add_embeddings = Mock(side_effect=ValueError("Invalid vector dimension")) - - with pytest.raises(ValueError, match="Invalid vector dimension"): - client.batch_add_embeddings(invalid_embeddings) - - def test_batch_add_with_missing_vector(self, mock_opensearch_store): - """Verify handling of missing vector data.""" - client = VectorBatchClient(vector_store=mock_opensearch_store) - - invalid_embeddings = [ - {'id': 'chunk_001', 'text': 'No vector'} + + def test_wraps_every_index_the_store_reports(self): + """Verify each of the store's indexes gets a batch wrapper, keyed by name.""" + client = batch_client('chunk', 'statement') + + assert set(client.indexes) == {'chunk', 'statement'} + assert all(isinstance(index, BatchVectorIndex) for index in client.indexes.values()) + + def test_starts_with_nothing_buffered(self): + """Verify no nodes are deferred before anything is written.""" + client = batch_client('chunk') + + assert client.all_nodes == [] + assert client.indexes['chunk'].nodes == [] + +class TestGetIndex: + """Tests for index lookup.""" + + @pytest.mark.parametrize('index_name', ALL_EMBEDDING_INDEXES) + def test_known_index_names_are_accepted(self, index_name): + """Verify every documented index name resolves to something usable.""" + client = batch_client(*ALL_EMBEDDING_INDEXES) + + assert client.get_index(index_name) is client.indexes[index_name] + + def test_unknown_index_name_raises_listing_the_valid_ones(self): + """Verify a typo says what was expected.""" + client = batch_client('chunk') + + with pytest.raises(ValueError) as excinfo: + client.get_index('chunks') + + message = str(excinfo.value) + assert 'chunks' in message + for index_name in ALL_EMBEDDING_INDEXES: + assert index_name in message + + def test_index_the_store_does_not_have_falls_back_to_a_dummy(self): + """Verify a valid but absent index writes nowhere rather than failing. + + A store configured with only some of the indexes should not make a build + that touches the others crash - the writes are simply discarded. + """ + client = batch_client('chunk') + + assert isinstance(client.get_index('topic'), DummyVectorIndex) + + def test_batch_writes_disabled_returns_the_underlying_index(self): + """Verify the wrapper is bypassed when batching is off.""" + client = batch_client('chunk', batch_writes_enabled=False) + + assert client.get_index('chunk') is client.indexes['chunk'].index + +class TestBatchVectorIndex: + """Tests for the per-index batch wrapper.""" + + def test_add_embeddings_defers_rather_than_writing(self): + """Verify nothing reaches the index until the batch is applied.""" + client = batch_client('chunk', batch_write_size=2) + index = client.get_index('chunk') + + index.add_embeddings(['n1', 'n2']) + + client.indexes['chunk'].index.add_embeddings.assert_not_called() + assert client.indexes['chunk'].nodes == ['n1', 'n2'] + + def test_write_embeddings_chunks_by_batch_write_size(self): + """Verify the buffer is written in `batch_write_size` slices, in order.""" + client = batch_client('chunk', batch_write_size=2) + client.get_index('chunk').add_embeddings(['n1', 'n2', 'n3', 'n4', 'n5']) + + client.apply_batch_operations() + + underlying = client.indexes['chunk'].index + assert [call.args[0] for call in underlying.add_embeddings.call_args_list] == [ + ['n1', 'n2'], ['n3', 'n4'], ['n5'] ] - - client.batch_add_embeddings = Mock(side_effect=KeyError("Missing vector")) - - with pytest.raises(KeyError, match="Missing vector"): - client.batch_add_embeddings(invalid_embeddings) + + def test_empty_buffer_writes_nothing(self): + """Verify applying an empty batch does not call the index at all.""" + client = batch_client('chunk') + + client.apply_batch_operations() + + client.indexes['chunk'].index.add_embeddings.assert_not_called() + +class TestAllowYield: + """Tests for node yielding under batching.""" + + def test_batching_defers_the_node(self): + """Verify a node is held back and returned by `apply_batch_operations`.""" + client = batch_client('chunk') + + assert client.allow_yield('n1') is False + assert client.allow_yield('n2') is False + assert client.apply_batch_operations() == ['n1', 'n2'] + + def test_no_batching_yields_immediately(self): + """Verify nothing is held back when batching is off.""" + client = batch_client('chunk', batch_writes_enabled=False) + + assert client.allow_yield('n1') is True + assert client.all_nodes == [] + +class TestContextManager: + """Tests for use as a context manager.""" + + def test_enter_returns_the_client_and_exit_does_not_apply(self): + """Verify leaving the block does not flush. + + `__exit__` deliberately does nothing: the caller decides when to apply, + because it needs the returned nodes. + """ + client = batch_client('chunk') + client.get_index('chunk').add_embeddings(['n1']) + + with client as entered: + assert entered is client + + client.indexes['chunk'].index.add_embeddings.assert_not_called() diff --git a/lexical-graph/tests/unit/indexing/build/test_vector_indexing.py b/lexical-graph/tests/unit/indexing/build/test_vector_indexing.py index a86ed778..4486e2c7 100644 --- a/lexical-graph/tests/unit/indexing/build/test_vector_indexing.py +++ b/lexical-graph/tests/unit/indexing/build/test_vector_indexing.py @@ -1,118 +1,207 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +"""Unit tests for `VectorIndexing`. + +Rewritten against the real API. The earlier version of this file called +`index_document_chunks`, `reindex_existing_chunks` and `delete_indexed_chunks`, +none of which exist - `VectorIndexing` is a `NodeHandler` whose entry point is +`accept` - and it depended on a `mock_opensearch_store` fixture that does not +exist, so every test errored at setup. Nothing noticed, because pytest's default +`norecursedirs` contains `build`, so this whole directory was skipped unless its +path was named explicitly. +""" + +import json +from unittest.mock import MagicMock, patch + import pytest -from unittest.mock import Mock + from graphrag_toolkit.lexical_graph.indexing.build.vector_indexing import VectorIndexing +from graphrag_toolkit.lexical_graph.storage.constants import INDEX_KEY +from graphrag_toolkit.lexical_graph.storage.vector import VectorStore + +from llama_index.core.schema import TextNode + +@pytest.fixture +def mock_vector_store(): + ''' + A vector store reporting one index per embedding index name. `VectorStore` is + a pydantic model, so its `__init__` is patched to allow a mock in the field. + ''' + with patch.object(VectorStore, '__init__', return_value=None): + store = MagicMock(spec=VectorStore) + + indexes = [] + for index_name in ('chunk', 'statement'): + index = MagicMock() + index.index_name = index_name + indexes.append(index) + store.all_indexes.return_value = indexes + return store + +def indexed_node(index_name:str, node_id:str='n-1', text:str='some text', metadata=None) -> TextNode: + node = TextNode(text=text, id_=node_id) + node.metadata = dict(metadata or {}) + node.metadata[INDEX_KEY] = {'index': index_name} + return node + +def added_to(store:MagicMock, index_name:str) -> list: + index = next(i for i in store.all_indexes.return_value if i.index_name == index_name) + return [call.args[0] for call in index.add_embeddings.call_args_list] class TestVectorIndexingInitialization: """Tests for VectorIndexing initialization.""" - - def test_initialization(self, mock_opensearch_store): - """Verify VectorIndexing initializes with vector store.""" - indexer = VectorIndexing(vector_store=mock_opensearch_store) - assert indexer is not None + def test_initialization(self, mock_vector_store): + """Verify VectorIndexing holds the store it was given.""" + indexing = VectorIndexing(vector_store=mock_vector_store) + + assert indexing.vector_store is mock_vector_store + + def test_for_vector_store_adopts_an_existing_store(self, mock_vector_store): + """Verify a caller who already has a store does not pay for the factory.""" + indexing = VectorIndexing.for_vector_store(mock_vector_store) + + assert indexing.vector_store is mock_vector_store class TestVectorIndexingOperations: - """Tests for vector indexing operations.""" - - def test_index_document_chunks(self, mock_opensearch_store): - """Verify indexing document chunks with embeddings.""" - indexer = VectorIndexing(vector_store=mock_opensearch_store) - - chunks = [ - {'id': 'chunk_001', 'text': 'First chunk of text', 'embedding': [0.1] * 384}, - {'id': 'chunk_002', 'text': 'Second chunk of text', 'embedding': [0.2] * 384} - ] - - indexer.index_chunks = Mock(return_value={'indexed': 2}) - result = indexer.index_chunks(chunks) - - assert result is not None - assert result['indexed'] == 2 - - def test_index_with_metadata(self, mock_opensearch_store): - """Verify indexing with metadata.""" - indexer = VectorIndexing(vector_store=mock_opensearch_store) - - chunks = [ - { - 'id': 'chunk_001', - 'text': 'Chunk with metadata', - 'embedding': [0.1] * 384, - 'metadata': {'source': 'doc_001', 'position': 0} - } - ] - - indexer.index_chunks = Mock(return_value={'indexed': 1}) - result = indexer.index_chunks(chunks) - - assert result is not None - assert result['indexed'] == 1 - - def test_reindex_existing_chunks(self, mock_opensearch_store): - """Verify reindexing existing chunks.""" - indexer = VectorIndexing(vector_store=mock_opensearch_store) - - chunks = [ - {'id': 'chunk_001', 'text': 'Updated text', 'embedding': [0.15] * 384} - ] - - indexer.reindex_chunks = Mock(return_value={'reindexed': 1}) - result = indexer.reindex_chunks(chunks) - - assert result is not None - assert result['reindexed'] == 1 - - def test_delete_indexed_chunks(self, mock_opensearch_store): - """Verify deleting indexed chunks.""" - indexer = VectorIndexing(vector_store=mock_opensearch_store) - - chunk_ids = ['chunk_001', 'chunk_002', 'chunk_003'] - - indexer.delete_chunks = Mock(return_value={'deleted': 3}) - result = indexer.delete_chunks(chunk_ids) - - assert result is not None - assert result['deleted'] == 3 - - -class TestVectorIndexingErrorHandling: - """Tests for vector indexing error handling.""" - - def test_index_with_empty_chunks(self, mock_opensearch_store): - """Verify handling of empty chunk list.""" - indexer = VectorIndexing(vector_store=mock_opensearch_store) - - indexer.index_chunks = Mock(return_value={'indexed': 0}) - result = indexer.index_chunks([]) - - assert result['indexed'] == 0 - - def test_index_with_missing_embeddings(self, mock_opensearch_store): - """Verify handling of chunks without embeddings.""" - indexer = VectorIndexing(vector_store=mock_opensearch_store) - - chunks = [ - {'id': 'chunk_001', 'text': 'Chunk without embedding'} - ] - - indexer.index_chunks = Mock(side_effect=ValueError("Missing embedding")) - - with pytest.raises(ValueError, match="Missing embedding"): - indexer.index_chunks(chunks) - - def test_index_with_invalid_embedding_dimension(self, mock_opensearch_store): - """Verify handling of invalid embedding dimensions.""" - indexer = VectorIndexing(vector_store=mock_opensearch_store) - - chunks = [ - {'id': 'chunk_001', 'text': 'Chunk', 'embedding': [0.1] * 100} - ] - - indexer.index_chunks = Mock(side_effect=ValueError("Invalid embedding dimension")) - - with pytest.raises(ValueError, match="Invalid embedding dimension"): - indexer.index_chunks(chunks) + """Tests for the accept() indexing path.""" + + def test_node_is_added_to_the_index_named_in_its_metadata(self, mock_vector_store): + """Verify routing is by `INDEX_KEY`, not by node type.""" + indexing = VectorIndexing(vector_store=mock_vector_store) + + list(indexing.accept( + [indexed_node('statement', node_id='s-1')], + batch_writes_enabled=False, + batch_write_size=10, + )) + + assert len(added_to(mock_vector_store, 'statement')) == 1 + assert added_to(mock_vector_store, 'chunk') == [] + + def test_node_without_index_metadata_is_yielded_but_not_indexed(self, mock_vector_store): + """Verify unrelated nodes pass through untouched.""" + indexing = VectorIndexing(vector_store=mock_vector_store) + + node = TextNode(text='some text', id_='n-1') + node.metadata = {} + + results = list(indexing.accept([node], batch_writes_enabled=False, batch_write_size=10)) + + assert results == [node] + assert added_to(mock_vector_store, 'chunk') == [] + + def test_unknown_index_name_is_ignored(self, mock_vector_store): + """Verify an index name outside the known set indexes nothing. + + Guarded rather than left to `get_index`, which would raise - a node + labelled with something unexpected should not fail the whole build. + """ + indexing = VectorIndexing(vector_store=mock_vector_store) + + results = list(indexing.accept( + [indexed_node('something_else')], + batch_writes_enabled=False, + batch_write_size=10, + )) + + assert len(results) == 1 + assert added_to(mock_vector_store, 'chunk') == [] + + def test_batching_defers_writes_and_yields_after_apply(self, mock_vector_store): + """Verify batched nodes are written and yielded once, at the end.""" + indexing = VectorIndexing(vector_store=mock_vector_store) + + nodes = [indexed_node('chunk', node_id=f'c-{i}') for i in range(3)] + + results = list(indexing.accept(nodes, batch_writes_enabled=True, batch_write_size=2)) + + assert [node.node_id for node in results] == ['c-0', 'c-1', 'c-2'] + assert [len(batch) for batch in added_to(mock_vector_store, 'chunk')] == [2, 1] + + def test_indexing_error_is_logged_and_reraised(self, mock_vector_store): + """Verify a failing index write is not swallowed.""" + indexing = VectorIndexing(vector_store=mock_vector_store) + + index = next(i for i in mock_vector_store.all_indexes.return_value if i.index_name == 'chunk') + index.add_embeddings.side_effect = RuntimeError('index unavailable') + + with pytest.raises(RuntimeError, match='index unavailable'): + list(indexing.accept( + [indexed_node('chunk')], + batch_writes_enabled=False, + batch_write_size=10, + )) + +class TestIndexableTransformations: + """Tests for the transformations applied before indexing.""" + + def test_json_content_is_rewritten_as_yaml(self, mock_vector_store): + """Verify JSON node content is indexed as YAML. + + Embedded text is what the retriever's similarity is computed over, so the + substitution is behaviour, not formatting. + """ + indexing = VectorIndexing(vector_store=mock_vector_store) + + node = indexed_node('chunk', text=json.dumps({'value': 'some statement'})) + + list(indexing.accept([node], batch_writes_enabled=False, batch_write_size=10)) + + (indexed,) = added_to(mock_vector_store, 'chunk')[0] + assert indexed.get_content() == 'value: some statement\n' + # The original node is not mutated - only the copy that gets indexed. + assert node.get_content() == json.dumps({'value': 'some statement'}) + + def test_non_json_content_is_left_alone(self, mock_vector_store): + """Verify plain text is indexed as written.""" + indexing = VectorIndexing(vector_store=mock_vector_store) + + list(indexing.accept( + [indexed_node('chunk', text='just prose')], + batch_writes_enabled=False, + batch_write_size=10, + )) + + (indexed,) = added_to(mock_vector_store, 'chunk')[0] + assert indexed.get_content() == 'just prose' + + def test_datetime_source_metadata_is_normalized(self, mock_vector_store): + """Verify a `_date`-suffixed source field is formatted before indexing.""" + indexing = VectorIndexing(vector_store=mock_vector_store) + + node = indexed_node('chunk', metadata={'source': {'metadata': {'publish_date': '2024-01-15'}}}) + + list(indexing.accept([node], batch_writes_enabled=False, batch_write_size=10)) + + (indexed,) = added_to(mock_vector_store, 'chunk')[0] + assert indexed.metadata['source']['metadata']['publish_date'].startswith('2024-01-15') + + def test_unparseable_datetime_source_metadata_is_dropped(self, mock_vector_store): + """Verify a `_date` field that cannot be parsed is removed, not indexed raw. + + Indexing the raw value would put an unparseable date into the embedded + text and into any metadata filter built over it. + """ + indexing = VectorIndexing(vector_store=mock_vector_store) + + node = indexed_node('chunk', metadata={'source': {'metadata': {'publish_date': 'not a date'}}}) + + list(indexing.accept([node], batch_writes_enabled=False, batch_write_size=10)) + + (indexed,) = added_to(mock_vector_store, 'chunk')[0] + assert 'publish_date' not in indexed.metadata['source']['metadata'] + + def test_non_datetime_source_metadata_is_untouched(self, mock_vector_store): + """Verify only `_date`-suffixed keys are rewritten.""" + indexing = VectorIndexing(vector_store=mock_vector_store) + + node = indexed_node('chunk', metadata={'source': {'metadata': {'title': 'A title'}}}) + + list(indexing.accept([node], batch_writes_enabled=False, batch_write_size=10)) + + (indexed,) = added_to(mock_vector_store, 'chunk')[0] + assert indexed.metadata['source']['metadata'] == {'title': 'A title'} diff --git a/lexical-graph/tests/unit/indexing/build/test_version_manager.py b/lexical-graph/tests/unit/indexing/build/test_version_manager.py index d6652f0d..86ca0080 100644 --- a/lexical-graph/tests/unit/indexing/build/test_version_manager.py +++ b/lexical-graph/tests/unit/indexing/build/test_version_manager.py @@ -1,145 +1,404 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +"""Unit tests for `VersionManager`. + +Rewritten against the real API. `VersionManager` versions *source documents* on a +`valid_from`/`valid_to` timeline; it has nothing to do with semantic versioning of +the library. The earlier version of this file tested `get_version`, +`increment_major`, `compare` and `is_compatible` - none of which exist - by +assigning each a `Mock` and then asserting the mock's own return value, so every +test passed on a class it never called. It also constructed `VersionManager()` and +`VersionManager(version=...)`, while the real class is a pydantic `NodeHandler` +with required `graph_store` and `vector_store` fields, so setup raised anyway. +Nothing noticed, because pytest's default `norecursedirs` contains `build`, so +this whole directory was skipped unless its path was named explicitly. +""" + +from unittest.mock import MagicMock, Mock, patch + import pytest -from unittest.mock import Mock + +from graphrag_toolkit.lexical_graph.errors import IndexError as GraphRAGIndexError from graphrag_toolkit.lexical_graph.indexing.build.version_manager import VersionManager +from graphrag_toolkit.lexical_graph.storage.constants import INDEX_KEY +from graphrag_toolkit.lexical_graph.storage.vector import DummyVectorIndex, VectorStore +from graphrag_toolkit.lexical_graph.versioning import ( + TIMESTAMP_UPPER_BOUND, + VALID_TO, + VERSION_INDEPENDENT_ID_FIELDS, +) + +from llama_index.core.schema import TextNode + +def vector_index(index_name:str) -> MagicMock: + index = MagicMock() + index.index_name = index_name + index.underlying_index_name.return_value = f'underlying_{index_name}' + # An empty list means "nothing failed"; the default MagicMock return is truthy + # and would send `_update_vector_store_versions` into its retry sleeps. + index.update_versioning.return_value = [] + return index +@pytest.fixture +def mock_vector_store(): + ''' + A vector store reporting a single chunk index. `VectorStore` is a pydantic + model, so its `__init__` is patched to allow a mock into the field. + ''' + with patch.object(VectorStore, '__init__', return_value=None): + store = MagicMock(spec=VectorStore) + store.all_indexes.return_value = [vector_index('chunk')] + return store + +@pytest.fixture +def version_manager(mock_neptune_store, mock_vector_store): + return VersionManager( + graph_store=mock_neptune_store, + vector_store=mock_vector_store, + show_progress=False, + ) + +def source_node(source_id:str='source-1', valid_from:int=200, id_fields=None, versioning=None) -> TextNode: + node = TextNode(text='', id_=source_id) + node.metadata = { + 'source': { + 'sourceId': source_id, + 'metadata': {'url': 'https://example.com/doc'}, + 'versioning': versioning if versioning is not None else {'valid_from': valid_from, 'id_fields': id_fields}, + }, + INDEX_KEY: {'index': 'source'}, + } + return node class TestVersionManagerInitialization: """Tests for VersionManager initialization.""" - - def test_initialization(self): - """Verify VersionManager initializes correctly.""" - manager = VersionManager() - assert manager is not None - - def test_initialization_with_version(self): - """Verify VersionManager initializes with specific version.""" - manager = VersionManager(version='1.0.0') - assert manager is not None - - -class TestVersionManagement: - """Tests for version management operations.""" - - def test_get_current_version(self): - """Verify getting current version.""" - manager = VersionManager(version='1.2.3') - - manager.get_version = Mock(return_value='1.2.3') - version = manager.get_version() - - assert version == '1.2.3' - - def test_set_version(self): - """Verify setting version.""" - manager = VersionManager() - - manager.set_version = Mock() - manager.set_version('2.0.0') - - manager.set_version.assert_called_once_with('2.0.0') - - def test_increment_major_version(self): - """Verify incrementing major version.""" - manager = VersionManager(version='1.2.3') - - manager.increment_major = Mock(return_value='2.0.0') - new_version = manager.increment_major() - - assert new_version == '2.0.0' - - def test_increment_minor_version(self): - """Verify incrementing minor version.""" - manager = VersionManager(version='1.2.3') - - manager.increment_minor = Mock(return_value='1.3.0') - new_version = manager.increment_minor() - - assert new_version == '1.3.0' - - def test_increment_patch_version(self): - """Verify incrementing patch version.""" - manager = VersionManager(version='1.2.3') - - manager.increment_patch = Mock(return_value='1.2.4') - new_version = manager.increment_patch() - - assert new_version == '1.2.4' - - def test_compare_versions(self): - """Verify version comparison.""" - manager = VersionManager(version='1.2.3') - - manager.compare = Mock(return_value=-1) - result = manager.compare('2.0.0') - - assert result == -1 # 1.2.3 < 2.0.0 - - def test_is_compatible(self): - """Verify version compatibility check.""" - manager = VersionManager(version='1.2.3') - - manager.is_compatible = Mock(return_value=True) - result = manager.is_compatible('1.2.0') - - assert result is True - - -class TestVersionManagerErrorHandling: - """Tests for version manager error handling.""" - - def test_set_invalid_version_format(self): - """Verify handling of invalid version format.""" - manager = VersionManager() - - manager.set_version = Mock(side_effect=ValueError("Invalid version format")) - - with pytest.raises(ValueError, match="Invalid version format"): - manager.set_version('invalid.version') - - def test_compare_with_invalid_version(self): - """Verify handling of invalid version in comparison.""" - manager = VersionManager(version='1.2.3') - - manager.compare = Mock(side_effect=ValueError("Invalid version")) - - with pytest.raises(ValueError, match="Invalid version"): - manager.compare('not.a.version') - - def test_initialization_with_invalid_version(self): - """Verify handling of invalid version during initialization.""" - with pytest.raises((ValueError, TypeError)): - VersionManager(version='invalid') - - -class TestVersionManagerEdgeCases: - """Tests for version manager edge cases.""" - - def test_version_with_prerelease(self): - """Verify handling of prerelease versions.""" - manager = VersionManager(version='1.2.3-alpha') - - manager.get_version = Mock(return_value='1.2.3-alpha') - version = manager.get_version() - - assert 'alpha' in version - - def test_version_with_build_metadata(self): - """Verify handling of build metadata.""" - manager = VersionManager(version='1.2.3+build.123') - - manager.get_version = Mock(return_value='1.2.3+build.123') - version = manager.get_version() - - assert 'build' in version - - def test_version_equality(self): - """Verify version equality check.""" - manager = VersionManager(version='1.2.3') - - manager.compare = Mock(return_value=0) - result = manager.compare('1.2.3') - - assert result == 0 # Equal versions + + def test_initialization(self, mock_neptune_store, mock_vector_store): + """Verify the manager holds the two stores it was given.""" + manager = VersionManager(graph_store=mock_neptune_store, vector_store=mock_vector_store) + + assert manager.graph_store is mock_neptune_store + assert manager.vector_store is mock_vector_store + + def test_both_stores_are_required(self, mock_neptune_store): + """Verify neither store defaults. + + Versioning has to read the graph and write the vector indexes, so a + manager missing either would fail partway through a build rather than at + construction. + """ + from pydantic import ValidationError + + with pytest.raises(ValidationError): + VersionManager(graph_store=mock_neptune_store) + + def test_for_graph_and_vector_store_adopts_existing_stores(self, mock_neptune_store, mock_vector_store): + """Verify a caller who already has stores does not go through the factories.""" + manager = VersionManager.for_graph_and_vector_store(mock_neptune_store, mock_vector_store) + + assert manager.graph_store is mock_neptune_store + assert manager.vector_store is mock_vector_store + +class TestGetUpdates: + """Tests for the timeline arithmetic in `_get_updates`.""" + + def test_first_version_of_a_source_is_open_ended(self, version_manager): + """Verify a source with no history is valid until the upper bound.""" + (source_node_result, adjustments) = version_manager._get_updates( + {'source_id': 's1', 'valid_from': 100, 'valid_to': None}, [] + ) + + assert source_node_result['valid_to'] == TIMESTAMP_UPPER_BOUND + assert adjustments == [] + + def test_newer_version_archives_the_previous_latest(self, version_manager): + """Verify the incoming version closes off the one it supersedes.""" + (new_node, adjustments) = version_manager._get_updates( + {'source_id': 's2', 'valid_from': 200, 'valid_to': None}, + [{'source_id': 's1', 'valid_from': 100, 'valid_to': TIMESTAMP_UPPER_BOUND}], + ) + + assert new_node['valid_to'] == TIMESTAMP_UPPER_BOUND + assert adjustments == [{'source_id': 's1', 'valid_from': 100, 'valid_to': 200}] + + def test_earliest_version_is_closed_by_the_existing_earliest(self, version_manager): + """Verify back-dated content is inserted below the timeline, not on top of it. + + A document arriving out of order must not claim to be current, and must + not disturb the version that already is. + """ + (new_node, adjustments) = version_manager._get_updates( + {'source_id': 's0', 'valid_from': 50, 'valid_to': None}, + [{'source_id': 's1', 'valid_from': 100, 'valid_to': TIMESTAMP_UPPER_BOUND}], + ) + + assert new_node['valid_to'] == 100 + assert adjustments == [] + + def test_same_valid_from_reuses_the_existing_valid_to(self, version_manager): + """Verify a re-ingest of the same version occupies the same interval.""" + (new_node, adjustments) = version_manager._get_updates( + {'source_id': 's1-again', 'valid_from': 100, 'valid_to': None}, + [{'source_id': 's1', 'valid_from': 100, 'valid_to': 300}], + ) + + assert new_node['valid_to'] == 300 + assert adjustments == [] + + def test_version_inserted_between_two_existing_versions(self, version_manager): + """Verify a mid-timeline insert takes the interval up to the next version.""" + (new_node, adjustments) = version_manager._get_updates( + {'source_id': 's-mid', 'valid_from': 150, 'valid_to': None}, + [ + {'source_id': 's1', 'valid_from': 100, 'valid_to': 200}, + {'source_id': 's2', 'valid_from': 200, 'valid_to': TIMESTAMP_UPPER_BOUND}, + ], + ) + + assert new_node['valid_to'] == 200 + assert adjustments == [{'source_id': 's1', 'valid_from': 100, 'valid_to': 150}] + + def test_existing_nodes_are_sorted_before_use(self, version_manager): + """Verify the result does not depend on the order the query returned rows in.""" + existing = [ + {'source_id': 's2', 'valid_from': 200, 'valid_to': TIMESTAMP_UPPER_BOUND}, + {'source_id': 's1', 'valid_from': 100, 'valid_to': 200}, + ] + + (ascending, _) = version_manager._get_updates( + {'source_id': 's3', 'valid_from': 300, 'valid_to': None}, list(reversed(existing)) + ) + (descending, _) = version_manager._get_updates( + {'source_id': 's3', 'valid_from': 300, 'valid_to': None}, existing + ) + + assert ascending == descending + +class TestGetExistingSourceNodes: + """Tests for the lookup of source nodes already in the graph.""" + + def test_no_id_fields_queries_nothing(self, version_manager): + """Verify versioning is opt-in via `id_fields`.""" + assert version_manager._get_existing_source_nodes(None, source_node()) == [] + assert not version_manager.graph_store.execute_query.called + + def test_id_field_absent_from_the_node_queries_nothing(self, version_manager): + """Verify a declared id field with no value does not match everything. + + Without the value there are no other filter criteria, and querying on the + id-fields marker alone would return every versioned source in the graph. + """ + node = source_node(id_fields=['url']) + node.metadata['source']['metadata'] = {} + + assert version_manager._get_existing_source_nodes(['url'], node) == [] + assert not version_manager.graph_store.execute_query.called + + def test_query_filters_on_the_id_field_and_its_marker(self, version_manager): + """Verify both the field value and the id-fields marker are in the query.""" + version_manager.graph_store.execute_query = Mock(return_value=[ + {'result': {'source_id': 's1', 'valid_from': 100, 'valid_to': TIMESTAMP_UPPER_BOUND}} + ]) + + results = version_manager._get_existing_source_nodes(['url'], source_node(id_fields=['url'])) + + (cypher, parameters) = version_manager.graph_store.execute_query.call_args.args + assert '__Source__' in cypher + assert VERSION_INDEPENDENT_ID_FIELDS in cypher + assert parameters == {'versionIndependentIdFields': 'url'} + assert results == [{'source_id': 's1', 'valid_from': 100, 'valid_to': TIMESTAMP_UPPER_BOUND}] + + def test_multiple_id_fields_are_joined_into_one_marker(self, version_manager): + """Verify the marker is the semicolon-joined field list.""" + version_manager.graph_store.execute_query = Mock(return_value=[]) + + version_manager._get_existing_source_nodes( + ['url', 'title'], + source_node(), + ) + + (_, parameters) = version_manager.graph_store.execute_query.call_args.args + assert parameters['versionIndependentIdFields'] == 'url;title' + +class TestGetNodeIds: + """Tests for resolving the graph nodes belonging to a source.""" + + @pytest.mark.parametrize( + ('index_name', 'id_field'), + [('chunk', 'chunkId'), ('topic', 'topicId'), ('statement', 'statementId')], + ) + def test_each_index_has_its_own_traversal(self, version_manager, index_name, id_field): + """Verify the query collects the ids of the index it was asked about.""" + version_manager.graph_store.execute_query = Mock(return_value=[ + {'result': {'sourceId': 's1', 'nodeIds': ['n1', 'n2']}} + ]) + + result = version_manager._get_node_ids(vector_index(index_name), ['s1']) + + (cypher, parameters) = version_manager.graph_store.execute_query.call_args.args + assert id_field in cypher + assert parameters == {'sourceIds': ['s1']} + assert result == {'s1': ['n1', 'n2']} + + def test_unknown_index_name_raises(self, version_manager): + """Verify an index with no traversal is an error, not a silent no-op.""" + with pytest.raises(ValueError, match='Invalid index name: fact'): + version_manager._get_node_ids(vector_index('fact'), ['s1']) + + def test_dummy_index_returns_empty(self, version_manager): + """Verify a store without a real index is skipped without querying. + + Returns an empty *list* where every other path returns a dict; callers + only iterate `.items()` on a non-empty result, so it happens to work. + Pinned so a change to either side is deliberate. + """ + assert version_manager._get_node_ids(DummyVectorIndex(index_name='chunk'), ['s1']) == [] + assert not version_manager.graph_store.execute_query.called + +class TestSetSourceNodeVersionInfo: + """Tests for closing off a source node in the graph.""" + + def test_valid_to_and_id_fields_are_bound_not_interpolated(self, version_manager): + """Verify the values travel as parameters.""" + version_manager._set_source_node_version_info('s1', 500, ['url', 'title']) + + (cypher, properties) = version_manager.graph_store.execute_query_with_retry.call_args.args + assert VALID_TO in cypher + assert properties == { + 'sourceId': 's1', + 'versioningTimestamp': 500, + 'versionIndependentIdFields': 'url;title', + } + +class TestUpdateVectorStoreVersions: + """Tests for propagating version info into a vector index.""" + + def test_node_ids_are_written_in_batches_of_100(self, version_manager): + """Verify a large id list is chunked rather than sent whole.""" + index = vector_index('chunk') + + version_manager._update_vector_store_versions('s1', [f'n{i}' for i in range(250)], 500, index) + + batches = [call.args[1] for call in index.update_versioning.call_args_list] + assert [len(batch) for batch in batches] == [100, 100, 50] + assert all(call.args[0] == 500 for call in index.update_versioning.call_args_list) + + def test_transient_failure_is_retried(self, version_manager): + """Verify a batch reporting failed ids is attempted again.""" + index = vector_index('chunk') + index.update_versioning.side_effect = [['n1'], []] + + with patch('graphrag_toolkit.lexical_graph.indexing.build.version_manager.time.sleep'): + version_manager._update_vector_store_versions('s1', ['n1', 'n2'], 500, index) + + assert index.update_versioning.call_count == 2 + + def test_persistent_failure_raises(self, version_manager): + """Verify ids that never succeed are surfaced, not dropped. + + A silently unversioned chunk would keep being returned by current-version + retrieval after its source had been superseded. + """ + index = vector_index('chunk') + index.update_versioning.return_value = ['n1'] + + with patch('graphrag_toolkit.lexical_graph.indexing.build.version_manager.time.sleep'): + with pytest.raises(GraphRAGIndexError, match='Failed to update valid_to version info'): + version_manager._update_vector_store_versions('s1', ['n1'], 500, index) + + assert index.update_versioning.call_count == 5 + +class TestAccept: + """Tests for the accept() entry point.""" + + def test_source_node_gets_its_version_window_written_into_metadata(self, version_manager): + """Verify downstream builders see the resolved window, not the raw input.""" + node = source_node(valid_from=200) + + results = list(version_manager.accept([node])) + + assert results == [node] + versioning = node.metadata['source']['versioning'] + assert versioning['valid_from'] == 200 + assert versioning['valid_to'] == TIMESTAMP_UPPER_BOUND + assert versioning['prev_versions'] == [] + + def test_superseded_source_is_closed_in_both_stores(self, version_manager): + """Verify an archived version is updated in the graph and the vector index.""" + version_manager.graph_store.execute_query = Mock(side_effect=[ + [{'result': {'source_id': 's-old', 'valid_from': 100, 'valid_to': TIMESTAMP_UPPER_BOUND}}], + [{'result': {'sourceId': 's-old', 'nodeIds': ['c1', 'c2']}}], + ]) + + node = source_node(source_id='s-new', valid_from=200, id_fields=['url']) + + list(version_manager.accept([node])) + + assert node.metadata['source']['versioning']['prev_versions'] == ['s-old'] + + (index,) = version_manager.vector_store.all_indexes.return_value + assert index.update_versioning.call_args.args == (200, ['c1', 'c2']) + + (_, properties) = version_manager.graph_store.execute_query_with_retry.call_args.args + assert properties == { + 'sourceId': 's-old', + 'versioningTimestamp': 200, + 'versionIndependentIdFields': 'url', + } + + def test_non_source_node_inherits_its_sources_window(self, version_manager): + """Verify a chunk is stamped with the window resolved for its source.""" + source = source_node(source_id='s1', valid_from=200) + + chunk = TextNode(text='some text', id_='c1') + chunk.metadata = { + 'source': {'sourceId': 's1', 'versioning': {}}, + INDEX_KEY: {'index': 'chunk'}, + } + + list(version_manager.accept([source, chunk])) + + assert chunk.metadata['source']['versioning'] == { + 'valid_from': 200, + 'valid_to': TIMESTAMP_UPPER_BOUND, + } + + def test_non_source_node_seen_before_its_source_is_left_alone(self, version_manager): + """Verify an out-of-order chunk is yielded rather than stamped or dropped. + + Source nodes precede their chunks in the pipeline, so this should not + happen; if it does, the chunk keeps whatever window it arrived with. + """ + chunk = TextNode(text='some text', id_='c1') + chunk.metadata = { + 'source': {'sourceId': 'unseen', 'versioning': {}}, + INDEX_KEY: {'index': 'chunk'}, + } + + results = list(version_manager.accept([chunk])) + + assert results == [chunk] + assert chunk.metadata['source']['versioning'] == {} + + def test_node_without_index_metadata_is_passed_through(self, version_manager): + """Verify unrelated nodes are neither versioned nor lost.""" + node = TextNode(text='some text', id_='n1') + node.metadata = {} + + assert list(version_manager.accept([node])) == [node] + assert not version_manager.graph_store.execute_query_with_retry.called + + def test_source_without_valid_from_falls_back_to_the_extract_timestamp_key(self, version_manager): + """Verify the documented fallback is a literal string, not a timestamp. + + `versioning.get('valid_from', 'extract_timestamp')` yields the *name* of + the field rather than a value, so a source built without a `valid_from` + ends up with a string where every comparison expects an int. Pinned as + current behaviour, not endorsed. + """ + node = source_node(versioning={}) + + list(version_manager.accept([node])) + + assert node.metadata['source']['versioning']['valid_from'] == 'extract_timestamp' diff --git a/lexical-graph/tests/unit/indexing/extract/test_ontology.py b/lexical-graph/tests/unit/indexing/extract/test_ontology.py new file mode 100644 index 00000000..ca866611 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ontology.py @@ -0,0 +1,391 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Loading an ontology, indexing it, and resolving emitted names against it. + +Correctness only. Nothing here asserts a rate, a count over model output, or a +comparison between two configurations - those belong in a measurement document, +not in a test that gates a build. +""" + +from pathlib import Path + +import pytest +from rdflib import Graph + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.naming import ( + camel_to_upper_snake, + resolution_key, + title_case_with_spaces, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology import ( + Ontology, + OntologyLoadError, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_index import ( + OntologyClass, + OntologyIndex, +) +from graphrag_toolkit.lexical_graph.indexing.utils.topic_utils import ( + format_classification, + format_value, +) + +FIXTURES = Path(__file__).parent.parent.parent.parent / 'fixtures' / 'ontologies' +COMPANY = FIXTURES / 'company.ttl' +NS = 'http://example.com/company#' +XSD = 'http://www.w3.org/2001/XMLSchema#' + +@pytest.fixture(scope='module') +def index(): + return Ontology.load(COMPANY).index() + +class TestNaming: + """`naming.py`'s three jobs, and the invariant that ties two of them together.""" + + @pytest.mark.parametrize('name,expected', [ + ('worksFor', 'works for'), + ('WORKS FOR', 'works for'), + ('WORKS_FOR', 'works for'), + ('works for', 'works for'), + ('SportsTeam', 'sports team'), + ('Sports Team', 'sports team'), + ('', ''), + (None, ''), + ]) + def test_resolution_key_folds_every_convention_onto_one(self, name, expected): + assert resolution_key(name) == expected + + def test_resolution_key_does_not_fold_word_boundaries(self): + """The documented limit of the filter's reach: one word is not two.""" + assert resolution_key('Sportsteam') != resolution_key('SportsTeam') + + @pytest.mark.parametrize('name,expected', [ + ('worksFor', 'WORKS_FOR'), + ('WORKS_FOR', 'WORKS_FOR'), + ('founded year', 'FOUNDED_YEAR'), + ('', ''), + ]) + def test_properties_render_as_upper_snake(self, name, expected): + assert camel_to_upper_snake(name) == expected + + @pytest.mark.parametrize('name,expected', [ + ('SportsTeam', 'Sports Team'), + ('SPORTS_TEAM', 'Sports Team'), + ('Sports Team', 'Sports Team'), + ('', ''), + ]) + def test_classes_render_as_title_case(self, name, expected): + assert title_case_with_spaces(name) == expected + + def test_a_rendered_name_survives_the_parser_and_still_resolves(self, index): + """The invariant the module exists for, over every term in the ontology. + + Render the term for the prompt, put it through the transform the response + parser applies, and it must still fold onto the term's own key. This is + what makes `title_case_with_spaces` rather than the local name the right + rendering: `SportsTeam` comes back from the parser as `Sportsteam`. + """ + for ontology_class in index.classes.values(): + rendered = title_case_with_spaces(ontology_class.local_name) + assert resolution_key(format_classification(rendered)) == \ + resolution_key(ontology_class.local_name) + + properties = list(index.object_properties.values()) + list(index.datatype_properties.values()) + for prop in properties: + rendered = camel_to_upper_snake(prop.local_name) + assert resolution_key(format_value(rendered)) == resolution_key(prop.local_name) + +class TestLoading: + """`Ontology.load` accepts what configuration accepts, and refuses the rest.""" + + def test_a_path_loads(self): + assert Ontology.load(COMPANY).index().classes + + def test_a_string_path_loads(self): + assert Ontology.load(str(COMPANY)).index().classes + + def test_a_graph_is_adopted(self): + graph = Graph() + graph.parse(source=str(COMPANY), format='turtle') + assert len(Ontology.load(graph).index().classes) == 5 + + def test_an_ontology_is_returned_unchanged(self): + ontology = Ontology.load(COMPANY) + assert Ontology.load(ontology) is ontology + + def test_a_turtle_string_needs_its_own_constructor(self): + turtle = f'@prefix : <{NS}> . @prefix owl: . :Thing a owl:Class .' + assert Ontology.from_turtle_string(turtle).index().classes + + @pytest.mark.parametrize('source', [42, None, ['company.ttl']]) + def test_an_unsupported_source_raises(self, source): + with pytest.raises(OntologyLoadError, match='Cannot load an ontology'): + Ontology.load(source) + + def test_a_non_turtle_suffix_raises(self, tmp_path): + path = tmp_path / 'company.owl' + path.write_text('') + with pytest.raises(OntologyLoadError, match='Unsupported ontology file format'): + Ontology.load(path) + + def test_a_missing_file_raises(self, tmp_path): + with pytest.raises(OntologyLoadError, match='not found'): + Ontology.load(tmp_path / 'absent.ttl') + + @pytest.mark.parametrize('name', [ + 'malformed_syntax.ttl', + 'missing_datatype_range.ttl', + 'non_xsd_range.ttl', + 'dual_typed_property.ttl', + 'dangling_subclass_reference.ttl', + 'dangling_domain_reference.ttl', + 'dangling_range_reference.ttl', + 'subclass_cycle.ttl', + ]) + def test_a_structurally_invalid_ontology_is_refused_at_load_time(self, name): + """Every failure is refused where it can still be reported to a user. + + Loading happens in the parent process at configuration time; the filter + runs in a spawn worker per node batch. Anything not caught here surfaces + as a warning nobody reads. + """ + with pytest.raises(OntologyLoadError): + Ontology.load(FIXTURES / 'malformed' / name) + + def test_the_message_names_the_term_and_the_rule(self): + """A load error a user cannot act on is a load error that wasted their time.""" + with pytest.raises(OntologyLoadError, match='not an XSD datatype'): + Ontology.load(FIXTURES / 'malformed' / 'non_xsd_range.ttl') + +class TestTheIndex: + """What `OntologyIndex` says about a loaded ontology.""" + + def test_every_declared_term_is_indexed_under_its_kind(self, index): + assert set(index.classes) == { + f'{NS}Agent', f'{NS}Person', f'{NS}Athlete', f'{NS}Company', f'{NS}SportsTeam', + } + assert set(index.object_properties) == { + f'{NS}worksFor', f'{NS}playsFor', f'{NS}subsidiaryOf', f'{NS}acquired', + } + assert len(index.datatype_properties) == 7 + + def test_a_declared_label_and_alias_are_carried(self, index): + sports_team = index.classes[f'{NS}SportsTeam'] + assert (sports_team.local_name, sports_team.label) == ('SportsTeam', 'Sports Team') + assert sports_team.aliases == ['Ball Club'] + + def test_ancestors_are_the_reflexive_transitive_closure(self, index): + assert index.classes[f'{NS}SportsTeam'].ancestors == { + f'{NS}SportsTeam', f'{NS}Company', f'{NS}Agent', + } + assert index.classes[f'{NS}Agent'].ancestors == {f'{NS}Agent'} + + @pytest.mark.parametrize('child,parent,expected', [ + ('SportsTeam', 'Company', True), + ('SportsTeam', 'Agent', True), + ('SportsTeam', 'SportsTeam', True), + ('Company', 'SportsTeam', False), + ('Athlete', 'Company', False), + ]) + def test_is_subclass_of_honours_the_closure(self, index, child, parent, expected): + assert index.is_subclass_of(f'{NS}{child}', f'{NS}{parent}') is expected + + def test_is_subclass_of_an_undeclared_iri_is_false_not_an_error(self, index): + """Callers pass unresolved classifications straight through.""" + assert index.is_subclass_of(f'{NS}Absent', f'{NS}Agent') is False + + def test_a_declared_domain_range_and_datatype_are_carried(self, index): + works_for = index.object_properties[f'{NS}worksFor'] + assert (works_for.domain, works_for.range) == (f'{NS}Person', f'{NS}Company') + assert index.datatype_properties[f'{NS}foundedYear'].datatype == f'{XSD}integer' + + def test_an_undeclared_domain_or_range_is_none_meaning_anything(self, index): + acquired = index.object_properties[f'{NS}acquired'] + assert (acquired.domain, acquired.range) == (None, None) + assert index.datatype_properties[f'{NS}officialName'].domain is None + +class TestResolution: + """Turning a name the model emitted into a declared term.""" + + @pytest.mark.parametrize('emitted', ['Sports Team', 'SPORTS_TEAM', 'SportsTeam', 'sports team']) + def test_a_class_resolves_from_any_convention(self, index, emitted): + assert index.resolve_class(emitted).iri == f'{NS}SportsTeam' + + @pytest.mark.parametrize('emitted,iri', [('Ball Club', 'SportsTeam'), ('Corporation', 'Company')]) + def test_a_class_resolves_from_a_declared_alias(self, index, emitted, iri): + assert index.resolve_class(emitted).iri == f'{NS}{iri}' + + @pytest.mark.parametrize('emitted', ['WORKS FOR', 'WORKS_FOR', 'worksFor', 'REQ_TO_HC']) + def test_an_object_predicate_resolves_from_any_convention_or_alias(self, index, emitted): + assert index.resolve_object_predicate(emitted).iri == f'{NS}worksFor' + + @pytest.mark.parametrize('emitted', ['FOUNDED YEAR', 'foundedYear', 'founded_year']) + def test_a_datatype_predicate_resolves_from_any_convention(self, index, emitted): + assert index.resolve_datatype_predicate(emitted).iri == f'{NS}foundedYear' + + @pytest.mark.parametrize('emitted', ['HIRED BY', 'EMPLOYER', '', 'Sportsteam']) + def test_a_name_the_ontology_does_not_carry_resolves_to_nothing(self, index, emitted): + assert index.resolve_class(emitted) is None + assert index.resolve_object_predicate(emitted) is None + assert index.resolve_datatype_predicate(emitted) is None + + def test_the_two_predicate_kinds_do_not_answer_for_each_other(self, index): + assert index.resolve_datatype_predicate('WORKS FOR') is None + assert index.resolve_object_predicate('FOUNDED YEAR') is None + + def test_a_local_name_outranks_another_terms_alias(self): + """Resolution precedence, which decides a genuine authoring collision. + + Two classes claim the key `company`: one owns it as its local name, the + other as an alias. The owner wins, so resolution is a function of the + ontology's content and not of dict iteration order. + """ + index = OntologyIndex(classes={ + 'urn:x#Company': OntologyClass(iri='urn:x#Company', local_name='Company'), + 'urn:x#Firm': OntologyClass(iri='urn:x#Firm', local_name='Firm', aliases=['Company']), + }) + + assert index.resolve_class('Company').iri == 'urn:x#Company' + + def test_the_key_maps_are_derived_and_survive_a_round_trip(self): + """The index is pickled into a spawn worker, so a rebuilt one must resolve. + + Keys passed in are discarded and recomputed, so an index cannot carry a + lookup table that disagrees with its terms. + """ + index = OntologyIndex( + classes={'urn:x#Company': OntologyClass(iri='urn:x#Company', local_name='Company')}, + class_by_key={'nonsense': 'urn:x#Absent'}, + ) + + assert index.resolve_class('Company').iri == 'urn:x#Company' + assert index.resolve_class('nonsense') is None + + rebuilt = OntologyIndex.model_validate(index.model_dump()) + assert rebuilt.resolve_class('COMPANY').iri == 'urn:x#Company' + +class TestOntologiesThatAreNotTheFixture: + """Shapes a real ontology has that `company.ttl` deliberately does not.""" + + PREFIXES = ( + '@prefix owl: . ' + '@prefix rdfs: . ' + '@prefix xsd: . ' + ) + + def load(self, body, prefix='@prefix : . '): + return Ontology.from_turtle_string(self.PREFIXES + prefix + body) + + def test_a_slash_iri_still_yields_a_local_name(self): + """Plenty of published vocabularies separate on `/` rather than `#`.""" + ontology = self.load( + ':Company a owl:Class .', prefix='@prefix : . ', + ) + [ontology_class] = ontology.index().classes.values() + + assert ontology_class.local_name == 'Company' + assert ontology.index().resolve_class('Company') is ontology_class + + def test_an_iri_with_no_separator_is_its_own_local_name(self): + graph = Graph() + graph.parse( + data=' a .', format='turtle', + ) + assert Ontology.load(graph).index().classes['urn:Company'].local_name == 'urn:Company' + + def test_an_anonymous_class_axiom_is_not_vocabulary(self): + """An `owl:Restriction` body is not a term a user can name in a domain or + range, and not something a model can be asked to emit.""" + ontology = self.load( + ':Company a owl:Class . ' + '[] a owl:Class, owl:ObjectProperty, owl:DatatypeProperty ; rdfs:range xsd:string .' + ) + index = ontology.index() + + assert list(index.classes) == ['urn:x#Company'] + assert index.object_properties == {} + assert index.datatype_properties == {} + + def test_a_second_domain_declaration_is_narrowed_with_a_warning(self, caplog): + """One domain and one range per property is what the rendering and the + domain/range check are defined over.""" + import logging + + with caplog.at_level(logging.WARNING): + ontology = self.load( + ':A a owl:Class . :B a owl:Class . ' + ':rel a owl:ObjectProperty ; rdfs:domain :A, :B ; rdfs:range :A .' + ) + + assert ontology.index().object_properties['urn:x#rel'].domain in ('urn:x#A', 'urn:x#B') + assert 'One domain and one range per property is supported' in caplog.text + + def test_owl_thing_in_a_slot_is_the_same_as_no_slot(self): + ontology = self.load( + ':A a owl:Class . :rel a owl:ObjectProperty ; rdfs:domain owl:Thing ; rdfs:range :A .' + ) + assert ontology.index().object_properties['urn:x#rel'].domain is None + + def test_a_dangling_datatype_property_domain_is_refused(self): + with pytest.raises(OntologyLoadError, match='dangling rdfs:domain'): + self.load(':attr a owl:DatatypeProperty ; rdfs:domain :Absent ; rdfs:range xsd:string .') + + def test_a_class_with_two_parents_is_rendered_once_and_names_the_others(self): + """Rendered under the first parent, with the rest named on its own line, so + nothing is lost and no subtree is duplicated.""" + ontology = self.load( + ':Agent a owl:Class . :Legal a owl:Class . ' + ':Company a owl:Class ; rdfs:subClassOf :Agent, :Legal .' + ) + block = ontology.format_as_prompt_constraint('align') + + assert block.count('Company') == 1 + assert 'also a kind of' in block + assert ontology.index().classes['urn:x#Company'].ancestors == { + 'urn:x#Company', 'urn:x#Agent', 'urn:x#Legal', + } + + def test_an_alias_that_renders_to_nothing_is_dropped(self): + """It would tell the model nothing, and an empty "also known as" is noise.""" + ontology = self.load( + '@prefix skos: . ' + ':Company a owl:Class ; skos:altLabel " ", "Corporation", "Corporation" .' + ) + block = ontology.format_as_prompt_constraint('align') + + assert block.count('Corporation') == 1 + assert 'also known as Corporation)' in block + + def test_an_alias_identical_to_the_primary_name_is_not_offered_twice(self): + """It tells the model nothing, and reads as two names for one thing.""" + ontology = self.load( + '@prefix skos: . ' + ':Company a owl:Class ; skos:altLabel "COMPANY" .' + ) + assert 'also known as' not in ontology.format_as_prompt_constraint('align') + + def test_a_graph_that_is_not_a_graph_is_refused(self): + with pytest.raises(OntologyLoadError): + Ontology.from_graph('@prefix : .') + + def test_turtle_that_does_not_parse_is_refused(self): + with pytest.raises(OntologyLoadError, match='Failed to parse Turtle string'): + Ontology.from_turtle_string(':Company a owl:Class') + + @pytest.mark.parametrize('turtle,expected', [ + ('@prefix : . a owl:Ontology .', 'urn:base#'), + (' a owl:Ontology . a owl:Class .', 'urn:header'), + ]) + def test_the_namespace_prefers_an_in_file_declaration(self, turtle, expected): + """In priority order: the `@prefix :` declaration, then rdflib's default + namespace, then the first `owl:Ontology` subject, then the supplied base - + last, so a declaration in the file always wins.""" + ontology = Ontology.from_turtle_string(self.PREFIXES + turtle, base_iri='urn:given#') + assert ontology.namespace == expected + + def test_the_supplied_base_is_used_when_the_file_declares_nothing(self): + graph = Graph() + graph.parse(data=' a .', format='turtle') + + assert Ontology(graph, base_iri='urn:given#').namespace == 'urn:given#' diff --git a/lexical-graph/tests/unit/indexing/extract/test_ontology_config.py b/lexical-graph/tests/unit/indexing/extract/test_ontology_config.py new file mode 100644 index 00000000..8eab557b --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ontology_config.py @@ -0,0 +1,538 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""`ontology_authority` as a level, its per-dimension overrides, and the wiring. + +The level is the user-facing knob and the six dimensions are what the code acts +on, so the mapping between them is a contract: this file pins the whole table +rather than sampling it, because a dimension that silently fails to turn on is +the one failure in this feature that looks like success. +""" + +from dataclasses import fields +from pathlib import Path + +import pytest + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology import Ontology +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_config import ( + DIMENSIONS, + ONTOLOGY_AUTHORITY_LEVELS, + OntologyConfig, + ResolvedDimensions, + to_ontology_config, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_filter import OntologyFilter +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.prompt_constraint import ( + PROMPT_CONSTRAINT_LEVELS, +) +from graphrag_toolkit.lexical_graph.lexical_graph_index import LexicalGraphIndex + +FIXTURES = Path(__file__).parent.parent.parent.parent / 'fixtures' / 'ontologies' +COMPANY = FIXTURES / 'company.ttl' + +# The level table, read out as data. `off` is the all-False default; `align` owns +# naming; `strict` alone decides membership. +LEVEL_TABLE = { + 'off': (), + 'align': ('normalize_names', 'drop_type_restatements'), + 'strict': DIMENSIONS, +} + +@pytest.fixture(scope='module') +def company(): + return Ontology.load(COMPANY) + +def turtle_declaring(property_name): + return Ontology.from_turtle_string( + '@prefix : . ' + '@prefix owl: . ' + '@prefix rdfs: . ' + '@prefix xsd: . ' + f':{property_name} a owl:DatatypeProperty ; rdfs:range xsd:string .' + ) + +class TestTheLevelResolvesToTheDimensions: + + def test_the_six_dimensions_are_derived_from_the_dataclass(self): + """So the defaults table and the override loop cannot fall out of step.""" + assert DIMENSIONS == tuple(field.name for field in fields(ResolvedDimensions)) + assert len(DIMENSIONS) == 6 + + @pytest.mark.parametrize('level', LEVEL_TABLE) + @pytest.mark.parametrize('dimension', DIMENSIONS) + def test_the_whole_table(self, company, level, dimension): + resolved = OntologyConfig(company, ontology_authority=level).resolved() + assert getattr(resolved, dimension) is (dimension in LEVEL_TABLE[level]) + + def test_the_default_level_is_align(self, company): + assert OntologyConfig(company).ontology_authority == 'align' + + def test_the_levels_are_the_renderers_levels_and_not_a_second_list(self): + """A level this module accepted but the renderer did not would raise from + inside pipeline setup, after configuration had already succeeded.""" + assert ONTOLOGY_AUTHORITY_LEVELS is PROMPT_CONSTRAINT_LEVELS + + def test_resolved_is_recomputed_rather_than_cached(self, company): + config = OntologyConfig(company, ontology_authority='strict') + config.ontology_authority = 'off' + assert config.resolved() == ResolvedDimensions() + +class TestOverrides: + """The level chooses defaults, it does not constrain them.""" + + @pytest.mark.parametrize('dimension', DIMENSIONS) + def test_any_dimension_can_be_turned_on_at_off(self, company, dimension): + config = OntologyConfig(company, ontology_authority='off', **{dimension: True}) + assert getattr(config.resolved(), dimension) is True + + @pytest.mark.parametrize('dimension', DIMENSIONS) + def test_any_dimension_can_be_turned_off_at_strict(self, company, dimension): + config = OntologyConfig(company, ontology_authority='strict', **{dimension: False}) + assert getattr(config.resolved(), dimension) is False + + @pytest.mark.parametrize('dimension', DIMENSIONS) + def test_none_means_follow_the_level(self, company, dimension): + config = OntologyConfig(company, ontology_authority='align', **{dimension: None}) + assert getattr(config.resolved(), dimension) is (dimension in LEVEL_TABLE['align']) + + def test_say_nothing_in_the_prompt_but_still_coerce_datatypes(self, company): + """The combination the override mechanism exists for, and not a contradiction.""" + resolved = OntologyConfig(company, ontology_authority='off', enforce_datatypes=True).resolved() + assert (resolved.enforce_datatypes, resolved.normalize_names) == (True, False) + +class TestValidation: + """Settings are checked before the ontology is parsed, except where they cannot be.""" + + @pytest.mark.parametrize('level', ['stricter', 'STRICT', 'guide', '', None]) + def test_an_unknown_authority_level_raises_and_lists_the_levels(self, company, level): + with pytest.raises(ValueError, match='Unknown ontology authority level'): + OntologyConfig(company, ontology_authority=level) + + def test_an_unknown_typed_properties_placement_raises(self, company): + with pytest.raises(ValueError, match='Unknown typed_properties placement'): + OntologyConfig(company, typed_properties='subjects') + + def test_an_unknown_vocabulary_format_raises(self, company): + with pytest.raises(ValueError, match='Unknown ontology vocabulary_format'): + OntologyConfig(company, vocabulary_format='ttl') + + def test_a_setting_is_validated_before_the_ontology_is_loaded(self): + """A typo should not first cost a parse of a file that does not exist.""" + with pytest.raises(ValueError, match='Unknown ontology authority level'): + OntologyConfig(FIXTURES / 'absent.ttl', ontology_authority='nonsense') + + @pytest.mark.parametrize('name', ['value', 'search_str', 'class']) + def test_a_property_that_would_overwrite_the_entity_is_refused(self, name): + """Under subject placement this key is the entity's own.""" + with pytest.raises(ValueError, match='the graph model already owns'): + OntologyConfig(turtle_declaring(name), typed_properties='subject') + + @pytest.mark.parametrize('name', ['typed_value', 'datatype']) + def test_the_complement_property_names_are_reserved_only_under_both(self, name): + """The check is scoped to the placement actually requested. + + Only subject placement keys a property from the ontology's vocabulary, so + `'complement'` alone has nothing to collide with - it widens the reserved + set only when subject placement is active too. + """ + OntologyConfig(turtle_declaring(name), typed_properties='complement') + + with pytest.raises(ValueError, match='the graph model already owns'): + OntologyConfig(turtle_declaring(name), typed_properties='both') + + def test_a_collision_is_not_raised_for_a_write_that_will_never_happen(self): + OntologyConfig(turtle_declaring('value'), typed_properties='off') + OntologyConfig(turtle_declaring('value'), typed_properties='complement') + +class TestNormalizingWhatTheUserConfigured: + + def test_the_ontology_is_loaded_whatever_form_it_arrived_in(self): + assert isinstance(OntologyConfig(COMPANY).ontology, Ontology) + assert isinstance(OntologyConfig(str(COMPANY)).ontology, Ontology) + + def test_an_existing_config_is_never_rebuilt_at_defaults(self, company): + """Settings a user made must survive normalization.""" + config = OntologyConfig(company, ontology_authority='strict', report_violations=True) + assert to_ontology_config(config) is config + + def test_a_bare_source_gets_the_default_level(self): + assert to_ontology_config(COMPANY).ontology_authority == 'align' + +class TestWhatTheBuildersAreTold: + + @pytest.mark.parametrize('placement,subject,complement', [ + ('off', False, False), + ('subject', True, False), + ('complement', False, True), + ('both', True, True), + ]) + def test_each_placement_writes_to_its_own_node(self, company, placement, subject, complement): + """`'both'` runs each write to its own node; it is not a fallback chain.""" + config = OntologyConfig(company, typed_properties=placement) + assert config.writes_subject_properties() is subject + assert config.writes_complement_properties() is complement + + def test_typed_properties_defaults_to_off(self, company): + assert OntologyConfig(company).typed_properties == 'off' + +class TestWhetherTheFilterIsBuiltAtAll: + """the rule that a filter with nothing to do is left out, and the reason `off` is not a filter with every flag off.""" + + @pytest.mark.parametrize('kwargs,required', [ + ({'ontology_authority': 'off'}, False), + ({'ontology_authority': 'align'}, True), + ({'ontology_authority': 'strict'}, True), + ({'ontology_authority': 'off', 'enforce_datatypes': True}, True), + ({'ontology_authority': 'off', 'typed_properties': 'subject'}, True), + ({'ontology_authority': 'off', 'typed_properties': 'complement'}, True), + ]) + def test_filter_required(self, company, kwargs, required): + assert OntologyConfig(company, **kwargs).filter_required() is required + + def test_no_ontology_means_no_filter(self): + assert LexicalGraphIndex._ontology_filter(None) is None + + def test_off_puts_no_component_in_the_pipeline(self, company): + """Not a no-op component: a filter would still round-trip TOPICS_KEY + through validate/dump and would still annotate, which `off` rules out.""" + assert LexicalGraphIndex._ontology_filter( + OntologyConfig(company, ontology_authority='off') + ) is None + + @pytest.mark.parametrize('level', ['align', 'strict']) + def test_every_resolved_dimension_reaches_the_filter(self, company, level): + """The `asdict` spread, asserted dimension by dimension. + + A hand-written argument list can omit one, and the symptom is a gate the + user asked for silently not running. + """ + config = OntologyConfig(company, ontology_authority=level) + filter_ = LexicalGraphIndex._ontology_filter(config) + + for dimension in DIMENSIONS: + assert getattr(filter_, dimension) is getattr(config.resolved(), dimension) + + def test_the_filter_accepts_every_dimension_by_name(self): + """Guards the spread from the other end: a new dimension with no matching + flag on the filter would raise at pipeline construction.""" + for dimension in DIMENSIONS: + assert dimension in OntologyFilter.model_fields + + def test_typed_properties_alone_builds_an_annotate_only_filter(self, company): + """`canonicalName` and `datatype` are written by the filter and nothing + else, so dropping it here would silently write no typed properties.""" + filter_ = LexicalGraphIndex._ontology_filter( + OntologyConfig(company, ontology_authority='off', typed_properties='subject') + ) + + assert filter_ is not None + assert not any(getattr(filter_, dimension) for dimension in DIMENSIONS) + + def test_report_violations_reaches_the_filter(self, company): + filter_ = LexicalGraphIndex._ontology_filter( + OntologyConfig(company, report_violations=True) + ) + assert filter_.report_violations is True + + def test_the_filter_carries_the_index_and_not_the_graph(self, company): + """What crosses the spawn boundary is plain data.""" + filter_ = LexicalGraphIndex._ontology_filter(OntologyConfig(company)) + assert filter_.index is company.index() + +class TestWhatEachPromptStageIsGiven: + """The two blocks are different blocks, and swapping them would be silent.""" + + def test_no_ontology_renders_two_empty_blocks(self): + assert LexicalGraphIndex._render_ontology_constraints(None) == ('', '') + + def test_the_topics_block_gets_the_vocabulary_and_the_propositions_block_the_classes(self, company): + constraints = LexicalGraphIndex._render_ontology_constraints(OntologyConfig(company)) + + assert 'WORKS_FOR' in constraints.topics + assert 'WORKS_FOR' not in constraints.propositions + assert 'Sports Team' in constraints.propositions + + def test_off_renders_neither_block(self, company): + constraints = LexicalGraphIndex._render_ontology_constraints( + OntologyConfig(company, ontology_authority='off') + ) + assert constraints == ('', '') + + def test_the_vocabulary_format_reaches_the_topics_block_only(self, company): + """That stage classifies the entities it names and extracts nothing else, + so there is no property vocabulary there for a serialization to present + differently.""" + prose = LexicalGraphIndex._render_ontology_constraints(OntologyConfig(company)) + turtle = LexicalGraphIndex._render_ontology_constraints( + OntologyConfig(company, vocabulary_format='turtle') + ) + + assert '```turtle' in turtle.topics + assert '```turtle' not in prose.topics + assert turtle.propositions == prose.propositions + +class TestWhatTheBuildPipelineIsTold: + + def index_for(self, ontology): + from graphrag_toolkit.lexical_graph.lexical_graph_index import ( + ExtractionConfig, + IndexingConfig, + LexicalGraphIndex, + ) + + index = LexicalGraphIndex.__new__(LexicalGraphIndex) + index.indexing_config = IndexingConfig(extraction=ExtractionConfig(ontology=ontology)) + return index + + def test_no_ontology_defers_rather_than_pinning_off(self): + """Returned as None so the `coalesce` chain behaves as it does for every + other setting: an unasked-for value defers to the layer below, and no + environment variable can reach a placement that writes.""" + assert self.index_for(None)._typed_properties() is None + + @pytest.mark.parametrize('placement', ['off', 'subject', 'complement', 'both']) + def test_a_configured_placement_is_passed_through(self, company, placement): + ontology_config = OntologyConfig(company, typed_properties=placement) + assert self.index_for(ontology_config)._typed_properties() == placement + +class TestSeedingThePreferredClassifications: + """Seeding, and the user's own list.""" + + def test_the_slot_is_seeded_from_the_ontology_when_the_user_said_nothing(self, company): + from graphrag_toolkit.lexical_graph.lexical_graph_index import ExtractionConfig + + config = ExtractionConfig(ontology=OntologyConfig(company)) + + assert LexicalGraphIndex._preferred_entity_classifications(config) == company.class_names() + + def test_a_user_list_is_kept_with_a_warning(self, company, caplog): + """Honouring the ontology instead would discard a setting the user made + deliberately; merging the two would produce a vocabulary neither asked for.""" + import logging + + from graphrag_toolkit.lexical_graph.lexical_graph_index import ExtractionConfig + + config = ExtractionConfig( + ontology=OntologyConfig(company), + preferred_entity_classifications=['Widget'], + ) + + with caplog.at_level(logging.WARNING): + assert LexicalGraphIndex._preferred_entity_classifications(config) == ['Widget'] + + assert 'Honouring preferred_entity_classifications' in caplog.text + + def test_without_an_ontology_the_users_value_is_returned_unchanged(self): + from graphrag_toolkit.lexical_graph.lexical_graph_index import ExtractionConfig + + config = ExtractionConfig(preferred_entity_classifications=['Widget']) + + assert LexicalGraphIndex._preferred_entity_classifications(config) == ['Widget'] + +# The pipeline is built with the store factories patched, so the store's *type* is +# the only thing the pipeline learns from it - which is what +# `_configure_extraction_pipeline` branches on. +MODULE = 'graphrag_toolkit.lexical_graph.lexical_graph_index' + +def build_pipeline(extraction, batch=False, graph_store=None, pre_processors=False): + from unittest.mock import Mock, patch + + from llama_index.core.llms.mock import MockLLM + + from graphrag_toolkit.lexical_graph.indexing.extract import BatchConfig + from graphrag_toolkit.lexical_graph.lexical_graph_index import IndexingConfig + + # The extractors fall back to `GraphRAGConfig.extraction_llm` when the config + # names none, and that constructs a real `BedrockConverse` - which needs a + # region, and so fails wherever AWS is not configured. Nothing here depends on + # which model it is. + if extraction.extraction_llm is None: + extraction.extraction_llm = MockLLM() + + config = IndexingConfig( + extraction=extraction, + batch_config=BatchConfig( + role_arn='arn:aws:iam::123456789012:role/test-batch-role', + region='us-east-1', + bucket_name='test-batch-bucket', + ) if batch else None, + ) + store = Mock() if graph_store is None else graph_store + + with ( + patch(f'{MODULE}.GraphStoreFactory.for_graph_store', return_value=store), + patch(f'{MODULE}.MultiTenantGraphStore.wrap', return_value=store), + patch(f'{MODULE}.VectorStoreFactory.for_vector_store', return_value=store), + patch(f'{MODULE}.MultiTenantVectorStore.wrap', return_value=store), + ): + index = LexicalGraphIndex( + graph_store='dummy://', vector_store='dummy://', indexing_config=config, + ) + + return index.extraction_pre_processors if pre_processors else index.extraction_components + +def only(components, component_type): + matches = [c for c in components if isinstance(c, component_type)] + assert len(matches) == 1, f'expected one {component_type.__name__}, found {len(matches)}' + return matches[0] + +class TestTheBlocksReachTheExtractors: + """Rendering the right block is not enough; each stage has to be handed its own. + + Swapping them would leave both prompts rendering and both stages extracting, + which is why this is asserted per extractor rather than inferred from + `_render_ontology_constraints`. + """ + + @pytest.mark.parametrize('batch', [False, True], ids=['non_batch', 'batch']) + def test_each_extractor_receives_its_own_block(self, company, batch): + from graphrag_toolkit.lexical_graph import ExtractionConfig + from graphrag_toolkit.lexical_graph.indexing.extract import ( + BatchLLMPropositionExtractorSync, + BatchTopicExtractorSync, + LLMPropositionExtractor, + TopicExtractor, + ) + + components = build_pipeline(ExtractionConfig(ontology=COMPANY), batch=batch) + propositions = BatchLLMPropositionExtractorSync if batch else LLMPropositionExtractor + topics = BatchTopicExtractorSync if batch else TopicExtractor + + assert only(components, propositions).ontology_constraints == \ + company.format_as_proposition_constraint('align') + assert only(components, topics).ontology_constraints == \ + company.format_as_prompt_constraint('align') + + @pytest.mark.parametrize('extraction_kwargs', [{}, {'ontology_authority': 'off'}]) + def test_no_ontology_and_off_both_leave_every_extractor_empty(self, company, extraction_kwargs): + """The empty string, not None: an empty block composes to no change.""" + from llama_index.core.node_parser import SentenceSplitter + + from graphrag_toolkit.lexical_graph import ExtractionConfig + + ontology = OntologyConfig(company, **extraction_kwargs) if extraction_kwargs else None + components = build_pipeline(ExtractionConfig(ontology=ontology)) + + for component in components: + if not isinstance(component, SentenceSplitter): + assert component.ontology_constraints == '' + + def test_the_filter_is_the_last_component_when_it_is_present(self, company): + """It reads what the topic extractor wrote, so ordering is not cosmetic.""" + from graphrag_toolkit.lexical_graph import ExtractionConfig + + components = build_pipeline(ExtractionConfig(ontology=COMPANY)) + + assert isinstance(components[-1], OntologyFilter) + + def test_no_filter_component_exists_at_off(self, company): + from graphrag_toolkit.lexical_graph import ExtractionConfig + + components = build_pipeline(ExtractionConfig( + ontology=OntologyConfig(company, ontology_authority='off'), + )) + + assert not [c for c in components if isinstance(c, OntologyFilter)] + +class TestOntologyAndClassificationInference: + """One combination has no coherent reading and is rejected at config time.""" + + def test_inference_alongside_an_ontology_is_fine(self): + from graphrag_toolkit.lexical_graph import ExtractionConfig + + assert ExtractionConfig( + ontology=COMPANY, infer_entity_classifications=True, + ).infer_entity_classifications is True + + def test_replacing_the_seeded_classifications_is_rejected_and_names_the_way_out(self): + """Inference that replaces the defaults would discard the very class names + the ontology seeded, leaving the vocabulary block and the preference slot + disagreeing.""" + from graphrag_toolkit.lexical_graph import ExtractionConfig + from graphrag_toolkit.lexical_graph.indexing.extract import InferClassificationsConfig + + with pytest.raises(ValueError) as error: + ExtractionConfig( + ontology=COMPANY, + infer_entity_classifications=InferClassificationsConfig( + replace_default_classifications=True, + ), + ) + + assert 'replace_default_classifications=False' in str(error.value) + + def test_replacing_without_an_ontology_is_still_fine(self): + """The rejection is about the combination, not about the setting.""" + from graphrag_toolkit.lexical_graph import ExtractionConfig + from graphrag_toolkit.lexical_graph.indexing.extract import InferClassificationsConfig + + config = ExtractionConfig(infer_entity_classifications=InferClassificationsConfig( + replace_default_classifications=True, + )) + + assert config.infer_entity_classifications.replace_default_classifications is True + +class TestSeedingReachesTheProvider: + """preference seeding delivered, not just decided. + + `_preferred_entity_classifications` picks the list; these assert the pipeline + then hands that list to whichever provider the configuration asks for. + """ + + def classifications_of(self, components): + from llama_index.core.schema import TextNode + + from graphrag_toolkit.lexical_graph.indexing.extract import TopicExtractor + + extractor = only(components, TopicExtractor) + return extractor.entity_classification_provider(TextNode(text='x', id_='chunk-1')) + + def test_the_ontologys_class_names_are_what_the_extractor_offers(self, company): + from graphrag_toolkit.lexical_graph import ExtractionConfig + + components = build_pipeline(ExtractionConfig(ontology=COMPANY)) + + assert self.classifications_of(components) == company.class_names() + + def test_a_users_own_provider_is_passed_through_untouched(self): + """A callable is not a list, so it cannot be merged with anything - and a + user who wrote one has said how the slot is filled.""" + from graphrag_toolkit.lexical_graph import ExtractionConfig + from graphrag_toolkit.lexical_graph.indexing.extract import PreferredValuesProvider + + class FixedProvider(PreferredValuesProvider): + def __call__(self, node): + return ['From Provider'] + + components = build_pipeline(ExtractionConfig( + ontology=COMPANY, preferred_entity_classifications=FixedProvider(), + )) + + assert self.classifications_of(components) == ['From Provider'] + + def test_inference_starts_from_the_ontologys_class_names(self, company): + """The seeded list becomes the inferencer's defaults rather than being + discarded, which is why replacing them is refused at config time.""" + from graphrag_toolkit.lexical_graph import ExtractionConfig + from graphrag_toolkit.lexical_graph.indexing.extract import InferClassifications + + pre_processors = build_pipeline(ExtractionConfig( + ontology=COMPANY, infer_entity_classifications=True, + ), pre_processors=True) + inferencer = only(pre_processors, InferClassifications) + + assert inferencer.default_classifications == company.class_names() + + def test_a_dummy_store_still_gets_the_prompt_block(self): + """Providers are forced empty on that path; the vocabulary block is not.""" + from graphrag_toolkit.lexical_graph import ExtractionConfig + from graphrag_toolkit.lexical_graph.indexing.extract import TopicExtractor + from graphrag_toolkit.lexical_graph.storage.graph import DummyGraphStore + + components = build_pipeline( + ExtractionConfig(ontology=COMPANY), graph_store=DummyGraphStore(), + ) + + assert only(components, TopicExtractor).ontology_constraints + assert self.classifications_of(components) == [] diff --git a/lexical-graph/tests/unit/indexing/extract/test_ontology_datatypes.py b/lexical-graph/tests/unit/indexing/extract/test_ontology_datatypes.py new file mode 100644 index 00000000..285d48b3 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ontology_datatypes.py @@ -0,0 +1,214 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Coercing an extracted literal to its declared XSD type. + +Two properties matter and everything here serves one of them: a value that +parses must be **JSON-serializable**, because it lands in node metadata and then +in Cypher parameters; and a value that would need *interpreting* rather than +parsing must be refused, because a wrong number stored under a typed key is +trusted by everything downstream. +""" + +import json + +import pytest + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.datatype_utils import ( + coerce_literal, + validate_literal_against_xsd, + validates_datatype, +) + +XSD = 'http://www.w3.org/2001/XMLSchema#' + +def xsd(name): + return f'{XSD}{name}' + +class TestAcceptedForms: + """Lexical latitude: the shapes a model actually emits for a right answer.""" + + @pytest.mark.parametrize('literal,expected', [ + ('1994', 1994), + (' 1994 ', 1994), + ('+1994', 1994), + ('-1994', -1994), + ('1994.0', 1994), + ('2,400,000', 2400000), + ]) + def test_integers(self, literal, expected): + assert coerce_literal(literal, xsd('integer')) == expected + + @pytest.mark.parametrize('literal,expected', [ + ('41500000', 41500000.0), + ('41,500,000', 41500000.0), + ('4.15e7', 41500000.0), + ('0', 0.0), + ('-2.5', -2.5), + ]) + def test_decimals(self, literal, expected): + assert coerce_literal(literal, xsd('double')) == expected + + @pytest.mark.parametrize('literal,expected', [ + ('true', True), ('TRUE', True), ('1', True), + ('false', False), ('False', False), ('0', False), + ]) + def test_booleans(self, literal, expected): + assert coerce_literal(literal, xsd('boolean')) is expected + + @pytest.mark.parametrize('literal', [ + '2020-03-03', 'March 3 2020', 'March 3rd 2020', '3 March 2020', + 'Mar 3, 2020', '2020-03-03Z', '2020-03-03+01:00', + ]) + def test_dates_reduce_to_one_iso_string(self, literal): + assert coerce_literal(literal, xsd('date')) == '2020-03-03' + + @pytest.mark.parametrize('name,literal,expected', [ + ('dateTime', '2020-03-03T09:30:00', '2020-03-03T09:30:00'), + ('time', '09:30', '09:30:00'), + ('time', '09:30:15', '09:30:15'), + ('anyURI', 'https://example.com/a', 'https://example.com/a'), + ('string', ' spaced ', 'spaced'), + ]) + def test_the_remaining_families(self, name, literal, expected): + assert coerce_literal(literal, xsd(name)) == expected + +class TestRefusedForms: + """Where parsing stops and interpreting would begin.""" + + @pytest.mark.parametrize('literal', [ + 'nineteen ninety four', + '1994.5', + '1,99,4', + '1994,5', + '', + ' ', + ]) + def test_a_value_that_is_not_an_integer_in_its_entirety(self, literal): + assert coerce_literal(literal, xsd('integer')) is None + + @pytest.mark.parametrize('literal', [ + '41,500,000 dollars', + '$41.5m', + 'nine hundred million', + 'nan', + 'inf', + ]) + def test_a_unit_or_a_magnitude_is_never_stripped(self, literal): + """Nothing distinguishes a harmless unit from a magnitude without reading it. + + `nan` and `inf` are refused for a different reason: `float()` accepts + both, and `json.dumps` then emits tokens no JSON parser must accept. + """ + assert coerce_literal(literal, xsd('double')) is None + + @pytest.mark.parametrize('literal', ['yes', 'no', 'Y', 'TRUE.']) + def test_a_boolean_synonym_nobody_was_offered(self, literal): + assert coerce_literal(literal, xsd('boolean')) is None + + @pytest.mark.parametrize('literal', [ + '01/03/1994', + '2015-02-29', + '15-2-9', + 'sometime in March', + ]) + def test_an_ambiguous_or_impossible_date(self, literal): + """`01/03/1994` is two different days depending on where the text came from.""" + assert coerce_literal(literal, xsd('date')) is None + + def test_a_sentence_where_a_uri_was_asked_for(self): + assert coerce_literal('see the company website', xsd('anyURI')) is None + + @pytest.mark.parametrize('literal,datatype', [(None, xsd('integer')), ('1994', None), (None, None)]) + def test_a_missing_literal_or_datatype(self, literal, datatype): + assert coerce_literal(literal, datatype) is None + + def test_a_non_xsd_range_is_refused_rather_than_guessed_at(self): + assert coerce_literal('anything', 'http://example.com/company#Money') is None + +class TestBounds: + """The reason the integer types are enumerated rather than treated alike.""" + + @pytest.mark.parametrize('name,literal,accepted', [ + ('nonNegativeInteger', '0', True), + ('nonNegativeInteger', '-5', False), + ('positiveInteger', '1', True), + ('positiveInteger', '0', False), + ('negativeInteger', '-1', True), + ('negativeInteger', '0', False), + ('unsignedByte', '255', True), + ('unsignedByte', '256', False), + ('byte', '-128', True), + ('byte', '-129', False), + ('short', '32767', True), + ('short', '32768', False), + ]) + def test_a_value_outside_its_declared_type_is_refused(self, name, literal, accepted): + assert (coerce_literal(literal, xsd(name)) is not None) is accepted + + def test_gyear_is_an_integer_because_that_is_what_the_prompt_promised(self): + assert coerce_literal('1994', xsd('gYear')) == 1994 + +class TestWhatTheCallersDependOn: + + @pytest.mark.parametrize('name,literal', [ + ('integer', '1994'), ('double', '2.5'), ('boolean', 'true'), + ('date', '2020-03-03'), ('dateTime', '2020-03-03T09:30:00'), + ('time', '09:30'), ('string', 'text'), ('anyURI', 'https://example.com'), + ]) + def test_every_successful_coercion_is_json_serializable(self, name, literal): + value = coerce_literal(literal, xsd(name)) + assert json.loads(json.dumps(value)) == value + + @pytest.mark.parametrize('name,literal', [('boolean', 'false'), ('integer', '0'), ('double', '0')]) + def test_a_falsy_result_is_still_a_successful_coercion(self, name, literal): + """Callers must test `is not None`. `False`, `0` and `0.0` all coerce.""" + assert coerce_literal(literal, xsd(name)) is not None + assert not coerce_literal(literal, xsd(name)) + + @pytest.mark.parametrize('name,expected', [ + ('integer', True), ('double', True), ('boolean', True), ('date', True), + ('string', True), ('token', True), + ('hexBinary', False), ('duration', False), ('gMonthDay', False), + ]) + def test_validates_datatype_reports_whether_a_check_actually_happens(self, name, expected): + assert validates_datatype(xsd(name)) is expected + + @pytest.mark.parametrize('datatype', [None, 'http://example.com/company#Money']) + def test_validates_datatype_is_false_for_a_non_xsd_range(self, datatype): + assert validates_datatype(datatype) is False + + def test_an_unimplemented_xsd_type_keeps_the_text_and_admits_it(self): + """The pair that lets `enforce_datatypes` warn instead of claiming a check. + + Coercion returns the trimmed text so nothing is lost, and + `validates_datatype` returns False so the caller can say out loud that + the declared type was not honoured. + """ + assert coerce_literal(' deadbeef ', xsd('hexBinary')) == 'deadbeef' + assert validates_datatype(xsd('hexBinary')) is False + + def test_validity_is_coercibility(self): + """One implementation, so the two can never disagree about a value.""" + for (literal, datatype) in [('1994', xsd('integer')), ('x', xsd('integer')), ('false', xsd('boolean'))]: + assert validate_literal_against_xsd(literal, datatype) is ( + coerce_literal(literal, datatype) is not None + ) + +class TestTheTimeFamilyEdges: + """`dateTime` and `time` accept only the ISO lexical form, and reject in two + different ways: the shape fails, or the calendar does.""" + + @pytest.mark.parametrize('literal', ['yesterday', '2020-13-01T00:00:00', 'March 3 2020']) + def test_a_datetime_that_is_not_iso(self, literal): + """Unlike `xsd:date`, a named month is not accepted here - no model in the + corpus produced one, and `fromisoformat` is the whole of the contract.""" + assert coerce_literal(literal, xsd('dateTime')) is None + + def test_a_space_separated_datetime_is_accepted(self): + """`datetime.fromisoformat` takes it, so this module does too.""" + assert coerce_literal('2020-03-03 09:30', xsd('dateTime')) == '2020-03-03T09:30:00' + + @pytest.mark.parametrize('literal', ['9:30', 'half past nine', '25:00', '09:99']) + def test_a_time_that_is_the_wrong_shape_or_not_on_the_clock(self, literal): + assert coerce_literal(literal, xsd('time')) is None diff --git a/lexical-graph/tests/unit/indexing/extract/test_ontology_filter.py b/lexical-graph/tests/unit/indexing/extract/test_ontology_filter.py new file mode 100644 index 00000000..661a8bac --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ontology_filter.py @@ -0,0 +1,608 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The filter: the one component in this feature that changes extracted data. + +Every test here runs the component the way the pipeline does - through +`__call__` on a node carrying `TOPICS_KEY` - so the `model_validate` / +`model_dump` round trip is exercised rather than bypassed. Facts are built by +hand: what is being asserted is what the code does with a given fact, which is +not a question about any model's output. + +The gates are independent by design, so each one is tested with the others off. +That is the whole point of the per-dimension escape hatches: `enforce_domain_range` +alone must not start rejecting unresolvable classifications. +""" + +import logging +import pickle +from pathlib import Path + +import pytest +from llama_index.core.schema import TextNode + +from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology import Ontology +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_filter import ( + FilterCounters, + OntologyFilter, + _warned_unvalidated_datatypes, + authored_name, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_index import ( + DatatypeProperty, + ObjectProperty, + OntologyClass, + OntologyIndex, +) +from graphrag_toolkit.lexical_graph.indexing.model import ( + Entity, + Fact, + Relation, + Statement, + Topic, + TopicCollection, +) + +FIXTURES = Path(__file__).parent.parent.parent.parent / 'fixtures' / 'ontologies' +NS = 'http://example.com/company#' +XSD = 'http://www.w3.org/2001/XMLSchema#' +LOCAL = '__Local_Entity__' + +@pytest.fixture(scope='module') +def index(): + return Ontology.load(FIXTURES / 'company.ttl').index() + +def attribute(predicate, value='1994', subject_class='Company', subject='Meridian Freight'): + """A subject-predicate-complement fact, as the parser leaves an attribute.""" + return Fact( + subject=Entity(value=subject, classification=subject_class), + predicate=Relation(value=predicate), + complement=Entity(value=value, classification=LOCAL), + ) + +def relation(predicate, subject_class='Person', object_class='Company'): + return Fact( + subject=Entity(value='Priya Raman', classification=subject_class), + predicate=Relation(value=predicate), + object=Entity(value='Meridian Freight', classification=object_class), + ) + +def run(filter_, *facts, entities=None): + """Put facts through the component the way the pipeline does.""" + topics = TopicCollection(topics=[Topic( + value='t', + entities=list(entities or []), + statements=[Statement(value='s', facts=list(facts))], + )]) + node = TextNode(text='x', metadata={TOPICS_KEY: topics.model_dump()}) + + returned = filter_([node]) + + assert returned[0] is node + return TopicCollection.model_validate(node.metadata[TOPICS_KEY]).topics[0] + +def facts_of(topic): + return [fact for statement in topic.statements for fact in statement.facts] + +def kept(index, fact, **flags): + return len(facts_of(run(OntologyFilter(index=index, **flags), fact))) == 1 + +def counted(index, *facts, **flags): + """The counters for one topic, which `__call__` accumulates and then discards.""" + counters = FilterCounters() + topic = Topic(value='t', statements=[Statement(value='s', facts=list(facts))]) + OntologyFilter(index=index, **flags)._filter_topic(topic, counters) + return counters + +class TestNormalization: + """A resolved name is stored as the ontology's author wrote it.""" + + def test_a_resolved_predicate_takes_the_authored_spelling(self, index): + [fact] = facts_of(run(OntologyFilter(index=index, normalize_names=True), relation('WORKS FOR'))) + assert fact.predicate.value == 'worksFor' + + def test_a_resolved_classification_takes_the_authored_spelling(self, index): + [fact] = facts_of(run(OntologyFilter(index=index, normalize_names=True), relation('WORKS FOR', subject_class='ATHLETE'))) + assert fact.subject.classification == 'Athlete' + + def test_a_declared_label_wins_over_the_local_name(self, index): + [fact] = facts_of(run( + OntologyFilter(index=index, normalize_names=True), + relation('WORKS FOR', object_class='SPORTS TEAM'), + )) + assert fact.object.classification == 'Sports Team' + + def test_the_authored_name_is_verbatim_and_gets_no_house_convention(self, index): + """The stored spelling is deliberately not the prompt's rendering. + + `:worksFor` renders as `WORKS_FOR` for the model and stores as `worksFor`. + Collapsing the two into one helper is the bug `naming.py` exists to + prevent. + """ + assert authored_name(index.object_properties[f'{NS}worksFor']) == 'worksFor' + assert authored_name(index.classes[f'{NS}SportsTeam']) == 'Sports Team' + assert authored_name(index.classes[f'{NS}Athlete']) == 'Athlete' + + def test_an_unresolved_name_is_left_exactly_as_the_parser_produced_it(self, index): + [fact] = facts_of(run( + OntologyFilter(index=index, normalize_names=True), + relation('HIRED BY', subject_class='BOARD MEMBER'), + )) + assert (fact.predicate.value, fact.subject.classification) == ('HIRED BY', 'BOARD MEMBER') + + def test_nothing_is_rewritten_with_normalize_names_off(self, index): + [fact] = facts_of(run(OntologyFilter(index=index), relation('WORKS FOR', subject_class='PERSON'))) + assert (fact.predicate.value, fact.subject.classification) == ('WORKS FOR', 'PERSON') + + def test_the_topic_entity_list_is_normalized_and_annotated_too(self, index): + """`topic.entities` is a separate view of the same names after the round + trip, so relying on object sharing to propagate a rewrite would pass in a + unit test and fail in the pipeline.""" + topic = run( + OntologyFilter(index=index, normalize_names=True), + relation('WORKS FOR', subject_class='PERSON'), + entities=[Entity(value='Priya Raman', classification='PERSON')], + ) + + assert topic.entities[0].classification == 'Person' + assert topic.entities[0].classIri == f'{NS}Person' + + def test_an_entity_is_never_pruned_from_the_topic(self, index): + """That list is not written to the graph, and pruning it + would change entity extraction rather than fact conformance.""" + topic = run( + OntologyFilter(index=index, normalize_names=True, enforce_entity_types=True), + entities=[Entity(value='A Trail', classification='TRAIL')], + ) + assert [entity.classification for entity in topic.entities] == ['TRAIL'] + + def test_normalization_is_idempotent(self, index): + """Resolution runs before normalization, which is only sound if resolving + an already-authored name returns the same term.""" + filter_ = OntologyFilter(index=index, normalize_names=True) + once = facts_of(run(filter_, relation('WORKS FOR', subject_class='PERSON'))) + twice = facts_of(run(filter_, once[0])) + + assert twice[0].model_dump() == once[0].model_dump() + +class TestEnforceEntityTypes: + + @pytest.mark.parametrize('subject_class,object_class,keep', [ + ('Person', 'Company', True), + ('Vehicle', 'Company', False), + ('Person', 'Vehicle', False), + ]) + def test_a_classification_that_does_not_resolve_is_dropped(self, index, subject_class, object_class, keep): + fact = relation('WORKS FOR', subject_class=subject_class, object_class=object_class) + assert kept(index, fact, enforce_entity_types=True) is keep + + def test_an_attribute_has_no_object_to_check(self, index): + assert kept(index, attribute('FOUNDED YEAR'), enforce_entity_types=True) + +class TestEnforceRelationshipTypes: + + @pytest.mark.parametrize('fact,keep', [ + ('WORKS FOR', True), + ('HIRED BY', False), + ]) + def test_a_predicate_that_does_not_resolve_is_dropped(self, index, fact, keep): + assert kept(index, relation(fact), enforce_relationship_types=True) is keep + + @pytest.mark.parametrize('predicate,keep', [('FOUNDED YEAR', True), ('CONDITION', False)]) + def test_the_attribute_shape_too(self, index, predicate, keep): + assert kept(index, attribute(predicate), enforce_relationship_types=True) is keep + + def test_a_relation_whose_employer_was_not_a_named_entity_still_resolves(self, index): + """The parser emits a complement whenever it could not match the object + text to an entity it had already seen, so a genuine object property + arrives attribute-shaped and must not be read as undeclared.""" + assert kept(index, attribute('WORKS FOR', value='Meridian Freight'), enforce_relationship_types=True) + +class TestEnforceDomainAndRange: + + @pytest.mark.parametrize('subject_class,object_class,keep', [ + ('Person', 'Company', True), + ('Athlete', 'Sports Team', True), + ('Company', 'Company', False), + ('Person', 'Person', False), + ]) + def test_a_declared_domain_and_range_bound_both_ends(self, index, subject_class, object_class, keep): + fact = relation('WORKS FOR', subject_class=subject_class, object_class=object_class) + assert kept(index, fact, enforce_domain_range=True) is keep + + def test_subclass_closure_is_honoured_in_both_directions(self, index): + """`:playsFor` is narrower than `:worksFor`, so a Person and a Company do + not satisfy it even though an Athlete and a Sports Team do.""" + assert kept(index, relation('PLAYS FOR', 'Athlete', 'Sports Team'), enforce_domain_range=True) + assert not kept(index, relation('PLAYS FOR', 'Person', 'Company'), enforce_domain_range=True) + + def test_an_undeclared_domain_or_range_constrains_nothing(self, index): + assert kept(index, relation('ACQUIRED', 'Person', 'Person'), enforce_domain_range=True) + + def test_a_datatype_property_has_only_a_domain_to_check(self, index): + assert kept(index, attribute('FOUNDED YEAR', subject_class='Sports Team'), enforce_domain_range=True) + assert not kept(index, attribute('FOUNDED YEAR', subject_class='Person'), enforce_domain_range=True) + assert kept(index, attribute('OFFICIAL NAME', value='x', subject_class='Person'), enforce_domain_range=True) + + def test_an_unresolved_class_is_unknown_rather_than_violating(self, index): + """A classification that resolves to nothing cannot be *shown* to breach a + declared domain, so this gate passes it and only `enforce_entity_types` + rejects it.""" + assert kept(index, relation('WORKS FOR', subject_class='Vehicle'), enforce_domain_range=True) + + def test_an_unresolved_predicate_has_no_domain_to_violate(self, index): + assert kept(index, relation('HIRED BY', subject_class='Company'), enforce_domain_range=True) + + def test_this_gate_checks_types_and_never_meaning(self, index): + """the types-not-meaning limit, stated as a test so nobody reads more into the gate. + + A board member recorded as an employee satisfies `:worksFor` exactly, and + is kept at every level. Domain and range bound which entities a property + may relate; they cannot see a wrong predicate whose endpoints fit. + """ + assert kept(index, relation('WORKS FOR', 'Person', 'Company'), enforce_domain_range=True) + +class TestEnforceDatatypes: + + @pytest.mark.parametrize('value,keep', [ + ('1994', True), + ('1,994', True), + ('nineteen ninety four', False), + ('1994.5', False), + ]) + def test_a_literal_that_does_not_parse_as_the_declared_type_is_dropped(self, index, value, keep): + assert kept(index, attribute('FOUNDED YEAR', value=value), enforce_datatypes=True) is keep + + def test_a_string_range_accepts_what_a_string_range_promises(self, index): + assert kept(index, attribute('TICKER SYMBOL', value='MFR'), enforce_datatypes=True) + + def test_a_relation_has_no_declared_datatype(self, index): + assert kept(index, relation('WORKS FOR'), enforce_datatypes=True) + + def test_an_unenforceable_declaration_keeps_the_fact_and_says_so(self, caplog): + """A value stored without validation must not look like a + validated one, so the WARN is unconditional on `report_violations`.""" + _warned_unvalidated_datatypes.clear() + index = OntologyIndex(datatype_properties={'urn:x#checksum': DatatypeProperty( + iri='urn:x#checksum', local_name='checksum', datatype=f'{XSD}hexBinary', + )}) + + with caplog.at_level(logging.WARNING): + assert kept(index, attribute('CHECKSUM', value='not hex at all'), enforce_datatypes=True) + + assert 'cannot validate' in caplog.text + assert f'{XSD}hexBinary' in caplog.text + + def test_the_warning_is_emitted_once_per_datatype(self, caplog): + _warned_unvalidated_datatypes.clear() + index = OntologyIndex(datatype_properties={'urn:x#checksum': DatatypeProperty( + iri='urn:x#checksum', local_name='checksum', datatype=f'{XSD}hexBinary', + )}) + filter_ = OntologyFilter(index=index, enforce_datatypes=True) + + with caplog.at_level(logging.WARNING): + run(filter_, attribute('CHECKSUM', value='a'), attribute('CHECKSUM', value='b')) + + assert caplog.text.count('cannot validate') == 1 + +class TestTheGatesAreIndependent: + + def test_one_gate_does_not_imply_another(self, index): + """A fact that breaches only entity types survives every other gate.""" + fact = relation('WORKS FOR', subject_class='Vehicle') + + assert not kept(index, fact, enforce_entity_types=True) + for gate in ('enforce_relationship_types', 'enforce_domain_range', 'enforce_datatypes'): + assert kept(index, fact, **{gate: True}) + + def test_nothing_is_dropped_with_every_gate_off(self, index): + assert kept(index, relation('HIRED BY', subject_class='Vehicle', object_class='Vehicle')) + + def test_a_fact_breaching_two_dimensions_is_dropped_once(self, index): + """The counters break down *drops*, so they must sum to facts lost.""" + counters = counted( + index, + relation('HIRED BY', subject_class='Vehicle'), + enforce_entity_types=True, + enforce_relationship_types=True, + ) + + assert counters.facts_dropped() == 1 + assert counters.facts_dropped_entity_type == 1 + assert counters.facts_dropped_relationship_type == 0 + +class TestDropTypeRestatements: + """The only dropping gate on at `align`, so what it drops has to be exact.""" + + @pytest.mark.parametrize('predicate', ['rdf:type', 'rdfs:subClassOf', 'owl:sameAs', 'skos:prefLabel']) + def test_ontology_language_is_dropped_whatever_the_value(self, index, predicate): + assert not kept(index, attribute(predicate, value='anything'), drop_type_restatements=True) + assert not kept(index, relation(predicate), drop_type_restatements=True) + + def test_a_bare_word_sharing_a_local_name_is_kept(self, index): + """Only the prefixed form is language. `RANGE` is also a real attribute of + a delivery van, and `DOMAIN` of a website.""" + assert kept(index, attribute('RANGE', value='400 miles'), drop_type_restatements=True) + assert kept(index, attribute('DOMAIN', value='meridian.example'), drop_type_restatements=True) + + @pytest.mark.parametrize('predicate', ['TYPE', 'CLASSIFICATION', 'IS A', 'is_a', 'DESCRIBED BY', 'CATEGORY']) + def test_a_value_restating_the_subjects_own_class_is_dropped(self, index, predicate): + assert not kept(index, attribute(predicate, value='Company'), drop_type_restatements=True) + + @pytest.mark.parametrize('value', ['company', 'COMPANY']) + def test_the_comparison_ignores_case_and_separators(self, index, value): + assert not kept(index, attribute('CLASSIFICATION', value=value), drop_type_restatements=True) + + def test_a_type_asserting_name_carrying_real_content_is_kept(self, index): + """Both conditions are necessary: a name-only rule would destroy these.""" + assert kept(index, attribute('CLASSIFICATION', value='football club'), drop_type_restatements=True) + assert kept(index, attribute('TYPE', value='haulage contractor'), drop_type_restatements=True) + + def test_a_predicate_that_is_not_type_asserting_is_kept(self, index): + """Even when the value repeats the class. Vagueness is not this gate's business.""" + assert kept(index, attribute('COMPETES WITH', value='Company'), drop_type_restatements=True) + assert kept(index, attribute('OCCUPATION', value='Company'), drop_type_restatements=True) + + def test_the_relation_shape_is_left_alone(self, index): + """The object then has its own identity in the graph, so dropping the edge + would not be the information-preserving move the attribute case is.""" + assert kept(index, relation('TYPE', subject_class='Company'), drop_type_restatements=True) + + @pytest.mark.parametrize('value', [' ', '']) + def test_an_empty_value_cannot_be_shown_to_duplicate_anything(self, index, value): + assert kept(index, attribute('TYPE', value=value), drop_type_restatements=True) + + def test_a_fact_with_neither_object_nor_complement_is_left_alone(self, index): + fact = Fact( + subject=Entity(value='Meridian Freight', classification='Company'), + predicate=Relation(value='TYPE'), + ) + assert kept(index, fact, drop_type_restatements=True) + + def test_a_declared_predicate_is_never_touched(self): + """An author who declares `:classification` has said what it means.""" + index = OntologyIndex( + classes={'urn:x#Company': OntologyClass(iri='urn:x#Company', local_name='Company')}, + datatype_properties={'urn:x#classification': DatatypeProperty( + iri='urn:x#classification', local_name='classification', datatype=f'{XSD}string', + )}, + ) + assert kept(index, attribute('CLASSIFICATION', value='Company'), drop_type_restatements=True) + + def test_nothing_is_dropped_with_the_flag_off(self, index): + assert kept(index, attribute('rdf:type', value='Company')) + assert kept(index, attribute('CLASSIFICATION', value='Company')) + +class TestAnnotation: + """Written once here so the build stage reads an answer.""" + + def test_a_resolved_class_records_its_iri(self, index): + [fact] = facts_of(run(OntologyFilter(index=index), relation('WORKS FOR'))) + assert (fact.subject.classIri, fact.object.classIri) == (f'{NS}Person', f'{NS}Company') + + def test_a_resolved_predicate_records_its_iri_and_canonical_name(self, index): + [fact] = facts_of(run(OntologyFilter(index=index), attribute('FOUNDED YEAR'))) + assert fact.predicate.propertyIri == f'{NS}foundedYear' + assert fact.predicate.canonicalName == 'foundedYear' + + def test_a_datatype_property_records_the_datatype_on_the_complement(self, index): + [fact] = facts_of(run(OntologyFilter(index=index), attribute('FOUNDED YEAR'))) + assert fact.complement.datatype == f'{XSD}integer' + + def test_an_object_property_leaves_no_datatype(self, index): + [fact] = facts_of(run(OntologyFilter(index=index), relation('WORKS FOR'))) + assert fact.object.datatype is None + + def test_annotation_is_not_gated_on_any_flag(self, index): + """This is how typed storage works at `off`.""" + [fact] = facts_of(run(OntologyFilter(index=index), attribute('FOUNDED YEAR'))) + assert fact.predicate.propertyIri is not None + + def test_an_unresolved_term_is_annotated_with_nothing(self, index): + [fact] = facts_of(run(OntologyFilter(index=index), relation('HIRED BY', subject_class='Vehicle'))) + assert (fact.predicate.propertyIri, fact.subject.classIri) == (None, None) + + def test_the_canonical_name_is_the_local_name_and_not_the_label(self): + """It is the key typed storage writes under, and an `rdfs:label` may + contain spaces.""" + index = OntologyIndex(datatype_properties={'urn:x#founded': DatatypeProperty( + iri='urn:x#founded', local_name='founded', label='Founded Year', datatype=f'{XSD}integer', + )}) + + [fact] = facts_of(run( + OntologyFilter(index=index, normalize_names=True), attribute('FOUNDED YEAR'), + )) + + assert fact.predicate.value == 'Founded Year' + assert fact.predicate.canonicalName == 'founded' + + def test_every_annotation_is_a_plain_string(self, index): + """The metadata is written through `json.dump` and revalidated strictly.""" + [fact] = facts_of(run(OntologyFilter(index=index), attribute('FOUNDED YEAR'))) + for value in (fact.subject.classIri, fact.predicate.propertyIri, + fact.predicate.canonicalName, fact.complement.datatype): + assert type(value) is str + +class TestCounters: + + def test_the_two_rewrite_kinds_are_counted_apart(self, index): + """A classification rewrite changes entity *identity*, since the id hashes + the classification in; a predicate rewrite only changes an edge label.""" + counters = counted(index, relation('WORKS FOR', subject_class='PERSON'), normalize_names=True) + + assert counters.predicates_rewritten == 1 + assert counters.classifications_rewritten == 1 + + @pytest.mark.parametrize('fact,flag,field', [ + (attribute('rdf:type'), 'drop_type_restatements', 'facts_dropped_type_restatement'), + (relation('WORKS FOR', subject_class='Vehicle'), 'enforce_entity_types', 'facts_dropped_entity_type'), + (relation('HIRED BY'), 'enforce_relationship_types', 'facts_dropped_relationship_type'), + (relation('WORKS FOR', subject_class='Company'), 'enforce_domain_range', 'facts_dropped_domain_range'), + (attribute('FOUNDED YEAR', value='x'), 'enforce_datatypes', 'facts_dropped_datatype'), + ]) + def test_each_drop_is_counted_under_the_setting_that_caused_it(self, index, fact, flag, field): + counters = counted(index, fact, **{flag: True}) + + assert getattr(counters, field) == 1 + assert counters.facts_dropped() == 1 + + def test_the_total_is_summed_from_the_dimensions(self, index): + """Derived rather than hand-written, so a new gate cannot arrive with a + counter the total omits.""" + counters = counted( + index, + attribute('rdf:type'), + relation('HIRED BY'), + drop_type_restatements=True, + enforce_relationship_types=True, + ) + assert counters.facts_dropped() == 2 + + def test_any_change_ignores_annotation(self, index): + """Annotation happens to every surviving fact at every level, so counting + it would make every call a change and the report carry no signal.""" + assert not counted(index, attribute('FOUNDED YEAR')).any_change() + assert counted(index, relation('WORKS FOR'), normalize_names=True).any_change() + + def test_the_summary_names_only_the_gates_that_dropped_something(self, index): + summary = counted(index, relation('HIRED BY'), enforce_relationship_types=True).summary() + + assert 'facts dropped: 1' in summary + assert 'enforce_relationship_types: 1' in summary + assert 'enforce_datatypes' not in summary + + def test_the_summary_always_carries_the_three_totals(self, index): + summary = counted(index, attribute('FOUNDED YEAR')).summary() + + assert 'classifications rewritten: 0' in summary + assert 'predicates rewritten: 0' in summary + assert 'facts dropped: 0' in summary + assert 'dropped by' not in summary + +class TestTheComponentContract: + + def test_a_node_without_topics_passes_through_untouched(self, index): + node = TextNode(text='x', metadata={'other': 1}) + assert OntologyFilter(index=index, normalize_names=True)([node])[0].metadata == {'other': 1} + + def test_the_topics_key_is_the_only_key_read_or_written(self, index): + topics = TopicCollection(topics=[Topic(value='t')]) + node = TextNode(text='x', metadata={TOPICS_KEY: topics.model_dump(), 'keep': 'me'}) + + OntologyFilter(index=index, normalize_names=True)([node]) + + assert node.metadata['keep'] == 'me' + assert set(node.metadata) == {TOPICS_KEY, 'keep'} + + def test_a_topic_is_never_dropped(self, index): + topics = TopicCollection(topics=[Topic(value='empty'), Topic(value='also empty')]) + node = TextNode(text='x', metadata={TOPICS_KEY: topics.model_dump()}) + + OntologyFilter(index=index, enforce_entity_types=True)([node]) + + assert len(TopicCollection.model_validate(node.metadata[TOPICS_KEY]).topics) == 2 + + def test_the_component_survives_the_spawn_boundary(self, index): + """Extraction pickles the component per node batch per worker, so a filter + that cannot round-trip is a filter that does not run.""" + filter_ = OntologyFilter(index=index, normalize_names=True, enforce_datatypes=True) + + revived = pickle.loads(pickle.dumps(filter_)) + + assert revived.normalize_names is True + assert revived.enforce_datatypes is True + assert revived.index.resolve_class('Company').iri == f'{NS}Company' + assert facts_of(run(revived, relation('WORKS FOR')))[0].predicate.value == 'worksFor' + + def test_the_serialization_name_is_stable(self, index): + """llama-index records it in a serialized pipeline, so it may not drift + with the class name.""" + assert OntologyFilter.class_name() == 'OntologyFilter' + + def test_a_bare_string_complement_is_read_the_same_as_an_entity(self, index): + """`Fact.complement` is `Union[Entity, str]` and both forms occur: the + parser builds an `Entity`, an older or hand-built payload may carry a str.""" + fact = Fact( + subject=Entity(value='Meridian Freight', classification='Company'), + predicate=Relation(value='TYPE'), + complement='Company', + ) + assert not kept(index, fact, drop_type_restatements=True) + + def test_the_index_holds_no_rdflib_term(self, index): + """What crosses the boundary is `str`, `list`, `dict` and `frozenset`.""" + for value in index.model_dump().values(): + assert type(value) in (dict, list, str) + +class TestReporting: + + def test_a_call_that_changed_something_logs_at_info(self, index, caplog): + with caplog.at_level(logging.INFO): + run( + OntologyFilter(index=index, normalize_names=True, report_violations=True), + relation('WORKS FOR', subject_class='PERSON'), + ) + + assert 'Ontology filter' in caplog.text + assert 'nodes: 1' in caplog.text + assert 'predicates rewritten: 1' in caplog.text + + def test_a_call_that_changed_nothing_does_not_log_at_info(self, index, caplog): + """Every gate is off by default, so a line of zeros per batch would bury + the batches that did something.""" + with caplog.at_level(logging.INFO): + run(OntologyFilter(index=index, report_violations=True), attribute('FOUNDED YEAR')) + + assert 'Ontology filter' not in caplog.text + + def test_nothing_is_logged_without_report_violations(self, index, caplog): + with caplog.at_level(logging.INFO): + run( + OntologyFilter(index=index, normalize_names=True), + relation('WORKS FOR', subject_class='PERSON'), + ) + + assert 'Ontology filter' not in caplog.text + + def test_the_line_says_how_many_nodes_it_covers(self, index, caplog): + """One `__call__` is the largest unit available on the far side of the + spawn boundary, so a reader has to be able to add the lines up.""" + topics = TopicCollection(topics=[Topic(value='t', statements=[Statement( + value='s', facts=[relation('WORKS FOR', subject_class='PERSON')], + )])]) + nodes = [ + TextNode(text='a', metadata={TOPICS_KEY: topics.model_dump()}), + TextNode(text='b', metadata={TOPICS_KEY: topics.model_dump()}), + TextNode(text='c', metadata={'other': 1}), + ] + + with caplog.at_level(logging.INFO): + OntologyFilter(index=index, normalize_names=True, report_violations=True)(nodes) + + assert 'nodes: 2' in caplog.text + +class TestAnIndexWithNoTerms: + """The degenerate ontology, which must be inert rather than fatal.""" + + def test_nothing_resolves_and_nothing_is_annotated(self): + [fact] = facts_of(run(OntologyFilter(index=OntologyIndex(), normalize_names=True), relation('WORKS FOR'))) + + assert fact.predicate.value == 'WORKS FOR' + assert fact.predicate.propertyIri is None + + def test_every_gate_drops_everything_it_can_be_asked_about(self): + index = OntologyIndex() + + assert not kept(index, relation('WORKS FOR'), enforce_entity_types=True) + assert not kept(index, relation('WORKS FOR'), enforce_relationship_types=True) + assert kept(index, relation('WORKS FOR'), enforce_domain_range=True) + assert kept(index, relation('WORKS FOR'), enforce_datatypes=True) + + def test_a_property_with_no_domain_declared_still_annotates(self): + index = OntologyIndex(object_properties={'urn:x#acquired': ObjectProperty( + iri='urn:x#acquired', local_name='acquired', + )}) + [fact] = facts_of(run(OntologyFilter(index=index, enforce_domain_range=True), relation('ACQUIRED'))) + + assert fact.predicate.propertyIri == 'urn:x#acquired' diff --git a/lexical-graph/tests/unit/indexing/extract/test_ontology_fixtures.py b/lexical-graph/tests/unit/indexing/extract/test_ontology_fixtures.py new file mode 100644 index 00000000..13acf28e --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ontology_fixtures.py @@ -0,0 +1,219 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Guards on the test ontologies. + +These are not tests of the loader - `test_ontology.py` does that with inline +Turtle. These assert that the shared fixtures still are what the rest of the +feature's tests assume they are, so a well-meaning edit to `company.ttl` cannot +quietly remove the shape some later test depends on: + +* every declared XSD range the datatype path claims to handle is present; +* `:SportsTeam`'s label differs from its local name, which is the whole reason + the naming contract exists; +* each malformed fixture still fails for the reason it was written for; +* the verification corpus and its manifest agree, and between them the + documents put every declared datatype property in play. +""" + +import json + +import pytest +from rdflib import Graph + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology import ( + XSD_NAMESPACE, + Ontology, + OntologyLoadError, +) + +NAMESPACE = 'http://example.com/company#' + +def iri(local_name:str) -> str: + """Expand a local name against the test ontology's namespace.""" + return NAMESPACE + local_name + +class TestCompanyOntology: + """`company.ttl` - the ontology these tests run on.""" + + def test_loads(self, company_ontology): + """It is valid: it loads without an `OntologyLoadError`.""" + assert company_ontology.namespace == NAMESPACE + + def test_declared_classes(self, company_ontology): + """The class vocabulary is fixed - later tests name these.""" + assert set(company_ontology.index().classes) == { + iri('Agent'), iri('Person'), iri('Athlete'), iri('Company'), + iri('SportsTeam'), + } + + def test_declared_properties(self, company_ontology): + """The property vocabulary is fixed too.""" + index = company_ontology.index() + + assert set(index.object_properties) == { + iri('worksFor'), iri('playsFor'), iri('subsidiaryOf'), iri('acquired'), + } + assert set(index.datatype_properties) == { + iri('foundedYear'), iri('revenue'), iri('isPubliclyTraded'), + iri('incorporatedOn'), iri('tickerSymbol'), iri('jobTitle'), + iri('officialName'), + } + + def test_hierarchy_is_three_deep(self, company_ontology): + """`SportsTeam < Company < Agent`, so the closure has something to close + over that a single subClassOf hop would not reach.""" + index = company_ontology.index() + + assert index.classes[iri('SportsTeam')].ancestors == frozenset( + {iri('SportsTeam'), iri('Company'), iri('Agent')} + ) + assert index.is_subclass_of(iri('SportsTeam'), iri('Agent')) is True + assert index.is_subclass_of(iri('Athlete'), iri('Agent')) is True + + def test_a_label_differs_from_its_local_name(self, company_ontology): + """`:SportsTeam` is labelled `Sports Team`. Without this the naming + contract's regression case has nothing to run against - `SportsTeam` + rendered directly comes back from `.title()` as `Sportsteam`.""" + sports_team = company_ontology.index().classes[iri('SportsTeam')] + + assert sports_team.local_name == 'SportsTeam' + assert sports_team.label == 'Sports Team' + assert sports_team.label != sports_team.local_name + + def test_alt_labels_are_present_on_a_class_and_a_property(self, company_ontology): + """Both halves of the alias path are exercised.""" + index = company_ontology.index() + + assert 'Ball Club' in index.classes[iri('SportsTeam')].aliases + assert 'REQ_TO_HC' in index.object_properties[iri('worksFor')].aliases + + def test_an_object_property_is_unconstrained(self, company_ontology): + """`:acquired` has no domain and no range, so it matches anything.""" + acquired = company_ontology.index().object_properties[iri('acquired')] + + assert acquired.domain is None + assert acquired.range is None + + def test_a_datatype_property_is_unconstrained(self, company_ontology): + """`:officialName` has no domain - any subject may carry it.""" + assert company_ontology.index().datatype_properties[iri('officialName')].domain is None + + def test_a_narrow_property_needs_the_subclass_closure(self, company_ontology): + """`:playsFor` runs `Athlete -> SportsTeam`, both of which are strict + subclasses. A fact stated of a `Person` and a `Company` only satisfies + it through the closure, which is what makes the domain/range check worth + testing.""" + index = company_ontology.index() + plays_for = index.object_properties[iri('playsFor')] + + assert plays_for.domain == iri('Athlete') + assert plays_for.range == iri('SportsTeam') + assert index.is_subclass_of(iri('Athlete'), iri('Person')) is True + assert index.is_subclass_of(iri('SportsTeam'), iri('Company')) is True + + @pytest.mark.parametrize( + 'datatype', ['integer', 'double', 'boolean', 'date', 'string'] + ) + def test_every_claimed_xsd_range_is_covered(self, company_ontology, datatype): + """The five ranges the datatype path claims to render and coerce each + appear at least once. This is the guard that stops the attribute tests + passing on a vocabulary that never exercises them.""" + declared = { + p.datatype for p in company_ontology.index().datatype_properties.values() + } + + assert XSD_NAMESPACE + datatype in declared + + def test_descriptions_are_present_but_not_universal(self, company_ontology): + """Some terms carry `rdfs:comment` and some deliberately do not, so the + rendering has to handle both.""" + index = company_ontology.index() + classes = index.classes.values() + properties = [ + *index.object_properties.values(), *index.datatype_properties.values() + ] + + assert any(c.description for c in classes) + assert any(p.description for p in properties) + assert any(p.description is None for p in properties) + + def test_comments_do_not_leak_into_terms(self, company_ontology): + """The file's `#` header is Turtle comment syntax, not content - no term + is named after it.""" + assert all( + i.startswith(NAMESPACE) for i in company_ontology.index().classes + ) + +class TestUpperSnakeOntology: + """`upper_snake_names.ttl` - the same vocabulary in the other convention.""" + + @pytest.fixture + def upper_snake_ontology(self, ontology_fixtures_dir): + return Ontology.from_turtle(ontology_fixtures_dir / 'upper_snake_names.ttl') + + def test_loads(self, upper_snake_ontology): + """It is a valid ontology in its own right.""" + assert upper_snake_ontology.namespace == 'http://example.com/upper#' + + def test_names_are_authored_upper_snake(self, upper_snake_ontology): + """The local names really are `WORKS_FOR` / `SPORTS_TEAM`. These must + resolve as `worksFor` and `SportsTeam` do, and that + cannot be tested against an ontology that spells them camelCase.""" + index = upper_snake_ontology.index() + + assert 'http://example.com/upper#WORKS_FOR' in index.object_properties + assert 'http://example.com/upper#SPORTS_TEAM' in index.classes + assert index.classes['http://example.com/upper#SPORTS_TEAM'].label == 'Sports Team' + + def test_it_does_not_collide_with_company_ttl( + self, upper_snake_ontology, company_ontology + ): + """The two ontologies are separate files on purpose: `:worksFor` and + `:WORKS_FOR` share a resolution key, so declaring both in one ontology + would make resolution deliberately ambiguous.""" + assert upper_snake_ontology.namespace != company_ontology.namespace + +class TestMalformedFixtures: + """Each malformed fixture still fails for the reason it was written for.""" + + # (fixture name, substring the message must contain) + CASES = [ + ('subclass_cycle', 'cycle'), + ('dangling_subclass_reference', 'dangling rdfs:subClassOf'), + ('dangling_domain_reference', 'dangling rdfs:domain'), + ('dangling_range_reference', 'dangling rdfs:range'), + ('dual_typed_property', 'both an owl:ObjectProperty'), + ('non_xsd_range', 'not an XSD datatype'), + ('missing_datatype_range', 'no rdfs:range'), + ('malformed_syntax', 'Failed to parse Turtle file'), + ] + + @pytest.mark.parametrize('name, expected', CASES) + def test_fixture_fails_for_its_stated_reason( + self, malformed_ttl_path, name, expected + ): + """The fixture exists, and the load error is the specific one it was + authored to trigger - not some unrelated fault that happens to raise.""" + path = malformed_ttl_path(name) + assert path.is_file(), f'missing malformed fixture: {path}' + + with pytest.raises(OntologyLoadError, match=expected): + Ontology.from_turtle(path) + + def test_only_the_syntax_fixture_is_unparseable(self, malformed_ttl_path): + """Every other malformed fixture is valid Turtle that is an invalid + ontology. If one of them stopped parsing, its structural case would + never be reached and the test above would pass for the wrong reason.""" + for name, _ in self.CASES: + if name == 'malformed_syntax': + continue + Graph().parse(source=str(malformed_ttl_path(name)), format='turtle') + + def test_every_malformed_fixture_on_disk_is_covered( + self, ontology_fixtures_dir + ): + """No fixture sits in the directory untested.""" + on_disk = {p.stem for p in (ontology_fixtures_dir / 'malformed').glob('*.ttl')} + + assert on_disk == {name for name, _ in self.CASES} diff --git a/lexical-graph/tests/unit/indexing/extract/test_ontology_pipeline.py b/lexical-graph/tests/unit/indexing/extract/test_ontology_pipeline.py new file mode 100644 index 00000000..f70c5571 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ontology_pipeline.py @@ -0,0 +1,555 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""`OntologyFilter` as a pipeline component. + +Two claims that only a running pipeline can settle, and one that only a log can. + +**It survives spawn.** `test_ontology_filter.py::TestItSurvivesTheProcessBoundary` +already pickles the component and asserts it classifies identically, holds only +plain data, and declares no custom `__init__`. Those are the preconditions, and +they are all in-process: a `pickle.dumps` round trip in the parent proves the +state is picklable, not that a *fresh interpreter* can rebuild the component. The +difference is real and it is where this feature's most plausible failure lives - +spawn re-imports every module from scratch, so an import that only resolves +because the parent imported something first, or a `GraphRAGConfig` value set +programmatically rather than from the environment, works in the parent and +vanishes in the worker. `run_pipeline` uses `ProcessPoolExecutor(mp_context= +'spawn')`, so the only way to know is to run one. + +`FixedTopicsExtractor` stands in for the topic extractor because the real one +calls an LLM, and a `Mock` LLM is the one thing that certainly does not cross a +process boundary. It is declared at module level so `spawn` can import it by +name in the child - which is also why it is a class rather than a closure. + +**Reporting.** `report_violations` is asserted in-process, because the claim is +about what the message says rather than about where it is emitted, and a log +record from a spawned worker does not reach `caplog` in the parent at all. That +limit is itself part of the contract: see +`test_it_reports_from_inside_the_worker`. +""" + +import logging +import pickle + +from pathlib import Path + +import pytest +from llama_index.core.schema import Document, TextNode, TransformComponent + +from graphrag_toolkit.lexical_graph.config import GraphRAGConfig +from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY +from graphrag_toolkit.lexical_graph.indexing.extract import ExtractionPipeline +from graphrag_toolkit.lexical_graph.indexing.utils.pipeline_utils import _init_worker +from graphrag_toolkit.lexical_graph import logging as graphrag_logging +from graphrag_toolkit.lexical_graph.logging import ( + apply_logging_config, + get_applied_logging_config, + set_logging_config, +) +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology import Ontology +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_filter import ( + FilterCounters, + OntologyFilter, +) +from graphrag_toolkit.lexical_graph.indexing.model import ( + Entity, + Fact, + Relation, + Statement, + Topic, + TopicCollection, +) + +FIXTURES = Path(__file__).parent.parent.parent.parent / 'fixtures' / 'ontologies' + +FILTER_MODULE = 'graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology_filter' + +@pytest.fixture(scope='module') +def company_index(): + return Ontology.load(FIXTURES / 'company.ttl').index() + +@pytest.fixture +def restore_logging(): + """Put global logging state back after a test that reconfigures it. + + `set_logging_config` calls `logging.config.dictConfig`, which replaces the + root logger's handlers process-wide. Without this, one test here would change + how every later test in the session logs - and, because `pytest`'s own + capturing sits on those handlers, could quietly break `caplog` assertions in + the classes below. + """ + root = logging.getLogger() + saved_handlers = root.handlers[:] + saved_level = root.level + saved_config = get_applied_logging_config() + + yield + + root.handlers[:] = saved_handlers + root.setLevel(saved_level) + + # Reset the remembered config through the module attribute rather than + # `apply_logging_config`, which cannot restore None - and a stale config left + # here would be propagated to the workers of every later spawned run. + graphrag_logging._applied_logging_config = saved_config + +def topics_the_model_might_emit(): + """One resolvable relation, one resolvable attribute, one of neither. + + Spelled the way `parse_extracted_topics` spells things - upper case with + spaces - because that is what the filter has to recognise, and the point of + the run is that it recognises it in a process that was not there when the + ontology was loaded. + """ + return TopicCollection(topics=[Topic( + value='Employment', + entities=[ + Entity(value='Amy Bell', classification='Person'), + Entity(value='Example Corp', classification='Company'), + ], + statements=[Statement( + value='Amy Bell works for Example Corp', + facts=[ + Fact( + subject=Entity(value='Amy Bell', classification='Person'), + predicate=Relation(value='WORKS FOR'), + object=Entity(value='Example Corp', classification='Company'), + ), + Fact( + subject=Entity(value='Example Corp', classification='Company'), + predicate=Relation(value='OFFICIAL NAME'), + complement=Entity(value='Example Corporation'), + ), + Fact( + subject=Entity(value='Amy Bell', classification='Person'), + predicate=Relation(value='HIRED BY'), + object=Entity(value='Example Corp', classification='Company'), + ), + ], + )], + )]) + +class FixedTopicsExtractor(TransformComponent): + """A stand-in for `TopicExtractor` that needs no LLM and pickles. + + Declared at module level and holding only a dict, so `spawn` can rebuild it + in a child from `(module, qualname)` plus its state. A lambda, a closure, or + a `Mock` could do none of that, which is the whole reason this class exists. + """ + + topics:dict = {} + + @classmethod + def class_name(cls) -> str: + return 'FixedTopicsExtractor' + + def __call__(self, nodes, **kwargs): + for node in nodes: + node.metadata[TOPICS_KEY] = self.topics + return nodes + +def topics_from(source_documents): + """Every `TopicCollection` the pipeline emitted, revalidated.""" + return [ + TopicCollection.model_validate(node.metadata[TOPICS_KEY]) + for source_document in source_documents + for node in source_document.nodes + if TOPICS_KEY in node.metadata + ] + +def facts_of(topics): + return [f for t in topics.topics for s in t.statements for f in s.facts] + +def predicates_of(topics): + return [f.predicate.value for f in facts_of(topics)] + +def run_spawned(components, documents, num_workers=2): + """Run the real `ExtractionPipeline`, which runs a real spawned pool. + + Constructed directly rather than through `ExtractionPipeline.create`, which + wraps `extract` in a `Pipe` for `|` chaining and so has no `extract` of its + own. Same object underneath. + """ + pipeline = ExtractionPipeline(components=components, num_workers=num_workers) + return list(pipeline.extract(documents)) + +def node_for(topics): + return TextNode(text='chunk text', metadata={TOPICS_KEY: topics.model_dump()}) + +def stand_in_extractor(): + return FixedTopicsExtractor(topics=topics_the_model_might_emit().model_dump()) + +# Each spawned run costs about four seconds, almost all of it interpreter startup, +# so the runs are module-scoped fixtures and the assertions read from them. Four +# runs rather than one because they are four different configurations, and a +# configuration is exactly what these tests are about. + +@pytest.fixture(scope='module') +def align_run(company_index): + """`align`: normalize, no gates. One document, so one worker does the work.""" + return topics_from(run_spawned( + [stand_in_extractor(), OntologyFilter(index=company_index, normalize_names=True)], + [Document(text='Amy Bell works for Example Corp.')], + )) + +@pytest.fixture(scope='module') +def strict_run(company_index): + """`strict`: normalize and all four gates.""" + return topics_from(run_spawned( + [ + stand_in_extractor(), + OntologyFilter( + index=company_index, + normalize_names=True, + enforce_entity_types=True, + enforce_relationship_types=True, + enforce_domain_range=True, + enforce_datatypes=True, + ), + ], + [Document(text='Amy Bell works for Example Corp.')], + )) + +@pytest.fixture(scope='module') +def unfiltered_run(): + """The control: the same spawned pipeline with no filter in it.""" + return topics_from(run_spawned( + [stand_in_extractor()], [Document(text='Amy Bell works for Example Corp.')] + )) + +@pytest.fixture(scope='module') +def many_document_run(company_index): + """Enough documents that `node_batcher` fills both workers.""" + return topics_from(run_spawned( + [stand_in_extractor(), OntologyFilter(index=company_index, normalize_names=True)], + [Document(text=f'Document {i}') for i in range(6)], + )) + +class TestItSurvivesASpawnedPipeline: + """The pipeline-wiring claims, run rather than argued.""" + + def test_normalization_happens_in_the_worker(self, align_run): + """`WORKS FOR` becomes `worksFor` on the far side of the boundary. + + The assertion is on the *output* of a spawned run, so it fails if the + component cannot be rebuilt in a child, if its index arrives empty, or if + `ontology_filter.py` grows an import that only resolves in the parent. + """ + assert align_run, 'the spawned pipeline returned no topics at all' + for topics in align_run: + assert 'worksFor' in predicates_of(topics) + assert 'WORKS FOR' not in predicates_of(topics) + + def test_two_workers_both_filter(self, many_document_run): + """`node_batcher` splits by `num_workers`, so with a single document only + one child ever runs and a filter that failed to rebuild in the *second* + interpreter would go unnoticed.""" + assert len(many_document_run) >= 6 + assert all('worksFor' in predicates_of(t) for t in many_document_run) + + def test_annotation_crosses_the_boundary_too(self, align_run): + """The annotation is written once, in extraction, and the + build stage reads it rather than an ontology - so it has to arrive on the + node that comes back out of the worker, not merely be computable.""" + for topics in align_run: + annotated = [f for f in facts_of(topics) if f.predicate.propertyIri] + assert annotated, 'nothing was annotated in the worker' + for fact in annotated: + assert fact.predicate.canonicalName + assert fact.subject.classIri + + def test_enforcement_drops_in_the_worker(self, strict_run): + """The gates, not just the rename. `HIRED BY` resolves to nothing the + ontology declares, so `strict`'s relationship-type gate must discard it - + in the child, where the decision actually gets made.""" + assert strict_run + for topics in strict_run: + assert 'HIRED BY' not in predicates_of(topics) + assert 'worksFor' in predicates_of(topics) + + def test_the_unfiltered_pipeline_is_the_control(self, unfiltered_run): + """Without the filter the same spawned run leaves `WORKS FOR` alone. + + Otherwise every assertion above could be satisfied by a stub extractor + that happened to emit the authored spelling, and the tests would prove + nothing about the filter at all. + """ + assert unfiltered_run + for topics in unfiltered_run: + assert 'WORKS FOR' in predicates_of(topics) + assert 'worksFor' not in predicates_of(topics) + + def test_the_stand_in_extractor_is_itself_picklable(self): + """Named so that a failure here reads as a problem with the test's own + scaffolding rather than with the filter.""" + extractor = stand_in_extractor() + + assert pickle.loads(pickle.dumps(extractor)).topics == extractor.topics + +class TestTheReportIsReachableFromAWorker: + """That the report is *emitted* is not that it is *seen*. + + Extraction components run only inside spawn-started workers, and until this + was fixed a worker's logging was never configured: `_init_worker` re-applied + the `GraphRAGConfig` snapshot but not `logging.config.dictConfig` state, which + is separate global state and does not travel in the snapshot. So a worker's + root logger sat at WARNING with no handler but `lastResort`, and every line + `report_violations=True` produced was discarded. The feature emitted nothing + in the only code path that runs it, and every in-process assertion about the + report's wording passed throughout. + + `caplog` cannot catch this: a child's log record never reaches the parent's + handlers. A file handler can, because the child opens the same path in append + mode, so this asserts on the file the parent asked for. + """ + + def test_the_report_reaches_a_configured_handler(self, company_index, tmp_path, restore_logging): + log_file = tmp_path / 'extraction.log' + set_logging_config('INFO', filename=str(log_file)) + + run_spawned( + [ + stand_in_extractor(), + OntologyFilter( + index=company_index, + normalize_names=True, + enforce_relationship_types=True, + report_violations=True, + ), + ], + [Document(text='Amy Bell works for Example Corp.')], + num_workers=1, + ) + + assert log_file.exists(), 'the worker never wrote to the configured log file' + written = log_file.read_text() + assert 'Ontology filter' in written, written + assert 'enforce_relationship_types: 1' in written, written + + def test_an_unconfigured_parent_leaves_the_worker_alone(self, restore_logging): + """None means "the parent never configured logging", and must not be + turned into a config in the worker - a library that starts emitting to + stdout because it spawned a process is worse than one that stays quiet.""" + root = logging.getLogger() + root.setLevel(logging.CRITICAL) + handlers_before = root.handlers[:] + + apply_logging_config(None) + + assert root.level == logging.CRITICAL + assert root.handlers == handlers_before + + def test_the_worker_initializer_applies_it(self, restore_logging): + """`_init_worker` is the seam, so it is asserted directly as well as + end-to-end: the end-to-end test above would also pass if some other part + of the stack happened to configure logging in the child.""" + set_logging_config('INFO') + config = get_applied_logging_config() + assert config is not None + + logging.getLogger().setLevel(logging.CRITICAL) + + _init_worker(GraphRAGConfig.get_config_snapshot(), config) + + assert logging.getLogger().level == logging.INFO + + def test_the_config_it_propagates_is_picklable(self, restore_logging): + """It crosses the boundary in `initargs`, so an unpicklable entry would + fail the pool's startup rather than degrade quietly.""" + set_logging_config('INFO') + + assert pickle.loads(pickle.dumps(get_applied_logging_config())) is not None + + def test_repeated_configuration_does_not_accumulate_handlers(self, tmp_path, restore_logging): + """`set_advanced_logging_config` deep-copies its base config. With the + shallow copy it used to make, a second call appended `file_handler` to a + list shared with the module-level base, so every later call inherited the + previous call's log file.""" + set_logging_config('INFO', filename=str(tmp_path / 'first.log')) + set_logging_config('INFO', filename=str(tmp_path / 'second.log')) + + handlers = get_applied_logging_config()['loggers']['']['handlers'] + + assert handlers.count('file_handler') == 1, handlers + + def test_configuring_without_a_filename_adds_no_file_handler(self, restore_logging): + set_logging_config('INFO') + + assert 'file_handler' not in get_applied_logging_config()['loggers']['']['handlers'] + +class TestTheViolationReport: + """What `report_violations=True` actually says.""" + + def strict_filter(self, index, report_violations=True): + return OntologyFilter( + index=index, + normalize_names=True, + enforce_entity_types=True, + enforce_relationship_types=True, + enforce_domain_range=True, + enforce_datatypes=True, + report_violations=report_violations, + ) + + def test_nothing_is_logged_when_reporting_is_off(self, company_index, caplog): + with caplog.at_level(logging.DEBUG, logger=FILTER_MODULE): + self.strict_filter(company_index, report_violations=False)( + [node_for(topics_the_model_might_emit())] + ) + + assert caplog.records == [] + + def test_it_reports_rewrites_and_drops_at_info(self, company_index, caplog): + with caplog.at_level(logging.INFO, logger=FILTER_MODULE): + self.strict_filter(company_index)([node_for(topics_the_model_might_emit())]) + + assert len(caplog.records) == 1 + message = caplog.records[0].getMessage() + assert 'predicates rewritten: ' in message + assert 'classifications rewritten: ' in message + assert 'facts dropped: 1' in message + + def test_it_names_the_setting_that_dropped_the_fact(self, company_index, caplog): + """The count alone is not actionable. `HIRED BY` is dropped by the + relationship-type gate, and the report has to say which gate, because that + is the setting a user would turn off.""" + with caplog.at_level(logging.INFO, logger=FILTER_MODULE): + self.strict_filter(company_index)([node_for(topics_the_model_might_emit())]) + + message = caplog.records[0].getMessage() + assert 'enforce_relationship_types: 1' in message + + def test_it_never_names_a_gate_that_did_not_run(self, company_index, caplog): + """A dimension the user left off has an unbreakable zero, so it must not + appear - a report listing `enforce_datatypes: 0` invites the reader to + conclude datatypes were checked.""" + ontology_filter = OntologyFilter( + index=company_index, + normalize_names=True, + enforce_relationship_types=True, + report_violations=True, + ) + + with caplog.at_level(logging.INFO, logger=FILTER_MODULE): + ontology_filter([node_for(topics_the_model_might_emit())]) + + message = caplog.records[0].getMessage() + assert 'enforce_relationship_types: 1' in message + for setting in ('enforce_entity_types', 'enforce_domain_range', 'enforce_datatypes'): + assert setting not in message, setting + + def test_a_call_that_changed_nothing_stays_below_info(self, company_index, caplog): + """Every gate is off at `align`, so on a corpus the ontology already + matches most batches change nothing. Those must not each emit a line of + zeros at INFO, or the batches that did something get buried.""" + already_normalized = OntologyFilter( + index=company_index, normalize_names=True, report_violations=True + ) + node = node_for(topics_the_model_might_emit()) + already_normalized([node]) + + with caplog.at_level(logging.INFO, logger=FILTER_MODULE): + already_normalized([node]) + + assert caplog.records == [] + + with caplog.at_level(logging.DEBUG, logger=FILTER_MODULE): + already_normalized([node]) + + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.DEBUG + + def test_it_says_how_many_nodes_the_counts_cover(self, company_index, caplog): + """One `__call__` is one node batch in one worker, so a reader has to be + able to add the lines up. Without the node count the line looks like a + run-wide total.""" + nodes = [node_for(topics_the_model_might_emit()) for _ in range(3)] + + with caplog.at_level(logging.INFO, logger=FILTER_MODULE): + self.strict_filter(company_index)(nodes) + + assert 'nodes: 3' in caplog.records[0].getMessage() + + def test_a_node_without_topics_is_not_counted(self, company_index, caplog): + """The count is of nodes the filter looked at, not of nodes it was + handed - otherwise a batch of chunks that never reached the extractor + reads as a batch that was filtered.""" + nodes = [node_for(topics_the_model_might_emit()), TextNode(text='no topics')] + + with caplog.at_level(logging.INFO, logger=FILTER_MODULE): + self.strict_filter(company_index)(nodes) + + assert 'nodes: 1' in caplog.records[0].getMessage() + + def test_one_line_per_call_and_no_state_between_calls(self, company_index, caplog): + """The counters are a local, so a second call reports its own work and not + the sum. Asserted because an instance attribute would pass every test + above and then read as a total that is really one worker's share of one + batch.""" + ontology_filter = self.strict_filter(company_index) + + with caplog.at_level(logging.INFO, logger=FILTER_MODULE): + ontology_filter([node_for(topics_the_model_might_emit())]) + ontology_filter([node_for(topics_the_model_might_emit())]) + + assert len(caplog.records) == 2 + assert caplog.records[0].getMessage() == caplog.records[1].getMessage() + + def test_it_reports_from_inside_the_worker(self, company_index): + """The report is emitted where the filtering happens, which is a child + process - so it reaches the configured logging handlers of that process + and not `caplog` in the parent. + + Recorded as a test rather than a comment because it is the reason every + other test in this class runs in-process, and because it is the thing a + user will notice: `report_violations=True` with `num_workers=4` produces + four processes' worth of lines, none of which are totals. + """ + source = OntologyFilter._report.__doc__ + + assert 'one node batch in one worker' in source + +class TestTheCounterSummary: + """`FilterCounters.summary` on its own, where the wording is cheap to pin.""" + + def test_the_three_totals_are_always_present(self): + summary = FilterCounters().summary() + + assert 'classifications rewritten: 0' in summary + assert 'predicates rewritten: 0' in summary + assert 'facts dropped: 0' in summary + + def test_no_breakdown_when_nothing_was_dropped(self): + assert 'dropped by' not in FilterCounters(predicates_rewritten=3).summary() + + def test_the_breakdown_follows_the_gate_order(self): + """Broadest gate first, matching `_enforce`, so the line reads in the + order the decisions were taken.""" + summary = FilterCounters( + facts_dropped_entity_type=1, + facts_dropped_relationship_type=2, + facts_dropped_domain_range=3, + facts_dropped_datatype=4, + ).summary() + + settings = [ + 'enforce_entity_types', 'enforce_relationship_types', + 'enforce_domain_range', 'enforce_datatypes', + ] + positions = [summary.index(setting) for setting in settings] + assert positions == sorted(positions), summary + assert 'facts dropped: 10' in summary + + def test_annotation_alone_is_not_a_change(self): + """Every surviving fact is annotated at every level, so counting it would + make `any_change` always true and the report would carry no signal.""" + assert FilterCounters().any_change() is False + + @pytest.mark.parametrize('field', [ + 'classifications_rewritten', 'predicates_rewritten', + 'facts_dropped_entity_type', 'facts_dropped_relationship_type', + 'facts_dropped_domain_range', 'facts_dropped_datatype', + ]) + def test_any_single_count_is_a_change(self, field): + assert FilterCounters(**{field: 1}).any_change() is True diff --git a/lexical-graph/tests/unit/indexing/extract/test_ontology_prompt.py b/lexical-graph/tests/unit/indexing/extract/test_ontology_prompt.py new file mode 100644 index 00000000..3f9db771 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ontology_prompt.py @@ -0,0 +1,312 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Rendering the ontology into the two prompt blocks. + +The block is prose fed to a model, so almost nothing about its wording is a +correctness claim and none of it is asserted here. What is asserted is the +handful of structural properties other code depends on: that `off` renders +nothing, that the two output channels stay separated, that a name rendered here +is a name the response parser can hand back, and that the level changes only the +closing guidance and never the vocabulary. +""" + +from pathlib import Path + +import pytest + +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.naming import resolution_key +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.ontology import Ontology +from graphrag_toolkit.lexical_graph.indexing.extract.ontology.prompt_constraint import ( + rendered_class_names, +) + +FIXTURES = Path(__file__).parent.parent.parent.parent / 'fixtures' / 'ontologies' + +CLASSES_HEADING = '## Entity types' +RELATIONSHIPS_HEADING = '## Relationships' +ATTRIBUTES_HEADING = '## Attributes' +PROTOCOL_HEADING = '## Using this vocabulary' + +@pytest.fixture(scope='module') +def company(): + return Ontology.load(FIXTURES / 'company.ttl') + +@pytest.fixture(scope='module') +def empty(): + return Ontology.from_turtle_string( + '@prefix owl: . ' + ' a owl:Ontology .' + ) + +def section(block, heading, headings=(CLASSES_HEADING, RELATIONSHIPS_HEADING, + ATTRIBUTES_HEADING, PROTOCOL_HEADING)): + """The text of one `##` section of a rendered block.""" + start = block.index(heading) + ends = [block.index(other) for other in headings if other in block and block.index(other) > start] + return block[start:min(ends)] if ends else block[start:] + +class TestOffRendersNothing: + """the `off`-changes-nothing guarantee: `off` says nothing, in either block or either format.""" + + @pytest.mark.parametrize('vocabulary_format', ['prose', 'turtle']) + def test_the_topics_block_is_empty(self, company, vocabulary_format): + assert company.format_as_prompt_constraint('off', vocabulary_format) == '' + + def test_the_propositions_block_is_empty(self, company): + assert company.format_as_proposition_constraint('off') == '' + + def test_the_class_names_are_still_available(self, company): + """`class_names` is independent of the level: the level decides how much + the prompt says *about* the vocabulary, not what the vocabulary is.""" + assert company.class_names() == rendered_class_names(company.index()) + assert len(company.class_names()) == 5 + +class TestAnOntologyWithNoTerms: + """A header with no vocabulary under it is worse than no block at all.""" + + @pytest.mark.parametrize('level', ['align', 'strict']) + @pytest.mark.parametrize('vocabulary_format', ['prose', 'turtle']) + def test_nothing_is_rendered(self, empty, level, vocabulary_format): + assert empty.format_as_prompt_constraint(level, vocabulary_format) == '' + + @pytest.mark.parametrize('level', ['align', 'strict']) + def test_the_propositions_block_is_empty_too(self, empty, level): + assert empty.format_as_proposition_constraint(level) == '' + +class TestValidation: + + @pytest.mark.parametrize('level', ['guide', 'STRICT', '', None]) + @pytest.mark.parametrize('vocabulary_format', ['prose', 'turtle']) + def test_an_unknown_level_raises_in_either_format(self, company, level, vocabulary_format): + with pytest.raises(ValueError, match='Unknown ontology authority level'): + company.format_as_prompt_constraint(level, vocabulary_format) + + @pytest.mark.parametrize('level', ['guide', 'STRICT', '', None]) + def test_an_unknown_level_raises_for_the_propositions_block(self, company, level): + with pytest.raises(ValueError, match='Unknown ontology authority level'): + company.format_as_proposition_constraint(level) + + def test_an_unknown_vocabulary_format_raises(self, company): + with pytest.raises(ValueError, match='Unknown ontology vocabulary format'): + company.format_as_prompt_constraint('align', 'json-ld') + +class TestTheChannelsStaySeparated: + """An attribute emitted where a relationship belongs is lost, + so which section a term appears under is load-bearing.""" + + @pytest.fixture(scope='class') + def block(self, company): + return company.format_as_prompt_constraint('align') + + def test_all_three_sections_and_the_protocol_are_present(self, block): + for heading in (CLASSES_HEADING, RELATIONSHIPS_HEADING, ATTRIBUTES_HEADING, PROTOCOL_HEADING): + assert heading in block + + def test_every_declared_class_appears_in_the_class_section(self, company, block): + classes = section(block, CLASSES_HEADING) + for name in company.class_names(): + assert name in classes + + def test_every_object_property_appears_only_on_the_relationship_channel(self, company, block): + relationships = section(block, RELATIONSHIPS_HEADING) + attributes = section(block, ATTRIBUTES_HEADING) + + for name in ('WORKS_FOR', 'PLAYS_FOR', 'SUBSIDIARY_OF', 'ACQUIRED'): + assert name in relationships + assert name not in attributes + + def test_every_datatype_property_appears_only_on_the_attribute_channel(self, company, block): + relationships = section(block, RELATIONSHIPS_HEADING) + attributes = section(block, ATTRIBUTES_HEADING) + + for name in ('FOUNDED_YEAR', 'REVENUE', 'IS_PUBLICLY_TRADED', 'TICKER_SYMBOL'): + assert name in attributes + assert name not in relationships + +class TestWhatTheRenderedNamesPromise: + + @pytest.fixture(scope='class') + def block(self, company): + return company.format_as_prompt_constraint('align') + + def test_every_name_the_block_shows_resolves_when_it_comes_back(self, company, block): + """The invariant `naming.py` documents, checked against the rendered text. + + Each name is taken out of the block itself and put back through the + parser's transform and then the index, so a rendering change that broke + the round trip fails here rather than at extraction time. + """ + index = company.index() + + for name in rendered_class_names(index): + assert name in block + assert index.resolve_class(name.title()) is not None + + for heading, resolve in ( + (RELATIONSHIPS_HEADING, index.resolve_object_predicate), + (ATTRIBUTES_HEADING, index.resolve_datatype_predicate), + ): + names = [ + word for line in section(block, heading).splitlines() + for word in [line.strip().split(' ')[0]] + if word.isupper() and word.replace('_', '').isalpha() + ] + assert names + for name in names: + assert resolve(name.replace('_', ' ')) is not None, name + + def test_the_class_names_the_block_shows_are_the_ones_seeded_as_preferences(self, company, block): + """The vocabulary block and the + `{preferred_entity_classifications}` slot cannot name a class two ways.""" + classes = section(block, CLASSES_HEADING) + for name in company.class_names(): + assert name in classes + + def test_a_declared_label_is_preferred_over_the_local_name(self, block): + assert 'Sports Team' in block + assert 'SportsTeam' not in block + + def test_a_declared_alias_is_offered_to_the_model(self, block): + assert 'Ball Club' in block + assert 'REQ_TO_HC' in block + + def test_an_alias_that_renders_to_the_primary_name_is_not_repeated(self, company): + """It would tell the model nothing. `:Company rdfs:label "Company"` with + `skos:altLabel "Corporation"` shows the alias but not itself twice.""" + block = company.format_as_prompt_constraint('align') + assert block.count('Corporation') == 1 + + def test_a_declared_comment_becomes_the_description(self, block): + assert 'An incorporated commercial organization.' in block + + def test_the_hierarchy_is_shown_by_indentation_parent_first(self, company): + """Alphabetical order would put `Athlete` above `Company` and lose the + structure the model is being shown.""" + classes = section(company.format_as_prompt_constraint('align'), CLASSES_HEADING) + lines = [line for line in classes.splitlines() if line.startswith(' ')] + depth = {line.strip().split(' ')[0]: len(line) - len(line.lstrip()) for line in lines} + + assert depth['Agent'] < depth['Company'] < depth['Sports Team'] + assert depth['Agent'] < depth['Person'] < depth['Athlete'] + + @pytest.mark.parametrize('property_name,rendered_type', [ + ('FOUNDED_YEAR', 'integer'), + ('REVENUE', 'decimal number'), + ('IS_PUBLICLY_TRADED', 'true/false'), + ('INCORPORATED_ON', 'date'), + ('TICKER_SYMBOL', 'text'), + ]) + def test_an_xsd_range_is_named_in_terms_the_model_can_act_on(self, block, property_name, rendered_type): + line = next(line for line in block.splitlines() if property_name in line) + assert rendered_type in line + + def test_no_xsd_iri_is_shown_in_the_prose_format(self, block): + """The prose rendering paraphrases the ranges; showing the IRI as well + reads as an instruction to convert.""" + assert 'XMLSchema#' not in block + + def test_an_unconstrained_domain_or_range_says_so(self, block): + assert 'anything' in section(block, RELATIONSHIPS_HEADING) + assert 'any entity' in section(block, ATTRIBUTES_HEADING) + +class TestTheLevelChangesOnlyTheClosing: + + def test_the_vocabulary_is_byte_identical_at_align_and_strict(self, company): + """How much authority the vocabulary has is not a fact about what the + vocabulary is, so the sections must not move between levels.""" + align = company.format_as_prompt_constraint('align') + strict = company.format_as_prompt_constraint('strict') + + assert align[:align.index(PROTOCOL_HEADING)] == strict[:strict.index(PROTOCOL_HEADING)] + assert align != strict + + def test_strict_tells_the_model_what_happens_to_an_unlisted_name(self, company): + """The claim that makes `strict` honest: the filter does discard them.""" + assert 'discarded' in company.format_as_prompt_constraint('strict') + assert 'discarded' not in company.format_as_prompt_constraint('align') + + def test_the_propositions_block_carries_the_class_names_at_both_levels(self, company): + for level in ('align', 'strict'): + block = company.format_as_proposition_constraint(level) + for name in company.class_names(): + assert name in block + + def test_the_propositions_block_names_no_property(self, company): + """That stage classifies the entities it names and extracts nothing else, + so a property vocabulary there is prompt spent on work it cannot do.""" + block = company.format_as_proposition_constraint('strict') + for name in ('WORKS_FOR', 'FOUNDED_YEAR', 'worksFor', 'foundedYear'): + assert name not in block + +class TestTheTurtleFormat: + """Experimental, and measured a net loss - but it ships, so it has to be sound.""" + + @pytest.fixture(scope='class') + def block(self, company): + return company.format_as_prompt_constraint('align', 'turtle') + + def test_it_shows_the_ontology_source_in_a_fenced_block(self, block): + assert '```turtle' in block + assert 'owl:Class' in block + assert 'rdfs:subClassOf' in block + + def test_it_keeps_the_same_protocol_section_as_the_prose_format(self, company, block): + """How to *use* a vocabulary is not a function of how it was written down, + and holding this constant is what makes the two formats comparable.""" + prose = company.format_as_prompt_constraint('align') + assert section(block, PROTOCOL_HEADING) == section(prose, PROTOCOL_HEADING) + + def test_it_still_maps_each_construct_to_a_channel(self, block): + """The prose layout separates the channels silently; Turtle interleaves + them, so the header has to say it in words.""" + assert 'owl:ObjectProperty' in block + assert 'owl:DatatypeProperty' in block + assert 'entity|RELATIONSHIP|entity' in block + assert 'entity|ATTRIBUTE_NAME|value' in block + + def test_the_level_closing_is_the_same_text_as_the_prose_format(self, company): + for level in ('align', 'strict'): + turtle = company.format_as_prompt_constraint(level, 'turtle') + prose = company.format_as_prompt_constraint(level, 'prose') + assert section(turtle, PROTOCOL_HEADING) == section(prose, PROTOCOL_HEADING) + + def test_the_propositions_block_has_no_turtle_in_it(self, company): + """There is no format argument on that method, by design: the stage renders + classes as a flat list either way, because no property declaration could + steer it.""" + block = company.format_as_proposition_constraint('align') + + assert '```' not in block + assert 'owl:' not in block + +class TestRenderingIsDeterministic: + """Rendered once at configuration time and compared across runs by the + recorded-prompt tests, so ordering may not depend on dict iteration.""" + + @pytest.mark.parametrize('vocabulary_format', ['prose', 'turtle']) + def test_the_same_ontology_renders_the_same_bytes(self, vocabulary_format): + first = Ontology.load(FIXTURES / 'company.ttl') + second = Ontology.load(FIXTURES / 'company.ttl') + + assert first.format_as_prompt_constraint('strict', vocabulary_format) == \ + second.format_as_prompt_constraint('strict', vocabulary_format) + + def test_the_class_names_are_sorted(self, company): + assert company.class_names() == sorted(company.class_names()) + + def test_a_name_the_parser_would_mangle_never_reaches_the_prompt(self, company): + """Underscores in a class name and spaces in a property name both survive + the parser as something else.""" + for name in company.class_names(): + assert '_' not in name + assert resolution_key(name) == resolution_key(name.title()) + +class TestTheTurtleWrapperOnItsOwn: + + def test_empty_turtle_renders_no_block(self): + """Reached when a caller serializes a graph that holds only prefixes.""" + from graphrag_toolkit.lexical_graph.indexing.extract.ontology.prompt_constraint import ( + format_turtle_vocabulary, + ) + assert format_turtle_vocabulary(' \n', 'align') == '' diff --git a/lexical-graph/tests/unit/indexing/extract/test_ontology_prompt_composition.py b/lexical-graph/tests/unit/indexing/extract/test_ontology_prompt_composition.py new file mode 100644 index 00000000..d3b3c2a7 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ontology_prompt_composition.py @@ -0,0 +1,547 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for composing an ontology's vocabulary into the prompts. + +Covers prompt composition: + +* The shipped templates carry no `{ontology_constraints}` placeholder, and + `with_ontology_constraints` returns its argument *itself* when there is nothing + to compose - so a no-ontology render cannot move. That matters because the + `LLMCache` key is a sha256 of the rendered prompt, and one added blank line + would silently invalidate every existing user's `cache/llm/` directory. + Otherwise the block is inserted at a placeholder, at the documented anchor, or + at the end. +* Braces arriving from an ontology survive `str.format`. Nothing asked for this; an `rdfs:comment` containing `{text}` would otherwise either + raise `KeyError` from inside the prompt renderer or quietly substitute the + chunk into the middle of the vocabulary block. +* The composition is applied at the render point in all four extractors, and the + batch pair carries it into `_get_json`'s request body and into the extractor + built by `_run_non_batch_extractor`. + +The extractor tests patch `LLMCache.predict` and read the prompt it was handed, +which is the same string `LLMCache` would have hashed - so they assert against +what the model would actually receive rather than against an intermediate. +""" + +import json +from hashlib import sha256 +from unittest.mock import Mock, patch + +import pytest + +from llama_index.core.llms import MockLLM +from llama_index.core.prompts import PromptTemplate +from llama_index.core.schema import TextNode + +from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY, TOPICS_KEY +from graphrag_toolkit.lexical_graph.indexing.extract.batch_config import BatchConfig +from graphrag_toolkit.lexical_graph.indexing.prompts import ( + EXTRACT_PROPOSITIONS_ANCHOR, + EXTRACT_PROPOSITIONS_PROMPT, + EXTRACT_TOPICS_ANCHOR, + EXTRACT_TOPICS_PROMPT, + ONTOLOGY_CONSTRAINTS_PLACEHOLDER, + with_ontology_constraints, +) +from graphrag_toolkit.lexical_graph.indexing.utils.topic_utils import format_list, format_text +from graphrag_toolkit.lexical_graph.utils import LLMCache +from graphrag_toolkit.lexical_graph.utils.llm_concurrency import shutdown as shutdown_llm_pool + +TEXT = 'Amy Bell works for Example Corp. Example Corp was founded in 1998.' + +VOCABULARY_HEADER = '# Vocabulary for this extraction' +ENTITY_TYPES_HEADER = '# Entity types for this extraction' + +# Stands in for llm.to_json() when comparing one cache key against another. The +# model half of the key is not a property of the prompt. +STUB_LLM_JSON = '{"stub": "llm"}' + +def cache_key_of(rendered, llm_json=STUB_LLM_JSON): + ''' + Reproduce LLMCache.predict's cache key for an already-rendered prompt. + ''' + return sha256(f'{llm_json},{rendered}'.encode('utf-8')).hexdigest() + +def render(template, constraints, **arguments): + ''' + Compose and format a template the way an extractor does. + ''' + return PromptTemplate( + template=with_ontology_constraints(template, constraints) + ).format(**arguments) + +def capture_prompts(): + ''' + Patch LLMCache.predict to record the prompt it was handed, formatted. + + The recorded string is what predict itself would have hashed, so a test can + read the model's view of the prompt without a model. + ''' + prompts = [] + + def predict(self, prompt, **prompt_args): + prompts.append(prompt.format(**prompt_args)) + return '' + + return (prompts, patch.object(LLMCache, 'predict', predict)) + +def mock_llm_cache(): + ''' + An LLMCache over MockLLM: no Bedrock client, no GraphRAGConfig. + ''' + return LLMCache(llm=MockLLM(max_tokens=16), enable_cache=False) + +def batch_config(): + ''' + A BatchConfig with the required fields. Nothing here reaches AWS. + ''' + return BatchConfig( + role_arn='arn:aws:iam::123456789012:role/test-role', + region='us-east-1', + bucket_name='test-bucket' + ) + +@pytest.fixture(autouse=True) +def fresh_pool(): + """ + Driving a real extractor creates the LLM call pool, which is module state - a + pool left behind decides what a later test sees, including tests in other + files that read its size. Same guard as `test_llm_concurrency.py`. + """ + shutdown_llm_pool() + yield + shutdown_llm_pool() + +@pytest.fixture +def company_topic_constraints(company_ontology): + ''' + Fixture for the align-level vocabulary block for the topics prompt. + ''' + return company_ontology.format_as_prompt_constraint('align') + +@pytest.fixture +def company_proposition_constraints(company_ontology): + ''' + Fixture for the align-level entity-type hint for the propositions prompt. + ''' + return company_ontology.format_as_proposition_constraint('align') + + +class TestShippedTemplatesAreUntouched: + """The templates gain no placeholder, and composing nothing changes nothing.""" + + @pytest.mark.parametrize('template', [EXTRACT_TOPICS_PROMPT, EXTRACT_PROPOSITIONS_PROMPT]) + def test_no_placeholder_in_the_shipped_templates(self, template): + """Verify neither shipped template carries an ontology placeholder.""" + assert ONTOLOGY_CONSTRAINTS_PLACEHOLDER not in template + assert 'ontology_constraints' not in template + + @pytest.mark.parametrize('template', [EXTRACT_TOPICS_PROMPT, EXTRACT_PROPOSITIONS_PROMPT]) + def test_composing_nothing_returns_the_template_itself(self, template): + """Verify an empty block leaves the shipped template untouched.""" + assert with_ontology_constraints(template, '') is template + + +class TestEmptyConstraints: + """An empty block is a no-op, on any template.""" + + def test_returns_the_same_object(self): + """Verify identity, not equality: nothing is rebuilt.""" + template = 'a template' + assert with_ontology_constraints(template, '') is template + + def test_a_template_with_a_placeholder_is_left_alone(self): + """Verify an empty block is not substituted into a placeholder. + + Substituting '' would leave the blank line the whole design avoids. + """ + template = f'before\n\n{ONTOLOGY_CONSTRAINTS_PLACEHOLDER}\n\nafter' + assert with_ontology_constraints(template, '') is template + + def test_a_template_with_an_anchor_is_left_alone(self): + """Verify an empty block does not disturb the anchor.""" + template = f'preamble\n\n{EXTRACT_TOPICS_ANCHOR}\n' + assert with_ontology_constraints(template, '') is template + + +class TestPlaceholderInsertion: + """A custom template may choose its own insertion point.""" + + def test_substitutes_at_the_placeholder(self): + """Verify the block replaces the placeholder in place.""" + template = f'before\n\n{ONTOLOGY_CONSTRAINTS_PLACEHOLDER}\n\nafter' + composed = with_ontology_constraints(template, 'BLOCK') + assert composed == 'before\n\nBLOCK\n\nafter' + assert ONTOLOGY_CONSTRAINTS_PLACEHOLDER not in composed + + def test_the_placeholder_wins_over_an_anchor(self): + """Verify an explicit placeholder takes precedence over the anchor.""" + template = ( + f'{ONTOLOGY_CONSTRAINTS_PLACEHOLDER}\n\nmiddle\n\n{EXTRACT_TOPICS_ANCHOR}' + ) + composed = with_ontology_constraints(template, 'BLOCK') + assert composed.startswith('BLOCK\n\nmiddle') + assert composed.endswith(EXTRACT_TOPICS_ANCHOR) + + def test_every_placeholder_is_substituted(self): + """Verify a template repeating the placeholder gets the block at each.""" + template = f'{ONTOLOGY_CONSTRAINTS_PLACEHOLDER}|{ONTOLOGY_CONSTRAINTS_PLACEHOLDER}' + assert with_ontology_constraints(template, 'BLOCK') == 'BLOCK|BLOCK' + + +class TestAnchorInsertion: + """The shipped templates are recognized by their documented anchors.""" + + @pytest.mark.parametrize('template,anchor', [ + (EXTRACT_TOPICS_PROMPT, EXTRACT_TOPICS_ANCHOR), + (EXTRACT_PROPOSITIONS_PROMPT, EXTRACT_PROPOSITIONS_ANCHOR), + ]) + def test_each_shipped_template_contains_its_anchor(self, template, anchor): + """Verify the anchors are real strings in the templates they name.""" + assert template.count(anchor) == 1 + + @pytest.mark.parametrize('template,anchor', [ + (EXTRACT_TOPICS_PROMPT, EXTRACT_TOPICS_ANCHOR), + (EXTRACT_PROPOSITIONS_PROMPT, EXTRACT_PROPOSITIONS_ANCHOR), + ]) + def test_the_block_lands_immediately_before_the_anchor(self, template, anchor): + """Verify the block is inserted ahead of the closing admonition.""" + composed = with_ontology_constraints(template, 'BLOCK') + assert f'BLOCK\n\n{anchor}' in composed + + @pytest.mark.parametrize('template', [EXTRACT_TOPICS_PROMPT, EXTRACT_PROPOSITIONS_PROMPT]) + def test_the_anchor_survives_the_insertion(self, template): + """Verify the instruction the block is placed against is still there.""" + composed = with_ontology_constraints(template, 'BLOCK') + for anchor in [EXTRACT_TOPICS_ANCHOR, EXTRACT_PROPOSITIONS_ANCHOR]: + assert composed.count(anchor) == template.count(anchor) + + def test_the_block_precedes_the_payload_in_the_topics_prompt(self): + """Verify the vocabulary is read before the propositions it applies to.""" + composed = with_ontology_constraints(EXTRACT_TOPICS_PROMPT, 'BLOCK') + assert composed.index('BLOCK') < composed.index('') + + def test_the_block_precedes_the_payload_in_the_propositions_prompt(self): + """Verify the hint is read before the text it applies to.""" + composed = with_ontology_constraints(EXTRACT_PROPOSITIONS_PROMPT, 'BLOCK') + assert composed.index('BLOCK') < composed.index('') + + def test_only_the_first_occurrence_of_an_anchor_is_used(self): + """Verify a template repeating an anchor receives the block once.""" + template = f'{EXTRACT_TOPICS_ANCHOR}\n\nand again: {EXTRACT_TOPICS_ANCHOR}' + composed = with_ontology_constraints(template, 'BLOCK') + assert composed.count('BLOCK') == 1 + assert composed.startswith(f'BLOCK\n\n{EXTRACT_TOPICS_ANCHOR}') + + +class TestAppendInsertion: + """A custom template with neither marker still receives the vocabulary.""" + + def test_appends_when_there_is_no_placeholder_and_no_anchor(self): + """Verify the block is appended rather than dropped.""" + template = 'my own prompt\n\n\n{text}\n\n' + composed = with_ontology_constraints(template, 'BLOCK') + assert composed == f'{template}\n\nBLOCK' + + def test_the_custom_template_is_preserved_verbatim(self): + """Verify nothing in the custom template is rewritten.""" + template = 'my own prompt' + composed = with_ontology_constraints(template, 'BLOCK') + assert composed.startswith(template) + + +class TestBraceHandling: + """A brace from the ontology must not be read as a prompt argument. + + `PromptTemplate.format` is llama-index's `SafeFormatter`, a regex over + `{name}`, not `str.format`: an unknown name is left exactly as written and + nothing is raised. So the case worth guarding is the narrow one - a brace + group that happens to name a prompt argument, which would otherwise pull the + chunk into the middle of the vocabulary block silently. + """ + + def test_a_prompt_argument_name_in_the_block_does_not_pull_in_the_chunk(self): + """Verify a description naming {text} does not receive the chunk.""" + constraints = 'MENTIONS text "as in {text}"' + rendered = render( + EXTRACT_PROPOSITIONS_PROMPT, constraints, text=TEXT, source_info='src' + ) + # Once for the payload, and not a second time inside the block. + assert rendered.count(TEXT) == 1 + assert 'as in { text }' in rendered + + def test_the_padded_name_still_reads_as_what_the_ontology_said(self): + """Verify neutralizing keeps the text legible rather than escaping it.""" + rendered = render( + EXTRACT_PROPOSITIONS_PROMPT, 'HAS_SOURCE text "see {source_info}"', + text=TEXT, source_info='example.txt' + ) + assert 'see { source_info }' in rendered + assert '{{' not in rendered + + def test_prose_braces_are_left_alone(self): + """Verify a JSON example in a description is not rewritten.""" + constraints = 'HAS_CONFIG text "a JSON object, e.g. {"a": 1}"' + rendered = render( + EXTRACT_PROPOSITIONS_PROMPT, constraints, text=TEXT, source_info='src' + ) + assert '{"a": 1}' in rendered + + def test_an_unbalanced_brace_is_left_alone(self): + """Verify a lone brace passes through untouched.""" + rendered = render( + EXTRACT_PROPOSITIONS_PROMPT, 'HAS_SET text "{"', text=TEXT, source_info='src' + ) + assert '"{"' in rendered + + def test_an_unknown_name_is_not_padded_away(self): + """Verify only the substitution hazard is neutralized, not the words.""" + rendered = render( + EXTRACT_PROPOSITIONS_PROMPT, 'HAS_SLOT text "{curly}"', + text=TEXT, source_info='src' + ) + # 'curly' is no prompt argument, so nothing could have substituted it - + # but the padding rule is applied by shape, not by knowing the argument + # names, which composition happens too early to know. + assert '"{ curly }"' in rendered + + +class TestTopicExtractorRenderPoint: + """TopicExtractor composes the block at the point it renders the prompt.""" + + def extractor(self, **kwargs): + from graphrag_toolkit.lexical_graph.indexing.extract.topic_extractor import ( + TopicExtractor, + ) + return TopicExtractor(llm=mock_llm_cache(), num_workers=1, **kwargs) + + def test_the_field_defaults_to_empty(self): + """Verify an extractor with no ontology carries no constraints.""" + assert self.extractor().ontology_constraints == '' + + @pytest.mark.asyncio + async def test_a_no_ontology_prompt_is_the_shipped_prompt(self): + """Verify the prompt with no ontology is the shipped template's.""" + (prompts, patched) = capture_prompts() + with patched: + await self.extractor().aextract([TextNode(text=TEXT, id_='chunk-1')]) + + expected = render( + EXTRACT_TOPICS_PROMPT, '', + text=format_text(TEXT), + preferred_entity_classifications=format_list([]), + preferred_topics=format_list([]), + ) + assert prompts == [expected] + + @pytest.mark.asyncio + async def test_the_vocabulary_reaches_the_prompt(self, company_topic_constraints): + """Verify a configured block is in the prompt the model is handed.""" + (prompts, patched) = capture_prompts() + with patched: + await self.extractor( + ontology_constraints=company_topic_constraints + ).aextract([TextNode(text=TEXT, id_='chunk-1')]) + + assert VOCABULARY_HEADER in prompts[0] + assert 'WORKS_FOR' in prompts[0] + assert f'{company_topic_constraints}\n\n{EXTRACT_TOPICS_ANCHOR}' in prompts[0] + + @pytest.mark.asyncio + async def test_an_ontology_changes_the_cache_key(self, company_topic_constraints): + """Verify the two runs would not share a cached response.""" + (prompts, patched) = capture_prompts() + node = TextNode(text=TEXT, id_='chunk-1') + with patched: + await self.extractor().aextract([node]) + await self.extractor( + ontology_constraints=company_topic_constraints + ).aextract([node]) + + (without, with_ontology) = (cache_key_of(prompt) for prompt in prompts) + assert without != with_ontology + + +class TestLLMPropositionExtractorRenderPoint: + """LLMPropositionExtractor composes the entity-type hint the same way.""" + + def extractor(self, **kwargs): + from graphrag_toolkit.lexical_graph.indexing.extract.llm_proposition_extractor import ( + LLMPropositionExtractor, + ) + return LLMPropositionExtractor(llm=mock_llm_cache(), num_workers=1, **kwargs) + + def test_the_field_defaults_to_empty(self): + """Verify an extractor with no ontology carries no constraints.""" + assert self.extractor().ontology_constraints == '' + + @pytest.mark.asyncio + async def test_a_no_ontology_prompt_is_the_shipped_prompt(self): + """Verify the prompt with no ontology is the shipped template's.""" + (prompts, patched) = capture_prompts() + with patched: + await self.extractor().aextract([TextNode(text=TEXT, id_='chunk-1')]) + + expected = render( + EXTRACT_PROPOSITIONS_PROMPT, '', + text=TEXT, + source_info='', + exclude_cache_keys=['source_info'], + ) + assert prompts == [expected] + + @pytest.mark.asyncio + async def test_the_entity_types_reach_the_prompt(self, company_proposition_constraints): + """Verify the entity-type hint is in the prompt, at the anchor.""" + (prompts, patched) = capture_prompts() + with patched: + await self.extractor( + ontology_constraints=company_proposition_constraints + ).aextract([TextNode(text=TEXT, id_='chunk-1')]) + + assert ENTITY_TYPES_HEADER in prompts[0] + assert f'{company_proposition_constraints}\n\n{EXTRACT_PROPOSITIONS_ANCHOR}' in prompts[0] + + @pytest.mark.asyncio + async def test_an_ontology_changes_the_cache_key(self, company_proposition_constraints): + """Verify the two runs would not share a cached response.""" + (prompts, patched) = capture_prompts() + node = TextNode(text=TEXT, id_='chunk-1') + with patched: + await self.extractor().aextract([node]) + await self.extractor( + ontology_constraints=company_proposition_constraints + ).aextract([node]) + + (without, with_ontology) = (cache_key_of(prompt) for prompt in prompts) + assert without != with_ontology + + +class TestBatchTopicExtractorSync: + """The batch topic path carries the block into the request body.""" + + MODULE = 'graphrag_toolkit.lexical_graph.indexing.extract.batch_topic_extractor_sync' + + def extractor(self, tmp_path, **kwargs): + from graphrag_toolkit.lexical_graph.indexing.extract.batch_topic_extractor_sync import ( + BatchTopicExtractorSync, + ) + return BatchTopicExtractorSync( + batch_config=batch_config(), + llm=mock_llm_cache(), + batch_inference_dir=str(tmp_path / 'batch-topics'), + **kwargs + ) + + def request_body(self, tmp_path, **kwargs): + ''' + Run _get_json with get_request_body reduced to the message text. + + The real get_request_body branches on the Bedrock model id, which has + nothing to do with what this test is about. + ''' + extractor = self.extractor(tmp_path, **kwargs) + with patch( + f'{self.MODULE}.get_request_body', + lambda llm, messages, parameters: [message.content for message in messages] + ): + return extractor._get_json( + TextNode(text=TEXT, id_='chunk-1'), MockLLM(max_tokens=16), {} + ) + + def test_the_field_defaults_to_empty(self, tmp_path): + """Verify the inherited field is there and empty by default.""" + assert self.extractor(tmp_path).ontology_constraints == '' + + def test_the_vocabulary_reaches_the_request_body(self, tmp_path, company_topic_constraints): + """Verify the rendered modelInput contains the vocabulary block.""" + body = self.request_body(tmp_path, ontology_constraints=company_topic_constraints) + rendered = '\n'.join(body['modelInput']) + assert VOCABULARY_HEADER in rendered + assert f'{company_topic_constraints}\n\n{EXTRACT_TOPICS_ANCHOR}' in rendered + + def test_no_ontology_leaves_the_request_body_alone(self, tmp_path): + """Verify the batch path adds nothing when nothing is configured.""" + body = self.request_body(tmp_path) + rendered = '\n'.join(body['modelInput']) + assert VOCABULARY_HEADER not in rendered + assert body['recordId'] == 'chunk-1' + + def test_run_non_batch_extractor_forwards_the_block( + self, tmp_path, company_topic_constraints + ): + """Verify the fallback extractor is built with the same constraints.""" + extractor = self.extractor(tmp_path, ontology_constraints=company_topic_constraints) + inner = Mock() + inner.extract.return_value = [{TOPICS_KEY: {'topics': []}}] + + with patch(f'{self.MODULE}.TopicExtractor', return_value=inner) as constructed: + extractor._run_non_batch_extractor([TextNode(text=TEXT, id_='chunk-1')]) + + assert constructed.call_args.kwargs['ontology_constraints'] == company_topic_constraints + + +class TestBatchLLMPropositionExtractorSync: + """The batch proposition path carries the hint into the request body.""" + + MODULE = 'graphrag_toolkit.lexical_graph.indexing.extract.batch_llm_proposition_extractor_sync' + + def extractor(self, tmp_path, **kwargs): + from graphrag_toolkit.lexical_graph.indexing.extract.batch_llm_proposition_extractor_sync import ( + BatchLLMPropositionExtractorSync, + ) + return BatchLLMPropositionExtractorSync( + batch_config=batch_config(), + llm=mock_llm_cache(), + batch_inference_dir=str(tmp_path / 'batch-propositions'), + **kwargs + ) + + def request_body(self, tmp_path, **kwargs): + ''' + Run _get_json with get_request_body reduced to the message text. + ''' + extractor = self.extractor(tmp_path, **kwargs) + with patch( + f'{self.MODULE}.get_request_body', + lambda llm, messages, parameters: [message.content for message in messages] + ): + return extractor._get_json( + TextNode(text=TEXT, id_='chunk-1'), MockLLM(max_tokens=16), {} + ) + + def test_the_field_defaults_to_empty(self, tmp_path): + """Verify the inherited field is there and empty by default.""" + assert self.extractor(tmp_path).ontology_constraints == '' + + def test_the_entity_types_reach_the_request_body( + self, tmp_path, company_proposition_constraints + ): + """Verify the rendered modelInput contains the entity-type hint.""" + body = self.request_body(tmp_path, ontology_constraints=company_proposition_constraints) + rendered = '\n'.join(body['modelInput']) + assert ENTITY_TYPES_HEADER in rendered + assert f'{company_proposition_constraints}\n\n{EXTRACT_PROPOSITIONS_ANCHOR}' in rendered + + def test_no_ontology_leaves_the_request_body_alone(self, tmp_path): + """Verify the batch path adds nothing when nothing is configured.""" + body = self.request_body(tmp_path) + rendered = '\n'.join(body['modelInput']) + assert ENTITY_TYPES_HEADER not in rendered + assert body['recordId'] == 'chunk-1' + + def test_run_non_batch_extractor_forwards_the_block( + self, tmp_path, company_proposition_constraints + ): + """Verify the fallback extractor is built with the same constraints.""" + extractor = self.extractor( + tmp_path, ontology_constraints=company_proposition_constraints + ) + inner = Mock() + inner.extract.return_value = [{PROPOSITIONS_KEY: []}] + + with patch(f'{self.MODULE}.LLMPropositionExtractor', return_value=inner) as constructed: + extractor._run_non_batch_extractor([TextNode(text=TEXT, id_='chunk-1')]) + + assert constructed.call_args.kwargs['ontology_constraints'] == company_proposition_constraints diff --git a/lexical-graph/tests/unit/indexing/utils/test_pipeline_utils.py b/lexical-graph/tests/unit/indexing/utils/test_pipeline_utils.py index 58a0f50b..f84935cd 100644 --- a/lexical-graph/tests/unit/indexing/utils/test_pipeline_utils.py +++ b/lexical-graph/tests/unit/indexing/utils/test_pipeline_utils.py @@ -1,6 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import logging import multiprocessing import pickle from concurrent.futures import ProcessPoolExecutor @@ -11,12 +12,33 @@ from llama_index.core.ingestion import IngestionPipeline from graphrag_toolkit.lexical_graph import GraphRAGConfig +from graphrag_toolkit.lexical_graph import logging as graphrag_logging from graphrag_toolkit.lexical_graph.indexing.utils.pipeline_utils import ( sink, run_pipeline, node_batcher, _init_worker, ) +from graphrag_toolkit.lexical_graph.logging import ( + get_applied_logging_config, + set_logging_config, +) + + +@pytest.fixture +def restore_logging(): + """`set_logging_config` replaces the root logger's handlers process-wide, and + pytest's capturing sits on those handlers.""" + root = logging.getLogger() + saved_handlers = root.handlers[:] + saved_level = root.level + saved_config = get_applied_logging_config() + + yield + + root.handlers[:] = saved_handlers + root.setLevel(saved_level) + graphrag_logging._applied_logging_config = saved_config def _worker_reads_config(_): @@ -213,11 +235,42 @@ def test_run_pipeline_passes_config_snapshot_to_workers(self, monkeypatch): _, call_kwargs = mock_executor.call_args assert call_kwargs['initializer'] is _init_worker - (snapshot,) = call_kwargs['initargs'] + (snapshot, _logging_config) = call_kwargs['initargs'] assert snapshot.get('_aws_profile') == "scoped-ingest" finally: GraphRAGConfig._aws_profile = orig + def test_run_pipeline_passes_the_logging_config_to_workers(self, restore_logging): + """The second initarg. `logging.config.dictConfig` state is not part of the + GraphRAGConfig snapshot, so without this a worker's root logger stays at + WARNING with no handler but `lastResort` - and extraction components, which + run only in workers, have their INFO logging discarded entirely. + + None here is correct and meaningful: it says the parent never configured + logging, so the worker is left at the interpreter default. + """ + set_logging_config('INFO') + + mock_pipeline = Mock(spec=IngestionPipeline) + mock_pipeline.transformations = [] + mock_pipeline.cache = None + mock_pipeline.disable_cache = True + node_batches = [[TextNode(text="Node 1", id_="1")]] + + with patch('graphrag_toolkit.lexical_graph.indexing.utils.pipeline_utils.run_transformations'): + with patch('graphrag_toolkit.lexical_graph.indexing.utils.pipeline_utils.ProcessPoolExecutor') as mock_executor: + mock_pool = MagicMock() + mock_pool.__enter__.return_value = mock_pool + mock_pool.map.return_value = [node_batches[0]] + mock_executor.return_value = mock_pool + + list(run_pipeline(mock_pipeline, node_batches, num_workers=2)) + + (_snapshot, logging_config) = mock_executor.call_args[1]['initargs'] + + assert logging_config is not None + assert logging_config['loggers']['']['level'] == 'INFO' + def test_run_pipeline_with_cache(self): """Verify run_pipeline uses cache when not disabled.""" mock_cache = Mock() diff --git a/lexical-graph/tests/unit/storage/graph/test_graph_utils.py b/lexical-graph/tests/unit/storage/graph/test_graph_utils.py index 2d9c379f..b94f1c7b 100644 --- a/lexical-graph/tests/unit/storage/graph/test_graph_utils.py +++ b/lexical-graph/tests/unit/storage/graph/test_graph_utils.py @@ -119,6 +119,30 @@ def test_replaces_non_alnum_with_underscore(self): def test_keeps_digits(self): assert relationship_name_from('rev2 of') == 'REV2_OF' + def test_splits_camel_case(self): + """An ontology-normalized predicate arrives authored, not parsed. Without + the split `worksFor` becomes `WORKSFOR`, and the domain-summary prompt + that `GraphSummary._get_paths` builds from this value reads + `(Person)-[WORKSFOR]->(Company)`.""" + assert relationship_name_from('worksFor') == 'WORKS_FOR' + assert relationship_name_from('hasRegisteredAddress') == 'HAS_REGISTERED_ADDRESS' + + def test_the_two_spellings_of_one_predicate_agree(self): + """The same relationship named with and without an ontology has to land on + the same summary-graph name, or turning `normalize_names` on silently + forks the summary graph.""" + assert relationship_name_from('worksFor') == relationship_name_from('WORKS FOR') + + def test_it_does_not_split_runs_of_capitals(self): + """The narrow rule: uppercase after *lowercase* only. Splitting after any + non-uppercase character would break `Company2X`, which is what `.title()` + makes of `Company2x`.""" + assert relationship_name_from('HTTPServer') == 'HTTPSERVER' + assert relationship_name_from('rev2X of') == 'REV2X_OF' + + def test_already_underscored_names_are_unchanged(self): + assert relationship_name_from('WORKS_FOR') == 'WORKS_FOR' + class TestNodeResult: def test_default_star_properties(self): diff --git a/lexical-graph/tests/unit/test_progress_monitor.py b/lexical-graph/tests/unit/test_progress_monitor.py index 4ce63b90..f81c281c 100644 --- a/lexical-graph/tests/unit/test_progress_monitor.py +++ b/lexical-graph/tests/unit/test_progress_monitor.py @@ -130,6 +130,7 @@ def _make_build_pipeline(self, monitor): pipeline.batch_write_size = 10 pipeline.include_domain_labels = False pipeline.include_local_entities = False + pipeline.typed_properties = 'off' pipeline.node_builders = MagicMock() pipeline.node_builders.return_value = [] pipeline.node_filter = MagicMock(return_value=[])