diff --git a/docs-site/astro.config.mjs b/docs-site/astro.config.mjs index d21f0bd6..4a27cdf0 100644 --- a/docs-site/astro.config.mjs +++ b/docs-site/astro.config.mjs @@ -64,6 +64,7 @@ export default defineConfig({ { label: 'Neptune Database', slug: 'lexical-graph/graph-store-neptune-db' }, { label: 'Neo4j', slug: 'lexical-graph/graph-store-neo4j' }, { label: 'FalkorDB', slug: 'lexical-graph/graph-store-falkor-db' }, + { label: 'RDF / SPARQL Stores', slug: 'lexical-graph/graph-store-sparql' }, ], }, { diff --git a/docs-site/src/content/docs/index.mdx b/docs-site/src/content/docs/index.mdx index 7692a772..df5c3292 100644 --- a/docs-site/src/content/docs/index.mdx +++ b/docs-site/src/content/docs/index.mdx @@ -29,7 +29,7 @@ import { Card, CardGrid, Code } from '@astrojs/starlight/components'; Bring your own knowledge graph. Plug an existing graph into a multi-strategy KGQA pipeline without re-extracting anything. - Graph stores: Amazon Neptune (DB and Analytics), Neo4j, FalkorDB. Vector stores: Neptune, OpenSearch, Postgres, S3 Vectors. + Graph stores: Amazon Neptune (DB and Analytics), Neo4j, FalkorDB, RDF/SPARQL stores. Vector stores: Neptune, OpenSearch, Postgres, S3 Vectors. Apache 2.0, developed in the open by AWS Labs. diff --git a/docs-site/src/content/docs/lexical-graph/graph-store-sparql.mdx b/docs-site/src/content/docs/lexical-graph/graph-store-sparql.mdx new file mode 100644 index 00000000..dd8cfb0e --- /dev/null +++ b/docs-site/src/content/docs/lexical-graph/graph-store-sparql.mdx @@ -0,0 +1,152 @@ +--- +title: RDF / SPARQL stores +--- + +The RDF / SPARQL contributor package stores the lexical graph in an existing +SPARQL 1.1 query/update endpoint. Generic standards-compatible endpoints are the +default; Amazon Neptune IAM authentication is an optional transport. + +The adapter executes native SPARQL. Builders and retrievers provide a +backend-neutral operation identifier and structured parameters. Property-graph +stores execute their native queries, while the RDF store selects the +corresponding native SPARQL operation. + +### Install and register + +Install the contributor package from a repository checkout: + +```bash +pip install ./lexical-graph-contrib/sparql +``` + +Register its factory before creating the graph store: + +```python +from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql import ( + SPARQLGraphStoreFactory, +) + +GraphStoreFactory.register(SPARQLGraphStoreFactory) +``` + +Connect the graph store to an existing repository or dataset. + +### Generic endpoint + +Provide the query URL and, if different, the update URL: + +```python +graph_store = GraphStoreFactory.for_graph_store( + 'sparql+https://rdf.example.com/query', + update_endpoint='https://rdf.example.com/update', +) +``` + +Supported schemes are: + +| Scheme | Endpoint transport | +|---|---| +| `sparql://host/path` | HTTP | +| `sparql+http://host/path` | HTTP | +| `sparql+s://host/path` | HTTPS | +| `sparql+https://host/path` | HTTPS | +| `sparql+neptune://host:8182` | HTTPS with Neptune IAM SigV4 | + +If `update_endpoint` is omitted, the query URL is used for both operations. + +HTTP Basic credentials can be passed in the URL, through `username` and +`password`, or through the `SPARQL_USER` and `SPARQL_PASSWORD` environment +variables. Custom headers support bearer tokens and endpoint-specific headers: + +```python +graph_store = GraphStoreFactory.for_graph_store( + 'sparql+https://rdf.example.com/query', + headers={'Authorization': 'Bearer token'}, +) +``` + +### Amazon Neptune IAM + +Install the optional botocore dependency: + +```bash +pip install './lexical-graph-contrib/sparql[neptune]' +``` + +Then use the Neptune scheme and AWS region: + +```python +graph_store = GraphStoreFactory.for_graph_store( + 'sparql+neptune://my-cluster.cluster-abcdefghijkl.eu-central-1.neptune.amazonaws.com:8182', + region_name='eu-central-1', +) +``` + +Each request is signed for the `neptune-db` service using credentials from +botocore's standard provider chain. Credentials are resolved on every request, +allowing botocore to refresh temporary role, web-identity, IAM Identity Center, +ECS, and EC2 credentials. IAM transport requires HTTPS. + +The same transport is available as a plain RDFLib graph: + +```python +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.neptune_iam import ( + neptune_iam_graph, +) + +graph = neptune_iam_graph( + 'https://my-cluster.cluster-abcdefghijkl.eu-central-1.neptune.amazonaws.com:8182', + region_name='eu-central-1', +) +try: + rows = graph.query('SELECT * WHERE { ?s ?p ?o } LIMIT 10') +finally: + graph.close() +``` + +### Namespaces and tenants + +The default schema namespace is +`https://awslabs.github.io/graphrag-toolkit/lexical#`; the default instance +namespace is `https://awslabs.github.io/graphrag-toolkit/lexical/`. + +Customize them only when creating a new dataset: + +```python +graph_store = GraphStoreFactory.for_graph_store( + 'sparql+https://rdf.example.com/query', + update_endpoint='https://rdf.example.com/update', + lexical_prefix='gt', + lexical_schema_namespace='https://example.com/graph/schema#', + lexical_instance_namespace='https://example.com/graph/data/', +) +``` + +Every tenant, including the default tenant, uses a deterministic named graph. +Writes target that graph explicitly, and reads select it with the standard +SPARQL Protocol `default-graph-uri` parameter. + +### RDF model + +Sources, chunks, topics, statements, facts, and entities retain their existing +lexical-graph identities. A node becomes a deterministic IRI with an RDF class +and `lg:id`; node properties become literal-valued predicates. Simple edges +become RDF predicates. + +Extracted facts record their subject, predicate, and object directly: + +```turtle + a lg:Fact ; + lg:subject ; + lg:predicate ; + lg:object ; + lg:supports ; + lg:value "Amazon PRODUCES EC2" . + + a lg:Relation ; + lg:value "PRODUCES" . +``` + +This is the toolkit's fact model, with the fact resource carrying its subject, +predicate, object, and supporting statement. diff --git a/docs-site/src/content/docs/lexical-graph/overview.mdx b/docs-site/src/content/docs/lexical-graph/overview.mdx index 665668a7..8ff2fc49 100644 --- a/docs-site/src/content/docs/lexical-graph/overview.mdx +++ b/docs-site/src/content/docs/lexical-graph/overview.mdx @@ -42,7 +42,7 @@ The graphrag-toolkit [lexical-graph](https://github.com/awslabs/graphrag-toolkit Source → chunk → topic → statement → fact → entity, all linked. Retrieval can hop between any of these levels. - Graph: Amazon Neptune (DB and Analytics), Neo4j, FalkorDB. Vectors: Neptune, OpenSearch, Postgres, S3 Vectors. + Graph: Amazon Neptune (DB and Analytics), Neo4j, FalkorDB, RDF/SPARQL stores. Vectors: Neptune, OpenSearch, Postgres, S3 Vectors. Extract and build run as separate micro-batched pipelines so ingest is continuous and resumable. diff --git a/docs-site/src/content/docs/lexical-graph/storage-model.mdx b/docs-site/src/content/docs/lexical-graph/storage-model.mdx index ec286481..aa7a3cbf 100644 --- a/docs-site/src/content/docs/lexical-graph/storage-model.mdx +++ b/docs-site/src/content/docs/lexical-graph/storage-model.mdx @@ -21,7 +21,7 @@ Graph stores and vector stores provide connectivity to an *existing* storage ins ### Graph store -Graph stores must support the [openCypher](https://opencypher.org/) property graph query language. Graph construction queries typically use an `UNWIND ... MERGE` idiom to create or update the graph for a [batch of inputs](https://docs.aws.amazon.com/neptune-analytics/latest/userguide/best-practices-content.html#best-practices-content-14). The Neptune graph store implementations override the `GraphStore.node_id()` method to ensure that node ids in the code (e.g. `chunkId`) are mapped to Neptune's `~id` reserved property. Alternative graph store implementations can leave the base implementation of `node_id()` as-is. This will result in node ids being mapped to a property of the same name (i.e. a reference to `chunkId` in the code will be mapped to a `chunkId` property of a node). +The built-in property-graph stores use [openCypher](https://opencypher.org/). Graph construction queries typically use an `UNWIND ... MERGE` idiom to create or update the graph for a [batch of inputs](https://docs.aws.amazon.com/neptune-analytics/latest/userguide/best-practices-content.html#best-practices-content-14). Graph builders and retrievers also identify their operation explicitly, allowing stores such as the RDF / SPARQL contributor to execute a native implementation. The Neptune graph store implementations override the `GraphStore.node_id()` method to ensure that node ids in the code (e.g. `chunkId`) are mapped to Neptune's `~id` reserved property. Alternative graph store implementations can leave the base implementation of `node_id()` as-is. This will result in node ids being mapped to a property of the same name (i.e. a reference to `chunkId` in the code will be mapped to a `chunkId` property of a node). You use the `GraphStoreFactory.for_graph_store()` static factory method to create a graph store. @@ -30,6 +30,7 @@ The lexical-graph supports the following graph databases: - [Amazon Neptune](/graphrag-toolkit/lexical-graph/graph-store-neptune-db/) - [Amazon Neptune Analytics](/graphrag-toolkit/lexical-graph/graph-store-neptune-analytics/) - [Neo4j](/graphrag-toolkit/lexical-graph/graph-store-neo4j/) + - [RDF / SPARQL Stores](/graphrag-toolkit/lexical-graph/graph-store-sparql/) #### Logging graph queries diff --git a/examples/lexical-graph-local-dev/docker/docker-compose-dev.yml b/examples/lexical-graph-local-dev/docker/docker-compose-dev.yml index c8f9435a..391f95c2 100644 --- a/examples/lexical-graph-local-dev/docker/docker-compose-dev.yml +++ b/examples/lexical-graph-local-dev/docker/docker-compose-dev.yml @@ -1,4 +1,4 @@ -name: local-dev +name: graphrag-toolkit-rdf-dev services: neo4j-local: image: neo4j:5.25-community diff --git a/examples/lexical-graph-local-dev/docker/jupyter/Dockerfile.dev b/examples/lexical-graph-local-dev/docker/jupyter/Dockerfile.dev index ac1a6a1f..63987931 100644 --- a/examples/lexical-graph-local-dev/docker/jupyter/Dockerfile.dev +++ b/examples/lexical-graph-local-dev/docker/jupyter/Dockerfile.dev @@ -32,6 +32,7 @@ RUN pip install --no-cache-dir \ plotly # LlamaIndex readers (hard imports in lexical-graph source) +USER root RUN pip install --no-cache-dir \ llama-index-readers-web \ llama-index-readers-file \ @@ -42,6 +43,8 @@ RUN pip install --no-cache-dir \ pymupdf # Build tools for packages requiring C compilation (e.g. lru-dict) -USER root -RUN apt-get update && apt-get install -y --no-install-recommends build-essential && rm -rf /var/lib/apt/lists/* +RUN apt-get update && apt-get install -y --no-install-recommends build-essential && \ + rm -rf /var/lib/apt/lists/* && \ + fix-permissions "${CONDA_DIR}" && \ + fix-permissions "/home/${NB_USER}" USER jovyan diff --git a/lexical-graph-contrib/sparql/README.md b/lexical-graph-contrib/sparql/README.md new file mode 100644 index 00000000..01c01ba2 --- /dev/null +++ b/lexical-graph-contrib/sparql/README.md @@ -0,0 +1,170 @@ +# RDF / SPARQL graph store + +This contributor package stores the AWS GraphRAG Toolkit lexical graph in an +existing SPARQL 1.1 query/update endpoint. It is endpoint-neutral by default and +has an optional Amazon Neptune IAM transport. + +Core graph builders and retrievers identify a backend-neutral `GraphOperation`. +Property-graph stores execute their native queries, while this package provides +the native SPARQL implementation of each operation. + +The update renderer batches related statements into a single request. Calling +RDFLib `Graph.add()` and `Graph.remove()` for every remote triple would be much +less efficient and would make multi-statement counters non-atomic. + +## Install + +From a repository checkout: + +```bash +pip install ./lexical-graph-contrib/sparql +``` + +Include the optional botocore dependency only when Neptune IAM authentication is +required: + +```bash +pip install './lexical-graph-contrib/sparql[neptune]' +``` + +## Generic SPARQL endpoint + +Register the contributor factory once before resolving a graph store: + +```python +from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql import ( + SPARQLGraphStoreFactory, +) + +GraphStoreFactory.register(SPARQLGraphStoreFactory) + +graph_store = GraphStoreFactory.for_graph_store( + 'sparql+https://rdf.example.com/query', + update_endpoint='https://rdf.example.com/update', +) +``` + +Connect the graph store to an existing repository or dataset. + +Supported connection schemes are: + +| Scheme | HTTP endpoint | +|---|---| +| `sparql://host/path` | `http://host/path` | +| `sparql+http://host/path` | `http://host/path` | +| `sparql+s://host/path` | `https://host/path` | +| `sparql+https://host/path` | `https://host/path` | +| `sparql+neptune://host:8182` | `https://host:8182/sparql`, signed with IAM | + +If `update_endpoint` is omitted, the query endpoint is used for both operations. +It can also be supplied as an encoded `update_endpoint` query parameter. + +### Generic authentication + +HTTP Basic credentials can be supplied in the connection URL, as keyword +arguments, or through `SPARQL_USER` and `SPARQL_PASSWORD`: + +```python +graph_store = GraphStoreFactory.for_graph_store( + 'sparql+https://rdf.example.com/query', + username='service-user', + password='secret', + headers={'X-Request-Origin': 'graphrag-toolkit'}, +) +``` + +Use `headers={'Authorization': 'Bearer ...'}` for bearer-token endpoints. Avoid +putting credentials directly in a URL when the URL may be logged by surrounding +application code. + +## Amazon Neptune with IAM + +The Neptune scheme changes only the transport. The RDF model, operations, and +SPARQL implementation remain the same as for every other endpoint. + +```python +graph_store = GraphStoreFactory.for_graph_store( + 'sparql+neptune://my-cluster.cluster-abcdefghijkl.eu-central-1.neptune.amazonaws.com:8182', + region_name='eu-central-1', +) +``` + +`NeptuneIAMAuth` uses botocore's standard credential provider chain and obtains +credentials for every request. This allows botocore to refresh temporary role, +web-identity, IAM Identity Center, ECS, and EC2 credentials. Requests are signed +for the `neptune-db` service with SigV4 and require HTTPS. + +The transport can also be used as a plain RDFLib graph: + +```python +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.neptune_iam import ( + neptune_iam_graph, +) + +graph = neptune_iam_graph( + 'https://my-cluster.cluster-abcdefghijkl.eu-central-1.neptune.amazonaws.com:8182', + region_name='eu-central-1', +) +try: + rows = graph.query('SELECT * WHERE { ?s ?p ?o } LIMIT 10') +finally: + graph.close() +``` + +## Namespaces and tenant isolation + +The default schema namespace is +`https://awslabs.github.io/graphrag-toolkit/lexical#`; the default instance +namespace is `https://awslabs.github.io/graphrag-toolkit/lexical/`. + +They can be changed when creating the store: + +```python +graph_store = GraphStoreFactory.for_graph_store( + 'sparql+https://rdf.example.com/query', + update_endpoint='https://rdf.example.com/update', + lexical_prefix='gt', + lexical_schema_namespace='https://example.com/graph/schema#', + lexical_instance_namespace='https://example.com/graph/data/', +) +``` + +Changing namespaces changes the IRIs written for new data. Use one configuration +consistently for a repository. + +Every tenant, including the default tenant, uses a deterministic named graph +below the instance namespace. Writes use `GRAPH` and reads pass the tenant graph +as the SPARQL Protocol `default-graph-uri`, keeping tenant selection independent +of endpoint-specific union-default-graph behavior. + +## RDF model + +Lexical nodes receive deterministic IRIs, an RDF class, and the original id as +`lg:id`. Simple edges become RDF predicates. Predicates whose property-graph +name is ambiguous are specialized by domain, for example +`lg:statementMentionedIn` versus `lg:topicMentionedIn`. + +Extracted facts are represented as statement resources: + +```turtle + a lg:Fact ; + lg:subject ; + lg:predicate ; + lg:object ; + lg:supports ; + lg:value "Amazon PRODUCES EC2" . + + a lg:Relation ; + lg:value "PRODUCES" . +``` + +This is the toolkit's fact model, with the fact resource carrying its subject, +predicate, object, and supporting statement. + +## Tests + +```bash +pytest lexical-graph-contrib/sparql/tests +pytest lexical-graph/tests +``` diff --git a/lexical-graph-contrib/sparql/pyproject.toml b/lexical-graph-contrib/sparql/pyproject.toml new file mode 100644 index 00000000..67cda5d0 --- /dev/null +++ b/lexical-graph-contrib/sparql/pyproject.toml @@ -0,0 +1,22 @@ +[build-system] +requires = ["hatchling", "hatch-requirements-txt"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/graphrag_toolkit_contrib"] + + +[project] +name = "graphrag-toolkit-lexical-graph-sparql" +version = "0.1.0" +description = "RDF/SPARQL endpoint support for the AWS GraphRAG Toolkit lexical graph" +readme = "README.md" +requires-python = ">=3.10" +dynamic = ["dependencies"] +license = "Apache-2.0" + +[tool.hatch.metadata.hooks.requirements_txt] +files = ["src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/requirements.txt"] + +[project.optional-dependencies] +neptune = ["botocore[crt]>=1.40.61,<2"] diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/__init__.py b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/__init__.py new file mode 100644 index 00000000..c33a9d7a --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/__init__.py @@ -0,0 +1,5 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .sparql_graph_store_factory import SPARQLGraphStoreFactory +from .sparql_graph_store import SPARQLDatabaseClient diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/neptune_iam.py b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/neptune_iam.py new file mode 100644 index 00000000..186c52f7 --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/neptune_iam.py @@ -0,0 +1,84 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.exceptions import NoCredentialsError +from botocore.session import Session +from rdflib import Graph +from requests.auth import AuthBase + +from .sparql_endpoint_client import RDFLibHTTPStore + + +class NeptuneIAMAuth(AuthBase): + """SigV4-sign each request with the current botocore credentials. + + Calling ``get_credentials`` for every request is intentional: botocore can + refresh temporary credentials supplied by roles, web identity, IAM Identity + Center, and container credential providers before they expire. + """ + + def __init__(self, region_name=None, aws_session=None): + self._session = aws_session or Session() + configured_region = getattr( + self._session, 'get_config_variable', lambda _: None + )('region') + self._region = ( + region_name + or getattr(self._session, 'region_name', None) + or configured_region + ) + if not self._region: + raise ValueError('AWS region is required for Neptune IAM authentication') + + def __call__(self, request): + credentials = self._session.get_credentials() + if credentials is None: + raise NoCredentialsError() + aws_request = AWSRequest( + method=request.method, + url=request.url, + data=request.body, + headers=dict(request.headers), + ) + SigV4Auth( + credentials.get_frozen_credentials(), 'neptune-db', self._region + ).add_auth(aws_request) + request.headers.update(dict(aws_request.headers.items())) + return request + + +class NeptuneIAMStore(RDFLibHTTPStore): + """Optional IAM transport for an Amazon Neptune SPARQL endpoint.""" + + def __init__(self, + endpoint_url, + region_name=None, + aws_session=None, + headers=None, + timeout=60.0): + endpoint = endpoint_url.rstrip('/') + if not endpoint.startswith('https://'): + raise ValueError('Neptune IAM requires an HTTPS endpoint') + if not endpoint.endswith('/sparql'): + endpoint += '/sparql' + super().__init__( + endpoint, + auth=NeptuneIAMAuth(region_name, aws_session), + headers=headers, + timeout=timeout, + ) + + +def neptune_iam_graph(endpoint_url, + region_name=None, + aws_session=None, + headers=None, + timeout=60.0): + """Create an RDFLib Graph backed by an IAM-enabled Neptune database.""" + return Graph( + store=NeptuneIAMStore( + endpoint_url, region_name, aws_session, headers, timeout, + ) + ) diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/ontology.py b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/ontology.py new file mode 100644 index 00000000..4ab00e57 --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/ontology.py @@ -0,0 +1,148 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import re +from dataclasses import dataclass, field +from typing import Mapping, Optional +from urllib.parse import quote, urlparse + +from rdflib import Literal, URIRef + +LEXICAL_SCHEMA = 'https://awslabs.github.io/graphrag-toolkit/lexical#' +LEXICAL_BASE = 'https://awslabs.github.io/graphrag-toolkit/lexical/' +LEXICAL_PREFIX = 'lg' + +RDF_TYPE = '' + +_PREFIX_RE = re.compile(r'^[A-Za-z_][A-Za-z0-9_-]*$') +_UNSAFE_IRI_RE = re.compile(r'[\x00-\x20<>"{}|^`\\]') + + +@dataclass(frozen=True) +class NamespaceConfig: + """Namespaces used when rendering lexical-graph RDF and SPARQL.""" + + prefix: str = LEXICAL_PREFIX + schema_namespace: str = LEXICAL_SCHEMA + instance_namespace: str = LEXICAL_BASE + extra_prefixes: Mapping[str, str] = field(default_factory=dict) + + def __post_init__(self): + schema = _namespace_with_separator(self.schema_namespace) + instance = _namespace_with_separator(self.instance_namespace, separator='/') + + if not _PREFIX_RE.match(self.prefix): + raise ValueError(f'Invalid SPARQL prefix name: {self.prefix!r}') + for prefix, namespace in self.extra_prefixes.items(): + if not _PREFIX_RE.match(prefix): + raise ValueError(f'Invalid SPARQL prefix name: {prefix!r}') + if prefix == self.prefix and namespace != schema: + raise ValueError( + f'Extra prefix {prefix!r} conflicts with lexical_schema_namespace' + ) + + for iri in (schema, instance, *self.extra_prefixes.values()): + if _UNSAFE_IRI_RE.search(iri): + raise ValueError(f'Invalid namespace IRI (unsafe characters): {iri!r}') + if not urlparse(iri).scheme: + raise ValueError(f'Namespace IRI must be absolute: {iri!r}') + + object.__setattr__(self, 'schema_namespace', schema) + object.__setattr__(self, 'instance_namespace', instance) + + @property + def prefix_ref(self) -> str: + return f'{self.prefix}:' + + def term(self, local_name: str) -> str: + return URIRef(f'{self.schema_namespace}{local_name}').n3() + + def instance_iri(self, kind: str, id_value) -> str: + value = quote(str(id_value), safe='') + return URIRef(f'{self.instance_namespace}{kind}/{value}').n3() + + def tenant_graph_iri(self, tenant_value) -> Optional[str]: + if not tenant_value: + return None + value = quote(str(tenant_value), safe='') + return URIRef(f'{self.instance_namespace}tenant/{value}').n3() + + def sparql_prefixes(self) -> str: + prefixes = [(self.prefix, self.schema_namespace)] + prefixes.extend( + (prefix, namespace) + for prefix, namespace in sorted(self.extra_prefixes.items()) + if prefix != self.prefix + ) + return '\n'.join(f'PREFIX {prefix}: <{namespace}>' for prefix, namespace in prefixes) + + +def _namespace_with_separator(namespace: str, separator: str = '#') -> str: + if namespace.endswith(('#', '/')): + return namespace + return f'{namespace}{separator}' + + +DEFAULT_NAMESPACE = NamespaceConfig() + +ID_KEY_TO_KIND = { + 'sourceId': ('source', 'Source'), + 'chunkId': ('chunk', 'Chunk'), + 'topicId': ('topic', 'Topic'), + 'statementId': ('statement', 'Statement'), + 'factId': ('fact', 'Fact'), + 'entityId': ('entity', 'Entity'), + 'sysClassId': ('sysclass', 'SysClass'), +} + +def term(local_name, namespace: Optional[NamespaceConfig] = None): + """Return a schema IRI in angle-bracket form.""" + return (namespace or DEFAULT_NAMESPACE).term(local_name) + + +def instance_iri(kind, id_value, namespace: Optional[NamespaceConfig] = None): + """Return a deterministic instance IRI for a node of the given kind. + + The id is percent-encoded so values such as ``aws::abc:def`` are legal IRIs. + """ + return (namespace or DEFAULT_NAMESPACE).instance_iri(kind, id_value) + + +def relation_iri(predicate_value, namespace: Optional[NamespaceConfig] = None): + """IRI for a shared predicate/relation resource, merged by normalised + (case-insensitive, space-insensitive) predicate value. + + So all facts with predicate "USES"/"uses" reference one lg:Relation node, + the same way entities are merged by their normalised value. + """ + key = str(predicate_value).lower().replace(' ', '_') + return instance_iri('relation', key, namespace) + + +def sys_relation_iri(subject_class_id, + predicate, + object_class_id, + namespace: Optional[NamespaceConfig] = None): + """Deterministic IRI for a sys-class relation node (edge metadata).""" + digest = hashlib.md5( + f'{subject_class_id}|{predicate}|{object_class_id}'.encode('utf-8'), + usedforsecurity=False, + ).hexdigest() + return instance_iri('sysrel', digest, namespace) + + +def tenant_graph_iri(tenant_value, namespace: Optional[NamespaceConfig] = None): + """Return the named-graph IRI for a tenant value.""" + return (namespace or DEFAULT_NAMESPACE).tenant_graph_iri(tenant_value) + + +def sparql_literal(value): + """Render a Python value as a SPARQL literal, or ``None`` to skip it. + + ``Literal.n3`` applies RDF escaping and datatype selection consistently + with endpoint result parsing. + """ + if value is None: + return None + return Literal(value).n3() diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/requirements.txt b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/requirements.txt new file mode 100644 index 00000000..4df620f5 --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/requirements.txt @@ -0,0 +1,2 @@ +rdflib>=7.3.0,<8 +requests>=2.31,<3 diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_endpoint_client.py b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_endpoint_client.py new file mode 100644 index 00000000..552eefa4 --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_endpoint_client.py @@ -0,0 +1,145 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from io import BytesIO +from typing import Any, Dict, List, Optional + +import requests +from rdflib import Graph +from rdflib.plugins.stores.sparqlstore import SPARQLUpdateStore +from rdflib.query import Result +from requests.auth import HTTPBasicAuth + +SPARQL_JSON = 'application/sparql-results+json' +FORM_URLENCODED = 'application/x-www-form-urlencoded' +_ACCEPT = f'{SPARQL_JSON}, text/turtle' +_SUPPORTED_QUERY_TYPES = frozenset({'ASK', 'SELECT'}) + + +class RDFLibHTTPStore(SPARQLUpdateStore): + """RDFLib remote store with a reusable, configurable HTTP session. + + The standard ``SPARQLUpdateStore`` provides query parsing and result + handling. This subclass configures its HTTP calls with ``requests`` + authentication, headers, and timeouts. + """ + + def __init__(self, + query_endpoint: str, + update_endpoint: Optional[str] = None, + auth=None, + headers: Optional[Dict[str, str]] = None, + timeout: float = 60.0): + super().__init__( + query_endpoint, + update_endpoint or query_endpoint, + context_aware=False, + returnFormat='json', + ) + if timeout <= 0: + raise ValueError('SPARQL endpoint timeout must be greater than zero') + self._headers = dict(headers or {}) + self._timeout = timeout + self._http = requests.Session() + self._http.auth = auth + + def _query(self, query, default_graph=None, named_graph=None): + self._queries += 1 + data = {'query': query} + if default_graph is not None: + data['default-graph-uri'] = default_graph + if named_graph is not None: + data['named-graph-uri'] = named_graph + response = self._post(self.query_endpoint, data, {'Accept': _ACCEPT}) + content_type = response.headers.get('Content-Type', SPARQL_JSON).split(';')[0] + return Result.parse(BytesIO(response.content), content_type=content_type) + + def query_with_default_graph(self, query: str, graph_uri: str): + """Execute against a protocol-defined default graph. + + RDFLib only forwards ``queryGraph`` for context-aware stores. This + transport deliberately is not context-aware, so tenant isolation uses + the standard SPARQL Protocol ``default-graph-uri`` parameter directly. + """ + return self._query(query, default_graph=graph_uri) + + def _update(self, update): + self._updates += 1 + self._post(self.update_endpoint, {'update': update}) + + def _post(self, endpoint, data, headers=None): + response = self._http.post( + endpoint, + data=data, + headers={ + **self._headers, + 'Content-Type': FORM_URLENCODED, + **(headers or {}), + }, + timeout=self._timeout, + ) + if response.status_code >= 400: + raise RuntimeError( + f'SPARQL endpoint request failed [status: {response.status_code}, ' + f'body: {response.text.strip()[:500]}]' + ) + return response + + def close(self, commit_pending_transaction=False): + if commit_pending_transaction: + self.commit() + self._http.close() + + +class SPARQLEndpointClient: + """Small SELECT/ASK and UPDATE facade over an RDFLib ``Graph``.""" + + def __init__(self, + query_endpoint: Optional[str] = None, + update_endpoint: Optional[str] = None, + username: Optional[str] = None, + password: Optional[str] = None, + headers: Optional[Dict[str, str]] = None, + timeout: float = 60.0, + store: Optional[SPARQLUpdateStore] = None): + if store is None: + if query_endpoint is None: + raise ValueError('SPARQL query endpoint is required') + auth = HTTPBasicAuth(username, password) if username is not None else None + store = RDFLibHTTPStore( + query_endpoint, update_endpoint, auth, headers, timeout, + ) + self._graph = Graph(store=store) + + @property + def store(self): + return self._graph.store + + def query(self, + sparql: str, + default_graph: Optional[str] = None) -> List[Dict[str, Any]]: + if default_graph: + result = self.store.query_with_default_graph(sparql, default_graph) + else: + result = self._graph.query(sparql) + if result.type not in _SUPPORTED_QUERY_TYPES: + raise ValueError( + f'Only SELECT and ASK queries are supported, got {result.type!r}' + ) + if result.type == 'ASK': + return [{'boolean': bool(result.askAnswer)}] + return [ + {str(var): self._coerce(value) for var, value in row.asdict().items()} + for row in result + ] + + def update(self, sparql: str) -> None: + self._graph.update(sparql) + + def close(self) -> None: + self._graph.close() + + @staticmethod + def _coerce(value: Any) -> Any: + value = value.toPython() if hasattr(value, 'toPython') else value + return str(value) if hasattr(value, 'n3') else value diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_graph_store.py b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_graph_store.py new file mode 100644 index 00000000..2a559c55 --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_graph_store.py @@ -0,0 +1,201 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import time +import uuid +from typing import Any, Dict, List, Optional + +from llama_index.core.bridge.pydantic import Field, PrivateAttr, SecretStr + +from graphrag_toolkit.lexical_graph.storage.graph import ( + GraphOperation, + GraphStore, + NodeId, + format_id, +) + +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.ontology import ( + LEXICAL_BASE, + LEXICAL_PREFIX, + LEXICAL_SCHEMA, + NamespaceConfig, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_endpoint_client import ( + SPARQLEndpointClient, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_queries import ( + QUERY_OPERATIONS, + run_query, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_updates import ( + UPDATE_OPERATIONS, + render_update, +) + +logger = logging.getLogger(__name__) + + +class SPARQLDatabaseClient(GraphStore): + """Lexical graph store for standards-compatible SPARQL endpoints. + + Backend-neutral ``GraphOperation`` values select native SPARQL read and + update implementations, which execute through the configured endpoint + transport. + """ + + query_endpoint: str + update_endpoint: Optional[str] = None + username: Optional[str] = None + password: Optional[SecretStr] = None + timeout: float = 60.0 + headers: Dict[str, str] = Field(default_factory=dict) + neptune_iam: bool = False + region_name: Optional[str] = None + lexical_prefix: str = LEXICAL_PREFIX + lexical_schema_namespace: str = LEXICAL_SCHEMA + lexical_instance_namespace: str = LEXICAL_BASE + sparql_prefixes: Dict[str, str] = Field(default_factory=dict) + + _client: Optional[Any] = PrivateAttr(default=None) + _namespace: Optional[NamespaceConfig] = PrivateAttr(default=None) + + def __init__(self, + query_endpoint: str, + update_endpoint: Optional[str] = None, + username: Optional[str] = None, + password: Optional[str] = None, + timeout: float = 60.0, + headers: Optional[Dict[str, str]] = None, + **kwargs) -> None: + super().__init__( + query_endpoint=query_endpoint, + update_endpoint=update_endpoint, + username=username, + password=password, + timeout=timeout, + headers=headers or {}, + **kwargs, + ) + + def __getstate__(self): + self._client = None + return super().__getstate__() + + @property + def client(self) -> SPARQLEndpointClient: + if self._client is None: + if self.neptune_iam: + from .neptune_iam import NeptuneIAMStore + + self._client = SPARQLEndpointClient(store=NeptuneIAMStore( + self.query_endpoint, + region_name=self.region_name, + headers=dict(self.headers), + timeout=self.timeout, + )) + else: + self._client = SPARQLEndpointClient( + query_endpoint=self.query_endpoint, + update_endpoint=self.update_endpoint, + username=self.username, + password=self.password.get_secret_value() if self.password is not None else None, + timeout=self.timeout, + headers=dict(self.headers), + ) + return self._client + + @property + def namespace(self) -> NamespaceConfig: + if self._namespace is None: + self._namespace = NamespaceConfig( + prefix=self.lexical_prefix, + schema_namespace=self.lexical_schema_namespace, + instance_namespace=self.lexical_instance_namespace, + extra_prefixes=dict(self.sparql_prefixes), + ) + return self._namespace + + def node_id(self, id_name: str) -> NodeId: + return format_id(id_name) + + def _execute_query(self, + sparql: str, + parameters: Optional[dict] = None, + correlation_id: Any = None) -> List[Any]: + """Execute caller-supplied native SPARQL.""" + parameters = parameters or {} + if parameters: + raise ValueError( + 'Raw SPARQL parameter binding is not supported. Serialize values ' + 'with RDFLib terms, or use a typed GraphOperation.' + ) + + query_id = uuid.uuid4().hex[:5] + log_entry = self.log_formatting.format_log_entry( + self._logging_prefix(query_id, correlation_id), sparql, parameters, + ) + logger.debug(f'[{log_entry.query_ref}] Query: [query: {log_entry.query}, ' + f'parameters: {log_entry.parameters}]') + + start = time.time() + results = self.client.query(sparql) + + if logger.isEnabledFor(logging.DEBUG): + elapsed = int((time.time() - start) * 1000) + logger.debug(f'[{log_entry.query_ref}] {elapsed}ms -> {len(results)} row(s)') + + return results + + def _execute_operation(self, + operation: GraphOperation, + query: str, + parameters: Dict[str, Any], + correlation_id=None, + **kwargs) -> List[Any]: + """Execute the native SPARQL implementation of a semantic operation.""" + del query # Retained in the method signature for GraphStore compatibility. + query_id = uuid.uuid4().hex[:5] + log_entry = self.log_formatting.format_log_entry( + self._logging_prefix(query_id, correlation_id), operation.value, parameters, + ) + logger.debug(f'[{log_entry.query_ref}] Operation: [{operation.value}, ' + f'parameters: {log_entry.parameters}]') + + start = time.time() + tenant_id = kwargs.get('tenant_id') or str(self.tenant_id) + if operation in UPDATE_OPERATIONS: + update = render_update( + operation, + parameters, + self.namespace, + tenant_id=tenant_id, + ) + if update: + self.client.update(update) + results: List[Any] = [] + elif operation in QUERY_OPERATIONS: + results = run_query( + operation, + self.client, + parameters, + self.namespace, + tenant_id=tenant_id, + ) + else: + raise NotImplementedError( + f'Operation {operation.value!r} is not implemented by the SPARQL backend' + ) + + if logger.isEnabledFor(logging.DEBUG): + elapsed = int((time.time() - start) * 1000) + logger.debug(f'[{log_entry.query_ref}] {elapsed}ms -> {len(results)} row(s)') + return results + + def __exit__(self, exception_type, exception_value, traceback): + if self._client is not None: + try: + self._client.close() + finally: + self._client = None + return False diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_graph_store_factory.py b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_graph_store_factory.py new file mode 100644 index 00000000..174aa542 --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_graph_store_factory.py @@ -0,0 +1,97 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import os +from urllib.parse import parse_qs, urlencode, unquote, urlparse, urlunparse + +from graphrag_toolkit.lexical_graph.storage.graph import ( + GraphStore, + GraphStoreFactoryMethod, + get_log_formatting, +) + +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_graph_store import ( + SPARQLDatabaseClient, +) + +logger = logging.getLogger(__name__) + +SPARQL_SCHEMES = ('sparql', 'sparql+s', 'sparql+http', 'sparql+https', 'sparql+neptune') + + +class SPARQLGraphStoreFactory(GraphStoreFactoryMethod): + """Create a graph store from SPARQL query/update endpoint configuration.""" + + def try_create(self, graph_info: str, **kwargs) -> GraphStore: + if not isinstance(graph_info, str): + return None + + parsed = urlparse(graph_info) + if parsed.scheme not in SPARQL_SCHEMES: + return None + + query_params = parse_qs(parsed.query) + query_endpoint = _endpoint_url(parsed, query_params) + update_endpoint = ( + kwargs.pop('update_endpoint', None) + or _single_query_value(query_params, 'update_endpoint') + or _single_query_value(query_params, 'update') + ) + if update_endpoint: + update_endpoint = unquote(update_endpoint) + + username_arg = kwargs.pop('username', None) + password_arg = kwargs.pop('password', None) + username = ( + unquote(parsed.username) if parsed.username else username_arg + ) or os.environ.get('SPARQL_USER') + password = ( + unquote(parsed.password) if parsed.password else password_arg + ) or os.environ.get('SPARQL_PASSWORD') + + kwargs.pop('config', None) + logger.debug( + 'Opening SPARQL graph store [scheme: %s, host: %s, path: %s]', + parsed.scheme, + parsed.hostname, + parsed.path, + ) + + return SPARQLDatabaseClient( + query_endpoint=query_endpoint, + update_endpoint=update_endpoint, + username=username, + password=password, + neptune_iam=(parsed.scheme == 'sparql+neptune'), + log_formatting=get_log_formatting(kwargs), + **kwargs, + ) + + +def _endpoint_url(parsed, query_params) -> str: + scheme = { + 'sparql': 'http', + 'sparql+s': 'https', + 'sparql+http': 'http', + 'sparql+https': 'https', + 'sparql+neptune': 'https', + }[parsed.scheme] + host = parsed.hostname or '' + if ':' in host and not host.startswith('['): + host = f'[{host}]' + netloc = host + if parsed.port: + netloc = f'{netloc}:{parsed.port}' + endpoint_query = urlencode([ + (name, value) + for name, values in query_params.items() + if name not in ('update_endpoint', 'update') + for value in values + ]) + return urlunparse((scheme, netloc, parsed.path, '', endpoint_query, '')) + + +def _single_query_value(query_params, name: str): + values = query_params.get(name) + return values[0] if values else None diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_queries.py b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_queries.py new file mode 100644 index 00000000..f17e3ac7 --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_queries.py @@ -0,0 +1,701 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native SPARQL read operations for the lexical graph. + +Each handler expresses one backend-neutral ``GraphOperation`` directly in +SPARQL and converts endpoint bindings into the result shape consumed by the +lexical-graph retrieval pipeline. + +Values that must be embedded in ``VALUES`` or ``FILTER`` clauses are serialized +with ``sparql_literal``; limits are converted to integers before interpolation. +Tenant scoping is applied once in ``run_query`` through the SPARQL Protocol's +``default-graph-uri`` parameter. +""" + +from typing import Any, Dict, List + +from graphrag_toolkit.lexical_graph.indexing.constants import LOCAL_ENTITY_CLASSIFICATION +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation +from graphrag_toolkit.lexical_graph.versioning import ( + VALID_FROM, + VALID_TO, + EXTRACT_TIMESTAMP, + BUILD_TIMESTAMP, + VERSION_INDEPENDENT_ID_FIELDS, + TIMESTAMP_LOWER_BOUND, + TIMESTAMP_UPPER_BOUND, +) + +from .ontology import DEFAULT_NAMESPACE, NamespaceConfig, sparql_literal, tenant_graph_iri + + +def run_query(operation: GraphOperation, + client, + parameters: Dict[str, Any], + namespace: NamespaceConfig = DEFAULT_NAMESPACE, + tenant_id: str = None) -> List[Dict[str, Any]]: + """Execute a lexical-graph operation with its native SPARQL query.""" + handler = _QUERY_HANDLERS.get(operation) + if handler is None: + raise NotImplementedError( + f'Operation {operation.value!r} is not implemented by the SPARQL backend' + ) + graph = tenant_graph_iri(tenant_id, namespace) + if graph: + client = _DefaultGraphClient(client, graph[1:-1]) + return handler(client, parameters, namespace) + + +class _DefaultGraphClient: + def __init__(self, client, graph_uri: str): + self._client = client + self._graph_uri = graph_uri + + def query(self, sparql): + return self._client.query(sparql, default_graph=self._graph_uri) + + +def _param_rows(parameters): + if 'params' in parameters: + return parameters['params'] or [] + return [parameters] if parameters else [] + + +def _positive_int(value, default: int, name: str) -> int: + """Validate numeric clauses before inserting them into SPARQL syntax.""" + result = int(value) if value is not None else default + if result < 1: + raise ValueError(f'{name} must be greater than zero') + return result + + +# Local-entity reconciliation ------------------------------------------------- + + +def _complements_matching_subject(client, parameters, namespace: NamespaceConfig): + """Local-entity rewrite lookup: real entities that have a local-entity twin + (same search_str). Empty when local entities are disabled.""" + out = [] + lg = namespace.prefix_ref + for row in _param_rows(parameters): + n_id = row.get('nId') + if not n_id: + continue + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?n_id ?c_id WHERE {{ + ?n {lg}id ?n_id ; {lg}search_str ?ss ; {lg}class ?ncls . + FILTER(?n_id = {sparql_literal(n_id)}) + FILTER(?ncls != "{LOCAL_ENTITY_CLASSIFICATION}") + ?c {lg}search_str ?ss ; {lg}class "{LOCAL_ENTITY_CLASSIFICATION}" ; {lg}id ?c_id . +}}''' + out.extend(client.query(sparql)) + return out + + +def _subjects_matching_complement(client, parameters, namespace: NamespaceConfig): + """Local-entity rewrite lookup: pair of nodes by id. Empty unless both + exist (i.e. a complement that also occurs as a real entity).""" + out = [] + lg = namespace.prefix_ref + for row in _param_rows(parameters): + n_id, c_id = row.get('nId'), row.get('cId') + if not n_id or not c_id: + continue + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?n_id ?c_id WHERE {{ + ?n {lg}id ?n_id . FILTER(?n_id = {sparql_literal(n_id)}) + ?c {lg}id ?c_id . FILTER(?c_id = {sparql_literal(c_id)}) +}}''' + out.extend(client.query(sparql)) + return out + + +def _single_entity_based_graph_search(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + start_id = parameters.get('startId') + if not start_id: + return [] + lg = namespace.prefix_ref + fact_pattern = f''' + ?entity {lg}id {sparql_literal(start_id)} . + ?fact {lg}subject ?entity .''' + return _statement_ids_for_fact_pattern(client, fact_pattern, parameters, namespace) + + +# Traversal searches ---------------------------------------------------------- + + +def _multiple_entity_based_graph_search(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + start_id = parameters.get('startId') + end_ids = parameters.get('endIds', []) or [] + if not start_id or not end_ids: + return [] + lg = namespace.prefix_ref + end_values = ' '.join(sparql_literal(e) for e in end_ids) + def connects(fact_var, e1, e2): + # a reified fact linking two entities, in either direction + return (f'{{ {{ {fact_var} {lg}subject {e1} ; {lg}object {e2} . }} ' + f'UNION {{ {fact_var} {lg}subject {e2} ; {lg}object {e1} . }} }}') + + fact_pattern = f''' + ?start {lg}id {sparql_literal(start_id)} . + VALUES ?endId {{ {end_values} }} + ?end {lg}id ?endId . + {{ + {connects('?fact', '?start', '?end')} + UNION + {{ ?mid a {lg}Entity . {connects('?fact', '?start', '?mid')} {connects('?fact2', '?mid', '?end')} }} + UNION + {{ ?mid a {lg}Entity . {connects('?fact1', '?start', '?mid')} {connects('?fact', '?mid', '?end')} }} + }}''' + return _statement_ids_for_fact_pattern(client, fact_pattern, parameters, namespace) + + +def _statement_ids_for_fact_pattern(client, + fact_pattern: str, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + limit = _positive_int(parameters.get('statementLimit'), 100, 'statementLimit') + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT DISTINCT ?l WHERE {{ +{fact_pattern} + ?fact {lg}supports ?statement . + {{ + {{ + ?statement {lg}id ?l . + }} + UNION + {{ + ?statement {lg}statementPrevious ?previous . + ?previous {lg}id ?l . + }} + UNION + {{ + ?next {lg}statementPrevious ?statement ; + {lg}id ?l . + }} + }} +}} LIMIT {limit}''' + return [{'l': row['l']} for row in client.query(sparql) if row.get('l') is not None] + + +def _facts_for_statements(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + statement_ids = parameters.get('statementIds', []) or [] + if not statement_ids: + return [] + values = ' '.join(sparql_literal(s) for s in statement_ids) + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?statementId ?factValue WHERE {{ + VALUES ?statementId {{ {values} }} + ?l {lg}id ?statementId . + ?f {lg}supports ?l . + OPTIONAL {{ ?f {lg}value ?factValue }} +}}''' + rows = client.query(sparql) + grouped: Dict[str, List[str]] = {} + for row in rows: + sid = row['statementId'] + grouped.setdefault(sid, []) + fact_value = row.get('factValue') + if fact_value is not None and fact_value not in grouped[sid]: + grouped[sid].append(fact_value) + return [{'statementId': sid, 'facts': facts} for sid, facts in grouped.items()] + + +# Retrieval content and entity scoring --------------------------------------- + + +def _chunk_content(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + chunk_ids = parameters.get('nodeIds', []) or [] + if not chunk_ids: + return [] + values = ' '.join(sparql_literal(chunk_id) for chunk_id in chunk_ids) + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?content WHERE {{ + VALUES ?chunkId {{ {values} }} + ?chunk a {lg}Chunk ; + {lg}id ?chunkId . + OPTIONAL {{ ?chunk {lg}value ?content }} +}}''' + return [{'content': row.get('content') or ''} for row in client.query(sparql)] + + +def _chunk_based_graph_search(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + return _statement_ids_for_chunk_id( + client, + parameters.get('chunkId') or parameters.get('nodeId'), + _positive_int(parameters.get('statementLimit'), 100, 'statementLimit'), + namespace, + ) + + +def _topic_based_entity_network_search(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + topic_id = parameters.get('nodeId') + if not topic_id: + return [] + lg = namespace.prefix_ref + limit = _positive_int(parameters.get('statementLimit'), 100, 'statementLimit') + sparql = f'''{namespace.sparql_prefixes()} +SELECT DISTINCT ?l WHERE {{ + ?topic a {lg}Topic ; + {lg}id {sparql_literal(topic_id)} . + ?statement a {lg}Statement ; + {lg}belongsTo ?topic ; + {lg}id ?l . +}} LIMIT {limit}''' + return [{'l': row['l']} for row in client.query(sparql) if row.get('l') is not None] + + +def _topic_content(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + topic_id = parameters.get('topicId') + if not topic_id: + return [] + lg = namespace.prefix_ref + limit = _positive_int(parameters.get('statementLimit'), 100, 'statementLimit') + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?statement ?details (COUNT(DISTINCT ?fact) AS ?score) WHERE {{ + ?topic a {lg}Topic ; + {lg}id {sparql_literal(topic_id)} . + ?statementNode a {lg}Statement ; + {lg}belongsTo ?topic . + OPTIONAL {{ ?statementNode {lg}value ?statement }} + OPTIONAL {{ ?statementNode {lg}details ?details }} + ?fact {lg}supports ?statementNode . +}} GROUP BY ?statement ?details +ORDER BY DESC(?score) +LIMIT {limit}''' + return [ + {'statement': row.get('statement') or '', 'details': row.get('details') or ''} + for row in client.query(sparql) + ] + + +def _statement_ids_for_chunk_id(client, + chunk_id, + limit: int, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + if not chunk_id: + return [] + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT DISTINCT ?l WHERE {{ + ?chunk a {lg}Chunk ; + {lg}id {sparql_literal(chunk_id)} . + ?statement a {lg}Statement ; + {lg}belongsTo ?topic ; + {lg}statementMentionedIn ?chunk ; + {lg}id ?l . +}} LIMIT {limit}''' + return [{'l': row['l']} for row in client.query(sparql) if row.get('l') is not None] + + +def _entities_for_keywords(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + keyword = parameters.get('keyword') + if not keyword: + return [] + classification = parameters.get('classification') + starts_with = parameters.get('_starts_with', False) + class_starts_with = parameters.get('_classification_starts_with', False) + lg = namespace.prefix_ref + + keyword_filter = ( + f'FILTER(STRSTARTS(?searchStr, {sparql_literal(keyword)}))' + if starts_with + else f'FILTER(?searchStr = {sparql_literal(keyword)})' + ) + if classification is not None: + class_filter = ( + f'FILTER(STRSTARTS(?class, {sparql_literal(classification)}))' + if class_starts_with + else f'FILTER(?class = {sparql_literal(classification)})' + ) + else: + class_filter = f'FILTER(?class != "{LOCAL_ENTITY_CLASSIFICATION}")' + + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?entityId ?value ?class (COUNT(?fact) AS ?score) WHERE {{ + ?entity a {lg}Entity ; + {lg}id ?entityId ; + {lg}search_str ?searchStr ; + {lg}class ?class . + OPTIONAL {{ ?entity {lg}value ?value }} + {keyword_filter} + {class_filter} + VALUES ?factPredicate {{ {lg}subject {lg}object }} + ?fact ?factPredicate ?entity . +}} GROUP BY ?entityId ?value ?class +ORDER BY DESC(?score)''' + return _entity_score_rows(client.query(sparql)) + + +def _entities_for_chunk_ids(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + node_ids = parameters.get('nodeIds', []) or [] + if not node_ids: + return [] + values = ' '.join(sparql_literal(node_id) for node_id in node_ids) + limit = _positive_int(parameters.get('limit'), 100, 'limit') + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?entityId ?value ?class (COUNT(?fact) AS ?score) WHERE {{ + {{ + SELECT DISTINCT ?entity WHERE {{ + VALUES ?chunkId {{ {values} }} + ?chunk a {lg}Chunk ; + {lg}id ?chunkId . + ?statement a {lg}Statement ; + {lg}statementMentionedIn ?chunk . + ?matchedFact {lg}supports ?statement ; + ?matchedFactPredicate ?entity . + VALUES ?matchedFactPredicate {{ {lg}subject {lg}object }} + }} + }} + ?entity a {lg}Entity ; + {lg}id ?entityId ; + {lg}class ?class . + OPTIONAL {{ ?entity {lg}value ?value }} + VALUES ?factPredicate {{ {lg}subject {lg}object }} + ?fact ?factPredicate ?entity . + FILTER(?class != "{LOCAL_ENTITY_CLASSIFICATION}") +}} GROUP BY ?entityId ?value ?class +ORDER BY DESC(?score) +LIMIT {limit}''' + return _entity_score_rows(client.query(sparql)) + + +def _entities_for_topic_ids(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + node_ids = parameters.get('nodeIds', []) or [] + if not node_ids: + return [] + values = ' '.join(sparql_literal(node_id) for node_id in node_ids) + limit = _positive_int(parameters.get('limit'), 100, 'limit') + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?entityId ?value ?class (COUNT(?fact) AS ?score) WHERE {{ + {{ + SELECT DISTINCT ?entity WHERE {{ + VALUES ?topicId {{ {values} }} + ?topic a {lg}Topic ; + {lg}id ?topicId . + ?statement a {lg}Statement ; + {lg}belongsTo ?topic . + ?matchedFact {lg}supports ?statement ; + ?matchedFactPredicate ?entity . + VALUES ?matchedFactPredicate {{ {lg}subject {lg}object }} + }} + }} + ?entity a {lg}Entity ; + {lg}id ?entityId ; + {lg}class ?class . + OPTIONAL {{ ?entity {lg}value ?value }} + VALUES ?factPredicate {{ {lg}subject {lg}object }} + ?fact ?factPredicate ?entity . + FILTER(?class != "{LOCAL_ENTITY_CLASSIFICATION}") +}} GROUP BY ?entityId ?value ?class +ORDER BY DESC(?score) +LIMIT {limit}''' + return _entity_score_rows(client.query(sparql)) + + +def _entity_score_rows(rows) -> List[Dict[str, Any]]: + out = [] + for row in rows: + entity_id = row.get('entityId') + if entity_id is None: + continue + out.append({ + 'result': { + 'entity': { + 'entityId': entity_id, + 'value': row.get('value') or '', + 'class': row.get('class') or '', + }, + 'score': float(row.get('score') or 0), + }, + }) + return out + + +def _next_level_in_tree(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + entity_ids = parameters.get('entityIds', []) or [] + if not entity_ids: + return [] + excluded = set(parameters.get('excludeEntityIds', []) or []) + num_neighbours = _positive_int( + parameters.get('numNeighbours'), 5, 'numNeighbours', + ) + values = ' '.join(sparql_literal(entity_id) for entity_id in entity_ids) + exclude_filter = '' + if excluded: + excluded_values = ', '.join(sparql_literal(entity_id) for entity_id in excluded) + exclude_filter = f'FILTER(?otherId NOT IN ({excluded_values}))' + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?entityId ?value ?class ?otherId (COUNT(?fact) AS ?score) WHERE {{ + VALUES ?entityId {{ {values} }} + ?entity a {lg}Entity ; + {lg}id ?entityId . + OPTIONAL {{ ?entity {lg}value ?value }} + OPTIONAL {{ ?entity {lg}class ?class }} + ?relFact {lg}subject ?entity ; + {lg}object ?other . + ?other a {lg}Entity ; + {lg}id ?otherId ; + {lg}class ?otherClass . + FILTER(?otherClass != "{LOCAL_ENTITY_CLASSIFICATION}") + {exclude_filter} + VALUES ?factPredicate {{ {lg}subject {lg}object }} + ?fact ?factPredicate ?other . +}} GROUP BY ?entityId ?value ?class ?otherId +ORDER BY ?entityId DESC(?score)''' + grouped: Dict[str, Dict[str, Any]] = {} + for row in client.query(sparql): + entity_id = row.get('entityId') + other_id = row.get('otherId') + if entity_id is None or other_id is None: + continue + result = grouped.setdefault(entity_id, { + 'result': { + 'entity': { + 'entityId': entity_id, + 'value': row.get('value') or '', + 'class': row.get('class') or '', + }, + 'others': [], + }, + }) + others = result['result']['others'] + if other_id not in others and len(others) < num_neighbours: + others.append(other_id) + return list(grouped.values()) + + +def _expand_entities(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + entity_ids = parameters.get('entityIds', []) or [] + if not entity_ids: + return [] + values = ' '.join(sparql_literal(entity_id) for entity_id in entity_ids) + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?entityId ?value ?class (COUNT(?fact) AS ?score) WHERE {{ + VALUES ?entityId {{ {values} }} + ?entity a {lg}Entity ; + {lg}id ?entityId ; + {lg}class ?class . + OPTIONAL {{ ?entity {lg}value ?value }} + VALUES ?factPredicate {{ {lg}subject {lg}object }} + ?fact ?factPredicate ?entity . +}} GROUP BY ?entityId ?value ?class''' + return _entity_score_rows(client.query(sparql)) + + +def _statements_grouped_by_topic_and_source(client, + parameters, + namespace: NamespaceConfig) -> List[Dict[str, Any]]: + statement_ids = parameters.get('statementIds', []) or [] + if not statement_ids: + return [] + + values = ' '.join(sparql_literal(s) for s in statement_ids) + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT DISTINCT ?statementId ?statementValue ?details ?chunkId ?topicId ?topicValue ?sourceId WHERE {{ + VALUES ?statementId {{ {values} }} + ?statement {lg}id ?statementId ; + {lg}belongsTo ?topic ; + {lg}statementMentionedIn ?chunk . + OPTIONAL {{ ?statement {lg}value ?statementValue }} + OPTIONAL {{ ?statement {lg}details ?details }} + ?topic {lg}id ?topicId . + OPTIONAL {{ ?topic {lg}value ?topicValue }} + ?chunk {lg}id ?chunkId ; + {lg}extractedFrom ?source . + ?source {lg}id ?sourceId . +}}''' + rows = client.query(sparql) + if not rows: + return [] + + source_ids = sorted({row['sourceId'] for row in rows if row.get('sourceId') is not None}) + chunk_ids = sorted({row['chunkId'] for row in rows if row.get('chunkId') is not None}) + source_props = _properties_by_id(client, 'Source', source_ids, namespace) + include_chunk_details = parameters.get('_include_chunk_details', False) + chunk_props = _properties_by_id(client, 'Chunk', chunk_ids, namespace) if include_chunk_details else {} + + grouped: Dict[str, Dict[str, Any]] = {} + topic_indexes: Dict[str, Dict[str, Dict[str, Any]]] = {} + chunk_seen: Dict[str, Dict[str, set]] = {} + statement_seen: Dict[str, Dict[str, set]] = {} + + for row in rows: + source_id = row['sourceId'] + source_metadata = dict(source_props.get(source_id, {})) + source_metadata.setdefault('sourceId', source_id) + + if source_id not in grouped: + grouped[source_id] = { + 'score': 0, + 'source': { + 'sourceId': source_id, + 'metadata': source_metadata, + 'versioning': _versioning_from(source_metadata), + }, + 'topics': [], + } + topic_indexes[source_id] = {} + chunk_seen[source_id] = {} + statement_seen[source_id] = {} + + topic_id = row['topicId'] + topics_by_id = topic_indexes[source_id] + if topic_id not in topics_by_id: + topic = { + 'topic': row.get('topicValue') or '', + 'topicId': topic_id, + 'chunks': [], + 'statements': [], + } + topics_by_id[topic_id] = topic + grouped[source_id]['topics'].append(topic) + chunk_seen[source_id][topic_id] = set() + statement_seen[source_id][topic_id] = set() + + topic = topics_by_id[topic_id] + chunk_id = row['chunkId'] + if chunk_id not in chunk_seen[source_id][topic_id]: + metadata = dict(chunk_props.get(chunk_id, {})) if include_chunk_details else {} + if include_chunk_details: + metadata.setdefault('chunkId', chunk_id) + topic['chunks'].append({ + 'chunkId': chunk_id, + 'value': None, + 'metadata': metadata, + }) + chunk_seen[source_id][topic_id].add(chunk_id) + + statement_id = row['statementId'] + if statement_id not in statement_seen[source_id][topic_id]: + topic['statements'].append({ + 'statementId': statement_id, + 'statement': row.get('statementValue') or '', + 'facts': [], + 'details': row.get('details'), + 'chunkId': chunk_id, + 'score': 0, + }) + statement_seen[source_id][topic_id].add(statement_id) + + results = [] + for result in grouped.values(): + result['score'] = sum( + len(topic['statements']) / len(topic['chunks']) + for topic in result['topics'] + if topic['chunks'] + ) + results.append({'result': result}) + + results.sort(key=lambda r: r['result']['score'], reverse=True) + limit = parameters.get('limit') + return results[:_positive_int(limit, 1, 'limit')] if limit is not None else results + + +def _properties_by_id(client, + cls: str, + ids: List[str], + namespace: NamespaceConfig) -> Dict[str, Dict[str, Any]]: + if not ids: + return {} + values = ' '.join(sparql_literal(i) for i in ids) + lg = namespace.prefix_ref + sparql = f'''{namespace.sparql_prefixes()} +SELECT ?id ?prop ?value WHERE {{ + VALUES ?id {{ {values} }} + ?node a {lg}{cls} ; + {lg}id ?id ; + ?prop ?value . + FILTER(STRSTARTS(STR(?prop), "{namespace.schema_namespace}")) + FILTER(?prop != {lg}id) + FILTER(isLiteral(?value)) +}}''' + out: Dict[str, Dict[str, Any]] = {i: {} for i in ids} + for row in client.query(sparql): + prop = _local_name(row.get('prop'), namespace) + if prop: + out.setdefault(row['id'], {})[prop] = row.get('value') + return out + + +def _versioning_from(metadata: Dict[str, Any]) -> Dict[str, Any]: + id_fields = metadata.get(VERSION_INDEPENDENT_ID_FIELDS, '') + return { + 'valid_from': _int_or_default(metadata.get(VALID_FROM), TIMESTAMP_LOWER_BOUND), + 'valid_to': _int_or_default(metadata.get(VALID_TO), TIMESTAMP_UPPER_BOUND), + 'extract_timestamp': _int_or_default(metadata.get(EXTRACT_TIMESTAMP), TIMESTAMP_LOWER_BOUND), + 'build_timestamp': _int_or_default(metadata.get(BUILD_TIMESTAMP), TIMESTAMP_LOWER_BOUND), + 'id_fields': id_fields.split(';') if isinstance(id_fields, str) and id_fields else [], + } + + +def _int_or_default(value, default: int) -> int: + if value is None: + return default + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _local_name(uri, namespace: NamespaceConfig) -> str: + if not uri: + return '' + text = str(uri) + if text.startswith(namespace.schema_namespace): + return text[len(namespace.schema_namespace):] + return text.rsplit('#', 1)[-1].rsplit('/', 1)[-1] + + +_QUERY_HANDLERS = { + GraphOperation.FIND_COMPLEMENTS: _complements_matching_subject, + GraphOperation.FIND_SUBJECTS: _subjects_matching_complement, + GraphOperation.GET_STATEMENTS: _statements_grouped_by_topic_and_source, + GraphOperation.GET_FACTS: _facts_for_statements, + GraphOperation.GET_CHUNKS: _chunk_content, + GraphOperation.GET_TOPIC: _topic_content, + GraphOperation.SEARCH_BY_CHUNK: _chunk_based_graph_search, + GraphOperation.SEARCH_BY_TOPIC: _topic_based_entity_network_search, + GraphOperation.FIND_ENTITIES_BY_KEYWORD: _entities_for_keywords, + GraphOperation.FIND_ENTITIES_BY_CHUNKS: _entities_for_chunk_ids, + GraphOperation.FIND_ENTITIES_BY_TOPICS: _entities_for_topic_ids, + GraphOperation.FIND_ENTITY_NEIGHBORS: _next_level_in_tree, + GraphOperation.SCORE_ENTITIES: _expand_entities, + GraphOperation.SEARCH_BY_ENTITY: _single_entity_based_graph_search, + GraphOperation.SEARCH_BY_ENTITIES: _multiple_entity_based_graph_search, +} + +QUERY_OPERATIONS = frozenset(_QUERY_HANDLERS) diff --git a/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_updates.py b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_updates.py new file mode 100644 index 00000000..c62aebe8 --- /dev/null +++ b/lexical-graph-contrib/sparql/src/graphrag_toolkit_contrib/lexical_graph/storage/graph/sparql/sparql_updates.py @@ -0,0 +1,378 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native SPARQL updates for lexical-graph build operations. + +Handlers consume structured GraphRAG parameters and render SPARQL Update +directly; no property-graph query is inspected. Related statements are batched +into one update request to avoid the many network round trips that RDFLib +``Graph.add``/``Graph.remove`` would otherwise require for a remote store. + +All identifiers and literals pass through the ontology helpers, and tenant data +is written to a deterministic named graph. +""" + +from typing import Any, Dict, List, Optional + +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation + +from .ontology import ( + DEFAULT_NAMESPACE, + ID_KEY_TO_KIND, + NamespaceConfig, + RDF_TYPE, + instance_iri, + relation_iri, + sparql_literal, + sys_relation_iri, + tenant_graph_iri, + term, +) + +_lit = sparql_literal + + +def render_update(operation: GraphOperation, + parameters: Dict[str, Any], + namespace: NamespaceConfig = DEFAULT_NAMESPACE, + tenant_id: Optional[str] = None) -> Optional[str]: + """Render a lexical-graph operation as a native SPARQL update.""" + rows = _rows(parameters) + if not rows: + return None + + handler = _UPDATE_HANDLERS.get(operation) + if handler is None: + raise NotImplementedError( + f'Operation {operation.value!r} is not implemented by the SPARQL backend' + ) + + graph = tenant_graph_iri(tenant_id, namespace) + ops: List[str] = [] + for row in rows: + update = handler(row, graph, namespace) + if update: + ops.append(update) + + return ' ;\n'.join(ops) if ops else None + + +def _source(row: Dict[str, Any], graph, namespace: NamespaceConfig) -> str: + source_id = row['_source_id'] + props = [] + for key, value in row.items(): + if key.startswith('_'): + continue + lit = _lit(value) + if lit is not None: + props.append((term(_safe_local(key), namespace), lit)) + return _node_upsert('sourceId', source_id, 'Source', props, graph, namespace) + + +def _chunk(row: Dict[str, Any], graph, namespace: NamespaceConfig) -> str: + props = [] + if row.get('text') is not None: + props.append((term('value', namespace), _lit(row['text']))) + for key, value in row.items(): + if key in ('chunk_id', 'text') or key.startswith('_'): + continue + lit = _lit(value) + if lit is not None: + props.append((term(_safe_local(key), namespace), lit)) + return _node_upsert('chunkId', row['chunk_id'], 'Chunk', props, graph, namespace) + + +def _topic(row: Dict[str, Any], graph, namespace: NamespaceConfig) -> str: + ops = [_node_upsert('topicId', row['topic_id'], 'Topic', + [(term('value', namespace), _lit(row.get('title')))] if row.get('title') is not None else [], + graph, namespace)] + topic_iri = instance_iri('topic', row['topic_id'], namespace) + for chunk_ref in row.get('chunk_ids', []) or []: + chunk_id = chunk_ref['chunk_id'] if isinstance(chunk_ref, dict) else chunk_ref + chunk_iri = instance_iri('chunk', chunk_id, namespace) + triples = [ + f'{chunk_iri} {RDF_TYPE} {term("Chunk", namespace)} .', + f'{chunk_iri} {term("id", namespace)} {_lit(chunk_id)} .', + f'{topic_iri} {term("topicMentionedIn", namespace)} {chunk_iri} .', + ] + ops.append(_insert_data('\n'.join(triples), graph)) + return ' ;\n'.join(ops) + + +def _statement(row: Dict[str, Any], graph, namespace: NamespaceConfig) -> str: + props = [] + if row.get('value') is not None: + props.append((term('value', namespace), _lit(row['value']))) + if row.get('details') is not None: + props.append((term('details', namespace), _lit(row['details']))) + return _node_upsert('statementId', row['statement_id'], 'Statement', props, graph, namespace) + + +def _fact(row: Dict[str, Any], graph, namespace: NamespaceConfig) -> str: + # Entity subjects and objects are linked separately; local values are literals. + props = [] + if row.get('fact') is not None: + props.append((term('value', namespace), _lit(row['fact']))) + subject_literal = row.get('_subject_literal') + if subject_literal is not None: + props.append((term('subject', namespace), _lit(subject_literal))) + object_literal = row.get('_object_literal') + if object_literal is not None: + props.append((term('object', namespace), _lit(object_literal))) + predicate_value = row.get('_predicate') + rel = relation_iri(predicate_value, namespace) if predicate_value is not None else None + if rel is not None: + props.append((term('predicate', namespace), rel)) + ops = [_node_upsert('factId', row['fact_id'], 'Fact', props, graph, namespace)] + if rel is not None: + ops.append(_insert_data('\n'.join([ + f'{rel} {RDF_TYPE} {term("Relation", namespace)} .', + f'{rel} {term("value", namespace)} {_lit(predicate_value)} .', + ]), graph)) + fact_iri = instance_iri('fact', row['fact_id'], namespace) + stmt_iri = instance_iri('statement', row['statement_id'], namespace) + triples = [ + f'{stmt_iri} {RDF_TYPE} {term("Statement", namespace)} .', + f'{stmt_iri} {term("id", namespace)} {_lit(row["statement_id"])} .', + f'{fact_iri} {term("supports", namespace)} {stmt_iri} .', + ] + ops.append(_insert_data('\n'.join(triples), graph)) + return ' ;\n'.join(ops) + + +def _entity(row: Dict[str, Any], graph, namespace: NamespaceConfig) -> str: + props = [] + for key, pred in (('v', 'value'), ('e_search_str', 'search_str'), ('ec', 'class')): + if row.get(key) is not None: + props.append((term(pred, namespace), _lit(row[key]))) + return _node_upsert('entityId', row['e_id'], 'Entity', props, graph, namespace) + + +def _edge(row, graph, namespace: NamespaceConfig, + a_key, a_param, b_key, b_param, predicate_name) -> Optional[str]: + if a_param not in row or b_param not in row: + return None + a_kind, a_cls = _kind_cls(a_key) + b_kind, b_cls = _kind_cls(b_key) + a_iri = instance_iri(a_kind, row[a_param], namespace) + b_iri = instance_iri(b_kind, row[b_param], namespace) + predicate = term(predicate_name, namespace) + triples = [ + f'{a_iri} {RDF_TYPE} {term(a_cls, namespace)} .', + f'{a_iri} {term("id", namespace)} {_lit(row[a_param])} .', + f'{b_iri} {RDF_TYPE} {term(b_cls, namespace)} .', + f'{b_iri} {term("id", namespace)} {_lit(row[b_param])} .', + f'{a_iri} {predicate} {b_iri} .', + ] + return _insert_data('\n'.join(triples), graph) + + +def _graph_summary(row: Dict[str, Any], graph, namespace: NamespaceConfig) -> str: + two_class = row['sc_id'] != row['oc_id'] + delta = 1 if two_class else 2 + + sc_iri = instance_iri('sysclass', row['sc_id'], namespace) + ops = [ + _insert_data('\n'.join([ + f'{sc_iri} {RDF_TYPE} {term("SysClass", namespace)} .', + f'{sc_iri} {term("id", namespace)} {_lit(row["sc_id"])} .', + f'{sc_iri} {term("value", namespace)} {_lit(row.get("sc"))} .', + ]), graph), + _increment(sc_iri, term('count', namespace), delta, graph), + ] + + object_class_id = row['oc_id'] if two_class else row['sc_id'] + if two_class: + oc_iri = instance_iri('sysclass', row['oc_id'], namespace) + ops.append(_insert_data('\n'.join([ + f'{oc_iri} {RDF_TYPE} {term("SysClass", namespace)} .', + f'{oc_iri} {term("id", namespace)} {_lit(row["oc_id"])} .', + f'{oc_iri} {term("value", namespace)} {_lit(row.get("oc"))} .', + ]), graph)) + ops.append(_increment(oc_iri, term('count', namespace), delta, graph)) + else: + oc_iri = sc_iri + + sysrel = sys_relation_iri(row['sc_id'], row.get('p'), object_class_id, namespace) + sysrel_triples = [ + f'{sysrel} {RDF_TYPE} {term("SysRelation", namespace)} .', + f'{sysrel} {term("sysRelSubject", namespace)} {sc_iri} .', + f'{sysrel} {term("sysRelObject", namespace)} {oc_iri} .', + ] + if row.get('p') is not None: + sysrel_triples.append(f'{sysrel} {term("value", namespace)} {_lit(row["p"])} .') + ops.append(_insert_data('\n'.join(sysrel_triples), graph)) + ops.append(_increment(sysrel, term('count', namespace), delta, graph)) + return ' ;\n'.join(ops) + + +def _domain_type(row: Dict[str, Any], graph, namespace: NamespaceConfig) -> Optional[str]: + entity_id = row.get('entityId') + classification = row.get('_classification') + if entity_id is None or not classification: + return None + entity_iri = instance_iri('entity', entity_id, namespace) + triples = [ + f'{entity_iri} {RDF_TYPE} {term("Entity", namespace)} .', + f'{entity_iri} {term("id", namespace)} {_lit(entity_id)} .', + f'{entity_iri} {RDF_TYPE} {term(_safe_local(classification), namespace)} .', + ] + return _insert_data('\n'.join(triples), graph) + + +def _node_upsert(id_key, id_value, cls, props, graph, namespace: NamespaceConfig) -> str: + ops = [_delete_prop(instance_iri(*_kind_cls_iri(id_key, id_value), namespace), pred, graph) + for pred, _ in props] + iri = instance_iri(*_kind_cls_iri(id_key, id_value), namespace) + lines = [f'{iri} {RDF_TYPE} {term(cls, namespace)} .', f'{iri} {term("id", namespace)} {_lit(id_value)} .'] + lines.extend(f'{iri} {pred} {lit} .' for pred, lit in props) + ops.append(_insert_data('\n'.join(lines), graph)) + return ' ;\n'.join(ops) + + +def _delete_prop(iri, predicate, graph) -> str: + return f'DELETE WHERE {{ {_wrap(f"{iri} {predicate} ?o", graph)} }}' + + +def _insert_data(triples, graph) -> str: + return f'INSERT DATA {{ {_wrap(triples, graph)} }}' + + +def _increment(iri, predicate, delta, graph) -> str: + del_block = _wrap(f'{iri} {predicate} ?c', graph) + ins_block = _wrap(f'{iri} {predicate} ?newc', graph) + where = (f'{_wrap(f"OPTIONAL {{ {iri} {predicate} ?c0 }}", graph)} ' + f'BIND(COALESCE(?c0, 0) + {delta} AS ?newc) BIND(?c0 AS ?c)') + return f'DELETE {{ {del_block} }} INSERT {{ {ins_block} }} WHERE {{ {where} }}' + + +def _wrap(pattern, graph) -> str: + if graph: + return f'GRAPH {graph} {{ {pattern} }}' + return pattern + + +def _rows(parameters: Dict[str, Any]) -> List[Dict[str, Any]]: + if parameters is None: + return [] + if 'params' in parameters: + return parameters['params'] or [] + return [parameters] if parameters else [] + + +def _kind_cls(id_key): + return ID_KEY_TO_KIND[id_key] + + +def _kind_cls_iri(id_key, id_value): + kind, _ = _kind_cls(id_key) + return kind, id_value + + +def _safe_local(key) -> str: + return ''.join(c if (c.isalnum() or c == '_') else '_' for c in str(key)) + + +def _link_chunk_source(row, graph, namespace): + return _edge( + row, graph, namespace, + 'chunkId', 'chunk_id', 'sourceId', 'source_id', 'extractedFrom', + ) + + +def _link_chunks(row, graph, namespace): + predicates = { + 'parent': 'parent', + 'child': 'child', + 'previous': 'chunkPrevious', + 'next': 'next', + } + relationship = str(row.get('_relationship_type', '')).lower() + if relationship not in predicates: + raise ValueError(f'Invalid chunk relationship type: {relationship!r}') + return _edge( + row, graph, namespace, + 'chunkId', 'chunk_id', 'chunkId', 'target_id', predicates[relationship], + ) + + +def _link_statement_chunk(row, graph, namespace): + return _edge( + row, graph, namespace, + 'statementId', 'statement_id', 'chunkId', 'chunk_id', 'statementMentionedIn', + ) + + +def _link_statement_topic(row, graph, namespace): + return _edge( + row, graph, namespace, + 'statementId', 'statement_id', 'topicId', 'topic_id', 'belongsTo', + ) + + +def _link_statements(row, graph, namespace): + return _edge( + row, graph, namespace, + 'statementId', 'statement_id', 'statementId', 'prev_statement_id', + 'statementPrevious', + ) + + +def _link_fact_entity(row, graph, namespace): + relationship = str(row.get('_relationship_type', '')).lower() + if relationship not in ('subject', 'object'): + if not row: + return None + raise ValueError(f'Invalid fact relationship type: {relationship!r}') + return _edge( + row, graph, namespace, + 'factId', 'fact_id', 'entityId', 'entity_id', relationship, + ) + + +def _fact_relation_is_reified(row, graph, namespace): + """The Fact resource already represents the entity relationship.""" + return None + + +def _copy_complement_relationships(row, graph, namespace): + complement = instance_iri('entity', row['c_id'], namespace) + resolved = instance_iri('entity', row['n_id'], namespace) + predicate = term('object', namespace) + existing = _wrap(f'?fact {predicate} {complement}', graph) + old = _wrap(f'?fact {predicate} {complement}', graph) + new = _wrap(f'?fact {predicate} {resolved}', graph) + incoming = _wrap(f'?subject ?predicate {complement}', graph) + return ( + f'DELETE {{ {old} }} INSERT {{ {new} }} WHERE {{ {existing} }} ;\n' + f'DELETE WHERE {{ {incoming} }}' + ) + + +def _delete_complement(row, graph, namespace): + complement = instance_iri('entity', row['c_id'], namespace) + outgoing = _wrap(f'{complement} ?predicate ?object', graph) + return f'DELETE WHERE {{ {outgoing} }}' + + +_UPDATE_HANDLERS = { + GraphOperation.UPSERT_SOURCE: _source, + GraphOperation.UPSERT_CHUNK: _chunk, + GraphOperation.LINK_CHUNK_SOURCE: _link_chunk_source, + GraphOperation.LINK_CHUNKS: _link_chunks, + GraphOperation.UPSERT_TOPIC: _topic, + GraphOperation.UPSERT_STATEMENT: _statement, + GraphOperation.LINK_STATEMENT_CHUNK: _link_statement_chunk, + GraphOperation.LINK_STATEMENT_TOPIC: _link_statement_topic, + GraphOperation.LINK_STATEMENTS: _link_statements, + GraphOperation.UPSERT_FACT: _fact, + GraphOperation.LINK_FACT_ENTITY: _link_fact_entity, + GraphOperation.UPSERT_ENTITY: _entity, + GraphOperation.LINK_ENTITIES: _fact_relation_is_reified, + GraphOperation.ADD_ENTITY_TYPE: _domain_type, + GraphOperation.UPDATE_GRAPH_SUMMARY: _graph_summary, + GraphOperation.COPY_COMPLEMENT_RELATIONSHIPS: _copy_complement_relationships, + GraphOperation.DELETE_COMPLEMENT: _delete_complement, +} + +UPDATE_OPERATIONS = frozenset(_UPDATE_HANDLERS) diff --git a/lexical-graph-contrib/sparql/tests/test_neptune_iam.py b/lexical-graph-contrib/sparql/tests/test_neptune_iam.py new file mode 100644 index 00000000..cbffaf3d --- /dev/null +++ b/lexical-graph-contrib/sparql/tests/test_neptune_iam.py @@ -0,0 +1,91 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import requests +from botocore.credentials import Credentials +from rdflib import Graph + +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.neptune_iam import ( + NeptuneIAMAuth, + NeptuneIAMStore, + neptune_iam_graph, +) + + +class _AWSSession: + region_name = 'eu-central-1' + + def get_credentials(self): + return Credentials('access-key', 'secret-key', 'session-token') + + +class _Response: + status_code = 200 + text = '' + content = b'{"head": {}, "boolean": true}' + headers = {'Content-Type': 'application/sparql-results+json'} + + def raise_for_status(self): + pass + + +class _HTTPSession: + def __init__(self): + self.calls = [] + self.closed = False + + def post(self, url, **kwargs): + self.calls.append((url, kwargs)) + return _Response() + + def close(self): + self.closed = True + + +def test_auth_signs_with_neptune_service_and_session_token(): + request = requests.Request( + 'POST', 'https://example.test/sparql', data={'query': 'ASK {}'} + ).prepare() + + NeptuneIAMAuth(aws_session=_AWSSession())(request) + + assert '/eu-central-1/neptune-db/aws4_request' in request.headers['Authorization'] + assert request.headers['X-Amz-Security-Token'] == 'session-token' + + +def test_store_queries_and_updates_neptune_sparql_endpoint(): + store = NeptuneIAMStore( + 'https://example.test:8182', + aws_session=_AWSSession(), + headers={'X-Application': 'graph-rag'}, + ) + store._http.close() + store._http = http = _HTTPSession() + graph = Graph(store=store) + + assert graph.query('ASK {}').askAnswer is True + graph.update('INSERT DATA { }') + graph.close() + + assert store.query_endpoint == 'https://example.test:8182/sparql' + assert http.calls[0][1]['data']['query'].endswith('ASK {}') + assert http.calls[1][1]['data']['update'].endswith( + 'INSERT DATA { }' + ) + assert http.calls[0][1]['headers']['X-Application'] == 'graph-rag' + assert http.closed is True + + +def test_store_rejects_unsigned_http_endpoint(): + with pytest.raises(ValueError, match='requires an HTTPS endpoint'): + NeptuneIAMStore('http://example.test:8182', aws_session=_AWSSession()) + + +def test_graph_helper_returns_rdflib_graph(): + graph = neptune_iam_graph( + 'https://example.test:8182/sparql/', aws_session=_AWSSession() + ) + assert isinstance(graph, Graph) + assert graph.store.query_endpoint == 'https://example.test:8182/sparql' + graph.close() diff --git a/lexical-graph-contrib/sparql/tests/test_ontology.py b/lexical-graph-contrib/sparql/tests/test_ontology.py new file mode 100644 index 00000000..cf1d6616 --- /dev/null +++ b/lexical-graph-contrib/sparql/tests/test_ontology.py @@ -0,0 +1,45 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the ontology helpers and NamespaceConfig validation.""" + +import pytest + +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.ontology import ( + NamespaceConfig, + sparql_literal, +) + + +def test_invalid_main_prefix_rejected(): + with pytest.raises(ValueError): + NamespaceConfig(prefix='1bad') + + +def test_invalid_extra_prefix_rejected(): + with pytest.raises(ValueError): + NamespaceConfig(extra_prefixes={'1bad': 'https://ex.test/schema#'}) + + +def test_extra_prefix_conflicting_with_main_rejected(): + with pytest.raises(ValueError): + NamespaceConfig(prefix='lg', extra_prefixes={'lg': 'https://different.test/schema#'}) + + +def test_namespace_without_separator_gets_one_appended(): + ns = NamespaceConfig(schema_namespace='https://ex.test/schema', + instance_namespace='https://ex.test/data') + assert ns.schema_namespace == 'https://ex.test/schema#' + assert ns.instance_namespace == 'https://ex.test/data/' + + +@pytest.mark.parametrize('value,expected', [ + (None, None), + (True, '"true"^^'), + (False, '"false"^^'), + (7, '"7"^^'), + (1.5, '"1.5"^^'), + ('x', '"x"'), +]) +def test_sparql_literal_renders_each_type(value, expected): + assert sparql_literal(value) == expected diff --git a/lexical-graph-contrib/sparql/tests/test_sparql_endpoint_client.py b/lexical-graph-contrib/sparql/tests/test_sparql_endpoint_client.py new file mode 100644 index 00000000..4cbc4928 --- /dev/null +++ b/lexical-graph-contrib/sparql/tests/test_sparql_endpoint_client.py @@ -0,0 +1,158 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +from rdflib import Graph +from requests.auth import HTTPBasicAuth + +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_endpoint_client import ( + FORM_URLENCODED, + SPARQL_JSON, + RDFLibHTTPStore, + SPARQLEndpointClient, +) + + +class _Response: + status_code = 200 + text = '' + content = b'{"head":{"vars":["value"]},"results":{"bindings":[{"value":{"type":"literal","value":"ok"}}]}}' + headers = {'Content-Type': SPARQL_JSON} + + +class _Session: + def __init__(self, response=None): + self.auth = None + self.calls = [] + self.closed = False + self.response = response or _Response() + + def post(self, url, data, headers, timeout): + self.calls.append({ + 'url': url, + 'data': data, + 'headers': headers, + 'timeout': timeout, + }) + return self.response + + def close(self): + self.closed = True + + +def _client(session=None, **kwargs): + client = SPARQLEndpointClient('http://example.test/query', **kwargs) + client.store._http.close() + client.store._http = session or _Session() + return client + + +def test_client_sends_form_encoded_query_and_update_requests(): + session = _Session() + client = _client( + session, + update_endpoint='http://example.test/update', + headers={'Authorization': 'Bearer token', 'Content-Type': 'text/plain'}, + timeout=12, + ) + + assert client.query('SELECT ?value WHERE { VALUES ?value { "ok" } }') == [ + {'value': 'ok'}, + ] + client.update('INSERT DATA { }') + + query_call, update_call = session.calls + assert query_call['url'] == 'http://example.test/query' + assert query_call['data']['query'].endswith( + 'SELECT ?value WHERE { VALUES ?value { "ok" } }' + ) + assert query_call['headers']['Content-Type'] == FORM_URLENCODED + assert query_call['headers']['Accept'].startswith(SPARQL_JSON) + assert 'text/turtle' in query_call['headers']['Accept'] + assert query_call['headers']['Authorization'] == 'Bearer token' + assert query_call['timeout'] == 12 + + assert update_call['url'] == 'http://example.test/update' + assert update_call['data']['update'].endswith( + 'INSERT DATA { }' + ) + assert update_call['headers']['Content-Type'] == FORM_URLENCODED + assert update_call['headers']['Authorization'] == 'Bearer token' + + +class _ErrorResponse(_Response): + status_code = 500 + text = 'kaboom' + + +def test_query_returns_ask_boolean(): + response = _Response() + response.content = b'{"head": {}, "boolean": true}' + assert _client(_Session(response)).query('ASK { ?s ?p ?o }') == [ + {'boolean': True}, + ] + + +def test_rdflib_graph_can_parse_construct_results(): + response = _Response() + response.content = b' .' + response.headers = {'Content-Type': 'text/turtle'} + store = RDFLibHTTPStore('http://example.test/query') + store._http.close() + store._http = _Session(response) + graph = Graph(store=store) + + result = graph.query('CONSTRUCT { ?s ?p ?o } WHERE { ?s ?p ?o }') + + assert result.type == 'CONSTRUCT' + assert len(result.graph) == 1 + graph.close() + + +def test_query_can_scope_the_protocol_default_graph(): + session = _Session() + client = _client(session) + + client.query('SELECT * WHERE { ?s ?p ?o }', default_graph='urn:tenant:acme') + + assert session.calls[0]['data']['default-graph-uri'] == 'urn:tenant:acme' + + +def test_query_defaults_update_endpoint_to_query_endpoint(): + client = _client() + assert client.store.update_endpoint == 'http://example.test/query' + + +def test_raise_for_status_raises_on_http_error(): + with pytest.raises(RuntimeError, match='kaboom') as error: + _client(_Session(_ErrorResponse())).query('SELECT * WHERE { ?s ?p ?o }') + assert 'SELECT' not in str(error.value) + + +def test_close_closes_the_session(): + session = _Session() + client = _client(session) + client.close() + assert session.closed is True + + +def test_client_uses_basic_auth_and_rdflib_store(): + client = SPARQLEndpointClient('http://example.test/query', username='u', password='pw') + assert isinstance(client.store, RDFLibHTTPStore) + assert isinstance(client.store._http.auth, HTTPBasicAuth) + client.close() + + +def test_client_rejects_non_tabular_query_results(): + client = _client() + client._graph.query = lambda _: SimpleNamespace(type='CONSTRUCT') + + with pytest.raises(ValueError, match='Only SELECT and ASK'): + client.query('CONSTRUCT { ?s ?p ?o } WHERE { ?s ?p ?o }') + + +def test_client_rejects_non_positive_timeout(): + with pytest.raises(ValueError, match='timeout must be greater than zero'): + SPARQLEndpointClient('http://example.test/query', timeout=0) diff --git a/lexical-graph-contrib/sparql/tests/test_sparql_graph_store.py b/lexical-graph-contrib/sparql/tests/test_sparql_graph_store.py new file mode 100644 index 00000000..8d07b202 --- /dev/null +++ b/lexical-graph-contrib/sparql/tests/test_sparql_graph_store.py @@ -0,0 +1,193 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for SPARQL operation dispatch and lifecycle (no server).""" + +import logging + +import pytest + +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_graph_store import ( + SPARQLDatabaseClient, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_endpoint_client import ( + RDFLibHTTPStore, + SPARQLEndpointClient, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.neptune_iam import ( + NeptuneIAMStore, +) + + +class _FakeClient: + def __init__(self): + self.updates = [] + self.queries = [] + self.default_graphs = [] + self.closed = False + + def update(self, sparql): + self.updates.append(sparql) + + def query(self, sparql, default_graph=None): + self.queries.append(sparql) + self.default_graphs.append(default_graph) + return [{'l': 'stmt-1'}] + + def close(self): + self.closed = True + + +def _store(**kwargs): + return SPARQLDatabaseClient(query_endpoint='http://ex.test/query', **kwargs) + + +def test_node_id_formats_id(): + assert _store().node_id('entityId') + + +def test_client_property_lazily_builds_and_caches_real_client(): + store = _store() + client = store.client + assert isinstance(client, SPARQLEndpointClient) + assert isinstance(client.store, RDFLibHTTPStore) + assert store.client is client # cached on the private attr + + +def test_client_property_unwraps_secret_password(): + store = _store(username='u', password='pw') + assert store.client.store._http.auth is not None + + +def test_client_property_uses_neptune_iam_store(): + store = SPARQLDatabaseClient( + query_endpoint='https://example.test:8182', + neptune_iam=True, + region_name='eu-central-1', + ) + assert isinstance(store.client.store, NeptuneIAMStore) + assert store.client.store.query_endpoint == 'https://example.test:8182/sparql' + store.client.close() + + +def test_execute_query_runs_caller_supplied_sparql_directly(): + store = _store() + fake = _FakeClient() + store._client = fake + + rows = store._execute_query('SELECT ?l WHERE { VALUES ?l { "stmt-1" } }') + + assert rows == [{'l': 'stmt-1'}] + assert fake.queries == ['SELECT ?l WHERE { VALUES ?l { "stmt-1" } }'] + assert fake.updates == [] + + +def test_execute_query_rejects_silently_ignored_parameters(): + store = _store() + store._client = _FakeClient() + + with pytest.raises(ValueError, match='parameter binding is not supported'): + store._execute_query('SELECT * WHERE { ?s ?p ?o }', {'value': 'ignored'}) + + +def test_semantic_operation_uses_native_update_without_inspecting_query(): + store = _store() + fake = _FakeClient() + store._client = fake + + store._execute_operation( + GraphOperation.UPSERT_ENTITY, + 'intentionally not parsed', + {'params': [{'e_id': 'e1', 'v': 'Alice'}]}, + ) + + assert len(fake.updates) == 1 + assert 'INSERT DATA' in fake.updates[0] + assert 'entity/e1' in fake.updates[0] + assert 'tenant/default_' in fake.updates[0] + + +def test_semantic_read_operation_uses_native_query(): + store = _store() + fake = _FakeClient() + store._client = fake + + rows = store._execute_operation( + GraphOperation.SEARCH_BY_CHUNK, + 'also not parsed', + {'chunkId': 'c1', 'statementLimit': 3}, + ) + + assert rows == [{'l': 'stmt-1'}] + assert 'SELECT DISTINCT ?l' in fake.queries[0] + assert fake.default_graphs == [ + 'https://awslabs.github.io/graphrag-toolkit/lexical/tenant/default_' + ] + + +def test_empty_update_parameters_send_nothing(): + store = _store() + fake = _FakeClient() + store._client = fake + store._execute_operation( + GraphOperation.UPSERT_ENTITY, + '', + {'params': []}, + ) + assert fake.updates == [] + + +def test_getstate_drops_client(): + store = _store() + store._client = _FakeClient() + state = store.__getstate__() + assert store._client is None + assert state is not None + + +def test_exit_closes_and_clears_client(): + store = _store() + fake = _FakeClient() + store._client = fake + assert store.__exit__(None, None, None) is False + assert fake.closed is True + assert store._client is None + + +def test_execute_query_emits_debug_timing(caplog): + store = _store() + store._client = _FakeClient() + logger_name = 'graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_graph_store' + with caplog.at_level(logging.DEBUG, logger=logger_name): + store._execute_operation( + GraphOperation.SEARCH_BY_CHUNK, + '', + {'chunkId': 'c1'}, + ) + assert any('search_by_chunk' in record.getMessage() for record in caplog.records) + + +def test_namespace_kwargs_thread_into_writes_and_reads(): + store = _store( + lexical_prefix='gt', + lexical_schema_namespace='https://example.test/schema#', + lexical_instance_namespace='https://example.test/data/', + sparql_prefixes={'xsd': 'http://www.w3.org/2001/XMLSchema#'}, + ) + fake = _FakeClient() + store._client = fake + + store._execute_operation( + GraphOperation.UPSERT_ENTITY, + '', + {'params': [{'e_id': 'e1', 'v': 'Alice'}]}, + ) + assert '' in fake.updates[0] + + store._execute_operation( + GraphOperation.SEARCH_BY_ENTITIES, + '', + {'startId': 'e1', 'endIds': ['e2'], 'statementLimit': 3}, + ) + assert 'PREFIX gt: ' in fake.queries[0] diff --git a/lexical-graph-contrib/sparql/tests/test_sparql_graph_store_factory.py b/lexical-graph-contrib/sparql/tests/test_sparql_graph_store_factory.py new file mode 100644 index 00000000..92f59f6c --- /dev/null +++ b/lexical-graph-contrib/sparql/tests/test_sparql_graph_store_factory.py @@ -0,0 +1,92 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory + +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_graph_store import ( + SPARQLDatabaseClient, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_graph_store_factory import ( + SPARQLGraphStoreFactory, +) + + +def test_sparql_factory_creates_generic_endpoint_store(): + store = SPARQLGraphStoreFactory().try_create( + 'sparql+https://alice:secret@example.test/sparql/query', + update_endpoint='https://example.test/sparql/update', + ) + + assert isinstance(store, SPARQLDatabaseClient) + assert store.query_endpoint == 'https://example.test/sparql/query' + assert store.update_endpoint == 'https://example.test/sparql/update' + assert store.username == 'alice' + assert store.password.get_secret_value() == 'secret' + + +def test_sparql_factory_decodes_uri_credentials(): + store = SPARQLGraphStoreFactory().try_create( + 'sparql+https://alice%40example.test:p%40ss%3Aword@example.test/query' + ) + + assert store.username == 'alice@example.test' + assert store.password.get_secret_value() == 'p@ss:word' + + +def test_sparql_factory_ignores_non_sparql_urls(): + assert SPARQLGraphStoreFactory().try_create('https://example.test/sparql/query') is None + + +def test_sparql_factory_accepts_https_scheme(): + store = SPARQLGraphStoreFactory().try_create('sparql+s://example.test/sparql/query') + assert store.query_endpoint == 'https://example.test/sparql/query' + + +def test_sparql_factory_enables_neptune_iam(): + store = SPARQLGraphStoreFactory().try_create( + 'sparql+neptune://example.test:8182', + region_name='eu-central-1', + ) + assert store.query_endpoint == 'https://example.test:8182' + assert store.neptune_iam is True + assert store.region_name == 'eu-central-1' + + +def test_registered_factory_handles_real_neptune_hostname(): + GraphStoreFactory.register(SPARQLGraphStoreFactory) + store = GraphStoreFactory.for_graph_store( + 'sparql+neptune://cluster.us-east-1.neptune.amazonaws.com:8182', + region_name='us-east-1', + ) + assert isinstance(store, SPARQLDatabaseClient) + assert store.neptune_iam is True + + +def test_sparql_factory_consumes_auth_kwargs_when_uri_has_credentials(): + store = SPARQLGraphStoreFactory().try_create( + 'sparql+https://alice:secret@example.test/sparql/query', + username='bob', + password='ignored', + ) + assert store.username == 'alice' + assert store.password.get_secret_value() == 'secret' + + +def test_sparql_factory_preserves_endpoint_query_params(): + store = SPARQLGraphStoreFactory().try_create( + 'sparql+https://example.test/sparql/query?default-graph-uri=http%3A%2F%2Fexample.test%2Fg' + '&update_endpoint=https%3A%2F%2Fexample.test%2Fsparql%2Fupdate' + ) + assert store.query_endpoint == ( + 'https://example.test/sparql/query?default-graph-uri=http%3A%2F%2Fexample.test%2Fg' + ) + assert store.update_endpoint == 'https://example.test/sparql/update' + + +def test_factory_returns_none_for_non_string_input(): + assert SPARQLGraphStoreFactory().try_create(12345) is None + + +def test_factory_handles_ipv6_host_and_explicit_port(): + store = SPARQLGraphStoreFactory().try_create('sparql+http://[::1]:7200/repositories/lg') + assert store.query_endpoint == 'http://[::1]:7200/repositories/lg' diff --git a/lexical-graph-contrib/sparql/tests/test_sparql_queries.py b/lexical-graph-contrib/sparql/tests/test_sparql_queries.py new file mode 100644 index 00000000..c47871b2 --- /dev/null +++ b/lexical-graph-contrib/sparql/tests/test_sparql_queries.py @@ -0,0 +1,317 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation +from graphrag_toolkit.lexical_graph.versioning import ( + BUILD_TIMESTAMP, + EXTRACT_TIMESTAMP, + VALID_FROM, + VALID_TO, + VERSION_INDEPENDENT_ID_FIELDS, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.ontology import ( + DEFAULT_NAMESPACE, + LEXICAL_SCHEMA, + NamespaceConfig, + sparql_literal, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_queries import ( + _entity_score_rows, + _int_or_default, + _local_name, + _properties_by_id, + run_query, +) + + +class FakeClient: + def __init__(self): + self.queries = [] + self.default_graphs = [] + + def query(self, sparql, default_graph=None): + self.queries.append(sparql) + self.default_graphs.append(default_graph) + if 'SELECT DISTINCT ?statementId' in sparql: + return [{ + 'statementId': 'stmt-1', + 'statementValue': 'Alice manages Bob', + 'details': 'detail', + 'chunkId': 'chunk-1', + 'topicId': 'topic-1', + 'topicValue': 'People', + 'sourceId': 'source-1', + }] + if 'SELECT DISTINCT ?l' in sparql: + return [{'l': 'stmt-1'}] + if 'SELECT ?content' in sparql: + return [{'content': 'Chunk text'}] + if 'SELECT ?statement ?details' in sparql: + return [{'statement': 'Alice manages Bob', 'details': 'detail'}] + if 'SELECT ?entityId ?value ?class ?otherId' in sparql: + return [{ + 'entityId': 'entity-1', + 'value': 'Alice', + 'class': 'Person', + 'otherId': 'entity-2', + 'score': 3, + }] + if 'SELECT ?entityId ?value ?class' in sparql: + return [{ + 'entityId': 'entity-1', + 'value': 'Alice', + 'class': 'Person', + 'score': 2, + }] + if 'a lg:Source' in sparql: + return [ + {'id': 'source-1', 'prop': f'{LEXICAL_SCHEMA}title', 'value': 'Source title'}, + {'id': 'source-1', 'prop': f'{LEXICAL_SCHEMA}{VALID_FROM}', 'value': 10}, + {'id': 'source-1', 'prop': f'{LEXICAL_SCHEMA}{VALID_TO}', 'value': 20}, + {'id': 'source-1', 'prop': f'{LEXICAL_SCHEMA}{EXTRACT_TIMESTAMP}', 'value': 11}, + {'id': 'source-1', 'prop': f'{LEXICAL_SCHEMA}{BUILD_TIMESTAMP}', 'value': 12}, + { + 'id': 'source-1', + 'prop': f'{LEXICAL_SCHEMA}{VERSION_INDEPENDENT_ID_FIELDS}', + 'value': 'doc;rev', + }, + ] + if 'a lg:Chunk' in sparql: + return [{ + 'id': 'chunk-1', + 'prop': f'{LEXICAL_SCHEMA}value', + 'value': 'Chunk text', + }] + if 'SELECT ?statementId ?factValue' in sparql: + return [ + {'statementId': 's1', 'factValue': 'fa'}, + {'statementId': 's1', 'factValue': 'fa'}, + {'statementId': 's1'}, + ] + return [] + + +def test_statements_are_grouped_into_retriever_shape(): + client = FakeClient() + rows = run_query(GraphOperation.GET_STATEMENTS, client, { + 'statementIds': ['stmt-1'], + 'limit': 5, + '_include_chunk_details': True, + }) + + result = rows[0]['result'] + assert result['score'] == 1 + assert result['source']['metadata']['title'] == 'Source title' + assert result['source']['versioning'] == { + 'valid_from': 10, + 'valid_to': 20, + 'extract_timestamp': 11, + 'build_timestamp': 12, + 'id_fields': ['doc', 'rev'], + } + assert result['topics'][0]['chunks'][0]['metadata']['value'] == 'Chunk text' + assert result['topics'][0]['statements'][0]['statement'] == 'Alice manages Bob' + + +def test_facts_are_grouped_and_deduplicated(): + rows = run_query( + GraphOperation.GET_FACTS, + FakeClient(), + {'statementIds': ['s1']}, + ) + assert rows == [{'statementId': 's1', 'facts': ['fa']}] + + +def test_chunk_and_topic_content_queries(): + chunk_client = FakeClient() + topic_client = FakeClient() + + assert run_query( + GraphOperation.GET_CHUNKS, chunk_client, {'nodeIds': ['chunk-1']}, + ) == [{'content': 'Chunk text'}] + assert run_query( + GraphOperation.GET_TOPIC, + topic_client, + {'topicId': 'topic-1', 'statementLimit': 3}, + ) == [{'statement': 'Alice manages Bob', 'details': 'detail'}] + assert 'VALUES ?chunkId { "chunk-1" }' in chunk_client.queries[0] + assert '?fact lg:supports ?statementNode .' in topic_client.queries[0] + + +@pytest.mark.parametrize(('operation', 'parameters', 'query_fragment'), [ + (GraphOperation.SEARCH_BY_CHUNK, + {'chunkId': 'chunk-1', 'statementLimit': 3}, + 'lg:statementMentionedIn ?chunk'), + (GraphOperation.SEARCH_BY_CHUNK, + {'nodeId': 'chunk-1', 'statementLimit': 3}, + 'lg:id "chunk-1"'), + (GraphOperation.SEARCH_BY_TOPIC, + {'nodeId': 'topic-1', 'statementLimit': 3}, + 'lg:belongsTo ?topic'), + (GraphOperation.SEARCH_BY_ENTITY, + {'startId': 'e1', 'statementLimit': 3}, + 'lg:subject ?entity'), + (GraphOperation.SEARCH_BY_ENTITIES, + {'startId': 'e1', 'endIds': ['e2'], 'statementLimit': 3}, + 'lg:object'), +]) +def test_search_operations_return_statement_ids(operation, parameters, query_fragment): + client = FakeClient() + assert run_query(operation, client, parameters) == [{'l': 'stmt-1'}] + assert query_fragment in client.queries[0] + + +def test_keyword_lookup_uses_explicit_match_modes(): + exact = FakeClient() + prefix = FakeClient() + + run_query(GraphOperation.FIND_ENTITIES_BY_KEYWORD, exact, {'keyword': 'alice'}) + run_query(GraphOperation.FIND_ENTITIES_BY_KEYWORD, prefix, { + 'keyword': 'ali', + 'classification': 'Per', + '_starts_with': True, + '_classification_starts_with': True, + }) + + assert 'FILTER(?searchStr = "alice")' in exact.queries[0] + assert 'FILTER(STRSTARTS(?searchStr, "ali"))' in prefix.queries[0] + assert 'FILTER(STRSTARTS(?class, "Per"))' in prefix.queries[0] + + +@pytest.mark.parametrize(('operation', 'value_name', 'query_fragment'), [ + (GraphOperation.FIND_ENTITIES_BY_CHUNKS, 'chunkId', 'lg:statementMentionedIn ?chunk'), + (GraphOperation.FIND_ENTITIES_BY_TOPICS, 'topicId', 'lg:belongsTo ?topic'), +]) +def test_entity_lookups_from_nodes(operation, value_name, query_fragment): + client = FakeClient() + rows = run_query(operation, client, {'nodeIds': ['n1'], 'limit': 5}) + + assert rows[0]['result']['entity']['entityId'] == 'entity-1' + assert f'VALUES ?{value_name} {{ "n1" }}' in client.queries[0] + assert query_fragment in client.queries[0] + + +def test_entity_neighbours_and_scores_use_reified_facts(): + neighbours = FakeClient() + scores = FakeClient() + + neighbour_rows = run_query(GraphOperation.FIND_ENTITY_NEIGHBORS, neighbours, { + 'entityIds': ['entity-1'], + 'excludeEntityIds': ['entity-1'], + 'numNeighbours': 2, + }) + score_rows = run_query( + GraphOperation.SCORE_ENTITIES, + scores, + {'entityIds': ['entity-1']}, + ) + + assert neighbour_rows[0]['result']['others'] == ['entity-2'] + assert score_rows[0]['result']['score'] == 2.0 + assert 'lg:subject ?entity' in neighbours.queries[0] + assert 'lg:object ?other' in neighbours.queries[0] + + +def test_local_entity_lookup_queries_skip_incomplete_rows(): + complements = FakeClient() + subjects = FakeClient() + + run_query(GraphOperation.FIND_COMPLEMENTS, complements, { + 'params': [{'nId': 'n1'}, {'nId': None}], + }) + run_query(GraphOperation.FIND_SUBJECTS, subjects, { + 'params': [{'nId': 'n1', 'cId': 'c1'}, {'nId': None, 'cId': None}], + }) + + assert len(complements.queries) == 1 and 'search_str' in complements.queries[0] + assert len(subjects.queries) == 1 and 'SELECT ?n_id ?c_id' in subjects.queries[0] + + +def test_custom_namespace_and_tenant_default_graph_are_forwarded(): + namespace = NamespaceConfig( + prefix='gt', + schema_namespace='https://example.test/schema#', + instance_namespace='https://example.test/data/', + extra_prefixes={'xsd': 'http://www.w3.org/2001/XMLSchema#'}, + ) + client = FakeClient() + + run_query( + GraphOperation.SEARCH_BY_ENTITY, + client, + {'startId': 'e1'}, + namespace=namespace, + tenant_id='acme', + ) + + assert 'PREFIX gt: ' in client.queries[0] + assert 'PREFIX xsd: ' in client.queries[0] + assert client.default_graphs == ['https://example.test/data/tenant/acme'] + + +@pytest.mark.parametrize(('operation', 'parameters'), [ + (GraphOperation.GET_FACTS, {'statementIds': []}), + (GraphOperation.GET_CHUNKS, {'nodeIds': []}), + (GraphOperation.GET_TOPIC, {}), + (GraphOperation.SEARCH_BY_CHUNK, {}), + (GraphOperation.SEARCH_BY_TOPIC, {}), + (GraphOperation.SEARCH_BY_ENTITY, {}), + (GraphOperation.SEARCH_BY_ENTITIES, {'startId': 'a', 'endIds': []}), + (GraphOperation.FIND_ENTITIES_BY_KEYWORD, {}), + (GraphOperation.FIND_ENTITIES_BY_CHUNKS, {'nodeIds': []}), + (GraphOperation.FIND_ENTITIES_BY_TOPICS, {'nodeIds': []}), + (GraphOperation.FIND_ENTITY_NEIGHBORS, {'entityIds': []}), + (GraphOperation.SCORE_ENTITIES, {'entityIds': []}), + (GraphOperation.GET_STATEMENTS, {'statementIds': []}), + (GraphOperation.FIND_COMPLEMENTS, {'params': []}), + (GraphOperation.FIND_SUBJECTS, {'params': []}), +]) +def test_empty_inputs_return_no_results(operation, parameters): + assert run_query(operation, FakeClient(), parameters) == [] + + +def test_update_operation_cannot_be_run_as_query(): + with pytest.raises(NotImplementedError, match='upsert_chunk'): + run_query(GraphOperation.UPSERT_CHUNK, FakeClient(), {'chunk_id': 'c1'}) + + +def test_sparql_values_are_escaped_and_unsafe_namespaces_rejected(): + assert sparql_literal('a" . } DELETE { ?s ?p ?o } #') == ( + '"a\\" . } DELETE { ?s ?p ?o } #"' + ) + with pytest.raises(ValueError): + NamespaceConfig(schema_namespace='https://x/ns#>\nINSERT DATA { } #') + with pytest.raises(ValueError, match='must be absolute'): + NamespaceConfig(schema_namespace='relative/schema') + + +@pytest.mark.parametrize(('operation', 'parameters'), [ + (GraphOperation.SEARCH_BY_CHUNK, {'chunkId': 'c1', 'statementLimit': 0}), + (GraphOperation.GET_TOPIC, {'topicId': 't1', 'statementLimit': -1}), + (GraphOperation.FIND_ENTITIES_BY_CHUNKS, {'nodeIds': ['c1'], 'limit': 0}), + (GraphOperation.FIND_ENTITY_NEIGHBORS, {'entityIds': ['e1'], 'numNeighbours': 0}), + (GraphOperation.GET_STATEMENTS, {'statementIds': ['s1'], 'limit': 0}), +]) +def test_numeric_query_clauses_must_be_positive(operation, parameters): + with pytest.raises(ValueError, match='must be greater than zero'): + run_query(operation, FakeClient(), parameters) + + +def test_result_helpers_handle_missing_and_invalid_values(): + assert _entity_score_rows([ + {'entityId': None}, + {'entityId': 'e1', 'value': 'v', 'class': 'c', 'score': 2}, + ]) == [{ + 'result': { + 'entity': {'entityId': 'e1', 'value': 'v', 'class': 'c'}, + 'score': 2.0, + }, + }] + assert _properties_by_id(FakeClient(), 'Source', [], DEFAULT_NAMESPACE) == {} + assert _int_or_default(None, 5) == 5 + assert _int_or_default('nope', 5) == 5 + assert _int_or_default('7', 0) == 7 + assert _local_name(None, DEFAULT_NAMESPACE) == '' + assert _local_name('http://other.test/foo#bar', DEFAULT_NAMESPACE) == 'bar' diff --git a/lexical-graph-contrib/sparql/tests/test_sparql_updates.py b/lexical-graph-contrib/sparql/tests/test_sparql_updates.py new file mode 100644 index 00000000..69d157e0 --- /dev/null +++ b/lexical-graph-contrib/sparql/tests/test_sparql_updates.py @@ -0,0 +1,223 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for native lexical-graph SPARQL updates.""" + +import re + +import pytest +from rdflib import Dataset, URIRef + +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.ontology import ( + NamespaceConfig, +) +from graphrag_toolkit_contrib.lexical_graph.storage.graph.sparql.sparql_updates import ( + render_update, +) + + +def _rows(row): + return {'params': [row]} + + +def test_source_uses_explicit_id_and_metadata(): + update = render_update( + GraphOperation.UPSERT_SOURCE, + _rows({'_source_id': 'aws::abc:def', 'url': 'https://x/y'}), + ) + + assert 'source/aws%3A%3Aabc%3Adef' in update + assert 'https://x/y' in update + assert 'DELETE WHERE' in update + assert '_source_id' not in update + + +@pytest.mark.parametrize(('operation', 'params', 'expected'), [ + (GraphOperation.UPSERT_CHUNK, + {'chunk_id': 'c1', 'text': 'hello', 'seq': 3}, + ('chunk/c1', '"hello"')), + (GraphOperation.UPSERT_TOPIC, + {'topic_id': 't1', 'title': 'T', 'chunk_ids': ['c1']}, + ('topic/t1', 'topicMentionedIn', 'chunk/c1')), + (GraphOperation.UPSERT_STATEMENT, + {'statement_id': 's1', 'value': 'v', 'details': 'd'}, + ('statement/s1', '"v"', '"d"')), + (GraphOperation.UPSERT_ENTITY, + {'e_id': 'e1', 'v': 'Alice', 'e_search_str': 'alice', 'ec': 'Person'}, + ('entity/e1', '"Alice"', '"alice"', '"Person"')), +]) +def test_node_updates(operation, params, expected): + update = render_update(operation, _rows(params)) + assert all(value in update for value in expected) + + +def test_chunk_internal_parameters_are_not_written_as_metadata(): + update = render_update( + GraphOperation.UPSERT_CHUNK, + _rows({'chunk_id': 'c1', 'text': 'hello', '_internal': 'private'}), + ) + + assert '_internal' not in update + + +@pytest.mark.parametrize(('operation', 'params', 'predicate'), [ + (GraphOperation.LINK_CHUNK_SOURCE, + {'chunk_id': 'c1', 'source_id': 's1'}, 'extractedFrom'), + (GraphOperation.LINK_CHUNKS, + {'chunk_id': 'c2', 'target_id': 'c1', '_relationship_type': 'previous'}, + 'chunkPrevious'), + (GraphOperation.LINK_STATEMENT_CHUNK, + {'statement_id': 's1', 'chunk_id': 'c1'}, 'statementMentionedIn'), + (GraphOperation.LINK_STATEMENT_TOPIC, + {'statement_id': 's1', 'topic_id': 't1'}, 'belongsTo'), + (GraphOperation.LINK_STATEMENTS, + {'statement_id': 's2', 'prev_statement_id': 's1'}, 'statementPrevious'), +]) +def test_link_updates(operation, params, predicate): + assert predicate in render_update(operation, _rows(params)) + + +def test_fact_is_reified_and_links_entities_from_fact(): + update = render_update(GraphOperation.UPSERT_FACT, _rows({ + 'statement_id': 's1', + 'fact_id': 'f1', + 'fact': 'a@x CONFIDENCE 95%', + '_predicate': 'CONFIDENCE', + '_subject_literal': 'a@x', + '_object_literal': '95%', + })) + edge = render_update(GraphOperation.LINK_FACT_ENTITY, _rows({ + 'fact_id': 'f1', + 'entity_id': 'e1', + '_relationship_type': 'subject', + })) + + assert 'Relation>' in update and 'relation/confidence>' in update + assert '#subject>' in update and '"a@x"' in update + assert '#object>' in update and '"95%"' in update + assert 'supports' in update and 'statement/s1' in update + assert re.search(r'fact/f1>\s+<[^>]*#subject>\s+<[^>]*entity/e1>', edge) + + +def test_lpg_convenience_relation_is_not_duplicated_in_rdf(): + assert render_update( + GraphOperation.LINK_ENTITIES, + _rows({'s_id': 'a', 'o_id': 'b', 'p': 'manages'}), + ) is None + + +def test_graph_summary_updates_distinct_class_counters(): + update = render_update(GraphOperation.UPDATE_GRAPH_SUMMARY, _rows({ + 'sc_id': 'sys::Person', + 'oc_id': 'sys::Company', + 'sc': 'Person', + 'oc': 'Company', + 'p': 'MANAGES', + })) + + assert 'SysRelation>' in update + assert 'sysRelSubject>' in update and 'sysRelObject>' in update + assert 'COALESCE' in update and update.count('DELETE') >= 2 + + +def test_graph_summary_same_class_uses_increment_of_two(): + update = render_update(GraphOperation.UPDATE_GRAPH_SUMMARY, _rows({ + 'sc_id': 'sys::Person', + 'oc_id': 'sys::Person', + 'sc': 'Person', + 'oc': 'Person', + 'p': 'KNOWS', + })) + + assert 'COALESCE(?c0, 0) + 2' in update + + +def test_domain_type_and_custom_namespace(): + namespace = NamespaceConfig( + prefix='gt', + schema_namespace='https://example.test/schema#', + instance_namespace='https://example.test/data/', + ) + update = render_update( + GraphOperation.ADD_ENTITY_TYPE, + {'entityId': 'e1', '_classification': 'Person'}, + namespace, + ) + + assert '' in update + assert '' in update + + +def test_tenant_routes_update_to_named_graph(): + update = render_update( + GraphOperation.UPSERT_ENTITY, + _rows({'e_id': 'e1', 'v': 'Alice'}), + tenant_id='acme', + ) + assert 'GRAPH ' in update + + +@pytest.mark.parametrize('parameters', [None, {}, {'params': []}]) +def test_empty_parameters_are_noop(parameters): + assert render_update(GraphOperation.UPSERT_CHUNK, parameters) is None + + +def test_missing_edge_parameter_is_noop(): + assert render_update( + GraphOperation.LINK_CHUNK_SOURCE, + _rows({'chunk_id': 'c1'}), + ) is None + + +def test_local_entity_rewrite_preserves_reified_fact(): + update = render_update( + GraphOperation.COPY_COMPLEMENT_RELATIONSHIPS, + _rows({'n_id': 'resolved', 'c_id': 'local'}), + ) + delete = render_update( + GraphOperation.DELETE_COMPLEMENT, + _rows({'c_id': 'local'}), + ) + + assert 'DELETE' in update and 'INSERT' in update + assert 'entity/local' in update and 'entity/resolved' in update + assert '#object>' in update + assert 'entity/local' in delete and 'DELETE WHERE' in delete + + +@pytest.mark.filterwarnings( + 'ignore:Dataset.default_context is deprecated:DeprecationWarning' +) +def test_local_entity_rewrite_moves_fact_object_when_delete_runs_first(): + dataset = Dataset() + tenant = 'acme' + for operation, row in ( + (GraphOperation.UPSERT_ENTITY, + {'e_id': 'resolved', 'v': 'Resolved', 'ec': 'Entity'}), + (GraphOperation.UPSERT_ENTITY, + {'e_id': 'local', 'v': 'Local', 'ec': '__LocalEntity__'}), + (GraphOperation.LINK_FACT_ENTITY, + {'fact_id': 'f1', 'entity_id': 'local', '_relationship_type': 'object'}), + (GraphOperation.DELETE_COMPLEMENT, {'c_id': 'local'}), + (GraphOperation.COPY_COMPLEMENT_RELATIONSHIPS, + {'n_id': 'resolved', 'c_id': 'local'}), + ): + dataset.update(render_update(operation, _rows(row), tenant_id=tenant)) + + graph = dataset.graph(URIRef( + 'https://awslabs.github.io/graphrag-toolkit/lexical/tenant/acme' + )) + fact = URIRef('https://awslabs.github.io/graphrag-toolkit/lexical/fact/f1') + predicate = URIRef('https://awslabs.github.io/graphrag-toolkit/lexical#object') + resolved = URIRef('https://awslabs.github.io/graphrag-toolkit/lexical/entity/resolved') + local = URIRef('https://awslabs.github.io/graphrag-toolkit/lexical/entity/local') + + assert (fact, predicate, resolved) in graph + assert not list(graph.triples((None, None, local))) + assert not list(graph.triples((local, None, None))) + + +def test_read_operation_cannot_be_rendered_as_update(): + with pytest.raises(NotImplementedError, match='get_facts'): + render_update(GraphOperation.GET_FACTS, {'statementIds': ['s1']}) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/chunk_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/chunk_graph_builder.py index 3ace9de5..4db37e00 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/chunk_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/chunk_graph_builder.py @@ -4,7 +4,7 @@ import logging from typing import Any -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder from llama_index.core.schema import BaseNode @@ -83,7 +83,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query_c = '\n'.join(statements_c) - graph_client.execute_query_with_retry(query_c, self._to_params(properties_c), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query_c, self._to_params(properties_c), max_attempts=5, max_wait=7, operation=GraphOperation.UPSERT_CHUNK) source_info = node.relationships.get(NodeRelationship.SOURCE, None) @@ -106,7 +106,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query_s = '\n'.join(statements_s) - graph_client.execute_query_with_retry(query_s, self._to_params(properties_s), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query_s, self._to_params(properties_s), max_attempts=5, max_wait=7, operation=GraphOperation.LINK_CHUNK_SOURCE) else: logger.warning(f'source_id missing from chunk node [node_id: {chunk_id}]') @@ -124,12 +124,13 @@ def insert_chunk_to_chunk_relationship(node_id:str, relationship_type:str): properties_c2c = { 'chunk_id': chunk_id, - 'target_id': node_id + 'target_id': node_id, + '_relationship_type': relationship_type, } query_c2c = '\n'.join(statements_c2c) - graph_client.execute_query_with_retry(query_c2c, self._to_params(properties_c2c), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query_c2c, self._to_params(properties_c2c), max_attempts=5, max_wait=7, operation=GraphOperation.LINK_CHUNKS) for node_relationship,relationship_info in node.relationships.items(): 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 2065239b..5acd6d96 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 @@ -5,7 +5,7 @@ from typing import Any from graphrag_toolkit.lexical_graph.indexing.model import Fact, Entity -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, 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 @@ -98,7 +98,7 @@ def insert_for_entity(entity:Entity): query = '\n'.join(statements) - graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphOperation.UPSERT_ENTITY) insert_for_entity(fact.subject) @@ -125,7 +125,7 @@ def insert_domain_entity(entity:Entity): e_label = escape_cypher_label(label_from(entity.classification or DEFAULT_CLASSIFICATION)) e_comment = f'// awsqid:{e_id}-{e_label}'.replace('\r', ' ').replace('\n', ' ') query_e = f"MERGE ({e_var}:`__Entity__`{{{graph_client.node_id('entityId')}: $entityId}}) SET {e_var} :`{e_label}` {e_comment}" - graph_client.execute_query_with_retry(query_e, {'entityId': e_id}, max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query_e, {'entityId': e_id, '_classification': entity.classification or DEFAULT_CLASSIFICATION}, max_attempts=5, max_wait=7, operation=GraphOperation.ADD_ENTITY_TYPE) insert_domain_entity(fact.subject) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_relation_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_relation_graph_builder.py index ad196c6d..b09be19e 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_relation_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_relation_graph_builder.py @@ -5,7 +5,7 @@ from typing import Any from graphrag_toolkit.lexical_graph.indexing.model import Fact -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import relationship_name_from, new_query_var 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 @@ -90,7 +90,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query = '\n'.join(statements) - graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphOperation.LINK_ENTITIES) # if include_domain_labels: @@ -137,7 +137,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query = '\n'.join(statements) - graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphOperation.LINK_ENTITIES) # if include_domain_labels: diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/fact_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/fact_graph_builder.py index 637cb1f7..80cebfa6 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/fact_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/fact_graph_builder.py @@ -5,7 +5,7 @@ from typing import Any, Optional from graphrag_toolkit.lexical_graph.indexing.model import Fact -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore, Query, QueryTree +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore, Query, QueryTree from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder from graphrag_toolkit.lexical_graph.indexing.constants import LOCAL_ENTITY_CLASSIFICATION from graphrag_toolkit.lexical_graph.indexing.utils.fact_utils import string_complement_to_entity @@ -82,15 +82,26 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): 'MERGE (fact)-[:`__SUPPORTS__`]->(statement)' ] + subject_literal = None + if fact.subject.classification == LOCAL_ENTITY_CLASSIFICATION and not include_local_entities: + subject_literal = fact.subject.value + + object_literal = None + if not fact.object and fact.complement and not include_local_entities: + object_literal = fact.complement.value + properties = { 'statement_id': fact.statementId, 'fact_id': fact.factId, - 'fact': node.text + 'fact': node.text, + '_predicate': fact.predicate.value, + '_subject_literal': subject_literal, + '_object_literal': object_literal } query = '\n'.join(statements) - graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphOperation.UPSERT_FACT) def insert_entity_fact_relationship(relationship_type:str, entity_id:Optional[str]=None): @@ -107,10 +118,11 @@ def insert_entity_fact_relationship(relationship_type:str, entity_id:Optional[st if entity_id: properties_e2f['fact_id'] = fact.factId properties_e2f['entity_id'] = entity_id + properties_e2f['_relationship_type'] = relationship_type query_e2f = '\n'.join(statements_e2f) - graph_client.execute_query_with_retry(query_e2f, self._to_params(properties_e2f), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query_e2f, self._to_params(properties_e2f), max_attempts=5, max_wait=7, operation=GraphOperation.LINK_FACT_ENTITY) insert_entity_fact_relationship('subject') diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_batch_client.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_batch_client.py index ba09e29f..9944a3c8 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_batch_client.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_batch_client.py @@ -46,6 +46,7 @@ def __init__(self, graph_client:GraphStore, batch_writes_enabled:bool, batch_wri self.batch_writes_enabled = batch_writes_enabled self.batch_write_size = batch_write_size self.batches:Dict[str, List] = {} + self.batch_operations:Dict[str, Any] = {} self.query_trees:Dict[str, QueryTree] = {} self.all_nodes = [] self.parameterless_queries:Dict[str, str] = {} @@ -122,8 +123,9 @@ def execute_query_with_retry(self, query:QueryTree, properties:Dict[str, Any], * if properties: if query not in self.batches: self.batches[query] = [] + self.batch_operations[query] = kwargs.get('operation') self.batches[query].extend(properties['params']) - else: + elif kwargs.get('operation') is None: self._add_parameterless_query(query) elif isinstance(query, QueryTree): properties = properties or {'params':[]} @@ -188,7 +190,7 @@ def _apply_batch_query(self, query, parameters): 'params': p } try: - self.graph_client.execute_query_with_retry(query, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT) + self._execute_batch(query, params) except Exception as e: logger.debug(f'Batch failed - queuing for retry: [query: {query}, params: {params}]') retry_batches.append((query, params)) @@ -197,12 +199,21 @@ def _apply_batch_query(self, query, parameters): for (query, params) in retry_batches: try: - self.graph_client.execute_query_with_retry(query, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT) + self._execute_batch(query, params) except Exception as e: logger.debug(f'Retry batch failed - queuing for return: [query: {query}, params: {params}]') failed_batches.append((query, params)) return failed_batches + + def _execute_batch(self, query, parameters): + return self.graph_client.execute_query_with_retry( + query, + parameters, + max_attempts=BATCH_MAX_ATTEMPTS, + max_wait=BATCH_MAX_WAIT, + operation=self.batch_operations.get(query), + ) def _retry_failed_batches(self, failed_batches): @@ -210,7 +221,7 @@ def _retry_failed_batches(self, failed_batches): for (query, params) in failed_batches: try: - self.graph_client.execute_query_with_retry(query, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT) + self._execute_batch(query, params) except Exception as e: logger.debug(f'Retry failed batch failed - queuing for individual writes retry: [query: {query}, params: {params}]') last_chance_batches.append((query, params)) @@ -223,7 +234,7 @@ def _retry_failed_batches(self, failed_batches): 'params': [p] } try: - self.graph_client.execute_query_with_retry(query, single_params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT) + self._execute_batch(query, single_params) except Exception as e: logger.error(f'Failed single write: [query: {query}, params: {params}, error: {str(e)}]') raise e @@ -233,7 +244,7 @@ def _apply_batch_query_tree(self, query_tree_id, parameters): query_tree = self.query_trees[query_tree_id] - def graph_store_op(q, p): + def graph_store_op(q, p, **kwargs): all_params = p['params'] @@ -248,7 +259,7 @@ def graph_store_op(q, p): 'params': chunk } - results = self.graph_client.execute_query_with_retry(q, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT) + results = self.graph_client.execute_query_with_retry(q, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT, **kwargs) for r in results: yield r diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_summary_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_summary_builder.py index 18c65806..7e662fc2 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_summary_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_summary_builder.py @@ -5,7 +5,7 @@ from typing import Any from graphrag_toolkit.lexical_graph.indexing.model import Fact -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import label_from, relationship_name_from from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder from graphrag_toolkit.lexical_graph.indexing.constants import DEFAULT_CLASSIFICATION, LOCAL_ENTITY_CLASSIFICATION @@ -116,7 +116,7 @@ def build(self, node:BaseNode, graph_client:GraphStore, **kwargs:Any): query = '\n'.join(statements) - graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=10, max_wait=10) + graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=10, max_wait=10, operation=GraphOperation.UPDATE_GRAPH_SUMMARY) else: logger.warning(f'fact_id missing from fact node [node_id: {node.node_id}]') \ No newline at end of file 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..74eea2a8 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 @@ -5,7 +5,7 @@ from typing import Any from graphrag_toolkit.lexical_graph.indexing.model import Fact -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore, Query, QueryTree +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, 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 @@ -43,7 +43,8 @@ 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) - """ + """, + operation=GraphOperation.COPY_COMPLEMENT_RELATIONSHIPS, ) delete_complement_relationships = Query( @@ -54,7 +55,8 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): DELETE r1 DELETE r2 DETACH DELETE c - """ + """, + operation=GraphOperation.DELETE_COMPLEMENT, ) if fact.subject: @@ -72,7 +74,8 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): child_queries=[ copy_complement_relationships_to_subject, delete_complement_relationships - ] + ], + operation=GraphOperation.FIND_COMPLEMENTS, ) params = { @@ -98,7 +101,8 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): child_queries=[ copy_complement_relationships_to_subject, delete_complement_relationships - ] + ], + operation=GraphOperation.FIND_SUBJECTS, ) params = { diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/source_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/source_graph_builder.py index 51ff59a8..89eb96c1 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/source_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/source_graph_builder.py @@ -4,7 +4,7 @@ import logging from typing import Any -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder from graphrag_toolkit.lexical_graph.versioning import VALID_FROM, VALID_TO, VERSION_INDEPENDENT_ID_FIELDS from graphrag_toolkit.lexical_graph.versioning import EXTRACT_TIMESTAMP, BUILD_TIMESTAMP, PREV_VERSIONS @@ -111,7 +111,8 @@ def format_assigment(key): query = '\n'.join(statements) - graph_client.execute_query_with_retry(query, self._to_params(clean_metadata)) + clean_metadata['_source_id'] = source_id + graph_client.execute_query_with_retry(query, self._to_params(clean_metadata), operation=GraphOperation.UPSERT_SOURCE) # prev_source_ids = source_metadata.get('prev_versions', []) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/statement_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/statement_graph_builder.py index f075dfd0..bc924088 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/statement_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/statement_graph_builder.py @@ -5,7 +5,7 @@ from typing import Any from graphrag_toolkit.lexical_graph.indexing.model import Statement -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder from llama_index.core.schema import BaseNode @@ -83,7 +83,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query = '\n'.join(statements) - graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphOperation.UPSERT_STATEMENT) if statement.chunkId: @@ -103,7 +103,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query_c = '\n'.join(statements_c) - graph_client.execute_query_with_retry(query_c, self._to_params(properties_c), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query_c, self._to_params(properties_c), max_attempts=5, max_wait=7, operation=GraphOperation.LINK_STATEMENT_CHUNK) if statement.topicId: @@ -122,7 +122,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query_t = '\n'.join(statements_t) - graph_client.execute_query_with_retry(query_t, self._to_params(properties_t), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query_t, self._to_params(properties_t), max_attempts=5, max_wait=7, operation=GraphOperation.LINK_STATEMENT_TOPIC) if prev_statement: @@ -141,7 +141,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query_p = '\n'.join(statements_p) - graph_client.execute_query_with_retry(query_p, self._to_params(properties_p), max_attempts=5, max_wait=7) + graph_client.execute_query_with_retry(query_p, self._to_params(properties_p), max_attempts=5, max_wait=7, operation=GraphOperation.LINK_STATEMENTS) else: logger.warning(f'statement_id missing from statement node [node_id: {node.node_id}]') \ No newline at end of file diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/topic_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/topic_graph_builder.py index bb588ab9..ff30ac4d 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/topic_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/topic_graph_builder.py @@ -5,7 +5,7 @@ from typing import Any from graphrag_toolkit.lexical_graph.indexing.model import Topic -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder from llama_index.core.schema import BaseNode @@ -88,7 +88,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any): query = '\n'.join(statements) - graph_client.execute_query_with_retry(query, self._to_params(properties)) + graph_client.execute_query_with_retry(query, self._to_params(properties), operation=GraphOperation.UPSERT_TOPIC) else: logger.warning(f'topic_id missing from topic node [node_id: {node.node_id}]') diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_context_provider.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_context_provider.py index 313dde94..2f343798 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_context_provider.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_context_provider.py @@ -6,7 +6,7 @@ import json from typing import List, Dict -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import node_result from graphrag_toolkit.lexical_graph.retrieval.model import ScoredEntity, EntityContexts, EntityContext from graphrag_toolkit.lexical_graph.retrieval.processors import ProcessorArgs @@ -69,7 +69,7 @@ def _get_entity_id_context_tree(self, entities:List[ScoredEntity]) -> Dict[str, 'numNeighbours': depth + 2 } - results = self.graph_store.execute_query(cypher, params) + results = self.graph_store.execute_query(cypher, params, operation=GraphOperation.FIND_ENTITY_NEIGHBORS) new_entity_id_contexts = {} @@ -138,7 +138,7 @@ def walk_tree(d): 'entityIds': list(neighbour_entity_ids) } - results = self.graph_store.execute_query(cypher, params) + results = self.graph_store.execute_query(cypher, params, operation=GraphOperation.SCORE_ENTITIES) neighbour_entities = [ ScoredEntity.model_validate(result['result']) for result in results diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_provider.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_provider.py index 88498b15..6b784e2b 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_provider.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_provider.py @@ -4,7 +4,7 @@ import logging from typing import List, Iterator, cast, Optional -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.metadata import FilterConfig from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import node_result, search_string_from, filter_config_to_opencypher_filters from graphrag_toolkit.lexical_graph.retrieval.model import ScoredEntity @@ -57,7 +57,7 @@ def _get_entities_for_keyword(self, keyword:str) -> List[ScoredEntity]: 'keyword': search_string_from(parts[0]) } - results = self.graph_store.execute_query(cypher, params) + results = self.graph_store.execute_query(cypher, params, operation=GraphOperation.FIND_ENTITIES_BY_KEYWORD) entities = [ ScoredEntity.model_validate(result['result']) @@ -80,7 +80,9 @@ def _get_entities_for_keyword(self, keyword:str) -> List[ScoredEntity]: params = { 'keyword': search_string_from(parts[0]), - 'classification': parts[1] + 'classification': parts[1], + '_starts_with': True, + '_classification_starts_with': True, } else: cypher = f""" @@ -95,10 +97,11 @@ def _get_entities_for_keyword(self, keyword:str) -> List[ScoredEntity]: }} AS result""" params = { - 'keyword': search_string_from(parts[0]) + 'keyword': search_string_from(parts[0]), + '_starts_with': True, } - results = self.graph_store.execute_query(cypher, params) + results = self.graph_store.execute_query(cypher, params, operation=GraphOperation.FIND_ENTITIES_BY_KEYWORD) entities = [ ScoredEntity.model_validate(result['result']) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_vss_provider.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_vss_provider.py index 9c7b7667..96703ba5 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_vss_provider.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/entity_vss_provider.py @@ -5,7 +5,7 @@ from typing import List, Optional from graphrag_toolkit.lexical_graph.metadata import FilterConfig -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.storage.vector import VectorStore from graphrag_toolkit.lexical_graph.storage.vector import DummyVectorIndex from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import node_result @@ -43,6 +43,7 @@ def _get_node_ids(self, keywords:List[str]) -> List[str]: def _get_entities_for_nodes(self, node_ids:List[str]) -> List[ScoredEntity]: if self.index_name == 'topic': + operation = GraphOperation.FIND_ENTITIES_BY_TOPICS cypher = f""" // get entities for topic ids MATCH (t:`__Topic__`)<-[:`__BELONGS_TO__`]-(:`__Statement__`) @@ -58,6 +59,7 @@ def _get_entities_for_nodes(self, node_ids:List[str]) -> List[ScoredEntity]: }} AS result """ else: + operation = GraphOperation.FIND_ENTITIES_BY_CHUNKS cypher = f""" // get entities for chunk ids MATCH (c:`__Chunk__`)<-[:`__MENTIONED_IN__`]-(:`__Statement__`) @@ -78,7 +80,7 @@ def _get_entities_for_nodes(self, node_ids:List[str]) -> List[ScoredEntity]: 'limit': self.args.intermediate_limit } - results = self.graph_store.execute_query(cypher, parameters) + results = self.graph_store.execute_query(cypher, parameters, operation=operation) scored_entities = [ ScoredEntity.model_validate(result['result']) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/keyword_vss_provider.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/keyword_vss_provider.py index 5a0e61cf..bbfbf063 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/keyword_vss_provider.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/query_context/keyword_vss_provider.py @@ -7,7 +7,7 @@ from graphrag_toolkit.lexical_graph.config import GraphRAGConfig from graphrag_toolkit.lexical_graph.utils import LLMCache, LLMCacheType from graphrag_toolkit.lexical_graph.metadata import FilterConfig -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.storage.vector import VectorStore from graphrag_toolkit.lexical_graph.storage.vector import DummyVectorIndex from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import node_result @@ -85,7 +85,7 @@ def _get_chunk_content(self, node_ids:List[str]) -> List[str]: 'nodeIds': node_ids } - results = self.graph_store.execute_query(cypher, parameters) + results = self.graph_store.execute_query(cypher, parameters, operation=GraphOperation.GET_CHUNKS) content = [result['content'] for result in results] @@ -116,7 +116,7 @@ def get_statements_for_topic(topic_id): 'statementLimit': self.args.intermediate_limit } - results = self.graph_store.execute_query(cypher, parameters) + results = self.graph_store.execute_query(cypher, parameters, operation=GraphOperation.GET_TOPIC) return '\n'.join(format_statement(r) for r in results) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/chunk_based_search.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/chunk_based_search.py index c470d0f6..2261f7de 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/chunk_based_search.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/chunk_based_search.py @@ -8,7 +8,7 @@ from graphrag_toolkit.lexical_graph.metadata import FilterConfig from graphrag_toolkit.lexical_graph.retrieval.model import SearchResultCollection from graphrag_toolkit.lexical_graph.storage.vector.vector_store import VectorStore -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.retrieval.processors import ProcessorBase, ProcessorArgs from graphrag_toolkit.lexical_graph.retrieval.retrievers.traversal_based_base_retriever import TraversalBasedBaseRetriever from graphrag_toolkit.lexical_graph.retrieval.utils.vector_utils import get_diverse_vss_elements @@ -96,7 +96,7 @@ def chunk_based_graph_search(self, chunk_id): 'statementLimit': self.args.intermediate_limit } - results = self.graph_store.execute_query(cypher, properties) + results = self.graph_store.execute_query(cypher, properties, operation=GraphOperation.SEARCH_BY_CHUNK) statement_ids = [r['l'] for r in results] return statement_ids diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/entity_based_search.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/entity_based_search.py index f71e9799..17b45efe 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/entity_based_search.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/entity_based_search.py @@ -7,7 +7,7 @@ from graphrag_toolkit.lexical_graph.metadata import FilterConfig from graphrag_toolkit.lexical_graph.retrieval.model import SearchResultCollection -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.storage.vector.vector_store import VectorStore from graphrag_toolkit.lexical_graph.retrieval.processors import ProcessorBase, ProcessorArgs from graphrag_toolkit.lexical_graph.retrieval.retrievers.traversal_based_base_retriever import TraversalBasedBaseRetriever @@ -164,7 +164,7 @@ def _multiple_entity_based_graph_search(self, start_id, end_ids, query:QueryBund 'statementLimit': self.args.intermediate_limit } - results = self.graph_store.execute_query(cypher, properties) + results = self.graph_store.execute_query(cypher, properties, operation=GraphOperation.SEARCH_BY_ENTITIES) statement_ids = [r['l'] for r in results] return statement_ids @@ -200,7 +200,7 @@ def _single_entity_based_graph_search(self, entity_id, query:QueryBundle): 'statementLimit': self.args.intermediate_limit } - results = self.graph_store.execute_query(cypher, properties) + results = self.graph_store.execute_query(cypher, properties, operation=GraphOperation.SEARCH_BY_ENTITY) statement_ids = [r['l'] for r in results] return statement_ids diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/entity_network_search.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/entity_network_search.py index 33d32a00..0c7d7673 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/entity_network_search.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/entity_network_search.py @@ -8,7 +8,7 @@ from graphrag_toolkit.lexical_graph.metadata import FilterConfig from graphrag_toolkit.lexical_graph.retrieval.model import SearchResultCollection -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.storage.vector.vector_store import VectorStore from graphrag_toolkit.lexical_graph.storage.vector.dummy_vector_index import DummyVectorIndex from graphrag_toolkit.lexical_graph.retrieval.processors import ProcessorBase, ProcessorArgs @@ -62,12 +62,14 @@ def __init__(self, def _graph_search(self, node_id): if self.index_name == 'topic': + operation = GraphOperation.SEARCH_BY_TOPIC cypher = f'''// topic-based entity network search MATCH (l)-[:`__BELONGS_TO__`]->(t:`__Topic__`) WHERE {self.graph_store.node_id("t.topicId")} = $nodeId RETURN DISTINCT {self.graph_store.node_id("l.statementId")} AS l LIMIT $statementLimit ''' else: + operation = GraphOperation.SEARCH_BY_CHUNK cypher = f'''// chunk-based entity network search MATCH (l)-[:`__BELONGS_TO__`]->()-[:`__MENTIONED_IN__`]->(c:`__Chunk__`) WHERE {self.graph_store.node_id("c.chunkId")} = $nodeId @@ -79,7 +81,7 @@ def _graph_search(self, node_id): 'statementLimit': self.args.intermediate_limit } - results = self.graph_store.execute_query(cypher, properties) + results = self.graph_store.execute_query(cypher, properties, operation=operation) statement_ids = [r['l'] for r in results] return statement_ids diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/traversal_based_base_retriever.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/traversal_based_base_retriever.py index 49158fff..4bf19ecc 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/traversal_based_base_retriever.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/retrievers/traversal_based_base_retriever.py @@ -9,7 +9,7 @@ from graphrag_toolkit.lexical_graph.metadata import FilterConfig from graphrag_toolkit.lexical_graph.versioning import VALID_FROM, VALID_TO, EXTRACT_TIMESTAMP, BUILD_TIMESTAMP, VERSION_INDEPENDENT_ID_FIELDS, TIMESTAMP_LOWER_BOUND, TIMESTAMP_UPPER_BOUND -from graphrag_toolkit.lexical_graph.storage.graph import GraphStore +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, GraphStore from graphrag_toolkit.lexical_graph.storage.vector.vector_store import VectorStore from graphrag_toolkit.lexical_graph.retrieval.query_context import KeywordProvider, KeywordVSSProvider, KeywordNLPProvider, KeywordProviderMode, PassThruKeywordProvider from graphrag_toolkit.lexical_graph.retrieval.query_context import EntityProvider, EntityVSSProvider, EntityContextProvider @@ -145,7 +145,8 @@ def get_statements_by_topic_and_source(self, statement_ids): statements_params = { 'statementLimit': self.args.intermediate_limit, 'limit': self.args.query_limit, - 'statementIds': statement_ids + 'statementIds': statement_ids, + '_include_chunk_details': self.args.include_chunk_details, } chunk_metadata = 'properties(c)' if self.args.include_chunk_details else '{}' @@ -188,7 +189,7 @@ def get_statements_by_topic_and_source(self, statement_ids): topics: topics }} as result ORDER BY result.score DESC LIMIT $limit''' - statements_results = self.graph_store.execute_query(statements_cypher, statements_params) + statements_results = self.graph_store.execute_query(statements_cypher, statements_params, operation=GraphOperation.GET_STATEMENTS) statement_facts_cypher = f'''// get facts for statements MATCH (f)-[:`__SUPPORTS__`]->(l:`__Statement__`) @@ -199,7 +200,7 @@ def get_statements_by_topic_and_source(self, statement_ids): 'statementIds': statement_ids } - statement_facts_results = self.graph_store.execute_query(statement_facts_cypher, statement_facts_params) + statement_facts_results = self.graph_store.execute_query(statement_facts_cypher, statement_facts_params, operation=GraphOperation.GET_FACTS) statement_facts = { r['statementId']:r['facts'] for r in statement_facts_results diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/__init__.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/__init__.py index efd1711c..649cd606 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/__init__.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/__init__.py @@ -5,4 +5,5 @@ from .graph_store_factory_method import GraphStoreFactoryMethod from .multi_tenant_graph_store import MultiTenantGraphStore from .dummy_graph_store import DummyGraphStore +from .graph_operation import GraphOperation from .query_tree import Query, QueryTree \ No newline at end of file diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_operation.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_operation.py new file mode 100644 index 00000000..d5ad50b0 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_operation.py @@ -0,0 +1,45 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from enum import Enum + + +class GraphOperation(str, Enum): + """Backend-neutral lexical-graph operations. + + Query-language-specific stores use these identifiers to select their native + implementation. + """ + + UPSERT_SOURCE = 'upsert_source' + UPSERT_CHUNK = 'upsert_chunk' + LINK_CHUNK_SOURCE = 'link_chunk_source' + LINK_CHUNKS = 'link_chunks' + UPSERT_TOPIC = 'upsert_topic' + UPSERT_STATEMENT = 'upsert_statement' + LINK_STATEMENT_CHUNK = 'link_statement_chunk' + LINK_STATEMENT_TOPIC = 'link_statement_topic' + LINK_STATEMENTS = 'link_statements' + UPSERT_FACT = 'upsert_fact' + LINK_FACT_ENTITY = 'link_fact_entity' + UPSERT_ENTITY = 'upsert_entity' + LINK_ENTITIES = 'link_entities' + ADD_ENTITY_TYPE = 'add_entity_type' + UPDATE_GRAPH_SUMMARY = 'update_graph_summary' + FIND_COMPLEMENTS = 'find_complements' + FIND_SUBJECTS = 'find_subjects' + COPY_COMPLEMENT_RELATIONSHIPS = 'copy_complement_relationships' + DELETE_COMPLEMENT = 'delete_complement' + GET_STATEMENTS = 'get_statements' + GET_FACTS = 'get_facts' + GET_CHUNKS = 'get_chunks' + GET_TOPIC = 'get_topic' + SEARCH_BY_CHUNK = 'search_by_chunk' + SEARCH_BY_TOPIC = 'search_by_topic' + FIND_ENTITIES_BY_KEYWORD = 'find_entities_by_keyword' + FIND_ENTITIES_BY_CHUNKS = 'find_entities_by_chunks' + FIND_ENTITIES_BY_TOPICS = 'find_entities_by_topics' + FIND_ENTITY_NEIGHBORS = 'find_entity_neighbors' + SCORE_ENTITIES = 'score_entities' + SEARCH_BY_ENTITY = 'search_by_entity' + SEARCH_BY_ENTITIES = 'search_by_entities' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_store.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_store.py index 61bd85ef..8a92e13e 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_store.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/graph_store.py @@ -11,6 +11,7 @@ from graphrag_toolkit.lexical_graph import TenantId, GraphQueryError from graphrag_toolkit.lexical_graph.storage.graph.query_tree import QueryTree +from graphrag_toolkit.lexical_graph.storage.graph.graph_operation import GraphOperation from llama_index.core.bridge.pydantic import BaseModel, Field @@ -390,7 +391,7 @@ def __exit__(self, exception_type, exception_value, traceback): def unretriable_exception_types(self) -> Tuple: return () - def execute_query_with_retry(self, query:str, parameters:Dict[str, Any], max_attempts=3, max_wait=5, **kwargs) -> Dict[str, Any]: + def execute_query_with_retry(self, query:str, parameters:Dict[str, Any], max_attempts=3, max_wait=5, operation:Optional[GraphOperation]=None, **kwargs) -> Dict[str, Any]: """ Executes a database query with a retry mechanism, allowing multiple attempts with delays between them. @@ -435,9 +436,11 @@ def execute_query_with_retry(self, query:str, parameters:Dict[str, Any], max_att attempt_number += 1 attempt.retry_state.attempt_number if isinstance(query, str): + if operation is not None: + return self._execute_operation(operation, query, parameters, **kwargs) return self._execute_query(query, parameters, **kwargs) elif isinstance(query, QueryTree): - return query.run(parameters, self.execute_query_with_retry) + return query.run(parameters, self.execute_query_with_retry, **kwargs) else: raise ValueError(f'Invalid query type. Expected string or Query Tree but received {type(query).__name__}.') @@ -496,23 +499,29 @@ def property_assigment_fn(self, key:str, value:Any) -> Callable[[str], str]: """ return lambda x: x - def execute_query(self, query:QUERY_TYPE, parameters={}, correlation_id:Optional[str]=None) -> List[Any]: + def execute_query(self, query:QUERY_TYPE, parameters={}, correlation_id:Optional[str]=None, operation:Optional[GraphOperation]=None) -> List[Any]: if correlation_id: return self.execute_query_with_retry( query=query, parameters=parameters, max_attempts=1, max_wait=0, - correlation_id=correlation_id + correlation_id=correlation_id, + operation=operation ) else: return self.execute_query_with_retry( query=query, parameters=parameters, max_attempts=1, - max_wait=0 + max_wait=0, + operation=operation ) + def _execute_operation(self, operation:GraphOperation, query:str, parameters:Dict[str, Any], correlation_id=None, **kwargs) -> List[Any]: + """Execute the existing native query for stores without an override.""" + return self._execute_query(query, parameters, correlation_id=correlation_id) + @abc.abstractmethod def _execute_query(self, cypher, parameters={}, correlation_id=None) -> List[Any]: diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/multi_tenant_graph_store.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/multi_tenant_graph_store.py index b2ef156a..61932e13 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/multi_tenant_graph_store.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/multi_tenant_graph_store.py @@ -6,6 +6,7 @@ from graphrag_toolkit.lexical_graph import TenantId from graphrag_toolkit.lexical_graph.storage.constants import LEXICAL_GRAPH_LABELS from graphrag_toolkit.lexical_graph.storage.graph import GraphStore, NodeId +from graphrag_toolkit.lexical_graph.storage.graph.query_tree import QueryTree class MultiTenantGraphStore(GraphStore): """ @@ -64,7 +65,12 @@ def execute_query_with_retry(self, query:str, parameters:Dict[str, Any], max_att **kwargs: Additional optional keyword arguments to be passed to the `execute_query_with_retry` method of the `inner` object. """ - return self.inner.execute_query_with_retry(query=self._rewrite_query(query), parameters=parameters, max_attempts=max_attempts, max_wait=max_wait) + if isinstance(query, QueryTree): + kwargs.setdefault('tenant_id', self.tenant_id.value) + return super().execute_query_with_retry(query=query, parameters=parameters, max_attempts=max_attempts, max_wait=max_wait, **kwargs) + if kwargs.get('operation') is not None: + kwargs.setdefault('tenant_id', self.tenant_id.value) + return self.inner.execute_query_with_retry(query=self._rewrite_query(query), parameters=parameters, max_attempts=max_attempts, max_wait=max_wait, **kwargs) def _logging_prefix(self, query_id:str, correlation_id:Optional[str]=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..2b3b5419 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 @@ -255,7 +255,7 @@ def try_create(self, graph_info:str, **kwargs) -> GraphStore: graph_endpoint = graph_info[len(NEPTUNE_DATABASE):] elif graph_info.endswith(NEPTUNE_DB_DNS): graph_endpoint = graph_info - elif NEPTUNE_DB_DNS in graph_info: + elif NEPTUNE_DB_DNS in graph_info and not graph_info.startswith('sparql+'): graph_endpoint = graph_info.replace('https://', '') if graph_endpoint: diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/query_tree.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/query_tree.py index 9f36f2db..b969d80f 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/query_tree.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/storage/graph/query_tree.py @@ -3,6 +3,8 @@ from typing import Any, List, Dict, Optional, Callable, Union, Iterable, Generator +from .graph_operation import GraphOperation + def _default_params_adapter(v): def _dedup(parameters:List): @@ -26,19 +28,25 @@ class Query(): def __init__(self, query:str, params_adapter:Optional[Callable[[Any], Dict]]=None, - child_queries:Optional[List]=None): + child_queries:Optional[List]=None, + operation:Optional[GraphOperation]=None): self.query = query self.params_adapter = params_adapter or DEFAULT_PARAMS_ADAPTER self.child_queries = child_queries or [] + self.operation = operation class Job(): def __init__(self, query:Query, params:Any): self.query = query self.params = params - def run(self, graph_store_fn:Callable[[str, Dict], List[Any]]): + def run(self, graph_store_fn:Callable[[str, Dict], List[Any]], **kwargs): parameters = self.query.params_adapter(self.params) - return graph_store_fn(self.query.query, parameters) + if self.query.operation is not None: + kwargs['operation'] = self.query.operation + else: + kwargs.pop('tenant_id', None) + return graph_store_fn(self.query.query, parameters, **kwargs) class QueryTree(): @@ -46,7 +54,7 @@ def __init__(self, name:str, root_query:Query): self.id = f'query-tree-{name}' self.root_query = root_query - def run(self, params, graph_store_fn:Callable[[str, Dict], List[Any]]) -> Iterable[Any]: + def run(self, params, graph_store_fn:Callable[[str, Dict], List[Any]], **kwargs) -> Iterable[Any]: job_queue = [] @@ -54,7 +62,7 @@ def run(self, params, graph_store_fn:Callable[[str, Dict], List[Any]]) -> Iterab while job: - results = job.run(graph_store_fn) + results = job.run(graph_store_fn, **kwargs) if job.query.child_queries: for q in job.query.child_queries: diff --git a/lexical-graph/tests/unit/indexing/build/test_graph_batch_client.py b/lexical-graph/tests/unit/indexing/build/test_graph_batch_client.py index 30a46798..8e8ff33c 100644 --- a/lexical-graph/tests/unit/indexing/build/test_graph_batch_client.py +++ b/lexical-graph/tests/unit/indexing/build/test_graph_batch_client.py @@ -4,6 +4,7 @@ import pytest from unittest.mock import Mock from graphrag_toolkit.lexical_graph.indexing.build.graph_batch_client import GraphBatchClient +from graphrag_toolkit.lexical_graph.storage.graph import GraphOperation, Query, QueryTree class TestGraphBatchClientInitialization: @@ -80,3 +81,43 @@ def test_node_id_delegates_to_graph_client(self, mock_neptune_store): result = client.node_id('entityId') assert result == 'params.entityId' mock_neptune_store.node_id.assert_called_once_with('entityId') + + +class TestGraphBatchClientOperations: + def test_batch_forwards_operation(self, mock_neptune_store): + client = GraphBatchClient(mock_neptune_store, True, 10) + + client.execute_query_with_retry( + 'UPSERT', + {'params': [{'chunk_id': 'c1'}]}, + operation=GraphOperation.UPSERT_CHUNK, + ) + client.apply_batch_operations() + + kwargs = mock_neptune_store.execute_query_with_retry.call_args.kwargs + assert kwargs['operation'] is GraphOperation.UPSERT_CHUNK + + def test_query_tree_batch_forwards_operation(self, mock_neptune_store): + client = GraphBatchClient(mock_neptune_store, True, 10) + tree = QueryTree( + 'lookup', + Query('SELECT', operation=GraphOperation.FIND_COMPLEMENTS), + ) + + client.execute_query_with_retry(tree, {'params': [{'nId': 'e1'}]}) + client.apply_batch_operations() + + kwargs = mock_neptune_store.execute_query_with_retry.call_args.kwargs + assert kwargs['operation'] is GraphOperation.FIND_COMPLEMENTS + + def test_empty_operation_is_a_noop(self, mock_neptune_store): + client = GraphBatchClient(mock_neptune_store, True, 10) + + client.execute_query_with_retry( + 'UNWIND $params AS params', + {'params': []}, + operation=GraphOperation.LINK_FACT_ENTITY, + ) + client.apply_batch_operations() + + mock_neptune_store.execute_query_with_retry.assert_not_called() diff --git a/lexical-graph/tests/unit/storage/graph/test_multi_tenant_graph_store.py b/lexical-graph/tests/unit/storage/graph/test_multi_tenant_graph_store.py index f74443b1..e0fad9f2 100644 --- a/lexical-graph/tests/unit/storage/graph/test_multi_tenant_graph_store.py +++ b/lexical-graph/tests/unit/storage/graph/test_multi_tenant_graph_store.py @@ -10,6 +10,8 @@ from graphrag_toolkit.lexical_graph.storage.graph.multi_tenant_graph_store import ( MultiTenantGraphStore, ) +from graphrag_toolkit.lexical_graph.storage.graph.graph_operation import GraphOperation +from graphrag_toolkit.lexical_graph.storage.graph.query_tree import Query, QueryTree def _wrap(tenant_value=None, labels=None): @@ -65,6 +67,34 @@ def test_execute_query_with_retry_rewrites_and_delegates(self): called_query = inner.execute_query_with_retry.call_args.kwargs['query'] assert '`Sourceacme__`' in called_query assert inner.execute_query_with_retry.call_args.kwargs['parameters'] == {'k': 1} + assert 'tenant_id' not in inner.execute_query_with_retry.call_args.kwargs + + def test_operation_receives_tenant_id(self): + store, inner = _wrap(tenant_value='acme', labels=['Source']) + + store.execute_query_with_retry( + 'MATCH (n:`Source`)', + {'k': 1}, + operation=GraphOperation.GET_FACTS, + ) + + kwargs = inner.execute_query_with_retry.call_args.kwargs + assert kwargs['operation'] is GraphOperation.GET_FACTS + assert kwargs['tenant_id'] == 'acme' + + def test_query_tree_operations_receive_tenant_id(self): + store, inner = _wrap(tenant_value='acme', labels=['Source']) + inner.execute_query_with_retry.return_value = [] + tree = QueryTree( + 'lookup', + Query('MATCH (n:`Source`)', operation=GraphOperation.GET_FACTS), + ) + + list(store.execute_query_with_retry(tree, {'statementIds': ['s1']})) + + kwargs = inner.execute_query_with_retry.call_args.kwargs + assert kwargs['operation'] is GraphOperation.GET_FACTS + assert kwargs['tenant_id'] == 'acme' def test_execute_query_rewrites_and_delegates(self): store, inner = _wrap(tenant_value='acme', labels=['Source']) diff --git a/lexical-graph/tests/unit/storage/graph/test_neptune_graph_stores.py b/lexical-graph/tests/unit/storage/graph/test_neptune_graph_stores.py index cd396c27..d5e65ba5 100644 --- a/lexical-graph/tests/unit/storage/graph/test_neptune_graph_stores.py +++ b/lexical-graph/tests/unit/storage/graph/test_neptune_graph_stores.py @@ -67,6 +67,12 @@ def test_try_create_extracts_graph_id(self): class TestNeptuneDatabaseGraphStoreFactory: """Tests for NeptuneDatabaseGraphStoreFactory.""" + def test_factory_does_not_claim_sparql_neptune_urls(self): + factory = NeptuneDatabaseGraphStoreFactory() + assert factory.try_create( + 'sparql+neptune://cluster.us-east-1.neptune.amazonaws.com:8182' + ) is None + def test_try_create_with_neptune_db_prefix(self): """Verify factory creates store for neptune-db:// prefix.""" factory = NeptuneDatabaseGraphStoreFactory() diff --git a/lexical-graph/tests/unit/storage/graph/test_query_tree.py b/lexical-graph/tests/unit/storage/graph/test_query_tree.py index 0d378039..812e020b 100644 --- a/lexical-graph/tests/unit/storage/graph/test_query_tree.py +++ b/lexical-graph/tests/unit/storage/graph/test_query_tree.py @@ -14,6 +14,7 @@ QueryTree, _default_params_adapter, ) +from graphrag_toolkit.lexical_graph.storage.graph.graph_operation import GraphOperation class TestDefaultParamsAdapter: @@ -88,6 +89,32 @@ def test_run_uses_custom_adapter(self): store.assert_called_once_with('MATCH (n)', {'wrapped': 'raw'}) + def test_run_drops_operation_context_for_native_query(self): + job = Job(Query('MATCH (n)'), params={}) + store = MagicMock(return_value=[]) + + job.run(store, tenant_id='acme', correlation_id='request-1') + + store.assert_called_once_with( + 'MATCH (n)', {}, correlation_id='request-1', + ) + + def test_run_forwards_semantic_operation_and_context(self): + job = Job( + Query('MATCH (n)', operation=GraphOperation.GET_FACTS), + params={'statementIds': ['s1']}, + ) + store = MagicMock(return_value=[]) + + job.run(store, tenant_id='acme') + + store.assert_called_once_with( + 'MATCH (n)', + {'statementIds': ['s1']}, + operation=GraphOperation.GET_FACTS, + tenant_id='acme', + ) + class TestQueryTree: def test_id_is_prefixed(self):