diff --git a/.github/workflows/codeproperty-graph-tests.yml b/.github/workflows/codeproperty-graph-tests.yml new file mode 100644 index 00000000..4a0ecb97 --- /dev/null +++ b/.github/workflows/codeproperty-graph-tests.yml @@ -0,0 +1,82 @@ +# Runs codeproperty-graph unit tests on push/PR to main. +name: Code Property Graph Unit Tests + +permissions: + contents: read + pull-requests: write + +on: + push: + branches: [main] + paths: + - "codeproperty-graph/**" + - ".github/workflows/codeproperty-graph-tests.yml" + pull_request: + branches: [main] + paths: + - "codeproperty-graph/**" + - ".github/workflows/codeproperty-graph-tests.yml" + +jobs: + license-header-check: + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v7 + + - name: Check license headers + run: | + MISSING=0 + HEADER_LINE1="# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved." + HEADER_LINE2="# SPDX-License-Identifier: Apache-2.0" + while IFS= read -r -d '' file; do + HEAD=$(head -5 "$file") + if ! echo "$HEAD" | grep -qF "$HEADER_LINE1" || \ + ! echo "$HEAD" | grep -qF "$HEADER_LINE2"; then + echo "MISSING: $file" + MISSING=$((MISSING + 1)) + fi + done < <(find codeproperty-graph/src -name '*.py' -print0) + if [ "$MISSING" -gt 0 ]; then + echo "" + echo "ERROR: $MISSING file(s) missing the license header." + exit 1 + fi + echo "All .py files have the required license header." + + test: + needs: license-header-check + runs-on: ubuntu-latest + timeout-minutes: 10 + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12"] + + defaults: + run: + working-directory: codeproperty-graph + + steps: + - uses: actions/checkout@v7 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Create virtual environment + run: uv venv .venv + + - name: Install dependencies + run: | + uv pip install --python .venv/bin/python \ + -e ".[dev]" + + - name: Run tests + run: | + PYTHONPATH=src .venv/bin/python -m pytest \ + -v --tb=short --maxfail=3 \ + tests/ diff --git a/.github/workflows/document-graph-tests.yml b/.github/workflows/document-graph-tests.yml new file mode 100644 index 00000000..e6c318a4 --- /dev/null +++ b/.github/workflows/document-graph-tests.yml @@ -0,0 +1,82 @@ +# Runs document-graph unit tests on push/PR to main. +name: Document Graph Unit Tests + +permissions: + contents: read + pull-requests: write + +on: + push: + branches: [main] + paths: + - "document-graph/**" + - ".github/workflows/document-graph-tests.yml" + pull_request: + branches: [main] + paths: + - "document-graph/**" + - ".github/workflows/document-graph-tests.yml" + +jobs: + license-header-check: + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v7 + + - name: Check license headers + run: | + MISSING=0 + HEADER_LINE1="# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved." + HEADER_LINE2="# SPDX-License-Identifier: Apache-2.0" + while IFS= read -r -d '' file; do + HEAD=$(head -5 "$file") + if ! echo "$HEAD" | grep -qF "$HEADER_LINE1" || \ + ! echo "$HEAD" | grep -qF "$HEADER_LINE2"; then + echo "MISSING: $file" + MISSING=$((MISSING + 1)) + fi + done < <(find document-graph/src -name '*.py' -print0) + if [ "$MISSING" -gt 0 ]; then + echo "" + echo "ERROR: $MISSING file(s) missing the license header." + exit 1 + fi + echo "All .py files have the required license header." + + test: + needs: license-header-check + runs-on: ubuntu-latest + timeout-minutes: 15 + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12"] + + defaults: + run: + working-directory: document-graph + + steps: + - uses: actions/checkout@v7 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Create virtual environment + run: uv venv .venv + + - name: Install dependencies + run: | + uv pip install --python .venv/bin/python \ + -e ".[dev,graphrag]" + + - name: Run tests + run: | + PYTHONPATH=src .venv/bin/python -m pytest \ + -v --tb=short --maxfail=3 \ + tests/ diff --git a/README.md b/README.md index 82dd3b44..3895572f 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,22 @@ The [lexical-graph](https://github.com/awslabs/graphrag-toolkit/tree/main/lexica [BYOKG-RAG](https://github.com/awslabs/graphrag-toolkit/tree/main/byokg-rag) is a novel approach to Knowledge Graph Question Answering (KGQA) that combines the power of Large Language Models (LLMs) with structured knowledge graphs. The system allows users to bring their own knowledge graph and perform complex question answering over it. +### Document Graph + +[Document Graph](https://github.com/awslabs/graphrag-toolkit/tree/main/document-graph) provides structured data ingestion into Neptune — a second ingestion route that complements lexical-graph. It transforms CSV, JSON, Parquet, Excel, XML, and YAML into typed property graphs with schema-driven construction, multi-tenancy, and batch-optimized Neptune writes. + +- **Domain Graph** in the [GraphRAG taxonomy](https://graphrag.com/) — schema-first, typed nodes, deterministic construction +- Shares tenant format with lexical-graph — both coexist in the same Neptune cluster +- Hybrid queries span structured (document-graph) and semantic (lexical-graph) layers + +### Code Property Graph + +[Code Property Graph](https://github.com/awslabs/graphrag-toolkit/tree/main/codeproperty-graph) provides delta-aware code analysis ingestion. It takes Joern CPG exports (20 node types, 14+ edge types, 9 languages) and writes them to Neptune with manifest-based change detection — skipping Neptune writes entirely when code hasn't changed. + +- Built on document-graph (DRY — ~200 lines of domain code, zero duplicated infrastructure) +- Delta logic: compare method signatures between builds, skip if unchanged (zero-cost for no-change commits) +- Full Joern schema support: Java, JavaScript/TypeScript, Python, C/C++, Go, PHP, Ruby, Kotlin, Swift + ## Security See [CONTRIBUTING](https://github.com/awslabs/graphrag-toolkit/tree/main/CONTRIBUTING.md#security-issue-notifications) for more information. diff --git a/READMORE.md b/READMORE.md new file mode 100644 index 00000000..d464c8f7 --- /dev/null +++ b/READMORE.md @@ -0,0 +1,165 @@ +# READMORE: Domain Graphs and the DRY Principle in graphrag-toolkit + +## The GraphRAG Pattern + +[GraphRAG](https://graphrag.com/) establishes a taxonomy of graph types for retrieval-augmented generation: + +| Graph Type | Input | Structure | Query Style | +|------------|-------|-----------|-------------| +| **Entity Graph** | Unstructured text | LLM-extracted entities + relationships | Semantic search, traversal | +| **Domain Graph** | Structured/domain-specific data | Typed nodes with schema | Cypher, exact match, domain queries | +| **Lexical Graph** | Text chunks | Chunk → entity → topic hierarchy | Hybrid (vector + graph) | + +graphrag-toolkit implements all three — and proves they compose into a single Neptune store: + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ Amazon Neptune │ +├─────────────────────────────────────────────────────────────────┤ +│ lexical-graph │ document-graph │ codeproperty-graph│ +│ (Entity/Lexical) │ (Domain Graph) │ (Domain Graph) │ +│ │ │ │ +│ Text → Entities │ CSV/JSON → Nodes │ Code → CPG Nodes │ +│ Chunks → Embeddings │ Schema → Edges │ Joern → Delta │ +│ │ │ │ +│ "What is X?" │ "Find all Y" │ "What changed?" │ +│ (semantic) │ (structured) │ (code analysis) │ +└─────────────────────────────────────────────────────────────────┘ +``` + +## document-graph: A Domain Graph for Structured Data + +`document-graph` is a **Domain Graph** in the GraphRAG taxonomy. It takes structured data (CSV, JSON, Parquet, Excel, XML, YAML) and builds a typed property graph with declared schema. + +What makes it a domain graph (vs. an entity graph): +- **Schema-first**: You declare node types, edge types, and property mappings +- **Typed nodes**: Each node has a specific label (`User`, `Account`, `Event`) — not generic "Entity" +- **Deterministic construction**: Same input → same graph (no LLM extraction variance) +- **Structured queries**: Cypher with known schema, not semantic similarity + +What makes it part of the toolkit (not standalone): +- Uses lexical-graph's `GraphStoreFactory` for Neptune connectivity +- Uses lexical-graph's `TenantId.format_label()` for multi-tenancy +- Shares the same Neptune cluster — hybrid queries span both graph types +- Follows the same namespace, packaging, and distribution pattern + +## codeproperty-graph: Domain Graph + DRY in Action + +`codeproperty-graph` is both a **Domain Graph** (code analysis domain) and the clearest demonstration of the **DRY principle** in this toolkit. + +### The Domain + +Code Property Graphs (CPGs) are the standard representation for static code analysis. Joern produces them from source code in 9 languages. codeproperty-graph ingests them into Neptune for: +- Call graph analysis ("what calls what?") +- Dependency tracking ("what changed between builds?") +- Security analysis ("trace data flow from input to output") +- Architecture visualization ("show me the class hierarchy") + +### The DRY Architecture + +codeproperty-graph is **~200 lines of domain-specific code**. Everything else is reused: + +``` +codeproperty-graph (domain layer — what's NEW) +├── schema.py — Joern type enums, property specs +├── models.py — CPGNode, CPGEdge, Manifest (with from_joern()) +├── graph_diff.py — Method signature comparison +├── manifest_manager.py — S3 state tracking +├── delta_ingestor.py — Skip-or-replace orchestration +└── tenant_ops.py — Tenant lifecycle + +document-graph (infrastructure — REUSED, not duplicated) +├── Node, Edge models +├── batch_nodes_unwind(), batch_edges_unwind() +├── CypherBuilder (tenant-scoped label generation) +└── Multi-tenancy (TenantId.format_label()) + +lexical-graph (foundation — REUSED, not duplicated) +├── GraphStoreFactory (Neptune connection) +├── VectorStoreFactory (OpenSearch connection) +├── TenantId (shared tenant format) +└── Graph write infrastructure +``` + +What codeproperty-graph does NOT contain: +- ❌ No Cypher generation code (uses document-graph's `batch_nodes_unwind`) +- ❌ No Neptune connection code (uses lexical-graph's `GraphStoreFactory`) +- ❌ No tenant label formatting (uses lexical-graph's `TenantId`) +- ❌ No graph write batching logic (uses document-graph's UNWIND patterns) +- ❌ No multi-tenancy implementation (inherits from the shared format) + +**This is what DRY looks like at the package level**: a domain-specific layer that adds *only* domain semantics (CPG types, delta logic, manifests) without reimplementing any infrastructure. + +### Why This Matters + +1. **New domain graphs are cheap to build.** If you have a new domain (security findings, infrastructure state, compliance records), you write ~200 lines of domain models + delta logic. The graph infra already exists. + +2. **All domain graphs are compatible.** Because they share tenant format and graph store, you can query across domains in a single Cypher traversal. A security domain graph can link to the code property graph via shared identifiers. + +3. **Infrastructure improvements benefit all domains.** When document-graph gets batch performance improvements or lexical-graph gets a new graph store backend, codeproperty-graph inherits them automatically — zero changes required. + +4. **The pattern is provably correct.** codeproperty-graph's 17 tests pass against the shared infrastructure. If the contract (TenantId format, Node/Edge models, GraphStore interface) holds, any domain graph built on it will work. + +## The Layered Architecture + +``` +┌─────────────────────────────────────────────────────┐ +│ Domain Graphs (your code) │ +│ codeproperty-graph │ your-domain-graph │ +│ ~200 lines │ ~200 lines │ +├─────────────────────────────────────────────────────┤ +│ document-graph (typed graph infra) │ +│ Node/Edge models, CypherBuilder, batch writes, │ +│ schema providers, transformers, constructors │ +├─────────────────────────────────────────────────────┤ +│ lexical-graph (foundation) │ +│ GraphStore, VectorStore, TenantId, Neptune/AOSS │ +│ connection, entity extraction, semantic indexing │ +└─────────────────────────────────────────────────────┘ +``` + +Each layer has a clear responsibility: +- **lexical-graph**: How to connect to and write to graph/vector stores +- **document-graph**: How to transform structured data into typed graphs +- **domain graphs**: What the domain-specific types and logic are + +## Building Your Own Domain Graph + +To add a new domain (e.g., infrastructure state, security findings, compliance records): + +```python +# 1. Define your domain models (~50 lines) +@dataclass +class SecurityFinding: + id: str + severity: str + resource_arn: str + ... + +# 2. Define your delta logic (~50 lines) +class FindingsDiff: + @staticmethod + def compare(prev, curr): ... + +# 3. Define your ingestor (~100 lines) +class FindingsIngestor: + async def ingest(self, findings, graph_store, tenant_id): ... + # Uses document-graph's batch_nodes_unwind + # Uses lexical-graph's GraphStore + # Uses shared TenantId format + +# That's it. ~200 lines. Full Neptune integration. +``` + +The toolkit pattern means you focus on *domain semantics* — what makes your domain unique — and inherit all the graph infrastructure for free. + +## Summary + +| Package | GraphRAG Type | Lines of Domain Code | Infrastructure Reused | +|---------|--------------|---------------------|-----------------------| +| lexical-graph | Entity + Lexical Graph | Foundation | — | +| document-graph | Domain Graph (generic) | ~5000 | lexical-graph | +| codeproperty-graph | Domain Graph (code) | ~200 | document-graph + lexical-graph | +| *your-domain-graph* | Domain Graph (yours) | ~200 | document-graph + lexical-graph | + +The architecture proves that GraphRAG domain graphs are composable, reusable, and cheap to build when the infrastructure layer is solid. diff --git a/codeproperty-graph/LICENSE b/codeproperty-graph/LICENSE new file mode 100644 index 00000000..67db8588 --- /dev/null +++ b/codeproperty-graph/LICENSE @@ -0,0 +1,175 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. diff --git a/codeproperty-graph/README.md b/codeproperty-graph/README.md new file mode 100644 index 00000000..e77cea14 --- /dev/null +++ b/codeproperty-graph/README.md @@ -0,0 +1,80 @@ +# codeproperty-graph + +CPG domain layer for Joern-based code property graph analysis with delta ingestion. + +## Architecture + +``` +┌─────────────────────────────────────────────────────┐ +│ codeproperty-graph (domain) │ +│ CPGNode, CPGEdge, GraphDiff, DeltaIngestor │ +│ ManifestManager, tenant lifecycle │ +├─────────────────────────────────────────────────────┤ +│ document-graph (infra) │ +│ Node, Edge, batch_nodes_unwind, batch_edges_unwind │ +│ CypherBuilder, PipelineExecutor, multi-tenancy │ +├─────────────────────────────────────────────────────┤ +│ graphrag-toolkit / lexical-graph (foundation)│ +│ GraphStore, Neptune writer, AOSS writer │ +│ Lexical indexing, entity resolution, retrieval │ +└─────────────────────────────────────────────────────┘ +``` + +This library sits on top of **document-graph** (typed property graph primitives) and +**lexical-graph** (graph storage foundation from AWS graphrag-toolkit). It adds +code-analysis-specific semantics: + +- **CPG Models** — Joern node/edge types with `full_name`, `hash`, `signature` +- **GraphDiff** — compare two CPG states by method signature, return adds/removes/modified +- **ManifestManager** — S3-backed state tracking per repository +- **DeltaIngestor** — skip-or-replace orchestration with old tenant purge +- **Tenant Ops** — clean lifecycle management (`delete_tenant`) + +## Usage + +```python +from codeproperty_graph import DeltaIngestor + +ingestor = DeltaIngestor(bucket="") + +result = await ingestor.ingest( + repo="amigo-core", + job_id="abc-123", + tenant_id="abc12345678901234567890", + nodes_data=nodes, # from Joern export + edges_data=edges, # from Joern export + nodes_path="s3://bucket/cpg-exports/amigo-core/abc-123/nodes.json", + edges_path="s3://bucket/cpg-exports/amigo-core/abc-123/edges.json", + graph_store=neptune, + write_fn=my_write_function, +) + +# result: {"status": "SKIPPED", ...} or {"status": "INGESTED", "delta": "+2 -1 ~3 =150", ...} +``` + +## Delta Logic + +1. Joern exports CPG → `nodes.json` + `edges.json` +2. Extract METHOD node signatures: `{full_name: hash}` +3. Compare against previous manifest in S3 +4. If identical → **SKIP** (no Neptune writes) +5. If changed → **INGEST** full graph under new tenant, purge old tenant, update manifest + +## Relationship to graphrag-toolkit + +| Layer | Purpose | Example | +|-------|---------|---------| +| graphrag-toolkit | Graph storage, Neptune/AOSS writers | `GraphStore.execute_query()` | +| lexical-graph | Document chunking, entity extraction | `LexicalGraphIndex` | +| document-graph | Typed nodes/edges, Cypher generation | `batch_nodes_unwind()` | +| **codeproperty-graph** | CPG delta, manifests, Joern semantics | `DeltaIngestor.ingest()` | + +## Install + +```bash +pip install codeproperty-graph +``` + +## License + +Apache-2.0 diff --git a/codeproperty-graph/docs/VERDICT-GRAPH-SPEC.md b/codeproperty-graph/docs/VERDICT-GRAPH-SPEC.md new file mode 100644 index 00000000..202a1e0d --- /dev/null +++ b/codeproperty-graph/docs/VERDICT-GRAPH-SPEC.md @@ -0,0 +1,243 @@ +# Verdict Graph: Specification + +> PRIVATE AND PROPRIETARY. Owned by Kanjani AI Research. See NOTICE.md. + +**Status:** Specification only — not implemented +**Layer:** Domain Graph (governance evidence domain) +**Depends on:** codeproperty-graph, document-graph, lexical-graph +**Related:** AIGP RFC-032 (Post-Hoc Evaluation Loop), RFC-035 (Mediation Vector Profile), D-DNA + +--- + +## 1. What Is the Verdict Graph? + +The Verdict Graph is a Domain Graph whose domain is **software governance evidence** — the accumulation of observations, verdicts, quality signals, and provenance about software artifacts over their lifecycle. + +``` +Code Property Graph = "What IS the code?" (structure, flow, calls) +Verdict Graph = "What do we KNOW about (verdicts, evidence, posture, + the code?" provenance, compliance) +``` + +It is the **Master Evidence Graph** described in the research papers — made executable as a queryable property graph in Neptune. + +--- + +## 2. Why It Exists + +Software governance today is stateless: +- A scan runs, produces a report, the report is filed +- Next scan runs, produces another report — no connection to the previous one +- Nobody can ask: "Show me the *trajectory* of this component's quality over time" + +The Verdict Graph makes governance **stateful**: +- Every verdict is a node with timestamp, evidence, observer, and result +- Verdicts connect to the artifact they evaluate (CPG nodes) +- History accumulates — you can query trajectories, not just snapshots +- Provenance is first-class: who observed, what evidence, under what authority + +--- + +## 3. Node Types + +| Node Type | Description | Key Properties | +|-----------|-------------|----------------| +| `Artifact` | The thing being evaluated (code unit, deployment, config) | artifact_id, type, version, source_repo | +| `Verdict` | A judgment rendered on an artifact | verdict_id, result (MATCH/MISMATCH/VIOLATION), timestamp, confidence | +| `Evidence` | A signed, provenance-bound piece of evidence | evidence_id, type, source, hash, admissibility_level | +| `Observer` | Who/what rendered the verdict | observer_id, type (human/machine/hybrid), calibration_status | +| `Concern` | The evaluative question being asked | concern_id, domain_of_concern, concern_class | +| `Policy` | The governance rule that triggered evaluation | policy_id, version, authority, scope | +| `MediationVector` | A point-in-time measurement snapshot (RFC-035) | vector_id, timestamp, stage, variable_values | +| `CouplingState` | Trajectory state for process-level concerns (ADD-035-001) | coupling_id, session_count, current_stage, progression_rate | + +--- + +## 4. Edge Types + +| Edge Type | From → To | Meaning | +|-----------|-----------|---------| +| `EVALUATES` | Verdict → Artifact | This verdict evaluates this artifact | +| `EVIDENCED_BY` | Verdict → Evidence | This verdict is supported by this evidence | +| `RENDERED_BY` | Verdict → Observer | This observer rendered this verdict | +| `CONCERNS` | Verdict → Concern | This verdict addresses this concern | +| `AUTHORIZED_BY` | Verdict → Policy | This policy authorized the evaluation | +| `SUPERSEDES` | Verdict → Verdict | This verdict replaces the previous one | +| `DERIVED_FROM` | Evidence → Evidence | Provenance chain | +| `OBSERVED_IN` | MediationVector → Artifact | This measurement observed this artifact | +| `PRODUCED_BY` | MediationVector → Observer | Who made this measurement | +| `PART_OF` | MediationVector → CouplingState | This measurement belongs to this trajectory | +| `REFERENCES` | Artifact → CPGNode | Links verdict subject to code structure | + +--- + +## 5. Relationship to Code Property Graph + +``` +Verdict Graph Code Property Graph +───────────── ──────────────────── +Artifact(id="parser-v2.1") ──REFERENCES──→ METHOD(full_name="com.example.Parser.parse") + │ + ├── EVALUATES ←── Verdict(result=MATCH, t=2026-06-30) + │ │ + │ ├── EVIDENCED_BY → Evidence(type=D-DNA, hash=...) + │ ├── RENDERED_BY → Observer(type=machine, id=sast-scanner) + │ └── CONCERNS → Concern(class=code_quality) + │ + └── EVALUATES ←── Verdict(result=MISMATCH, t=2026-06-15) + │ + └── SUPERSEDES → Verdict(result=MATCH, t=2026-06-01) +``` + +The REFERENCES edge is the **bridge** between the Verdict Graph and the Code Property Graph. It links governance observations to the code structure they observe. + +--- + +## 6. Relationship to AIGP Protocol + +| AIGP Concept | Verdict Graph Node/Edge | +|--------------|------------------------| +| RFC-032 VERIFY verdict | `Verdict` node with result, timestamp, evidence | +| RFC-035 Mediation Vector | `MediationVector` node with variable_values | +| D-DNA signed evidence | `Evidence` node with hash, provenance, admissibility_level | +| RFC-034 Domain of Concern | `Concern` node with domain, class | +| RFC-037 Observer | `Observer` node with type, calibration, accreditation | +| ADD-035-001 Coupling | `CouplingState` node with stage, progression_rate | +| ADD-035-004 Progressive Degradation | Trajectory of `MediationVector` nodes over time | + +The Verdict Graph is **AIGP made queryable** — the protocol's evidence and verdicts stored as graph state that can be traversed, aggregated, and analyzed. + +--- + +## 7. DRY Architecture (Projected) + +``` +verdict-graph (domain layer — ~200 lines) +├── schema.py — Verdict/Evidence/Observer/Concern types +├── models.py — VerdictNode, EvidenceNode, etc. with from_aigp() +├── verdict_tracker.py — Accumulate verdicts over time, link to artifacts +├── posture_query.py — Query concern posture trajectories +└── provenance.py — Evidence chain traversal + +document-graph (infrastructure — REUSED) +├── Node, Edge, batch_nodes_unwind +├── CypherBuilder, tenant-scoped labels +└── Multi-tenancy + +lexical-graph (foundation — REUSED) +├── GraphStoreFactory, Neptune connection +└── TenantId +``` + +--- + +## 8. Key Queries the Verdict Graph Enables + +```cypher +-- "What is the current quality posture of this artifact?" +MATCH (a:Artifact {artifact_id: $id})<-[:EVALUATES]-(v:Verdict) +WHERE NOT (v)<-[:SUPERSEDES]-() +RETURN v.result, v.timestamp, v.confidence + +-- "Show me the verdict trajectory for this component over 6 months" +MATCH (a:Artifact {artifact_id: $id})<-[:EVALUATES]-(v:Verdict) +WHERE v.timestamp > datetime() - duration('P6M') +RETURN v.result, v.timestamp ORDER BY v.timestamp + +-- "Which artifacts have unresolved MISMATCH verdicts?" +MATCH (v:Verdict {result: 'MISMATCH'})-[:EVALUATES]->(a:Artifact) +WHERE NOT (v)<-[:SUPERSEDES]-() +RETURN a.artifact_id, v.timestamp, v.concern_class + +-- "What evidence supports this verdict?" +MATCH (v:Verdict {verdict_id: $id})-[:EVIDENCED_BY]->(e:Evidence) +RETURN e.type, e.source, e.admissibility_level, e.hash + +-- "Show the progressive degradation trajectory for this coupling" +MATCH (cs:CouplingState {coupling_id: $id})<-[:PART_OF]-(mv:MediationVector) +RETURN mv.timestamp, mv.stage, mv.variable_values +ORDER BY mv.timestamp + +-- "Which code methods have never been evaluated?" +MATCH (m:METHOD) +WHERE NOT (:Artifact)-[:REFERENCES]->(m) +RETURN m.full_name, m.filename +``` + +--- + +## 9. Delta Logic + +Like codeproperty-graph, the Verdict Graph has delta semantics — but **additive**, not replacement: + +| codeproperty-graph | verdict-graph | +|--------------------|---------------| +| Full replacement on change | Append-only (verdicts accumulate) | +| Old tenant purged | Old verdicts stay (linked via SUPERSEDES) | +| Manifest = method signatures | No manifest needed (timestamps are natural ordering) | +| Skip if no code change | Always write new verdicts (observation is always new) | + +The Verdict Graph never deletes — it only supersedes. History is preserved for audit trail and trajectory analysis. + +--- + +## 10. Temporal Model + +``` +t=0: First scan + Verdict(MATCH) → Artifact(parser-v2.0) + +t=1: Code changes, new scan + Verdict(MISMATCH) → Artifact(parser-v2.1) + └── SUPERSEDES → Verdict(MATCH) at t=0 + +t=2: Fix applied, re-scan + Verdict(MATCH) → Artifact(parser-v2.2) + └── SUPERSEDES → Verdict(MISMATCH) at t=1 + +Query at any time: "What's the CURRENT posture?" +→ Find verdicts with no incoming SUPERSEDES edge = latest state +``` + +--- + +## 11. Integration Points + +| System | Integration | +|--------|-------------| +| CI/CD pipeline | Writes Verdict nodes after each scan/test | +| SAST scanner (Joern, Semgrep) | Observer; Evidence is the scan output | +| AIGP Governance Server | Writes CHECK/VERIFY results as Verdicts | +| SBOM (CycloneDX) | Evidence node linked to dependency verdicts | +| D-DNA signer | Signs Evidence nodes; hash + signature stored | +| Human reviewer | Observer(type=human); manual Verdict | +| Quality Moderator (RFC-032) | Reads verdict trajectory; adapts posture | + +--- + +## 12. When to Build + +The Verdict Graph should be built when: +1. codeproperty-graph is stable in production (code structure is queryable) +2. AIGP RFC-036 (calculation semantics) is drafted (defines how vectors combine) +3. A governance pipeline exists that produces verdicts to store + +It should NOT be built prematurely — without real verdicts flowing, the graph is empty and adds no value. The spec exists so the architecture is ready when the pipeline is. + +--- + +## 13. Summary + +| Aspect | Value | +|--------|-------| +| GraphRAG type | Domain Graph (governance evidence domain) | +| Domain | Software verdicts, evidence, provenance, quality posture | +| Enables | Stateful governance — trajectories, not just snapshots | +| Architecture | ~200 lines domain code + document-graph + lexical-graph (DRY) | +| Key innovation | SUPERSEDES edge makes history queryable without deletion | +| AIGP alignment | Makes RFC-032/035/037 concepts queryable in Neptune | +| Prerequisite | Real governance pipeline producing verdicts | + +--- + +© 2024-2026 Kanjani AI Research. All rights reserved. diff --git a/codeproperty-graph/pyproject.toml b/codeproperty-graph/pyproject.toml new file mode 100644 index 00000000..39f2e18a --- /dev/null +++ b/codeproperty-graph/pyproject.toml @@ -0,0 +1,34 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/graphrag_toolkit"] + +[project] +name = "graphrag-toolkit-codeproperty-graph" +version = "0.4.2" +description = "AWS GraphRAG Toolkit, Code Property Graph — delta-aware CPG ingestion" +readme = "README.md" +requires-python = ">=3.10" +license = "Apache-2.0" +dependencies = [ + "graphrag-lexical-graph>=3.18.0", + "boto3>=1.26.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0", + "pytest-cov>=4.0", + "pytest-mock>=3.0", + "pytest-asyncio>=0.21.0", +] + +[project.urls] +Homepage = "https://github.com/awslabs/graphrag-toolkit" +Repository = "https://github.com/awslabs/graphrag-toolkit" + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src"] diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/__init__.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/__init__.py new file mode 100644 index 00000000..d5ab3fe3 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/__init__.py @@ -0,0 +1,59 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Code Property Graph — domain layer for Joern CPG delta ingestion. + +Built on document-graph for typed property graph primitives. +Adds CPG-specific models, delta comparison, manifest management, +and tenant lifecycle for incremental code analysis. + +Extended with: +- Vector loading (precomputed embeddings → OpenSearch) +- Summary overlay (LLM summaries → Neptune node properties) +- Code slice storage (source evidence → Neptune node properties) +- Artifact validation (schema compliance before Build) +- Full pipeline orchestration (ingest_artifact) + +v0.4.2: GraphLoader now uses toolkit's GraphBatchClient + NeptuneDatabaseClient +(boto3, 32-pool, 600s timeout, UNWIND dedup, retry cascade). + +Supported Joern frontends: Java, JavaScript/TypeScript, Python, C/C++, +Go, PHP, Ruby, Kotlin, Swift. + +Full CPG schema: 20 node types, 14+ edge types. +See schema.py for complete enumeration. +""" + +__version__ = "0.4.2" + +from .models import CPGNode, CPGEdge, Manifest, VectorRecord, SummaryRecord, CodeSliceRecord +from .schema import NodeType, EdgeType, DELTA_RELEVANT_TYPES, SUPPORTED_LANGUAGES, joern_export_command +from .graph_diff import GraphDiff +from .manifest_manager import ManifestManager, ManifestConflictError +from .delta_ingestor import DeltaIngestor +from .tenant_ops import delete_tenant, delete_domain, list_domains +from .artifact_validator import ArtifactValidator, ValidationResult +from .graph_loader import GraphLoader +from .vector_loader import VectorLoader, LoadResult +from .summary_overlay import SummaryOverlay +from .code_slice_store import CodeSliceStore +from .graphson_converter import GraphSONConverter +from .artifact_reader import ArtifactReader, ArtifactContents + +__all__ = [ + # Models + "CPGNode", "CPGEdge", "Manifest", + "VectorRecord", "SummaryRecord", "CodeSliceRecord", + # Schema + "NodeType", "EdgeType", "DELTA_RELEVANT_TYPES", "SUPPORTED_LANGUAGES", + "joern_export_command", + # Pipeline + "GraphDiff", "ManifestManager", "ManifestConflictError", "DeltaIngestor", + "delete_tenant", "delete_domain", "list_domains", + # Artifact pipeline + "ArtifactValidator", "ValidationResult", + "GraphLoader", "VectorLoader", "LoadResult", + "SummaryOverlay", "CodeSliceStore", + # Conversion + I/O + "GraphSONConverter", "ArtifactReader", "ArtifactContents", +] diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/artifact_reader.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/artifact_reader.py new file mode 100644 index 00000000..7556678b --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/artifact_reader.py @@ -0,0 +1,179 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Artifact Reader — reads CPG artifact JSONL files from S3 or local filesystem. + +Reads the standard CPG artifact layout: + /// + manifest.json + nodes.jsonl + edges.jsonl + vectors.jsonl + summaries.jsonl + code_slices.jsonl + lineage.jsonl (optional) + findings.jsonl (optional) +""" + +import json +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +import boto3 + +logger = logging.getLogger(__name__) + + +@dataclass +class ArtifactContents: + """Parsed contents of a CPG artifact.""" + manifest: dict = field(default_factory=dict) + nodes: list[dict] = field(default_factory=list) + edges: list[dict] = field(default_factory=list) + vectors: list[dict] = field(default_factory=list) + summaries: list[dict] = field(default_factory=list) + code_slices: list[dict] = field(default_factory=list) + lineage: list[dict] = field(default_factory=list) + findings: list[dict] = field(default_factory=list) + + +class ArtifactReader: + """Reads a CPG artifact from S3 or local filesystem. + + Usage (S3): + reader = ArtifactReader(region="us-east-1") + artifact = reader.read_s3("s3://bucket/cpg-exports/payments-api/job-001/") + + Usage (local): + artifact = ArtifactReader.read_local("/path/to/artifact/") + """ + + def __init__(self, region: str = "us-east-1"): + self._s3 = boto3.client("s3", region_name=region) + + def read_s3(self, s3_uri: str) -> ArtifactContents: + """Read a CPG artifact from an S3 prefix. + + Args: + s3_uri: S3 URI prefix (e.g., "s3://bucket/cpg-exports/repo/job/") + + Returns: + ArtifactContents with all parsed files. + """ + bucket, prefix = self._parse_s3_uri(s3_uri) + if not prefix.endswith("/"): + prefix += "/" + + artifact = ArtifactContents() + + # Read manifest + artifact.manifest = self._read_s3_json(bucket, f"{prefix}manifest.json") + + # Read JSONL files + artifact.nodes = self._read_s3_jsonl(bucket, f"{prefix}nodes.jsonl") + artifact.edges = self._read_s3_jsonl(bucket, f"{prefix}edges.jsonl") + artifact.vectors = self._read_s3_jsonl(bucket, f"{prefix}vectors.jsonl") + artifact.summaries = self._read_s3_jsonl(bucket, f"{prefix}summaries.jsonl") + artifact.code_slices = self._read_s3_jsonl(bucket, f"{prefix}code_slices.jsonl") + + # Optional files + artifact.lineage = self._read_s3_jsonl(bucket, f"{prefix}lineage.jsonl", required=False) + artifact.findings = self._read_s3_jsonl(bucket, f"{prefix}findings.jsonl", required=False) + + logger.info( + f"Artifact read from s3://{bucket}/{prefix}: " + f"nodes={len(artifact.nodes)}, edges={len(artifact.edges)}, " + f"vectors={len(artifact.vectors)}, summaries={len(artifact.summaries)}" + ) + return artifact + + @staticmethod + def read_local(path: str) -> ArtifactContents: + """Read a CPG artifact from a local directory. + + Args: + path: Local directory path containing the artifact files. + + Returns: + ArtifactContents with all parsed files. + """ + directory = Path(path) + artifact = ArtifactContents() + + # Read manifest + manifest_path = directory / "manifest.json" + if manifest_path.exists(): + with open(manifest_path) as f: + artifact.manifest = json.load(f) + + # Read JSONL files + artifact.nodes = ArtifactReader._read_local_jsonl(directory / "nodes.jsonl") + artifact.edges = ArtifactReader._read_local_jsonl(directory / "edges.jsonl") + artifact.vectors = ArtifactReader._read_local_jsonl(directory / "vectors.jsonl") + artifact.summaries = ArtifactReader._read_local_jsonl(directory / "summaries.jsonl") + artifact.code_slices = ArtifactReader._read_local_jsonl(directory / "code_slices.jsonl") + artifact.lineage = ArtifactReader._read_local_jsonl(directory / "lineage.jsonl") + artifact.findings = ArtifactReader._read_local_jsonl(directory / "findings.jsonl") + + logger.info( + f"Artifact read from {path}: " + f"nodes={len(artifact.nodes)}, edges={len(artifact.edges)}, " + f"vectors={len(artifact.vectors)}, summaries={len(artifact.summaries)}" + ) + return artifact + + def _read_s3_json(self, bucket: str, key: str) -> dict: + """Read a single JSON file from S3.""" + try: + resp = self._s3.get_object(Bucket=bucket, Key=key) + return json.loads(resp["Body"].read()) + except self._s3.exceptions.NoSuchKey: + logger.warning(f"File not found: s3://{bucket}/{key}") + return {} + except Exception as e: + logger.error(f"Failed to read s3://{bucket}/{key}: {e}") + return {} + + def _read_s3_jsonl(self, bucket: str, key: str, required: bool = True) -> list[dict]: + """Read a JSONL file from S3 (one JSON object per line).""" + try: + resp = self._s3.get_object(Bucket=bucket, Key=key) + body = resp["Body"].read().decode("utf-8") + records = [] + for line in body.strip().split("\n"): + line = line.strip() + if line: + records.append(json.loads(line)) + return records + except self._s3.exceptions.NoSuchKey: + if required: + logger.warning(f"Required file not found: s3://{bucket}/{key}") + return [] + except Exception as e: + logger.error(f"Failed to read s3://{bucket}/{key}: {e}") + return [] + + @staticmethod + def _read_local_jsonl(path: Path) -> list[dict]: + """Read a local JSONL file.""" + if not path.exists(): + return [] + records = [] + with open(path) as f: + for line in f: + line = line.strip() + if line: + records.append(json.loads(line)) + return records + + @staticmethod + def _parse_s3_uri(uri: str) -> tuple[str, str]: + """Parse s3://bucket/prefix into (bucket, prefix).""" + if uri.startswith("s3://"): + uri = uri[5:] + parts = uri.split("/", 1) + bucket = parts[0] + prefix = parts[1] if len(parts) > 1 else "" + return bucket, prefix diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/artifact_validator.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/artifact_validator.py new file mode 100644 index 00000000..be41fe3c --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/artifact_validator.py @@ -0,0 +1,162 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validates a CPG artifact before Build processes it.""" + +import logging +from dataclasses import dataclass, field +from typing import Any + +logger = logging.getLogger(__name__) + +REQUIRED_MANIFEST_FIELDS = { + "artifact_type", + "schema_version", + "repo_id", + "commit_sha", + "embedding_model", + "embedding_dimensions", + "similarity_function", + "counts", +} + +REQUIRED_NODE_FIELDS = {"cpg_node_id", "node_type", "properties"} +REQUIRED_VECTOR_FIELDS = {"cpg_node_id", "embedding_target", "vector", "embedding_dimensions"} +REQUIRED_SUMMARY_FIELDS = {"cpg_node_id", "summary_type", "text"} + + +@dataclass +class ValidationResult: + """Result of artifact validation.""" + + valid: bool = True + errors: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + +class ArtifactValidator: + """Validates CPG artifact manifests and sample records. + + Designed to fail fast on critical errors (manifest structure) + but collect all warnings from record-level checks. + """ + + def validate( + self, + manifest: dict[str, Any], + sample_records: dict[str, list[dict[str, Any]]] | None = None, + ) -> ValidationResult: + """Validate a CPG artifact manifest and optional sample records. + + Args: + manifest: The artifact manifest dictionary. + sample_records: Optional dict with keys 'nodes', 'vectors', 'summaries', + each mapping to a list of sample JSONL records. + + Returns: + ValidationResult with valid flag, errors, and warnings. + """ + result = ValidationResult() + + # --- Manifest-level validation (fail fast on critical errors) --- + self._validate_manifest(manifest, result) + if result.errors: + result.valid = False + logger.error("Manifest validation failed: %s", result.errors) + return result + + # --- Record-level validation (collect all warnings) --- + if sample_records: + expected_dimensions = manifest.get("embedding_dimensions") + self._validate_node_records(sample_records.get("nodes", []), result) + self._validate_vector_records( + sample_records.get("vectors", []), expected_dimensions, result + ) + self._validate_summary_records(sample_records.get("summaries", []), result) + + if result.errors: + result.valid = False + + return result + + def _validate_manifest(self, manifest: dict[str, Any], result: ValidationResult) -> None: + """Validate required manifest fields and their values.""" + missing = REQUIRED_MANIFEST_FIELDS - set(manifest.keys()) + if missing: + result.errors.append(f"Missing required manifest fields: {sorted(missing)}") + return + + # artifact_type must be 'cpg' + if manifest["artifact_type"] != "cpg": + result.errors.append( + f"Invalid artifact_type: expected 'cpg', got '{manifest['artifact_type']}'" + ) + + # Validate counts + counts = manifest.get("counts") + if not isinstance(counts, dict): + result.errors.append("'counts' must be a dictionary") + return + + for count_field in ("nodes", "vectors", "summaries"): + value = counts.get(count_field) + if value is None: + result.errors.append(f"'counts.{count_field}' is missing") + elif not isinstance(value, int) or value <= 0: + result.errors.append( + f"'counts.{count_field}' must be a positive integer, got {value!r}" + ) + + def _validate_node_records( + self, records: list[dict[str, Any]], result: ValidationResult + ) -> None: + """Validate sample node records.""" + for i, record in enumerate(records): + missing = REQUIRED_NODE_FIELDS - set(record.keys()) + if missing: + result.warnings.append( + f"Node record [{i}]: missing fields {sorted(missing)}" + ) + + def _validate_vector_records( + self, + records: list[dict[str, Any]], + expected_dimensions: int | None, + result: ValidationResult, + ) -> None: + """Validate sample vector records and embedding dimensions.""" + for i, record in enumerate(records): + missing = REQUIRED_VECTOR_FIELDS - set(record.keys()) + if missing: + result.warnings.append( + f"Vector record [{i}]: missing fields {sorted(missing)}" + ) + continue + + # Check embedding_dimensions match manifest + record_dims = record.get("embedding_dimensions") + if expected_dimensions is not None and record_dims != expected_dimensions: + result.errors.append( + f"Vector record [{i}]: embedding_dimensions mismatch — " + f"record has {record_dims}, manifest specifies {expected_dimensions}" + ) + + # Check vector length matches declared dimensions + vector = record.get("vector") + if isinstance(vector, list) and expected_dimensions is not None: + if len(vector) != expected_dimensions: + result.warnings.append( + f"Vector record [{i}]: vector length {len(vector)} " + f"does not match embedding_dimensions {expected_dimensions}" + ) + + def _validate_summary_records( + self, records: list[dict[str, Any]], result: ValidationResult + ) -> None: + """Validate sample summary records.""" + for i, record in enumerate(records): + missing = REQUIRED_SUMMARY_FIELDS - set(record.keys()) + if missing: + result.warnings.append( + f"Summary record [{i}]: missing fields {sorted(missing)}" + ) diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/code_slice_store.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/code_slice_store.py new file mode 100644 index 00000000..68fc6093 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/code_slice_store.py @@ -0,0 +1,209 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Code Slice Store — attaches extracted code slices to CPG nodes in Neptune. + +Reads CodeSliceRecord dicts (from code_slices.jsonl) and stores them as +properties on their corresponding Neptune nodes. Each slice becomes a set +of sub-properties keyed by slice_type (e.g. full_method, signature_only). + +Property layout on a node: + code_slice_full_method = "" + code_slice_full_method_lang = "java" + code_slice_full_method_line_start = 42 + code_slice_full_method_line_end = 67 + code_slice_full_method_file_path = "src/Main.java" + +Multiple slice types per node are supported (full_method + signature_only). +""" + +import asyncio +import logging +from typing import Any, Dict, List, Protocol, TypedDict, runtime_checkable + +logger = logging.getLogger(__name__) + + +class CodeSliceRecord(TypedDict, total=False): + """Schema for a single code slice record from code_slices.jsonl.""" + + cpg_node_id: str + slice_type: str # e.g. "full_method", "signature_only" + code: str + language: str + line_start: int + line_end: int + file_path: str + + +@runtime_checkable +class GraphStoreProtocol(Protocol): + """Protocol for Neptune graph store — matches the interface used by delta_ingestor.""" + + def execute_query(self, query: str, parameters: Dict[str, Any]) -> Any: + ... + + +class CodeSliceStore: + """Stores code slices as properties on existing CPG nodes in Neptune. + + Usage: + store = CodeSliceStore(graph_store) + count = await store.store(slices) + """ + + def __init__(self, graph_store: GraphStoreProtocol) -> None: + """Initialize with a Neptune graph store. + + Args: + graph_store: Neptune graph store instance with execute_query method. + """ + self._graph_store = graph_store + + async def store(self, slices: List[CodeSliceRecord]) -> int: + """Attach code slices to their corresponding CPG nodes. + + For each record, sets properties on the node identified by cpg_node_id: + - code_slice_{slice_type}: the code text + - code_slice_{slice_type}_lang: language + - code_slice_{slice_type}_line_start: start line number + - code_slice_{slice_type}_line_end: end line number + - code_slice_{slice_type}_file_path: source file path + + Args: + slices: List of CodeSliceRecord dicts parsed from code_slices.jsonl. + + Returns: + Number of slices successfully stored. + """ + if not slices: + return 0 + + stored = 0 + + for record in slices: + cpg_node_id = record.get("cpg_node_id") + slice_type = record.get("slice_type") + + if not cpg_node_id or not slice_type: + logger.warning( + "Skipping slice record with missing cpg_node_id or slice_type: %s", + record, + ) + continue + + try: + await self._store_single(record, cpg_node_id, slice_type) + stored += 1 + except Exception as e: + logger.error( + "Failed to store slice (node=%s, type=%s): %s", + cpg_node_id, + slice_type, + e, + ) + + logger.info("Stored %d/%d code slices", stored, len(slices)) + return stored + + async def _store_single( + self, record: CodeSliceRecord, cpg_node_id: str, slice_type: str + ) -> None: + """Store a single code slice as properties on a Neptune node. + + Args: + record: The code slice record dict. + cpg_node_id: Neptune node identifier. + slice_type: Slice type key (e.g. "full_method"). + """ + prefix = f"code_slice_{slice_type}" + + # Build the SET clause with parameterized values + query = ( + f"MATCH (n) WHERE n.id = $node_id " + f"SET n.`{prefix}` = $code, " + f"n.`{prefix}_lang` = $language, " + f"n.`{prefix}_line_start` = $line_start, " + f"n.`{prefix}_line_end` = $line_end, " + f"n.`{prefix}_file_path` = $file_path" + ) + + parameters: Dict[str, Any] = { + "node_id": cpg_node_id, + "code": record.get("code", ""), + "language": record.get("language", ""), + "line_start": record.get("line_start", 0), + "line_end": record.get("line_end", 0), + "file_path": record.get("file_path", ""), + } + + await asyncio.to_thread(self._graph_store.execute_query, query, parameters) + + async def store_batch(self, slices: List[CodeSliceRecord], batch_size: int = 50) -> int: + """Store slices in batches for improved throughput on large datasets. + + Groups slices by cpg_node_id to minimize round-trips when a single node + has multiple slice types (e.g. full_method + signature_only). + + Args: + slices: List of CodeSliceRecord dicts. + batch_size: Number of slices per batch query. + + Returns: + Number of slices successfully stored. + """ + if not slices: + return 0 + + # Group slices by node to coalesce updates + node_slices: Dict[str, List[CodeSliceRecord]] = {} + for record in slices: + node_id = record.get("cpg_node_id", "") + if node_id: + node_slices.setdefault(node_id, []).append(record) + + stored = 0 + + for node_id, node_records in node_slices.items(): + try: + set_clauses: List[str] = [] + parameters: Dict[str, Any] = {"node_id": node_id} + + for idx, record in enumerate(node_records): + slice_type = record.get("slice_type", "") + if not slice_type: + continue + + prefix = f"code_slice_{slice_type}" + param_suffix = f"_{idx}" + + set_clauses.extend([ + f"n.`{prefix}` = $code{param_suffix}", + f"n.`{prefix}_lang` = $language{param_suffix}", + f"n.`{prefix}_line_start` = $line_start{param_suffix}", + f"n.`{prefix}_line_end` = $line_end{param_suffix}", + f"n.`{prefix}_file_path` = $file_path{param_suffix}", + ]) + + parameters[f"code{param_suffix}"] = record.get("code", "") + parameters[f"language{param_suffix}"] = record.get("language", "") + parameters[f"line_start{param_suffix}"] = record.get("line_start", 0) + parameters[f"line_end{param_suffix}"] = record.get("line_end", 0) + parameters[f"file_path{param_suffix}"] = record.get("file_path", "") + + if not set_clauses: + continue + + query = ( + f"MATCH (n) WHERE n.id = $node_id " + f"SET {', '.join(set_clauses)}" + ) + + await asyncio.to_thread(self._graph_store.execute_query, query, parameters) + stored += len(node_records) + + except Exception as e: + logger.error("Batch store failed for node %s: %s", node_id, e) + + logger.info("Batch-stored %d/%d code slices", stored, len(slices)) + return stored diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/delta_ingestor.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/delta_ingestor.py new file mode 100644 index 00000000..451f8e89 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/delta_ingestor.py @@ -0,0 +1,388 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Delta Ingestor — orchestrates skip-or-replace CPG ingestion. + +Extended to handle the full CPG artifact pipeline: + 1. Delta check (method signatures) + 2. Validate artifact (manifest + sample records) + 3. Load graph (nodes + edges → Neptune) + 4. Load vectors (vectors.jsonl → OpenSearch) + 5. Apply summaries (summaries.jsonl → Neptune node properties) + 6. Store code slices (code_slices.jsonl → Neptune node properties) + 7. Purge old tenant + 8. Update manifest +""" + +import logging +from datetime import datetime, timezone +from typing import Any, Optional + +from .graph_diff import GraphDiff +from .manifest_manager import ManifestManager +from .models import Manifest +from .tenant_ops import delete_tenant, delete_domain +from .artifact_validator import ArtifactValidator +from .graph_loader import GraphLoader +from .vector_loader import VectorLoader +from .summary_overlay import SummaryOverlay +from .code_slice_store import CodeSliceStore +from .artifact_reader import ArtifactReader + +logger = logging.getLogger(__name__) + + +class DeltaIngestor: + """Orchestrates CPG delta ingestion: compare → validate → load → enrich → purge → manifest. + + Usage (full artifact pipeline): + ingestor = DeltaIngestor(bucket="") + result = await ingestor.ingest_artifact( + manifest=manifest_dict, + nodes_data=[...], + edges_data=[...], + vectors_data=[...], + summaries_data=[...], + code_slices_data=[...], + graph_store=neptune_store, + vector_store=opensearch_store, + ) + + Usage (legacy — graph only): + result = await ingestor.ingest( + repo="amigo-core", + job_id="uuid", + nodes_data=[...], + edges_data=[...], + graph_store=neptune_store, + write_fn=my_write_function, + ) + """ + + def __init__(self, bucket: str, prefix: str = "cpg-exports", region: str = "us-east-1"): + self._manifest_mgr = ManifestManager(bucket, prefix, region) + + async def ingest_artifact( + self, + manifest: dict, + nodes_data: list[dict], + edges_data: list[dict], + vectors_data: list[dict], + summaries_data: list[dict], + code_slices_data: list[dict], + graph_store: Any, + vector_store: Any, + tenant_id: Optional[str] = None, + validate: bool = True, + batch_size: int = 500, + ) -> dict: + """Execute the full CPG artifact ingestion pipeline. + + Args: + manifest: Parsed manifest.json dict + nodes_data: Parsed records from nodes.jsonl + edges_data: Parsed records from edges.jsonl + vectors_data: Parsed records from vectors.jsonl + summaries_data: Parsed records from summaries.jsonl + code_slices_data: Parsed records from code_slices.jsonl + graph_store: Neptune graph store (supports write_nodes, write_edges, update_node_properties, execute_query) + vector_store: OpenSearch vector store (supports put_vectors) + tenant_id: Override tenant ID (derived from repo if not provided) + validate: Whether to validate the artifact before ingestion (default True) + + Returns: + Dict with status, counts, delta summary, and any validation warnings. + """ + repo = manifest.get("repo_id", "") + job_id = manifest.get("analysis_run_id", "") + tenant_id = tenant_id or self._derive_tenant_id(repo) + + # Step 1: Validate artifact + if validate: + validator = ArtifactValidator() + validation = validator.validate( + manifest=manifest, + sample_records={ + "nodes": nodes_data[:10], + "vectors": vectors_data[:10], + "summaries": summaries_data[:10], + }, + ) + if not validation.valid: + logger.error(f"Artifact validation failed for {repo}: {validation.errors}") + return { + "status": "REJECTED", + "reason": "validation_failed", + "errors": validation.errors, + "warnings": validation.warnings, + } + if validation.warnings: + logger.warning(f"Artifact validation warnings for {repo}: {validation.warnings}") + + # Step 2: Delta check + method_sigs = { + n.get("fully_qualified_name", n.get("full_name", "")): n.get("code_hash", n.get("hash", "")) + for n in nodes_data + if n.get("node_type") == "METHOD" and (n.get("fully_qualified_name") or n.get("full_name")) + } + + changed, previous = self._manifest_mgr.has_changes(repo, method_sigs) + + if not changed: + logger.info(f"Delta check: no changes for {repo}, skipping ingest") + return { + "status": "SKIPPED", + "reason": "no_changes", + "tenant_id": previous.tenant_id, + "previous_job_id": previous.job_id, + } + + if previous: + diff = GraphDiff.compare(previous.method_signatures, method_sigs) + logger.info(f"Delta: {diff.summary} for {repo}") + + # Step 3: Load graph (nodes + edges → Neptune) + graph_loader = GraphLoader(graph_store=graph_store, tenant_id=tenant_id, batch_size=batch_size) + graph_result = await graph_loader.load(nodes_data, edges_data) + logger.info(f"Graph loaded for {repo}: {graph_result}") + + # Step 4: Load vectors → OpenSearch (enriched with summary text for full-text search) + # Join summary text into vector records so OpenSearch indexes both vector + text + summary_by_node = {} + for s in summaries_data: + node_id = s.get("cpg_node_id", "") + summary_by_node.setdefault(node_id, {})[s.get("summary_type", "")] = s.get("text", "") + + enriched_vectors = [] + for v in vectors_data: + enriched = dict(v) + node_id = v.get("cpg_node_id", "") + target = v.get("embedding_target", "") + # If this vector's embedding_target matches a summary_type, include the text + if node_id in summary_by_node and target in summary_by_node[node_id]: + enriched["text"] = summary_by_node[node_id][target] + enriched_vectors.append(enriched) + + vector_loader = VectorLoader( + vector_store=vector_store, + expected_dimensions=manifest.get("embedding_dimensions", 0), + tenant_id=tenant_id, + ) + vector_result = await vector_loader.load(enriched_vectors) + logger.info(f"Vectors loaded for {repo}: {vector_result.total_loaded} vectors") + + # Step 5: Apply summaries → Neptune node properties + summary_overlay = SummaryOverlay(graph_store=graph_store) + summaries_applied = await summary_overlay.apply(summaries_data) + logger.info(f"Summaries applied for {repo}: {summaries_applied}") + + # Step 6: Store code slices → Neptune node properties + slice_store = CodeSliceStore(graph_store=graph_store) + slices_stored = await slice_store.store(code_slices_data) + logger.info(f"Code slices stored for {repo}: {slices_stored}") + + # Step 7: Purge old tenant (domain-scoped — only CPG nodes) + if previous and previous.tenant_id != tenant_id: + try: + await delete_domain(previous.tenant_id, "cpg", graph_store) + logger.info(f"Purged old CPG domain for tenant {previous.tenant_id}") + except Exception as e: + logger.warning(f"Failed to purge old CPG domain {previous.tenant_id}: {e}") + + # Step 8: Update manifest + new_manifest = Manifest( + repo=repo, + signature=self._manifest_mgr.compute_signature(method_sigs), + job_id=job_id, + tenant_id=tenant_id, + exported_at=datetime.now(timezone.utc).isoformat(), + nodes_path=manifest.get("nodes_path", ""), + edges_path=manifest.get("edges_path", ""), + node_count=graph_result.get("nodes_written", 0), + edge_count=graph_result.get("edges_written", 0), + language=manifest.get("language", ""), + method_signatures=method_sigs, + ) + self._manifest_mgr.put(new_manifest) + + return { + "status": "INGESTED", + "tenant_id": tenant_id, + "repo": repo, + "job_id": job_id, + "nodes_written": graph_result.get("nodes_written", 0), + "edges_written": graph_result.get("edges_written", 0), + "vectors_loaded": vector_result.total_loaded, + "summaries_applied": summaries_applied, + "code_slices_stored": slices_stored, + "delta": GraphDiff.compare( + previous.method_signatures if previous else {}, method_sigs + ).summary, + "warnings": validation.warnings if validate else [], + } + + async def ingest( + self, + repo: str, + job_id: str, + tenant_id: str, + nodes_data: list[dict], + edges_data: list[dict], + nodes_path: str, + edges_path: str, + graph_store: Any, + write_fn=None, + ) -> dict: + """Execute delta-aware ingestion (legacy — graph only, no vectors/summaries). + + Args: + repo: Repository name + job_id: Current job UUID + tenant_id: Derived tenant for this job + nodes_data: Parsed node records from Joern export + edges_data: Parsed edge records from Joern export + nodes_path: S3 URI to nodes.json + edges_path: S3 URI to edges.json + graph_store: Neptune graph store instance + write_fn: async callable(nodes_data, edges_data, tenant_id, graph_store) → dict + + Returns: + Dict with status, nodes_written, etc. + """ + # Extract method signatures from nodes + method_sigs = { + n["full_name"]: n.get("hash", "") + for n in nodes_data + if n.get("node_type") == "METHOD" and n.get("full_name") + } + + # Check manifest + changed, previous = self._manifest_mgr.has_changes(repo, method_sigs) + + if not changed: + logger.info(f"Delta check: no changes for {repo}, skipping ingest") + return { + "status": "SKIPPED", + "reason": "no_changes", + "tenant_id": previous.tenant_id, + "previous_job_id": previous.job_id, + } + + # Log diff if previous exists + if previous: + diff = GraphDiff.compare(previous.method_signatures, method_sigs) + logger.info(f"Delta: {diff.summary} for {repo}") + + # Perform full ingest + if write_fn: + result = await write_fn(nodes_data, edges_data, tenant_id, graph_store) + else: + result = {"nodes_written": len(nodes_data), "edges_written": len(edges_data)} + + # Purge old tenant + if previous and previous.tenant_id != tenant_id: + try: + await delete_tenant(previous.tenant_id, graph_store) + except Exception: + pass # logged inside delete_tenant + + # Update manifest + new_manifest = Manifest( + repo=repo, + signature=self._manifest_mgr.compute_signature(method_sigs), + job_id=job_id, + tenant_id=tenant_id, + exported_at=datetime.now(timezone.utc).isoformat(), + nodes_path=nodes_path, + edges_path=edges_path, + method_signatures=method_sigs, + ) + self._manifest_mgr.put(new_manifest) + + result["status"] = "INGESTED" + result["tenant_id"] = tenant_id + result["delta"] = GraphDiff.compare( + previous.method_signatures if previous else {}, method_sigs + ).summary + + return result + + async def rollback(self, repo: str, graph_store: Any, vector_store: Any) -> dict: + """Rollback to the previous artifact version. + + Reads the previous manifest, fetches the previous artifact from S3, + and re-ingests it (replacing the current bad data). + + Args: + repo: Repository name to rollback + graph_store: Neptune graph store + vector_store: OpenSearch vector store + + Returns: + Dict with rollback status and details. + """ + previous = self._manifest_mgr.get_previous(repo) + if not previous: + return { + "status": "FAILED", + "reason": "no_previous_manifest", + "message": f"No previous manifest found for {repo}. Cannot rollback.", + } + + # Determine S3 URI of previous artifact + s3_prefix = f"s3://{self._manifest_mgr._bucket}/{self._manifest_mgr._prefix}/{repo}/{previous.job_id}/" + + # Read previous artifact from S3 + reader = ArtifactReader(region=self._manifest_mgr._s3.meta.region_name) + try: + artifact = reader.read_s3(s3_prefix) + except Exception as e: + return { + "status": "FAILED", + "reason": "artifact_read_failed", + "message": f"Failed to read previous artifact from {s3_prefix}: {e}", + } + + if not artifact.manifest: + artifact.manifest = { + "repo_id": previous.repo, + "analysis_run_id": previous.job_id, + "embedding_dimensions": 0, + } + + # Re-ingest the previous artifact (this triggers domain-scoped delete + reload) + logger.info(f"Rolling back {repo} to job_id={previous.job_id}") + result = await self.ingest_artifact( + manifest=artifact.manifest, + nodes_data=artifact.nodes, + edges_data=artifact.edges, + vectors_data=artifact.vectors, + summaries_data=artifact.summaries, + code_slices_data=artifact.code_slices, + graph_store=graph_store, + vector_store=vector_store, + tenant_id=previous.tenant_id, + validate=False, # Previous artifact was already validated + ) + + result["rollback_from"] = self._manifest_mgr.get(repo).job_id if self._manifest_mgr.get(repo) else "unknown" + result["rollback_to"] = previous.job_id + return result + + @staticmethod + def _derive_tenant_id(repo: str) -> str: + """Derive a valid tenant_id from a repository name. + + Tenant IDs must be 1-25 lowercase chars (letters, numbers, periods — not at start/end). + """ + # Lowercase, replace non-alphanumeric with periods, trim to 25 chars + tid = repo.lower().replace("/", ".").replace("-", ".").replace("_", ".") + # Remove leading/trailing periods + tid = tid.strip(".") + # Collapse consecutive periods + while ".." in tid: + tid = tid.replace("..", ".") + # Truncate to 25 chars + tid = tid[:25].rstrip(".") + # Ensure at least 1 char + return tid or "default" diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graph_diff.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graph_diff.py new file mode 100644 index 00000000..b18a9db0 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graph_diff.py @@ -0,0 +1,64 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Graph Diff — compare two CPG states and compute delta.""" + +from dataclasses import dataclass, field +from typing import Dict + + +@dataclass +class DiffResult: + """Result of comparing two method signature sets.""" + + added: Dict[str, str] = field(default_factory=dict) # full_name → hash (new methods) + removed: Dict[str, str] = field(default_factory=dict) # full_name → hash (deleted methods) + modified: Dict[str, str] = field(default_factory=dict) # full_name → new_hash (body changed) + unchanged: int = 0 + + @property + def has_changes(self) -> bool: + return bool(self.added or self.removed or self.modified) + + @property + def summary(self) -> str: + return f"+{len(self.added)} -{len(self.removed)} ~{len(self.modified)} ={self.unchanged}" + + +class GraphDiff: + """Compare method signatures between two CPG exports.""" + + @staticmethod + def compare( + previous: Dict[str, str], + current: Dict[str, str], + ) -> DiffResult: + """Compare previous vs current method_signatures dicts. + + Args: + previous: {full_name: hash} from manifest + current: {full_name: hash} from new export + + Returns: + DiffResult with added/removed/modified/unchanged counts + """ + prev_keys = set(previous.keys()) + curr_keys = set(current.keys()) + + added = {k: current[k] for k in curr_keys - prev_keys} + removed = {k: previous[k] for k in prev_keys - curr_keys} + + modified = {} + unchanged = 0 + for k in prev_keys & curr_keys: + if previous[k] != current[k]: + modified[k] = current[k] + else: + unchanged += 1 + + return DiffResult( + added=added, + removed=removed, + modified=modified, + unchanged=unchanged, + ) diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graph_loader.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graph_loader.py new file mode 100644 index 00000000..cb988676 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graph_loader.py @@ -0,0 +1,248 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Graph Loader — loads CPG nodes and edges from artifact data into Neptune. + +Converts raw artifact dicts (from nodes.jsonl / edges.jsonl) into typed CPG +models and writes them in batches to a graph store, applying tenant isolation. +""" + +import logging +from typing import Any, Protocol, runtime_checkable + +from .models import CPGNode, CPGEdge + +logger = logging.getLogger(__name__) + + +@runtime_checkable +class GraphStoreProtocol(Protocol): + """Protocol for graph store backends (e.g., Neptune).""" + + async def write_nodes(self, nodes: list[dict]) -> int: + """Write a batch of node dicts to the graph store. + + Each dict has keys: 'id', 'label', 'properties'. + + Returns: + Number of nodes written. + """ + ... + + async def write_edges(self, edges: list[dict]) -> int: + """Write a batch of edge dicts to the graph store. + + Each dict has keys: 'source_id', 'target_id', 'edge_type', 'properties'. + + Returns: + Number of edges written. + """ + ... + + +class GraphLoader: + """Loads CPG nodes and edges from artifact data into a graph store. + + Converts raw artifact dicts into CPGNode/CPGEdge models, applies tenant + isolation via tenant_id, and writes in configurable batches. + + Usage: + loader = GraphLoader( + graph_store=neptune_store, + tenant_id="amigo-core__abc123", + batch_size=200, + ) + result = await loader.load(nodes_data, edges_data) + # result = {"nodes_written": 1500, "edges_written": 3200} + """ + + def __init__( + self, + graph_store: GraphStoreProtocol, + tenant_id: str, + batch_size: int = 500, + ) -> None: + """Initialize the GraphLoader. + + Args: + graph_store: Backend implementing GraphStoreProtocol. + tenant_id: Tenant identifier applied to all nodes/edges for isolation. + batch_size: Number of nodes/edges per write batch. Default 100. + """ + self._graph_store = graph_store + self._tenant_id = tenant_id + self._batch_size = batch_size + + async def load( + self, + nodes_data: list[dict], + edges_data: list[dict], + ) -> dict[str, Any]: + """Load nodes and edges into the graph store. + + Converts raw artifact dicts to typed CPG models, applies tenant_id, + batches writes, and returns summary counts. + + Args: + nodes_data: List of node dicts from nodes.jsonl artifact. + edges_data: List of edge dicts from edges.jsonl artifact. + + Returns: + Dict with 'nodes_written' and 'edges_written' counts. + """ + logger.info( + "GraphLoader.load: tenant_id=%s, nodes=%d, edges=%d, batch_size=%d", + self._tenant_id, + len(nodes_data), + len(edges_data), + self._batch_size, + ) + + nodes_written = await self._load_nodes(nodes_data) + edges_written = await self._load_edges(edges_data) + + logger.info( + "GraphLoader.load complete: nodes_written=%d, edges_written=%d", + nodes_written, + edges_written, + ) + + return { + "nodes_written": nodes_written, + "edges_written": edges_written, + } + + async def _load_nodes(self, nodes_data: list[dict]) -> int: + """Convert and write nodes in batches. + + Args: + nodes_data: Raw node dicts from artifact. + + Returns: + Total number of nodes written. + """ + total_written = 0 + total = len(nodes_data) + + for batch_start in range(0, total, self._batch_size): + batch_end = min(batch_start + self._batch_size, total) + raw_batch = nodes_data[batch_start:batch_end] + + write_batch = [] + for raw in raw_batch: + node = CPGNode.from_artifact(raw) + write_batch.append(self._node_to_write_dict(node)) + + count = await self._graph_store.write_nodes(write_batch) + total_written += count + + logger.debug( + "Nodes batch [%d-%d] of %d: wrote %d", + batch_start, + batch_end, + total, + count, + ) + + return total_written + + async def _load_edges(self, edges_data: list[dict]) -> int: + """Convert and write edges in batches. + + Args: + edges_data: Raw edge dicts from artifact. + + Returns: + Total number of edges written. + """ + total_written = 0 + total = len(edges_data) + + for batch_start in range(0, total, self._batch_size): + batch_end = min(batch_start + self._batch_size, total) + raw_batch = edges_data[batch_start:batch_end] + + write_batch = [] + for raw in raw_batch: + edge = CPGEdge.from_artifact(raw) + write_batch.append(self._edge_to_write_dict(edge)) + + count = await self._graph_store.write_edges(write_batch) + total_written += count + + logger.debug( + "Edges batch [%d-%d] of %d: wrote %d", + batch_start, + batch_end, + total, + count, + ) + + return total_written + + def _node_to_write_dict(self, node: CPGNode) -> dict: + """Convert a CPGNode to the write dict format expected by GraphStoreProtocol. + + Applies tenant_id as both a label prefix and a property for isolation. + + Args: + node: Typed CPGNode instance. + + Returns: + Dict with 'id', 'label', 'properties' keys. + """ + properties = { + "tenant_id": self._tenant_id, + "domain": "cpg", + "full_name": node.full_name, + "hash": node.hash, + "filename": node.filename, + "name": node.name, + "code": node.code, + "signature": node.signature, + "type_full_name": node.type_full_name, + "is_external": node.is_external, + } + + # Include optional fields only when present + if node.line_number is not None: + properties["line_number"] = node.line_number + if node.line_number_end is not None: + properties["line_number_end"] = node.line_number_end + if node.order is not None: + properties["order"] = node.order + + # Merge any extra properties from the raw artifact + for key, value in node.properties.items(): + if key not in properties: + properties[key] = value + + return { + "id": node.id, + "label": f"{self._tenant_id}__{node.node_type}", + "properties": properties, + } + + def _edge_to_write_dict(self, edge: CPGEdge) -> dict: + """Convert a CPGEdge to the write dict format expected by GraphStoreProtocol. + + Applies tenant_id as a property for isolation. + + Args: + edge: Typed CPGEdge instance. + + Returns: + Dict with 'source_id', 'target_id', 'edge_type', 'properties' keys. + """ + properties = { + "tenant_id": self._tenant_id, + "domain": "cpg", + **edge.properties, + } + + return { + "source_id": edge.source_id, + "target_id": edge.target_id, + "edge_type": edge.edge_type, + "properties": properties, + } diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graphson_converter.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graphson_converter.py new file mode 100644 index 00000000..6c636a71 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/graphson_converter.py @@ -0,0 +1,243 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GraphSON → JSONL Converter — transforms Joern GraphSON exports to portable CPG JSONL. + +Joern's export produces GraphSON format (JSON with nodes[] and edges[] arrays). +This converter reads that format and outputs our portable JSONL schema with +stable cpg_node_id identifiers (replacing Joern's internal numeric IDs). + +Usage: + converter = GraphSONConverter(repo_id="payments-api", commit_sha="abc123") + nodes, edges = converter.convert_file("path/to/joern-export.json") + converter.write_jsonl(nodes, "output/nodes.jsonl") + converter.write_jsonl(edges, "output/edges.jsonl") +""" + +import json +import logging +from pathlib import Path +from typing import IO + +from .models import CPGNode, CPGEdge + +logger = logging.getLogger(__name__) + + +class GraphSONConverter: + """Converts Joern GraphSON exports to portable CPG JSONL format. + + The conversion: + 1. Reads GraphSON (JSON array or {nodes:[], edges:[]}) + 2. Parses each record via CPGNode.from_joern() / CPGEdge.from_joern() + 3. Generates stable cpg_node_id (replacing numeric IDs) + 4. Outputs JSONL with portable identifiers + + The cpg_node_id format: :::: + """ + + def __init__(self, repo_id: str, commit_sha: str = ""): + self._repo_id = repo_id + self._commit_sha = commit_sha + self._id_map: dict[str, str] = {} # joern_id → cpg_node_id + + def convert_file(self, path: str) -> tuple[list[dict], list[dict]]: + """Convert a GraphSON file to portable JSONL records. + + Accepts either: + - A single JSON file with {"nodes": [...], "edges": [...]} + - Separate nodes.json and edges.json (pass nodes path, infers edges) + + Args: + path: Path to the GraphSON file. + + Returns: + Tuple of (node_records, edge_records) ready for JSONL output. + """ + path = Path(path) + with open(path) as f: + data = json.load(f) + + # Handle TinkerPop GraphSON format: {"@type": "tinker:graph", "@value": {"vertices": [...], "edges": [...]}} + if isinstance(data, dict) and data.get("@type") == "tinker:graph": + graph_value = data.get("@value", {}) + raw_nodes = [self._unwrap_vertex(v) for v in graph_value.get("vertices", [])] + raw_edges = [self._unwrap_edge(e) for e in graph_value.get("edges", [])] + elif isinstance(data, dict): + raw_nodes = data.get("nodes", data.get("vertices", [])) + raw_edges = data.get("edges", []) + elif isinstance(data, list): + # Heuristic: if items have "src"/"dst", they're edges; otherwise nodes + if data and ("src" in data[0] or "dst" in data[0]): + raw_nodes = [] + raw_edges = data + else: + raw_nodes = data + raw_edges = [] + # Try to find edges file alongside + edges_path = path.parent / "edges.json" + if edges_path.exists(): + with open(edges_path) as f: + raw_edges = json.load(f) + else: + raw_nodes, raw_edges = [], [] + + nodes = self._convert_nodes(raw_nodes) + edges = self._convert_edges(raw_edges) + + logger.info( + f"GraphSON converted: {len(nodes)} nodes, {len(edges)} edges " + f"(repo={self._repo_id}, commit={self._commit_sha})" + ) + return nodes, edges + + def convert_nodes_stream(self, stream: IO) -> list[dict]: + """Convert a stream of GraphSON nodes (JSON array).""" + raw_nodes = json.load(stream) + if isinstance(raw_nodes, dict): + raw_nodes = raw_nodes.get("nodes", raw_nodes.get("vertices", [])) + return self._convert_nodes(raw_nodes) + + def convert_edges_stream(self, stream: IO) -> list[dict]: + """Convert a stream of GraphSON edges (JSON array).""" + raw_edges = json.load(stream) + if isinstance(raw_edges, dict): + raw_edges = raw_edges.get("edges", []) + return self._convert_edges(raw_edges) + + @staticmethod + def _unwrap_value(val): + """Unwrap a TinkerPop typed value: {"@type": "g:Int64", "@value": 123} → 123.""" + if isinstance(val, dict) and "@value" in val: + inner = val["@value"] + # Handle nested lists: {"@type": "g:List", "@value": ["foo"]} + if isinstance(inner, list): + return inner[0] if len(inner) == 1 else inner + return inner + return val + + def _unwrap_vertex(self, vertex: dict) -> dict: + """Convert a TinkerPop g:Vertex to a flat Joern-style dict for from_joern(). + + TinkerPop format: + {"@type": "g:Vertex", "id": {"@type": "g:Int64", "@value": 123}, + "label": "METHOD", "properties": {"FULL_NAME": {"@type": "g:VertexProperty", "@value": {...}}}} + + Joern format (what from_joern expects): + {"id": "123", "label": "METHOD", "properties": {"FULL_NAME": "pkg.Foo.bar"}} + """ + raw_id = self._unwrap_value(vertex.get("id", "")) + label = vertex.get("label", "UNKNOWN") + + # Unwrap properties from TinkerPop VertexProperty format + props = {} + raw_props = vertex.get("properties", {}) + for key, prop_val in raw_props.items(): + if isinstance(prop_val, dict): + # {"@type": "g:VertexProperty", "@value": {"@type": "g:List", "@value": ["actual_value"]}} + inner = prop_val.get("@value", prop_val) + if isinstance(inner, dict) and "@value" in inner: + inner = inner["@value"] + if isinstance(inner, list): + val = inner[0] if len(inner) == 1 else inner + else: + val = inner + # Final unwrap if value is still a typed wrapper (e.g., {"@type": "g:Int32", "@value": 4}) + if isinstance(val, dict) and "@value" in val: + val = val["@value"] + props[key] = val + else: + props[key] = prop_val + + return {"id": str(raw_id), "label": label, "properties": props} + + def _unwrap_edge(self, edge: dict) -> dict: + """Convert a TinkerPop g:Edge to a flat Joern-style dict for from_joern(). + + TinkerPop format: + {"@type": "g:Edge", "id": {...}, "inV": {"@type": "g:Int64", "@value": 456}, + "outV": {"@type": "g:Int64", "@value": 123}, "label": "AST", "properties": {}} + + Joern format: + {"src": "123", "dst": "456", "label": "AST", "properties": {}} + """ + src = str(self._unwrap_value(edge.get("outV", ""))) + dst = str(self._unwrap_value(edge.get("inV", ""))) + label = edge.get("label", "UNKNOWN") + + # Unwrap edge properties if any + props = {} + for key, val in edge.get("properties", {}).items(): + props[key] = self._unwrap_value(val) if isinstance(val, dict) else val + + return {"src": src, "dst": dst, "label": label, "properties": props} + + def _convert_nodes(self, raw_nodes: list[dict]) -> list[dict]: + """Convert raw Joern nodes to portable JSONL records.""" + records = [] + for raw in raw_nodes: + node = CPGNode.from_joern(raw) + cpg_node_id = self._make_cpg_node_id(node) + self._id_map[node.id] = cpg_node_id + + records.append({ + "cpg_node_id": cpg_node_id, + "node_type": node.node_type, + "labels": [node.node_type], + "repo_id": self._repo_id, + "commit_sha": self._commit_sha, + "file_path": node.filename, + "fully_qualified_name": node.full_name, + "name": node.name, + "line_start": node.line_number, + "line_end": node.line_number_end, + "code_hash": node.hash, + "properties": node.properties, + }) + return records + + def _convert_edges(self, raw_edges: list[dict]) -> list[dict]: + """Convert raw Joern edges to portable JSONL records.""" + records = [] + for raw in raw_edges: + edge = CPGEdge.from_joern(raw) + source_cpg_id = self._id_map.get(edge.source_id, edge.source_id) + target_cpg_id = self._id_map.get(edge.target_id, edge.target_id) + + records.append({ + "source_cpg_node_id": source_cpg_id, + "target_cpg_node_id": target_cpg_id, + "edge_type": edge.edge_type, + "properties": edge.properties, + }) + return records + + def _make_cpg_node_id(self, node: CPGNode) -> str: + """Generate a stable cpg_node_id for a Joern node. + + Format: :::: + """ + identity = node.full_name or node.name or node.id + hash_part = node.hash or "" + parts = [self._repo_id, self._commit_sha, node.node_type, identity] + if hash_part: + parts.append(hash_part) + return ":".join(parts) + + @staticmethod + def write_jsonl(records: list[dict], output_path: str) -> int: + """Write records to a JSONL file (one JSON object per line). + + Args: + records: List of dicts to write. + output_path: Path to output file. + + Returns: + Number of records written. + """ + path = Path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + for record in records: + f.write(json.dumps(record) + "\n") + return len(records) diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/manifest_manager.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/manifest_manager.py new file mode 100644 index 00000000..f1821eb0 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/manifest_manager.py @@ -0,0 +1,174 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Manifest Manager — S3-backed CPG state tracking with optimistic locking. + +Provides: +- Manifest read/write with ETag-based optimistic locking +- Change detection (signature comparison) +- Rollback support (maintains previous manifest reference) +""" + +import hashlib +import json +import logging +from typing import Optional + +import boto3 +from botocore.exceptions import ClientError + +from .models import Manifest + +logger = logging.getLogger(__name__) + + +class ManifestConflictError(Exception): + """Raised when a concurrent write conflicts with our manifest update.""" + pass + + +class ManifestManager: + """Read/write/compare CPG manifests on S3 with optimistic locking.""" + + def __init__(self, bucket: str, prefix: str = "cpg-exports", region: str = "us-east-1"): + self._bucket = bucket + self._prefix = prefix + self._s3 = boto3.client("s3", region_name=region) + self._etags: dict[str, str] = {} # repo → last-known ETag + + def compute_signature(self, method_signatures: dict[str, str]) -> str: + """Compute sha256 signature from method full_name:hash pairs.""" + payload = json.dumps(method_signatures, sort_keys=True) + return "sha256:" + hashlib.sha256(payload.encode()).hexdigest() + + def get(self, repo: str) -> Optional[Manifest]: + """Read manifest for a repo from S3. Returns None if not found. + + Stores the ETag for subsequent optimistic locking on put(). + """ + key = self._manifest_key(repo) + try: + resp = self._s3.get_object(Bucket=self._bucket, Key=key) + self._etags[repo] = resp.get("ETag", "") + data = json.loads(resp["Body"].read()) + return Manifest(**data) + except ClientError as e: + error_code = e.response["Error"]["Code"] + if error_code in ("NoSuchKey", "AccessDenied", "403"): + # NoSuchKey = first-time ingest (no manifest yet) + # AccessDenied = S3/KMS permissions not yet configured (treat as first-time) + if error_code != "NoSuchKey": + logger.warning(f"Manifest read got {error_code} for {repo} — treating as first-time ingest") + return None + raise + except Exception: + return None + + def put(self, manifest: Manifest) -> None: + """Write manifest to S3 with optimistic locking. + + If the manifest was previously read (ETag stored), uses If-Match + to prevent concurrent overwrites. If the ETag doesn't match + (another writer updated it), raises ManifestConflictError. + + Args: + manifest: The manifest to write. + + Raises: + ManifestConflictError: If a concurrent write detected. + """ + key = self._manifest_key(manifest.repo) + body = json.dumps({ + "repo": manifest.repo, + "signature": manifest.signature, + "job_id": manifest.job_id, + "tenant_id": manifest.tenant_id, + "exported_at": manifest.exported_at, + "nodes_path": manifest.nodes_path, + "edges_path": manifest.edges_path, + "node_count": manifest.node_count, + "edge_count": manifest.edge_count, + "language": manifest.language, + "method_signatures": manifest.method_signatures, + }) + + put_kwargs = { + "Bucket": self._bucket, + "Key": key, + "Body": body.encode(), + "ContentType": "application/json", + } + + # Optimistic locking: if we have a previous ETag, require it matches + etag = self._etags.get(manifest.repo) + if etag: + put_kwargs["IfMatch"] = etag + + try: + resp = self._s3.put_object(**put_kwargs) + # Update stored ETag for next operation + self._etags[manifest.repo] = resp.get("ETag", "") + logger.info(f"Manifest written: s3://{self._bucket}/{key}") + except ClientError as e: + error_code = e.response["Error"]["Code"] + if error_code in ("PreconditionFailed", "412"): + raise ManifestConflictError( + f"Concurrent manifest update detected for {manifest.repo}. " + f"Another process wrote a newer version. Retry the ingest." + ) from e + if error_code in ("AccessDenied", "403"): + logger.warning(f"Manifest write AccessDenied for {manifest.repo} — S3/KMS permissions not configured. Ingest succeeded but manifest not persisted (delta detection disabled until fixed).") + return + raise + + def get_previous(self, repo: str) -> Optional[Manifest]: + """Read the previous manifest (for rollback). + + The previous manifest is stored at //manifest.previous.json. + """ + key = f"{self._prefix}/{repo}/manifest.previous.json" + try: + resp = self._s3.get_object(Bucket=self._bucket, Key=key) + data = json.loads(resp["Body"].read()) + return Manifest(**data) + except Exception: + return None + + def put_with_history(self, manifest: Manifest) -> None: + """Write manifest to S3, preserving current as previous (for rollback). + + 1. Copies current manifest.json → manifest.previous.json + 2. Writes new manifest.json (with optimistic lock) + """ + key = self._manifest_key(manifest.repo) + prev_key = f"{self._prefix}/{manifest.repo}/manifest.previous.json" + + # Copy current to previous (if it exists) + try: + self._s3.copy_object( + Bucket=self._bucket, + CopySource={"Bucket": self._bucket, "Key": key}, + Key=prev_key, + ) + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchKey": + logger.warning(f"Failed to archive previous manifest: {e}") + + # Write new manifest with lock + self.put(manifest) + + def has_changes(self, repo: str, current_signatures: dict[str, str]) -> tuple[bool, Optional[Manifest]]: + """Check if current export differs from stored manifest. + + Returns: + (has_changes: bool, previous_manifest: Optional[Manifest]) + """ + previous = self.get(repo) + if not previous: + return True, None + new_sig = self.compute_signature(current_signatures) + return new_sig != previous.signature, previous + + def _manifest_key(self, repo: str) -> str: + """Get the S3 key for a repo's manifest.""" + return f"{self._prefix}/{repo}/manifest.json" diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/models.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/models.py new file mode 100644 index 00000000..4a62959b --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/models.py @@ -0,0 +1,313 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPG domain models — typed representations of Joern code property graph output. + +These models accept ALL Joern node/edge types (the full CPG schema). +See schema.py for the complete type enumeration and property specs. + +Node types (20): METHOD, CALL, IDENTIFIER, LITERAL, BLOCK, COMMENT, FILE, + LOCAL, MEMBER, META_DATA, METHOD_REF, METHOD_RETURN, MODIFIER, NAMESPACE, + NAMESPACE_BLOCK, PARAMETER, RETURN, TAG, TYPE_DECL, TYPE + +Edge types (14): AST, CFG, CDG, REACHING_DEF, CALL, ARGUMENT, DOMINATE, + POST_DOMINATE, CONTAINS, BINDS_TO, REF, INHERITS_FROM, CONDITION, TAGGED_BY + +Supported languages: Java, JavaScript/TypeScript, Python, C/C++, Go, PHP, + Ruby, Kotlin, Swift (via Joern frontends) +""" + +from dataclasses import dataclass, field +from typing import Dict, Any, List, Optional + + +@dataclass +class CPGNode: + """A code property graph node from Joern export. + + Accepts any node_type — the full Joern schema (20 types) flows through. + See schema.NodeType for the complete enumeration. + + Identity: full_name (stable across line shifts — used for delta comparison) + Change detection: hash (content fingerprint computed by Joern) + + Key properties by type: + METHOD: full_name, signature, hash, is_external, line_number + CALL: full_name (callee), code, line_number + TYPE_DECL: full_name, name, is_external + IDENTIFIER: name, code, line_number + LITERAL: code, line_number + PARAMETER: name, type_full_name, order + FILE: name (path) + """ + + id: str + node_type: str # Any of the 20 Joern node types (see schema.NodeType) + full_name: str = "" + hash: str = "" + filename: str = "" + name: str = "" + code: str = "" + signature: str = "" + line_number: Optional[int] = None + line_number_end: Optional[int] = None + order: Optional[int] = None + type_full_name: str = "" + is_external: bool = False + properties: Dict[str, Any] = field(default_factory=dict) + + @property + def stable_id(self) -> str: + """Content-addressed identity for delta comparison. + + For METHOD nodes: full_name is stable across line shifts. + For other nodes: falls back to Joern's numeric id. + """ + return self.full_name or self.id + + @classmethod + def from_joern(cls, raw: dict) -> "CPGNode": + """Create a CPGNode from a raw Joern JSON export entry. + + Args: + raw: Dict from Joern's nodes.json export with 'id', 'label', 'properties' keys. + + Returns: + CPGNode with Joern properties mapped to typed fields. + """ + props = raw.get("properties", {}) + return cls( + id=str(raw.get("id", "")), + node_type=raw.get("label", "UNKNOWN"), + full_name=props.get("FULL_NAME", props.get("fullName", "")), + hash=props.get("HASH", props.get("hash", "")), + filename=props.get("FILENAME", props.get("filename", "")), + name=props.get("NAME", props.get("name", "")), + code=props.get("CODE", props.get("code", "")), + signature=props.get("SIGNATURE", props.get("signature", "")), + line_number=props.get("LINE_NUMBER", props.get("lineNumber")), + line_number_end=props.get("LINE_NUMBER_END", props.get("lineNumberEnd")), + order=props.get("ORDER", props.get("order")), + type_full_name=props.get("TYPE_FULL_NAME", props.get("typeFullName", "")), + is_external=props.get("IS_EXTERNAL", props.get("isExternal", False)), + properties=props, + ) + + @classmethod + def from_artifact(cls, raw: dict) -> "CPGNode": + """Create a CPGNode from a portable CPG artifact schema entry. + + Args: + raw: Dict with keys: cpg_node_id, node_type, labels, repo_id, + commit_sha, file_path, fully_qualified_name, name, + line_start, line_end, code_hash, properties. + + Returns: + CPGNode with artifact fields mapped to typed fields. + """ + return cls( + id=raw.get("cpg_node_id", ""), + node_type=raw.get("node_type", "UNKNOWN"), + full_name=raw.get("fully_qualified_name", ""), + hash=raw.get("code_hash", ""), + filename=raw.get("file_path", ""), + name=raw.get("name", ""), + line_number=raw.get("line_start"), + line_number_end=raw.get("line_end"), + properties=raw.get("properties", {}), + ) + + +@dataclass +class CPGEdge: + """A code property graph edge from Joern export. + + Accepts any edge_type — the full Joern schema (14+ types) flows through. + See schema.EdgeType for the complete enumeration. + + Edge layers: + AST: Syntax tree structure (parent → child) + CFG: Control flow (execution order) + CDG: Control dependence (branching influence) + REACHING_DEF: Data flow (definition → use) + CALL: Call graph (caller → callee) + ARGUMENT: Call → argument mapping + CONTAINS: File/Type → contained elements + INHERITS_FROM: Type hierarchy + """ + + source_id: str + target_id: str + edge_type: str # Any of the 14+ Joern edge types (see schema.EdgeType) + properties: Dict[str, Any] = field(default_factory=dict) + + @property + def key(self) -> str: + """Unique edge identity for diff.""" + return f"{self.source_id}->{self.edge_type}->{self.target_id}" + + @classmethod + def from_joern(cls, raw: dict) -> "CPGEdge": + """Create a CPGEdge from a raw Joern JSON export entry. + + Args: + raw: Dict from Joern's edges.json export with 'src', 'dst', 'label' keys. + + Returns: + CPGEdge with Joern properties mapped. + """ + return cls( + source_id=str(raw.get("src", raw.get("source_id", ""))), + target_id=str(raw.get("dst", raw.get("target_id", ""))), + edge_type=raw.get("label", raw.get("edge_type", "UNKNOWN")), + properties=raw.get("properties", {}), + ) + + @classmethod + def from_artifact(cls, raw: dict) -> "CPGEdge": + """Create a CPGEdge from a portable CPG artifact schema entry. + + Args: + raw: Dict with keys: edge_id, source_cpg_node_id, + target_cpg_node_id, edge_type, properties. + + Returns: + CPGEdge with artifact fields mapped. + """ + return cls( + source_id=raw.get("source_cpg_node_id", ""), + target_id=raw.get("target_cpg_node_id", ""), + edge_type=raw.get("edge_type", "UNKNOWN"), + properties=raw.get("properties", {}), + ) + + +@dataclass +class Manifest: + """CPG extraction manifest — tracks the current graph state for a repo. + + The manifest stores method signatures ({full_name: hash}) as the + fingerprint for delta comparison. Only METHOD nodes participate in + change detection because method body changes represent meaningful + code behavior changes. Other changes (comments, formatting, metadata) + don't affect program behavior. + """ + + repo: str + signature: str # sha256 of sorted method full_name:hash pairs + job_id: str + tenant_id: str + exported_at: str + nodes_path: str + edges_path: str + node_count: int = 0 + edge_count: int = 0 + language: str = "" + method_signatures: Dict[str, str] = field(default_factory=dict) # full_name → hash + + +@dataclass +class VectorRecord: + """An embedding vector associated with a CPG node. + + Used for semantic search over code elements (methods, types, etc.). + """ + + cpg_node_id: str + embedding_target: str + embedding_model: str + embedding_dimensions: int + similarity_function: str + embedding_text_hash: str + vector: List[float] = field(default_factory=list) + + @classmethod + def from_artifact(cls, raw: dict) -> "VectorRecord": + """Create a VectorRecord from a portable artifact schema entry. + + Args: + raw: Dict with keys: cpg_node_id, embedding_target, embedding_model, + embedding_dimensions, similarity_function, embedding_text_hash, vector. + + Returns: + VectorRecord with artifact fields mapped. + """ + return cls( + cpg_node_id=raw.get("cpg_node_id", ""), + embedding_target=raw.get("embedding_target", ""), + embedding_model=raw.get("embedding_model", ""), + embedding_dimensions=raw.get("embedding_dimensions", 0), + similarity_function=raw.get("similarity_function", ""), + embedding_text_hash=raw.get("embedding_text_hash", ""), + vector=raw.get("vector", []), + ) + + +@dataclass +class SummaryRecord: + """An LLM-generated summary associated with a CPG node. + + Used for natural-language descriptions of code elements. + """ + + cpg_node_id: str + summary_type: str + text: str + generation_model: str + generation_prompt_version: str + + @classmethod + def from_artifact(cls, raw: dict) -> "SummaryRecord": + """Create a SummaryRecord from a portable artifact schema entry. + + Args: + raw: Dict with keys: cpg_node_id, summary_type, text, + generation_model, generation_prompt_version. + + Returns: + SummaryRecord with artifact fields mapped. + """ + return cls( + cpg_node_id=raw.get("cpg_node_id", ""), + summary_type=raw.get("summary_type", ""), + text=raw.get("text", ""), + generation_model=raw.get("generation_model", ""), + generation_prompt_version=raw.get("generation_prompt_version", ""), + ) + + +@dataclass +class CodeSliceRecord: + """A code slice (excerpt) associated with a CPG node. + + Represents extracted code fragments for analysis or display. + """ + + cpg_node_id: str + slice_type: str + code: str + language: str + line_start: Optional[int] = None + line_end: Optional[int] = None + file_path: str = "" + + @classmethod + def from_artifact(cls, raw: dict) -> "CodeSliceRecord": + """Create a CodeSliceRecord from a portable artifact schema entry. + + Args: + raw: Dict with keys: cpg_node_id, slice_type, code, language, + line_start, line_end, file_path. + + Returns: + CodeSliceRecord with artifact fields mapped. + """ + return cls( + cpg_node_id=raw.get("cpg_node_id", ""), + slice_type=raw.get("slice_type", ""), + code=raw.get("code", ""), + language=raw.get("language", ""), + line_start=raw.get("line_start"), + line_end=raw.get("line_end"), + file_path=raw.get("file_path", ""), + ) diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/schema.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/schema.py new file mode 100644 index 00000000..e0dc5c80 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/schema.py @@ -0,0 +1,197 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPG Schema — Complete Joern Code Property Graph node and edge type definitions. + +This module defines the canonical node types, edge types, and their properties +as specified by the Joern CPG schema (https://docs.joern.io/). + +codeproperty-graph ingests ALL Joern types — the schema is not restrictive. +These enums and property specs serve as: +1. Documentation — what types exist and what they mean +2. Validation — optional strict mode to reject unknown types +3. Query helpers — typed constants for building Cypher queries + +Supported Joern frontends (languages): +- Java (javasrc, jimple) +- JavaScript/TypeScript (jssrc) +- Python (pysrc) +- C/C++ (c2cpg) +- Go, PHP, Ruby, Kotlin, Swift (community frontends) +""" + +from enum import Enum +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Set + + +class NodeType(str, Enum): + """All CPG node types as defined by Joern's schema. + + Reference: https://docs.joern.io/cpgql/node-type-steps/ + """ + + # === Structure === + FILE = "FILE" # Source file + NAMESPACE = "NAMESPACE" # Package/module namespace + NAMESPACE_BLOCK = "NAMESPACE_BLOCK" # Namespace scope block + TYPE_DECL = "TYPE_DECL" # Class/struct/interface declaration + TYPE = "TYPE" # Type reference + MEMBER = "MEMBER" # Class/struct member (field) + META_DATA = "META_DATA" # Graph metadata (language, version) + + # === Methods === + METHOD = "METHOD" # Function/method definition + METHOD_REF = "METHOD_REF" # Reference to a method (lambda, callback) + METHOD_RETURN = "METHOD_RETURN" # Formal return parameter + PARAMETER = "PARAMETER" # Formal parameter + MODIFIER = "MODIFIER" # Access modifier (public, private, static) + + # === Statements & Expressions === + BLOCK = "BLOCK" # Code block (scope) + CALL = "CALL" # Call site (method/function invocation) + IDENTIFIER = "IDENTIFIER" # Variable/member reference + LITERAL = "LITERAL" # Constant value (number, string) + LOCAL = "LOCAL" # Local variable declaration + RETURN = "RETURN" # Return statement + + # === Metadata === + COMMENT = "COMMENT" # Source code comment + TAG = "TAG" # User-defined tag/annotation + + +class EdgeType(str, Enum): + """All CPG edge types as defined by Joern's schema. + + Joern produces multiple graph layers overlaid on the same nodes: + - AST: Abstract Syntax Tree (parent-child structure) + - CFG: Control Flow Graph (execution order) + - CDG: Control Dependence Graph (branching influences) + - PDG: Program Dependence Graph (CDG + data deps) + - DDG: Data Dependence Graph (REACHING_DEF) + - CALL: Call graph (caller → callee) + """ + + # === AST Layer === + AST = "AST" # Parent → child in syntax tree + ARGUMENT = "ARGUMENT" # Call → argument expression + CONTAINS = "CONTAINS" # File/Type/Method → contained elements + BINDS_TO = "BINDS_TO" # Type parameter binding + + # === Control Flow Layer === + CFG = "CFG" # Control flow edge (execution order) + CONDITION = "CONDITION" # Conditional branch + + # === Data Flow Layer === + REACHING_DEF = "REACHING_DEF" # Data dependency (definition reaches use) + REF = "REF" # Reference edge (identifier → declaration) + + # === Dependence Layers === + CDG = "CDG" # Control dependence + DOMINATE = "DOMINATE" # Dominance relation + POST_DOMINATE = "POST_DOMINATE" # Post-dominance relation + + # === Call Graph === + CALL = "CALL" # Caller → callee method + + # === Type System === + INHERITS_FROM = "INHERITS_FROM" # Type inheritance/implementation + ALIAS_OF = "ALIAS_OF" # Type alias + + # === Annotations === + TAGGED_BY = "TAGGED_BY" # Node → Tag association + + +@dataclass +class NodePropertySpec: + """Expected properties for each node type per Joern's schema.""" + + common: List[str] = field(default_factory=lambda: [ + "id", "label", "order" + ]) + + by_type: Dict[str, List[str]] = field(default_factory=lambda: { + "METHOD": ["fullName", "name", "signature", "lineNumber", "lineNumberEnd", + "isExternal", "code", "filename", "hash"], + "CALL": ["name", "code", "lineNumber", "methodFullName", "signature", + "dispatchType", "argumentIndex"], + "IDENTIFIER": ["name", "code", "lineNumber", "order", "argumentIndex"], + "LITERAL": ["code", "lineNumber", "order", "argumentIndex"], + "LOCAL": ["name", "code", "lineNumber", "typeFullName"], + "PARAMETER": ["name", "code", "lineNumber", "order", "typeFullName"], + "TYPE_DECL": ["fullName", "name", "isExternal", "order"], + "TYPE": ["fullName", "name"], + "MEMBER": ["name", "code", "order", "typeFullName"], + "FILE": ["name", "order"], + "NAMESPACE": ["name", "order"], + "NAMESPACE_BLOCK": ["fullName", "name", "order"], + "BLOCK": ["lineNumber", "order", "argumentIndex"], + "METHOD_RETURN": ["code", "lineNumber", "order", "typeFullName"], + "METHOD_REF": ["code", "lineNumber", "order", "methodFullName"], + "MODIFIER": ["modifierType", "order"], + "RETURN": ["code", "lineNumber", "order", "argumentIndex"], + "COMMENT": ["code", "lineNumber"], + "TAG": ["name", "value"], + "META_DATA": ["language", "version"], + }) + + +# Delta-relevant node types: only these participate in change detection. +# Rationale: method body changes are the meaningful code changes. +# Other changes (comments, formatting, metadata) don't affect behavior. +DELTA_RELEVANT_TYPES: Set[str] = { + NodeType.METHOD, +} + +# Structure-relevant types: used for understanding code organization +# (not for delta, but for graph queries about architecture). +STRUCTURE_TYPES: Set[str] = { + NodeType.FILE, + NodeType.NAMESPACE, + NodeType.NAMESPACE_BLOCK, + NodeType.TYPE_DECL, + NodeType.TYPE, + NodeType.MEMBER, +} + +# Data-flow-relevant types: used for security analysis, taint tracking. +DATAFLOW_TYPES: Set[str] = { + NodeType.CALL, + NodeType.IDENTIFIER, + NodeType.LITERAL, + NodeType.PARAMETER, + NodeType.LOCAL, + NodeType.RETURN, +} + +# Supported Joern frontends (languages) +SUPPORTED_LANGUAGES = { + "java": "Java (javasrc frontend)", + "javascript": "JavaScript/TypeScript (jssrc frontend)", + "python": "Python (pysrc frontend)", + "c": "C/C++ (c2cpg frontend)", + "go": "Go (gosrc frontend)", + "php": "PHP (php2cpg frontend)", + "ruby": "Ruby (rubysrc frontend)", + "kotlin": "Kotlin (kotlin2cpg frontend)", + "swift": "Swift (swiftsrc frontend)", +} + + +def joern_export_command(source_path: str, output_dir: str = "cpg-export", + language: str = None) -> str: + """Generate the Joern CLI command to export a CPG. + + Args: + source_path: Path to the source code directory. + output_dir: Directory for the exported JSON files. + language: Optional language hint (auto-detected if omitted). + + Returns: + Shell command string to run Joern export. + """ + cmd = f"joern-export --repr cpg14 --format json --out {output_dir}" + if language: + cmd += f" --language {language}" + cmd += f" {source_path}" + return cmd diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/summary_overlay.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/summary_overlay.py new file mode 100644 index 00000000..07938d3f --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/summary_overlay.py @@ -0,0 +1,149 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Attaches summaries from summaries.jsonl to Neptune graph nodes as properties.""" + +import asyncio +import logging +from typing import Any, Dict, List, Protocol, runtime_checkable + +logger = logging.getLogger(__name__) + + +class SummaryRecord: + """Type alias documentation for summary record dicts. + + Expected keys: + cpg_node_id: str - The Neptune node identifier. + summary_type: str - The type of summary (e.g., "method_summary"). + summary_text: str - The generated summary content. + generation_model: str - The model used for generation. + generation_prompt_version: str - The prompt version used. + """ + + +@runtime_checkable +class GraphStoreProtocol(Protocol): + """Protocol for graph store implementations that support node property updates.""" + + async def update_node_properties(self, node_id: str, properties: dict) -> bool: + """Update properties on a graph node. + + Args: + node_id: The unique identifier of the node to update. + properties: A dictionary of property names to values to set. + + Returns: + True if the update was successful, False otherwise. + """ + ... + + +class SummaryOverlay: + """Applies summary records as properties on Neptune graph nodes. + + This class takes a collection of summary records and attaches them + to their corresponding nodes in the graph store. Summaries are stored + as node properties with the naming convention `summary_{summary_type}`. + + Args: + graph_store: An implementation of GraphStoreProtocol. + batch_size: Number of updates to process concurrently. Defaults to 50. + """ + + DEFAULT_BATCH_SIZE = 50 + + def __init__( + self, + graph_store: GraphStoreProtocol, + batch_size: int = DEFAULT_BATCH_SIZE, + ) -> None: + self._graph_store = graph_store + self._batch_size = batch_size + + async def apply(self, records: List[Dict[str, Any]]) -> int: + """Apply summary records to graph nodes as properties. + + For each record, sets the following properties on the node: + - summary_{summary_type}: The summary text. + - generation_model: The model used for generation. + - generation_prompt_version: The prompt version used. + + Args: + records: A list of summary record dicts. Each dict must contain + keys: cpg_node_id, summary_type, summary_text, + generation_model, generation_prompt_version. + + Returns: + The count of summaries successfully applied. + """ + if not records: + logger.info("No summary records to apply.") + return 0 + + logger.info("Applying %d summary records in batches of %d.", len(records), self._batch_size) + + applied_count = 0 + + for batch_start in range(0, len(records), self._batch_size): + batch = records[batch_start : batch_start + self._batch_size] + results = await self._apply_batch(batch) + applied_count += sum(results) + + logger.debug( + "Batch %d-%d: %d/%d applied.", + batch_start, + batch_start + len(batch), + sum(results), + len(batch), + ) + + logger.info("Summary overlay complete: %d/%d records applied.", applied_count, len(records)) + return applied_count + + async def _apply_batch(self, batch: List[Dict[str, Any]]) -> List[bool]: + """Process a batch of summary records concurrently. + + Args: + batch: A subset of summary records to process. + + Returns: + A list of booleans indicating success/failure for each record. + """ + tasks = [self._apply_single(record) for record in batch] + return await asyncio.gather(*tasks) + + async def _apply_single(self, record: Dict[str, Any]) -> bool: + """Apply a single summary record to a graph node. + + Args: + record: A summary record dict. + + Returns: + True if the update succeeded, False otherwise. + """ + try: + node_id = record["cpg_node_id"] + summary_type = record["summary_type"] + summary_text = record["summary_text"] + generation_model = record["generation_model"] + generation_prompt_version = record["generation_prompt_version"] + except KeyError as e: + logger.warning("Skipping record with missing key %s: %s", e, record) + return False + + property_name = f"summary_{summary_type}" + properties = { + property_name: summary_text, + "generation_model": generation_model, + "generation_prompt_version": generation_prompt_version, + } + + try: + success = await self._graph_store.update_node_properties(node_id, properties) + if not success: + logger.warning("Failed to update node %s (store returned False).", node_id) + return success + except Exception: + logger.exception("Error updating node %s.", node_id) + return False diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/tenant_ops.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/tenant_ops.py new file mode 100644 index 00000000..d8e2b332 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/tenant_ops.py @@ -0,0 +1,86 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tenant Operations — lifecycle management for CPG tenants in Neptune.""" + +import asyncio +import logging +from typing import Optional + +logger = logging.getLogger(__name__) + + +async def delete_tenant(tenant_id: str, graph_store) -> int: + """Purge ALL nodes and edges for a tenant from Neptune. + + WARNING: This deletes everything for the tenant across all domains. + For domain-scoped deletion (preserving non-CPG data), use delete_domain(). + + Args: + tenant_id: The tenant scope to delete + graph_store: Neptune graph store with execute_query method + + Returns: + Number of nodes deleted (approximate — Neptune doesn't return count) + """ + query = "MATCH (n {tenant_id: $tenant_id}) DETACH DELETE n" + try: + await asyncio.to_thread(graph_store.execute_query, query, {"tenant_id": tenant_id}) + logger.info(f"Tenant purged (all domains): {tenant_id}") + return -1 + except Exception as e: + logger.warning(f"Tenant purge failed for {tenant_id}: {e}") + raise + + +async def delete_domain(tenant_id: str, domain: str, graph_store) -> int: + """Delete only nodes with a specific domain label for a tenant. + + This preserves nodes from other domains (QA, Perf, Security, Lexical) + while replacing the specified domain (e.g., CPG). + + The client's requirement: "When a code change triggers a CPG reload, + the update must only delete/update CPG-labeled nodes and not touch + other app-level data." + + Args: + tenant_id: The tenant (app) scope + domain: The domain label to delete (e.g., "cpg", "lexical", "qa") + graph_store: Neptune graph store with execute_query method + + Returns: + Number of nodes deleted (approximate) + """ + query = "MATCH (n {tenant_id: $tenant_id, domain: $domain}) DETACH DELETE n" + try: + await asyncio.to_thread( + graph_store.execute_query, query, {"tenant_id": tenant_id, "domain": domain} + ) + logger.info(f"Domain purged: tenant={tenant_id}, domain={domain}") + return -1 + except Exception as e: + logger.warning(f"Domain purge failed for {tenant_id}/{domain}: {e}") + raise + + +async def list_domains(tenant_id: str, graph_store) -> list[str]: + """List all distinct domain labels for a tenant. + + Useful for inspecting what data exists before a scoped delete. + + Args: + tenant_id: The tenant scope + graph_store: Neptune graph store with execute_query method + + Returns: + List of domain strings (e.g., ["cpg", "lexical", "qa"]) + """ + query = "MATCH (n {tenant_id: $tenant_id}) RETURN DISTINCT n.domain AS domain" + try: + result = await asyncio.to_thread( + graph_store.execute_query, query, {"tenant_id": tenant_id} + ) + return [r["domain"] for r in result if r.get("domain")] + except Exception as e: + logger.warning(f"List domains failed for {tenant_id}: {e}") + return [] diff --git a/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/vector_loader.py b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/vector_loader.py new file mode 100644 index 00000000..fedaac04 --- /dev/null +++ b/codeproperty-graph/src/graphrag_toolkit/codeproperty_graph/vector_loader.py @@ -0,0 +1,307 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Vector Loader — loads CPG embedding vectors into OpenSearch via a VectorStoreProtocol. + +Reads VectorRecord dicts (from vectors.jsonl or in-memory list) and transforms +them into the format expected by the vector store, grouping by embedding_target +and generating compound document IDs for multi-vector-per-node support. + +Document ID scheme: f"{cpg_node_id}::{embedding_target}" + Allows a single CPG node to have multiple embeddings (e.g. code, docstring, signature). +""" + +import json +import logging +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Protocol, runtime_checkable + +logger = logging.getLogger(__name__) + + +@runtime_checkable +class VectorStoreProtocol(Protocol): + """Protocol for vector store backends (OpenSearch, etc.). + + Each document dict passed to put_vectors has: + id: str — compound key (cpg_node_id::embedding_target) + vector: list[float] + text: str — the source text that was embedded + metadata: dict — arbitrary metadata (node_type, filename, etc.) + """ + + async def put_vectors(self, documents: list[dict]) -> int: + """Bulk-index vector documents. Returns count of documents indexed.""" + ... + + +@dataclass +class VectorRecord: + """A single embedding vector produced by the CPG embedding stage. + + Fields align with vectors.jsonl output format: + cpg_node_id: The source CPG node ID this vector represents + embedding_target: What was embedded (e.g. 'code', 'docstring', 'signature') + vector: The embedding vector (list of floats) + text: The source text that was embedded + metadata: Additional context (node_type, filename, line_number, etc.) + embedding_model: Model identifier that produced this vector + embedding_dimensions: Declared dimensionality of the vector + """ + + cpg_node_id: str + embedding_target: str + vector: list[float] + text: str + metadata: dict[str, Any] = field(default_factory=dict) + embedding_model: str = "" + embedding_dimensions: int = 0 + + @property + def document_id(self) -> str: + """Compound key for OpenSearch: allows multiple vectors per CPG node.""" + return f"{self.cpg_node_id}::{self.embedding_target}" + + @classmethod + def from_dict(cls, raw: dict) -> "VectorRecord": + """Create a VectorRecord from a raw dict (e.g. parsed JSONL line).""" + return cls( + cpg_node_id=raw["cpg_node_id"], + embedding_target=raw["embedding_target"], + vector=raw["vector"], + text=raw.get("text", ""), + metadata=raw.get("metadata", {}), + embedding_model=raw.get("embedding_model", ""), + embedding_dimensions=raw.get("embedding_dimensions", 0), + ) + + +@dataclass +class LoadResult: + """Result of a vector load operation.""" + + total_loaded: int = 0 + by_target: dict[str, int] = field(default_factory=dict) + errors: list[str] = field(default_factory=list) + + +class VectorLoader: + """Loads CPG embedding vectors into a vector store (OpenSearch). + + Usage: + loader = VectorLoader( + vector_store=opensearch_store, + expected_dimensions=1024, + batch_size=500, + ) + result = await loader.load(records) + # or + result = await loader.load_from_file("s3://bucket/vectors.jsonl") + + Args: + vector_store: Any object implementing VectorStoreProtocol + expected_dimensions: Expected embedding dimensions (for validation) + batch_size: Number of documents per bulk indexing call + """ + + def __init__( + self, + vector_store: VectorStoreProtocol, + expected_dimensions: int = 0, + batch_size: int = 500, + graph_store=None, + tenant_id: str = "", + ): + self._store = vector_store + self._expected_dimensions = expected_dimensions + self._batch_size = batch_size + self._graph_store = graph_store + self._tenant_id = tenant_id + + async def load(self, records: list[dict | VectorRecord]) -> LoadResult: + """Load vector records into the vector store. + + Args: + records: List of VectorRecord instances or raw dicts (vectors.jsonl format). + + Returns: + LoadResult with counts and any validation errors. + """ + result = LoadResult() + grouped: dict[str, list[dict]] = defaultdict(list) + + for raw in records: + record = raw if isinstance(raw, VectorRecord) else VectorRecord.from_dict(raw) + + # Validate dimensions + error = self._validate(record) + if error: + result.errors.append(error) + continue + + # Transform to vector store document format + doc = { + "id": record.document_id, + "vector": record.vector, + "text": record.text, + "metadata": { + **record.metadata, + "cpg_node_id": record.cpg_node_id, + "embedding_target": record.embedding_target, + "embedding_model": record.embedding_model, + "tenant_id": self._tenant_id, + }, + } + grouped[record.embedding_target].append(doc) + + # Bulk-index by embedding_target group + for target, documents in grouped.items(): + count = await self._bulk_index(documents) + result.by_target[target] = count + result.total_loaded += count + logger.info(f"Loaded {count} vectors for target '{target}'") + + # Also write embedding as a node property on Neptune (client schema conformance) + if self._graph_store and records: + written = 0 + for raw in records: + record = raw if isinstance(raw, VectorRecord) else VectorRecord.from_dict(raw) + try: + await self._graph_store.update_node_properties( + record.cpg_node_id, + {"embedding": record.vector, "embeddedTime": record.metadata.get("embedded_time", "")}, + ) + written += 1 + except Exception as e: + logger.debug(f"Failed to write embedding property for {record.cpg_node_id}: {e}") + if written: + logger.info(f"Wrote embedding property to {written} Neptune nodes") + + if result.errors: + logger.warning(f"Encountered {len(result.errors)} validation errors during load") + + return result + + async def load_from_file(self, path: str) -> LoadResult: + """Load vectors from a JSONL file (local path or S3 URI). + + Args: + path: Local filesystem path or S3 URI (s3://bucket/key) to vectors.jsonl. + + Returns: + LoadResult with counts and any validation errors. + """ + records = await self._read_jsonl(path) + logger.info(f"Read {len(records)} vector records from {path}") + return await self.load(records) + + def _validate(self, record: VectorRecord) -> str | None: + """Validate a vector record. Returns error string or None if valid.""" + actual_dims = len(record.vector) + + # Check declared dimensions match actual vector length + if record.embedding_dimensions and actual_dims != record.embedding_dimensions: + return ( + f"Vector {record.document_id}: declared dimensions " + f"({record.embedding_dimensions}) != actual ({actual_dims})" + ) + + # Check against expected model dimensions if configured + if self._expected_dimensions and actual_dims != self._expected_dimensions: + return ( + f"Vector {record.document_id}: actual dimensions " + f"({actual_dims}) != expected ({self._expected_dimensions})" + ) + + if not record.vector: + return f"Vector {record.document_id}: empty vector" + + return None + + async def _bulk_index(self, documents: list[dict]) -> int: + """Index documents in batches via the vector store protocol.""" + total_indexed = 0 + + for i in range(0, len(documents), self._batch_size): + batch = documents[i : i + self._batch_size] + try: + count = await self._store.put_vectors(batch) + total_indexed += count + except Exception as e: + logger.error(f"Bulk index error on batch {i // self._batch_size}: {e}") + # Continue with remaining batches + continue + + return total_indexed + + async def _read_jsonl(self, path: str) -> list[dict]: + """Read vector records from a JSONL file (local or S3). + + Args: + path: Local path or s3://bucket/key URI. + + Returns: + List of raw dicts parsed from JSONL lines. + """ + if path.startswith("s3://"): + return await self._read_s3_jsonl(path) + return self._read_local_jsonl(path) + + @staticmethod + def _read_local_jsonl(path: str) -> list[dict]: + """Read JSONL from local filesystem.""" + records: list[dict] = [] + file_path = Path(path) + + if not file_path.exists(): + raise FileNotFoundError(f"Vector file not found: {path}") + + with file_path.open("r", encoding="utf-8") as f: + for line_num, line in enumerate(f, 1): + line = line.strip() + if not line: + continue + try: + records.append(json.loads(line)) + except json.JSONDecodeError as e: + logger.warning(f"Skipping malformed line {line_num} in {path}: {e}") + + return records + + @staticmethod + async def _read_s3_jsonl(uri: str) -> list[dict]: + """Read JSONL from S3 using aiobotocore/boto3. + + Args: + uri: S3 URI in format s3://bucket/key + + Returns: + List of raw dicts parsed from JSONL lines. + """ + import aiobotocore.session # type: ignore[import-untyped] + + # Parse s3://bucket/key + parts = uri.replace("s3://", "").split("/", 1) + bucket = parts[0] + key = parts[1] if len(parts) > 1 else "" + + session = aiobotocore.session.get_session() + records: list[dict] = [] + + async with session.create_client("s3") as client: + response = await client.get_object(Bucket=bucket, Key=key) + async with response["Body"] as stream: + content = await stream.read() + + for line_num, line in enumerate(content.decode("utf-8").splitlines(), 1): + line = line.strip() + if not line: + continue + try: + records.append(json.loads(line)) + except json.JSONDecodeError as e: + logger.warning(f"Skipping malformed line {line_num} in {uri}: {e}") + + return records diff --git a/codeproperty-graph/tests/__init__.py b/codeproperty-graph/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/codeproperty-graph/tests/test_codeproperty_graph.py b/codeproperty-graph/tests/test_codeproperty_graph.py new file mode 100644 index 00000000..bc2af8fb --- /dev/null +++ b/codeproperty-graph/tests/test_codeproperty_graph.py @@ -0,0 +1,165 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for codeproperty-graph core components.""" + +from graphrag_toolkit.codeproperty_graph.graph_diff import GraphDiff, DiffResult +from graphrag_toolkit.codeproperty_graph.manifest_manager import ManifestManager +from graphrag_toolkit.codeproperty_graph.models import CPGNode, CPGEdge, Manifest + + +def test_graph_diff_no_changes(): + prev = {"pkg.Foo.bar": "hash1", "pkg.Foo.baz": "hash2"} + curr = {"pkg.Foo.bar": "hash1", "pkg.Foo.baz": "hash2"} + result = GraphDiff.compare(prev, curr) + assert not result.has_changes + assert result.unchanged == 2 + assert result.summary == "+0 -0 ~0 =2" + + +def test_graph_diff_added(): + prev = {"pkg.Foo.bar": "h1"} + curr = {"pkg.Foo.bar": "h1", "pkg.Foo.new_method": "h2"} + result = GraphDiff.compare(prev, curr) + assert result.has_changes + assert "pkg.Foo.new_method" in result.added + + +def test_graph_diff_removed(): + prev = {"pkg.Foo.bar": "h1", "pkg.Foo.old": "h2"} + curr = {"pkg.Foo.bar": "h1"} + result = GraphDiff.compare(prev, curr) + assert "pkg.Foo.old" in result.removed + + +def test_graph_diff_modified(): + prev = {"pkg.Foo.bar": "hash_v1"} + curr = {"pkg.Foo.bar": "hash_v2"} + result = GraphDiff.compare(prev, curr) + assert "pkg.Foo.bar" in result.modified + assert result.modified["pkg.Foo.bar"] == "hash_v2" + + +def test_graph_diff_mixed(): + prev = {"a": "1", "b": "2", "c": "3"} + curr = {"a": "1", "b": "CHANGED", "d": "4"} + result = GraphDiff.compare(prev, curr) + assert result.added == {"d": "4"} + assert result.removed == {"c": "3"} + assert result.modified == {"b": "CHANGED"} + assert result.unchanged == 1 + assert result.summary == "+1 -1 ~1 =1" + + +def test_cpg_node_stable_id(): + node = CPGNode(id="123", node_type="METHOD", full_name="pkg.MyClass.run") + assert node.stable_id == "pkg.MyClass.run" + + +def test_cpg_node_stable_id_fallback(): + node = CPGNode(id="456", node_type="LITERAL") + assert node.stable_id == "456" + + +def test_cpg_edge_key(): + edge = CPGEdge(source_id="1", target_id="2", edge_type="AST") + assert edge.key == "1->AST->2" + + +def test_manifest_signature(): + mgr = ManifestManager.__new__(ManifestManager) + mgr._bucket = "test" + mgr._prefix = "cpg-exports" + mgr._s3 = None + sig = mgr.compute_signature({"a": "1", "b": "2"}) + assert sig.startswith("sha256:") + # Deterministic + assert sig == mgr.compute_signature({"b": "2", "a": "1"}) + + +# === Schema tests === + +def test_node_type_enum(): + from graphrag_toolkit.codeproperty_graph.schema import NodeType + assert NodeType.METHOD == "METHOD" + assert NodeType.CALL == "CALL" + assert NodeType.TYPE_DECL == "TYPE_DECL" + assert len(NodeType) == 20 + + +def test_edge_type_enum(): + from graphrag_toolkit.codeproperty_graph.schema import EdgeType + assert EdgeType.AST == "AST" + assert EdgeType.CFG == "CFG" + assert EdgeType.REACHING_DEF == "REACHING_DEF" + assert len(EdgeType) >= 14 + + +def test_delta_relevant_types(): + from graphrag_toolkit.codeproperty_graph.schema import DELTA_RELEVANT_TYPES, NodeType + assert NodeType.METHOD in DELTA_RELEVANT_TYPES + assert NodeType.COMMENT not in DELTA_RELEVANT_TYPES + + +def test_joern_export_command(): + from graphrag_toolkit.codeproperty_graph.schema import joern_export_command + cmd = joern_export_command("/src/myapp", output_dir="out", language="java") + assert "joern-export" in cmd + assert "--format json" in cmd + assert "--language java" in cmd + assert "/src/myapp" in cmd + + +# === from_joern factory tests === + +def test_cpg_node_from_joern(): + raw = { + "id": "1001", + "label": "METHOD", + "properties": { + "FULL_NAME": "com.example.Main.main", + "SIGNATURE": "void(String[])", + "HASH": "abc123", + "FILENAME": "src/Main.java", + "LINE_NUMBER": 5, + "CODE": "public static void main(String[] args) {}", + "IS_EXTERNAL": False, + } + } + node = CPGNode.from_joern(raw) + assert node.id == "1001" + assert node.node_type == "METHOD" + assert node.full_name == "com.example.Main.main" + assert node.signature == "void(String[])" + assert node.hash == "abc123" + assert node.filename == "src/Main.java" + assert node.line_number == 5 + assert node.is_external is False + + +def test_cpg_node_from_joern_identifier(): + raw = { + "id": "2001", + "label": "IDENTIFIER", + "properties": {"NAME": "myVar", "CODE": "myVar", "LINE_NUMBER": 10} + } + node = CPGNode.from_joern(raw) + assert node.node_type == "IDENTIFIER" + assert node.name == "myVar" + assert node.line_number == 10 + + +def test_cpg_edge_from_joern(): + raw = {"src": "1001", "dst": "2001", "label": "AST", "properties": {"ORDER": 1}} + edge = CPGEdge.from_joern(raw) + assert edge.source_id == "1001" + assert edge.target_id == "2001" + assert edge.edge_type == "AST" + assert edge.properties == {"ORDER": 1} + + +def test_cpg_edge_from_joern_cfg(): + raw = {"src": "100", "dst": "200", "label": "CFG", "properties": {}} + edge = CPGEdge.from_joern(raw) + assert edge.edge_type == "CFG" + assert edge.key == "100->CFG->200" diff --git a/docs-site/astro.config.mjs b/docs-site/astro.config.mjs index d21f0bd6..0e8ce4a1 100644 --- a/docs-site/astro.config.mjs +++ b/docs-site/astro.config.mjs @@ -97,6 +97,24 @@ export default defineConfig({ { label: 'FAQ', slug: 'byokg-rag/faq' }, ], }, + { + label: 'Document Graph', + items: [ + { label: 'Overview', slug: 'document-graph/overview' }, + { label: 'Pipeline', slug: 'document-graph/pipeline' }, + { label: 'Schema Providers', slug: 'document-graph/schema-providers' }, + { label: 'Hybrid Graph', slug: 'document-graph/hybrid-graph' }, + { label: 'Configuration', slug: 'document-graph/configuration' }, + ], + }, + { + label: 'Code Property Graph', + items: [ + { label: 'Overview', slug: 'codeproperty-graph/overview' }, + { label: 'Delta Ingestion', slug: 'codeproperty-graph/delta-ingestion' }, + { label: 'Configuration', slug: 'codeproperty-graph/configuration' }, + ], + }, ], }), ], diff --git a/docs-site/src/content/docs/codeproperty-graph/configuration.mdx b/docs-site/src/content/docs/codeproperty-graph/configuration.mdx new file mode 100644 index 00000000..0a9e0826 --- /dev/null +++ b/docs-site/src/content/docs/codeproperty-graph/configuration.mdx @@ -0,0 +1,140 @@ +--- +title: Configuration +--- + +## Installation + +```bash +pip install codeproperty-graph +``` + +## Dependencies + +codeproperty-graph depends on document-graph for typed graph primitives and optionally on lexical-graph for the graph store: + +| Package | Purpose | +|---------|---------| +| `document-graph` | CPGNode/CPGEdge → Node/Edge mapping, batch Cypher | +| `graphrag-toolkit-lexical-graph` (optional) | GraphStore, Neptune connection | +| `boto3` | S3 manifest storage | + +## S3 Configuration + +The ManifestManager stores and retrieves manifests from S3: + +```python +from codeproperty_graph import DeltaIngestor + +# Manifests stored at: s3://{bucket}/cpg-manifests/{repo}/manifest.json +ingestor = DeltaIngestor(bucket="my-artifacts-bucket") +``` + +Required IAM permissions: +- `s3:GetObject` on `cpg-manifests/*` +- `s3:PutObject` on `cpg-manifests/*` + +## Neptune Configuration + +Uses the same GraphStore as document-graph and lexical-graph: + +```python +from graphrag_toolkit.lexical_graph import GraphStoreFactory + +graph_store = GraphStoreFactory.for_graph_store( + "neptune-analytics://my-graph-id" +) +``` + +## Tenant ID Format + +Tenant IDs must match lexical-graph's TenantId validation: +- 1–25 characters +- Lowercase letters, numbers, periods (not at start/end) + +```python +# Good: +tenant_id = "abc12345678901234567890" +tenant_id = "build.123" + +# Bad: +tenant_id = "ABC" # uppercase +tenant_id = ".leading-dot" # starts with period +tenant_id = "way-too-long-tenant-id-string-exceeds-limit" # >25 chars +``` + +## Joern Export Format + +codeproperty-graph expects Joern's JSON export format: + +### Generating the Export + +```bash +# Install Joern: https://joern.io/docs/installation/ +# Supports: Java, JavaScript/TypeScript, Python, C/C++, Go, PHP, Ruby, Kotlin, Swift + +# Export CPG as JSON (auto-detects language) +joern-export --repr cpg14 --format json --out cpg-export/ /path/to/source + +# Or specify language explicitly +joern-export --repr cpg14 --format json --out cpg-export/ --language java /path/to/source + +# Output: +# cpg-export/nodes.json (all 20 node types) +# cpg-export/edges.json (all edge types: AST, CFG, CDG, CALL, etc.) +``` + +### Programmatic Export Command + +```python +from graphrag_toolkit.codeproperty_graph import joern_export_command + +cmd = joern_export_command("/src/my-service", output_dir="cpg-out", language="java") +print(cmd) +# joern-export --repr cpg14 --format json --out cpg-out --language java /src/my-service +``` + +### Node Format (nodes.json) + +```json +[ + { + "id": "1234", + "label": "METHOD", + "properties": { + "FULL_NAME": "com.example.Main.main", + "SIGNATURE": "void(String[])", + "HASH": "abc123def456", + "FILENAME": "src/Main.java", + "LINE_NUMBER": 10, + "LINE_NUMBER_END": 25, + "CODE": "public static void main(String[] args) {...}", + "IS_EXTERNAL": false + } + } +] +``` + +### Edge Format (edges.json) + +```json +[ + {"src": "1234", "dst": "5678", "label": "AST", "properties": {"ORDER": 1}}, + {"src": "1234", "dst": "9012", "label": "CFG", "properties": {}}, + {"src": "1234", "dst": "3456", "label": "CALL", "properties": {}} +] +``` + +### Parsing with from_joern() + +```python +import json +from graphrag_toolkit.codeproperty_graph import CPGNode, CPGEdge + +with open("cpg-export/nodes.json") as f: + nodes = [CPGNode.from_joern(raw) for raw in json.load(f)] + +with open("cpg-export/edges.json") as f: + edges = [CPGEdge.from_joern(raw) for raw in json.load(f)] + +print(f"Loaded {len(nodes)} nodes, {len(edges)} edges") +``` diff --git a/docs-site/src/content/docs/codeproperty-graph/delta-ingestion.mdx b/docs-site/src/content/docs/codeproperty-graph/delta-ingestion.mdx new file mode 100644 index 00000000..a556bb50 --- /dev/null +++ b/docs-site/src/content/docs/codeproperty-graph/delta-ingestion.mdx @@ -0,0 +1,88 @@ +--- +title: Delta Ingestion +--- + +The delta ingestion pipeline avoids unnecessary Neptune writes by comparing method signatures between exports. + +## How It Works + +``` +Previous Manifest (S3) Current Export (Joern) +───────────────────── ───────────────────── +full_name → hash full_name → hash +main() → abc123 main() → abc123 (unchanged) +parse() → def456 parse() → xyz789 (MODIFIED) +validate()→ ghi012 - (REMOVED) +- check() → jkl345 (ADDED) +``` + +**GraphDiff** compares these two dictionaries and produces: + +```python +DiffResult( + added={"check()": "jkl345"}, + removed={"validate()": "ghi012"}, + modified={"parse()": "xyz789"}, + unchanged=1 # main() +) +``` + +## Decision Logic + +```python +if not diff.has_changes: + return {"status": "SKIPPED"} # Zero Neptune cost +else: + # Full graph replacement under new tenant + new_tenant = generate_tenant_id() + await write_graph(nodes, edges, tenant_id=new_tenant) + await delete_tenant(old_tenant, graph_store) + await update_manifest(repo, new_signatures, new_tenant) + return {"status": "INGESTED", "delta": diff.summary} +``` + +## Why Full Replacement? + +CPG edges depend on node positions in the AST. A single method body change can cascade through AST, CFG, and CDG edges. Rather than computing fine-grained edge diffs (expensive and error-prone), codeproperty-graph replaces the entire graph when any change is detected and relies on the skip optimization for unchanged code. + +This is safe because: +- Tenant isolation means the new graph doesn't conflict with the old +- Old tenant is purged atomically after new graph is confirmed +- Manifest ensures the next run compares against the latest state + +## Manifest Structure + +```json +{ + "repo": "my-service", + "signature": "sha256:...", + "job_id": "build-123", + "tenant_id": "abc123", + "exported_at": "2026-06-30T07:00:00Z", + "nodes_path": "s3://bucket/cpg-exports/my-service/build-123/nodes.json", + "edges_path": "s3://bucket/cpg-exports/my-service/build-123/edges.json", + "method_signatures": { + "com.example.Main.main": "abc123def456", + "com.example.Parser.parse": "789xyz012abc" + } +} +``` + +## CI/CD Integration + +In a typical pipeline: + +```yaml +# GitHub Actions / CodePipeline step: +- name: CPG Delta Ingest + run: | + joern-export --out nodes.json edges.json + python -c " + from codeproperty_graph import DeltaIngestor + ingestor = DeltaIngestor(bucket='$ARTIFACTS_BUCKET') + result = asyncio.run(ingestor.ingest(...)) + print(result['status']) + " +``` + +On most commits (no code change to method bodies), this completes in <1 second with zero Neptune writes. diff --git a/docs-site/src/content/docs/codeproperty-graph/overview.mdx b/docs-site/src/content/docs/codeproperty-graph/overview.mdx new file mode 100644 index 00000000..1f548507 --- /dev/null +++ b/docs-site/src/content/docs/codeproperty-graph/overview.mdx @@ -0,0 +1,93 @@ +--- +title: Overview +--- + +The graphrag-toolkit [codeproperty-graph](https://github.com/awslabs/graphrag-toolkit/blob/main/codeproperty-graph/) library provides code-analysis-specific graph ingestion with delta detection. It sits on top of document-graph and lexical-graph, adding Joern-based code property graph (CPG) semantics, manifest-driven change detection, and incremental ingestion. + +- [Architecture](#architecture) +- [Core concepts](#core-concepts) +- [Delta logic](#delta-logic) +- [Getting started](#getting-started) + +### Architecture + +codeproperty-graph is the top layer in the graphrag-toolkit stack for code analysis: + +``` +┌─────────────────────────────────────────────────────┐ +│ codeproperty-graph (domain) │ +│ CPGNode, CPGEdge, GraphDiff, DeltaIngestor │ +│ ManifestManager, tenant lifecycle │ +├─────────────────────────────────────────────────────┤ +│ document-graph (infra) │ +│ Node, Edge, batch_nodes_unwind, batch_edges_unwind │ +│ CypherBuilder, PipelineExecutor, multi-tenancy │ +├─────────────────────────────────────────────────────┤ +│ graphrag-toolkit / lexical-graph (foundation) │ +│ GraphStore, Neptune writer, AOSS writer │ +│ Lexical indexing, entity resolution, retrieval │ +└─────────────────────────────────────────────────────┘ +``` + +### Core concepts + +**CPG Schema** — Complete Joern type system: 20 node types, 14+ edge types, multi-language support. See `schema.py` for enums and property specs. + +``` +Node Types (20): METHOD, CALL, IDENTIFIER, LITERAL, BLOCK, COMMENT, + FILE, LOCAL, MEMBER, META_DATA, METHOD_REF, + METHOD_RETURN, MODIFIER, NAMESPACE, NAMESPACE_BLOCK, + PARAMETER, RETURN, TAG, TYPE_DECL, TYPE + +Edge Types (14+): AST, CFG, CDG, REACHING_DEF, CALL, ARGUMENT, + DOMINATE, POST_DOMINATE, CONTAINS, BINDS_TO, + REF, INHERITS_FROM, CONDITION, TAGGED_BY + +Languages: Java, JavaScript/TypeScript, Python, C/C++, + Go, PHP, Ruby, Kotlin, Swift (via Joern frontends) +``` + +**CPG Models** — Typed representations with `from_joern()` factory methods that parse Joern's JSON export format directly. All 20 node types flow through — the model is schema-flexible. + +**GraphDiff** — Compare two CPG states by method signature (`full_name → hash`). Returns adds, removes, modifications, and unchanged counts. + +**ManifestManager** — S3-backed state tracking per repository. Stores the method signature fingerprint so subsequent exports can be compared. + +**DeltaIngestor** — Orchestrates the skip-or-replace decision: if the CPG hasn't changed, skip Neptune writes entirely. If changed, ingest the full graph under a new tenant and purge the old one. + +**Tenant lifecycle** — Clean tenant management with `delete_tenant()` for atomic graph replacement. + +### Delta logic + +``` +1. Joern exports CPG → nodes.json + edges.json +2. Extract METHOD node signatures: {full_name: hash} +3. Compare against previous manifest in S3 +4. If identical → SKIP (no Neptune writes, no cost) +5. If changed → INGEST full graph under new tenant, + purge old tenant, update manifest +``` + +This ensures Neptune only receives writes when code actually changes — critical for CI/CD pipelines that run on every commit. + +### Getting started + +```bash +pip install codeproperty-graph +``` + +```python +from codeproperty_graph import DeltaIngestor + +ingestor = DeltaIngestor(bucket="my-artifacts-bucket") +result = await ingestor.ingest( + repo="my-service", + job_id="build-123", + tenant_id="abc123", + nodes_data=nodes, + edges_data=edges, + graph_store=neptune, + write_fn=my_write_function, +) +print(result["status"]) # "SKIPPED" or "INGESTED" +``` diff --git a/docs-site/src/content/docs/document-graph/configuration.mdx b/docs-site/src/content/docs/document-graph/configuration.mdx new file mode 100644 index 00000000..31ca2620 --- /dev/null +++ b/docs-site/src/content/docs/document-graph/configuration.mdx @@ -0,0 +1,69 @@ +--- +title: Configuration +--- + +## Installation + +```bash +# Standalone (pure ETL, no lexical-graph dependency): +pip install graphrag-toolkit-document-graph + +# With Neptune support: +pip install graphrag-toolkit-document-graph[neptune] + +# With lexical-graph for hybrid graphs: +pip install graphrag-toolkit-document-graph[graphrag] + +# Full install: +pip install graphrag-toolkit-document-graph[graphrag,neptune,neo4j] +``` + +## Dependencies + +| Extra | Packages | Purpose | +|-------|----------|---------| +| (base) | pydantic, networkx, boto3, pandas, pluggy | Core pipeline | +| `[graphrag]` | graphrag-toolkit-lexical-graph | Hybrid graph, shared TenantId | +| `[neptune]` | boto3 | Neptune graph store | +| `[neo4j]` | neo4j | Neo4j graph store | +| `[dev]` | pytest, pytest-cov, pytest-mock | Development | + +## Graph Store Connection + +document-graph uses the same `GraphStoreFactory` as lexical-graph: + +```python +from graphrag_toolkit.lexical_graph import GraphStoreFactory + +# Neptune Analytics: +graph_store = GraphStoreFactory.for_graph_store( + "neptune-analytics://my-graph-id" +) + +# Neptune Database: +graph_store = GraphStoreFactory.for_graph_store( + "neptune-db://my-cluster.region.neptune.amazonaws.com" +) +``` + +## Multi-Tenancy Configuration + +```python +from graphrag_toolkit.document_graph import PipelineExecutor + +# Each tenant's data is label-scoped in the same Neptune cluster: +executor = PipelineExecutor( + source_path="data/tenant_a.csv", + schema_path="schemas/accounts.json", + graph_store=graph_store, + tenant_id="tenant_a" # Labels become: `Accounttenant_a__` +) +``` + +## Environment Variables + +| Variable | Purpose | Default | +|----------|---------|---------| +| `AWS_REGION` | AWS region for Neptune/Bedrock | From AWS config | +| `NEPTUNE_ENDPOINT` | Neptune cluster endpoint | None | +| `NEPTUNE_GRAPH_ID` | Neptune Analytics graph ID | None | diff --git a/docs-site/src/content/docs/document-graph/hybrid-graph.mdx b/docs-site/src/content/docs/document-graph/hybrid-graph.mdx new file mode 100644 index 00000000..d91488dc --- /dev/null +++ b/docs-site/src/content/docs/document-graph/hybrid-graph.mdx @@ -0,0 +1,68 @@ +--- +title: Hybrid Graph +--- + +document-graph and lexical-graph can share a single Neptune cluster, creating a **hybrid graph** that combines structured records with unstructured text entities in a unified queryable knowledge graph. + +## Architecture + +``` +┌──────────────────────────────────────────────────────────┐ +│ Amazon Neptune │ +├──────────────────────┬───────────────────────────────────┤ +│ lexical-graph │ document-graph │ +│ (unstructured) │ (structured) │ +│ │ │ +│ Document ──→ Chunk │ CSV Row ──→ Person │ +│ Chunk ──→ Entity │ Person ──→ Organization │ +│ Entity ──→ Entity │ Organization ──→ Location │ +├──────────────────────┴───────────────────────────────────┤ +│ Shared edges (hybrid path) │ +│ Entity ←──→ Person (same real-world entity) │ +└──────────────────────────────────────────────────────────┘ +``` + +## Tenant Compatibility + +Both packages use the same tenant label format. document-graph imports lexical-graph's `TenantId.format_label()` when the `[graphrag]` extra is installed: + +```python +# Both produce identical label format: +# Default tenant: `Person` +# Custom tenant: `Personacme__` + +# This means nodes from both packages coexist cleanly +# and can be queried across in a single Cypher traversal. +``` + +## Setup + +```bash +# Install both packages: +pip install graphrag-toolkit-lexical-graph +pip install graphrag-toolkit-document-graph[graphrag] +``` + +## Building a Hybrid Graph + +1. **Index unstructured content** with lexical-graph (documents, PDFs) +2. **Ingest structured data** with document-graph (CSV, JSON, databases) +3. **Link entities** across both graphs using shared identifiers or entity resolution + +See the example notebooks: +- `07a-Lexical-Integration-Data-Processing.ipynb` +- `07b-Lexical-Integration-Lexical-Setup.ipynb` +- `08a-Nelson-Mandela-Hybrid-Build.ipynb` +- `08b-Nelson-Mandela-Hybrid-Query.ipynb` + +## Multi-Tenant Coexistence + +Multiple tenants can share a single Neptune cluster. Both document-graph and lexical-graph scope their labels identically, so tenant isolation works across both graph types: + +```python +# Tenant "research" sees only its own nodes from both packages: +# lexical-graph: `Chunkresearch__`, `Entityresearch__` +# document-graph: `Personresearch__`, `Organizationresearch__` +``` + +See `09-Multi-Tenant-Coexistence-Demo.ipynb` for a working example. diff --git a/docs-site/src/content/docs/document-graph/overview.mdx b/docs-site/src/content/docs/document-graph/overview.mdx new file mode 100644 index 00000000..b25b00d8 --- /dev/null +++ b/docs-site/src/content/docs/document-graph/overview.mdx @@ -0,0 +1,65 @@ +--- +title: Overview +--- + +The graphrag-toolkit [document-graph](https://github.com/awslabs/graphrag-toolkit/blob/main/document-graph/) library provides a structured data ingestion route into Neptune that complements lexical-graph. While lexical-graph handles unstructured text (documents, PDFs, web pages), document-graph handles structured and semi-structured data (CSV, JSON, Parquet, Excel, XML, YAML) — transforming tabular and hierarchical records into typed property graphs. + +- [Architecture](#architecture) +- [Relationship to lexical-graph](#relationship-to-lexical-graph) +- [Core concepts](#core-concepts) +- [Getting started](#getting-started) + +### Architecture + +document-graph is a second ingestion route into the same Neptune store, not a replacement for lexical-graph. It declares lexical-graph as an optional dependency, enabling hybrid graphs that combine structured records with unstructured text in a single queryable knowledge graph. + +``` +┌─────────────────┐ ┌─────────────────┐ +│ Unstructured │ │ Structured │ +│ (PDFs, docs) │ │ (CSV, JSON) │ +└────────┬────────┘ └────────┬────────┘ + │ │ + lexical-graph document-graph + │ │ + └───────────┬───────────┘ + │ + Amazon Neptune + (shared graph) +``` + +### Relationship to lexical-graph + +| Concern | lexical-graph | document-graph | +|---------|---------------|----------------| +| Input data | Unstructured text | Structured/semi-structured records | +| Schema | Inferred from extraction | Declared or discovered from source | +| Graph construction | LLM-extracted entities + relationships | Schema-driven node/edge mapping | +| Tenant model | `TenantId.format_label()` | Same (shared via optional import) | +| Neptune target | Same cluster | Same cluster | +| Install | `graphrag-toolkit-lexical-graph` | `graphrag-toolkit-document-graph` | + +### Core concepts + +**Pipeline:** Extract → Transform → Load → Build. Each stage is pluggable with provider registries. + +**Schema providers:** Discover or declare the schema of your source data (CSV headers, JSON keys, Parquet columns, Glue catalog). + +**Transformers:** Field-level (JSON flattening, UUID generation), document-level (chunking, PII redaction), and graph-level (row-to-node, edge inference). + +**Constructors:** Typed graph construction patterns — 1:N, N:M, schema-driven, batch-optimized with Neptune's native `~id` index. + +**Multi-tenancy:** Shared Neptune cluster with tenant-scoped labels, compatible with lexical-graph's tenant format. + +### Getting started + +```bash +pip install graphrag-toolkit-document-graph +``` + +For hybrid graphs with lexical-graph: + +```bash +pip install graphrag-toolkit-document-graph[graphrag] +``` + +See the [examples](https://github.com/awslabs/graphrag-toolkit/blob/main/examples/document-graph/) for notebooks demonstrating end-to-end pipelines. diff --git a/docs-site/src/content/docs/document-graph/pipeline.mdx b/docs-site/src/content/docs/document-graph/pipeline.mdx new file mode 100644 index 00000000..a4ad8962 --- /dev/null +++ b/docs-site/src/content/docs/document-graph/pipeline.mdx @@ -0,0 +1,78 @@ +--- +title: Pipeline +--- + +document-graph uses a four-stage pipeline to transform structured data into a typed property graph in Neptune. + +## Stages + +``` +Extract → Transform → Load → Build +``` + +### Extract + +Reads source data and produces a normalized DataFrame. Supports multiple formats via pluggable extract providers: + +| Provider | Input | Description | +|----------|-------|-------------| +| CSV | `.csv` files | Column-oriented tabular data | +| JSON | `.json` files | Nested or flat JSON documents | +| Parquet | `.parquet` files | Columnar binary format | +| Excel | `.xlsx` files | Spreadsheet data | +| CPG | Code property graph | Code analysis output (from codeproperty-graph) | + +### Transform + +Applies transformations to the extracted data. Transformers are composable and ordered: + +**Field transformers:** Operate on individual columns +- JSON flattening, array expansion, regex cleaning +- UUID generation, timestamp normalization +- Comma splitting, enum standardization + +**Document transformers:** Operate on entire records +- Text chunking, PII redaction, JSON-to-rows + +**Normalizers:** Data quality +- Whitespace, null handling, case normalization, spelling + +**Filter transformers:** Row/column selection +- Row filtering, column pruning + +**Graph transformers:** Structure mapping +- Row-to-node conversion, edge inference + +### Load + +Writes the transformed data to the graph store. Uses batched UNWIND queries with Neptune's native `~id` for O(1) vertex lookups. + +### Build + +Constructs the typed property graph using constructors: + +| Constructor | Pattern | Description | +|-------------|---------|-------------| +| Node | 1:1 | One row → one node | +| Edge | 1:1 | One relationship per matched pair | +| Schema-driven | N/A | Graph structure from declared schema | +| One-to-many | 1:N | One parent → N children | +| Many-to-many | N:M | Junction table → edges | +| Batch | Bulk | Optimized bulk write | +| Deduplication | Merge | Merge duplicate nodes | + +## Configuration + +Pipelines are configured via schema files (JSON or YAML) or programmatically: + +```python +from graphrag_toolkit.document_graph import PipelineExecutor + +executor = PipelineExecutor( + source_path="data/accounts.csv", + schema_path="schemas/accounts.json", + graph_store=graph_store, + tenant_id="my_tenant" +) +result = executor.run() +``` diff --git a/docs-site/src/content/docs/document-graph/schema-providers.mdx b/docs-site/src/content/docs/document-graph/schema-providers.mdx new file mode 100644 index 00000000..4db5db92 --- /dev/null +++ b/docs-site/src/content/docs/document-graph/schema-providers.mdx @@ -0,0 +1,93 @@ +--- +title: Schema Providers +--- + +Schema providers discover or load the schema definition that drives document-graph's pipeline. The schema defines what nodes and edges to create from your source data. + +## Provider Types + +| Provider | Source | Use Case | +|----------|--------|----------| +| CSV | `.csv` file headers | Quick start with tabular data | +| JSON | `.json` structure | Nested document ingestion | +| Parquet | `.parquet` column metadata | Data lake integration | +| Excel | `.xlsx` sheet structure | Spreadsheet-based workflows | +| YAML | `.yaml` schema definition | Declarative schema-first approach | +| XML | `.xml` document structure | Legacy data integration | +| Glue | AWS Glue Data Catalog | Enterprise data catalog integration | +| S3 | Schema file on S3 | Remote schema storage | +| File | Local file system | Development and testing | +| Static | Programmatic definition | Code-first schema declaration | + +## Schema Discovery + +For tabular sources (CSV, Parquet, Excel), document-graph can **discover** the schema automatically from the data: + +```python +from graphrag_toolkit.document_graph.schema.providers import SchemaProviderFactory + +provider = SchemaProviderFactory.create("csv", path="data/accounts.csv") +schema = provider.discover() +# Returns an ETLSchema with inferred node types, properties, and relationships +``` + +## Declarative Schema + +For production use, declare your schema explicitly: + +```json +{ + "extract": { + "source": "data/accounts.csv", + "format": "csv" + }, + "transform": { + "nodes": [ + { + "label": "Account", + "id_field": "account_id", + "properties": ["name", "type", "created_date"] + }, + { + "label": "Owner", + "id_field": "owner_id", + "properties": ["owner_name", "email"] + } + ], + "relationships": [ + { + "type": "OWNED_BY", + "source": "Account", + "target": "Owner", + "source_field": "owner_id", + "target_field": "owner_id" + } + ] + } +} +``` + +## Static Schema Provider + +For maximum control, define schemas programmatically: + +```python +from graphrag_toolkit.document_graph.schema.static_schema_provider import StaticSchemaProvider +from graphrag_toolkit.document_graph.schema.etl_schema_model import ( + ETLSchema, ExtractConfig, TransformConfig, NodeDefinition, RelationshipDefinition +) + +schema = ETLSchema( + extract=ExtractConfig(source="data/accounts.csv", format="csv"), + transform=TransformConfig( + nodes=[ + NodeDefinition(label="Account", id_field="account_id", properties=["name", "type"]), + ], + relationships=[ + RelationshipDefinition(type="OWNED_BY", source="Account", target="Owner"), + ] + ) +) + +provider = StaticSchemaProvider(schema) +``` diff --git a/docs-site/src/content/docs/index.mdx b/docs-site/src/content/docs/index.mdx index 7692a772..5aacbaab 100644 --- a/docs-site/src/content/docs/index.mdx +++ b/docs-site/src/content/docs/index.mdx @@ -28,7 +28,13 @@ 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. - + + Structured data ingestion — CSV, JSON, Parquet, Excel into typed property graphs. Schema-driven pipelines with multi-tenancy and batch-optimized Neptune writes. + + + Delta-aware code analysis ingestion from Joern. Full CPG schema (20 node types, 9 languages) with skip-on-no-change semantics for CI/CD. + + Graph stores: Amazon Neptune (DB and Analytics), Neo4j, FalkorDB. Vector stores: Neptune, OpenSearch, Postgres, S3 Vectors. diff --git a/document-graph/CHANGELOG.md b/document-graph/CHANGELOG.md new file mode 100644 index 00000000..a87881dd --- /dev/null +++ b/document-graph/CHANGELOG.md @@ -0,0 +1,58 @@ +# Changelog + +## 3.0.4 (2026-06-26) + +### Added +- **100% docstring coverage** across all 163 source files and 141 classes +- **59 tests passing** — schema providers, transformers, ETL model, cypher edge cases, CSV discovery +- **Edge relationships** — AUTHORED_BY, CITES edge generation in hybrid pipeline +- **Lineage preservation** — source metadata structure for graphrag-toolkit (`source.sourceId` + `source.metadata`) +- **Cleanup notebook** (00-Cleanup) — remove old test tenants from Neptune + +### Fixed +- `S3SchemaProvider` — removed dead `DocumentGraphConfig` references, uses `boto3.Session()` directly +- `StaticSchemaProvider` — added missing `from_config()` classmethod +- `PIIRedactorProvider` — lazy import for `sanitary` (optional dependency) +- MockStore tests — `query()` → `execute_query()` to match actual API + +## 3.0.3 (2026-06-25) + +### Added +- **Sphinx docs** scaffolding +- **Hybrid notebooks** — 07a (data processing), 07b (lexical-graph query + correlation) +- **Mandela demo** — 08a (build), 08b (query) +- **Multi-tenant demo** — 09 (isolation proven) +- **Schema notebook** — 005 (providers, integration, validation) +- **Transformer notebook** — 006 (all 7 categories) + +### Changed +- `graphrag-lexical-graph` moved to optional dependency (avoids pip resolver conflicts) +- `push-to-sagemaker.sh` — S3 cleanup before upload, only latest wheel +- Neptune upgraded to 1.4.7.0 (graphrag-toolkit nested UNWIND compatibility) +- Batch writes enabled (Neptune 1.4.x supports it) + +### Fixed +- All corrupt notebooks (ai4triage removal artifacts) — JSON repaired +- Notebook numbering consolidated: 00–09 sequential + +## 3.0.0 (2026-06-24) + +### Breaking Changes +- **Removed `root_id`** — use `tenant_id` only (aligned with graphrag-toolkit) +- **Removed vendored storage** — uses graphrag-toolkit as dependency +- **Removed `tenant_id.py`, `versioning.py`, `root_id.py`** — sourced from graphrag-toolkit + +### Added +- `storage/` — thin extension layer (ReadOnlyGraphStore, factory extensions, complementary drivers) +- `graph_build/` — cypher_builder (node_to_cypher, edge_to_cypher, batch) +- `query/` — DocumentGraphQueryEngine +- Schema providers: CSV, JSON, S3, Static, File, Glue, Parquet, Excel, XML, YAML +- Schema discovery: auto-infer ETLSchema from data files +- 20+ transformers across 6 categories + +## 2.0.0 (2026-06-24) + +### Breaking Changes +- Pydantic V2 migration (all validators, ConfigDict) +- Removed visualisation, resilient_client, graph_id +- graphrag-toolkit as dependency (not vendored) diff --git a/document-graph/LICENSE b/document-graph/LICENSE new file mode 100644 index 00000000..67db8588 --- /dev/null +++ b/document-graph/LICENSE @@ -0,0 +1,175 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. diff --git a/document-graph/README.md b/document-graph/README.md new file mode 100644 index 00000000..efabfa88 --- /dev/null +++ b/document-graph/README.md @@ -0,0 +1,205 @@ +# Document Graph + +> Structured data ETL for graph-enhanced GenAI applications. Extends [graphrag-toolkit](https://github.com/awslabs/graphrag-toolkit) with typed node support, schema providers, and hybrid search. + +## What it does + +| Input | Process | Output | +|-------|---------|--------| +| CSV, Excel, JSON, Parquet, XML | Deterministic ETL (no LLM) | Typed Neptune nodes + edges | + +**Hybrid with Lexical Graph:** Same Neptune cluster holds both structured (document-graph) and unstructured (lexical-graph) data. Query both with a single semantic search. + +## Features + +| Feature | Module | Description | +|---------|--------|-------------| +| **Schema Providers** | `schema.providers` | CSV, JSON, S3, Static, File, Glue — auto-discover or define schemas | +| **Schema Discovery** | `schema.discovery` | Infer ETL schema from data files (CSV, Excel, JSON, Parquet, XML, YAML) | +| **Transformers** | `transform` | Normalizers, field, document, filter, graph, truncators (20+ built-in) | +| **Graph Build** | `graph_build` | Cypher generation (node_to_cypher, edge_to_cypher) with tenant scoping | +| **Query Engine** | `query` | Typed node queries with tenant isolation | +| **Multi-Tenancy** | Labels | `__Type__tenant_id__` — complete data isolation per tenant | +| **Hybrid Search** | Integration | Lexical-graph indexes document-graph content for semantic retrieval | + +## Architecture + +``` +graphrag-toolkit (PyPI: graphrag-lexical-graph) +├── GraphStore, VectorStore, Factories +├── LexicalGraphIndex (unstructured → lexical graph) +└── LexicalGraphQueryEngine (semantic search) + +document-graph (this package) +├── schema/ → ETL schema model, providers (CSV/JSON/S3/Static), discovery +├── transform/ → 20+ transformers (normalizers, field, document, filter, graph, truncators) +├── graph_build/ → Cypher generation with tenant-scoped labels +├── query/ → DocumentGraphQueryEngine (typed node queries) +├── pipeline/ → Extract providers (CSV, Excel, JSON, Parquet) +├── ingest/ → Column selectors, row filters, renamers +└── storage/ → ReadOnlyGraphStore, factory extensions +``` + +## Infrastructure + +This project relies on [graphrag-toolkit](https://github.com/awslabs/graphrag-toolkit) for AWS infrastructure. Deploy the CloudFormation stack: + +> https://github.com/awslabs/graphrag-toolkit/tree/main/examples/lexical-graph/cloudformation-templates + +This provisions Neptune, OpenSearch Serverless, and a SageMaker notebook instance. All example notebooks (`examples/cloud/notebooks/`) must run on **SageMaker** within the provisioned VPC. + +## Install + +```bash +pip install document-graph # core (no heavy deps) +pip install document-graph[graphrag] # with lexical-graph integration +pip install graphrag-lexical-graph # for hybrid search +``` + +## Quick Start + +### 1. Write nodes to Neptune + +```python +from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory +from document_graph.graph_build import node_to_cypher +from document_graph import Node + +gs = GraphStoreFactory.for_graph_store('neptune-db://endpoint:8182').__enter__() + +node = Node(id='u1', labels=['User'], properties={'name': 'Alice', 'role': 'admin'}) +cypher, params = node_to_cypher(node, tenant_id='my_app') +gs.execute_query(cypher, params) +``` + +### 2. Query nodes + +```python +from document_graph.query import DocumentGraphQueryEngine + +engine = DocumentGraphQueryEngine(gs, tenant_id='my_app') +users = engine.get_nodes('User') +admins = engine.find_by_property('User', 'role', 'admin') +``` + +### 3. Schema-driven pipeline + +```python +from document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from document_graph.schema.providers.csv_schema_provider import CSVSchemaProvider + +config = SchemaProviderConfig(type='csv', connection_config={'path': 'data/users.csv'}) +provider = CSVSchemaProvider(config) +schema = provider.load_schema() # Returns ETLSchema with discovered fields +``` + +### 4. Transform data + +```python +from document_graph.transform.transformer_provider_config import TransformerProviderConfig +from document_graph.transform.normalizers.normalize_whitespace_provider import NormalizeWhitespaceProvider +from document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer + +# Normalize +records = NormalizeWhitespaceProvider(TransformerProviderConfig(name='ws', args={})).transform(records) + +# Convert to nodes +nodes = RowToNodeTransformer(TransformerProviderConfig(name='r2n', args={'type': 'User'})).transform(records) +``` + +### 5. Hybrid search (document-graph + lexical-graph) + +```python +from graphrag_toolkit.lexical_graph import LexicalGraphIndex, LexicalGraphQueryEngine +from graphrag_toolkit.lexical_graph.storage import VectorStoreFactory +from llama_index.core.schema import Document + +# Index document-graph data into lexical-graph +docs = [Document(text=f'[User | {r["id"]} | my_app]\nname: {r["name"]}', metadata={...}) for r in records] +vector_store = VectorStoreFactory.for_vector_store('aoss://endpoint') +graph_index = LexicalGraphIndex(graph_store, vector_store) +graph_index.extract_and_build(docs, show_progress=True) + +# Semantic query +query_engine = LexicalGraphQueryEngine.for_traversal_based_search(graph_store, vector_store) +results = query_engine.retrieve('Who are the admin users?') +``` + +## Lexical-Graph Integration + +Document-graph and lexical-graph coexist in the same Neptune database: + +| | Lexical Graph | Document Graph | +|--|---------------|----------------| +| **Input** | Unstructured text (PDF, web) | Structured data (CSV, Excel, JSON) | +| **Extraction** | LLM-based (Claude) | Deterministic ETL (pandas) | +| **Graph Model** | Source → Chunk → Topic → Statement → Entity | Row → Typed Node + Edges | +| **Query** | Traversal-based semantic search | Cypher over typed nodes | +| **Labels** | `__Source__`, `__Chunk__`, `__Topic__` | `__User__tenant__`, `__Account__tenant__` | +| **Coexistence** | Same Neptune, different labels | Complete isolation | + +### Lineage + +When indexing document-graph data into lexical-graph, embed lineage in the text: + +```python +text = f'[{node_type} | {node_id} | {tenant}]\n...' +``` + +This survives the lexical-graph chunking pipeline and enables correlation back to the original nodes. + +## Multi-Tenancy + +All operations use tenant-scoped labels: `__Type__tenant_id__` + +```python +# Write to tenant A +node_to_cypher(node, tenant_id='acme_corp') # → MERGE (n:`__User__acme_corp__` ...) + +# Query tenant A only +engine = DocumentGraphQueryEngine(gs, tenant_id='acme_corp') +engine.get_nodes('User') # Only sees acme_corp's users +``` + +## Schema Providers + +| Provider | Source | Usage | +|----------|--------|-------| +| `CSVSchemaProvider` | CSV file | Auto-discover schema | +| `JSONSchemaProvider` | JSON file | Load pre-defined schema | +| `S3SchemaProvider` | S3 bucket | Shared schemas across environments | +| `StaticSchemaProvider` | Code | Programmatic definition | +| `SchemaProviderFactory` | Config dict | Unified creation | + +## Transformers + +| Category | Examples | Purpose | +|----------|----------|---------| +| Normalizers | whitespace, nulls, case, enum, timestamp | Clean & standardize | +| Field | json_flattener, uuid_gen, regex_clean | Reshape fields | +| Document | json_to_rows, text_chunker, pii_redactor | Split/transform docs | +| Filter | row_filter, column_pruner | Remove unwanted data | +| Graph | row_to_node, infer_edges | Create graph structures | +| Truncators | length, field_count, token | Limit data size | + +All follow: `TransformerProvider(config).transform(records) → records` + +## Testing + +```bash +pip install pytest +pytest tests/ -q # 59 tests +``` + +## Requirements + +- Python >= 3.10 +- `pydantic >= 2.0` +- `pandas >= 2.0` +- `boto3 >= 1.26` +- Optional: `graphrag-lexical-graph >= 3.18.0` (for hybrid search) + +## License + +Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +SPDX-License-Identifier: Apache-2.0 diff --git a/document-graph/pyproject.toml b/document-graph/pyproject.toml new file mode 100644 index 00000000..d12e250d --- /dev/null +++ b/document-graph/pyproject.toml @@ -0,0 +1,45 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/graphrag_toolkit"] + +[project] +name = "graphrag-toolkit-document-graph" +version = "3.0.7" +description = "Graph-enhanced document processing for AWS GraphRAG Toolkit" +readme = "README.md" +requires-python = ">=3.10" +license = "Apache-2.0" +dependencies = [ + "pydantic>=2.0.0", + "networkx>=3.0", + "boto3>=1.26.0", + "pandas>=2.0.0", + "pluggy>=1.0.0", +] + +[project.optional-dependencies] +graphrag = [ + "graphrag-toolkit-lexical-graph>=3.18.0", +] +neptune = [ + "boto3>=1.26.0", +] +neo4j = [ + "neo4j>=5.0.0", +] +dev = [ + "pytest>=7.0", + "pytest-cov>=4.0", + "pytest-mock>=3.0", +] + +[project.urls] +Homepage = "https://github.com/aws/graphrag-toolkit" +Repository = "https://github.com/aws/graphrag-toolkit" + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src"] diff --git a/document-graph/src/graphrag_toolkit/document_graph/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/__init__.py new file mode 100644 index 00000000..2bd1b880 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/__init__.py @@ -0,0 +1,33 @@ +"""Document Graph — structured data → typed graph nodes. + +A transformation library that extends graphrag-toolkit with typed node support +for structured formats (CSV, Excel, JSON, Parquet). + +Architecture: + Toolkit Reader → [ingest → transform → build] → Toolkit GraphStore + +Usage: + from graphrag_toolkit.document_graph import NodeModel, EdgeModel + from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNode + from graphrag_toolkit.document_graph.graph_build.cypher_builder import CypherBuilder +""" + +__version__ = "3.0.3" + +from .model import NodeModel, EdgeModel +from .model_elements import Node, Edge +from .pipeline_executor import PipelineExecutor +from .errors import ( + ModelError, + DatabaseConnectionError, + ConfigurationError, + QueryExecutionError, +) + +__all__ = [ + "PipelineExecutor", + "NodeModel", "EdgeModel", + "Node", "Edge", + "ModelError", "DatabaseConnectionError", + "ConfigurationError", "QueryExecutionError", +] diff --git a/document-graph/src/graphrag_toolkit/document_graph/config.py b/document-graph/src/graphrag_toolkit/document_graph/config.py new file mode 100644 index 00000000..cf0a12d8 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/config.py @@ -0,0 +1,6 @@ +"""Document Graph Config — minimal wrapper.""" +class DocumentGraphConfig: + """Minimal config. AWS operations go through graphrag-toolkit.""" + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) diff --git a/document-graph/src/graphrag_toolkit/document_graph/errors.py b/document-graph/src/graphrag_toolkit/document_graph/errors.py new file mode 100644 index 00000000..91146d2c --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/errors.py @@ -0,0 +1,131 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Document Graph Exception Classes and Error Handling. + +This module defines a comprehensive set of custom exceptions for document graph operations, +providing specific error types for different failure scenarios. These exceptions enable +precise error handling and debugging throughout the document graph toolkit. + +Exception Hierarchy: + - ModelError: Data model validation and processing failures + - BatchJobError: Batch processing and ETL pipeline failures + - IndexError: Graph indexing and search index failures + - DatabaseConnectionError: Database connection and communication failures + - ConfigurationError: Configuration validation and parameter failures + - QueryExecutionError: Graph query execution and syntax failures + +Error Handling Patterns: + All exceptions follow consistent patterns for error reporting and debugging: + - Clear error messages with context information + - Preservation of original exception details when wrapping errors + - Integration with logging system for error tracking + - Support for error recovery and retry mechanisms + +Usage: + from graphrag_toolkit.document_graph.errors import ( + ModelError, DatabaseConnectionError, QueryExecutionError + ) + + try: + result = execute_graph_operation() + except DatabaseConnectionError as e: + logger.error(f"Database connection failed: {e}") + # Implement retry logic + except QueryExecutionError as e: + logger.error(f"Query failed: {e}") + # Handle query errors + except ModelError as e: + logger.error(f"Model validation failed: {e}") + # Handle data model issues +""" + +import logging + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +class ModelError(Exception): + """Raised when model validation or processing fails. + + This exception indicates issues with data model validation, + schema mismatches, or model processing errors. + + Examples: + >>> try: + ... validate_document_model(invalid_data) + ... except ModelError as e: + ... print(f"Model validation failed: {e}") + """ + pass + +class BatchJobError(Exception): + """Raised when batch job operations fail. + + This exception indicates failures in batch processing operations, + such as ETL pipeline errors or bulk data processing issues. + + Examples: + >>> try: + ... run_batch_job(job_config) + ... except BatchJobError as e: + ... print(f"Batch job failed: {e}") + """ + pass + +class IndexError(Exception): + """Raised when graph indexing operations fail. + + This exception indicates failures in graph database indexing, + search index creation, or index maintenance operations. + + Examples: + >>> try: + ... create_graph_index(index_config) + ... except IndexError as e: + ... print(f"Index creation failed: {e}") + """ + pass + +class DatabaseConnectionError(Exception): + """Raised when database connection fails. + + This exception indicates failures in establishing or maintaining + connections to graph databases ( Neptune, Neo4j). + + Examples: + >>> try: + ... connect_to_database(connection_config) + ... except DatabaseConnectionError as e: + ... print(f"Database connection failed: {e}") + """ + pass + +class ConfigurationError(Exception): + """Raised when configuration is invalid or missing. + + This exception indicates issues with provider configuration, + missing required parameters, or invalid configuration values. + + Examples: + >>> try: + ... validate_config(user_config) + ... except ConfigurationError as e: + ... print(f"Configuration error: {e}") + """ + pass + +class QueryExecutionError(Exception): + """Raised when graph query execution fails. + + This exception indicates failures in executing Cypher or Gremlin + queries against graph databases. + + Examples: + >>> try: + ... execute_query("MATCH (n) RETURN n LIMIT 10") + ... except QueryExecutionError as e: + ... print(f"Query execution failed: {e}") + """ + pass diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/__init__.py new file mode 100644 index 00000000..5849870b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/__init__.py @@ -0,0 +1,3 @@ +"""Build stage — Node/Edge → Cypher → GraphStore.""" + +from .cypher_builder import node_to_cypher, edge_to_cypher, batch_nodes_to_cypher, batch_edges_to_cypher diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/build_result.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/build_result.py new file mode 100644 index 00000000..5c62e1c2 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/build_result.py @@ -0,0 +1,263 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build Result module for document graph operations. + +This module defines the BuildResult class, which represents the complete +output of a build operation. It includes information about the build status, +output files, graph statistics, and metadata. It also provides factory methods +for creating success, empty, and failed results. +""" + +import logging +from typing import Optional, Dict, Any, List +from datetime import datetime, timezone +from pydantic import BaseModel, ConfigDict, computed_field +from graphrag_toolkit.lexical_graph import TenantId, to_tenant_id + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +class BuildResult(BaseModel): + """ + Represents the result of a build operation. + + The BuildResult class models the outcome of a build process, providing details + such as the document identifier, tenant information, build mode, and status. It + is designed to encapsulate various states including successful, empty, or + failed build results. This class includes helper methods to simplify the + creation of instances based on different build scenarios. + + Attributes: + document_id (str): The unique identifier of the document associated with the + build operation. + tenant_id (TenantId): The tenant information to which this build result + belongs. + build_mode (str): The mode in which the build operation was performed. Can + be one of: "graph", "code", "schema", or "sample". + build_time (str): The timestamp indicating when the build was created. + status (str): The status of the build operation, default is "success". Other + valid statuses include "empty_input" and "failed". + output_files (Optional[Dict[str, str]]): A dictionary containing the output + files resulting from the build operation. + graph_stats (Optional[Dict[str, Any]]): Optional metrics related to the + generated graph during the build. + metadata (Optional[Dict[str, Any]]): Additional metadata about the build + operation. + """ + + #Add this to allow custom TenantId types + model_config = ConfigDict(arbitrary_types_allowed=True) + + document_id: str + tenant_id: TenantId + build_mode: str # graph, code, schema, sample + build_time: str + status: str = "success" + output_files: Optional[Dict[str, str]] = None + graph_stats: Optional[Dict[str, Any]] = None + metadata: Optional[Dict[str, Any]] = None + tenant_id: Optional[str] = None + + @computed_field + @property + def scope_id(self) -> Optional[str]: + """Returns 'tenant_id.tenant_id' when both present, None otherwise.""" + if self.tenant_id is not None and self.tenant_id is not None: + return f"{self.tenant_id}.{self.tenant_id}" + return None + + def __init__(self, **data): + """ + [Brief one-line description of the function's purpose] + + [Detailed explanation of what the function does and how it works] + + Args: + [param_name] ([param_type]): [Description of parameter] + [param_name] ([param_type], optional): [Description of parameter]. Defaults to [default_value]. + + Returns: + [return_type]: [Description of return value] + + Raises: + [exception_type]: [Conditions under which exception is raised] + """ + if 'build_time' not in data: + data['build_time'] = datetime.now(timezone.utc).isoformat() + if 'tenant_id' in data: + data['tenant_id'] = to_tenant_id(data['tenant_id']) + super().__init__(**data) + logger.debug(f"Created BuildResult for document: {self.document_id} (mode: {self.build_mode})") + + def __str__(self) -> str: + return f"BuildResult(document_id={self.document_id}, mode={self.build_mode}, status={self.status})" + + def is_successful(self) -> bool: + """ + Determines whether the operation was successful based on the status attribute. + + Returns + ------- + bool + True if the status attribute equals "success", otherwise False. + """ + return self.status == "success" + + def is_empty(self) -> bool: + """ + Determines if the current instance represents an empty state. + + This method checks whether the instance is in an empty state based on its + `status` attribute or the lack of output files. + + Returns: + bool: True if the instance is in an empty state, False otherwise. + """ + return self.status == "empty_input" or not self.output_files + + @classmethod + def create_success( + cls, + document_id: str, + build_mode: str, + output_files: Dict[str, str], + tenant_id: Optional[str] = None, + graph_stats: Optional[Dict[str, Any]] = None, + metadata: Optional[Dict[str, Any]] = None, + tenant_id: Optional[str] = None + ) -> "BuildResult": + """ + Creates a successful build result object for a document build process. + + Summary: + This class method is responsible for generating a `BuildResult` instance + representing a successful document build operation. The method logs the process + and initializes the `BuildResult` object with relevant details such as document + ID, build mode, output files, and optional metadata or stats. + + Args: + document_id (str): The unique identifier of the document being built. + build_mode (str): The mode used during the build process (e.g., production, + draft). + output_files (Dict[str, str]): A mapping of output file types to their + corresponding paths or locations. + tenant_id (Optional[str]): The identifier of the tenant, if applicable. + graph_stats (Optional[Dict[str, Any]]): Statistics about graph data, if + available. + metadata (Optional[Dict[str, Any]]): Additional metadata associated with + the build. + tenant_id (Optional[str]): The root identifier for scoped operation. + + Returns: + BuildResult: A `BuildResult` instance initialized with the provided data. + + Raises: + None + """ + logger.info(f"Creating successful build result: {document_id} (mode: {build_mode})") + return cls( + document_id=document_id, + tenant_id=to_tenant_id(tenant_id), + build_mode=build_mode, + status="success", + output_files=output_files, + graph_stats=graph_stats or {}, + metadata=metadata or {}, + ) + + @classmethod + def create_empty( + cls, + document_id: str, + build_mode: str, + tenant_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + tenant_id: Optional[str] = None + ) -> "BuildResult": + """ + Creates an empty BuildResult object with a predefined status and optional metadata. + + This class method serves as a utility to instantiate a BuildResult object in situations + where no input is provided or where an empty initialization is required. It initializes + the object with the provided `document_id`, `build_mode`, `tenant_id`, and optional + metadata. The status of the created object is set to "empty_input" by default. + + Arguments: + document_id: str + Unique identifier of the document associated with the build result. + build_mode: str + Specifies the mode in which the build is executed. + tenant_id: Optional[str] + Identifier for the tenant, if applicable. If not provided, + a default tenant ID will be used. + metadata: Optional[Dict[str, Any]] + Additional information related to the build result, provided + as a dictionary. Defaults to an empty dictionary if not provided. + tenant_id: Optional[str] + The root identifier for scoped operation. + + Returns: + BuildResult + A new instance of BuildResult with the specified attributes and + a predefined status "empty_input". + + Additional Notes: + This method logs a warning message indicating that an empty build + result is being created for the specified document. + """ + logger.warning(f"Creating empty build result: {document_id}") + return cls( + document_id=document_id, + tenant_id=to_tenant_id(tenant_id), + build_mode=build_mode, + status="empty_input", + metadata=metadata or {}, + ) + + @classmethod + def create_failed( + cls, + document_id: str, + build_mode: str, + error_message: str, + tenant_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + tenant_id: Optional[str] = None + ) -> "BuildResult": + """ + Creates a failed build result with the provided details. + + This method initializes a failed build result instance for the specified document, + document creation/build mode, and error message. + + Parameters: + document_id: str + Identifier of the document for which the build has failed. + build_mode: str + Mode in which the build was attempted. + error_message: str + Detailed error message explaining the failure. + tenant_id: Optional[str] + Optional identifier of the tenant. Defaults to None. + metadata: Optional[Dict[str, Any]] + Optional additional metadata associated with the failure. Defaults to None. + tenant_id: Optional[str] + The root identifier for scoped operation. + + Returns: + BuildResult + An instance of the BuildResult class in a failed state with the provided information. + """ + logger.error(f"Creating failed build result: {document_id} - {error_message}") + metadata = metadata or {} + metadata["error_message"] = error_message + return cls( + document_id=document_id, + tenant_id=to_tenant_id(tenant_id), + build_mode=build_mode, + status="failed", + metadata=metadata, + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/__init__.py new file mode 100644 index 00000000..7cc1d815 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/__init__.py @@ -0,0 +1,12 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph build constructors — registry, factory, and base classes for graph construction.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_registry import ConstructorProviderRegistry +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_config import ConstructorProviderConfig +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_factory import ConstructorProviderFactory +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_plan import ConstructorPlan + +# Auto-register all constructors +from graphrag_toolkit.document_graph.graph_build.constructors import register_constructors + diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_plan.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_plan.py new file mode 100644 index 00000000..43f9fb63 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_plan.py @@ -0,0 +1,94 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Constructors plan — orchestrates execution of constructor providers on data.""" + +import logging +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_config import ConstructorProviderConfig +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_factory import ConstructorProviderFactory +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +# Use document-graph logging +logger = logging.getLogger(__name__) + +class ConstructorPlan: + """ + Manages a construction plan for executing multiple constructor providers. + + The `ConstructorPlan` class is responsible for handling a list of constructor + provider configurations. It executes all constructors in sequence, processes + an input DataFrame, and combines the resulting nodes and edges. Appropriate + logging and error handling are applied throughout the execution of the + constructor pipeline. + """ + + def __init__(self, configs: List[ConstructorProviderConfig]): + """ + Initializes an instance of the ConstructorPlan with a given list of constructor + configurations. Logs the number of configurations provided during + initialization. + + Attributes: + configs: List of ConstructorProviderConfig + A list of configuration objects provided for the ConstructorPlan + instance. + + Parameters: + configs: List[ConstructorProviderConfig] + List of constructor configuration objects passed during the + initialization of the class. + """ + self.configs = configs + logger.info(f"ConstructorPlan initialized with {len(configs)} constructors") + + def execute(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """ + Executes a pipeline of constructors to generate nodes and edges from a given data. + + This method processes an input DataFrame and iteratively applies a series of constructors + defined in the instance configuration. Each constructor generates a set of nodes and edges + based on its specific implementation logic and the provided data. The pipeline continues + even if individual constructors fail, logging warnings or errors at each step. At the end + of the pipeline, the method aggregates all the nodes and edges created by the constructors + and returns them as a tuple. + + Raises: + Logger warnings and errors for an empty DataFrame or for failures during the execution + of individual constructors. + + Parameters: + data (pd.DataFrame): The input data to be processed. + + Returns: + Tuple[List[Node], List[Edge]]: A tuple containing two lists: + - A list of all nodes created by the constructors in the pipeline. + - A list of all edges created by the constructors in the pipeline. + """ + if data.empty: + logger.warning("Input DataFrame is empty") + return [], [] + + all_nodes = [] + all_edges = [] + + logger.info(f"Starting constructor pipeline: {len(data)} rows") + + for i, config in enumerate(self.configs): + try: + logger.info(f"Executing constructor {i+1}/{len(self.configs)}: {config.name} ({config.type})") + + constructor = ConstructorProviderFactory.create(config) + nodes, edges = constructor.construct(data) + + all_nodes.extend(nodes) + all_edges.extend(edges) + + logger.info(f" Completed: +{len(nodes)} nodes, +{len(edges)} edges") + + except Exception as e: + logger.error(f"Constructor {config.name} failed: {e}") + # Continue with other constructors + + logger.info(f"Constructor pipeline completed: {len(all_nodes)} total nodes, {len(all_edges)} total edges") + return all_nodes, all_edges \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_base.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_base.py new file mode 100644 index 00000000..10fd1c99 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_base.py @@ -0,0 +1,71 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Constructors provider base — abstract base class for constructor providers.""" + +import logging +from abc import ABC, abstractmethod +from typing import List, Tuple + +import pandas as pd + +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +# Use document-graph logging +logger = logging.getLogger(__name__) + +class ConstructorProvider(ABC): + """ + Provides an interface for constructing graph nodes and edges from data. + + This is an abstract base class designed to standardize the construction + of graph components (nodes and edges) from a given dataset. By defining + a common interface, it ensures a consistent approach and validation for + turning raw data into graph representations. Extend this class and provide + implementations for the abstract methods to use it effectively. + """ + + def __init__(self, config): + """Initialize constructor with configuration.""" + self.config = config + self.args = config.args + logger.debug(f"Initialized {self.__class__.__name__} with config: {config.type}") + + @abstractmethod + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """ + Defines an abstract method for constructing graph nodes and edges from a given + data source. This method is intended to be implemented by subclasses to provide + specific logic for parsing data and generating corresponding graph structures. + + Parameters: + data (pd.DataFrame): A pandas DataFrame that serves as the input data + for constructing the graph structure. + + Returns: + Tuple[List[Node], List[Edge]]: A tuple containing a list of Node objects + and a list of Edge objects. These represent the constructed graph's nodes + and edges respectively. + + Raises: + None + """ + pass + + def validate_input(self, data: pd.DataFrame) -> bool: + """ + Validates the input DataFrame to ensure it is not empty. + + This method checks whether the provided DataFrame is valid by verifying that + it is neither None nor empty. It returns a boolean indicating the validity + of the DataFrame. If the DataFrame is invalid, a warning will be logged. + + Parameters: + data (pd.DataFrame): The DataFrame to be validated. + + Returns: + bool: True if the DataFrame is valid (not None and not empty), otherwise False. + """ + if data is None or data.empty: + logger.warning(f"{self.__class__.__name__}: Input DataFrame is empty") + return False + return True \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_config.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_config.py new file mode 100644 index 00000000..407bbf80 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_config.py @@ -0,0 +1,52 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Constructors provider config — configuration dataclass for constructor providers.""" + +from dataclasses import dataclass, field +from typing import Dict, Any, Optional, List + +@dataclass +class ConstructorProviderConfig: + """ + Represents the configuration for a constructor provider. + + This class defines the parameters required to initialize and validate a constructor provider. It ensures + that mandatory fields are present and allows for optional customization of additional configurations. + + Attributes: + name: A string representing the name of the constructor. + type: A string representing the type of the constructor. + args: A dictionary containing additional arguments for the constructor configuration. Defaults to + an empty dictionary. + description: An optional string providing a description for the constructor. Defaults to None. + required_columns: An optional list of strings specifying the required columns for the constructor. + Defaults to None. + batch_size: An optional integer specifying the batch size. Defaults to None. + """ + name: str + type: str + args: Dict[str, Any] = field(default_factory=dict) + description: Optional[str] = None + required_columns: Optional[List[str]] = None + batch_size: Optional[int] = None + + def __post_init__(self): + """ + Validates and initializes post-construction logic for the object. + + Ensures that the `name` and `type` attributes are not empty upon initialization. + If the `description` attribute is not provided, it sets a default description + using the `type` and `name` attributes. + + Raises: + ValueError: If `name` is empty. + ValueError: If `type` is empty. + """ + if not self.name: + raise ValueError("Constructor name cannot be empty") + if not self.type: + raise ValueError("Constructor type cannot be empty") + + # Set default description if not provided + if not self.description: + self.description = f"{self.type} constructor: {self.name}" \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_factory.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_factory.py new file mode 100644 index 00000000..79bb28f4 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_factory.py @@ -0,0 +1,33 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Constructors provider factory — creates constructor provider instances from config.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_registry import ConstructorProviderRegistry +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_config import ConstructorProviderConfig + +class ConstructorProviderFactory: + """ + Factory class for creating instances of constructor providers. + + This class provides a method to create a specific constructor provider + instance based on the configuration provided. It utilizes a registry to + look up the appropriate constructor provider based on the type specified + in the configuration. + """ + + @staticmethod + def create(config: ConstructorProviderConfig): + """ + This static method is responsible for creating an instance of a provider class based on the provided + configuration. It retrieves the appropriate class from the ConstructorProviderRegistry based on the + type specified in the configuration and initializes it using the given configuration. + + Parameters: + config (ConstructorProviderConfig): The configuration object used to determine the provider + type and to initialize the provider instance. + + Returns: + Any: An instance of the provider class corresponding to the type defined in the configuration. + """ + provider_class = ConstructorProviderRegistry.get(config.type) + return provider_class(config) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_registry.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_registry.py new file mode 100644 index 00000000..821fd6fb --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/constructors_provider_registry.py @@ -0,0 +1,75 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Constructors provider registry — centralized registry of constructor provider classes.""" + +from typing import Dict, Type +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_base import ConstructorProvider + +class ConstructorProviderRegistry: + """ + Manages the registration and access of constructor providers. + + This class serves as a centralized registry for constructor provider classes, + allowing registration, retrieval, and listing of available providers. It is + commonly used to dynamically manage and access different provider classes + by their unique names. + """ + + _providers: Dict[str, Type[ConstructorProvider]] = {} + + @classmethod + def register(cls, name: str, provider_class: Type[ConstructorProvider]): + """ + Registers a provider with the given name. + + This class method allows registering a new provider by associating a name with + a provider class. + + Parameters: + name: str + The unique name used to identify the provider. + provider_class: Type[ConstructorProvider] + A class representing the provider to be registered, typically inheriting + from ConstructorProvider. + + Raises: + KeyError + If the provided name is already registered in the providers. + """ + cls._providers[name] = provider_class + + @classmethod + def get(cls, name: str) -> Type[ConstructorProvider]: + """ + Retrieves a constructor provider based on the provided name. + + This method accesses a registry of constructor providers and retrieves the provider + corresponding to the given name. If the name is not found within the registry, an + exception is raised. + + Args: + name (str): The name of the desired constructor provider. + + Returns: + Type[ConstructorProvider]: The constructor provider associated with the given name. + + Raises: + ValueError: If the specified name is not found in the registry of constructor providers. + """ + if name not in cls._providers: + raise ValueError(f"Unknown constructor provider: {name}") + return cls._providers[name] + + @classmethod + def list_providers(cls) -> list: + """ + Provides functionality to retrieve the list of available providers. + + This method is a class-level method that retrieves and returns a list of + all the available providers registered within the class. + + Returns: + list: A list containing the keys of the registered providers within + the class. + """ + return list(cls._providers.keys()) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/__init__.py new file mode 100644 index 00000000..1d105863 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/__init__.py @@ -0,0 +1,17 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Core constructors — basic graph building operations for nodes, edges, and properties.""" + +# Core constructors for basic graph building operations + +from .schema_driven_constructor import SchemaDrivenConstructor +from .node_constructor import NodeConstructor +from .edge_constructor import EdgeConstructor +from .property_constructor import PropertyConstructor + +__all__ = [ + 'SchemaDrivenConstructor', + 'NodeConstructor', + 'EdgeConstructor', + 'PropertyConstructor' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/edge_constructor.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/edge_constructor.py new file mode 100644 index 00000000..a0f94f29 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/edge_constructor.py @@ -0,0 +1,48 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Edge constructor — generates graph edges from DataFrame relationships.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_base import ConstructorProvider +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +class EdgeConstructor(ConstructorProvider): + """ + Provides functionality for constructing edges from DataFrame relationships. + + This class is responsible for generating a list of nodes and edges based on the + relationships defined within a given pandas DataFrame. It extends functionality from + a ConstructorProvider class, allowing it to process data for graph-based representations + or analysis. + + Methods + ------- + construct(data) + Constructs edges and nodes from the provided DataFrame. + """ + + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """ + Converts a given DataFrame into nodes and edges. + + This method processes the input DataFrame and triage_constructs a list of nodes + and edges based on its data. The structure of the DataFrame is expected to + conform to the necessary format suitable for the transformation. + + Parameters: + data: pd.DataFrame + The input DataFrame containing the data to be converted + into nodes and edges. The DataFrame must contain the + required information for defining both the nodes and their connections. + + Returns: + Tuple[List[Node], List[Edge]] + A tuple where the first element is a list of nodes and the second + element is a list of edges derived from the input DataFrame. + """ + # TODO: Implement DataFrame to edges conversion + nodes = [] + edges = [] + return nodes, edges + diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/node_constructor.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/node_constructor.py new file mode 100644 index 00000000..5245ca5f --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/node_constructor.py @@ -0,0 +1,47 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Node constructor — converts DataFrame rows into graph nodes.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_base import ConstructorProvider +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +class NodeConstructor(ConstructorProvider): + """ + Represents a node constructor that converts data into nodes and edges. + + This class is used to construct nodes and edges from the input data provided + as a DataFrame. It inherits from the `ConstructorProvider` base class and + implements the `construct` method, which processes the data to generate a + list of `Node` objects and a list of `Edge` objects. + + Methods + ------- + construct(data: pd.DataFrame) -> Tuple[List[Node], List[Edge]] + Constructs and returns the nodes and edges based on the provided DataFrame. + """ + + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """ + Converts a DataFrame into nodes and edges. + + This method takes a pandas DataFrame, processes its content, and produces a list + of nodes and edges based on the data. Each row in the DataFrame may represent + relationships or entities that are used to construct nodes and edges. + + Parameters: + data : pd.DataFrame + A pandas DataFrame containing raw data from which nodes and edges are to be + constructed. + + Returns: + Tuple[List[Node], List[Edge]] + A tuple containing two lists: + - nodes: A list of Node objects created from the DataFrame. + - edges: A list of Edge objects defining relationships between the nodes. + """ + # TODO: Implement DataFrame to nodes conversion + nodes = [] + edges = [] + return nodes, edges \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/property_constructor.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/property_constructor.py new file mode 100644 index 00000000..e4eb374a --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/property_constructor.py @@ -0,0 +1,40 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Property constructor — constructs node and edge properties from data.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_base import ConstructorProvider +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +class PropertyConstructor(ConstructorProvider): + """ + Handles the construction of properties from given data. + + This class is specifically designed to generate properties, represented as nodes + and edges, from the provided data encapsulated in a DataFrame. It serves as a + provider that processes this data and triage_constructs meaningful structures that can + be used in further processing or analysis. + """ + + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """ + Construct nodes and edges from a given DataFrame. + + This function is intended to map the properties of a DataFrame to a list of + nodes and edges. The details of the mapping process need to be implemented. + It will create and return a list of nodes and edges based on the provided + data structure. + + Args: + data (pd.DataFrame): The input DataFrame that contains the necessary data + for constructing nodes and edges. + + Returns: + Tuple[List[Node], List[Edge]]: A tuple containing a list of nodes and a + list of edges. + """ + # TODO: Implement DataFrame columns to properties mapping + nodes = [] + edges = [] + return nodes, edges \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/schema_driven_constructor.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/schema_driven_constructor.py new file mode 100644 index 00000000..f9c93366 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/core/schema_driven_constructor.py @@ -0,0 +1,43 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Schema-driven constructor — creates nodes and edges based on schema definitions.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_base import ConstructorProvider +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +class SchemaDrivenConstructor(ConstructorProvider): + """ + Represents a schema-driven constructor for creating nodes and edges. + + This class provides functionality to construct nodes and edges + from a schema-driven configuration. It works based on the input + data and uses specific logic defined in its implementation to + derive graph elements. Designed to be extended or instantiated + in systems where schema-driven graph construction is required. + """ + + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """ + Construct a schema-driven graph from a given DataFrame. + + This method processes a pandas DataFrame and creates a graph representation + composed of nodes and edges based on a predefined schema. The exact schema + for constructing the graph must be implemented in this method. This function + returns two lists: one containing Node objects and another containing Edge + objects that define relationships between the nodes. + + Args: + data: The input pandas DataFrame containing the data to construct the graph + from. + + Returns: + A tuple where the first element is a list of Node objects representing + the entities in the graph, and the second element is a list of Edge + objects representing the relationships between the entities. + """ + # TODO: Implement schema-driven graph construction + nodes = [] + edges = [] + return nodes, edges \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/list_constructors.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/list_constructors.py new file mode 100644 index 00000000..ba8a3bf3 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/list_constructors.py @@ -0,0 +1,73 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""API to list available constructors.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_registry import ConstructorProviderRegistry + +# Import registration module to ensure constructors are registered +from . import register_constructors + +def list_available_constructors(): + """ + Lists all registered constructor providers available in the ConstructorProviderRegistry. + + Returns: + list: A list containing all registered providers in the + ConstructorProviderRegistry. + """ + return ConstructorProviderRegistry.list_providers() + +def list_core_constructors(): + """ + Filters and returns a list of core constructors from all available constructors. + + This function retrieves a list of all constructors from the + ConstructorProviderRegistry and filters them to match a predefined + list of core constructor names. The filtered list contains + only the constructors that are considered core. + + Returns: + list: A list of constructor names that match the predefined + list of core constructors. + """ + all_constructors = ConstructorProviderRegistry.list_providers() + core = ["schema_driven", "node_constructor", "edge_constructor", "property_constructor"] + return [const for const in all_constructors if const in core] + +def list_pattern_constructors(): + """ + Lists all constructors that match specific patterns. + + This function retrieves all available constructor providers and filters + them based on predefined patterns. + + Returns: + list: A list of constructor providers that match the specified patterns. + """ + all_constructors = ConstructorProviderRegistry.list_providers() + patterns = ["one_to_many", "many_to_many"] + return [const for const in all_constructors if const in patterns] + +def list_optimization_constructors(): + """ + Filters and returns a list of constructors that match specified optimization features. + + This function retrieves all available constructors from the + `ConstructorProviderRegistry`, compares them against a set of predefined + optimization features, and returns only the ones that are part of the + optimization list. + + Returns: + list[str]: A list of constructors matching the specified optimization + features. + """ + all_constructors = ConstructorProviderRegistry.list_providers() + optimizations = ["batch_constructor", "deduplication"] + return [const for const in all_constructors if const in optimizations] + +if __name__ == "__main__": + print("Available Constructors:") + print("Core:", list_core_constructors()) + print("Patterns:", list_pattern_constructors()) + print("Optimizations:", list_optimization_constructors()) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/__init__.py new file mode 100644 index 00000000..59da13bf --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/__init__.py @@ -0,0 +1,13 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Optimization constructors — batch and deduplication strategies for graph building.""" + +# Optimization constructors for performance and efficiency + +from .batch_constructor import BatchConstructor +from .deduplication_constructor import DeduplicationConstructor + +__all__ = [ + 'BatchConstructor', + 'DeduplicationConstructor' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/batch_constructor.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/batch_constructor.py new file mode 100644 index 00000000..2af2b3c2 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/batch_constructor.py @@ -0,0 +1,37 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Batch constructor — groups graph operations into batches for efficiency.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_base import ConstructorProvider +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +class BatchConstructor(ConstructorProvider): + """ + + """ + + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """ + Construct nodes and edges from a given DataFrame. + + This method processes the given DataFrame and triage_constructs a list of nodes and edges + based on its content. The method is designed to be extended with additional processing + to handle large DataFrames in batches. + + Raises: + ValueError: If the DataFrame is invalid or does not meet the required format. + + Args: + data (pd.DataFrame): Input DataFrame containing information to construct nodes + and edges. + + Returns: + Tuple[List[Node], List[Edge]]: A tuple containing a list of constructed nodes + and edges. + """ + # TODO: Implement batch processing for large DataFrames + nodes = [] + edges = [] + return nodes, edges \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/deduplication_constructor.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/deduplication_constructor.py new file mode 100644 index 00000000..946cefab --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/optimizations/deduplication_constructor.py @@ -0,0 +1,42 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deduplication constructor — ensures unique graph objects during construction.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_base import ConstructorProvider +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +class DeduplicationConstructor(ConstructorProvider): + """ + DeduplicationConstructor is responsible for constructing unique graph objects. + + This class provides implementations for the construction of deduplicated nodes + and edges from the given input data. It ensures that no duplicate graph elements + are created, facilitating the creation of a cleaner graph structure. + + Attributes + ---------- + None + """ + + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """ + Construct a graph representation from the given data. + + This method processes a DataFrame to generate a list of unique nodes and + edges, suitable for representing a graph structure. Deduplication logic + for nodes and edges is currently to be implemented. + + Args: + data (pd.DataFrame): The input data containing information required to + construct nodes and edges. + + Returns: + Tuple[List[Node], List[Edge]]: A tuple containing a list of unique nodes + and a list of unique edges derived from the input data. + """ + # TODO: Implement deduplication logic for nodes and edges + nodes = [] + edges = [] + return nodes, edges \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/__init__.py new file mode 100644 index 00000000..175ff1f7 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/__init__.py @@ -0,0 +1,13 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pattern constructors — common relationship patterns like one-to-many and many-to-many.""" + +# Pattern-based constructors for common relationship patterns + +from .one_to_many_constructor import OneToManyConstructor +from .many_to_many_constructor import ManyToManyConstructor + +__all__ = [ + 'OneToManyConstructor', + 'ManyToManyConstructor' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/many_to_many_constructor.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/many_to_many_constructor.py new file mode 100644 index 00000000..d708b941 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/many_to_many_constructor.py @@ -0,0 +1,18 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Many-to-many constructor — handles M:N relationships via junction patterns.""" + +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_base import ConstructorProvider +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +class ManyToManyConstructor(ConstructorProvider): + """Handle M:N relationships via junction patterns.""" + + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """Construct M:N relationship patterns.""" + # TODO: Implement many-to-many relationship construction + nodes = [] + edges = [] + return nodes, edges diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/one_to_many_constructor.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/one_to_many_constructor.py new file mode 100644 index 00000000..48b209b4 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/patterns/one_to_many_constructor.py @@ -0,0 +1,18 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-to-many constructor — handles 1:N parent-child relationships.""" + +from ..constructors_provider_base import ConstructorProvider +from typing import List, Tuple +import pandas as pd +from graphrag_toolkit.document_graph.model_elements import Node, Edge + +class OneToManyConstructor(ConstructorProvider): + """Handle 1:N relationships (one parent to many children).""" + + def construct(self, data: pd.DataFrame) -> Tuple[List[Node], List[Edge]]: + """Construct 1:N relationship patterns.""" + # TODO: Implement one-to-many relationship construction + nodes = [] + edges = [] + return nodes, edges \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/register_constructors.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/register_constructors.py new file mode 100644 index 00000000..d5985a3c --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/constructors/register_constructors.py @@ -0,0 +1,34 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Register constructors — auto-registers all constructor providers with the registry.""" + +# Register all constructor providers +from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_registry import ConstructorProviderRegistry + +# Core constructors +from graphrag_toolkit.document_graph.graph_build.constructors.core.schema_driven_constructor import SchemaDrivenConstructor +from graphrag_toolkit.document_graph.graph_build.constructors.core.node_constructor import NodeConstructor +from graphrag_toolkit.document_graph.graph_build.constructors.core.edge_constructor import EdgeConstructor +from graphrag_toolkit.document_graph.graph_build.constructors.core.property_constructor import PropertyConstructor + +# Pattern constructors +from graphrag_toolkit.document_graph.graph_build.constructors.patterns.one_to_many_constructor import OneToManyConstructor +from graphrag_toolkit.document_graph.graph_build.constructors.patterns.many_to_many_constructor import ManyToManyConstructor + +# Optimization constructors +from graphrag_toolkit.document_graph.graph_build.constructors.optimizations.batch_constructor import BatchConstructor +from graphrag_toolkit.document_graph.graph_build.constructors.optimizations.deduplication_constructor import DeduplicationConstructor + +# Register core constructors +ConstructorProviderRegistry.register("schema_driven", SchemaDrivenConstructor) +ConstructorProviderRegistry.register("node_constructor", NodeConstructor) +ConstructorProviderRegistry.register("edge_constructor", EdgeConstructor) +ConstructorProviderRegistry.register("property_constructor", PropertyConstructor) + +# Register pattern constructors +ConstructorProviderRegistry.register("one_to_many", OneToManyConstructor) +ConstructorProviderRegistry.register("many_to_many", ManyToManyConstructor) + +# Register optimization constructors +ConstructorProviderRegistry.register("batch_constructor", BatchConstructor) +ConstructorProviderRegistry.register("deduplication", DeduplicationConstructor) diff --git a/document-graph/src/graphrag_toolkit/document_graph/graph_build/cypher_builder.py b/document-graph/src/graphrag_toolkit/document_graph/graph_build/cypher_builder.py new file mode 100644 index 00000000..79ff9101 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/graph_build/cypher_builder.py @@ -0,0 +1,67 @@ +"""Cypher Builder — generates openCypher MERGE statements from Node/Edge models. + +Uses graphrag-toolkit's GraphStore.query() for execution. +""" + +from typing import Optional, Union +from ..model_elements import Node, Edge + +try: + from graphrag_toolkit.lexical_graph import TenantId, to_tenant_id + HAS_LEXICAL = True +except ImportError: + HAS_LEXICAL = False + + +def _format_label(label: str, tenant_id: Optional[str] = None) -> str: + """Format a label using TenantId.format_label() for consistency with lexical-graph. + + Produces backtick-quoted labels: `Label` (default tenant) or `Label{tenant}__` (scoped). + Falls back to manual formatting if lexical-graph is not installed. + """ + if not tenant_id: + return f'`{label}`' + if HAS_LEXICAL: + tid = to_tenant_id(tenant_id) + return tid.format_label(label) + # Fallback when lexical-graph not installed — same format + return f'`{label}{tenant_id}__`' + + +def node_to_cypher(node: Node, tenant_id: Optional[str] = None) -> tuple[str, dict]: + """Generate MERGE statement for a typed node.""" + labels = node.labels or ["Node"] + labels = [_format_label(l, tenant_id) for l in labels] + label_str = ":".join(labels) + props = node.properties or {} + id_val = node.id + + query = f"MERGE (n:{label_str} {{id: $id_val}}) SET n += $props" + params = {"id_val": id_val, "props": props} + return query, params + + +def edge_to_cypher(edge: Edge, tenant_id: Optional[str] = None) -> tuple[str, dict]: + """Generate MERGE statement for a typed edge.""" + rel_type = edge.label + + query = ( + f"MATCH (a {{id: $src_id}}), (b {{id: $tgt_id}}) " + f"MERGE (a)-[r:{rel_type}]->(b) SET r += $props" + ) + params = { + "src_id": edge.source_id, + "tgt_id": edge.target_id, + "props": edge.properties or {}, + } + return query, params + + +def batch_nodes_to_cypher(nodes: list[Node], tenant_id: Optional[str] = None) -> list[tuple[str, dict]]: + """Generate MERGE statements for a batch of nodes.""" + return [node_to_cypher(n, tenant_id) for n in nodes] + + +def batch_edges_to_cypher(edges: list[Edge], tenant_id: Optional[str] = None) -> list[tuple[str, dict]]: + """Generate MERGE statements for a batch of edges.""" + return [edge_to_cypher(e, tenant_id) for e in edges] diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/__init__.py new file mode 100644 index 00000000..b53db66f --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/__init__.py @@ -0,0 +1 @@ +"""Ingest stage — data preparation (column/row/field operations).""" diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/column/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/__init__.py new file mode 100644 index 00000000..8a4aa3d1 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/__init__.py @@ -0,0 +1,15 @@ +"""Column-level ingestors — renaming, reordering, selection, and type conversion.""" +# Column-level ingestors + +from graphrag_toolkit.document_graph.ingest.column.column_renamer import ColumnRenamerProvider + +from graphrag_toolkit.document_graph.ingest.column.column_reorder import ColumnReorderProvider +from graphrag_toolkit.document_graph.ingest.column.column_selector import ColumnSelectorProvider +from graphrag_toolkit.document_graph.ingest.column.column_type_converter import ColumnTypeConverterProvider + +__all__ = [ + 'ColumnRenamerProvider', + 'ColumnReorderProvider', + 'ColumnSelectorProvider', + 'ColumnTypeConverterProvider' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_renamer.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_renamer.py new file mode 100644 index 00000000..404b25bc --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_renamer.py @@ -0,0 +1,41 @@ +import pandas as pd +from graphrag_toolkit.document_graph.ingest.ingestors_provider_base import IngestorProvider + +class ColumnRenamerProvider(IngestorProvider): + """ + A provider class for column renaming in a DataFrame. + + The ColumnRenamerProvider allows ingestion of a DataFrame with column + renaming applied based on a provided mapping. The mapping should be + supplied via the arguments dictionary and contains the old column names + as keys and the new column names as values. Only columns that exist in + the DataFrame will be renamed. + + Methods: + ingest(data: pd.DataFrame) -> pd.DataFrame + Renames DataFrame columns based on the provided mapping. + + """ + + def ingest(self, data: pd.DataFrame) -> pd.DataFrame: + """ + Processes and renames columns in the given DataFrame based on a mapping + configuration. If a mapping is provided, only the columns that exist + in the DataFrame and are included in the mapping will be renamed. + + Parameters: + data (pd.DataFrame): The input DataFrame to process. + + Returns: + pd.DataFrame: A new DataFrame with renamed columns according to the + provided mapping. + """ + mapping = self.args.get("mapping", {}) + + if not mapping: + return data + + # Only rename columns that exist + existing_mapping = {old: new for old, new in mapping.items() if old in data.columns} + + return data.rename(columns=existing_mapping) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_reorder.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_reorder.py new file mode 100644 index 00000000..b91a9b48 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_reorder.py @@ -0,0 +1,53 @@ +import pandas as pd +from graphrag_toolkit.document_graph.ingest.ingestors_provider_base import IngestorProvider + +class ColumnReorderProvider(IngestorProvider): + """ + Provides functionality to reorder the columns of a given pandas DataFrame + based on a specified column order. + + The class is designed to take a pandas DataFrame and rearrange its columns + based on a predefined order specified in the `args` attribute. If no column + order is specified, the original DataFrame is returned unchanged. The primary + purpose of the class is to modify column arrangements while preserving columns + that are not explicitly mentioned in the provided order. This ensures all + columns in the DataFrame are retained, even if not reordered. + + Methods + ------- + ingest(data: pd.DataFrame) -> pd.DataFrame + Reorders the columns of the provided DataFrame as per the specified + column order or retains the original order if no preference is provided. + """ + + def ingest(self, data: pd.DataFrame) -> pd.DataFrame: + """ + Reorders the columns of a DataFrame based on the specified order provided in the arguments. + + If a specific column order is provided in the arguments (`column_order`), the method will + rearrange the DataFrame's columns to match the specified order. Columns not present in the + `column_order` list will retain their positions and be appended after the ordered columns. + + Parameters: + data (pd.DataFrame): The DataFrame to be ingested and reordered. The input DataFrame + should contain columns that may partially or fully match the `column_order` list in + the instance's arguments. + + Returns: + pd.DataFrame: A DataFrame with columns reordered based on the `column_order` provided + in the arguments. If `column_order` is not specified or empty, the original DataFrame + is returned without any modifications. + """ + column_order = self.args.get("column_order", []) + + if not column_order: + return data + + # Keep existing columns that aren't in the order list + existing_columns = [col for col in column_order if col in data.columns] + remaining_columns = [col for col in data.columns if col not in column_order] + + # Combine ordered columns with remaining columns + new_order = existing_columns + remaining_columns + + return data[new_order] diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_selector.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_selector.py new file mode 100644 index 00000000..616d099e --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_selector.py @@ -0,0 +1,52 @@ +import pandas as pd +from graphrag_toolkit.document_graph.ingest.ingestors_provider_base import IngestorProvider + +class ColumnSelectorProvider(IngestorProvider): + """ + Provides functionality to either select specific columns or drop specific columns + from a given pandas DataFrame. + + The ColumnSelectorProvider class processes incoming data based on the specified + action ("select" or "drop") and a list of columns. It is useful when there is a + need to streamline DataFrame operations by filtering or excluding certain + columns dynamically during data ingestion. + + Methods + ------- + ingest(data: pd.DataFrame) -> pd.DataFrame + Processes the DataFrame according to the specified action and columns. + """ + + def ingest(self, data: pd.DataFrame) -> pd.DataFrame: + """ + Processes a DataFrame by either selecting or dropping specified columns based on + the provided action. The method utilizes the arguments passed to dynamically + filter columns of the DataFrame. This is useful for cleaning or transforming + datasets by retaining or removing specific columns. + + Parameters: + data (pd.DataFrame): The input DataFrame that needs to be processed. + + Returns: + pd.DataFrame: The modified DataFrame with columns either selected or dropped, + based on the specified action. + + Raises: + ValueError: Raised when the provided action is unsupported. Supported values + are 'select' or 'drop'. + """ + action = self.args.get("action", "select") # "select" or "drop" + columns = self.args.get("columns", []) + + if not columns: + return data + + if action == "select": + # Keep only specified columns + available_columns = [col for col in columns if col in data.columns] + return data[available_columns] + elif action == "drop": + # Drop specified columns + return data.drop(columns=[col for col in columns if col in data.columns]) + else: + raise ValueError(f"Unsupported action: {action}. Use 'select' or 'drop'.") \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_type_converter.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_type_converter.py new file mode 100644 index 00000000..7a893dc6 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/column_type_converter.py @@ -0,0 +1,48 @@ +import pandas as pd +from graphrag_toolkit.document_graph.ingest.ingestors_provider_base import IngestorProvider + +class ColumnTypeConverterProvider(IngestorProvider): + """ + Provides functionality to convert dataframe column types based on a specified mapping. + + This class extends the base IngestorProvider and allows for dynamic conversion of column data + types in a pandas DataFrame. It is particularly useful when ingesting data with columns that + need to adhere to a specific type structure for further processing or analysis. + + Methods + ------- + ingest(data: pd.DataFrame) -> pd.DataFrame + Converts column types in the provided DataFrame according to the 'type_mapping' argument. + + """ + + def ingest(self, data: pd.DataFrame) -> pd.DataFrame: + type_mapping = self.args.get("type_mapping", {}) + + if not type_mapping: + return data + + result = data.copy() + + for column, target_type in type_mapping.items(): + if column not in result.columns: + continue + + try: + if target_type == "string": + result[column] = result[column].astype(str) + elif target_type == "int": + result[column] = pd.to_numeric(result[column], errors='coerce').astype('Int64') + elif target_type == "float": + result[column] = pd.to_numeric(result[column], errors='coerce') + elif target_type == "datetime": + result[column] = pd.to_datetime(result[column], errors='coerce') + elif target_type == "bool": + result[column] = result[column].astype(bool) + else: + # Direct pandas dtype + result[column] = result[column].astype(target_type) + except Exception as e: + print(f"Warning: Could not convert column '{column}' to {target_type}: {e}") + + return result \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/column/register_ingestors.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/register_ingestors.py new file mode 100644 index 00000000..eee8eeb0 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/column/register_ingestors.py @@ -0,0 +1,14 @@ +"""Register column ingestors — registers all column-level ingestor providers.""" +# Register all column_level ingestors +from graphrag_toolkit.document_graph.ingest.ingestors_provider_registry import IngestorProviderRegistry +from graphrag_toolkit.document_graph.ingest.column.column_selector import ColumnSelectorProvider +from graphrag_toolkit.document_graph.ingest.column.column_renamer import ColumnRenamerProvider +from graphrag_toolkit.document_graph.ingest.column.column_reorder import ColumnReorderProvider +from graphrag_toolkit.document_graph.ingest.column.column_type_converter import ColumnTypeConverterProvider + +# Register column_level ingestors +IngestorProviderRegistry.register("column_selector", ColumnSelectorProvider) +IngestorProviderRegistry.register("column_renamer", ColumnRenamerProvider) +IngestorProviderRegistry.register("column_reorder", ColumnReorderProvider) +IngestorProviderRegistry.register("column_type_converter", ColumnTypeConverterProvider) + diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/field/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/field/__init__.py new file mode 100644 index 00000000..1e80fcd9 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/field/__init__.py @@ -0,0 +1,14 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Field-level ingestors for data transformation and cleanup. + +This module contains ingestors that operate on individual fields or field-level +transformations during the data extraction phase. +""" + +from .numeric_id_cleanup_ingestor import NumericIdCleanupIngestor + +__all__ = [ + 'NumericIdCleanupIngestor', +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/field/numeric_id_cleanup_ingestor.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/field/numeric_id_cleanup_ingestor.py new file mode 100644 index 00000000..c6934f23 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/field/numeric_id_cleanup_ingestor.py @@ -0,0 +1,40 @@ +"""Numeric ID cleanup ingestor — removes trailing decimals from numeric ID fields.""" +import pandas as pd +from graphrag_toolkit.document_graph.ingest.ingestors_provider_base import IngestorProvider + + +class NumericIdCleanupIngestor(IngestorProvider): + """ + Ingestor provider for cleaning up numeric IDs with specific formatting. + + This class processes a pandas DataFrame to clean up numeric ID fields, + specifically by converting them to string format and removing any `.0` + suffixes. It ensures the field specified in the arguments is cleaned if + it exists in the DataFrame. Intended for use in data ingestion workflows + where numeric IDs might be improperly formatted or need standardization. + + Methods: + ingest: Processes and cleans a DataFrame based on the specified field. + + """ + + def ingest(self, data: pd.DataFrame) -> pd.DataFrame: + """ + Processes a DataFrame to modify a specified field by converting its values to strings + and removing ".0" suffix if present. The operation is applied only when the field + is specified and exists in the DataFrame columns. + + Parameters: + data (pd.DataFrame): The input DataFrame containing the data to be processed. + + Returns: + pd.DataFrame: The processed DataFrame with the specified field modified, if applicable. + """ + field = self.args.get("field") + if not field or field not in data.columns: + return data + + # Convert to string and remove .0 suffix + data[field] = data[field].astype(str).str.replace(r'\.0$', '', regex=True) + + return data \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/field/register_ingestors.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/field/register_ingestors.py new file mode 100644 index 00000000..5e500959 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/field/register_ingestors.py @@ -0,0 +1,7 @@ +"""Register field ingestors — registers all field-level ingestor providers.""" +# Register all field_level ingestors +from ..ingestors_provider_registry import IngestorProviderRegistry +from .numeric_id_cleanup_ingestor import NumericIdCleanupIngestor + +# Register field_level ingestors +IngestorProviderRegistry.register("numeric_id_cleanup", NumericIdCleanupIngestor) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_error_handler.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_error_handler.py new file mode 100644 index 00000000..164f533a --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_error_handler.py @@ -0,0 +1,112 @@ +"""Ingestors error handler — error handling and validation utilities for ingestion.""" +import logging +import pandas as pd +from typing import List, Dict, Any + +# Use document-graph logging +logger = logging.getLogger(__name__) + +class IngestorErrorHandler: + """ + Provides error handling and validation utilities for ingestor processes. + + IngestorErrorHandler contains static methods that validate the presence of required + columns in a DataFrame, ensure configuration fields are complete, and handle failures + by rolling back to a safe state. This class is designed to support robust error + management in data ingestion pipelines, promoting reliability and easier debugging. + """ + + @staticmethod + def validate_column_exists(df: pd.DataFrame, column: str, ingestor_name: str): + """ + Validates the existence of a specified column in a DataFrame for the given ingestor. + + If the specified column does not exist in the DataFrame, an error message will be + logged and a ValueError will be raised. + + Parameters: + df (pd.DataFrame): The input DataFrame in which to check the availability of the column. + column (str): The name of the column to verify existence. + ingestor_name (str): The name of the ingestor to include in error logging. + + Raises: + ValueError: If the specified column is not found in the DataFrame's columns. + """ + if column not in df.columns: + error_msg = f"{ingestor_name}: Column '{column}' not found. Available columns: {list(df.columns)}" + logger.error(error_msg) + raise ValueError(error_msg) + + @staticmethod + def validate_columns_exist(df: pd.DataFrame, columns: List[str], ingestor_name: str): + """ + Validates if the specified columns exist in the provided DataFrame. This method checks + for missing columns in the DataFrame compared to the list of required columns. If any + columns are missing, an error is logged and an exception is raised. + + Args: + df (pd.DataFrame): The DataFrame to validate. + columns (List[str]): A list of column names to check for existence in the DataFrame. + ingestor_name (str): The name of the ingestor for logging purposes. + + Raises: + ValueError: If any of the specified columns are missing in the DataFrame. + """ + if missing := [col for col in columns if col not in df.columns]: + error_msg = f"{ingestor_name}: Columns not found: {missing}. Available: {list(df.columns)}" + logger.error(error_msg) + raise ValueError(error_msg) + + @staticmethod + def handle_ingestor_failure(ingestor_name: str, error: Exception, df: pd.DataFrame) -> pd.DataFrame: + """ + Handles the failure of an ingestor operation and reverts to the original DataFrame. + + The method logs an error and a warning message whenever an ingestor fails. It ensures + that in case of any error during ingestion, the original DataFrame is retained and + returned unchanged. + + Parameters: + ingestor_name: str + The name of the ingestor that encountered a failure. + error: Exception + The exception instance describing the failure. + df: pd.DataFrame + The original DataFrame prior to the failed ingestion attempt. + + Returns: + pd.DataFrame + The original DataFrame unchanged, ensuring data consistency. + """ + logger.error(f"{ingestor_name} failed: {error}") + logger.warning(f"Rolling back to original DataFrame with {len(df)} rows") + return df + + @staticmethod + def validate_config(config_dict: Dict[str, Any], required_fields: List[str], ingestor_name: str): + """ + Validates a configuration dictionary against a list of required fields. + + This method checks that all required fields are present in the provided + configuration dictionary. If any required fields are missing, an error + is logged, and a ValueError is raised. + + Parameters: + config_dict : Dict[str, Any] + The configuration dictionary to validate. + required_fields : List[str] + A list of field names required to be in the configuration dictionary. + ingestor_name : str + The name of the ingestor, included in error messages for better + context. + + Raises: + ValueError + Raised when one or more required fields are missing from the + configuration dictionary. + """ + missing = [field for field in required_fields if field not in config_dict] + if missing: + error_msg = f"{ingestor_name}: Missing required config fields: {missing}" + logger.error(error_msg) + raise ValueError(error_msg) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_plan.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_plan.py new file mode 100644 index 00000000..9e44bec8 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_plan.py @@ -0,0 +1,171 @@ +"""Ingestors plan — orchestrates execution of ingestor providers on DataFrames.""" +import logging +from typing import List +import pandas as pd +from .ingestors_provider_config import IngestorProviderConfig +from .ingestors_provider_factory import IngestorProviderFactory +from .ingestors_validator import IngestorConfigValidator +from .ingestors_schema import IngestorSchemaTracker +from .ingestors_error_handler import IngestorErrorHandler + +# Use document-graph logging +logger = logging.getLogger(__name__) + +class IngestorPlan: + """ + Represents a plan for executing a sequence of ingestor operations on a dataset. + + The IngestorPlan class is used to manage and execute a series of data ingestion steps + defined by a set of configurations. Each ingestor performs a specific processing step, + and the class ensures that errors are handled and schema changes are tracked throughout + the execution. The class also provides functionalities for validating configurations + upon initialization and tracking schema modifications to aid in debugging or pipeline analysis. + """ + + def __init__(self, configs: List[IngestorProviderConfig], validate: bool = True): + """ + Initialize an IngestorPlan instance. + + The IngestorPlan class manages the initialization and validation of ingestor + configurations. It tracks the state of ingestor schemas and ensures the + provided configurations are valid based on the validation rules. + + Arguments: + configs: List of IngestorProviderConfig objects containing the configurations + for initializing ingestors. + validate: Boolean flag to specify whether to validate the configurations. + Default is True. + + Raises: + ValueError: If validation is enabled and any configuration errors are detected. + """ + self.configs = configs + self.schema_tracker = None + + if validate: + errors = IngestorConfigValidator.validate_configs(configs) + if errors: + error_msg = "Ingestor configuration errors:\n" + "\n".join(errors) + logger.error(error_msg) + raise ValueError(error_msg) + + logger.info(f"IngestorPlan initialized with {len(configs)} ingestors") + + def execute(self, data: pd.DataFrame) -> pd.DataFrame: + """ + Executes a pipeline of ingestors on the input DataFrame. + + This method processes the provided DataFrame by sequentially applying + configured ingestors in the pipeline. Each ingestor modifies the DataFrame + and tracks schema changes using an internal schema tracker. If any ingestor + fails, the failure is logged, and an error handler processes the issue. + + Parameters: + data (pd.DataFrame): The input DataFrame to process. + + Returns: + pd.DataFrame: The resulting DataFrame after applying the pipeline of ingestors. + + Raises: + Exception: If an unexpected error occurs during the execution of the pipeline. + """ + if data.empty: + logger.warning("Input DataFrame is empty") + return data + + # Initialize schema tracking + self.schema_tracker = IngestorSchemaTracker(data) + result = data.copy() + + logger.info(f"Starting ingestor pipeline: {len(result)} rows, {len(result.columns)} columns") + + for i, config in enumerate(self.configs): + try: + logger.info(f"Executing ingestor {i+1}/{len(self.configs)}: {config.name} ({config.type})") + + # Store state before ingestor + before_df = result.copy() + + # Execute ingestor + ingestor = IngestorProviderFactory.create(config) + result = ingestor.ingest(result) + + # Track schema changes + change_type = self._determine_change_type(config.type, before_df, result) + self.schema_tracker.track_change( + ingestor_name=config.name, + change_type=change_type, + details=config.args, + df_before=before_df, + df_after=result + ) + + logger.info(f" Completed: {len(before_df)} → {len(result)} rows") + + except Exception as e: + logger.error(f"Ingestor {config.name} failed: {e}") + result = IngestorErrorHandler.handle_ingestor_failure(config.name, e, result) + + # Print final summary + self.schema_tracker.print_summary() + + logger.info(f"Ingestor pipeline completed: {len(result)} rows, {len(result.columns)} columns") + return result + + def get_schema_changes(self): + """ + Retrieves schema changes tracked by the schema tracker. + + This method checks if a schema tracker exists and, if so, returns the + changes tracked by it. If no schema tracker is available, an empty list + is returned. + + Returns: + list: A list containing the tracked schema changes, or an empty list + if no schema tracker is present. + """ + if self.schema_tracker: + return self.schema_tracker.changes + return [] + + def _determine_change_type(self, ingestor_type: str, before_df: pd.DataFrame, after_df: pd.DataFrame) -> str: + """ + Determines the type of change between two dataframes given their structure and content. + + This method compares two dataframes to determine the nature of the changes + between them, such as removal or addition of rows, modification of data, or + alteration of column structure. It also includes logic to detect specific + column-related changes like renaming, addition, or removal. The determination + of the change type relies entirely on structural and content differences + of the dataframes. + + Parameters: + ingestor_type: str + Identifier for the type of ingestor or process responsible for + generating the dataframes. + before_df: pd.DataFrame + The initial dataframe before the change was applied. + after_df: pd.DataFrame + The modified dataframe after the change was applied. + + Returns: + str: A string indicating the type of change detected. Possible values are: + - "rows_filtered": If the number of rows has changed. + - "columns_removed": If fewer columns exist after the change. + - "columns_added": If additional columns exist after the change. + - "columns_renamed": If the column names differ without a change + in the column count. + - "data_modified": If neither the rows nor columns changed, but + data content was altered. + """ + if len(before_df) != len(after_df): + return "rows_filtered" + elif list(before_df.columns) != list(after_df.columns): + if len(before_df.columns) > len(after_df.columns): + return "columns_removed" + elif len(before_df.columns) < len(after_df.columns): + return "columns_added" + else: + return "columns_renamed" + else: + return "data_modified" \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_base.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_base.py new file mode 100644 index 00000000..cba3cf99 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_base.py @@ -0,0 +1,70 @@ +"""Ingestors provider base — abstract base class for ingestor providers.""" +import logging +from abc import ABC, abstractmethod +from typing import List, Dict, Any +import pandas as pd + +# Use document-graph logging +logger = logging.getLogger(__name__) + +class IngestorProvider(ABC): + """ + Abstract base class that defines the interface and behavior for an ingestor provider. + + This class is intended to be inherited and implemented by subclasses to perform + specific ingestion tasks. It provides the basic blueprint and utility methods for + ingesting and validating data. Subclasses must implement the `ingest` method to + define the ingestion logic. The class also includes a utility method for input + validation. + + Attributes: + config: Configuration object for the ingestor, which provides necessary + parameters and settings. + + Methods must be overridden or used directly where applicable in the defined + workflow for ingestion. + """ + + def __init__(self, config): + """ + Initializes an instance of the class with the provided configuration. + + Attributes: + config: The configuration object associated with the instance. + args: Parsed arguments extracted from the provided configuration. + + Args: + config: The configuration object used to initialize the instance. + """ + self.config = config + self.args = config.args + logger.debug(f"Initialized {self.__class__.__name__} with config: {config.type}") + + @abstractmethod + def ingest(self, data: pd.DataFrame) -> pd.DataFrame: + """ + Abstract method that must be implemented by derived classes to process and + transform the input data. + + Parameters: + data (pd.DataFrame): Input data in the form of a Pandas DataFrame. + + Returns: + pd.DataFrame: Processed and transformed data as a Pandas DataFrame. + """ + pass + + def validate_input(self, data: pd.DataFrame) -> bool: + """ + Validates the input DataFrame to ensure it is not None or empty. + + This method checks if the input DataFrame is valid by ensuring it is not + None and contains data. + + Returns: + bool: True if the input DataFrame is valid, otherwise False. + """ + if data is None or data.empty: + logger.warning(f"{self.__class__.__name__}: Input DataFrame is empty") + return False + return True \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_config.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_config.py new file mode 100644 index 00000000..7eb5c9bc --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_config.py @@ -0,0 +1,52 @@ +"""Ingestors provider config — configuration dataclass for ingestor providers.""" +from dataclasses import dataclass, field +from typing import Dict, Any, Optional, List + +@dataclass +class IngestorProviderConfig: + """ + Represents the configuration for an ingestor provider. + + This class encapsulates the necessary attributes and validation rules for + defining an ingestor provider's configuration. It is mainly used to store + the metadata and validation requirements for the ingestor. + + Attributes: + name: A unique name identifying the ingestor. + type: The type or category of the ingestor. + args: A dictionary containing additional parameters for the ingestor. + description: An optional description for the ingestor; defaults to a + combination of type and name if not provided. + required_columns: An optional list of column names required by the + ingestor. + + Raises: + ValueError: Raised if `name` or `type` is empty during initialization. + + """ + name: str + type: str + args: Dict[str, Any] = field(default_factory=dict) + description: Optional[str] = None + required_columns: Optional[List[str]] = None + + def __post_init__(self): + """ + Performs post-initialization validation and sets default values for some + attributes of the ingestor instance after the object is instantiated. + + Raises + ------ + ValueError + If the 'name' attribute is empty or None. + ValueError + If the 'type' attribute is empty or None. + """ + if not self.name: + raise ValueError("Ingestor name cannot be empty") + if not self.type: + raise ValueError("Ingestor type cannot be empty") + + # Set default description if not provided + if not self.description: + self.description = f"{self.type} ingestor: {self.name}" \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_factory.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_factory.py new file mode 100644 index 00000000..912e7291 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_factory.py @@ -0,0 +1,34 @@ +from graphrag_toolkit.document_graph.ingest.ingestors_provider_registry import IngestorProviderRegistry +from graphrag_toolkit.document_graph.ingest.ingestors_provider_config import IngestorProviderConfig + +class IngestorProviderFactory: + """ + Provides a factory for creating instances of ingestor providers based on a + given configuration. + + This class serves as a factory for instantiating ingestor providers that + are registered in the `IngestorProviderRegistry`. It uses the type specified + in the configuration to determine the appropriate provider class and returns + an instance of it. + """ + + @staticmethod + def create(config: IngestorProviderConfig): + """ + Creates an instance of an ingestor provider based on the provided configuration. + + This method retrieves the appropriate ingestor provider class from the + IngestorProviderRegistry based on the type specified in the given configuration. + It then initializes an instance of this class using the configuration and + returns the created instance. + + Args: + config (IngestorProviderConfig): Configuration object containing + details necessary for selecting and initializing the ingestor provider. + + Returns: + IngestorProvider: An instance of the provider class initialized with + the given configuration. + """ + provider_class = IngestorProviderRegistry.get(config.type) + return provider_class(config) diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_registry.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_registry.py new file mode 100644 index 00000000..8e88e9c4 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_provider_registry.py @@ -0,0 +1,77 @@ +"""Ingestors provider registry — centralized registry of ingestor provider classes.""" +from typing import Dict, Type +from graphrag_toolkit.document_graph.ingest.ingestors_provider_base import IngestorProvider +from graphrag_toolkit.document_graph.ingest.field.numeric_id_cleanup_ingestor import NumericIdCleanupIngestor + +class IngestorProviderRegistry: + """ + Provides a registry for managing ingestor providers. + + The IngestorProviderRegistry class allows for the registration, retrieval, and + listing of ingestor providers by name. This serves as a central management system for + providers, enabling dynamic customization and extension of functionality. + + Methods: + register: Registers a new ingestor provider with a specified name. + get: Retrieves a registered ingestor provider by its name. + list_providers: Lists all currently registered provider names. + """ + + _providers: Dict[str, Type[IngestorProvider]] = { + 'numeric_id_cleanup': NumericIdCleanupIngestor, + } + + @classmethod + def register(cls, name: str, provider_class: Type[IngestorProvider]): + """ + Registers a new ingestor provider with a specified name. + + This class method allows adding a provider to the internal registry by associating it + with a unique name. The registered ingestor provider can be used in the system for + handling specific ingestion tasks. The uniqueness of the provider name must be ensured + by the caller to avoid accidental overwrites. + + Parameters: + name (str): The unique identifier for the ingestor provider. + provider_class (Type[IngestorProvider]): The class implementing the ingestor + provider functionality. + + Raises: + KeyError: Raises an error if the provider name already exists in the registry. + """ + cls._providers[name] = provider_class + + @classmethod + def get(cls, name: str) -> Type[IngestorProvider]: + """ + Retrieves an ingestor provider class by its name. + + Raises a ValueError if the provided name does not exist in the registered + ingestor providers. + + Args: + name: The name of the ingestor provider to retrieve. + + Returns: + The ingestor provider class associated with the given name. + + Raises: + ValueError: If the given name is not a valid provider. + """ + if name not in cls._providers: + raise ValueError(f"Unknown ingestor provider: {name}") + return cls._providers[name] + + @classmethod + def list_providers(cls) -> list: + """ + Returns a list of available providers for the class. + + This method retrieves the keys from the internal `_providers` dictionary + and returns them as a list. It is used to view all registered providers + accessible for the class. + + Returns: + list: A list containing the names of all available providers. + """ + return list(cls._providers.keys()) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_schema.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_schema.py new file mode 100644 index 00000000..6617702d --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_schema.py @@ -0,0 +1,223 @@ +"""Ingestors schema — tracks schema changes during ingestion pipeline execution.""" +import logging +import pandas as pd +from typing import Dict, List, Any, Optional +from dataclasses import dataclass + +# Use document-graph logging +logger = logging.getLogger(__name__) + +@dataclass +class IngestorSchemaChange: + """ + Represents a schema change event in an ingestor. + + This data structure is used to track schema-related changes occurring + during data ingestion. It stores the details of the change, including + the type of change, the affected rows, and the affected columns. + + Attributes: + ingestor_name: Name of the ingestor where the schema change occurred. + change_type: Type of schema change. Possible values include "column_added", + "column_removed", "column_renamed", "rows_filtered". + details: A dictionary containing additional information about the schema change. + rows_before: Number of rows before the change occurred. + rows_after: Number of rows after the change occurred. + columns_before: List of column names present before the change occurred. + columns_after: List of column names present after the change occurred. + """ + ingestor_name: str + change_type: str # "column_added", "column_removed", "column_renamed", "rows_filtered" + details: Dict[str, Any] + rows_before: int + rows_after: int + columns_before: List[str] + columns_after: List[str] + +@dataclass +class SourceSchema: + """ + Represents a source schema with information about dataset structure. + + This class provides metadata and details about the structure of a dataset, + including its columns, data types, the number of rows, and sample data + for initial validation or analysis. + + Attributes: + columns: List of column names in the dataset. + dtypes: Mapping of column names to their respective data types. + row_count: Total number of rows in the dataset. + sample_data: A sample set of data values from the dataset, useful + for verification or inspection. + """ + columns: List[str] + dtypes: Dict[str, str] + row_count: int + sample_data: Dict[str, Any] + +@dataclass +class IngestorSchema: + """ + Represents the schema for an ingestor. + + The IngestorSchema class defines the necessary schema required for data ingestion + operations. It encapsulates the source schema, final columns and their data types, + final row count, and the list of changes applied during ingestion. + + Attributes: + source_schema: SourceSchema + The original schema from the source data being ingested. + final_columns: List[str] + The list of column names that will exist in the final ingested data. + final_dtypes: Dict[str, str] + A dictionary mapping each column name to its data type in the final data. + final_row_count: int + The total row count in the final ingested dataset. + changes: List[IngestorSchemaChange] + A list of changes or transformations made to the schema during the ingestion + process. + """ + source_schema: SourceSchema + final_columns: List[str] + final_dtypes: Dict[str, str] + final_row_count: int + changes: List[IngestorSchemaChange] + +class IngestorSchemaTracker: + """ + Tracks schema changes for a source DataFrame during data ingestion processes. + + This class is designed to monitor and document changes to a source DataFrame's schema + caused by data ingestors. It keeps track of initial schema details, lists all changes made, + and provides a mechanism to summarize these changes. The class also offers a way to + retrieve the final schema after all transformations. + + Attributes: + source_schema: An object representing the schema of the source DataFrame, including + column names, data types, row count, and an optional sample of data. + changes: A list of schema change events recorded during data ingestion. + + Methods: + __init__(source_df) + Initializes the schema tracker with the structure of the source DataFrame. + track_change(ingestor_name, change_type, details, df_before, df_after) + Records a schema change caused by an ingestor. + get_final_schema(final_df) + Constructs and retrieves the final schema after all data ingestion operations. + print_summary() + Displays a logged summary of all schema changes. + """ + + def __init__(self, source_df: pd.DataFrame): + """ + Initializes a schema tracker for a source DataFrame. + + Summary: + The constructor method initializes the schema tracking by capturing details + about the column names, data types, row count, and sample data from the provided + DataFrame. It also sets up an empty list to track schema changes and logs an + initialization message. + + Parameters: + source_df (pd.DataFrame): The DataFrame for which schema tracking is to be + initialized. + + Attributes: + source_schema (SourceSchema): Contains information about the structure of the + source DataFrame, including columns, data types, row count, and a sample record. + changes (List[IngestorSchemaChange]): A list intended to track schema changes + that occur during operations. + """ + self.source_schema = SourceSchema( + columns=list(source_df.columns), + dtypes={col: str(dtype) for col, dtype in source_df.dtypes.items()}, + row_count=len(source_df), + sample_data=source_df.head(1).to_dict('records')[0] if not source_df.empty else {} + ) + self.changes: List[IngestorSchemaChange] = [] + logger.info(f"Schema tracker initialized: {len(self.source_schema.columns)} columns, {self.source_schema.row_count} rows") + + def track_change(self, ingestor_name: str, change_type: str, details: Dict[str, Any], + df_before: pd.DataFrame, df_after: pd.DataFrame): + """ + Tracks changes in the ingested data and logs the details of the changes, including the + number of rows and columns before and after the change. + + Parameters: + ingestor_name: str + The name of the ingestor responsible for the dataset change. + change_type: str + A description of the type of change applied (e.g., "schema_update"). + details: Dict[str, Any] + Additional information or metadata about the change. + df_before: pd.DataFrame + The dataframe representing the state of the data before the change. + df_after: pd.DataFrame + The dataframe representing the state of the data after the change. + """ + change = IngestorSchemaChange( + ingestor_name=ingestor_name, + change_type=change_type, + details=details, + rows_before=len(df_before), + rows_after=len(df_after), + columns_before=list(df_before.columns), + columns_after=list(df_after.columns) + ) + self.changes.append(change) + + logger.info(f"{ingestor_name}: {change_type} - {change.rows_before}→{change.rows_after} rows, " + f"{len(change.columns_before)}→{len(change.columns_after)} columns") + + def get_final_schema(self, final_df: pd.DataFrame) -> IngestorSchema: + """ + Generate the final schema for the given DataFrame. + + This method prepares an instance of IngestorSchema by collecting metadata + from the provided DataFrame, including column names, data types, row count, + and any registered changes. It helps to encapsulate the final schema details + for use in data ingestion or transformation processes. + + Args: + final_df (pd.DataFrame): A DataFrame which contains the finalized data structure. + + Returns: + IngestorSchema: An object encapsulating the schema details including + source schema, final column names, data types, row count, and applied changes. + """ + return IngestorSchema( + source_schema=self.source_schema, + final_columns=list(final_df.columns), + final_dtypes={col: str(dtype) for col, dtype in final_df.dtypes.items()}, + final_row_count=len(final_df), + changes=self.changes + ) + + def print_summary(self): + """ + Logs a summary of the schema changes during the ingestion process. + + The method provides a detailed summary of the source schema and all + the changes applied during the ingestion. Each change is logged with + its name, type of change, number of rows and columns before and after + the change, and additional details if provided. The final state of the + data after all changes is also logged. + + Raises: + AttributeError: If required attributes like `source_schema` or + `changes` are missing or improperly formatted. + """ + logger.info("=== INGESTOR SCHEMA SUMMARY ===") + logger.info(f"Source: {len(self.source_schema.columns)} columns, {self.source_schema.row_count} rows") + + for change in self.changes: + logger.info(f" {change.ingestor_name}: {change.change_type}") + logger.info(f" Rows: {change.rows_before} → {change.rows_after}") + logger.info(f" Columns: {len(change.columns_before)} → {len(change.columns_after)}") + if change.details: + logger.info(f" Details: {change.details}") + + if self.changes: + final_change = self.changes[-1] + logger.info(f"Final: {len(final_change.columns_after)} columns, {final_change.rows_after} rows") + logger.info("=== END SUMMARY ===") \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_validator.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_validator.py new file mode 100644 index 00000000..dbf1b967 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/ingestors_validator.py @@ -0,0 +1,246 @@ +"""Ingestors validator — validates ingestor configurations before execution.""" +import logging +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.ingest.ingestors_provider_config import IngestorProviderConfig +from graphrag_toolkit.document_graph.ingest.ingestors_provider_registry import IngestorProviderRegistry + +# Use document-graph logging +logger = logging.getLogger(__name__) + +class IngestorConfigValidator: + """ + This class provides methods to validate ingestor configurations for an ingestion pipeline. + + The class aims to ensure that the configuration settings for different types of ingestors are + compatible and meet the required standards. It validates individual configuration parameters, + checks for missing required fields, and ensures there are no dependency conflicts between + multiple ingestors. This helps ensure a robust and error-free ingestion process. + """ + + # Required fields for each ingestor type + REQUIRED_FIELDS = { + "skip_row": ["conditions"], + "date_range_filter": ["date_field"], + "column_selector": ["columns"], + "column_renamer": ["mapping"], + "column_reorder": ["column_order"], + "column_type_converter": ["type_mapping"] + } + + @classmethod + def validate_config(cls, config: IngestorProviderConfig) -> List[str]: + """ + Validates the configuration for an ingestor type. + + This method checks if the provided ingestor type is registered, ensures that + all required fields for the specified ingestor type are present, and performs + type-specific validations based on the type of the ingestor. Errors found + during the validation process are collected and returned as a list of strings. + + Args: + config (IngestorProviderConfig): Configuration object containing the type + of the ingestor and its arguments. + + Returns: + List[str]: A list of error messages, if any validation issues are found. + """ + errors = [] + + # Check if ingestor type is registered + try: + IngestorProviderRegistry.get(config.type) + except ValueError: + errors.append(f"Unknown ingestor type: {config.type}") + return errors + + # Check required fields + required = cls.REQUIRED_FIELDS.get(config.type, []) + for field in required: + if field not in config.args: + errors.append(f"{config.type}: Missing required field '{field}'") + + # Type-specific validation + if config.type == "skip_row": + errors.extend(cls._validate_skip_row(config.args)) + elif config.type == "date_range_filter": + errors.extend(cls._validate_date_range_filter(config.args)) + elif config.type == "column_selector": + errors.extend(cls._validate_column_selector(config.args)) + + return errors + + @classmethod + def validate_configs(cls, configs: List[IngestorProviderConfig]) -> List[str]: + """ + Validates a list of ingestor configuration objects and checks for possible issues. + + This method performs validation on each provided configuration by using + the `validate_config` method. It also checks for dependency conflicts + among configurations and collects any validation errors. + + Args: + configs (List[IngestorProviderConfig]): A list of configuration objects + to be validated. + + Returns: + List[str]: A list of error messages for invalid configurations, or an + empty list if all configurations are valid. + """ + all_errors = [] + + for i, config in enumerate(configs): + errors = cls.validate_config(config) + for error in errors: + all_errors.append(f"Ingestor {i+1} ({config.name}): {error}") + + # Check for dependency conflicts + all_errors.extend(cls._check_dependencies(configs)) + + return all_errors + + @classmethod + def _validate_skip_row(cls, args: Dict[str, Any]) -> List[str]: + """ + Validates the provided conditions for skipping rows against the required + format and applicable operator requirements. + + The method checks if the 'conditions' in `args` is a list and validates each + condition entry to ensure it is a dictionary containing required keys such as + 'field' and 'operator'. It also ensures that specific operators have associated + 'required' values. Returns a list of error messages detailing any validation + failures. + + Arguments: + args (Dict[str, Any]): A dictionary containing the conditions to validate. + + Returns: + List[str]: A list of error messages resulting from the validation process. + """ + errors = [] + conditions = args.get("conditions", []) + + if not isinstance(conditions, list): + errors.append("conditions must be a list") + return errors + + for i, condition in enumerate(conditions): + if not isinstance(condition, dict): + errors.append(f"Condition {i+1} must be a dictionary") + continue + + if "field" not in condition: + errors.append(f"Condition {i+1}: Missing 'field'") + if "operator" not in condition: + errors.append(f"Condition {i+1}: Missing 'operator'") + + operator = condition.get("operator") + if operator in ["eq", "ne", "lt", "le", "gt", "ge", "in", "not_in", "contains", "not_contains", "startswith", "endswith", "regex", "length_eq", "length_gt", "length_lt"]: + if "value" not in condition: + errors.append(f"Condition {i+1}: Operator '{operator}' requires 'value'") + + return errors + + @classmethod + def _validate_date_range_filter(cls, args: Dict[str, Any]) -> List[str]: + """ + Validates the presence of at least one date constraint in the given arguments. + + This method ensures that the arguments dictionary contains at least one of the + specified date constraint keys: "start_date", "end_date", "days_back", + "weeks_back", or "months_back". If none of these keys are present, it + accumulates an error message. + + Args: + args (Dict[str, Any]): The dictionary of arguments to be validated. + + Returns: + List[str]: A list of error messages. An empty list indicates that the + validation was successful. + """ + errors = [] + + # Must have at least one date constraint + date_constraints = ["start_date", "end_date", "days_back", "weeks_back", "months_back"] + if not any(constraint in args for constraint in date_constraints): + errors.append("Must specify at least one date constraint") + + return errors + + @classmethod + def _validate_column_selector(cls, args: Dict[str, Any]) -> List[str]: + """ + Validates the column selector arguments and returns a list of any detected + validation errors. This method ensures that the provided arguments for the + column selection or dropping process conform to the expected format and + requirements. + + Parameters: + args (Dict[str, Any]): A dictionary containing the arguments for the + column selector. Must include 'action' and 'columns'. + + Returns: + List[str]: A list of errors as strings, if any. An empty list is returned + if no errors are detected. + """ + errors = [] + + action = args.get("action", "select") + if action not in ["select", "drop"]: + errors.append(f"Invalid action '{action}'. Must be 'select' or 'drop'") + + columns = args.get("columns", []) + if not isinstance(columns, list): + errors.append("columns must be a list") + elif not columns: + errors.append("columns list cannot be empty") + + return errors + + @classmethod + def _check_dependencies(cls, configs: List[IngestorProviderConfig]) -> List[str]: + """ + Checks for dependencies and conflicts among a list of ingestor configurations. + + This method ensures that column operations, such as renaming and filtering, do not + depend on columns that have been removed by prior operations. It analyzes the sequence + of configurations and identifies any conflicting dependencies which might cause runtime + issues. + + Parameters: + configs : List[IngestorProviderConfig] + List of configurations for ingestors that define various operations such as + column selection, renaming, filtering, etc. + + Returns: + List[str] + A list of error messages describing any conflicts detected between column operations + performed by the ingestors. Each message specifies the index and type of conflict + in the configurations. + """ + errors = [] + + # Check for column operations after column removal + column_removers = [] + column_users = [] + + for i, config in enumerate(configs): + if config.type == "column_selector" and config.args.get("action") == "drop": + column_removers.append((i, config.args.get("columns", []))) + elif config.type in ["column_renamer", "skip_row", "date_range_filter"]: + if config.type == "column_renamer": + column_users.append((i, list(config.args.get("mapping", {}).keys()))) + elif config.type == "skip_row": + fields = [cond.get("field") for cond in config.args.get("conditions", [])] + column_users.append((i, fields)) + elif config.type == "date_range_filter": + column_users.append((i, [config.args.get("date_field")])) + + # Check if any column user comes after a column remover that removes its columns + for user_idx, user_columns in column_users: + for remover_idx, removed_columns in column_removers: + if remover_idx < user_idx: # Remover comes before user + conflicts = set(user_columns) & set(removed_columns) + if conflicts: + errors.append(f"Ingestor {user_idx+1} uses columns {list(conflicts)} that were removed by ingestor {remover_idx+1}") + + return errors \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/list_ingestors.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/list_ingestors.py new file mode 100644 index 00000000..88e09d63 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/list_ingestors.py @@ -0,0 +1,89 @@ +"""API to list available ingestors.""" + +from .ingestors_provider_registry import IngestorProviderRegistry + +# Import registration modules to ensure ingestors are registered +try: + from .row_level.register_ingestors import * +except ImportError: + # Handle hyphenated directory names + import importlib + row_module = importlib.import_module('.row_level.register_ingestors', package=__package__) + +try: + from .column_level.register_ingestors import * +except ImportError: + # Handle hyphenated directory names + import importlib + col_module = importlib.import_module('.column_level.register_ingestors', package=__package__) + +# Import field_level ingestors (no register_ingestors file needed, they're in registry) +# The numeric_id_cleanup is already registered in ingestors_provider_registry.py + +def list_available_ingestors(): + """ + Lists all available ingestor providers. + + This function retrieves and returns a list of all registered ingestor + providers from the IngestorProviderRegistry. + + Returns: + list: A list of available ingestor providers. + """ + return IngestorProviderRegistry.list_providers() + +def list_row_level_ingestors(): + """ + Lists all row-level ingestors. + + This function retrieves the list of all ingestors available in the + IngestorProviderRegistry and filters only those that operate at the + row level. It checks against predefined row-level ingestors and + returns a filtered list. + + Returns: + list[str]: A list of names of ingestors that process data at + the row level. + """ + all_ingestors = IngestorProviderRegistry.list_providers() + row_level = ["skip_row", "date_range_filter"] + return [ing for ing in all_ingestors if ing in row_level] + +def list_column_level_ingestors(): + """ + Filters and retrieves column-level ingestors from a list of all available ingestors. + + This function identifies specific ingestors operating at the column level, such as + selectors, renamers, reorders, and type converters, from the general list of + available ingestor providers. + + Returns + ------- + list[str] + A list of ingestor names corresponding to column-level operations. + """ + all_ingestors = IngestorProviderRegistry.list_providers() + column_level = ["column_selector", "column_renamer", "column_reorder", "column_type_converter"] + return [ing for ing in all_ingestors if ing in column_level] + +def list_field_level_ingestors(): + """ + Lists all field-level ingestors available in the IngestorProviderRegistry. + + This function retrieves a list of all ingestors registered in the + IngestorProviderRegistry and filters them to return only those that + operate at the field level. Field-level ingestors are predefined + and identified within the function. + + Returns: + list: A list of ingestor names that operate at the field level. + """ + all_ingestors = IngestorProviderRegistry.list_providers() + field_level = ["numeric_id_cleanup"] + return [ing for ing in all_ingestors if ing in field_level] + +if __name__ == "__main__": + print("Available Ingestors:") + print("Row-level:", list_row_level_ingestors()) + print("Column-level:", list_column_level_ingestors()) + print("Field-level:", list_field_level_ingestors()) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/row/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/row/__init__.py new file mode 100644 index 00000000..99bd21bd --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/row/__init__.py @@ -0,0 +1,10 @@ +"""Row-level ingestors — filtering and skipping rows during ingestion.""" +# Row-level ingestors + +from graphrag_toolkit.document_graph.ingest.row.date_range_filter import DateRangeFilterProvider +from .skip_row import SkipRowProvider + +__all__ = [ + 'DateRangeFilterProvider', + 'SkipRowProvider' +] diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/row/date_range_filter.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/row/date_range_filter.py new file mode 100644 index 00000000..28bc7ce9 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/row/date_range_filter.py @@ -0,0 +1,82 @@ +"""Date range filter — filters DataFrame rows based on date range criteria.""" +import pandas as pd +from datetime import datetime, timedelta, timezone +from typing import Optional +from graphrag_toolkit.document_graph.ingest.ingestors_provider_base import IngestorProvider + + +class DateRangeFilterProvider(IngestorProvider): + """Filter rows based on date ranges.""" + + def ingest(self, data: pd.DataFrame) -> pd.DataFrame: + date_field = self.args.get("date_field") + if not date_field or date_field not in data.columns: + return data + + # Convert to datetime if not already + date_column = pd.to_datetime(data[date_field], errors='coerce') + + # Start with all rows included + mask = pd.Series([True] * len(data), index=data.index) + + # Handle different date range specifications + start_date = self._parse_date(self.args.get("start_date")) + end_date = self._parse_date(self.args.get("end_date")) + + # Relative date ranges + days_back = self.args.get("days_back") + weeks_back = self.args.get("weeks_back") + months_back = self.args.get("months_back") + + # Auto-detect timezone handling - match data timezone + data_has_tz = date_column.dt.tz is not None + + # Calculate relative dates matching data timezone + if data_has_tz: + now = datetime.now(timezone.utc) # Use UTC for timezone-aware data + else: + now = datetime.now() # Use naive datetime for naive data + + if days_back: + start_date = now - timedelta(days=days_back) + elif weeks_back: + start_date = now - timedelta(weeks=weeks_back) + elif months_back: + start_date = now - timedelta(days=months_back * 30) # Approximate + + # Apply date filters - ensure timezone compatibility + def make_compatible(dt, target_has_tz): + if target_has_tz and dt.tzinfo is None: + return dt.replace(tzinfo=timezone.utc) + elif not target_has_tz and dt.tzinfo is not None: + return dt.replace(tzinfo=None) + return dt + + if start_date: + start_date = make_compatible(start_date, data_has_tz) + mask &= (date_column >= start_date) + if end_date: + end_date = make_compatible(end_date, data_has_tz) + mask &= (date_column <= end_date) + + # Handle null dates + include_null = self.args.get("include_null", False) + if not include_null: + mask &= date_column.notna() + + return data[mask] + + def _parse_date(self, date_value) -> Optional[datetime]: + """Parse various date formats.""" + if not date_value: + return None + + if isinstance(date_value, str): + try: + return pd.to_datetime(date_value) + except: + return None + elif isinstance(date_value, datetime): + return date_value + + return None \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/row/register_ingestors.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/row/register_ingestors.py new file mode 100644 index 00000000..0363fc03 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/row/register_ingestors.py @@ -0,0 +1,9 @@ +"""Register row ingestors — registers all row-level ingestor providers.""" +# Register all row_level ingestors +from ..ingestors_provider_registry import IngestorProviderRegistry +from .skip_row import SkipRowProvider +from .date_range_filter import DateRangeFilterProvider + +# Register row_level ingestors +IngestorProviderRegistry.register("skip_row", SkipRowProvider) +IngestorProviderRegistry.register("date_range_filter", DateRangeFilterProvider) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/ingest/row/skip_row.py b/document-graph/src/graphrag_toolkit/document_graph/ingest/row/skip_row.py new file mode 100644 index 00000000..867ca241 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/ingest/row/skip_row.py @@ -0,0 +1,71 @@ +"""Skip row — skips DataFrame rows based on flexible conditions.""" +import pandas as pd +from typing import List, Dict, Any, Union +from ..ingestors_provider_base import IngestorProvider + +class SkipRowProvider(IngestorProvider): + """Skip rows based on flexible conditions.""" + + def ingest(self, data: pd.DataFrame) -> pd.DataFrame: + conditions = self.args.get("conditions", []) + if not conditions: + return data + + # Start with all rows included + mask = pd.Series([True] * len(data), index=data.index) + + for condition in conditions: + field = condition["field"] + operator = condition["operator"] + value = condition.get("value") + + # Apply condition logic + if operator == "eq": # equals + mask &= (data[field] == value) + elif operator == "ne": # not equals + mask &= (data[field] != value) + elif operator == "lt": # less than + mask &= (data[field] < value) + elif operator == "le": # less than or equal + mask &= (data[field] <= value) + elif operator == "gt": # greater than + mask &= (data[field] > value) + elif operator == "ge": # greater than or equal + mask &= (data[field] >= value) + elif operator == "in": # in list + mask &= data[field].isin(value) + elif operator == "not_in": # not in list + mask &= ~data[field].isin(value) + elif operator == "contains": # string contains + mask &= data[field].str.contains(str(value), na=False) + elif operator == "not_contains": # string doesn't contain + mask &= ~data[field].str.contains(str(value), na=False) + elif operator == "startswith": # string starts with + mask &= data[field].str.startswith(str(value), na=False) + elif operator == "endswith": # string ends with + mask &= data[field].str.endswith(str(value), na=False) + elif operator == "is_null": # field is null/NaN + mask &= data[field].isna() + elif operator == "not_null": # field is not null/NaN + mask &= data[field].notna() + elif operator == "is_empty": # field is empty string + mask &= (data[field] == "") + elif operator == "not_empty": # field is not empty string + mask &= (data[field] != "") + elif operator == "regex": # regex match + mask &= data[field].str.match(str(value), na=False) + elif operator == "length_eq": # string length equals + mask &= (data[field].str.len() == value) + elif operator == "length_gt": # string length greater than + mask &= (data[field].str.len() > value) + elif operator == "length_lt": # string length less than + mask &= (data[field].str.len() < value) + else: + raise ValueError(f"Unsupported operator: {operator}") + + # Return rows that match the conditions (keep=True) or don't match (keep=False) + keep_matching = self.args.get("keep_matching", True) + if keep_matching: + return data[mask] + else: + return data[~mask] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/model.py b/document-graph/src/graphrag_toolkit/document_graph/model.py new file mode 100644 index 00000000..a3006ee6 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/model.py @@ -0,0 +1,67 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Document Graph Data Models. + +This module defines Pydantic models for graph nodes and edges. +""" + +import logging +from pydantic import BaseModel +from typing import Dict + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +class NodeModel(BaseModel): + """Pydantic model for graph nodes. + + Attributes: + id: Unique identifier for the node + type: Node type/label + properties: Dictionary of node properties + + Examples: + >>> node = NodeModel( + ... id="doc-123", + ... type="Document", + ... properties={"title": "Example Document", "created_date": "2025-07-25"} + ... ) + >>> node.id + 'doc-123' + >>> node.type + 'Document' + >>> node.properties["title"] + 'Example Document' + """ + id: str + type: str + properties: Dict[str, object] + +class EdgeModel(BaseModel): + """Pydantic model for graph edges. + + Attributes: + source: Source node identifier + target: Target node identifier + type: Edge type/relationship + properties: Dictionary of edge properties + + Examples: + >>> edge = EdgeModel( + ... source="doc-123", + ... target="doc-456", + ... type="REFERENCES", + ... properties={"weight": 0.8, "created_date": "2025-07-25"} + ... ) + >>> edge.source + 'doc-123' + >>> edge.target + 'doc-456' + >>> edge.type + 'REFERENCES' + """ + source: str + target: str + type: str + properties: Dict[str, object] diff --git a/document-graph/src/graphrag_toolkit/document_graph/model_elements.py b/document-graph/src/graphrag_toolkit/document_graph/model_elements.py new file mode 100644 index 00000000..a14e95bb --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/model_elements.py @@ -0,0 +1,163 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""[Module Name] Module for Document Graph Plugin System. + +This module [provides/defines/implements] [brief description of module purpose]. + +The module [includes/contains/offers] [key components or features]: +- [Component/Feature 1]: [Brief description] +- [Component/Feature 2]: [Brief description] +- [Component/Feature 3]: [Brief description] + +[Additional context about design principles, architecture, or implementation details] + +Usage: + # [Example usage code] + [explanation of example] +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +"""Graph Elements Module for Document Graph Storage. + +This module defines the fundamental data structures used to represent graph elements +(nodes and edges) in the document graph storage system. These classes serve as the +core data transfer objects between different components of the storage layer. + +The module provides two main dataclasses: +- Node: Represents a vertex in the graph with an ID, labels, and properties +- Edge: Represents a relationship between two nodes with source and target IDs, + a label, and properties + +These classes are designed to be simple, serializable data containers that can be +easily converted to and from various graph database formats. + +Usage: + + # Create a node + document_node = Node( + id="doc123", + labels=["Document"], + properties={"title": "Example Document", "created_at": "2023-01-01"} + ) + + # Create an edge + contains_edge = Edge( + id="edge456", + source_id="doc123", + target_id="section789", + label="CONTAINS", + properties={"order": 1} + ) +""" + +from typing import Dict, Any, Optional +from dataclasses import dataclass + + +@dataclass +class Node: + """ + Represents a node (vertex) in the document graph. + + A node is a fundamental element in a graph that represents an entity or concept. + In the document graph context, nodes typically represent documents, sections, + paragraphs, or other document elements. + + This class is implemented as a dataclass for simplicity and serialization support. + It provides default initialization for labels and properties if they are not provided. + + Attributes: + id (str): Unique identifier for the node. + labels (list): List of labels/types associated with this node. + Examples: ["Document", "Section", "Paragraph"]. + Defaults to an empty list if not provided. + properties (Dict[str, Any]): Dictionary of key-value properties for the node. + Examples: {"title": "Introduction", "created_at": "2023-01-01"}. + Defaults to an empty dictionary if not provided. + + Example: + >>> node = Node(id="doc123", labels=["Document"], properties={"title": "Example"}) + >>> print(node.id, node.labels) + doc123 ['Document'] + """ + + id: str + labels: list = None + properties: Dict[str, Any] = None + + def __post_init__(self): + """ + Initialize default values for properties and labels after instance creation. + + This method is automatically called by the dataclass after the object is created. + It ensures that properties and labels are never None by setting them to empty + collections when they are not provided during initialization. + + - Sets properties to an empty dict if None + - Sets labels to an empty list if None + """ + if self.properties is None: + self.properties = {} + if self.labels is None: + self.labels = [] + + +@dataclass +class Edge: + """ + Represents an edge (relationship) in the document graph. + + An edge is a connection between two nodes in a graph that represents a relationship + or association. In the document graph context, edges typically represent relationships + like "CONTAINS", "REFERENCES", "SIMILAR_TO", etc. between document elements. + + This class is implemented as a dataclass for simplicity and serialization support. + It provides default initialization for properties if they are not provided. + + Attributes: + id (str): Unique identifier for the edge. + source_id (str): Identifier of the source/from node. + target_id (str): Identifier of the target/to node. + label (str): Type of relationship this edge represents. + Examples: "CONTAINS", "REFERENCES", "SIMILAR_TO". + properties (Dict[str, Any]): Dictionary of key-value properties for the edge. + Examples: {"weight": 0.8, "created_at": "2023-01-01"}. + Defaults to an empty dictionary if not provided. + + Example: + >>> edge = Edge( + ... id="edge123", + ... source_id="doc1", + ... target_id="section2", + ... label="CONTAINS", + ... properties={"order": 1} + ... ) + >>> print(edge.source_id, edge.label, edge.target_id) + doc1 CONTAINS section2 + """ + + id: str + source_id: str + target_id: str + label: str + properties: Dict[str, Any] = None + + def __post_init__(self): + """ + Initialize default values for properties after instance creation. + + This method is automatically called by the dataclass after the object is created. + It ensures that properties is never None by setting it to an empty + dictionary when it is not provided during initialization. + + - Sets properties to an empty dict if None + """ + if self.properties is None: + self.properties = {} \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/__init__.py new file mode 100644 index 00000000..635a9759 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/__init__.py @@ -0,0 +1 @@ +"""Pipeline — extract, transform, and load pipeline components.""" diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/__init__.py new file mode 100644 index 00000000..ff034b0b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/__init__.py @@ -0,0 +1 @@ +"""Extract — data extraction providers for various file formats.""" diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_base.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_base.py new file mode 100644 index 00000000..7191b0dd --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_base.py @@ -0,0 +1,77 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Extract Provider Base module for document graph operations.""" + +import logging +from abc import ABC, abstractmethod +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_config import ExtractProviderConfig +from graphrag_toolkit.document_graph.config import DocumentGraphConfig +from graphrag_toolkit.document_graph.pipeline.extract.extraction_result import ExtractionResult +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + + +class ExtractProvider(ABC): + """ + Abstract base class for an extract provider. + + This class serves as a blueprint for creating extract providers that perform + ETL (Extract, Transform, Load) operations on structured document nodes. + The main purpose of this class is to define the interface and shared + functionality for concrete implementations. Subclasses are required to + implement the `extract` method. + + Attributes: + config: Extract provider configuration. + aws_config: AWS configuration for cloud operations. + """ + + def __init__(self, config: ExtractProviderConfig, aws_config: DocumentGraphConfig): + """Initialize extract provider with configuration. + + Args: + config: Extract provider configuration + aws_config: AWS configuration for cloud operations + """ + self.config = config + self.aws_config = aws_config + logger.debug(f"Initialized {self.__class__.__name__} with config: {config.type}") + + @abstractmethod + def extract(self, source: str, **kwargs) -> ExtractionResult: + """ + Represents an abstract base class for data extraction. + + This class defines an interface for extracting data from a given source string. + Subclasses are required to implement the `extract` method. + + Methods: + extract(source: str, **kwargs) -> ExtractionResult + + Abstract method intended to be implemented by subclasses. This method + extracts specific data from a provided string source using potential + additional arguments. + + Raises: + NotImplementedError: If the method is not implemented in a deriving class. + + """ + pass + + def get_etl_schema(self) -> ETLSchema: + """ + Returns the ETL schema for the current configuration. + + The method retrieves the ETL schema object by loading it from the + config if it hasn't been loaded already. + + Returns: + ETLSchema: The schema object representing the ETL schema. + """ + return self.config.load_schema_if_needed() \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_config.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_config.py new file mode 100644 index 00000000..67568c22 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_config.py @@ -0,0 +1,257 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Extract Provider Config module for document graph operations.""" + +import logging +from pydantic import BaseModel, Field +from typing import Literal, Dict, Any, Optional, TYPE_CHECKING + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema, ExtractConfig +from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import SchemaProviderFactory + +if TYPE_CHECKING: + pass + + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + + +class ExtractProviderConfig(BaseModel): + """ + Represents a configuration for an Extract Provider. + + This class provides a detailed configuration structure for extract providers, + facilitating integration with document sources of various types. It serves + as a foundational element for managing the extraction process, including + the ETL schema, pipeline arguments, transformers, and factory integrations. + + Attributes: + type (str): Type of extract provider. It can be one of the following values: + "csv", "json", "parquet", or "excel". + source (str): Path or URI to the document source. + document_id (str): Document ID used in node construction. + etl_schema (Optional[ETLSchema]): Complete ETL schema configuration. + schema_provider_config (Optional[Dict[str, Any]]): Schema provider configuration. + args (Dict[str, Any]): Arguments for Pandas file reader, such as separators + or encoding specifications. + pipeline_args (Dict[str, Any]): Configurations affecting pipeline behavior, + such as logging settings or skipping conditions. + transformers (Optional[list]): List of transformer configurations. + factory (Optional[S3ArtifactFactory]): Factory used to load sources or pipelines. + model_config (dict): Arbitrary types allowance configuration. + + Methods: + parameters -> Dict[str, Any]: + Retrieves backward-compatible alias for 'args'. + + schema_path -> Optional[str]: + Provides backward-compatible alias for the schema path. + + get_extract_config -> ExtractConfig: + Derives the extract configuration from ETL schema or defaults to the + type configuration. + + load_schema_if_needed -> Optional[ETLSchema]: + Loads ETL schema configuration dynamically if not already available. + + from_factory_source(cls, factory, source_name, pipeline_name=None, **kwargs) -> ExtractProviderConfig: + Creates configuration using a factory source artifact with optional + pipeline arguments. + + from_factory_pipeline(cls, factory, pipeline_name) -> ExtractProviderConfig: + Creates configuration using a factory pipeline artifact. + """ + type: Literal["csv", "json", "parquet", "excel"] = Field(..., description="Type of extract provider") + source: str = Field(..., description="Path or URI to the document source") + document_id: str = Field(..., description="Document ID used in node construction") + + # Optional schema integration + etl_schema: Optional[ETLSchema] = Field(None, description="Complete ETL schema configuration") + schema_provider_config: Optional[Dict[str, Any]] = Field(None, description="Schema provider configuration") + + # Clean separation of arguments + args: Dict[str, Any] = Field(default_factory=dict, + description="Arguments for Pandas file reader (e.g., sep, encoding, on_bad_lines)") + pipeline_args: Dict[str, Any] = Field(default_factory=dict, + description="Pipeline behavior (e.g., logging, skipping)") + transformers: Optional[list] = Field(default_factory=list, + description="List of transformer configurations") + + # Factory support + factory: Optional[Any] = Field(None, description="Artifact factory for loading sources/pipelines") + + model_config = {"arbitrary_types_allowed": True} + + @property + def parameters(self) -> Dict[str, Any]: + """ + Provides access to the parameters of an object. + + This property allows retrieval of the object's parameters in the form of a dictionary. + + Returns: + Dict[str, Any]: A dictionary containing the parameters of the object. + """ + return self.args + + @property + def schema_path(self) -> Optional[str]: + """ + Gets the schema path of the object. + + The `schema_path` property is intended to provide the schema location + for this specific instance, if available. If no schema path is available, + it will return `None`. + + Returns: + Optional[str]: The path to the schema or `None` if no schema path exists. + """ + return None + + def get_extract_config(self) -> ExtractConfig: + """ + Gets the extraction configuration. + + This method retrieves the extraction configuration based on the current ETL schema. + If no specific ETL schema is defined, it returns a default minimal configuration. + + Returns: + ExtractConfig: The extraction configuration object. + + """ + if self.etl_schema: + return self.etl_schema.extract + + # Default minimal config + return ExtractConfig( + source_type="file", + file_type=self.type + ) + + def load_schema_if_needed(self) -> Optional[ETLSchema]: + """ + Loads the ETL schema if it is not already loaded. + + This method attempts to load and return the ETL schema. If the schema is already + available in the `etl_schema` attribute, it will return that. If not, it will try to + load the schema using the provided schema provider configuration. In the event of an + error while loading the schema, it logs a warning and returns None. + + Raises: + Logs a warning if the schema loading process fails due to an invalid configuration + or other errors. + + Returns: + Optional[ETLSchema]: The loaded ETL schema, or None if the schema could not be loaded. + """ + if self.etl_schema: + return self.etl_schema + + if self.schema_provider_config: + try: + schema_provider = SchemaProviderFactory.create(self.schema_provider_config) + return schema_provider.load_schema() + except Exception as e: + # ErrorHandler removed + raise ValueError(f"Schema load failed: {e}") from e + + + + + logger.warning(f"Failed to load schema from provider: {error.message}") + return None + + return None + + @classmethod + def from_factory_source(cls, factory: Any, source_name: str, pipeline_name: str = None, **kwargs) -> "ExtractProviderConfig": + """ + Creates an instance of ExtractProviderConfig from a factory source. + + Detailed Summary: + This class method initializes an ExtractProviderConfig instance using a provided + S3ArtifactFactory and a source name. It pulls metadata and arguments from the specified + source artifact or associated pipeline if applicable. Additional parameters can be + customized using keyword arguments, with a fallback to source-defined attributes. + + Parameters: + factory (S3ArtifactFactory): The factory responsible for managing artifact storage + and retrieval. + source_name (str): The name of the source artifact within the factory to load data from. + pipeline_name (str, optional): The name of the pipeline artifact to extract additional + configuration arguments, if applicable. + **kwargs: Additional keyword arguments for custom configuration, with document_id + being extracted separately, defaulting to the source's name if not provided. + + Raises: + ValueError: If the specified source artifact cannot be found within the provided factory. + + Returns: + ExtractProviderConfig: An initialized configuration object for the extract provider. + """ + source_artifact = factory.load_source(source_name) + if not source_artifact: + raise ValueError(f"Source artifact '{source_name}' not found in factory") + + # Load args from pipeline or source metadata + args = {} + if pipeline_name: + pipeline_artifact = factory.load_pipeline(pipeline_name) + if pipeline_artifact: + args = pipeline_artifact.pipeline_config.get('source', {}).get('args', {}) + + # Fallback to source metadata processing_args + if not args and source_artifact.metadata: + args = source_artifact.metadata.get('processing_args', {}) + + # Extract document_id from kwargs to avoid duplicate + document_id = kwargs.pop('document_id', source_artifact.name) + + return cls( + factory=factory, + type=source_artifact.source_type, + source=f"factory://{source_name}", + document_id=document_id, + args=args, + **kwargs + ) + + @classmethod + def from_factory_pipeline(cls, factory: Any, pipeline_name: str) -> "ExtractProviderConfig": + """ + Creates an `ExtractProviderConfig` instance from a factory and pipeline name. + + This class method allows loading a pipeline artifact using the provided factory + and pipeline name, extracting configuration data, and constructing an instance + of `ExtractProviderConfig` with it. + + Parameters: + factory (S3ArtifactFactory): The factory instance used to load the pipeline artifact. + pipeline_name (str): The identifier of the pipeline whose artifact needs to + be loaded. + + Returns: + ExtractProviderConfig: A new instance of `ExtractProviderConfig` initialized with the + extracted configuration data. + + Raises: + ValueError: If the specified pipeline artifact is not found in the factory. + """ + pipeline_artifact = factory.load_pipeline(pipeline_name) + if not pipeline_artifact: + raise ValueError(f"Pipeline artifact '{pipeline_name}' not found in factory") + + config_data = pipeline_artifact.pipeline_config + return cls( + factory=factory, + type=config_data['source']['type'], + source=config_data['source']['path'], + document_id=config_data['extract']['document_id'], + args=config_data['source'].get('args', {}), + **config_data.get('extract', {}) + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_csv.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_csv.py new file mode 100644 index 00000000..de3d2ced --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_csv.py @@ -0,0 +1,413 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Extract Provider Csv module for document graph operations.""" + +import logging +from pathlib import Path +from typing import List, Tuple + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 +from graphrag_toolkit.document_graph.config import DocumentGraphConfig +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_base import ExtractProvider +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_config import ExtractProviderConfig +from graphrag_toolkit.document_graph.schema.providers.csv_schema_provider import CSVSchemaProvider +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema +from graphrag_toolkit.document_graph.pipeline.extract.utils.dataframe_reader import DataFrameReader +# PathResolver removed + + + +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig + + + +class CSVExtractProvider(ExtractProvider): + """ + Provide functionality to extract data from a CSV file based on a specified ETL configuration. + + This class is responsible for handling CSV files, supporting local file paths or S3 URIs, + and extracting structured information while adhering to ETL schema configurations. It can + apply data transformations, infer schema, save skipped rows if necessary, and stream process + data when ingestors are provided. + + Attributes: + config (ExtractProviderConfig): Configuration details for the extraction provider. + aws_config (DocumentGraphConfig): AWS settings for operations like S3 file handling. + """ + + def __init__(self, config: ExtractProviderConfig, aws_config: DocumentGraphConfig): + """ + Initializes an instance of CSVExtractProvider. + + The initializer sets up the CSVExtractProvider with the provided configuration + objects. It logs the configurations and calls the parent class initializer + to ensure proper initialization. + + Parameters: + config (ExtractProviderConfig): Configuration specific to the extraction provider. + aws_config (DocumentGraphConfig): Configuration related to AWS Document Graph services. + """ + logger.debug(f"Initializing CSVExtractProvider with config: {vars(config)}") + logger.debug(f"AWS config: {vars(aws_config)}") + super().__init__(config, aws_config) + logger.debug(f"[CSVExtractProvider] args: {self.config.args}") + logger.debug(f"[CSVExtractProvider] pipeline_args: {self.config.pipeline_args}") + + def extract(self, source: str, ingestors=None, **kwargs) -> None: + """ + Extracts data from a CSV file at the given source and processes it into a structured format. + The extraction process includes validating the file's existence, handling any parsing errors, + and applying configurations for data transformation. Optionally, it integrates with ETL schemas, + infers the schema from the data, and generates document nodes for downstream processes. + + The function achieves robust CSV parsing and ensures compatibility with specific configurations, + including handling row limits and bad line strategies. + + Parameters: + source (str): The source path of the CSV file to extract data from. + ingestors: Optional list of ingestors to process data in a streaming fashion. + kwargs: Arbitrary keyword arguments for additional extraction configurations. + + Raises: + Exception: If errors occur during file parsing or data inference. + + Returns: + None + """ + logger.debug(f"[extract] Entry point kwargs: {kwargs}") + logger.debug(f"[extract] config.args received: {self.config.args}") + logger.debug(f"[extract] config.pipeline_args received: {self.config.pipeline_args}") + + logger.info(f"Starting CSV extraction from source: {source}") + logger.debug(f"extract called with kwargs: {kwargs}") + resolved_path = source + resolved_path = Path(resolved_path) + logger.debug(f"Resolved path: {resolved_path}") + + # Use centralized error handling + # ErrorHandler removed + assert Path(resolved_path).exists(), f"File not found: {resolved_path}" + + try: + # Filter out custom parameters before passing to pandas + logger.debug(f"Original config.args: {self.config.args}") + pandas_args = { + k: v for k, v in self.config.args.items() if k != 'output_skip_lines' + } + + # Add header config from transformers but never row limits + logger.debug(f"Transformers config: {self.config.transformers}") + if self.config.transformers: + for transformer in self.config.transformers: + if transformer.get('type') == 'row_to_node': + transformer_config = transformer.get('config', {}) + # Keep header override if explicitly provided + if 'header' in transformer_config: + pandas_args['header'] = transformer_config['header'] + + # Never let Transformers set row limits for Extract + pandas_args.pop('nrows', None) + pandas_args.pop('skiprows', None) + + logger.info(f"Final pandas_args: {pandas_args}") + + # Ensure pandas uses the correct engine for robust CSV parsing + if pandas_args.get("on_bad_lines") == "skip": + current_engine = pandas_args.get("engine", "c").lower() + if current_engine != "python": + logger.warning( + "on_bad_lines='skip' requires engine='python'; overriding engine to ensure compatibility." + ) + pandas_args["engine"] = "python" + logger.debug(f"Overridden pandas_args with engine='python': {pandas_args}") + + # Use streaming if ingestors provided + logger.debug(f"Ingestors parameter: {ingestors}") + logger.debug(f"Ingestors type: {type(ingestors)}") + if ingestors: + logger.info("Using streaming extraction with ingestors") + df = self._stream_with_ingestors(resolved_path, pandas_args, ingestors) + else: + logger.info("No ingestors provided, using regular extraction") + logger.debug(f"Reading CSV with final pandas args: {pandas_args}") + df = DataFrameReader(path=resolved_path, **pandas_args).read() + + logger.info(f"Successfully read CSV with shape: {df.shape}") + logger.debug(f"DataFrame dtypes: {df.dtypes.to_dict()}") + logger.debug(f"DataFrame head(2): {df.head(2).to_dict(orient='records')}") + + # Save skipped lines if requested + if self.config.args.get('output_skip_lines', False) and self.config.args.get('on_bad_lines') == 'skip': + self._save_skipped_lines(resolved_path) + + except Exception as e: + logger.error(f"Error during CSV parsing: {e}", exc_info=True) + raise RuntimeError(f"CSV parse error: {resolved_path}: {e}") from e + + if df.empty: + strict_mode = kwargs.get("strict", True) + logger.warning(f"DataFrame is empty, strict_mode: {strict_mode}") + error = ValueError(f"Empty data: {resolved_path}") if strict_mode else None + if error: + raise error + return [], {} + + logger.debug("Inferring schema from DataFrame") + schema = {col: str(df[col].dtype) for col in df.columns} + + logger.debug(f"Inferred schema with {len(schema)} fields: {schema}") + + logger.debug(f"Converting {len(df)} rows to document nodes") + nodes = [] + for i, row in df.iterrows(): + content = ", ".join(f"{k}: {v}" for k, v in row.items()) + node = dict( + id=f"{self.config.document_id}-row-{i}", + type="Row", + content=content, + metadata={"row_index": str(i)}) + nodes.append(node) + if nodes: + logger.debug(f"Sample node: id={nodes[0]['id']}, type={nodes[0]['type']}, metadata={nodes[0]['metadata']}") + + # Load ETL schema if configured + etl_schema = self.config.load_schema_if_needed() + if not etl_schema: + # Dynamically generate schema_id from document_id + document_id = self.config.document_id + schema_id = f"{document_id}-schema" if document_id else None + logger.debug(f"Generated schema_id: {schema_id}") + + # Separate args for schema discovery (exclude batching parameters) + schema_args = {k: v for k, v in pandas_args.items() + if k not in ['skiprows', 'nrows']} + + # Create schema from CSV file using SchemaProviderConfig with args + schema_config = SchemaProviderConfig( + type="csv", + schema_id=schema_id, + connection_config={"path": str(resolved_path), "args": schema_args} + ) + logger.debug(f"SchemaProviderConfig: {vars(schema_config)}") + csv_schema_provider = CSVSchemaProvider(schema_config) + etl_schema = csv_schema_provider.load_schema() + logger.info(f"Generated ETL schema from CSV: {etl_schema.schema_id}") + logger.debug(f"ETL schema details: {vars(etl_schema)}") + + # Create extraction result with ETL schema integration + from graphrag_toolkit.document_graph.pipeline.extract.extraction_result import ExtractionResult + + metadata = { + "source": source, + "extraction_type": "csv", + "etl_schema_id": etl_schema.schema_id if etl_schema else None, + "rows_processed": str(len(df)), + "schema_fields": str(len(schema)) + } + logger.debug(f"Extraction metadata: {metadata}") + + result = ExtractionResult( + document_id=self.config.document_id, + extracted_schema=schema, + nodes=nodes, + dataframe=df, + metadata=metadata + ) + + logger.info( + f"CSV extraction completed: {len(nodes)} nodes, ETL schema: {etl_schema.schema_id if etl_schema else 'None'}") + return result + + def _save_skipped_lines(self, csv_path: Path): + """ + Saves lines from a CSV file that were skipped during parsing to a separate file. + + This method processes a given CSV file, identifies lines that do not match the + expected field count based on the file's header, and writes the invalid lines to a + new file. The new file contains information about the skipped lines and their + discrepancies. + + Attributes + ---------- + None + + Parameters + ---------- + csv_path : Path + Path to the CSV file to be checked for skipped lines. + + Raises + ------ + Exception + If an exception occurs during reading, processing, or writing skipped lines, + it logs the warning including an exception traceback. + + Note + ---- + The output file containing the skipped lines is named with a `.skip` suffix before + the original file's extension. + """ + import pandas as pd + import csv + + skip_file = csv_path.with_suffix('.skip' + csv_path.suffix) + logger.debug(f"Checking for skipped lines, output: {skip_file}") + + try: + # Read all lines + with open(csv_path, 'r', encoding='utf-8') as f: + all_lines = f.readlines() + logger.debug(f"Total lines in file: {len(all_lines)}") + + # Read successfully parsed lines + pandas_args = {k: v for k, v in self.config.args.items() if k != 'output_skip_lines'} + logger.debug(f"pandas_args for reading good_df: {pandas_args}") + good_df = pd.read_csv(csv_path, **pandas_args) + logger.debug(f"Good DataFrame shape: {good_df.shape}") + + # Detect delimiter from first line + sniffer = csv.Sniffer() + delimiter = sniffer.sniff(all_lines[0]).delimiter if all_lines else ',' + logger.debug(f"Detected delimiter: '{delimiter}'") + + # Find lines with wrong field count + header_fields = len(all_lines[0].split(delimiter)) if all_lines else 0 + logger.debug(f"Header fields count: {header_fields}") + skipped_lines = [] + + for i, line in enumerate(all_lines[1:], 1): # Skip header + field_count = len(line.split(delimiter)) + if field_count != header_fields: + skipped_lines.append(f"# Line {i + 1}: Expected {header_fields} fields, got {field_count}\n") + skipped_lines.append(line) + + # Save skipped lines if any found + if skipped_lines: + with open(skip_file, 'w', encoding='utf-8') as f: + f.write(f"# Skipped lines from {csv_path.name}\n") + f.write(f"# Header: {all_lines[0]}") + f.writelines(skipped_lines) + logger.info(f"Saved {len(skipped_lines) // 2} skipped lines to: {skip_file}") + else: + logger.debug("No skipped lines detected") + + except Exception as e: + logger.warning(f"Could not save skipped lines: {e}", exc_info=True) + + def _stream_with_ingestors(self, csv_path, pandas_args, ingestors): + """ + Handles reading and processing a CSV file with error tolerance and memory handling. + + The method attempts to read the entire file first, applying custom error handling on bad lines. + If the full read fails, it falls back to reading the file in chunks, leveraging chunked processing + to handle memory constraints. The method integrates the use of ingestors for data filtering + and transformation, which is applied to data either fully or in chunks. + + Attributes: + csv_path (str): The path to the CSV file to be read. + pandas_args (dict): Additional arguments passed to pandas read_csv for configuration. + ingestors: An object or utility providing the `execute` method for applying transformations or + filtering to the loaded data. + + Parameters and Configuration: + csv_path: The file path of the CSV to process. + pandas_args: Contains any additional arguments needed for `pd.read_csv`. The function utilizes + these for fallback operations if required. + ingestors: Utilized to process and filter the read DataFrame. + + Raises: + Various exceptions may be generated on initial full read-pandas arguments under directly/during + ```""" + import pandas as pd + + # First try: Load entire file with robust error handling + logger.info(f"Attempting full file read: {csv_path}") + try: + # Custom bad line handler for debugging + def bad_line_handler(line): + """ + Handles lines considered as "bad" by logging a warning. + + This function is used to process lines that are identified as problematic + or invalid for further operations. It logs a warning with detailed + information about the bad line and then skips it by returning None. + + Parameters: + line (str): The input line identified as bad. The first 200 characters + of this line will be included in the warning message for debugging purposes. + + Returns: + None: Indicates that the bad line has been skipped. + """ + logger.warning(f"Skipping bad line: {repr(line[:200])}...") + return None + + # Read entire file with maximum error tolerance + full_df = pd.read_csv( + csv_path, + encoding='utf-8', + engine='python', + on_bad_lines=bad_line_handler, + low_memory=False, + quoting=1, # QUOTE_ALL + skipinitialspace=True + ) + logger.info(f"Full file loaded: {len(full_df)} rows") + + # Apply ingestors to full dataset + filtered_df = ingestors.execute(full_df) + logger.info(f"Ingestors applied: {len(full_df)} -> {len(filtered_df)} rows") + return filtered_df + + except Exception as e: + logger.warning(f"Full file read failed: {e}, trying chunked approach") + + # Fallback: Try chunked reading + chunk_size = 10000 + filtered_chunks = [] + chunk_args = {'encoding': pandas_args.get('encoding', 'utf-8')} + + logger.info(f"Pandas chunked reading: {csv_path}") + logger.info(f"Chunk args: {chunk_args}") + + try: + chunk_reader = pd.read_csv(csv_path, chunksize=chunk_size, iterator=True, **chunk_args) + + chunk_count = 0 + total_rows_read = 0 + + while True: + try: + chunk = next(chunk_reader) + chunk_count += 1 + total_rows_read += len(chunk) + logger.info(f"Processing chunk {chunk_count}: {len(chunk)} rows (total: {total_rows_read})") + + filtered_chunk = ingestors.execute(chunk) + if not filtered_chunk.empty: + filtered_chunks.append(filtered_chunk) + + except StopIteration: + logger.info("Reached end of file - no more chunks") + break + + logger.info(f"Pandas chunked reading completed: {chunk_count} chunks, {total_rows_read} total rows") + + if filtered_chunks: + result = pd.concat(filtered_chunks, ignore_index=True) + logger.info(f"Combined {len(filtered_chunks)} chunks -> {len(result)} rows") + return result + else: + return pd.DataFrame() + + except Exception as e: + logger.error(f"Pandas chunked reading failed: {e}", exc_info=True) + logger.warning("Falling back to pure pandas read_csv") + return pd.read_csv(csv_path, **pandas_args) + + + diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_excel.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_excel.py new file mode 100644 index 00000000..da615c4b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_excel.py @@ -0,0 +1,165 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Extract Provider Excel module for document graph operations.""" + +import logging +from pathlib import Path +from typing import List + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 +from graphrag_toolkit.document_graph.config import DocumentGraphConfig +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_base import ExtractProvider +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_config import ExtractProviderConfig +from graphrag_toolkit.document_graph.pipeline.extract.utils.dataframe_reader import DataFrameReader +# PathResolver removed +from graphrag_toolkit.document_graph.schema.document_node import DocumentNode +from graphrag_toolkit.document_graph.schema.schema_inference_utils import SchemaInferenceEngine +from graphrag_toolkit.document_graph.schema.schema_io import SchemaWriter +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.providers.excel_schema_provider import ExcelSchemaProvider + + + +class ExcelExtractProvider(ExtractProvider): + """ + Provides functionality for extracting, processing, and analyzing Excel data. + + This class is responsible for managing the extraction of Excel data into structured formats including + DataFrames and schema-based nodes. It supports resolving file paths, reading Excel data, applying + ingestors for preprocessing, inferring schemas, saving schemas with mock data, and generating document + nodes based on the content of each row. It also ensures handling of errors and validation during the + entire extraction process. + + Attributes: + config (ExtractProviderConfig): Configuration object for the extraction provider that stores + various extraction-related settings and details. + aws_config (DocumentGraphConfig): AWS-specific configuration for managing data storage, schema + saving, and other AWS-related operations. + """ + def __init__(self, config: ExtractProviderConfig, aws_config: DocumentGraphConfig): + """ + [Brief one-line description of the function's purpose] + + [Detailed explanation of what the function does and how it works] + + Args: + [param_name] ([param_type]): [Description of parameter] + [param_name] ([param_type], optional): [Description of parameter]. Defaults to [default_value]. + + Returns: + [return_type]: [Description of return value] + + Raises: + [exception_type]: [Conditions under which exception is raised] + """ + super().__init__(config, aws_config) + + def extract(self, source: str, ingestors=None, **kwargs): + """ + Extracts data from an Excel file, processes it, and returns a comprehensive + extraction result, including inferred schema, document nodes, and metadata. + + This method performs several actions, including resolving the file path, + verifying the file's existence, reading the Excel file into a DataFrame, + and applying specified ingestors if provided. It further processes the data + to infer its schema, saves the schema with mock data, and creates document nodes + from its rows. If an ETL schema is configured, it uses the schema; otherwise, it + generates an ETL schema dynamically. + + Parameters: + source (str): Source path to the Excel file to be extracted. + ingestors (optional): An optional ingestors instance used to process the DataFrame. + **kwargs: Arbitrary keyword arguments, e.g., `strict` for handling empty data. + + Raises: + Exception: If an error occurs while reading the Excel file or applying the + ingestors. + Error: When the DataFrame is empty, depending on the strict mode configuration. + + Returns: + ExtractionResult: An object encapsulating the document ID, extracted schema, + document nodes, processed DataFrame, and associated metadata. + + """ + logger.info(f"Starting Excel extraction from source: {source}") + resolved_path = PathResolver(self.aws_config).resolve(source) + resolved_path = Path(resolved_path) + logger.debug(f"Resolved path: {resolved_path}") + + # ErrorHandler removed + ErrorHandler.check_file_exists(str(resolved_path), "Excel extraction") + + try: + logger.debug(f"Reading Excel with args: {self.config.args}") + df = DataFrameReader(path=resolved_path, **self.config.args).read() + logger.info(f"Successfully read Excel with shape: {df.shape}") + + # Apply ingestors if provided + if ingestors: + logger.info("Applying ingestors to Excel data") + df = ingestors.execute(df) + except Exception as e: + raise ErrorHandler.file_parse_error(str(resolved_path), "excel", e) from e + + if df.empty: + strict_mode = kwargs.get("strict", True) + error = ErrorHandler.empty_data(str(resolved_path), strict_mode) + if error: + raise error + return [], {} + + logger.debug("Inferring schema from DataFrame") + schema_engine = SchemaInferenceEngine(df) + schema = schema_engine.infer_basic_schema() + logger.debug(f"Inferred schema with {len(schema)} fields") + + logger.debug("Saving schema with mock data") + schema_path = SchemaWriter(self.aws_config).save_schema_with_mock_data( + schema_dict=schema, + df=df, + source=source, + output_dir=Path("schemas"), + document_id=self.config.document_id, + ) + logger.info(f"Schema saved to: {schema_path}") + + logger.debug(f"Converting {len(df)} rows to document nodes") + nodes = [] + for i, row in df.iterrows(): + content = ", ".join(f"{k}: {v}" for k, v in row.items()) + nodes.append(DocumentNode( + id=f"{self.config.document_id}-row-{i}", + type="Row", + content=content, + metadata={"row_index": str(i)} + )) + + # Load ETL schema if provided in config, otherwise fall back to dynamic schema generation + etl_schema = self.config.load_schema_if_needed() + if not etl_schema: + schema_config = SchemaProviderConfig( + type="excel", + connection_config={"path": str(resolved_path)}, + ) + excel_schema_provider = ExcelSchemaProvider(schema_config) + etl_schema = excel_schema_provider.load_schema() + logger.info(f"Generated ETL schema from Excel: {etl_schema.schema_id}") + + from graphrag_toolkit.document_graph.pipeline.extract.extraction_result import ExtractionResult + return ExtractionResult( + document_id=self.config.document_id, + extracted_schema=schema, + nodes=nodes, + dataframe=df, + metadata={ + "source": source, + "extraction_type": "excel", + "etl_schema_id": etl_schema.schema_id if etl_schema else None, + "rows_processed": str(len(df)), + "schema_fields": str(len(schema)) + } + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_factory.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_factory.py new file mode 100644 index 00000000..4c8c9b70 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_factory.py @@ -0,0 +1,77 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Extract Provider Factory module for document graph operations.""" + +import logging +from typing import Type + +from .extract_provider_base import ExtractProvider +from .extract_provider_config import ExtractProviderConfig +from .extract_provider_registry import extract_provider_registry +from ...document_graph_config import DocumentGraphConfig + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + + +class ExtractProviderFactory: + """ + Factory class for creating and retrieving extract provider instances. + + This class provides a mechanism for creating extract providers based on their + configuration and type. It utilizes a registry of available providers and + performs validation to ensure the requested provider type is supported. + The created provider is configured according to the supplied configurations. + """ + + @classmethod + def get_provider( + cls, + config: ExtractProviderConfig, + aws_config: DocumentGraphConfig + ) -> ExtractProvider: + """ + Creates and returns an instance of the appropriate `ExtractProvider` subclass based on the given + configuration and type. The provider's class is determined dynamically by looking up the + `extract_provider_registry` with the given type. If the type is not registered, an appropriate + validation error is raised. Logs detailed debug and informational messages for the creation process. + + Args: + config (ExtractProviderConfig): Configuration object specifying the type of extract provider + to create and any additional provider-specific configuration details. + aws_config (DocumentGraphConfig): Configuration object containing AWS-related settings for + document processing. + + Raises: + ErrorHandler.ValidationError: If the given `provider_type` is not recognized or not registered + in the `extract_provider_registry`. + + Returns: + ExtractProvider: An instance of the dynamically determined subclass of `ExtractProvider`. + """ + provider_type = config.type + logger.debug(f"Creating extract provider: type={provider_type}, config={config.dict()}") + logger.debug(f"Looking up provider for type: {provider_type}") + + try: + provider_cls: Type[ExtractProvider] = extract_provider_registry.get(provider_type) + logger.debug(f"Found provider class: {provider_cls.__name__}") + except KeyError: + from ...shared.error_handler import ErrorHandler + available_types = extract_provider_registry.list_providers() + logger.error(f"Unknown extract provider type: {provider_type}. Available: {available_types}") + raise ErrorHandler.validation_error( + "provider_type", + provider_type, + f"one of {available_types}" + ) + + logger.debug(f"Instantiating provider: {provider_cls.__name__}") + logger.info(f"Creating extract provider of type: {provider_type}") + provider = provider_cls(config=config, aws_config=aws_config) + logger.debug(f"Successfully created provider: {provider.__class__.__name__}") + return provider diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_json.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_json.py new file mode 100644 index 00000000..8c843a12 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_json.py @@ -0,0 +1,235 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Extract Provider Json module for document graph operations.""" + +import logging +from pathlib import Path +from typing import List + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 +from graphrag_toolkit.document_graph.config import DocumentGraphConfig +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_base import ExtractProvider +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_config import ExtractProviderConfig +from graphrag_toolkit.document_graph.pipeline.extract.utils.dataframe_reader import DataFrameReader +# PathResolver removed +from graphrag_toolkit.document_graph.schema.document_node import DocumentNode +from graphrag_toolkit.document_graph.schema.schema_inference_utils import SchemaInferenceEngine +from graphrag_toolkit.document_graph.schema.schema_io import SchemaWriter +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.providers.json_schema_provider import JSONSchemaProvider + + + +class JSONExtractProvider(ExtractProvider): + """ + Handles JSON data extraction and schema inference for a given source. + + This class extends the ExtractProvider base class and provides functionality + to read data from a JSON source, infer its schema, and generate document + nodes for downstream processing. It also supports integration with external + schema factories and provides mechanisms for ingesting additional data. + + Attributes: + config (ExtractProviderConfig): Configuration for the extract provider, which includes + document ID, extraction arguments, and schema handling details. + aws_config (DocumentGraphConfig): AWS-specific configuration, including factory settings and + schema management. + + Methods: + __init__(config: ExtractProviderConfig, aws_config: DocumentGraphConfig) + Initializes the JSONExtractProvider with the required configuration details. + extract(source: str, ingestors=None, **kwargs) + Extracts data from the JSON source, infers schema, and creates document nodes. + """ + + def __init__(self, config: ExtractProviderConfig, aws_config: DocumentGraphConfig): + """ + [Brief one-line description of the function's purpose] + + [Detailed explanation of what the function does and how it works] + + Args: + [param_name] ([param_type]): [Description of parameter] + [param_name] ([param_type], optional): [Description of parameter]. Defaults to [default_value]. + + Returns: + [return_type]: [Description of return value] + + Raises: + [exception_type]: [Conditions under which exception is raised] + """ + super().__init__(config, aws_config) + + def extract(self, source: str, ingestors=None, **kwargs): + """ + Extracts data from a JSON source, processes it, and returns an ExtractionResult object. + + This method performs the following actions: + 1. Resolves the path to the given source. + 2. Validates the existence of the file. + 3. Reads the JSON data into a DataFrame. + 4. Applies optional ingestors for data transformation. + 5. Infers and/or generates schemas for the data. + 6. Converts the DataFrame rows into document nodes. + 7. Returns an ExtractionResult containing the extracted information, + nodes, and metadata. + + The method supports strict mode for handling empty data. It also integrates + with schema factories for retrieving or storing schemas. + + Parameters: + source: str + Path to the JSON file to be extracted. + ingestors + Optional ingestors to be applied for data transformation. + **kwargs + Additional options for configuration. Supports "strict" flag + to determine behavior for empty data. + + Returns: + ExtractionResult + Contains the document ID, extracted schema, document nodes, + DataFrame, and metadata. + + Raises: + Exception + If an error occurs during reading, parsing, or schema handling. + ErrorHandler.empty_data + If the data is empty and strict mode is enabled. + """ + logger.info(f"Starting JSON extraction from source: {source}") + resolved_path = PathResolver(self.aws_config).resolve(source) + resolved_path = Path(resolved_path) + logger.debug(f"Resolved path: {resolved_path}") + + # ErrorHandler removed + ErrorHandler.check_file_exists(str(resolved_path), "JSON extraction") + + try: + logger.debug(f"Reading JSON with args: {self.config.args}") + df = DataFrameReader(path=resolved_path, **self.config.args).read() + logger.info(f"Successfully read JSON with shape: {df.shape}") + + # Apply ingestors if provided + if ingestors: + logger.info("Applying ingestors to JSON data") + df = ingestors.execute(df) + except Exception as e: + raise ErrorHandler.file_parse_error(str(resolved_path), "json", e) from e + + if df.empty: + strict_mode = kwargs.get("strict", True) + error = ErrorHandler.empty_data(str(resolved_path), strict_mode) + if error: + raise error + return [], {} + + logger.debug("Inferring schema from DataFrame") + schema_engine = SchemaInferenceEngine(df) + schema = schema_engine.infer_basic_schema() + logger.debug(f"Inferred schema with {len(schema)} fields") + + # Check factory for existing schema first (before local generation) + factory_schema_found = False + if hasattr(self.aws_config, 'factory') and self.aws_config.factory: + factory_schema = self.aws_config.factory.find_schema_for_source( + self.config.document_id, "1.0.0" + ) + if factory_schema: + logger.info(f"Using existing factory schema: {factory_schema.name}@{factory_schema.version}") + factory_schema_found = True + + # Only generate local schema if not found in factory + if not factory_schema_found: + logger.debug("Saving schema with mock data") + schema_path = SchemaWriter(self.aws_config).save_schema_with_mock_data( + schema_dict=schema, + df=df, + source=source, + output_dir=Path("schemas"), + document_id=self.config.document_id, + ) + logger.info(f"Schema saved to: {schema_path}") + else: + logger.info("Skipping local schema generation - using factory schema") + + logger.debug(f"Converting {len(df)} rows to document nodes") + nodes = [] + for i, row in df.iterrows(): + content = ", ".join(f"{k}: {v}" for k, v in row.items()) + nodes.append(DocumentNode( + id=f"{self.config.document_id}-row-{i}", + type="Row", + content=content, + metadata={"row_index": str(i)} + )) + + # Load ETL schema if configured + etl_schema = self.config.load_schema_if_needed() + if not etl_schema: + # Check factory for existing schema first + factory_schema = None + if hasattr(self.aws_config, 'factory') and self.aws_config.factory: + factory_schema = self.aws_config.factory.find_schema_for_source( + self.config.document_id, "1.0.0" + ) + if factory_schema: + logger.info(f"Using factory schema: {factory_schema.name}@{factory_schema.version}") + # Convert factory schema to ETL schema format if needed + etl_schema = factory_schema.schema_data + + # Generate new schema if not found in factory + if not factory_schema: + schema_config = SchemaProviderConfig( + type="json", + connection_config={"path": str(resolved_path)}, + ) + json_schema_provider = JSONSchemaProvider(schema_config) + etl_schema = json_schema_provider.load_schema() + logger.info(f"Generated ETL schema from JSON: {etl_schema.schema_id if hasattr(etl_schema, 'schema_id') else 'new-schema'}") + + # Store generated schema in factory + if hasattr(self.aws_config, 'factory') and self.aws_config.factory: + from graphrag_toolkit.document_graph.factory.schema_artifact import SchemaArtifact + import json + + # Convert ETLSchema to simple dict for JSON serialization + try: + schema_dict = { + "schema_id": getattr(etl_schema, 'schema_id', 'generated'), + "fields": schema, # Use the inferred basic schema + "document_id": self.config.document_id, + "source_type": "json" + } + schema_artifact = SchemaArtifact( + name=f"{self.config.document_id}-schema", + version="1.0.0", + description=f"Auto-generated schema for {self.config.document_id}", + schema_type="etl", + schema_data=schema_dict, + source_ref=f"{self.config.document_id}@1.0.0" + ) + schema_id = self.aws_config.factory.store_schema(schema_artifact) + logger.info(f"Stored schema in factory: {schema_id}") + except Exception as e: + logger.warning(f"Failed to store schema in factory: {e}") + + + from graphrag_toolkit.document_graph.pipeline.extract.extraction_result import ExtractionResult + return ExtractionResult( + document_id=self.config.document_id, + extracted_schema=schema, + nodes=nodes, + dataframe=df, + metadata={ + "source": source, + "extraction_type": "json", + "etl_schema_id": etl_schema.get('schema_id') if isinstance(etl_schema, dict) else (etl_schema.schema_id if etl_schema and hasattr(etl_schema, 'schema_id') else None), + "rows_processed": str(len(df)), + "schema_fields": str(len(schema)) + } + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_parquet.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_parquet.py new file mode 100644 index 00000000..c45d90eb --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_parquet.py @@ -0,0 +1,160 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Extract Provider Parquet module for document graph operations.""" + +import logging +from pathlib import Path +from typing import List + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 +from graphrag_toolkit.document_graph.config import DocumentGraphConfig +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_base import ExtractProvider +from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_config import ExtractProviderConfig +from graphrag_toolkit.document_graph.pipeline.extract.utils.dataframe_reader import DataFrameReader +# PathResolver removed +from graphrag_toolkit.document_graph.schema.document_node import DocumentNode +from graphrag_toolkit.document_graph.schema.schema_inference_utils import SchemaInferenceEngine +from graphrag_toolkit.document_graph.schema.schema_io import SchemaWriter +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.providers.parquet_schema_provider import ParquetSchemaProvider + + + +class ParquetExtractProvider(ExtractProvider): + """ + A provider for extracting data from Parquet files. + + This class is designed to handle the process of reading, processing, and + extracting data from Parquet files. It leverages configurable components + to provide flexibility in handling various data extraction and schema + inference scenarios. The main functionalities include resolving file paths, + reading Parquet files, schema inference, and metadata preparation. + + Attributes: + config (ExtractProviderConfig): Configuration for the extraction process. + aws_config (DocumentGraphConfig): Configuration for AWS-related operations. + """ + + def __init__(self, config: ExtractProviderConfig, aws_config: DocumentGraphConfig): + """ + Initializes an instance of the class. + + This constructor initializes the class instance by calling the parent class + constructor with the provided configuration objects. It ensures that the + necessary configurations are passed for operation. + + Args: + config (ExtractProviderConfig): The configuration object containing + settings for the extract provider. + aws_config (DocumentGraphConfig): The configuration object specific + to the document graph service or AWS-related settings. + """ + super().__init__(config, aws_config) + + def extract(self, source: str, ingestors=None, **kwargs): + """ + Extracts data from a source in Parquet format, processes it, and infers schema and + document nodes. + + This method handles different aspects of the Parquet extraction pipeline including + path resolution, Parquet file reading, schema inference, and document node creation. + Additionally, it supports applying custom ingestors to transform the data. The method + saves the inferred schema along with mock data and returns all results in an + ExtractionResult object. + + Raises: + FileNotFoundError: If the specified Parquet source file does not exist. + RuntimeError: If the schema inference or document node creation fails. + ValueError: If the extracted data is empty and strict mode is enabled. + + Args: + source (str): The source path of the Parquet file. + ingestors (optional): Object containing ingestors to process the DataFrame. Defaults to None. + **kwargs: Additional keyword arguments for configuring extraction behavior. + + Returns: + ExtractionResult: Result of the extraction process, including document nodes, + inferred schema, processed DataFrame, and metadata. + """ + logger.info(f"Starting Parquet extraction from source: {source}") + resolved_path = PathResolver(self.aws_config).resolve(source) + resolved_path = Path(resolved_path) + logger.debug(f"Resolved path: {resolved_path}") + + # ErrorHandler removed + ErrorHandler.check_file_exists(str(resolved_path), "Parquet extraction") + + try: + logger.debug(f"Reading Parquet with args: {self.config.args}") + df = DataFrameReader(path=resolved_path, **self.config.args).read() + logger.info(f"Successfully read Parquet with shape: {df.shape}") + + # Apply ingestors if provided + if ingestors: + logger.info("Applying ingestors to Parquet data") + df = ingestors.execute(df) + except Exception as e: + raise ErrorHandler.csv_parse_error(str(resolved_path), e) from e + + if df.empty: + strict_mode = kwargs.get("strict", True) + error = ErrorHandler.empty_data(str(resolved_path), strict_mode) + if error: + raise error + return [], {} + + logger.debug("Inferring schema from DataFrame") + schema_engine = SchemaInferenceEngine(df) + schema = schema_engine.infer_basic_schema() + logger.debug(f"Inferred schema with {len(schema)} fields") + + logger.debug("Saving schema with mock data") + schema_path = SchemaWriter(self.aws_config).save_schema_with_mock_data( + schema_dict=schema, + df=df, + source=source, + output_dir=Path("schemas"), + document_id=self.config.document_id, + ) + logger.info(f"Schema saved to: {schema_path}") + + logger.debug(f"Converting {len(df)} rows to document nodes") + nodes = [] + for i, row in df.iterrows(): + content = ", ".join(f"{k}: {v}" for k, v in row.items()) + nodes.append(DocumentNode( + id=f"{self.config.document_id}-row-{i}", + type="Row", + content=content, + metadata={"row_index": str(i)} + )) + + # Load ETL schema if configured + etl_schema = self.config.load_schema_if_needed() + if not etl_schema: + schema_config = SchemaProviderConfig( + type="parquet", + connection_config={"path": str(resolved_path)}, + ) + parquet_schema_provider = ParquetSchemaProvider(schema_config) + etl_schema = parquet_schema_provider.load_schema() + logger.info(f"Generated ETL schema from Parquet: {etl_schema.schema_id}") + + from graphrag_toolkit.document_graph.pipeline.extract.extraction_result import ExtractionResult + return ExtractionResult( + document_id=self.config.document_id, + extracted_schema=schema, + nodes=nodes, + dataframe=df, + metadata={ + "source": source, + "extraction_type": "parquet", + "etl_schema_id": etl_schema.schema_id if etl_schema else None, + "rows_processed": str(len(df)), + "schema_fields": str(len(schema)) + } + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_registry.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_registry.py new file mode 100644 index 00000000..ae9670e5 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extract_provider_registry.py @@ -0,0 +1,98 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Registry for extract provider implementations. + +Maintains a mapping of provider type names to their concrete implementations. +New extract providers should be registered here to be available through +the ExtractProviderFactory. + +This module provides a registry for extract provider classes, allowing them to be +dynamically registered and retrieved by name or type. The registry is a central +component of the extract framework, enabling the factory to instantiate providers +without direct dependencies on specific implementations. + +The registry follows a singleton pattern, with a single instance (extract_provider_registry) +that is used throughout the application to register and retrieve provider classes. +""" + +import logging +from typing import Dict, Type + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + +from .extract_provider_base import ExtractProvider + + + +class ExtractProviderRegistry: + """ + Provides a registry for managing extract provider classes. + + The class enables registering, retrieving, and listing extract provider classes by their type + names. It is case-insensitive and supports runtime replacement of registered providers. This is + commonly used to handle dynamic provider configurations in a flexible and modular way. + """ + + _registry: Dict[str, Type[ExtractProvider]] = {} + + @classmethod + def register(cls, type_: str, provider_cls: Type[ExtractProvider]) -> None: + """ + Registers an extract provider with a specific type. + + This method associates a provider class with a specific `type_` and + stores it in the `_registry`. The type is stored in lowercase form + to ensure case-insensitive lookups. + + Parameters: + type_ (str): The case-insensitive type associated with the provider. + provider_cls (Type[ExtractProvider]): The class of the provider to be registered. + + Returns: + None + """ + logger.debug(f"Registering extract provider: {type_} -> {provider_cls.__name__}") + cls._registry[type_.lower()] = provider_cls + + @classmethod + def get(cls, type_: str) -> Type[ExtractProvider]: + """ + Retrieves an extract provider class registered for the specified type. + + This method is responsible for looking up an extract provider in the + registry using the specified type. If no provider is found for the given + type, a KeyError is raised. + + Args: + type_: The type of extract provider to retrieve. The type should be + provided as a string. + + Raises: + KeyError: If no extract provider is registered for the specified type. + + Returns: + A registered extract provider class corresponding to the specified type. + """ + key = type_.lower() + if key not in cls._registry: + raise KeyError(f"No extract provider registered for type='{type_}'") + return cls._registry[key] + + @classmethod + def list_providers(cls) -> list[str]: + """ + Provides a method to list all registered providers. + + Returns: + list[str]: A list containing the names of all providers currently + registered in the class-level registry. + """ + return list(cls._registry.keys()) + + +# Initialize registry instance +extract_provider_registry = ExtractProviderRegistry() diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extraction_result.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extraction_result.py new file mode 100644 index 00000000..6e4a179c --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/extraction_result.py @@ -0,0 +1,14 @@ +"""Extraction result — output of the extract stage.""" +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional +import pandas as pd + +@dataclass +class ExtractionResult: + """Holds the output of an extraction stage including schema, nodes, and data.""" + + document_id: str = "" + extracted_schema: dict = field(default_factory=dict) + nodes: list = field(default_factory=list) + dataframe: Optional[pd.DataFrame] = None + metadata: Dict[str, Any] = field(default_factory=dict) diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/__init__.py new file mode 100644 index 00000000..058253e2 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/__init__.py @@ -0,0 +1 @@ +"""Extract utilities — helper functions for data extraction.""" diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_fixer.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_fixer.py new file mode 100644 index 00000000..ba0d5d49 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_fixer.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +"""CSV Fixer utility using document-graph schema for validation and repair.""" + +import json +import csv +import logging +from pathlib import Path +from typing import List, Dict, Any, Optional + +logger = logging.getLogger(__name__) + +class CSVFixer: + """ + Class to handle fixing malformed CSV files. + + This class supports fixing and processing CSV files based on a predefined schema. It uses a two-stage + process to handle the input, ensuring proper formatting of the records. The process identifies and corrects + issues such as incomplete records caused by multiline fields and provides logging for errors and corrections. + + Attributes + ---------- + schema : dict + Dictionary extracting column metadata and optional fields from the given schema file. + expected_columns : list[str] + List of expected column names as per the schema. + expected_count : int + Total count of expected columns based on the schema. + optional_fields : set[str] + Set of column names identified as optional. + """ + + def __init__(self, schema_path: str): + """Initialize with schema file.""" + with open(schema_path, 'r') as f: + self.schema = json.load(f) + + self.expected_columns = list(self.schema['columns'].keys()) + self.expected_count = len(self.expected_columns) + self.optional_fields = set(self.schema.get('optional_fields', [])) + + logger.info(f"Schema loaded: {self.expected_count} columns expected") + + def _stage_lines(self, input_path: str, staging_path: Path): + """ + Stages the lines from an input file by copying them to a staging file. + + This method reads the contents of the specified input file line by line and + writes them to a designated staging file. Errors encountered during input file + reading are replaced to ensure the process completes without interruption. + + Parameters: + input_path (str): The path to the input file to be staged. + staging_path (Path): The path to the output staging file. + + Returns: + None + """ + with open(input_path, 'r', encoding='utf-8', errors='replace') as infile: + with open(staging_path, 'w', encoding='utf-8') as stagefile: + for line in infile: + stagefile.write(line) + + def fix_csv(self, input_path: str, output_path: str, badlog_path: Optional[str] = None) -> Dict[str, Any]: + """ + Fixes and processes a CSV file, identifying and resolving errors or inconsistencies such as incomplete records, + excess fields, or multiline fields. The method can optionally log bad lines to a specified file. + + Parameters: + input_path: str + Path to the input CSV file. + output_path: str + Path to the corrected output CSV file. + badlog_path: Optional[str] + Path to the log file where bad/incomplete lines will be recorded, if provided. + + Returns: + Dict[str, Any] + A dictionary containing processing statistics: + - 'total_lines': Total lines read from the input file. + - 'valid_records': Number of valid records written to the output file without modification. + - 'fixed_records': Number of records corrected and written. + - 'skipped_lines': Number of invalid or irrecoverable lines that were skipped. + - 'errors': List of error messages encountered during processing. + """ + # Stage 1: Read all lines into staging + staging_path = Path(output_path).with_suffix('.staging') + self._stage_lines(input_path, staging_path) + + # Stage 2: Process staged lines + stats = { + 'total_lines': 0, + 'valid_records': 0, + 'fixed_records': 0, + 'skipped_lines': 0, + 'errors': [] + } + + badlog_file = open(badlog_path, 'w', encoding='utf-8') if badlog_path else None + + with (open(staging_path, 'r', encoding='utf-8') as infile, open(output_path, 'w', encoding='utf-8', newline='') as outfile): + + writer = csv.writer(outfile, quoting=csv.QUOTE_MINIMAL) + + # Write header + writer.writerow(self.expected_columns) + + # Process line by line + current_record = [] + in_multiline_field = False + multiline_buffer = "" + multiline_line_count = 0 + MAX_MULTILINE_LINES = 10 + + for line_num, line in enumerate(infile, 1): + stats['total_lines'] += 1 + line = line.rstrip('\n\r') + + if line_num % 10000 == 0 or line_num > 78000: + print(f"Processing line {line_num}, records output: {stats['valid_records'] + stats['fixed_records']}, multiline: {in_multiline_field}") + + if line_num == 1: # Skip original header + continue + + try: + # Parse line as CSV + parsed = list(csv.reader([line]))[0] + + if not in_multiline_field: + # Start new record + if len(parsed) == self.expected_count: + # Perfect record + writer.writerow(parsed) + stats['valid_records'] += 1 + elif len(parsed) < self.expected_count: + # Potentially incomplete (multiline field) + current_record = parsed + in_multiline_field = True + multiline_buffer = line + multiline_line_count = 1 + else: + # Too many fields - log and skip + error_msg = f"Line {line_num}: Too many fields ({len(parsed)}) expected {self.expected_count}" + stats['errors'].append(error_msg) + if badlog_file: + badlog_file.write(error_msg + "\n") + stats['skipped_lines'] += 1 + else: + # Continue multiline field + multiline_line_count += 1 + + # Safety check: reset if too many lines accumulated + if multiline_line_count > MAX_MULTILINE_LINES: + # Output partial record with padding + while len(current_record) < self.expected_count: + current_record.append("") + writer.writerow(current_record[:self.expected_count]) + stats['fixed_records'] += 1 + + error_msg = f"Line {line_num}: Multiline field exceeded {MAX_MULTILINE_LINES} lines, outputting partial record" + stats['errors'].append(error_msg) + if badlog_file: + badlog_file.write(error_msg + "\n") + in_multiline_field = False + current_record = [] + multiline_buffer = "" + multiline_line_count = 0 + continue + + multiline_buffer += "\n" + line + + # Try to parse accumulated buffer + try: + accumulated = list(csv.reader([multiline_buffer]))[0] + if len(accumulated) == self.expected_count: + # Complete record found + writer.writerow(accumulated) + stats['fixed_records'] += 1 + in_multiline_field = False + current_record = [] + multiline_buffer = "" + multiline_line_count = 0 + elif len(accumulated) > self.expected_count: + # Overshot - reset and try current line as new record + error_msg = f"Line {line_num}: Multiline field overshot, resetting" + stats['errors'].append(error_msg) + if badlog_file: + badlog_file.write(error_msg + "\n") + in_multiline_field = False + current_record = [] + multiline_buffer = "" + multiline_line_count = 0 + # Retry current line as new record + if len(parsed) == self.expected_count: + writer.writerow(parsed) + stats['valid_records'] += 1 + except Exception: + # Continue accumulating + pass + + except Exception as e: + if in_multiline_field: + # Add to multiline buffer + multiline_buffer += "\n" + line + else: + error_msg = f"Line {line_num}: Parse error - {e}" + stats['errors'].append(error_msg) + if badlog_file: + badlog_file.write(error_msg + "\n") + stats['skipped_lines'] += 1 + + # Handle any remaining multiline buffer at end of file + if in_multiline_field and multiline_buffer: + try: + accumulated = list(csv.reader([multiline_buffer]))[0] + if len(accumulated) <= self.expected_count: + # Pad with empty strings if needed + while len(accumulated) < self.expected_count: + accumulated.append("") + writer.writerow(accumulated[:self.expected_count]) + stats['fixed_records'] += 1 + if badlog_file: + badlog_file.write(f"EOF: Flushed incomplete multiline record with {len(accumulated)} fields\n") + except Exception as e: + if badlog_file: + badlog_file.write(f"EOF: Failed to flush multiline buffer: {e}\n") + stats['skipped_lines'] += 1 + + if badlog_file: + badlog_file.close() + + # Cleanup staging file + try: + staging_path.unlink() + except Exception: + pass + + logger.info(f"CSV fix complete: {stats['valid_records']} valid, {stats['fixed_records']} fixed, {stats['skipped_lines']} skipped") + return stats + +def main(): + """ + Fixes a malformed CSV file using a provided JSON schema, saving the corrected + file and optionally logging errors. + + Parameters: + input_csv: str + The path to the input malformed CSV file. + schema_json: str + The path to the JSON file containing the schema to validate and fix the + CSV. + output_csv: str + The path to the output file for the fixed CSV. + badlog: str, optional + The path to the file where errors will be logged during the processing. + + Returns: + None + + Raises: + None + """ + import argparse + + parser = argparse.ArgumentParser(description="Fix malformed CSV using schema") + parser.add_argument("input_csv", help="Input CSV file") + parser.add_argument("schema_json", help="Schema JSON file") + parser.add_argument("output_csv", help="Output fixed CSV file") + parser.add_argument("--badlog", help="Write error log to file") + + args = parser.parse_args() + + fixer = CSVFixer(args.schema_json) + stats = fixer.fix_csv(args.input_csv, args.output_csv, args.badlog) + + print(f"Fixed CSV saved to: {args.output_csv}") + print(f"Stats: {stats}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_line_normalizer.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_line_normalizer.py new file mode 100644 index 00000000..c38f38ba --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_line_normalizer.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Binary-level CSV line normalizer to fix corrupted line endings.""" + +import argparse +from pathlib import Path + +def normalize_csv_lines(input_path: str, output_path: str, expected_fields: int = 39): + """ + Normalizes lines in a CSV file to ensure proper field structure and formatting. + + This function reads a CSV file in binary mode, processes its lines to ensure + each record has the expected number of fields, and writes normalized lines + to an output file. Lines with incomplete records are adjusted to maintain + proper field continuity by replacing line breaks with spaces within fields. + + Attributes: + None + + Args: + input_path: str + Path to the input CSV file to be normalized. + output_path: str + Path to the output CSV file where normalized content will be saved. + expected_fields: int, optional + The number of fields expected in each record. Defaults to 39. + + Raises: + None + + Returns: + None + """ + + with open(input_path, 'rb') as infile, open(output_path, 'wb') as outfile: + buffer = b'' + field_count = 0 + in_quotes = False + line_count = 0 + + while True: + chunk = infile.read(8192) + if not chunk: + break + + buffer += chunk + + i = 0 + while i < len(buffer): + byte = buffer[i:i+1] + + if byte == b'"': + in_quotes = not in_quotes + elif byte == b',' and not in_quotes: + field_count += 1 + elif byte in (b'\n', b'\r') and not in_quotes: + # Check if we have the right number of fields + if field_count == expected_fields - 1: # -1 because last field doesn't have comma + # Proper record end - write LF + outfile.write(b'\n') + field_count = 0 + line_count += 1 + if line_count % 10000 == 0: + print(f"Processed {line_count} records") + else: + # Incomplete record - replace with space to continue field + outfile.write(b' ') + + # Skip any additional CR/LF bytes + while i + 1 < len(buffer) and buffer[i+1:i+2] in (b'\n', b'\r'): + i += 1 + else: + # Regular byte - pass through + outfile.write(byte) + + i += 1 + + # Keep last incomplete line in buffer + buffer = b'' + + print(f"Normalized {line_count} records") + +def main(): + """ + Parses command-line arguments to normalize CSV line endings at the binary level + and invokes the function for processing the CSV file. + + Summary: + This script provides functionality to take an input corrupted CSV file and an output file location, + normalize the corrupted line endings by considering the provided number of fields, + and save the normalized version into the output file. + + Args: + input_csv (str): Path to the input corrupted CSV file. + output_csv (str): Path where the output normalized CSV file should be saved. + fields (int, optional): The expected number of fields per record. Defaults to 39. + """ + parser = argparse.ArgumentParser(description="Normalize CSV line endings at binary level") + parser.add_argument("input_csv", help="Input corrupted CSV file") + parser.add_argument("output_csv", help="Output normalized CSV file") + parser.add_argument("--fields", type=int, default=39, help="Expected number of fields per record") + + args = parser.parse_args() + + normalize_csv_lines(args.input_csv, args.output_csv, args.fields) + print(f"Normalized CSV saved to: {args.output_csv}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_to_json.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_to_json.py new file mode 100644 index 00000000..dfa3c41d --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/csv_to_json.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Convert CSV to JSON Lines format.""" + +import pandas as pd +import json +import argparse + +def csv_to_jsonl(csv_path: str, json_path: str): + """ + Converts a CSV file to a JSON Lines file with additional data processing and cleaning. + + This function attempts to read a CSV file using multiple strategies to ensure resilience + against format inconsistencies. Once successfully read, the CSV data is processed + row-by-row to handle NaN values, clean problematic characters in string fields, + and parse predefined JSON fields. The processed data is then written to the specified + JSON Lines file. + + Parameters: + csv_path: str + The file path to the input CSV file. + json_path: str + The file path where the output JSON Lines file will be written. + + Raises: + ValueError + If all CSV reading strategies fail. + """ + # Try multiple CSV reading strategies + df = None + strategies = [ + {'sep': ',', 'quotechar': '"', 'quoting': 1}, # QUOTE_ALL + {'sep': ',', 'quotechar': '"', 'quoting': 0}, # QUOTE_MINIMAL + {'sep': ',', 'quotechar': '"', 'quoting': 3, 'skipinitialspace': True}, # QUOTE_NONE + {'sep': ',', 'on_bad_lines': 'skip'}, # Skip bad lines + ] + + for i, strategy in enumerate(strategies): + try: + print(f"Trying CSV read strategy {i+1}...") + df = pd.read_csv(csv_path, **strategy) + print(f"✅ Success with strategy {i+1}") + break + except Exception as e: + print(f"❌ Strategy {i+1} failed: {e}") + continue + + if df is None: + raise ValueError("All CSV reading strategies failed") + + print(f"Successfully read CSV with {len(df)} rows and {len(df.columns)} columns") + + with open(json_path, 'w', encoding='utf-8') as f: + for _, row in df.iterrows(): + # Convert pandas Series to dict, handling NaN values + record = row.to_dict() + # Replace NaN with None for proper JSON serialization + record = {k: (None if pd.isna(v) else v) for k, v in record.items()} + + # Clean string values - remove problematic characters + for k, v in record.items(): + if isinstance(v, str): + # Clean up common problematic characters + v = v.replace('\r\n', '\n').replace('\r', '\n') + # Ensure proper encoding + record[k] = v + + # Handle special JSON fields + json_fields = ['Resource original JSON', 'Evidence'] + for field in json_fields: + if field in record and record[field] and record[field] != 'NaN': + try: + # Try to parse as JSON + if isinstance(record[field], str): + record[field] = json.loads(record[field]) + except (json.JSONDecodeError, TypeError): + # Keep as string if not valid JSON + pass + + # Use json.dumps with ensure_ascii=False for proper encoding + f.write(json.dumps(record, ensure_ascii=False, separators=(',', ':')) + '\n') + + print(f"Converted {len(df)} records from CSV to JSON Lines") + +def main(): + """ + Main function for converting a CSV file into a JSON Lines file. + + This function initializes an argument parser to handle the input and output + file paths. It processes the command line arguments, and executes the + conversion using the csv_to_jsonl function. If the conversion is successful, + it displays a confirmation message; otherwise, it displays an error message + and raises the encountered exception. + + Raises: + Exception: Re-raises any exception encountered during the conversion + process. + """ + parser = argparse.ArgumentParser(description="Convert CSV to JSON Lines") + parser.add_argument("csv_file", help="Input CSV file") + parser.add_argument("json_file", help="Output JSON Lines file") + + args = parser.parse_args() + + try: + csv_to_jsonl(args.csv_file, args.json_file) + print(f"✅ Successfully converted {args.csv_file} to {args.json_file}") + except Exception as e: + print(f"❌ Conversion failed: {e}") + raise + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/dataframe_reader.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/dataframe_reader.py new file mode 100644 index 00000000..fb91b8ed --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/dataframe_reader.py @@ -0,0 +1,13 @@ +"""DataFrame reader — wraps pandas read_csv with config.""" +import pandas as pd +from pathlib import Path + +class DataFrameReader: + """Reads data files into pandas DataFrames with configurable options.""" + + def __init__(self, path, **kwargs): + self.path = Path(path) + self.kwargs = kwargs + + def read(self) -> pd.DataFrame: + return pd.read_csv(self.path, **self.kwargs) diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/graph_element_utils.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/graph_element_utils.py new file mode 100644 index 00000000..61aa87d8 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/graph_element_utils.py @@ -0,0 +1,57 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""[Module Name] Module for Document Graph Plugin System. + +This module [provides/defines/implements] [brief description of module purpose]. + +The module [includes/contains/offers] [key components or features]: +- [Component/Feature 1]: [Brief description] +- [Component/Feature 2]: [Brief description] +- [Component/Feature 3]: [Brief description] + +[Additional context about design principles, architecture, or implementation details] + +Usage: + # [Example usage code] + [explanation of example] +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Graph element utilities for creating nodes, links, and edges. + +This module provides functions for creating graph elements such as text nodes and links/edges +between nodes. It delegates to the shared implementation in the graph module. +""" + +# Import from the authoritative shared location +from graphrag_toolkit.document_graph.shared.graph.graph_element_utils import ( + create_text_node, + create_link +) + +# Alias for backward compatibility +def create_edge(source: str, target: str, edge_type: str): + """Create an edge between two nodes in the graph. + + This is an alias for create_link, maintained for backward compatibility. + + Args: + source: The ID of the source node + target: The ID of the target node + edge_type: The type of edge to create + + Returns: + A dictionary representing the edge/link between the nodes + """ + return create_link(source, target, edge_type) diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/parsing_utils.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/parsing_utils.py new file mode 100644 index 00000000..a9553eec --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/extract/utils/parsing_utils.py @@ -0,0 +1,48 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Parsing utilities for text processing in document graph operations. + +This module provides utility functions for cleaning and parsing text data +before it is processed further in the document graph pipeline. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + +import re +from typing import List + +def clean_text(text: str) -> str: + """Clean text by removing excess whitespace. + + This function strips leading and trailing whitespace and normalizes + all internal whitespace sequences (spaces, tabs, newlines) to a single space. + + Args: + text: The input text to clean + + Returns: + The cleaned text with normalized whitespace + """ + return re.sub(r'\s+', ' ', text.strip()) + +def split_into_paragraphs(text: str) -> List[str]: + """Split text into logical paragraphs. + + This function divides text into paragraphs by looking for double line breaks + (i.e., an empty line between paragraphs). It strips whitespace from each + paragraph and filters out empty paragraphs. + + Args: + text: The input text to split into paragraphs + + Returns: + A list of paragraphs, with each paragraph as a string + """ + return [p.strip() for p in re.split(r'\n\s*\n', text) if p.strip()] diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/__init__.py new file mode 100644 index 00000000..70b02766 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/__init__.py @@ -0,0 +1 @@ +"""Load — data loading providers for graph databases.""" diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_context.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_context.py new file mode 100644 index 00000000..ba34774a --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_context.py @@ -0,0 +1,93 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load Context module for document graph operations.""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + +from typing import Optional +from graphrag_toolkit.document_graph.pipeline.extract.extracted_data import ExtractedData +from graphrag_toolkit.document_graph.storage.graph.storage_provider_factory import StorageProviderFactory +from graphrag_toolkit.document_graph.storage.graph.validated_config import StorageProviderConfig + + +class LoadContext: + """ + Represents a context for loading and managing extracted data. + + This class is responsible for handling the management of data extracted from a source. + It provides utilities to access or create relevant storage mechanisms as well as + to interact with the data source registration. The class ensures that the data needed + for processing or analysis is properly handled and stored. + + Attributes: + extracted_data (ExtractedData): The extracted data object that includes information + about the data source and registration. + """ + + def __init__(self, extracted_data: ExtractedData): + """ + Initializes a new instance of the class with provided extracted data. + + Attributes: + extracted_data: ExtractedData + The data extracted and provided for initialization. + + Arguments: + extracted_data: ExtractedData + The extracted data to initialize the instance. + """ + self.extracted_data = extracted_data + self._storage = None + + @property + def storage(self): + """ + Provides accessor for the `storage` property of the class. The property ensures + that a storage provider instance is initialized and configured properly for + the class instance. If the storage has already been initialized, the existing + instance is returned. Otherwise, the storage configuration is built using the + graph information and a new storage provider instance is created. + + Raises: + ValueError: If the graph information is not found in the data registration. + + Returns: + StorageProvider: An instance of a storage provider initialized with the + corresponding configuration. + """ + if self._storage is None: + graph_info = self.extracted_data.source.get_graph_info() + if not graph_info: + raise ValueError(f"Graph {self.extracted_data.graph_id} not found in registration") + + # Create storage config from registration + # This would be configured based on the data plane setup + storage_config = StorageProviderConfig( + provider_type="graphjson", # Default, could be from registration + connection_config={"path": f"/tmp/{self.extracted_data.graph_id}.json"}, + tenant_id=self.extracted_data.tenant_id + ) + + self._storage = StorageProviderFactory.for_writer(storage_config) + + return self._storage + + @property + def registration(self): + """ + Provides access to the 'registration' property. + + This property retrieves and returns the value of the 'registration' attribute + from the 'extracted_data' object. + + Returns: + The registration value fetched from the 'extracted_data' object. + """ + return self.extracted_data.registration \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_provider_base.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_provider_base.py new file mode 100644 index 00000000..66c1ba23 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_provider_base.py @@ -0,0 +1,170 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load Provider Base module for document graph operations.""" + +import logging +from abc import ABC, abstractmethod +from typing import Dict, Any, List, Optional +from pydantic import BaseModel +from graphrag_toolkit.document_graph.pipeline.extract.extraction_result import ExtractionResult +from graphrag_toolkit.document_graph.pipeline.load.load_result import LoadResult +from graphrag_toolkit.document_graph.pipeline.load.transformers.transformation_plan import TransformationPlan + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + + +class LoadProviderConfig(BaseModel): + """ + Represents the configuration needed to load a provider. + + This class is used to define the necessary information for loading a + specific provider, including its name, type, optional parameters, and + optional configurations such as transformation plans or AWS-specific + settings. It extends the BaseModel class for structured data handling. + """ + name: str + type: str + parameters: Dict[str, Any] = {} + transformation_plan: Optional[TransformationPlan] = None + aws_config: Optional[Any] = None + + +class LoadProvider(ABC): + """ + Abstract base class for implementing load providers. + + The LoadProvider class defines the blueprint for any load provider that processes + and optionally transforms extracted data before persisting it in a target system. + It allows customization through a configuration and manages the lifecycle of loading + and transformation operations. + + Attributes: + config (LoadProviderConfig): Configuration for the load provider, including the transformation plan. + transformation_plan (TransformationPlan): Transformation plan detailing the steps to be applied + before persisting the data. + """ + + def __init__(self, config: LoadProviderConfig): + """Initialize load provider with configuration.""" + self.config = config + self.transformation_plan = config.transformation_plan + logger.debug(f"Initialized {self.__class__.__name__}: name={config.name}, type={config.type}") + logger.debug(f"Load provider parameters: {config.parameters}") + if self.transformation_plan: + logger.debug(f"Transformation plan configured with {len(self.transformation_plan.steps)} steps") + + def load(self, extraction_result: ExtractionResult) -> LoadResult: + """ + Load data into the target destination after applying necessary transformations. + + This method orchestrates the load operation, which involves optional data + transformations and persistence into the targeted destination. Logging is + employed to track each stage of the operation for debugging and informational + purposes. In case of failure, a failed load result is created and returned. + + Parameters: + extraction_result (ExtractionResult): The input data to be loaded, + containing the document ID and other necessary information. + + Returns: + LoadResult: Indicates the success or failure of the load operation, + including any relevant metadata or error details. + """ + logger.debug(f"Starting load operation: document={extraction_result.document_id}, provider={self.__class__.__name__}") + logger.info(f"Starting load operation for document: {extraction_result.document_id}") + + try: + if self.transformation_plan: + logger.debug(f"Applying {len(self.transformation_plan.steps)} transformation steps") + extraction_result = self._apply_transformations(extraction_result) + logger.debug(f"Transformations completed successfully") + + logger.debug(f"Persisting data using {self.__class__.__name__}") + result = self._persist(extraction_result) + logger.debug(f"Persistence completed: success={result.success}") + logger.info(f"Load operation completed for document: {extraction_result.document_id}") + return result + except Exception as e: + logger.error(f"Load operation failed for document {extraction_result.document_id}: {e}") + return LoadResult.create_failed( + document_id=extraction_result.document_id, + error_message=str(e) + ) + + def _apply_transformations(self, extraction_result: ExtractionResult) -> ExtractionResult: + """ + Applies a series of transformations to the given extraction result based on a predefined + transformation plan. Each transformation step is executed sequentially, modifying the + nodes or dataframe of the extraction result as specified. + + Parameters: + extraction_result (ExtractionResult): The initial extraction result to be transformed. + + Returns: + ExtractionResult: The transformed extraction result after applying all transformation + steps. + + Raises: + Exception: If any transformation step fails during the process. + """ + from graphrag_toolkit.document_graph.transformers.transformer_provider_factory import TransformerProviderFactory + from graphrag_toolkit.document_graph.transformers.transformer_provider_config import TransformerProviderConfig + + transformed_result = extraction_result + + for i, step in enumerate(self.transformation_plan.steps): + logger.debug(f"Applying transformation step {i + 1}/{len(self.transformation_plan.steps)}: {step.name}") + try: + transformer_config = TransformerProviderConfig( + name=step.name, + type=step.type or "transformer", + args=step.args or {} + ) + transformer = TransformerProviderFactory.get_provider(transformer_config) + + if transformed_result.nodes: + transformed_result.nodes = [ + transformer.transform(node, **(step.args or {})) + for node in transformed_result.nodes + ] + + if transformed_result.dataframe is not None: + transformed_result.dataframe = transformer.transform( + transformed_result.dataframe, **(step.args or {}) + ) + + logger.debug(f"Successfully applied transformation: {step.name}") + except Exception as e: + logger.error(f"Failed to apply transformation {step.name}: {e}") + raise + + return transformed_result + + @abstractmethod + def _persist(self, extraction_result: ExtractionResult) -> LoadResult: + """ + Represents an abstract method for persistence of extraction results. This method + must be implemented by subclasses to define how to persist the given + extraction result and return the corresponding load result. + + Parameters: + extraction_result: ExtractionResult + The result of some extraction process that needs to be persisted. + + Returns: + LoadResult + The result of the persistence operation. + + Raises: + NotImplementedError + If this method is not implemented in a subclass. + """ + pass + + def __repr__(self): + return f"<{self.__class__.__name__} name={self.config.name} type={self.config.type}>" diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_provider_config.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_provider_config.py new file mode 100644 index 00000000..91700fe6 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_provider_config.py @@ -0,0 +1,208 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load Provider Configuration Module. + +This module defines configuration classes for load providers in the +document graph pipeline. +""" + +import logging +from typing import Dict, Any, Optional +from pydantic import BaseModel, Field + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + + +class LoadProviderConfig(BaseModel): + """ + Represents the configuration for a load provider. + + Provides detailed setup, validation, and pre-defined methods to configure + load providers based on type. These providers can be file-based, database-based, + or S3 configurations. The class supports defining and customizing configurations + with specific attributes and capabilities depending on the provider type. + + Attributes: + name (str): Unique name for this load provider. + type (str): Load provider type (e.g., 'file', 'database', 's3'). + output_dir (Optional[str]): Output directory for file-based providers. + output_format (str): Output format (json, csv, parquet). + database_config (Optional[Dict[str, Any]]): Database connection configuration. + file_config (Optional[Dict[str, Any]]): File output configuration. + s3_config (Optional[Dict[str, Any]]): S3 output configuration. + parameters (Dict[str, Any]): Additional provider parameters. + batch_size (int): Batch size for processing. + overwrite (bool): Whether to overwrite existing output. + """ + + name: str = Field(..., description="Unique name for this load provider") + type: str = Field(..., description="Load provider type (e.g., 'file', 'database', 's3')") + output_dir: Optional[str] = Field(None, description="Output directory for file-based providers") + output_format: str = Field(default="json", description="Output format (json, csv, parquet)") + + # Database-specific configuration + database_config: Optional[Dict[str, Any]] = Field(None, description="Database connection configuration") + + # File-specific configuration + file_config: Optional[Dict[str, Any]] = Field(None, description="File output configuration") + + # S3-specific configuration + s3_config: Optional[Dict[str, Any]] = Field(None, description="S3 output configuration") + + # General parameters + parameters: Dict[str, Any] = Field(default_factory=dict, description="Additional provider parameters") + + # Processing options + batch_size: int = Field(default=1000, description="Batch size for processing") + overwrite: bool = Field(default=False, description="Whether to overwrite existing output") + + class Config: + """Pydantic configuration.""" + extra = "allow" # Allow additional fields + + def __post_init__(self): + """ + Logs initialization details of the LoadProviderConfig object. + + This method is automatically called after the object is fully + initialized, providing a debug log entry with the name and type + of the configuration. + + Raises: + None + """ + logger.debug(f"Initialized LoadProviderConfig: {self.name} ({self.type})") + + @classmethod + def create_file_config( + cls, + name: str, + output_dir: str, + output_format: str = "json", + **kwargs + ) -> "LoadProviderConfig": + """ + Creates a configuration for a file-based load provider. + + This class method generates and returns a LoadProviderConfig instance + with configuration tailored for file-based data load. It includes specific + parameters such as the file output directory and format and encapsulates + additional configuration options if provided. + + Parameters: + name: str + The name of the configuration. + output_dir: str + The directory where files will be output. + output_format: str, optional + The format of the output files. Defaults to "json". + kwargs: + Additional keyword arguments for custom configurations. + + Returns: + LoadProviderConfig + An instance of LoadProviderConfig initialized with the specified + parameters. + + """ + return cls( + name=name, + type="file", + output_dir=output_dir, + output_format=output_format, + file_config={ + "output_dir": output_dir, + "format": output_format + }, + **kwargs + ) + + @classmethod + def create_database_config( + cls, + name: str, + database_type: str, + connection_config: Dict[str, Any], + **kwargs + ) -> "LoadProviderConfig": + """ + Creates a LoadProviderConfig instance configured for a database load provider. + + This class method initializes a LoadProviderConfig object with parameters suited + to a database setup. The method associates the instance with a specific + database type and a set of connection configurations. Additional keyword arguments + can be supplied to further customize the configuration. + + Parameters: + name: str + The name of the load provider configuration. + database_type: str + The type of database being configured, such as 'PostgreSQL', 'MySQL', etc. + connection_config: Dict[str, Any] + A dictionary containing the database connection configuration settings. + **kwargs + Additional keyword arguments to customize the configuration. + + Returns: + LoadProviderConfig + An instance of LoadProviderConfig initialized with database-related settings. + """ + return cls( + name=name, + type="database", + database_config={ + "type": database_type, + "connection": connection_config + }, + **kwargs + ) + + @classmethod + def create_s3_config( + cls, + name: str, + bucket: str, + prefix: str = "", + output_format: str = "json", + **kwargs + ) -> "LoadProviderConfig": + """ + Creates a LoadProviderConfig instance for an s3 load provider. + + The method acts as a factory method to generate configurations specifically for + an s3-based data load provider. The s3 configuration includes the bucket name, + optional prefix for object paths, and the desired output format for processing. + + Parameters: + name: str + The name of the load provider configuration. + bucket: str + The bucket name of the s3 storage where data resides. + prefix: str, optional + The prefix of object paths within the s3 bucket. Default is an empty string. + output_format: str, optional + The format of output data (e.g., "json"). Default is "json". + **kwargs: dict + Additional arguments which will be passed to the LoadProviderConfig + constructor. + + Returns: + LoadProviderConfig + A configured instance of LoadProviderConfig representing an s3 load provider. + """ + return cls( + name=name, + type="s3", + output_format=output_format, + s3_config={ + "bucket": bucket, + "prefix": prefix, + "format": output_format + }, + **kwargs + ) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_result.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_result.py new file mode 100644 index 00000000..5a57d805 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline/load/load_result.py @@ -0,0 +1,189 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load Result module for document graph operations.""" + +import logging +from typing import Optional, Dict, Any +from datetime import datetime, timezone +from pydantic import BaseModel + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +# logging via graphrag-toolkit # noqa: F401 + + + +class LoadResult(BaseModel): + """ + Represents the outcome of a document load operation. + + This class is used to encapsulate the details and status of a load operation, + including the document identifier, records loaded, load time, overall status, + and additional metadata related to the load. It provides utility methods to + check the success or emptiness of the load result and class methods to create + specific types of results such as successful, empty, or failed load outcomes. + + Attributes: + document_id (str): Identifier of the document being loaded. + records_loaded (int): Number of records successfully loaded. Defaults to 0. + load_time (str): Timestamp when the load operation occurred. If not provided + during initialization, it will be automatically generated in UTC. + status (str): Status of the load operation, defaulting to "success". + transformations_applied (bool): Indicates whether transformations were applied + during the load process. Defaults to False. + output_files (Optional[Dict[str, str]]): Dictionary containing file paths for + generated output files. Defaults to None if not provided. + metadata (Optional[Dict[str, Any]]): Additional metadata related to the load + operation. Defaults to None if not provided. + """ + + document_id: str + records_loaded: int = 0 + load_time: str + status: str = "success" + transformations_applied: bool = False + output_files: Optional[Dict[str, str]] = None + metadata: Optional[Dict[str, Any]] = None + + def __init__(self, **data): + """ + Initializes a LoadResult object. + + This method sets the load_time for the object upon creation if it is not explicitly + provided in the input data. It then invokes the parent class initializer with the + given data. A debug log entry is generated for the initialization of a LoadResult + instance. + + Attributes: + load_time (str): A UTC timestamp indicating the object's creation time. If not + provided, it defaults to the current time. + + Parameters: + data (dict): Arbitrary keyword arguments representing the properties of the + LoadResult object. + + Returns: + None + """ + if 'load_time' not in data: + data['load_time'] = datetime.now(timezone.utc).isoformat() + super().__init__(**data) + logger.debug(f"Created LoadResult for document: {self.document_id}") + + def __str__(self) -> str: + """ + Converts the object to its string representation. + + Returns + ------- + str + A string that includes the document ID, the number of records loaded, and + the status of the operation. + """ + return f"LoadResult(document_id={self.document_id}, records_loaded={self.records_loaded}, status={self.status})" + + def is_successful(self) -> bool: + """ + Determines whether the operation was successful based on its status + and the number of records loaded. + + Returns: + bool: True if the status is "success" and at least one record was + loaded; otherwise, False. + """ + return self.status == "success" and self.records_loaded > 0 + + def is_empty(self) -> bool: + """ + Determines whether the collection is empty. + + This method evaluates whether the number of loaded records equals zero, + indicating that the collection is currently empty. + + Returns: + bool: True if the collection has no loaded records; otherwise, False. + """ + return self.records_loaded == 0 + + @classmethod + def create_success( + cls, + document_id: str, + records_loaded: int, + output_files: Optional[Dict[str, str]] = None, + transformations_applied: bool = False, + metadata: Optional[Dict[str, Any]] = None + ) -> "LoadResult": + """Create a successful load result. + + Args: + document_id: Document identifier + records_loaded: Number of records loaded + output_files: Dictionary of output file paths + transformations_applied: Whether transformations were applied + metadata: Additional metadata + + Returns: + LoadResult: Successful load result + """ + logger.info(f"Creating successful load result: {document_id} ({records_loaded} records)") + return cls( + document_id=document_id, + records_loaded=records_loaded, + status="success", + transformations_applied=transformations_applied, + output_files=output_files or {}, + metadata=metadata or {} + ) + + @classmethod + def create_empty( + cls, + document_id: str, + metadata: Optional[Dict[str, Any]] = None + ) -> "LoadResult": + """Create an empty load result. + + Args: + document_id: Document identifier + metadata: Additional metadata + + Returns: + LoadResult: Empty load result + """ + logger.warning(f"Creating empty load result: {document_id}") + return cls( + document_id=document_id, + records_loaded=0, + status="empty_result", + metadata=metadata or {} + ) + + @classmethod + def create_failed( + cls, + document_id: str, + error_message: str, + metadata: Optional[Dict[str, Any]] = None + ) -> "LoadResult": + """Create a failed load result. + + Args: + document_id: Document identifier + error_message: Error message describing the failure + metadata: Additional metadata + + Returns: + LoadResult: Failed load result + """ + logger.error(f"Creating failed load result: {document_id} - {error_message}") + metadata = metadata or {} + metadata["error_message"] = error_message + return cls( + document_id=document_id, + records_loaded=0, + status="failed", + metadata=metadata + ) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/pipeline_executor.py b/document-graph/src/graphrag_toolkit/document_graph/pipeline_executor.py new file mode 100644 index 00000000..033fd066 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/pipeline_executor.py @@ -0,0 +1,224 @@ +"""Pipeline Executor — schema-driven orchestration of ingest → transform → build. + +Reads an ETLSchema (from JSON/mapping file) and executes the full pipeline: +1. Read source data (via graphrag-toolkit readers) +2. Ingest: column/row/field prep (pandas) +3. Transform: rows → typed Nodes/Edges +4. Build: Cypher generation + write to Neptune via GraphStore +""" + +import json +import logging +from dataclasses import dataclass, field +from typing import Any, Optional + +import pandas as pd + +from graphrag_toolkit.document_graph.model_elements import Node, Edge +from graphrag_toolkit.document_graph.graph_build.cypher_builder import node_to_cypher, edge_to_cypher + +logger = logging.getLogger(__name__) + + +@dataclass +class PipelineResult: + """Result of a pipeline execution.""" + nodes_created: int = 0 + edges_created: int = 0 + records_processed: int = 0 + errors: list = field(default_factory=list) + + +class PipelineExecutor: + """Schema-driven pipeline executor. + + Reads a mapping/schema JSON and orchestrates: + source → ingest (DataFrame ops) → transform (row→node) → build (Cypher→Neptune) + + Usage: + from graphrag_toolkit.document_graph.pipeline import PipelineExecutor + + executor = PipelineExecutor(graph_store, tenant_id='scim') + result = executor.run_from_schema('mapping.json', source_df=df) + """ + + def __init__(self, graph_store, tenant_id: str = "default"): + self._store = graph_store + self._tenant_id = tenant_id + + def run_from_schema(self, schema_path: str, source_df: pd.DataFrame) -> PipelineResult: + """Execute the full pipeline from a schema file. + + Args: + schema_path: Path to mapping.json / ETL schema file + source_df: Source data as pandas DataFrame + """ + with open(schema_path) as f: + schema = json.load(f) + + result = PipelineResult(records_processed=len(source_df)) + + # Extract graph definition from schema + graph_def = schema.get("graph", {}) + node_defs = graph_def.get("nodes", []) + edge_defs = graph_def.get("edges", []) + mappings = schema.get("mappings", schema.get("properties", [])) + + # Build property map: csv_column → (node_type, property_name) + property_map = {} + for m in mappings: + csv_col = m.get("csv_property_name", "") + prop_name = m.get("arrow_property_name", csv_col) + parents = m.get("parents", [m.get("name", "")]) + map_type = m.get("type", "property") + if csv_col: + property_map[csv_col] = { + "property_name": prop_name, + "parents": parents, + "type": map_type, + } + + # Build node type definitions + node_types = {} + for node_def in node_defs: + for label, config in node_def.items(): + id_field = config.get("property", {}).get("csv_map", "id") + node_types[label] = {"id_field": id_field, "properties": {}} + + # Map properties to their node types + for csv_col, mapping in property_map.items(): + if mapping["type"] == "property": + for parent in mapping["parents"]: + if parent in node_types: + node_types[parent]["properties"][csv_col] = mapping["property_name"] + + # Build edge definitions + edge_types = [] + for edge_def in edge_defs: + for source_label, rels in edge_def.items(): + for rel_type, target_label in rels.items(): + edge_types.append({ + "source": source_label, + "target": target_label, + "type": rel_type, + }) + + # Execute: create nodes + nodes = [] + for _, row in source_df.iterrows(): + for label, type_def in node_types.items(): + id_col = type_def["id_field"] + if id_col in row and pd.notna(row[id_col]): + props = {} + for csv_col, prop_name in type_def["properties"].items(): + if csv_col in row and pd.notna(row[csv_col]): + props[prop_name] = str(row[csv_col]) + props["id"] = str(row[id_col]) + props[id_col] = str(row[id_col]) + + node = Node(id=str(row[id_col]), labels=[label], properties=props) + nodes.append(node) + + # Deduplicate nodes by id+label + seen = set() + unique_nodes = [] + for n in nodes: + key = (n.id, tuple(n.labels)) + if key not in seen: + seen.add(key) + unique_nodes.append(n) + + # Write nodes + for node in unique_nodes: + cypher, params = node_to_cypher(node, tenant_id=self._tenant_id) + try: + self._store.execute_query(cypher, params) + result.nodes_created += 1 + except Exception as e: + result.errors.append(f"Node {node.id}: {e}") + + # Write edges (based on shared fields in the data) + for edge_def in edge_types: + src_label = edge_def["source"] + tgt_label = edge_def["target"] + rel_type = edge_def["type"] + + src_id_field = node_types.get(src_label, {}).get("id_field", "") + tgt_id_field = node_types.get(tgt_label, {}).get("id_field", "") + + if src_id_field in source_df.columns and tgt_id_field in source_df.columns: + for _, row in source_df.iterrows(): + src_id = str(row[src_id_field]) if pd.notna(row.get(src_id_field)) else None + tgt_id = str(row[tgt_id_field]) if pd.notna(row.get(tgt_id_field)) else None + if src_id and tgt_id: + edge = Edge( + id=f"{src_id}-{rel_type}-{tgt_id}", + source_id=src_id, + target_id=tgt_id, + label=rel_type, + ) + cypher, params = edge_to_cypher(edge, tenant_id=self._tenant_id) + try: + self._store.execute_query(cypher, params) + result.edges_created += 1 + except Exception as e: + result.errors.append(f"Edge {src_id}->{tgt_id}: {e}") + + logger.info("Pipeline complete: %d nodes, %d edges, %d errors", + result.nodes_created, result.edges_created, len(result.errors)) + return result + + def run_from_dataframe(self, df: pd.DataFrame, node_label: str, + id_field: str = "id", + edge_field: Optional[str] = None, + edge_type: str = "RELATED_TO") -> PipelineResult: + """Simple pipeline: DataFrame → typed nodes (+ optional edges). + + For when you don't have a full schema, just a label and DataFrame. + """ + from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer + + result = PipelineResult(records_processed=len(df)) + records = df.to_dict("records") + + # Transform + config = TransformerProviderConfig(name="r2n", args={"type": node_label}) + nodes_data = RowToNodeTransformer(config).transform(records) + + # Build + write nodes + for n in nodes_data: + node = Node( + id=n.get(id_field, n.get("_id", "")), + labels=[n.get("node_type", node_label)], + properties=n, + ) + cypher, params = node_to_cypher(node, tenant_id=self._tenant_id) + try: + self._store.execute_query(cypher, params) + result.nodes_created += 1 + except Exception as e: + result.errors.append(str(e)) + + # Optional: infer edges + if edge_field and edge_field in df.columns: + from graphrag_toolkit.document_graph.transform.graph_transformers.infer_edges import EdgeInferencer + config = TransformerProviderConfig(name="edges", args={ + "source_field": edge_field, "edge_type": edge_type + }) + edges_data = EdgeInferencer(config).transform(records) + for ed in edges_data: + edge = Edge( + id=ed.get("id", ""), + source_id=ed.get("source_id", ""), + target_id=ed.get("target_id", ""), + label=ed.get("edge_type", edge_type), + ) + cypher, params = edge_to_cypher(edge, tenant_id=self._tenant_id) + try: + self._store.execute_query(cypher, params) + result.edges_created += 1 + except Exception as e: + result.errors.append(str(e)) + + return result diff --git a/document-graph/src/graphrag_toolkit/document_graph/plugins/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/plugins/__init__.py new file mode 100644 index 00000000..a264c5c4 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/plugins/__init__.py @@ -0,0 +1,84 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Plugins — extension hooks for custom transform and validation steps. + +document-graph uses pluggy to allow external packages to register custom +transformers, validators, and post-build hooks without modifying core code. + +Usage: + 1. Define a plugin class with @hookimpl methods + 2. Register it with the PluginManager + 3. The pipeline calls hooks at each extension point + +Extension Points: + - pre_transform: modify records before transformation + - post_transform: modify records after transformation + - validate_node: validate a node before graph write + - post_build: run logic after graph construction (e.g., metrics, alerts) +""" + +import pluggy + +PROJECT_NAME = "document_graph" + +hookspec = pluggy.HookspecMarker(PROJECT_NAME) +hookimpl = pluggy.HookimplMarker(PROJECT_NAME) + + +class DocumentGraphHookSpec: + """Hook specifications for document-graph plugin extension points.""" + + @hookspec + def pre_transform(self, records: list) -> list: + """Called before transformation. Can modify or filter records. + + Args: + records: List of record dicts from extraction. + + Returns: + Modified list of records (or original if unchanged). + """ + + @hookspec + def post_transform(self, records: list) -> list: + """Called after transformation. Can enrich or validate transformed records. + + Args: + records: List of transformed record dicts. + + Returns: + Modified list of records. + """ + + @hookspec + def validate_node(self, node_id: str, labels: list, properties: dict) -> bool: + """Called before writing a node to the graph. Return False to skip. + + Args: + node_id: The node's identifier. + labels: List of node labels. + properties: Node properties dict. + + Returns: + True to allow write, False to skip this node. + """ + + @hookspec + def post_build(self, stats: dict) -> None: + """Called after graph build completes. For metrics, logging, alerts. + + Args: + stats: Build statistics (nodes_written, edges_written, duration, etc.) + """ + + +def get_plugin_manager() -> pluggy.PluginManager: + """Create and return a configured plugin manager. + + Returns: + PluginManager with document-graph hookspecs registered. + """ + pm = pluggy.PluginManager(PROJECT_NAME) + pm.add_hookspecs(DocumentGraphHookSpec) + return pm diff --git a/document-graph/src/graphrag_toolkit/document_graph/plugins/samples.py b/document-graph/src/graphrag_toolkit/document_graph/plugins/samples.py new file mode 100644 index 00000000..97b472b9 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/plugins/samples.py @@ -0,0 +1,113 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Sample plugins demonstrating document-graph extension points. + +These plugins show how to hook into the pipeline without modifying core code. +Use them as templates for your own custom plugins. +""" + +from datetime import datetime, timezone +from graphrag_toolkit.document_graph.plugins import hookimpl + + +class AuditTrailPlugin: + """Adds audit metadata to every record passing through the pipeline. + + Stamps each record with: + - _ingested_at: ISO timestamp of when the record was processed + - _pipeline_version: version tag for lineage tracking + """ + + def __init__(self, pipeline_version: str = "1.0.0"): + self.pipeline_version = pipeline_version + + @hookimpl + def pre_transform(self, records: list) -> list: + now = datetime.now(timezone.utc).isoformat() + for record in records: + record['_ingested_at'] = now + record['_pipeline_version'] = self.pipeline_version + return records + + +class DataQualityPlugin: + """Validates records and rejects those that don't meet quality criteria. + + Checks: + - Required fields are present and non-empty + - No record exceeds max property count (prevents graph bloat) + """ + + def __init__(self, required_fields: list = None, max_properties: int = 50): + self.required_fields = required_fields or [] + self.max_properties = max_properties + self.rejected = [] + + @hookimpl + def pre_transform(self, records: list) -> list: + valid = [] + for record in records: + # Check required fields + missing = [f for f in self.required_fields if not record.get(f)] + if missing: + self.rejected.append({'record': record, 'reason': f'missing fields: {missing}'}) + continue + + # Check property count + if len(record) > self.max_properties: + self.rejected.append({'record': record, 'reason': f'too many properties: {len(record)}'}) + continue + + valid.append(record) + return valid + + +class NodeValidatorPlugin: + """Validates nodes before graph write — reject nodes with empty IDs or blocked labels.""" + + def __init__(self, blocked_labels: list = None): + self.blocked_labels = blocked_labels or ['_Internal', '_System', '_Temp'] + self.skipped = [] + + @hookimpl + def validate_node(self, node_id: str, labels: list, properties: dict) -> bool: + if not node_id or node_id.strip() == '': + self.skipped.append({'node_id': node_id, 'reason': 'empty ID'}) + return False + + for label in labels: + if label in self.blocked_labels: + self.skipped.append({'node_id': node_id, 'reason': f'blocked label: {label}'}) + return False + + return True + + +class MetricsPlugin: + """Collects build metrics for monitoring and alerting.""" + + def __init__(self): + self.build_history = [] + + @hookimpl + def post_build(self, stats: dict) -> None: + stats['recorded_at'] = datetime.now(timezone.utc).isoformat() + self.build_history.append(stats) + + # Alert on large builds + nodes = stats.get('nodes_written', 0) + if nodes > 10000: + print(f"⚠️ Large build detected: {nodes} nodes written") + + def get_summary(self) -> dict: + if not self.build_history: + return {'builds': 0} + total_nodes = sum(b.get('nodes_written', 0) for b in self.build_history) + total_edges = sum(b.get('edges_written', 0) for b in self.build_history) + return { + 'builds': len(self.build_history), + 'total_nodes': total_nodes, + 'total_edges': total_edges, + 'last_build': self.build_history[-1], + } diff --git a/document-graph/src/graphrag_toolkit/document_graph/query/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/query/__init__.py new file mode 100644 index 00000000..7b385453 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/query/__init__.py @@ -0,0 +1,3 @@ +"""Query stage — Cypher queries over typed document graph nodes.""" + +from .query_engine import DocumentGraphQueryEngine diff --git a/document-graph/src/graphrag_toolkit/document_graph/query/query_engine.py b/document-graph/src/graphrag_toolkit/document_graph/query/query_engine.py new file mode 100644 index 00000000..b2149f03 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/query/query_engine.py @@ -0,0 +1,43 @@ +"""Document Graph Query Engine — Cypher queries over typed nodes. + +Uses graphrag-toolkit's GraphStore for execution. +""" + +from typing import Optional + + +class DocumentGraphQueryEngine: + """Query typed document graph nodes via GraphStore.""" + + def __init__(self, graph_store, tenant_id: Optional[str] = None): + """Initialize with a graphrag-toolkit GraphStore. + + Args: + graph_store: GraphStore from GraphStoreFactory.for_graph_store() + tenant_id: Optional tenant scope for queries + """ + self._store = graph_store + self._tenant_id = tenant_id + + def query(self, cypher: str, params: dict = None) -> list[dict]: + """Execute a Cypher query against the typed document graph.""" + return self._store.execute_query(cypher, params or {}) + + def get_nodes(self, label: str, limit: int = 100) -> list[dict]: + """Get all nodes of a given type.""" + scoped_label = f"__{label}__{self._tenant_id}__" if self._tenant_id else label + return self.query(f"MATCH (n:{scoped_label}) RETURN n LIMIT $limit", {"limit": limit}) + + def get_relationships(self, source_label: str, rel_type: str, target_label: str, limit: int = 100) -> list[dict]: + """Get relationships between typed nodes.""" + src = f"__{source_label}__{self._tenant_id}__" if self._tenant_id else source_label + tgt = f"__{target_label}__{self._tenant_id}__" if self._tenant_id else target_label + return self.query( + f"MATCH (a:{src})-[r:{rel_type}]->(b:{tgt}) RETURN a, r, b LIMIT $limit", + {"limit": limit}, + ) + + def find_by_property(self, label: str, key: str, value) -> list[dict]: + """Find nodes by property value.""" + scoped_label = f"__{label}__{self._tenant_id}__" if self._tenant_id else label + return self.query(f"MATCH (n:{scoped_label} {{{key}: $val}}) RETURN n", {"val": value}) diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/schema/__init__.py new file mode 100644 index 00000000..cb40c36f --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/__init__.py @@ -0,0 +1,62 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Package for Document Graph Operations. + +This package provides the core schema definitions and utilities for working with +document graph ETL (Extract, Transform, Load) operations. It includes models for +defining extraction sources, transformation rules, and graph loading configurations. + +The schema package contains: +- ETL schema models for defining document processing pipelines +- Schema providers for loading schemas from various sources +- Schema discovery for automatic schema inference +- Utilities for schema serialization and deserialization +- Static schema definitions for testing and default configurations + +These components work together to provide a flexible and extensible way to define +how documents are processed and loaded into graph databases. +""" + +# Core schema models +from .etl_schema_model import ETLSchema, ExtractConfig, TransformConfig, LoadConfig +from .document_graph_schema import ( + ChunkingConfig, MetadataMapping, EntityExtractionConfig, NormalizeConfig, + NodeDefinition, RelationshipDefinition +) + +# Schema I/O utilities +from .schema_io import save_schema, load_schema + +# Static schema provider +from .static_schema_provider import StaticSchemaProvider + +# Import submodules for convenience +from . import providers +from . import discovery + +__all__ = [ + # Core schema models + 'ETLSchema', + 'ExtractConfig', + 'TransformConfig', + 'LoadConfig', + 'ChunkingConfig', + 'MetadataMapping', + 'EntityExtractionConfig', + 'NormalizeConfig', + 'NodeDefinition', + 'RelationshipDefinition', + + # Schema I/O utilities + 'save_schema', + 'load_schema', + + # Utilities + 'StaticSchemaProvider', + + # Submodules + 'providers', + 'discovery' +] + diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/__init__.py new file mode 100644 index 00000000..7e3ace48 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/__init__.py @@ -0,0 +1,42 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Discovery Module for Document Graph Operations. + +This module provides schema discovery capabilities for various file formats, +allowing automatic inference of ETL schemas from source data files. +""" + +from .schema_discovery_base import SchemaDiscoveryProvider +from .tabular_discovery_base import TabularSchemaDiscoveryProvider +from .csv_discovery_provider import CSVSchemaDiscoveryProvider +from .excel_discovery_provider import ExcelSchemaDiscoveryProvider +from .json_discovery_provider import JSONSchemaDiscoveryProvider +from .parquet_discovery_provider import ParquetSchemaDiscoveryProvider +from .xml_discovery_provider import XMLSchemaDiscoveryProvider +from .yaml_discovery_provider import YAMLSchemaDiscoveryProvider +from .schema_discovery_registry import get_discovery_provider, DISCOVERY_REGISTRY +from .schema_discovery_registry_class import SchemaDiscoveryRegistry + +# Aliases for notebook 24 compatibility +CSVDiscoveryProvider = CSVSchemaDiscoveryProvider +JSONDiscoveryProvider = JSONSchemaDiscoveryProvider +ExcelDiscoveryProvider = ExcelSchemaDiscoveryProvider + +__all__ = [ + 'SchemaDiscoveryProvider', + 'TabularSchemaDiscoveryProvider', + 'CSVSchemaDiscoveryProvider', + 'ExcelSchemaDiscoveryProvider', + 'JSONSchemaDiscoveryProvider', + 'ParquetSchemaDiscoveryProvider', + 'XMLSchemaDiscoveryProvider', + 'YAMLSchemaDiscoveryProvider', + 'get_discovery_provider', + 'DISCOVERY_REGISTRY', + # Notebook 24 compatibility + 'CSVDiscoveryProvider', + 'JSONDiscoveryProvider', + 'ExcelDiscoveryProvider', + 'SchemaDiscoveryRegistry' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/csv_discovery_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/csv_discovery_provider.py new file mode 100644 index 00000000..5963e9dc --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/csv_discovery_provider.py @@ -0,0 +1,145 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CSV discovery provider — discovers schema from CSV file structure.""" + +import logging +from typing import List, Dict, Any +import pandas as pd + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ( + ETLSchema, ExtractConfig, TransformConfig, LoadConfig, + ChunkingConfig, MetadataMapping, EntityExtractionConfig, + NormalizeConfig, NodeDefinition, RelationshipDefinition +) +from graphrag_toolkit.document_graph.schema.discovery.tabular_discovery_base import TabularSchemaDiscoveryProvider + +logger = logging.getLogger(__name__) + + +class CSVSchemaDiscoveryProvider(TabularSchemaDiscoveryProvider): + """ + Stable CSV schema discovery that ignores batching and reads only headers. + + Key design choices: + - Remove batching args (skiprows/nrows) so schema isn't batch-dependent. + - Read only header row (nrows=0) to discover columns quickly & deterministically. + - If header=None, sample one row to count fields and synthesize column names. + - Keep on_bad_lines='skip' resilient by forcing engine='python' when needed. + """ + + _BATCH_ARG_KEYS = {"skiprows", "nrows"} # anything that could make schema batch-specific + + def _sanitize_read_args(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Return a copy of args safe for schema discovery (no batching).""" + clean = dict(args or {}) + # Pull out discovery flags + clean.pop("log_on_extract", None) + + # Remove batch-affecting args + for k in list(clean.keys()): + if k in self._BATCH_ARG_KEYS: + clean.pop(k, None) + + # Be resilient by default + clean.setdefault("on_bad_lines", "skip") + + # Ensure pandas engine compatibility for on_bad_lines='skip' + if str(clean.get("on_bad_lines")).lower() == "skip": + # Pandas requires engine='python' for on_bad_lines='skip' unless already set + engine = clean.get("engine", "python") + if engine.lower() != "python": + logger.warning( + "on_bad_lines='skip' requires engine='python' for reliable behavior during discovery; " + "overriding engine to 'python'." + ) + clean["engine"] = "python" + + # Avoid dtype surprises during discovery + # (Don't set dtype here; let pandas leave them as object when nrows=0) + return clean + + def _discover_headers(self, read_args: Dict[str, Any]) -> List[str]: + """ + Determine headers deterministically: + - If header is provided (e.g., 0), read with nrows=0 and use df.columns. + - If header=None, read a single row with header=None to count fields and synthesize names. + """ + # Case 1: header row exists (header is 0 or an int/list) + header_arg = read_args.get("header", 0) + try: + # Fast path: read only the header row structure + discovery_args = dict(read_args) + discovery_args["nrows"] = 0 + df = pd.read_csv(self.source, **discovery_args) + headers = list(df.columns) + + # If user explicitly said header=None, columns will be RangeIndex([]) with nrows=0. + # In that case, synthesize names from a one-row sample. + if header_arg is None or len(headers) == 0: + sample_args = dict(read_args) + sample_args["header"] = None + sample_args["nrows"] = 1 + sample_df = pd.read_csv(self.source, **sample_args) + n_cols = sample_df.shape[1] + headers = [f"column_{i}" for i in range(n_cols)] + except Exception as e: + raise ErrorHandler.csv_parse_error(str(self.source), e) + + if not headers: + raise ValueError("Validation error") + + return headers + + def discover_schema(self) -> ETLSchema: + """ + Discover and return an ETL schema from a CSV file, ignoring batching. + """ + if not self.source.exists(): + raise FileNotFoundError(f"File not found") + + raw_args = dict(self.args or {}) + log_enabled = bool(raw_args.get("log_on_extract", False)) + + # Build safe args for discovery + read_args = self._sanitize_read_args(raw_args) + + # Optional log file + log_path = self.source.with_suffix(self.source.suffix + ".log") if log_enabled else None + + try: + headers: List[str] = self._discover_headers(read_args) + logger.info(f"Discovered columns from CSV: {headers}") + + if log_path: + with log_path.open("w", encoding="utf-8") as log_file: + log_file.write(f"CSV Discovery Log for {self.source.name}\n") + log_file.write(f"Discovered columns: {headers}\n") + # Show the sanitized args we actually used (without batching) + log_file.write(f"Applied pandas.read_csv args (sanitized): {read_args}\n") + + except Exception as e: + if log_path: + with log_path.open("a", encoding="utf-8") as log_file: + log_file.write("\nException encountered during CSV parsing:\n") + log_file.write(str(e)) + raise + + # Construct a stable ETL schema using discovered headers only + return ETLSchema( + schema_id=f"discovered-{self.source.stem}", + description=f"Discovered ETL schema from CSV file: {self.source.name}", + extract=ExtractConfig(source_type="file", file_type="csv"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="fixed_length", min_length=100), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig() + ), + load=LoadConfig( + document_node=NodeDefinition(type="Document", fields=headers), + section_node=NodeDefinition(type="Row", fields=headers), + relationships=[ + RelationshipDefinition(type="contains", source="document_id", target="row_id") + ] + ) + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/excel_discovery_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/excel_discovery_provider.py new file mode 100644 index 00000000..1a3c7a7b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/excel_discovery_provider.py @@ -0,0 +1,134 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Excel Schema Discovery Provider Module for Document Graph Operations. + +This module provides a schema discovery provider for Excel files (.xlsx, .xls). +It analyzes Excel files to infer their structure and generate appropriate ETL schemas +for processing and loading the data into a document graph. + +The module includes the following components: +- ExcelSchemaDiscoveryProvider: Schema discovery provider for Excel files + +The Excel schema discovery provider uses pandas to read and analyze Excel files, +extracting column headers and other metadata to create an ETL schema. It supports +various Excel parsing options through the args parameter, which is passed directly +to pandas.read_excel. + +Usage: + # Get an Excel discovery provider for a specific file + from pathlib import Path + from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_registry import get_discovery_provider + + file_path = Path("data/sample.xlsx") + provider = get_discovery_provider(file_path) + + # Or create it directly with custom arguments + from graphrag_toolkit.document_graph.schema.discovery.excel_discovery_provider import ExcelSchemaDiscoveryProvider + + provider = ExcelSchemaDiscoveryProvider( + source=file_path, + args={"sheet_name": "Sheet1", "header": 0} + ) + + # Discover the schema + schema = provider.discover_schema() + + # Use the schema for ETL operations + print(f"Discovered schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + + +# Import custom logging to ensure configuration is available + + +import logging +from typing import List +import pandas as pd + +from graphrag_toolkit.document_graph.schema.etl_schema_model import * +from graphrag_toolkit.document_graph.schema.discovery.tabular_discovery_base import TabularSchemaDiscoveryProvider + + +logger = logging.getLogger(__name__) + +class ExcelSchemaDiscoveryProvider(TabularSchemaDiscoveryProvider): + """ + Schema discovery provider for Excel files (.xlsx, .xls). + + This class analyzes Excel files to infer their structure and generate appropriate + ETL schemas for processing and loading the data into a document graph. It uses + pandas to read and analyze Excel files, extracting column headers and other metadata. + + The provider supports various Excel parsing options through the args parameter, + which is passed directly to pandas.read_excel. This allows for customization of + the Excel parsing process, such as specifying sheet names, header rows, and + other pandas.read_excel options. + + Attributes: + source (Path): The path to the Excel file to analyze. + args (Dict[str, Any]): Optional arguments that control the discovery process. + These arguments are passed directly to pandas.read_excel. + + Config options: + - sheet_name (str or int): Name or index of the sheet to read. Defaults to 0. + - header (int): Row to use for the column labels. Defaults to 0. + - Any other argument accepted by pandas.read_excel. + """ + + def discover_schema(self) -> ETLSchema: + """ + Discover and return an ETL schema from an Excel file. + + This method reads the Excel file using pandas, extracts the column headers, + and generates an appropriate ETL schema. It includes error handling to + manage issues that might arise during the Excel parsing process. + + The method performs the following steps: + 1. Validates that the source file exists + 2. Reads the Excel file using pandas.read_excel with the provided arguments + 3. Extracts the column headers + 4. Validates that the headers are not empty + 5. Logs the discovered columns + 6. Constructs and returns an ETLSchema object + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the Excel file and how it should be processed. + + Raises: + FileNotFoundError: If the specified source file does not exist. + ValueError: If the Excel file has no columns or cannot be parsed. + Exception: Any exception raised by pandas.read_excel. + """ + if not self.source.exists(): + raise FileNotFoundError(f"File not found") + + try: + df = pd.read_excel(self.source, **self.args) + headers: List[str] = list(df.columns) + if not headers: + raise ValueError("Validation error") + logger.info(f"Discovered columns from Excel: {headers}") + except Exception as e: + raise ErrorHandler.csv_parse_error(str(self.source), e) + + return ETLSchema( + schema_id=f"discovered-{self.source.stem}", + description=f"Discovered ETL schema from Excel file: {self.source.name}", + extract=ExtractConfig(source_type="file", file_type="excel"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="fixed_length", min_length=100), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig() + ), + load=LoadConfig( + document_node=NodeDefinition(type="Document", fields=headers), + section_node=NodeDefinition(type="Row", fields=headers), + relationships=[ + RelationshipDefinition(type="contains", source="document_id", target="row_id") + ] + ) + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/json_discovery_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/json_discovery_provider.py new file mode 100644 index 00000000..280486d9 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/json_discovery_provider.py @@ -0,0 +1,144 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON Schema Discovery Provider Module for Document Graph Operations. + +This module provides a schema discovery provider for JSON (JavaScript Object Notation) files. +It analyzes JSON files to infer their structure and generate appropriate ETL schemas +for processing and loading the data into a document graph. + +The module includes the following components: +- JSONSchemaDiscoveryProvider: Schema discovery provider for JSON files + +The JSON schema discovery provider reads and analyzes JSON files, extracting object keys +from the first record (or the entire object if it's not a list) to create an ETL schema. +It handles both array-based JSON files (containing multiple records) and single-object +JSON files. + +Usage: + # Get a JSON discovery provider for a specific file + from pathlib import Path + from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_registry import get_discovery_provider + + file_path = Path("data/sample.json") + provider = get_discovery_provider(file_path) + + # Or create it directly + from graphrag_toolkit.document_graph.schema.discovery.json_discovery_provider import JSONSchemaDiscoveryProvider + + provider = JSONSchemaDiscoveryProvider(source=file_path) + + # Discover the schema + schema = provider.discover_schema() + + # Use the schema for ETL operations + print(f"Discovered schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + + +# Import custom logging to ensure configuration is available + + +import json +import logging +from typing import List + +from graphrag_toolkit.document_graph.schema.etl_schema_model import * +from graphrag_toolkit.document_graph.schema.discovery.tabular_discovery_base import TabularSchemaDiscoveryProvider + + +logger = logging.getLogger(__name__) + +class JSONSchemaDiscoveryProvider(TabularSchemaDiscoveryProvider): + """ + Schema discovery provider for JSON (JavaScript Object Notation) files. + + This class analyzes JSON files to infer their structure and generate appropriate + ETL schemas for processing and loading the data into a document graph. It reads + JSON files and extracts object keys to determine the fields that should be included + in the ETL schema. + + The provider handles both array-based JSON files (containing multiple records) + and single-object JSON files. For array-based files, it examines the first record + to determine the schema. For single-object files, it uses the keys of the object. + + Attributes: + source (Path): The path to the JSON file to analyze. + args (Dict[str, Any]): Optional arguments that control the discovery process. + Currently, no specific arguments are used for JSON discovery. + """ + + def discover_schema(self) -> ETLSchema: + """ + Discover and return an ETL schema from a JSON file. + + This method reads the JSON file, extracts the keys from the first record + (or the entire object if it's not a list), and generates an appropriate + ETL schema. It includes error handling to manage issues that might arise + during the JSON parsing process. + + The method performs the following steps: + 1. Validates that the source file exists + 2. Reads the JSON file using json.load + 3. Determines if the data is an array or a single object + 4. Extracts the keys from the first record or the entire object + 5. Validates that the keys are not empty + 6. Logs the discovered fields + 7. Constructs and returns an ETLSchema object + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the JSON file and how it should be processed. + + Raises: + FileNotFoundError: If the specified source file does not exist. + ValueError: If the JSON file has no keys or cannot be parsed. + Exception: Any exception raised during JSON parsing. + """ + if not self.source.exists(): + raise FileNotFoundError(f"File not found") + + try: + with self.source.open("r", encoding="utf-8") as f: + # Try regular JSON first + try: + data = json.load(f) + first_record = data[0] if isinstance(data, list) else data + except json.JSONDecodeError: + # If regular JSON fails, try JSONL (newline-delimited JSON) + logger.debug("Regular JSON failed, trying JSONL format") + f.seek(0) # Reset file pointer + first_line = f.readline().strip() + if first_line: + first_record = json.loads(first_line) + else: + raise ValueError("Empty JSONL file") + + headers: List[str] = list(first_record.keys()) + + if not headers: + raise ValueError("Validation error") + + logger.info(f"Discovered fields from JSON/JSONL: {headers}") + except Exception as e: + raise ErrorHandler.csv_parse_error(str(self.source), e) + + return ETLSchema( + schema_id=f"discovered-{self.source.stem}", + description=f"Discovered ETL schema from JSON file: {self.source.name}", + extract=ExtractConfig(source_type="file", file_type="json"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="fixed_length", min_length=100), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig() + ), + load=LoadConfig( + document_node=NodeDefinition(type="Document", fields=headers), + section_node=NodeDefinition(type="Row", fields=headers), + relationships=[ + RelationshipDefinition(type="contains", source="document_id", target="row_id") + ] + ) + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/parquet_discovery_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/parquet_discovery_provider.py new file mode 100644 index 00000000..e8cb4561 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/parquet_discovery_provider.py @@ -0,0 +1,138 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Parquet Schema Discovery Provider Module for Document Graph Operations. + +This module provides a schema discovery provider for Parquet files, a columnar storage +format commonly used in big data processing. It analyzes Parquet files to infer their +structure and generate appropriate ETL schemas for processing and loading the data +into a document graph. + +The module includes the following components: +- ParquetSchemaDiscoveryProvider: Schema discovery provider for Parquet files + +The Parquet schema discovery provider uses pandas to read and analyze Parquet files, +extracting column headers and other metadata to create an ETL schema. It supports +various Parquet parsing options through the args parameter, which is passed directly +to pandas.read_parquet. + +Usage: + # Get a Parquet discovery provider for a specific file + from pathlib import Path + from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_registry import get_discovery_provider + + file_path = Path("data/sample.parquet") + provider = get_discovery_provider(file_path) + + # Or create it directly with custom arguments + from graphrag_toolkit.document_graph.schema.discovery.parquet_discovery_provider import ParquetSchemaDiscoveryProvider + + provider = ParquetSchemaDiscoveryProvider( + source=file_path, + args={"columns": ["id", "name", "value"]} + ) + + # Discover the schema + schema = provider.discover_schema() + + # Use the schema for ETL operations + print(f"Discovered schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + + +# Import custom logging to ensure configuration is available + + +import logging +from typing import List +import pandas as pd + +from graphrag_toolkit.document_graph.schema.etl_schema_model import * +from graphrag_toolkit.document_graph.schema.discovery.tabular_discovery_base import TabularSchemaDiscoveryProvider + + +logger = logging.getLogger(__name__) + +class ParquetSchemaDiscoveryProvider(TabularSchemaDiscoveryProvider): + """ + Schema discovery provider for Parquet files. + + This class analyzes Parquet files to infer their structure and generate appropriate + ETL schemas for processing and loading the data into a document graph. It uses + pandas to read and analyze Parquet files, extracting column headers and other metadata. + + Parquet is a columnar storage format commonly used in big data processing. It is + designed for efficient data compression and encoding schemes, making it particularly + well-suited for query performance and minimizing I/O operations. + + The provider supports various Parquet parsing options through the args parameter, + which is passed directly to pandas.read_parquet. This allows for customization of + the Parquet parsing process, such as specifying which columns to read. + + Attributes: + source (Path): The path to the Parquet file to analyze. + args (Dict[str, Any]): Optional arguments that control the discovery process. + These arguments are passed directly to pandas.read_parquet. + + Config options: + - columns (List[str]): List of columns to read. If not provided, all columns are read. + - Any other argument accepted by pandas.read_parquet. + """ + + def discover_schema(self) -> ETLSchema: + """ + Discover and return an ETL schema from a Parquet file. + + This method reads the Parquet file using pandas, extracts the column headers, + and generates an appropriate ETL schema. It includes error handling to + manage issues that might arise during the Parquet parsing process. + + The method performs the following steps: + 1. Validates that the source file exists + 2. Reads the Parquet file using pandas.read_parquet with the provided arguments + 3. Extracts the column headers + 4. Validates that the headers are not empty + 5. Logs the discovered columns + 6. Constructs and returns an ETLSchema object + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the Parquet file and how it should be processed. + + Raises: + FileNotFoundError: If the specified source file does not exist. + ValueError: If the Parquet file has no columns or cannot be parsed. + Exception: Any exception raised by pandas.read_parquet. + """ + if not self.source.exists(): + raise FileNotFoundError(f"File not found") + + try: + df = pd.read_parquet(self.source, **self.args) + headers: List[str] = list(df.columns) + if not headers: + raise ValueError("Validation error") + + logger.info(f"Discovered columns from Parquet: {headers}") + except Exception as e: + raise ErrorHandler.csv_parse_error(str(self.source), e) + + return ETLSchema( + schema_id=f"discovered-{self.source.stem}", + description=f"Discovered ETL schema from Parquet file: {self.source.name}", + extract=ExtractConfig(source_type="file", file_type="parquet"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="fixed_length", min_length=100), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig() + ), + load=LoadConfig( + document_node=NodeDefinition(type="Document", fields=headers), + section_node=NodeDefinition(type="Row", fields=headers), + relationships=[ + RelationshipDefinition(type="contains", source="document_id", target="row_id") + ] + ) + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_base.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_base.py new file mode 100644 index 00000000..76aaf379 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_base.py @@ -0,0 +1,115 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Discovery Base Module for Document Graph Operations. + +This module provides the foundation for schema discovery in the document graph system. +It defines the abstract base class that all schema discovery providers must implement. +Schema discovery is the process of inferring a structured ETL schema from raw source data, +which is essential for properly processing and loading data into the document graph. + +The module includes the following components: +- SchemaDiscoveryProvider: Abstract base class for all schema discovery providers + +Schema discovery providers are responsible for examining source data files and +determining their structure, field types, and other metadata needed to create +a valid ETL schema. This schema is then used to guide the extraction, transformation, +and loading of data into the document graph. + +Usage: + # Get a discovery provider for a specific file + from pathlib import Path + from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_registry import get_discovery_provider + + file_path = Path("data/sample.csv") + provider = get_discovery_provider(file_path) + + # Discover the schema + schema = provider.discover_schema() + + # Use the schema for ETL operations + print(f"Discovered schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Union, Dict, Any, Optional + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema + + +class SchemaDiscoveryProvider(ABC): + """ + Base class for all schema discovery providers. + + This abstract base class defines the interface that all schema discovery providers + must implement. Schema discovery providers are responsible for examining source + data files and inferring a structured ETL schema that can be used for document + graph operations. + + The provider analyzes the structure, field types, and other metadata of the source + data to create an appropriate ETL schema. This schema is then used to guide the + extraction, transformation, and loading of data into the document graph. + + Subclasses should implement the discover_schema method to provide format-specific + schema discovery logic. + + Attributes: + source (Path): The path to the source data file to analyze. + args (Dict[str, Any]): Optional arguments that control the discovery process. + These arguments are format-specific and are passed to the underlying + data reading functions. + """ + + def __init__(self, source: Union[str, Path], args: Optional[Dict[str, Any]] = None): + """ + Initialize the schema discovery provider. + + This constructor sets up the provider with the source file path and any + format-specific arguments needed for the discovery process. It validates + that the source file exists before proceeding. + + Args: + source (Union[str, Path]): The path to the source data file to analyze. + Can be provided as a string or a Path object. + args (Optional[Dict[str, Any]]): Optional arguments that control the + discovery process. These arguments are format-specific and are + passed to the underlying data reading functions. Defaults to None. + + Raises: + FileNotFoundError: If the specified source file does not exist. + """ + self.source = Path(source) + self.args = args or {} + + if not self.source.exists(): + raise FileNotFoundError(f"File not found") + + @abstractmethod + def discover_schema(self) -> ETLSchema: + """ + Discover and return an ETL schema from the source data. + + This abstract method must be implemented by subclasses to provide + format-specific schema discovery logic. The implementation should + analyze the source data file and infer an appropriate ETL schema + that captures its structure and field types. + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the source data and how it should be processed. + + Raises: + Various exceptions: Implementations may raise format-specific exceptions + when the source data cannot be parsed or does not meet expected + format requirements. + """ + pass diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_registry.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_registry.py new file mode 100644 index 00000000..e9271164 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_registry.py @@ -0,0 +1,118 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Discovery Registry Module for Document Graph Operations. + +This module provides a registry system for schema discovery providers, allowing +the appropriate provider to be selected based on file extensions. It serves as +the entry point for schema discovery operations, abstracting away the details of +which provider to use for a given file type. + +The module includes the following components: +- DISCOVERY_REGISTRY: A dictionary mapping file extensions to provider classes +- get_discovery_provider: Function to get the appropriate provider for a file + +The registry system simplifies the process of discovering ETL schemas from various +file formats by automatically selecting the correct provider based on the file extension. +This allows client code to work with a consistent interface regardless of the +underlying file format. + +Supported file formats: +- CSV (.csv): Comma-separated values files +- Parquet (.parquet): Columnar storage format files +- JSON (.json): JavaScript Object Notation files +- Excel (.xlsx, .xls): Microsoft Excel spreadsheet files +- XML (.xml): eXtensible Markup Language files +- YAML (.yaml, .yml): YAML Ain't Markup Language files + +Usage: + # Get a discovery provider for a specific file + from pathlib import Path + from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_registry import get_discovery_provider + + # The appropriate provider is selected based on the file extension + csv_provider = get_discovery_provider("data/sample.csv") + excel_provider = get_discovery_provider("data/sample.xlsx") + json_provider = get_discovery_provider("data/sample.json") + parquet_provider = get_discovery_provider("data/sample.parquet") + + # Discover the schema using the provider + schema = csv_provider.discover_schema() + + # Use the schema for ETL operations + print(f"Discovered schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +from pathlib import Path +from typing import Type, Dict, Optional, Any + +from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_base import SchemaDiscoveryProvider +from graphrag_toolkit.document_graph.schema.discovery.csv_discovery_provider import CSVSchemaDiscoveryProvider +from graphrag_toolkit.document_graph.schema.discovery.parquet_discovery_provider import ParquetSchemaDiscoveryProvider +from graphrag_toolkit.document_graph.schema.discovery.json_discovery_provider import JSONSchemaDiscoveryProvider +from graphrag_toolkit.document_graph.schema.discovery.excel_discovery_provider import ExcelSchemaDiscoveryProvider +from graphrag_toolkit.document_graph.schema.discovery.xml_discovery_provider import XMLSchemaDiscoveryProvider +from graphrag_toolkit.document_graph.schema.discovery.yaml_discovery_provider import YAMLSchemaDiscoveryProvider + +DISCOVERY_REGISTRY: Dict[str, Type[SchemaDiscoveryProvider]] = { + "csv": CSVSchemaDiscoveryProvider, + "parquet": ParquetSchemaDiscoveryProvider, + "json": JSONSchemaDiscoveryProvider, + "xlsx": ExcelSchemaDiscoveryProvider, + "xls": ExcelSchemaDiscoveryProvider, + "xml": XMLSchemaDiscoveryProvider, + "yaml": YAMLSchemaDiscoveryProvider, + "yml": YAMLSchemaDiscoveryProvider, +} + + +def get_discovery_provider(file_path: Path, args: Optional[Dict[str, Any]] = None) -> SchemaDiscoveryProvider: + """ + Get the appropriate schema discovery provider for a given file. + + This function selects the appropriate schema discovery provider based on the + file extension of the provided file path. It uses the DISCOVERY_REGISTRY + dictionary to map file extensions to provider classes. + + The function handles the instantiation of the provider with the file path + and any provided arguments, returning a ready-to-use provider instance. + + Args: + file_path (Path): The path to the file for which to get a discovery provider. + The file extension determines which provider is selected. + args (Optional[Dict[str, Any]]): Optional arguments to pass to the provider. + These arguments are format-specific and are passed to the underlying + data reading functions. Defaults to None. + + Returns: + SchemaDiscoveryProvider: An instance of the appropriate schema discovery + provider for the given file type, initialized with the provided file + path and arguments. + + Raises: + ValueError: If no discovery provider is registered for the file extension. + + Examples: + >>> from pathlib import Path + >>> provider = get_discovery_provider(Path("data/sample.csv")) + >>> schema = provider.discover_schema() + + >>> # With custom arguments + >>> provider = get_discovery_provider( + ... Path("data/sample.csv"), + ... args={"delimiter": ",", "encoding": "utf-8"} + ... ) + >>> schema = provider.discover_schema() + """ + ext = file_path.suffix.lower().lstrip(".") + provider_class = DISCOVERY_REGISTRY.get(ext) + if not provider_class: + raise ValueError(f"No discovery provider registered for extension: {ext}") + return provider_class(source=file_path, args=args or {}) diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_registry_class.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_registry_class.py new file mode 100644 index 00000000..3ebf0a5b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/schema_discovery_registry_class.py @@ -0,0 +1,67 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Discovery Registry Class for notebook 24 compatibility.""" + +from typing import Dict, Type +from .schema_discovery_base import SchemaDiscoveryProvider + + +class SchemaDiscoveryRegistry: + """ + Manages the registration and retrieval of schema discovery providers. + + This class provides functionality for registering schema discovery providers, + retrieving providers by name, and listing all available providers. It serves as + a central registry for accessing different schema discovery mechanisms. + """ + + def __init__(self): + self._providers: Dict[str, Type[SchemaDiscoveryProvider]] = {} + + def register_provider(self, name: str, provider_class: Type[SchemaDiscoveryProvider]): + """ + Registers a schema discovery provider with the specified name. + + This method associates a provider class with a given name, allowing the + provider to be referenced and used later. + + Args: + name: The name to register the provider under. + provider_class: The class of the schema discovery provider to register. + """ + self._providers[name] = provider_class + + def get_provider(self, name: str) -> Type[SchemaDiscoveryProvider]: + """ + Retrieve a schema discovery provider by its name. + + This method looks up and returns a schema discovery provider associated with + the provided name. If the name does not correspond to any registered provider, + a KeyError is raised. + + Parameters: + name (str): The name of the schema discovery provider to retrieve. + + Raises: + KeyError: If the specified provider name is not found among registered providers. + + Returns: + Type[SchemaDiscoveryProvider]: The schema discovery provider associated with + the specified name. + """ + if name not in self._providers: + raise KeyError(f"Unknown discovery provider: {name}") + return self._providers[name] + + def list_providers(self) -> list: + """ + Returns a list of registered providers. + + This method retrieves all the registered providers from the internal + providers dictionary and returns their keys as a list. + + Returns: + list: List containing the names of all registered providers. + """ + return list(self._providers.keys()) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/tabular_discovery_base.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/tabular_discovery_base.py new file mode 100644 index 00000000..3fb81234 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/tabular_discovery_base.py @@ -0,0 +1,86 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tabular Discovery Base Module for Document Graph Operations. + +This module provides a base class for schema discovery providers that work with +tabular data formats such as CSV, Excel, Parquet, and similar structured data. +It extends the core schema discovery functionality with tabular-specific features. + +The module includes the following components: +- TabularSchemaDiscoveryProvider: Base class for tabular schema discovery providers + +Tabular schema discovery providers are specialized for handling data formats that +have a row-column structure. They share common patterns for extracting column headers, +inferring data types, and creating appropriate ETL schemas for tabular data. + +This module serves as an intermediate layer between the abstract SchemaDiscoveryProvider +and concrete implementations for specific tabular formats like CSV, Excel, and Parquet. + +Usage: + # This class is not typically used directly, but through its subclasses + from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_registry import get_discovery_provider + + # Get a provider for a specific tabular format + csv_provider = get_discovery_provider("data/sample.csv") + excel_provider = get_discovery_provider("data/sample.xlsx") + parquet_provider = get_discovery_provider("data/sample.parquet") + + # All these providers inherit from TabularSchemaDiscoveryProvider + # and share common behavior for tabular data +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from pathlib import Path +from typing import Optional, Dict, Any + +from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_base import SchemaDiscoveryProvider + + +class TabularSchemaDiscoveryProvider(SchemaDiscoveryProvider): + """ + Base class for tabular schema discovery providers like CSV, Excel, and Parquet. + + This class extends the SchemaDiscoveryProvider to provide common functionality + for tabular data formats. Tabular data is characterized by a row-column structure + where each row represents a record and each column represents a field. + + Concrete subclasses implement format-specific logic for reading and analyzing + tabular data sources like CSV, Excel, and Parquet files. These subclasses share + common patterns for extracting column headers, inferring data types, and creating + appropriate ETL schemas. + + The TabularSchemaDiscoveryProvider simplifies the implementation of concrete + providers by handling common aspects of tabular data processing, allowing + subclasses to focus on format-specific details. + + Attributes: + source (Path): The path to the tabular data file to analyze. + args (Dict[str, Any]): Optional arguments that control the discovery process. + These arguments are format-specific and are passed to the underlying + data reading functions (e.g., pandas.read_csv, pandas.read_excel). + """ + def __init__(self, source: Path, args: Optional[Dict[str, Any]] = None): + """ + Initialize the tabular schema discovery provider. + + This constructor sets up the provider with the source file path and any + format-specific arguments needed for the discovery process. Unlike the + parent class, it does not validate that the source file exists, leaving + that responsibility to the concrete subclasses. + + Args: + source (Path): The path to the tabular data file to analyze. + args (Optional[Dict[str, Any]]): Optional arguments that control the + discovery process. These arguments are format-specific and are + passed to the underlying data reading functions (e.g., pandas.read_csv, + pandas.read_excel). Defaults to None. + """ + self.source = Path(source) + self.args = args or {} diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/xml_discovery_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/xml_discovery_provider.py new file mode 100644 index 00000000..3cb2eadf --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/xml_discovery_provider.py @@ -0,0 +1,172 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""XML Schema Discovery Provider Module for Document Graph Operations. + +This module provides a schema discovery provider for XML (eXtensible Markup Language) files. +It analyzes XML files to infer their structure and generate appropriate ETL schemas +for processing and loading the data into a document graph. + +The module includes the following components: +- XMLSchemaDiscoveryProvider: Schema discovery provider for XML files + +The XML schema discovery provider reads and analyzes XML files, extracting element names +and attributes from the root elements to create an ETL schema. It handles both simple +XML structures and more complex nested structures by flattening them into field names. + +Usage: + # Get an XML discovery provider for a specific file + from pathlib import Path + from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_registry import get_discovery_provider + + file_path = Path("data/sample.xml") + provider = get_discovery_provider(file_path) + + # Or create it directly + from graphrag_toolkit.document_graph.schema.discovery.xml_discovery_provider import XMLSchemaDiscoveryProvider + + provider = XMLSchemaDiscoveryProvider(source=file_path) + + # Discover the schema + schema = provider.discover_schema() + + # Use the schema for ETL operations + print(f"Discovered schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging +import xml.etree.ElementTree as ET +from typing import List, Set + +# Import custom logging to ensure configuration is available + +from graphrag_toolkit.document_graph.schema.etl_schema_model import * +from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_base import SchemaDiscoveryProvider + +logger = logging.getLogger(__name__) + + +class XMLSchemaDiscoveryProvider(SchemaDiscoveryProvider): + """ + Schema discovery provider for XML (eXtensible Markup Language) files. + + This class analyzes XML files to infer their structure and generate appropriate + ETL schemas for processing and loading the data into a document graph. It reads + XML files and extracts element names and attributes to determine the fields that + should be included in the ETL schema. + + The provider handles XML structures by examining the first few elements to + determine the schema. It flattens nested structures into field names using + dot notation (e.g., "parent.child"). + + Attributes: + source (Path): The path to the XML file to analyze. + args (Dict[str, Any]): Optional arguments that control the discovery process. + Currently, no specific arguments are used for XML discovery. + """ + + def discover_schema(self) -> ETLSchema: + """ + Discover and return an ETL schema from an XML file. + + This method reads the XML file, extracts element names and attributes + from the structure, and generates an appropriate ETL schema. It includes + error handling to manage issues that might arise during XML parsing. + + The method performs the following steps: + 1. Validates that the source file exists + 2. Parses the XML file using ElementTree + 3. Extracts field names from elements and attributes + 4. Handles nested structures by flattening them + 5. Validates that fields were discovered + 6. Logs the discovered fields + 7. Constructs and returns an ETLSchema object + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the XML file and how it should be processed. + + Raises: + FileNotFoundError: If the specified source file does not exist. + ValueError: If the XML file has no discoverable fields. + Exception: Any exception raised during XML parsing. + """ + if not self.source.exists(): + raise FileNotFoundError(f"File not found") + + try: + tree = ET.parse(self.source) + root = tree.getroot() + + # Extract field names from XML structure + fields = self._extract_fields_from_element(root) + + if not fields: + raise ValueError("Validation error") + + logger.info(f"Discovered fields from XML: {sorted(fields)}") + + except ET.ParseError as e: + raise ErrorHandler.file_parse_error(str(self.source), "xml", e) + except Exception as e: + raise ErrorHandler.file_parse_error(str(self.source), "xml", e) + + return ETLSchema( + schema_id=f"discovered-{self.source.stem}", + description=f"Discovered ETL schema from XML file: {self.source.name}", + extract=ExtractConfig(source_type="file", file_type="xml"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="fixed_length", min_length=100), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig() + ), + load=LoadConfig( + document_node=NodeDefinition(type="Document", fields=sorted(fields)), + section_node=NodeDefinition(type="Element", fields=sorted(fields)), + relationships=[ + RelationshipDefinition(type="contains", source="document_id", target="element_id") + ] + ) + ) + + def _extract_fields_from_element(self, element: ET.Element, prefix: str = "") -> Set[str]: + """ + Recursively extract field names from an XML element. + + This helper method traverses the XML structure and extracts field names + from element tags and attributes. It handles nested structures by using + dot notation to create hierarchical field names. + + Args: + element (ET.Element): The XML element to analyze. + prefix (str): The prefix to use for nested field names. + + Returns: + Set[str]: A set of field names discovered in the XML structure. + """ + fields = set() + + # Add the element's tag name + tag_name = element.tag.split('}')[-1] if '}' in element.tag else element.tag + field_name = f"{prefix}.{tag_name}" if prefix else tag_name + fields.add(field_name) + + # Add attributes + for attr_name in element.attrib: + attr_field = f"{field_name}.{attr_name}" + fields.add(attr_field) + + # Add text content if present + if element.text and element.text.strip(): + text_field = f"{field_name}.text" + fields.add(text_field) + + # Recursively process child elements (limit depth to avoid excessive nesting) + if len(field_name.split('.')) < 3: # Limit nesting depth + for child in element: + child_fields = self._extract_fields_from_element(child, field_name) + fields.update(child_fields) + + return fields \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/yaml_discovery_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/yaml_discovery_provider.py new file mode 100644 index 00000000..693a8f47 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/discovery/yaml_discovery_provider.py @@ -0,0 +1,182 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""YAML Schema Discovery Provider Module for Document Graph Operations. + +This module provides a schema discovery provider for YAML (YAML Ain't Markup Language) files. +It analyzes YAML files to infer their structure and generate appropriate ETL schemas +for processing and loading the data into a document graph. + +The module includes the following components: +- YAMLSchemaDiscoveryProvider: Schema discovery provider for YAML files + +The YAML schema discovery provider reads and analyzes YAML files, extracting keys +from the structure to create an ETL schema. It handles both simple YAML structures +and more complex nested structures by flattening them into field names. + +Usage: + # Get a YAML discovery provider for a specific file + from pathlib import Path + from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_registry import get_discovery_provider + + file_path = Path("data/sample.yaml") + provider = get_discovery_provider(file_path) + + # Or create it directly + from graphrag_toolkit.document_graph.schema.discovery.yaml_discovery_provider import YAMLSchemaDiscoveryProvider + + provider = YAMLSchemaDiscoveryProvider(source=file_path) + + # Discover the schema + schema = provider.discover_schema() + + # Use the schema for ETL operations + print(f"Discovered schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging +from typing import List, Set, Dict, Any, Union + +# Import custom logging to ensure configuration is available + +from graphrag_toolkit.document_graph.schema.etl_schema_model import * +from graphrag_toolkit.document_graph.schema.discovery.schema_discovery_base import SchemaDiscoveryProvider + +logger = logging.getLogger(__name__) + +try: + import yaml +except ImportError: + yaml = None + + +class YAMLSchemaDiscoveryProvider(SchemaDiscoveryProvider): + """ + Schema discovery provider for YAML (YAML Ain't Markup Language) files. + + This class analyzes YAML files to infer their structure and generate appropriate + ETL schemas for processing and loading the data into a document graph. It reads + YAML files and extracts keys from the structure to determine the fields that + should be included in the ETL schema. + + The provider handles both simple YAML structures and complex nested structures + by flattening them into field names using dot notation (e.g., "parent.child"). + It supports both single-document and multi-document YAML files. + + Attributes: + source (Path): The path to the YAML file to analyze. + args (Dict[str, Any]): Optional arguments that control the discovery process. + Currently, no specific arguments are used for YAML discovery. + """ + + def discover_schema(self) -> ETLSchema: + """ + Discover and return an ETL schema from a YAML file. + + This method reads the YAML file, extracts keys from the structure, + and generates an appropriate ETL schema. It includes error handling + to manage issues that might arise during YAML parsing. + + The method performs the following steps: + 1. Validates that the source file exists and PyYAML is available + 2. Parses the YAML file using yaml.safe_load + 3. Extracts field names from the structure + 4. Handles nested structures by flattening them + 5. Validates that fields were discovered + 6. Logs the discovered fields + 7. Constructs and returns an ETLSchema object + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the YAML file and how it should be processed. + + Raises: + ImportError: If PyYAML is not installed. + FileNotFoundError: If the specified source file does not exist. + ValueError: If the YAML file has no discoverable fields. + Exception: Any exception raised during YAML parsing. + """ + if yaml is None: + raise ImportError("PyYAML is required for YAML schema discovery. Install with: pip install PyYAML") + + if not self.source.exists(): + raise FileNotFoundError(f"File not found") + + try: + with self.source.open("r", encoding="utf-8") as f: + data = yaml.safe_load(f) + + # Extract field names from YAML structure + fields = self._extract_fields_from_data(data) + + if not fields: + raise ValueError("Validation error") + + logger.info(f"Discovered fields from YAML: {sorted(fields)}") + + except yaml.YAMLError as e: + raise ErrorHandler.file_parse_error(str(self.source), "yaml", e) + except Exception as e: + raise ErrorHandler.file_parse_error(str(self.source), "yaml", e) + + return ETLSchema( + schema_id=f"discovered-{self.source.stem}", + description=f"Discovered ETL schema from YAML file: {self.source.name}", + extract=ExtractConfig(source_type="file", file_type="yaml"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="fixed_length", min_length=100), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig() + ), + load=LoadConfig( + document_node=NodeDefinition(type="Document", fields=sorted(fields)), + section_node=NodeDefinition(type="Section", fields=sorted(fields)), + relationships=[ + RelationshipDefinition(type="contains", source="document_id", target="section_id") + ] + ) + ) + + def _extract_fields_from_data(self, data: Any, prefix: str = "") -> Set[str]: + """ + Recursively extract field names from YAML data structure. + + This helper method traverses the YAML data structure and extracts field names + from dictionaries, lists, and scalar values. It handles nested structures by + using dot notation to create hierarchical field names. + + Args: + data (Any): The YAML data to analyze (dict, list, or scalar). + prefix (str): The prefix to use for nested field names. + + Returns: + Set[str]: A set of field names discovered in the YAML structure. + """ + fields = set() + + if isinstance(data, dict): + for key, value in data.items(): + field_name = f"{prefix}.{key}" if prefix else key + fields.add(field_name) + + # Recursively process nested structures (limit depth to avoid excessive nesting) + if len(field_name.split('.')) < 3: # Limit nesting depth + nested_fields = self._extract_fields_from_data(value, field_name) + fields.update(nested_fields) + + elif isinstance(data, list) and data: + # For lists, analyze the first item to determine structure + first_item = data[0] + if isinstance(first_item, dict): + # If it's a list of dictionaries, extract fields from the first dictionary + nested_fields = self._extract_fields_from_data(first_item, prefix) + fields.update(nested_fields) + elif prefix: + fields.add(prefix) + + if prefix: + fields.add(prefix) + + return fields \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/document_graph_schema.py b/document-graph/src/graphrag_toolkit/document_graph/schema/document_graph_schema.py new file mode 100644 index 00000000..77aeab26 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/document_graph_schema.py @@ -0,0 +1,66 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Document Graph Schema Definitions. + +This module provides a centralized import location for all ETL schema models +used in document graph processing. It re-exports the core schema classes from +the etl_schema_model module to provide a clean, unified interface for working +with document graph schemas. + +The exported classes include: +- ETLSchema: The top-level schema definition for ETL operations +- ExtractConfig: Configuration for data extraction from various sources +- TransformConfig: Configuration for document transformation operations +- LoadConfig: Configuration for loading processed documents into a graph +- Supporting classes for specific aspects of the ETL process + +This module simplifies imports by allowing users to import all schema-related +classes from a single location rather than from multiple modules. + +Example: + from graphrag_toolkit.document_graph.schema.document_graph_schema import ( + ETLSchema, ExtractConfig, LoadConfig + ) + + # Create a schema using the imported classes + schema = ETLSchema( + schema_id="example-schema", + extract=ExtractConfig(...), + transform=TransformConfig(...), + load=LoadConfig(...) + ) +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +from typing import List +from .etl_schema_model import ( + ETLSchema, + ExtractConfig, + TransformConfig, + LoadConfig, + ChunkingConfig, + MetadataMapping, + EntityExtractionConfig, + NormalizeConfig, + NodeDefinition, + RelationshipDefinition, +) + +__all__: List[str] = [ + "ETLSchema", + "ExtractConfig", + "TransformConfig", + "LoadConfig", + "ChunkingConfig", + "MetadataMapping", + "EntityExtractionConfig", + "NormalizeConfig", + "NodeDefinition", + "RelationshipDefinition", +] diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/etl_schema_model.py b/document-graph/src/graphrag_toolkit/document_graph/schema/etl_schema_model.py new file mode 100644 index 00000000..2cf652ed --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/etl_schema_model.py @@ -0,0 +1,316 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ETL Schema Model Module for Document Graph Operations. + +This module defines the data models used for ETL (Extract, Transform, Load) operations +in the document graph system. It provides a comprehensive set of Pydantic models that +represent different aspects of the ETL process: + +1. Extract: Models for configuring data extraction from various sources +2. Transform: Models for text chunking, metadata mapping, entity extraction, and normalization +3. Load: Models for defining graph nodes, relationships, and loading configurations + +These models are used throughout the document graph system to define how documents +are processed, transformed, and loaded into graph databases. The models are designed +to be extensible, allowing for custom configurations through the use of Pydantic's +extra fields feature. + +The top-level ETLSchema class combines all these components into a complete schema +that can be serialized to/from JSON and used to configure the entire ETL pipeline. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Optional +from pydantic import BaseModel, ConfigDict + + +# Extract section +class ExtractConfig(BaseModel): + """ + Configuration for data extraction from various sources. + + This class defines the parameters needed to extract data from different source types, + such as S3 buckets, local files, or APIs. It specifies where to find the data and + how to read it. + + Attributes: + source_type (str): The type of data source (e.g., "s3", "html", "csv", "api"). + bucket (Optional[str]): The S3 bucket name for S3 sources. + prefix (Optional[str]): The S3 prefix/folder path for S3 sources. + key (Optional[str]): The specific S3 key or file path. + file_type (Optional[str]): The type of file to extract (e.g., "pdf", "docx"). + reader (Optional[str]): The reader to use for extraction (e.g., "pymupdf"). + delimiter (Optional[str]): The delimiter for CSV or similar files. + encoding (Optional[str]): The character encoding of the source file. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to add source-specific parameters like API tokens or headers. + """ + source_type: str # Allow user-defined types (e.g., "html", "csv", "api") + bucket: Optional[str] = None + prefix: Optional[str] = None + key: Optional[str] = None + file_type: Optional[str] = None + reader: Optional[str] = None + delimiter: Optional[str] = None + encoding: Optional[str] = None + + model_config = ConfigDict(extra="allow") + + +# Transform section +class ChunkingConfig(BaseModel): + """ + Configuration for chunking text during document transformation. + + This class defines how documents should be divided into smaller, manageable chunks + for processing and storage. Different chunking strategies can be used depending on + the document structure and requirements. + + Attributes: + strategy (str): The chunking strategy to use (e.g., "by_heading", "fixed_length", + "by_customer"). + min_length (Optional[int]): The minimum length of a chunk in characters or tokens, + defaulting to 100. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to add strategy-specific parameters like maximum chunk size or + overlap settings. + """ + strategy: str # e.g., "by_heading", "fixed_length", "by_customer" + min_length: Optional[int] = 100 + + model_config = ConfigDict(extra="allow") + + +class MetadataMapping(BaseModel): + """ + Field-level mapping of document metadata to standardized fields. + + This class defines how metadata fields from source documents should be mapped + to standardized fields in the document graph. It allows for consistent metadata + representation across different document types and sources. + + Attributes: + title (Optional[str]): Field path or expression to extract the document title. + author (Optional[str]): Field path or expression to extract the document author. + created_date (Optional[str]): Field path or expression to extract the document creation date. + category (Optional[str]): Field path or expression to extract the document category. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to map any custom metadata fields specific to their document types. + Field values can be direct field names or dot-notation paths (e.g., "metadata.author"). + """ + title: Optional[str] = None + author: Optional[str] = None + created_date: Optional[str] = None + category: Optional[str] = None # e.g., for classification purposes + + model_config = ConfigDict(extra="allow") + + +class EntityExtractionConfig(BaseModel): + """ + Configuration for extracting named entities from document text. + + This class defines how named entities (such as people, organizations, locations) + should be extracted from document text during the transformation phase. It specifies + the extraction method, model to use, and which entity types to extract. + + Attributes: + method (str): The entity extraction method to use (e.g., "ner", "regex", "rule_based"). + model (Optional[str]): The specific model to use for extraction (e.g., "spacy_en_core_web_lg"). + extract_emails (Optional[bool]): Whether to extract email addresses, defaults to False. + extract_dates (Optional[bool]): Whether to extract date references, defaults to False. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to specify additional entity types to extract or method-specific + parameters like confidence thresholds. + """ + method: str # e.g., "ner", "regex", "rule_based" + model: Optional[str] = None + extract_emails: Optional[bool] = False + extract_dates: Optional[bool] = False + + model_config = ConfigDict(extra="allow") + + +class NormalizeConfig(BaseModel): + """ + Configuration for text normalization during document processing. + + This class defines how document text should be normalized during the transformation + phase. Normalization can include removing headers, standardizing whitespace, + standardizing date formats, and other text cleaning operations. + + Attributes: + remove_headers (bool): Whether to remove headers from the text, defaults to True. + collapse_whitespace (bool): Whether to normalize whitespace (removing extra spaces, + newlines, etc.), defaults to True. + standardize_dates (Optional[bool]): Whether to convert dates to a standard format, + defaults to False. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to specify additional normalization operations specific to their + document types or requirements. + """ + remove_headers: bool = True + collapse_whitespace: bool = True + standardize_dates: Optional[bool] = False + + model_config = ConfigDict(extra="allow") + + +class TransformConfig(BaseModel): + """ + Complete configuration for document transformation operations. + + This class combines all the transformation-related configurations into a single + comprehensive configuration for the transform phase of the ETL process. It includes + settings for chunking, metadata mapping, entity extraction, and text normalization. + + Attributes: + chunking (ChunkingConfig): Configuration for how to chunk the document text. + metadata_mapping (MetadataMapping): Configuration for mapping document metadata fields. + entity_extraction (EntityExtractionConfig): Configuration for extracting named entities. + normalize (NormalizeConfig): Configuration for text normalization operations. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to add custom transformation configurations not covered by the + standard components. + """ + chunking: ChunkingConfig + metadata_mapping: MetadataMapping + entity_extraction: EntityExtractionConfig + normalize: NormalizeConfig + + model_config = ConfigDict(extra="allow") + + +# Load section +class NodeDefinition(BaseModel): + """ + Definition of a graph node type and its associated fields. + + This class defines a type of node in the document graph and specifies which fields + should be included in nodes of this type. It is used in the load phase to determine + how to create nodes in the graph database. + + Attributes: + type (str): The type or label of the node (e.g., "DocumentNode", "SectionNode"). + fields (List[str]): List of field names to include in nodes of this type. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to specify additional node properties like indexes or constraints. + """ + type: str + fields: List[str] + + model_config = ConfigDict(extra="allow") + + +class RelationshipDefinition(BaseModel): + """ + Defines a relationship (edge) between two node types in the document graph. + + This class specifies how nodes in the graph should be connected to each other, + defining the type of relationship and which nodes should be the source and target + of the relationship. It is used in the load phase to create edges in the graph database. + + Attributes: + type (str): The type or label of the relationship (e.g., "has_section", "references"). + source (str): The identifier for the source node of the relationship. + target (str): The identifier for the target node of the relationship. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to specify additional relationship properties like weights or + directional constraints. + """ + type: str + source: str + target: str + + model_config = ConfigDict(extra="allow") + + +class LoadConfig(BaseModel): + """ + Configuration for loading processed documents into a graph database. + + This class defines how processed document data should be loaded into a graph database, + specifying the node types, their fields, and the relationships between them. It is + used in the load phase of the ETL process to create the document graph structure. + + Attributes: + document_node (NodeDefinition): Definition of the document node type and its fields. + section_node (NodeDefinition): Definition of the section node type and its fields. + relationships (List[RelationshipDefinition]): List of relationships between nodes + in the graph. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to specify additional loading configurations like batch sizes, + transaction settings, or custom node types beyond the standard document and + section nodes. + """ + document_node: NodeDefinition + section_node: NodeDefinition + relationships: List[RelationshipDefinition] + + model_config = ConfigDict(extra="allow") + + +# Top-level schema +class ETLSchema(BaseModel): + """ + Complete configuration for ETL (Extract, Transform, Load) processing of document graphs. + + This is the top-level schema class that combines all aspects of the ETL process into + a single comprehensive configuration. It includes configurations for extracting data + from sources, transforming the data through various operations, and loading the + processed data into a graph database. + + Attributes: + schema_id (str): A unique identifier for this schema. + description (Optional[str]): A human-readable description of the schema's purpose. + extract (ExtractConfig): Configuration for the extract phase. + transform (TransformConfig): Configuration for the transform phase. + load (LoadConfig): Configuration for the load phase. + + Note: + This class allows additional custom fields through the `extra = "allow"` config, + enabling users to add custom configurations not covered by the standard components. + + Example: + ```python + schema = ETLSchema( + schema_id="pdf-document-schema", + description="Schema for processing PDF documents", + extract=ExtractConfig(source_type="s3", bucket="documents", file_type="pdf"), + transform=TransformConfig(...), + load=LoadConfig(...) + ) + ``` + """ + schema_id: str + description: Optional[str] = None + extract: ExtractConfig + transform: TransformConfig + load: LoadConfig + + model_config = ConfigDict(extra="allow") diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/__init__.py new file mode 100644 index 00000000..5f9443f0 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/__init__.py @@ -0,0 +1,101 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Providers Package for Document Graph Operations. + +This package provides a collection of schema providers for loading and managing ETL schemas +used in document graph operations. Schema providers are responsible for loading schemas +from various sources such as files, databases, and cloud services. + +The package includes the following components: + +Base Classes: +- SchemaProviderBase: Abstract base class for all schema providers +- SchemaProviderConfig: Base configuration model for schema providers +- AWSSchemaProviderConfig: Base configuration model for AWS-backed schema providers + +Factory and Registry: +- SchemaProviderFactory: Factory for creating schema providers based on configuration +- SchemaProviderRegistry: Registry for managing schema provider classes + +Provider Implementations: +- FileSchemaProvider: Provider for loading schemas from local files +- S3SchemaProvider: Provider for loading schemas from Amazon S3 +- CSVSchemaProvider: Provider for generating schemas from CSV files +- JSONSchemaProvider: Provider for generating schemas from JSON files +- ExcelSchemaProvider: Provider for generating schemas from Excel files +- GlueSchemaProvider: Provider for loading schemas from AWS Glue +- ParquetSchemaProvider: Provider for generating schemas from Parquet files +- YAMLSchemaProvider: Provider for generating schemas from YAML files +- XMLSchemaProvider: Provider for generating schemas from XML files + +Usage: + # Create a schema provider using the factory + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "type": "file", + "schema_id": "my_schema", + "connection_config": { + "path": "/path/to/schema.yaml" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") +""" + +# Base classes +from .schema_provider_base import SchemaProviderBase +from .schema_provider_config import SchemaProviderConfig +from .schema_provider_config_aws_base import AWSSchemaProviderConfig + +# Factory and registry +from .schema_provider_factory import SchemaProviderFactory +from .schema_provider_registry import SchemaProviderRegistry + +# Provider implementations +from .file_schema_provider import FileSchemaProvider +from .s3_schema_provider import S3SchemaProvider +from .csv_schema_provider import CSVSchemaProvider +from .json_schema_provider import JSONSchemaProvider +from .excel_schema_provider import ExcelSchemaProvider +from .glue_schema_provider import GlueSchemaProvider +from .parquet_schema_provider import ParquetSchemaProvider +from .yaml_schema_provider import YAMLSchemaProvider +from .xml_schema_provider import XMLSchemaProvider + +__all__ = [ + # Base classes + 'SchemaProviderBase', + 'SchemaProviderConfig', + 'AWSSchemaProviderConfig', + + # Factory and registry + 'SchemaProviderFactory', + 'SchemaProviderRegistry', + + # Provider implementations + 'FileSchemaProvider', + 'S3SchemaProvider', + 'CSVSchemaProvider', + 'JSONSchemaProvider', + 'ExcelSchemaProvider', + 'GlueSchemaProvider', + 'ParquetSchemaProvider', + 'YAMLSchemaProvider', + 'XMLSchemaProvider', + + # Aliases for backward compatibility + 'CSVSchemaProvider', + 'JSONSchemaProvider' +] + +# Aliases +CSVSchemaProvider = CSVSchemaProvider # Already correct name +JSONSchemaProvider = JSONSchemaProvider # Already correct name \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/csv_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/csv_schema_provider.py new file mode 100644 index 00000000..72ff64da --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/csv_schema_provider.py @@ -0,0 +1,158 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CSV Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for CSV (Comma-Separated Values) files. +It discovers and generates ETL schemas from CSV files for processing and loading +the data into a document graph. + +The module includes the following components: +- CSVSchemaProvider: Schema provider for CSV files + +The CSV schema provider uses the CSVSchemaDiscoveryProvider to analyze CSV files +and generate appropriate ETL schemas. It supports loading schemas from CSV files +and optionally saving the discovered schemas to JSON files. + +Usage: + # Get a CSV schema provider for a specific file + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "csv", + "schema_id": "my_csv_schema", + "connection_config": { + "path": "/path/to/data.csv" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging +import csv +import json +from pathlib import Path +from typing import List, Optional, Any + +# Import custom logging to ensure configuration is available + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.discovery.csv_discovery_provider import CSVSchemaDiscoveryProvider + +logger = logging.getLogger(__name__) + + +class CSVSchemaProvider(SchemaProviderBase): + """ + Schema provider for CSV (Comma-Separated Values) files. + + This class provides functionality to discover and generate ETL schemas from CSV files. + It uses the CSVSchemaDiscoveryProvider to analyze CSV files and extract field information + to create appropriate ETL schemas for processing and loading the data into a document graph. + + The provider supports loading schemas directly from CSV files and optionally saving + the discovered schemas to JSON files for later use. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + csv_path (Path): The path to the CSV file to analyze. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize a CSV schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the path to the CSV file in the connection_config. + + Raises: + ValueError: If the path is not provided in the connection_config. + FileNotFoundError: If the specified CSV file does not exist. + """ + self.config = config + path = config.connection_config.get("path") + if not path: + raise ValueError("Validation error") + self.csv_path = Path(path) + + if not self.csv_path.exists(): + raise FileNotFoundError(f"File not found") + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "CSVSchemaProvider": + """ + Factory method to construct a CSV provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + CSVSchemaProvider: A new instance of the CSV schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load (or generate) an ETL schema from a CSV file using schema discovery. + + This method uses the CSVSchemaDiscoveryProvider to analyze the CSV file and + generate an appropriate ETL schema based on the structure and content of the file. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters (not used in this provider). + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the CSV file and how it should be processed. + + Raises: + FileNotFoundError: If the CSV file does not exist. + ValueError: If the CSV file cannot be parsed or has no discoverable fields. + """ + discovery = CSVSchemaDiscoveryProvider(source=self.csv_path) + return discovery.discover_schema() + + def save_schema(self, output_path: Path) -> None: + """ + Save the inferred ETL schema to a JSON file at the given output path. + + This method discovers the schema from the CSV file and saves it to a JSON file + at the specified path. The schema can then be loaded directly using a FileSchemaProvider. + + Args: + output_path (Path): The path where the schema JSON file should be saved. + + Raises: + IOError: If the schema cannot be saved to the output path. + """ + schema = self.load_schema() + try: + with output_path.open("w", encoding="utf-8") as f: + json.dump(schema.model_dump(mode="json", exclude_unset=True), f, indent=2) + logger.info(f"Saved CSV schema to file: {output_path}") + except Exception as e: + raise IOError(f"IO error: {output_path, e}") + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the CSV file name. + + Returns: + str: A unique identifier for the schema. + """ + return self.config.schema_id or f"csv-schema-{self.csv_path.stem}" diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/excel_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/excel_schema_provider.py new file mode 100644 index 00000000..4cd0141f --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/excel_schema_provider.py @@ -0,0 +1,162 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Excel Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for Excel files. +It discovers and generates ETL schemas from Excel files for processing and loading +the data into a document graph. + +The module includes the following components: +- ExcelSchemaProvider: Schema provider for Excel files + +The Excel schema provider uses the ExcelSchemaDiscoveryProvider to analyze Excel files +and generate appropriate ETL schemas. It supports loading schemas from Excel files +and optionally saving the discovered schemas to JSON files. + +Usage: + # Get an Excel schema provider for a specific file + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "excel", + "schema_id": "my_excel_schema", + "connection_config": { + "path": "/path/to/data.xlsx" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging + +# Import custom logging to ensure configuration is available + +import json +from pathlib import Path +from typing import Optional, Any + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig + +from graphrag_toolkit.document_graph.schema.discovery.excel_discovery_provider import ( + ExcelSchemaDiscoveryProvider +) + +logger = logging.getLogger(__name__) + + +class ExcelSchemaProvider(SchemaProviderBase): + """ + Schema provider for Excel files. + + This class provides functionality to discover and generate ETL schemas from Excel files. + It uses the ExcelSchemaDiscoveryProvider to analyze Excel files and extract field information + to create appropriate ETL schemas for processing and loading the data into a document graph. + + The provider supports loading schemas directly from Excel files and optionally saving + the discovered schemas to JSON files for later use. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + excel_path (Path): The path to the Excel file to analyze. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize an Excel schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the path to the Excel file in the connection_config. + + Raises: + ValueError: If the path is not provided in the connection_config. + FileNotFoundError: If the specified Excel file does not exist. + """ + self.config = config + path = config.connection_config.get("path") + if not path: + raise ValueError("Validation error") + self.excel_path = Path(path) + + if not self.excel_path.exists(): + raise FileNotFoundError(f"File not found") + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "ExcelSchemaProvider": + """ + Factory method to construct an Excel provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + ExcelSchemaProvider: A new instance of the Excel schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load (or generate) an ETL schema from an Excel file using schema discovery. + + This method uses the ExcelSchemaDiscoveryProvider to analyze the Excel file and + generate an appropriate ETL schema based on the structure and content of the file. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters (not used in this provider). + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the Excel file and how it should be processed. + + Raises: + FileNotFoundError: If the Excel file does not exist. + ValueError: If the Excel file cannot be parsed or has no discoverable fields. + ImportError: If the required Excel libraries are not installed. + """ + discovery = ExcelSchemaDiscoveryProvider(source=self.excel_path) + return discovery.discover_schema() + + def save_schema(self, output_path: Path) -> None: + """ + Save the discovered ETL schema to a JSON file at the given path. + + This method discovers the schema from the Excel file and saves it to a JSON file + at the specified path. The schema can then be loaded directly using a FileSchemaProvider. + + Args: + output_path (Path): The path where the schema JSON file should be saved. + + Raises: + IOError: If the schema cannot be saved to the output path. + """ + schema = self.load_schema() + try: + with output_path.open("w", encoding="utf-8") as f: + json.dump(schema.model_dump(mode="json", exclude_unset=True), f, indent=2) + logger.info(f"Saved Excel schema to file: {output_path}") + except Exception as e: + raise IOError(f"IO error: {output_path, e}") + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the Excel file name. + + Returns: + str: A unique identifier for the schema. + """ + return self.config.schema_id or f"excel-schema-{self.excel_path.stem}" diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/file_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/file_schema_provider.py new file mode 100644 index 00000000..2f438888 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/file_schema_provider.py @@ -0,0 +1,175 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""File Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for loading and saving ETL schemas from/to local JSON files. +Unlike discovery-based providers, this provider works with pre-defined schema files rather than +inferring schemas from data sources. + +The module includes the following components: +- FileSchemaProvider: Schema provider for JSON schema files + +The file schema provider loads ETL schemas directly from JSON files and can also save +schemas back to files. It's commonly used to store and retrieve pre-defined or previously +discovered schemas. + +Usage: + # Get a file schema provider for a specific schema file + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "file", + "schema_id": "my_schema", + "connection_config": { + "path": "/path/to/schema.json" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") + + # Save a modified schema + provider.save_schema(schema) +""" + +# Import custom logging to ensure configuration is available + +import json +import logging +from pathlib import Path +from typing import Any, Optional + +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema + +logger = logging.getLogger(__name__) + + +class FileSchemaProvider(SchemaProviderBase): + """ + Schema provider for loading and saving ETL schemas from/to local JSON files. + + This class provides functionality to load pre-defined ETL schemas from JSON files + and save schemas back to files. Unlike discovery-based providers, this provider + works with existing schema files rather than inferring schemas from data sources. + + The provider is commonly used to store and retrieve pre-defined or previously + discovered schemas for reuse in ETL pipelines. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + path (Path): The path to the JSON schema file. + _schema_id (str): The unique identifier for the schema. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize a file schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the path to the JSON schema file in the connection_config. + + Raises: + ValueError: If the path is not provided in the connection_config. + FileNotFoundError: If the specified JSON file does not exist. + """ + self.config = config + path_str = config.connection_config.get("path") + + if not path_str: + raise ErrorHandler.validation_error( + "schema_config_path", + "None", + "valid file path in config.connection_config['path']" + ) + + self.path: Path = Path(path_str) + if not self.path.exists(): + raise FileNotFoundError(f"File not found: {self.path}") + + self._schema_id: str = config.schema_id or self.path.stem + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "FileSchemaProvider": + """ + Factory method to construct a file provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + FileSchemaProvider: A new instance of the file schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load the ETL schema from the JSON file and parse it into an ETLSchema object. + + This method reads the JSON schema file, parses it, and returns an ETLSchema object. + Unlike discovery-based providers, this method loads a pre-defined schema rather + than inferring one from a data source. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters (not used in this provider). + + Returns: + ETLSchema: The ETL schema loaded from the JSON file. + + Raises: + FileNotFoundError: If the JSON file does not exist. + ValueError: If the JSON file cannot be parsed or is not a valid ETL schema. + IOError: If there's an error reading the file. + """ + try: + with self.path.open("r", encoding="utf-8") as f: + schema_json = json.load(f) + logger.info(f"Loaded schema from file: {self.path}") + return ETLSchema(**schema_json) + except Exception as e: + logger.error(f"Failed to load ETL schema from file {self.path}: {e}") + raise ErrorHandler.database_error("file schema load", e) + + def save_schema(self, schema: ETLSchema) -> None: + """ + Save the ETL schema to the file path specified in the configuration. + + This method serializes the ETL schema to JSON and writes it to the file + specified in the provider's configuration. The schema can then be loaded + later using this or another FileSchemaProvider. + + Args: + schema (ETLSchema): The ETL schema to save. + + Raises: + IOError: If there's an error writing to the file. + """ + try: + with self.path.open("w", encoding="utf-8") as f: + json.dump(schema.model_dump(mode="json", exclude_unset=True), f, indent=2) + logger.info(f"Saved schema to file: {self.path}") + except Exception as e: + logger.error(f"Failed to save ETL schema to file {self.path}: {e}") + raise IOError(f"IO error: {self.path, e}") + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the JSON file name. + + Returns: + str: A unique identifier for the schema. + """ + return self._schema_id diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/glue_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/glue_schema_provider.py new file mode 100644 index 00000000..368311c6 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/glue_schema_provider.py @@ -0,0 +1,247 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AWS Glue Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for loading ETL schemas from AWS Glue Schema Registry. +It retrieves AVRO schemas from Glue and converts them into ETL schemas for processing +and loading data into a document graph. + +The module includes the following components: +- GlueSchemaProvider: Schema provider for AWS Glue Schema Registry + +The Glue schema provider connects to AWS Glue Schema Registry, retrieves the latest +version of a specified schema, and converts it into an ETL schema. It extracts field +information from AVRO schemas and builds appropriate ETL schemas for document graph +operations. + +Usage: + # Get a Glue schema provider for a specific schema + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "glue", + "schema_id": "my_glue_schema", + "connection_config": { + "registry_name": "my-registry", + "schema_name": "my-schema", + "region": "us-east-1" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +# Import custom logging to ensure configuration is available + +import json +import logging +from typing import Any, List, Optional + +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.etl_schema_model import ( + ETLSchema, ExtractConfig, TransformConfig, LoadConfig, + ChunkingConfig, MetadataMapping, EntityExtractionConfig, NormalizeConfig, + NodeDefinition, RelationshipDefinition +) +from graphrag_toolkit.document_graph.config import DocumentGraphConfig + +logger = logging.getLogger(__name__) + + +class GlueSchemaProvider(SchemaProviderBase): + """ + Schema provider for AWS Glue Schema Registry. + + This class provides functionality to load ETL schemas from AWS Glue Schema Registry. + It connects to the registry, retrieves the latest version of a specified schema, + and converts it into an ETL schema for document graph operations. + + The provider extracts field information from AVRO schemas in the Glue registry + and builds appropriate ETL schemas with those fields. It supports AWS authentication + through boto3 sessions and can be configured with different AWS regions. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + registry_name (str): The name of the Glue schema registry. + schema_name (str): The name of the schema in the registry. + region (str): The AWS region where the registry is located. + session: The boto3 session for AWS authentication. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize a Glue schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the registry name and schema name in the connection_config. + + Raises: + ValueError: If the registry name or schema name is not provided in the connection_config. + """ + self.config = config + conn = config.connection_config + + self.registry_name = conn.get("registry_name") + self.schema_name = conn.get("schema_name") + self.region = conn.get("region", DocumentGraphConfig.aws_region) + self.session = conn.get("session") or DocumentGraphConfig.session + + if not self.registry_name or not self.schema_name: + raise ErrorHandler.validation_error( + "glue.connection_config", + conn, + "registry_name and schema_name must be defined" + ) + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "GlueSchemaProvider": + """ + Factory method to construct a Glue provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + GlueSchemaProvider: A new instance of the Glue schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load an ETL schema from AWS Glue Schema Registry. + + This method connects to AWS Glue Schema Registry, retrieves the latest version + of the specified schema, extracts field information from the AVRO schema, + and builds an appropriate ETL schema for document graph operations. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters (not used in this provider). + + Returns: + ETLSchema: A complete ETL schema object built from the Glue schema. + + Raises: + ValueError: If the schema does not exist in the registry or cannot be parsed. + Exception: If there's an error connecting to AWS or retrieving the schema. + """ + glue = self.session.client("glue", region_name=self.region) + try: + response = glue.get_schema_version( + SchemaId={ + "SchemaName": self.schema_name, + "RegistryName": self.registry_name + }, + SchemaVersionNumber={"LatestVersion": True} + ) + schema_definition = response["SchemaDefinition"] + fields = self._parse_avro_fields(schema_definition) + + logger.info(f"Loaded schema from Glue registry: {self.registry_name}/{self.schema_name}") + return self._build_minimal_etl_schema(fields) + + except glue.exceptions.EntityNotFoundException: + raise ErrorHandler.validation_error( + "glue_schema", + f"{self.registry_name}/{self.schema_name}", + "existing schema in Glue registry" + ) + except Exception as e: + logger.error(f"Failed to load Glue schema: {e}") + raise ErrorHandler.database_error("Glue schema load", e) + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the registry and schema names. + + Returns: + str: A unique identifier for the schema. + """ + return self.config.schema_id or f"{self.registry_name}:{self.schema_name}" + + def _parse_avro_fields(self, schema_def: str) -> List[str]: + """ + Parse field names from an AVRO schema definition. + + This method parses the AVRO schema definition to extract field names, + safely handling union types and defaults. It validates that the schema + is of AVRO record type and contains fields. + + Args: + schema_def (str): The AVRO schema definition as a JSON string. + + Returns: + List[str]: A list of field names extracted from the AVRO schema. + + Raises: + ValueError: If the schema is not of AVRO record type, is missing fields, + or no fields could be extracted. + Exception: If there's an error parsing the schema. + """ + try: + schema_json = json.loads(schema_def) + + # Handle AVRO record type at root + if schema_json.get("type") != "record" or "fields" not in schema_json: + raise ValueError("Schema is not of AVRO record type or missing 'fields'") + + field_names = [] + for field in schema_json["fields"]: + name = field.get("name") + if not name: + logger.warning(f"Skipping unnamed field in AVRO schema: {field}") + continue + field_names.append(name) + + if not field_names: + raise ValueError("No fields extracted from AVRO schema") + + return field_names + except Exception as e: + raise ErrorHandler.schema_error("parse", "glue_avro", e) + + def _build_minimal_etl_schema(self, field_names: List[str]) -> ETLSchema: + """ + Build a minimal ETL schema from the extracted field names. + + This method creates a basic ETL schema with default configurations for + extraction, transformation, and loading, using the field names extracted + from the AVRO schema. + + Args: + field_names (List[str]): The list of field names to include in the schema. + + Returns: + ETLSchema: A complete ETL schema object with the specified fields. + """ + return ETLSchema( + schema_id=self.get_schema_id(), + description=f"Autogenerated from Glue schema {self.registry_name}/{self.schema_name}", + extract=ExtractConfig(source_type="s3"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="fixed_length"), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig() + ), + load=LoadConfig( + document_node=NodeDefinition(type="Document", fields=field_names), + section_node=NodeDefinition(type="Section", fields=[]), + relationships=[ + RelationshipDefinition(type="contains", source="document_id", target="section_id") + ] + ) + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/json_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/json_schema_provider.py new file mode 100644 index 00000000..a6e58627 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/json_schema_provider.py @@ -0,0 +1,161 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for JSON (JavaScript Object Notation) files. +It discovers and generates ETL schemas from JSON files for processing and loading +the data into a document graph. + +The module includes the following components: +- JSONSchemaProvider: Schema provider for JSON files + +The JSON schema provider uses the JSONSchemaDiscoveryProvider to analyze JSON files +and generate appropriate ETL schemas. It supports loading schemas from JSON files +and optionally saving the discovered schemas to JSON files. + +Usage: + # Get a JSON schema provider for a specific file + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "json", + "schema_id": "my_json_schema", + "connection_config": { + "path": "/path/to/data.json" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging +import json +from pathlib import Path +from typing import Optional, Any, Dict + +# Import custom logging to ensure configuration is available + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.discovery.json_discovery_provider import JSONSchemaDiscoveryProvider + +logger = logging.getLogger(__name__) + + +class JSONSchemaProvider(SchemaProviderBase): + """ + Schema provider for JSON (JavaScript Object Notation) files. + + This class provides functionality to discover and generate ETL schemas from JSON files. + It uses the JSONSchemaDiscoveryProvider to analyze JSON files and extract field + information to create appropriate ETL schemas for processing and loading the data into + a document graph. + + The provider supports loading schemas directly from JSON files and optionally saving + the discovered schemas to JSON files for later use. It handles both array-based JSON + files (containing multiple records) and single-object JSON files. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + json_path (Path): The path to the JSON file to analyze. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize a JSON schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the path to the JSON file in the connection_config. + + Raises: + ValueError: If the path is not provided in the connection_config. + FileNotFoundError: If the specified JSON file does not exist. + """ + self.config = config + path_str = config.connection_config.get("path") + if not path_str: + raise ValueError("Validation error") + self.json_path = Path(path_str) + + if not self.json_path.exists(): + raise FileNotFoundError(f"File not found") + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "JSONSchemaProvider": + """ + Factory method to construct a JSON provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + JSONSchemaProvider: A new instance of the JSON schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load (or generate) an ETL schema from a JSON file using schema discovery. + + This method uses the JSONSchemaDiscoveryProvider to analyze the JSON file and + generate an appropriate ETL schema based on the structure and content of the file. + It handles both array-based JSON files and single-object JSON files. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters (not used in this provider). + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the JSON file and how it should be processed. + + Raises: + FileNotFoundError: If the JSON file does not exist. + ValueError: If the JSON file cannot be parsed or has no discoverable fields. + Exception: Any exception raised during JSON parsing. + """ + discovery = JSONSchemaDiscoveryProvider(source=self.json_path) + return discovery.discover_schema() + + def save_schema(self, output_path: Path) -> None: + """ + Save the inferred ETL schema to a JSON file at the given output path. + + This method discovers the schema from the JSON file and saves it to a JSON file + at the specified path. The schema can then be loaded directly using a FileSchemaProvider. + + Args: + output_path (Path): The path where the schema JSON file should be saved. + + Raises: + IOError: If the schema cannot be saved to the output path. + """ + schema = self.load_schema() + try: + with output_path.open("w", encoding="utf-8") as f: + json.dump(schema.model_dump(mode="json", exclude_unset=True), f, indent=2) + logger.info(f"Saved JSON schema to file: {output_path}") + except Exception as e: + raise IOError(f"IO error: {output_path, e}") + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the JSON file name. + + Returns: + str: A unique identifier for the schema. + """ + return self.config.schema_id or f"json-schema-{self.json_path.stem}" diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/parquet_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/parquet_schema_provider.py new file mode 100644 index 00000000..23bcb6c9 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/parquet_schema_provider.py @@ -0,0 +1,160 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Parquet Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for Parquet files, a columnar storage format +commonly used in big data processing. It discovers and generates ETL schemas from +Parquet files for processing and loading the data into a document graph. + +The module includes the following components: +- ParquetSchemaProvider: Schema provider for Parquet files + +The Parquet schema provider uses the ParquetSchemaDiscoveryProvider to analyze Parquet +files and generate appropriate ETL schemas. It supports loading schemas from Parquet +files and optionally saving the discovered schemas to JSON files. + +Usage: + # Get a Parquet schema provider for a specific file + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "parquet", + "schema_id": "my_parquet_schema", + "connection_config": { + "path": "/path/to/data.parquet" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging +import json +from pathlib import Path +from typing import Optional, Any, Dict + +# Import custom logging to ensure configuration is available + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.discovery.parquet_discovery_provider import ParquetSchemaDiscoveryProvider + +logger = logging.getLogger(__name__) + + +class ParquetSchemaProvider(SchemaProviderBase): + """ + Schema provider for Parquet files. + + This class provides functionality to discover and generate ETL schemas from Parquet files. + It uses the ParquetSchemaDiscoveryProvider to analyze Parquet files and extract field + information to create appropriate ETL schemas for processing and loading the data into + a document graph. + + The provider supports loading schemas directly from Parquet files and optionally saving + the discovered schemas to JSON files for later use. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + parquet_path (Path): The path to the Parquet file to analyze. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize a Parquet schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the path to the Parquet file in the connection_config. + + Raises: + ValueError: If the path is not provided in the connection_config. + FileNotFoundError: If the specified Parquet file does not exist. + """ + self.config = config + path = config.connection_config.get("path") + if not path: + raise ValueError("Validation error") + self.parquet_path = Path(path) + + if not self.parquet_path.exists(): + raise FileNotFoundError(f"File not found") + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "ParquetSchemaProvider": + """ + Factory method to construct a Parquet provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + ParquetSchemaProvider: A new instance of the Parquet schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load (or generate) an ETL schema from a Parquet file using schema discovery. + + This method uses the ParquetSchemaDiscoveryProvider to analyze the Parquet file and + generate an appropriate ETL schema based on the structure and content of the file. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters that are passed to the + discovery provider, such as specific columns to include. + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the Parquet file and how it should be processed. + + Raises: + FileNotFoundError: If the Parquet file does not exist. + ValueError: If the Parquet file cannot be parsed or has no discoverable fields. + Exception: Any exception raised during Parquet parsing. + """ + discovery = ParquetSchemaDiscoveryProvider(source=self.parquet_path, args=kwargs) + return discovery.discover_schema() + + def save_schema(self, output_path: Path) -> None: + """ + Save the inferred ETL schema to a JSON file at the given output path. + + This method discovers the schema from the Parquet file and saves it to a JSON file + at the specified path. The schema can then be loaded directly using a FileSchemaProvider. + + Args: + output_path (Path): The path where the schema JSON file should be saved. + + Raises: + IOError: If the schema cannot be saved to the output path. + """ + schema = self.load_schema() + try: + with output_path.open("w", encoding="utf-8") as f: + json.dump(schema.model_dump(mode="json", exclude_unset=True), f, indent=2) + logger.info(f"Saved Parquet schema to file: {output_path}") + except Exception as e: + raise IOError(f"IO error: {output_path, e}") + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the Parquet file name. + + Returns: + str: A unique identifier for the schema. + """ + return self.config.schema_id or f"parquet-schema-{self.parquet_path.stem}" diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/s3_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/s3_schema_provider.py new file mode 100644 index 00000000..004cc32a --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/s3_schema_provider.py @@ -0,0 +1,204 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""S3 Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for loading and saving ETL schemas from/to Amazon S3. +It retrieves schema files from S3 buckets and can also save schemas back to S3. + +The module includes the following components: +- S3SchemaProvider: Schema provider for Amazon S3 + +The S3 schema provider loads ETL schemas directly from JSON files stored in S3 buckets +and can also save schemas back to S3. It's commonly used to store and retrieve pre-defined +or previously discovered schemas in cloud environments. + +Usage: + # Get an S3 schema provider for a specific schema file in S3 + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "s3", + "schema_id": "my_s3_schema", + "connection_config": { + "bucket": "my-schema-bucket", + "key": "schemas/my-schema.json", + "region": "us-east-1" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") + + # Save a modified schema back to S3 + provider.save_schema() +""" + +# Import custom logging to ensure configuration is available + +import json +import logging +from typing import Any, Optional + +import boto3 +from botocore.config import Config as BotoConfig + +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema + +logger = logging.getLogger(__name__) + + +class S3SchemaProvider(SchemaProviderBase): + """ + Schema provider for loading and saving ETL schemas from/to Amazon S3. + + This class provides functionality to load pre-defined ETL schemas from JSON files + stored in S3 buckets and save schemas back to S3. It handles AWS authentication + and S3 operations to retrieve and store schema files. + + The provider is commonly used in cloud environments to store and retrieve + pre-defined or previously discovered schemas for reuse in ETL pipelines. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + bucket (str): The name of the S3 bucket containing the schema file. + key (str): The S3 key (path) to the schema JSON file. + region (str): The AWS region where the S3 bucket is located. + session: The boto3 session for AWS authentication. + s3: The boto3 S3 client used for S3 operations. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize an S3 schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the bucket and key in the connection_config. + + Raises: + ValueError: If the bucket or key is not provided in the connection_config, + or if the key does not point to a JSON file. + """ + self.config = config + conn = config.connection_config + + self.bucket = conn.get("bucket") + self.key = conn.get("key") + self.region = conn.get("region", "us-east-1") + self.session = conn.get("session") or boto3.Session() + + if not self.bucket or not self.key: + raise ValueError( + "s3.connection_config", + conn, + "bucket and key must be provided" + ) + + if not self.key.endswith(".json"): + raise ValueError( + "s3.key", + self.key, + "S3 key must point to a .json file" + ) + + self.s3 = self.session.client( + "s3", + region_name=self.region, + config=BotoConfig(retries={"max_attempts": 3}) + ) + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "S3SchemaProvider": + """ + Factory method to construct an S3 provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + S3SchemaProvider: A new instance of the S3 schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load the ETL schema from the S3 bucket and parse it into an ETLSchema object. + + This method retrieves the JSON schema file from S3, parses it, and returns + an ETLSchema object. It handles AWS authentication and S3 operations to + retrieve the schema file. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters (not used in this provider). + + Returns: + ETLSchema: The ETL schema loaded from the S3 file. + + Raises: + FileNotFoundError: If the S3 object does not exist. + ValueError: If the S3 object cannot be parsed or is not a valid ETL schema. + Exception: If there's an error connecting to AWS or retrieving the object. + """ + s3_path = f"s3://{self.bucket}/{self.key}" + try: + response = self.s3.get_object(Bucket=self.bucket, Key=self.key) + schema_str = response["Body"].read().decode("utf-8") + schema_json = json.loads(schema_str) + logger.info(f"Successfully loaded schema from {s3_path}") + return ETLSchema(**schema_json) + except self.s3.exceptions.NoSuchKey: + raise FileNotFoundError(f"File not found") + except Exception as e: + logger.error(f"Error loading schema from {s3_path}: {e}") + raise RuntimeError(f"S3 schema load failed: {e}") from e + + def save_schema(self, output_key: Optional[str] = None) -> None: + """ + Save the ETL schema to the S3 bucket under the given key. + + This method retrieves the current schema, serializes it to JSON, and saves + it to the S3 bucket. If no output key is provided, the original key is used. + + Args: + output_key (Optional[str]): The S3 key where the schema should be saved. + If not provided, the original key is used. + + Raises: + IOError: If there's an error writing to S3. + Exception: If there's an error connecting to AWS or retrieving the schema. + """ + schema = self.load_schema() + key_to_use = output_key or self.key + s3_path = f"s3://{self.bucket}/{key_to_use}" + try: + self.s3.put_object( + Bucket=self.bucket, + Key=key_to_use, + Body=json.dumps(schema.model_dump(mode="json", exclude_unset=True), indent=2).encode("utf-8"), + ContentType="application/json" + ) + logger.info(f"Schema saved to {s3_path}") + except Exception as e: + raise IOError(f"IO error: {s3_path, e}") + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the S3 key. + + Returns: + str: A unique identifier for the schema. + """ + return self.config.schema_id or self.key.split("/")[-1].replace(".json", "") diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_base.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_base.py new file mode 100644 index 00000000..a7a530d7 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_base.py @@ -0,0 +1,149 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Provider Base Module for Document Graph Operations. + +This module defines the abstract base class that all schema providers must implement. +Schema providers are responsible for loading ETL schemas from various sources such as +files, databases, and cloud services. + +The module includes the following components: +- SchemaProviderBase: Abstract base class for all schema providers + +Schema providers are used to load ETL schemas that define how data should be extracted, +transformed, and loaded into a document graph. Different providers handle different +source types (e.g., files, S3, databases) and formats (e.g., YAML, JSON, CSV). + +Usage: + # Get a schema provider for a specific configuration + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "file", + "schema_id": "my_schema", + "connection_config": { + "path": "/path/to/schema.yaml" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") +""" + +import logging +from abc import ABC, abstractmethod +from typing import Dict, Any, Type + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema + + +class SchemaProviderBase(ABC): + """ + Abstract base class for all ETL schema providers. + + This class defines the interface that all schema providers must implement. + Schema providers are responsible for loading ETL schemas from various sources + including files, databases, and cloud services. Each provider implementation + handles a specific source type or format. + + The SchemaProviderBase class enforces a consistent interface across all providers, + ensuring that they can be used interchangeably in the document graph ETL pipeline. + Concrete implementations must provide methods for loading schemas and retrieving + schema identifiers. + + Attributes: + None directly in the base class. Subclasses typically have: + config: Configuration for the provider, usually a SchemaProviderConfig instance. + schema_id: A unique identifier for the schema being provided. + """ + + @classmethod + @abstractmethod + def from_config(cls: Type["SchemaProviderBase"], config: Dict[str, Any]) -> "SchemaProviderBase": + """ + Factory method to construct a provider from a raw config dictionary. + + This method is responsible for parsing the configuration dictionary and creating + an instance of the appropriate schema provider. It validates the configuration + parameters and raises appropriate exceptions if the configuration is invalid. + + The configuration dictionary typically includes: + - provider_type: The type of provider (e.g., "file", "s3", "glue") + - schema_id: A unique identifier for the schema + - connection_config: Provider-specific connection parameters + + Args: + config: A dictionary containing configuration parameters for the provider. + The exact structure depends on the provider implementation. + + Returns: + An instance of the implementing SchemaProviderBase subclass, properly + configured according to the provided configuration. + + Raises: + ValueError: If the configuration is invalid or missing required parameters. + TypeError: If the configuration contains parameters of the wrong type. + Other exceptions may be raised by specific provider implementations. + """ + pass + + @abstractmethod + def load_schema(self, source = None, **kwargs) -> ETLSchema: + """ + Load and return the ETL schema as a Pydantic model. + + This method is responsible for loading the ETL schema from the source specified + in the provider's configuration. It reads the schema definition, validates it, + and returns it as an ETLSchema object. The schema defines how data should be + extracted, transformed, and loaded into a document graph. + + The method may perform additional processing such as: + - Validating the schema against a schema definition + - Enriching the schema with additional information + - Converting between different schema formats + + Args: + source: Optional DocumentGraphSource containing data source and registration info. + If provided, this may override the source specified in the provider's + configuration. + **kwargs: Additional provider-specific parameters that control how the schema + is loaded and processed. + + Returns: + ETLSchema: A complete ETL schema object that defines how data should be + extracted, transformed, and loaded into a document graph. + + Raises: + FileNotFoundError: If the schema source file does not exist. + ValueError: If the schema is invalid or cannot be parsed. + Other exceptions may be raised by specific provider implementations. + """ + pass + + @abstractmethod + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema that this provider + is responsible for. The schema ID is used to reference the schema in the + document graph ETL pipeline and to identify the schema in storage systems. + + The schema ID should be unique within the context of the application and + should be consistent across multiple invocations of the provider for the + same schema source. + + Returns: + str: A unique identifier for the schema. This is typically derived from + the schema source (e.g., filename, database table) and may include + additional information such as version numbers or timestamps. + """ + pass diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_config.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_config.py new file mode 100644 index 00000000..fa622c04 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_config.py @@ -0,0 +1,70 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Provider Config Module for Document Graph Operations. + +This module defines the configuration model for schema providers used in document graph +operations. It provides a standardized way to configure different types of schema providers +through a Pydantic model. + +The module includes the following components: +- SchemaProviderConfig: Base configuration model for all schema providers + +The SchemaProviderConfig class is used to configure schema providers with information +about the provider type, schema ID, and connection-specific parameters. This configuration +is used by the schema provider factory to instantiate the appropriate provider. + +Usage: + # Create a configuration for a file-based schema provider + from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig + + config = SchemaProviderConfig( + type="file", + schema_id="my_schema", + connection_config={ + "path": "/path/to/schema.yaml" + } + ) + + # Use the configuration to create a schema provider + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + provider = get_schema_provider(config) +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +from typing import Dict, Any, Literal, Optional +from pydantic import BaseModel, Field + +class SchemaProviderConfig(BaseModel): + """ + Base configuration for schema providers. + + This Pydantic model defines the configuration parameters for schema providers + used in document graph operations. It provides a standardized way to configure + different types of schema providers with information about the provider type, + schema ID, and connection-specific parameters. + + The SchemaProviderConfig is used by the schema provider factory to instantiate + the appropriate provider based on the specified type. Each provider type has + its own specific connection configuration requirements. + + Attributes: + type (str): The type of schema provider. Must be one of the supported types: + "file", "s3", "static", "csv", "json", "excel", "glue", "parquet", "yaml", "xml". + schema_id (Optional[str]): An optional unique identifier for the schema. + If not provided, the provider will generate one based on the source. + connection_config (Dict[str, Any]): A dictionary containing connection-specific + configuration parameters. The exact structure depends on the provider type. + For example: + - "file" provider: {"path": "/path/to/schema.yaml"} + - "s3" provider: {"bucket": "my-bucket", "key": "path/to/schema.yaml"} + """ + type: Literal["file", "s3", "static", "csv", "json", "excel", "glue", "parquet", "yaml", "xml"] = Field(..., description="Type of schema provider") + schema_id: Optional[str] = Field(None, description="Unique ID for the schema (optional override)") + connection_config: Dict[str, Any] = Field(default_factory=dict, description="Connection-specific configuration") diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_config_aws_base.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_config_aws_base.py new file mode 100644 index 00000000..cfab9559 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_config_aws_base.py @@ -0,0 +1,64 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Provider Config AWS Base Module for Document Graph Operations. + +This module defines the base configuration model for AWS-backed schema providers +used in document graph operations. It provides a standardized way to configure +schema providers that interact with AWS services like S3 and AWS Glue. + +The module includes the following components: +- AWSSchemaProviderConfig: Base configuration model for AWS-backed schema providers + +The AWSSchemaProviderConfig class is used as a base for more specific AWS provider +configurations, providing common AWS connection parameters such as bucket, key, +region, and profile name. This configuration is extended by specific AWS provider +implementations like S3SchemaProvider and GlueSchemaProvider. + +Usage: + # Create a configuration for an S3-based schema provider + from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig + + config = SchemaProviderConfig( + type="s3", + schema_id="my_s3_schema", + connection_config={ + "bucket": "my-schema-bucket", + "key": "schemas/my-schema.json", + "region": "us-west-2", + "profile_name": "my-aws-profile" + } + ) + + # Use the configuration to create a schema provider + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + provider = get_schema_provider(config) +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from pydantic import BaseModel, Field +from typing import Optional + + +class AWSSchemaProviderConfig(BaseModel): + """ + Configuration base class for AWS-backed schema providers (e.g., S3, Glue). + + Attributes: + bucket (str): S3 bucket name where the schema is stored. + key (str): S3 key (file path) to the schema JSON file. + region (Optional[str]): AWS region of the S3 bucket (default: inferred from environment or session). + profile_name (Optional[str]): Named AWS CLI profile to use for credentials (optional). + """ + + bucket: str = Field(..., description="S3 bucket name containing the schema") + key: str = Field(..., description="S3 key (path) to the schema JSON file") + region: Optional[str] = Field(None, description="AWS region for the S3 bucket") + profile_name: Optional[str] = Field(None, description="AWS CLI profile to use for credentials (optional)") diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_factory.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_factory.py new file mode 100644 index 00000000..a1612adf --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_factory.py @@ -0,0 +1,194 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Provider Factory Module for Document Graph Operations. + +This module provides a factory for creating schema providers based on configuration +parameters. It centralizes the instantiation of different schema provider types and +ensures that the appropriate provider is created based on the specified type. + +The module includes the following components: +- SchemaProviderFactory: Factory class for creating schema providers + +The SchemaProviderFactory maintains a registry of available schema provider types +and provides methods for registering new provider types and creating provider instances. +It supports various provider types including file-based, S3-based, and database-based +providers. + +Usage: + # Create a schema provider using a configuration dictionary + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import SchemaProviderFactory + + config = { + "type": "file", + "schema_id": "my_schema", + "connection_config": { + "path": "/path/to/schema.yaml" + } + } + + provider = SchemaProviderFactory.create(config) + + # Or use the convenience function + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import Dict, Any, Type +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase + +from graphrag_toolkit.document_graph.schema.providers.file_schema_provider import FileSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.s3_schema_provider import S3SchemaProvider +from graphrag_toolkit.document_graph.schema.providers.csv_schema_provider import CSVSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.glue_schema_provider import GlueSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.parquet_schema_provider import ParquetSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.json_schema_provider import JSONSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.excel_schema_provider import ExcelSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.yaml_schema_provider import YAMLSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.xml_schema_provider import XMLSchemaProvider + + +class SchemaProviderFactory: + """ + Factory class to instantiate schema providers based on a configuration dictionary. + + This class provides a centralized way to create schema provider instances based on + a configuration dictionary. It maintains a registry of available provider types and + maps them to their corresponding provider classes. The factory ensures that the + appropriate provider is created based on the specified type in the configuration. + + The factory supports various provider types including: + - file: For loading schemas from local files + - s3: For loading schemas from Amazon S3 + - static: For using predefined schemas + - csv: For generating schemas from CSV files + - json: For generating schemas from JSON files + - excel: For generating schemas from Excel files + - glue: For loading schemas from AWS Glue + - parquet: For generating schemas from Parquet files + - yaml: For generating schemas from YAML files + - xml: For generating schemas from XML files + + New provider types can be registered using the register_provider method. + + Attributes: + _registry (Dict[str, Type[SchemaProviderBase]]): A mapping of provider type names + to their corresponding provider classes. + """ + + _registry: Dict[str, Type[SchemaProviderBase]] = { + "file": FileSchemaProvider, + "s3": S3SchemaProvider, + "csv": CSVSchemaProvider, + "json": JSONSchemaProvider, + "excel": ExcelSchemaProvider, + "glue": GlueSchemaProvider, + "parquet": ParquetSchemaProvider, + "yaml": YAMLSchemaProvider, + "xml": XMLSchemaProvider, + } + + @classmethod + def _get_static_provider(cls): + """Late import for StaticSchemaProvider to avoid circular imports.""" + from graphrag_toolkit.document_graph.schema.static_schema_provider import StaticSchemaProvider + return StaticSchemaProvider + + @classmethod + def register_provider(cls, type_name: str, provider_class: Type[SchemaProviderBase]) -> None: + """ + Register a new schema provider class. + + This method allows for the registration of custom schema provider classes + that are not included in the default registry. Once registered, the provider + can be created using the factory's create method by specifying the registered + type name in the configuration. + + Args: + type_name: A unique string identifier for the provider type. This will be + used as the "type" value in configuration dictionaries to specify + that this provider should be used. + provider_class: The schema provider class to register. This must be a subclass + of SchemaProviderBase and implement the required interface, + including the from_config class method. + + Returns: + None + + Example: + # Register a custom provider + from my_package import MyCustomProvider + SchemaProviderFactory.register_provider("custom", MyCustomProvider) + + # Use the custom provider + config = {"type": "custom", "schema_id": "my_schema", "connection_config": {...}} + provider = SchemaProviderFactory.create(config) + """ + cls._registry[type_name] = provider_class + + @classmethod + def create(cls, config: Dict[str, Any]) -> SchemaProviderBase: + """ + Instantiate a schema provider based on the `type` in the config. + + This method creates an instance of the appropriate schema provider based on + the "type" specified in the configuration dictionary. It validates that the + type is registered and that the provider class implements the required interface. + + The method performs the following steps: + 1. Extracts the "type" from the configuration + 2. Validates that the type is registered in the factory + 3. Gets the provider class from the registry + 4. Validates that the provider class implements the required interface + 5. Converts the configuration dictionary to a SchemaProviderConfig object + 6. Creates and returns an instance of the provider using the from_config method + + Args: + config: A dictionary containing the configuration for the schema provider. + Must include a "type" key with a value that is registered in the factory. + Other required keys depend on the specific provider type. + + Returns: + An instance of a SchemaProviderBase subclass, configured according to the + provided configuration. + + Raises: + ValueError: If the "type" is missing or not registered in the factory. + TypeError: If the provider class does not implement the required interface. + Other exceptions may be raised by the provider's from_config method. + """ + type_name = config.get("type") + + # Handle static provider with late import to avoid circular dependency + if type_name == "static": + provider_class = cls._get_static_provider() + elif type_name and type_name in cls._registry: + provider_class = cls._registry[type_name] + else: + available_types = list(cls._registry.keys()) + ["static"] + raise ErrorHandler.validation_error( + "schema_provider_type", + type_name or "None", + f"one of {available_types}" + ) + + # provider_class is already set above + + if not hasattr(provider_class, "from_config"): + raise TypeError(f"{provider_class.__name__} is missing required `from_config(config)` method.") + + # Convert dict to SchemaProviderConfig object + from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig + config_obj = SchemaProviderConfig(**config) + return provider_class.from_config(config_obj) diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_registry.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_registry.py new file mode 100644 index 00000000..d725a87e --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/schema_provider_registry.py @@ -0,0 +1,170 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema Provider Registry Module for Document Graph Operations. + +This module provides a registry system for schema providers, allowing for runtime +registration and lookup of provider classes by name. It centralizes the management +of available schema provider types and ensures that providers can be dynamically +added and retrieved. + +The module includes the following components: +- SchemaProviderRegistry: Registry class for managing schema provider classes + +The SchemaProviderRegistry maintains a dictionary of provider names to provider classes +and provides methods for registering new providers, retrieving providers by name, and +listing all available providers. This registry is used by the schema provider factory +to look up provider classes based on configuration. + +Usage: + # Create a registry and register providers + from graphrag_toolkit.document_graph.schema.providers.schema_provider_registry import SchemaProviderRegistry + from graphrag_toolkit.document_graph.schema.providers.file_schema_provider import FileSchemaProvider + from graphrag_toolkit.document_graph.schema.providers.s3_schema_provider import S3SchemaProvider + + registry = SchemaProviderRegistry() + registry.register("file", FileSchemaProvider) + registry.register("s3", S3SchemaProvider) + + # Look up a provider by name + provider_class = registry.get("file") + + # List all registered providers + available_providers = registry.list() +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import Dict, Type +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase + + +class SchemaProviderRegistry: + """ + Central registry for all schema provider classes. + + This class provides a registry system for schema provider classes, allowing for + runtime registration and lookup of providers by name. It centralizes the management + of available schema provider types and ensures that providers can be dynamically + added and retrieved. + + The registry maintains a dictionary mapping provider names to provider classes and + provides methods for registering new providers, retrieving providers by name, and + listing all available providers. This registry is used by the schema provider factory + to look up provider classes based on configuration. + + Attributes: + _registry (Dict[str, Type[SchemaProviderBase]]): A dictionary mapping provider + names to provider classes. + """ + + def __init__(self): + """ + [Brief one-line description of the function's purpose] + + [Detailed explanation of what the function does and how it works] + + Returns: + [return_type]: [Description of return value] + """ + self._registry: Dict[str, Type[SchemaProviderBase]] = {} + + def register(self, name: str, provider_cls: Type[SchemaProviderBase], overwrite: bool = False) -> None: + """ + Register a new schema provider class. + + This method adds a new schema provider class to the registry, associating it + with the specified name. The provider class must be a subclass of SchemaProviderBase. + By default, attempting to register a provider with a name that already exists + in the registry will raise an error, but this behavior can be overridden with + the overwrite parameter. + + Args: + name: Provider type name (e.g., 'csv', 's3'). This is the identifier that + will be used to look up the provider class. + provider_cls: Provider implementation class. Must be a subclass of + SchemaProviderBase and implement the required interface. + overwrite: Allow replacing an existing provider (default: False). If True, + an existing provider with the same name will be replaced. + + Raises: + ValueError: If a provider with the same name already exists and overwrite + is False. + + Example: + # Register a custom provider + registry = SchemaProviderRegistry() + registry.register("custom", MyCustomProvider) + """ + if name in self._registry and not overwrite: + raise ErrorHandler.validation_error( + "schema_provider_registration", + name, + "a unique provider name not already registered" + ) + self._registry[name] = provider_cls + + def get(self, name: str) -> Type[SchemaProviderBase]: + """ + Retrieve a registered schema provider class by name. + + This method looks up a schema provider class in the registry by its registered + name. If the name is not found in the registry, an error is raised with a list + of available provider names. + + Args: + name: The name of the provider to retrieve. This must be a name that has + been previously registered with the register method. + + Returns: + Type[SchemaProviderBase]: The schema provider class associated with the + specified name. + + Raises: + ValueError: If the specified name is not found in the registry. + + Example: + # Get a provider class by name + registry = SchemaProviderRegistry() + provider_class = registry.get("file") + + # Create an instance of the provider + provider = provider_class.from_config(config) + """ + if name not in self._registry: + available = list(self._registry.keys()) + raise ErrorHandler.validation_error( + "schema_provider_lookup", + name, + f"one of: {available}" + ) + return self._registry[name] + + def list(self) -> Dict[str, Type[SchemaProviderBase]]: + """ + Return a copy of all registered providers. + + This method returns a dictionary containing all the schema provider classes + currently registered in the registry. The dictionary maps provider names to + their corresponding provider classes. A copy of the internal registry is + returned to prevent modification of the registry through the returned dictionary. + + Returns: + Dict[str, Type[SchemaProviderBase]]: A dictionary mapping provider names + to their corresponding provider classes. + + Example: + # List all registered providers + registry = SchemaProviderRegistry() + providers = registry.list() + + # Print the names of all registered providers + for name in providers.keys(): + print(f"Provider: {name}") + """ + return self._registry.copy() diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/xml_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/xml_schema_provider.py new file mode 100644 index 00000000..69db6af6 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/xml_schema_provider.py @@ -0,0 +1,157 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""XML Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for XML (eXtensible Markup Language) files. +It discovers and generates ETL schemas from XML files for processing and loading +the data into a document graph. + +The module includes the following components: +- XMLSchemaProvider: Schema provider for XML files + +The XML schema provider uses the XMLSchemaDiscoveryProvider to analyze XML files +and generate appropriate ETL schemas. It supports loading schemas from XML files +and optionally saving the discovered schemas to JSON files. + +Usage: + # Get an XML schema provider for a specific file + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "xml", + "schema_id": "my_xml_schema", + "connection_config": { + "path": "/path/to/data.xml" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging +import json +from pathlib import Path +from typing import Optional, Any + +# Import custom logging to ensure configuration is available + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.discovery.xml_discovery_provider import XMLSchemaDiscoveryProvider + +logger = logging.getLogger(__name__) + + +class XMLSchemaProvider(SchemaProviderBase): + """ + Schema provider for XML (eXtensible Markup Language) files. + + This class provides functionality to discover and generate ETL schemas from XML files. + It uses the XMLSchemaDiscoveryProvider to analyze XML files and extract field information + to create appropriate ETL schemas for processing and loading the data into a document graph. + + The provider supports loading schemas directly from XML files and optionally saving + the discovered schemas to JSON files for later use. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + xml_path (Path): The path to the XML file to analyze. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize an XML schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the path to the XML file in the connection_config. + + Raises: + ValueError: If the path is not provided in the connection_config. + FileNotFoundError: If the specified XML file does not exist. + """ + self.config = config + path = config.connection_config.get("path") + if not path: + raise ValueError("Validation error") + self.xml_path = Path(path) + + if not self.xml_path.exists(): + raise FileNotFoundError(f"File not found") + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "XMLSchemaProvider": + """ + Factory method to construct an XML provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + XMLSchemaProvider: A new instance of the XML schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load (or generate) an ETL schema from an XML file using schema discovery. + + This method uses the XMLSchemaDiscoveryProvider to analyze the XML file and + generate an appropriate ETL schema based on the structure and content of the file. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters (not used in this provider). + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the XML file and how it should be processed. + + Raises: + FileNotFoundError: If the XML file does not exist. + ValueError: If the XML file cannot be parsed or has no discoverable fields. + """ + discovery = XMLSchemaDiscoveryProvider(source=self.xml_path) + return discovery.discover_schema() + + def save_schema(self, output_path: Path) -> None: + """ + Save the inferred ETL schema to a JSON file at the given output path. + + This method discovers the schema from the XML file and saves it to a JSON file + at the specified path. The schema can then be loaded directly using a FileSchemaProvider. + + Args: + output_path (Path): The path where the schema JSON file should be saved. + + Raises: + IOError: If the schema cannot be saved to the output path. + """ + schema = self.load_schema() + try: + with output_path.open("w", encoding="utf-8") as f: + json.dump(schema.model_dump(mode="json", exclude_unset=True), f, indent=2) + logger.info(f"Saved XML schema to file: {output_path}") + except Exception as e: + raise IOError(f"IO error: {output_path, e}") + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the XML file name. + + Returns: + str: A unique identifier for the schema. + """ + return self.config.schema_id or f"xml-schema-{self.xml_path.stem}" \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/providers/yaml_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/yaml_schema_provider.py new file mode 100644 index 00000000..1ba32021 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/providers/yaml_schema_provider.py @@ -0,0 +1,158 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""YAML Schema Provider Module for Document Graph Operations. + +This module provides a schema provider for YAML (YAML Ain't Markup Language) files. +It discovers and generates ETL schemas from YAML files for processing and loading +the data into a document graph. + +The module includes the following components: +- YAMLSchemaProvider: Schema provider for YAML files + +The YAML schema provider uses the YAMLSchemaDiscoveryProvider to analyze YAML files +and generate appropriate ETL schemas. It supports loading schemas from YAML files +and optionally saving the discovered schemas to JSON files. + +Usage: + # Get a YAML schema provider for a specific file + from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider + + config = { + "provider_type": "yaml", + "schema_id": "my_yaml_schema", + "connection_config": { + "path": "/path/to/data.yaml" + } + } + + provider = get_schema_provider(config) + + # Load the schema + schema = provider.pipeline_schema() + + # Use the schema for ETL operations + print(f"Loaded schema ID: {schema.schema_id}") + print(f"Fields: {schema.load.document_node.fields}") +""" + +import logging +import json +from pathlib import Path +from typing import Optional, Any + +# Import custom logging to ensure configuration is available + +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.discovery.yaml_discovery_provider import YAMLSchemaDiscoveryProvider + +logger = logging.getLogger(__name__) + + +class YAMLSchemaProvider(SchemaProviderBase): + """ + Schema provider for YAML (YAML Ain't Markup Language) files. + + This class provides functionality to discover and generate ETL schemas from YAML files. + It uses the YAMLSchemaDiscoveryProvider to analyze YAML files and extract field information + to create appropriate ETL schemas for processing and loading the data into a document graph. + + The provider supports loading schemas directly from YAML files and optionally saving + the discovered schemas to JSON files for later use. + + Attributes: + config (SchemaProviderConfig): Configuration for the provider. + yaml_path (Path): The path to the YAML file to analyze. + """ + + def __init__(self, config: SchemaProviderConfig): + """ + Initialize a YAML schema provider with the given configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider, including + the path to the YAML file in the connection_config. + + Raises: + ValueError: If the path is not provided in the connection_config. + FileNotFoundError: If the specified YAML file does not exist. + """ + self.config = config + path = config.connection_config.get("path") + if not path: + raise ValueError("Validation error") + self.yaml_path = Path(path) + + if not self.yaml_path.exists(): + raise FileNotFoundError(f"File not found") + + @classmethod + def from_config(cls, config: SchemaProviderConfig) -> "YAMLSchemaProvider": + """ + Factory method to construct a YAML provider from a configuration. + + Args: + config (SchemaProviderConfig): Configuration for the provider. + + Returns: + YAMLSchemaProvider: A new instance of the YAML schema provider. + """ + return cls(config) + + def load_schema(self, source=None, **kwargs) -> ETLSchema: + """ + Load (or generate) an ETL schema from a YAML file using schema discovery. + + This method uses the YAMLSchemaDiscoveryProvider to analyze the YAML file and + generate an appropriate ETL schema based on the structure and content of the file. + + Args: + source: Optional source override (not used in this provider). + **kwargs: Additional provider-specific parameters (not used in this provider). + + Returns: + ETLSchema: A complete ETL schema object that describes the structure + of the YAML file and how it should be processed. + + Raises: + FileNotFoundError: If the YAML file does not exist. + ValueError: If the YAML file cannot be parsed or has no discoverable fields. + ImportError: If PyYAML is not installed. + """ + discovery = YAMLSchemaDiscoveryProvider(source=self.yaml_path) + return discovery.discover_schema() + + def save_schema(self, output_path: Path) -> None: + """ + Save the inferred ETL schema to a JSON file at the given output path. + + This method discovers the schema from the YAML file and saves it to a JSON file + at the specified path. The schema can then be loaded directly using a FileSchemaProvider. + + Args: + output_path (Path): The path where the schema JSON file should be saved. + + Raises: + IOError: If the schema cannot be saved to the output path. + """ + schema = self.load_schema() + try: + with output_path.open("w", encoding="utf-8") as f: + json.dump(schema.model_dump(mode="json", exclude_unset=True), f, indent=2) + logger.info(f"Saved YAML schema to file: {output_path}") + except Exception as e: + raise IOError(f"IO error: {output_path, e}") + + def get_schema_id(self) -> str: + """ + Return the unique identifier for the schema. + + This method returns a unique identifier for the schema, either from the + configuration or generated from the YAML file name. + + Returns: + str: A unique identifier for the schema. + """ + return self.config.schema_id or f"yaml-schema-{self.yaml_path.stem}" \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/schema_io.py b/document-graph/src/graphrag_toolkit/document_graph/schema/schema_io.py new file mode 100644 index 00000000..95518df1 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/schema_io.py @@ -0,0 +1,113 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema IO Module for Document Graph Operations. + +This module provides utility functions for saving and loading ETL schemas to and from +various storage formats. It handles serialization and deserialization of ETL schema +objects, making it easy to persist schemas and share them between different components +of the document graph system. + +The module includes functions for: +- Saving ETL schemas to JSON files or file-like objects +- Loading ETL schemas from JSON files or file-like objects + +These functions handle the conversion between ETLSchema objects and their JSON +representations, ensuring that all schema components are properly serialized and +deserialized. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +import json +from pathlib import Path +from typing import Union, IO +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema + +def save_schema(schema: ETLSchema, output: Union[str, Path, IO], indent: int = 2) -> None: + """ + Save an ETL schema to a file or file-like object in JSON format. + + This function serializes an ETLSchema object to JSON and writes it to the specified + output. The output can be a file path (as a string or Path object) or a file-like + object that supports the write method. + + Args: + schema (ETLSchema): The ETL schema to save. + output (Union[str, Path, IO]): The output destination, which can be a file path + (as a string or Path object) or a file-like object. + indent (int, optional): The indentation level for the JSON output. Defaults to 2. + + Raises: + IOError: If there's an error writing to the output. + + Example: + ```python + from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema + from graphrag_toolkit.document_graph.schema.schema_io import save_schema + + # Create a schema + schema = ETLSchema(...) + + # Save to a file + save_schema(schema, "path/to/schema.json") + + # Or save to a file-like object + with open("path/to/schema.json", "w") as f: + save_schema(schema, f) + ``` + """ + try: + if isinstance(output, (str, Path)): + with Path(output).open("w", encoding="utf-8") as f: + json.dump(schema.model_dump(mode="json", exclude_unset=True), f, indent=indent) + else: + json.dump(schema.model_dump(mode="json", exclude_unset=True), output, indent=indent) + except Exception as e: + raise IOError(f"IO error: {str(output), e}") + + +def load_schema(input_: Union[str, Path, IO]) -> ETLSchema: + """ + Load an ETL schema from a file or file-like object. + + This function deserializes an ETLSchema object from JSON read from the specified + input. The input can be a file path (as a string or Path object) or a file-like + object that supports the read method. + + Args: + input_ (Union[str, Path, IO]): The input source, which can be a file path + (as a string or Path object) or a file-like object. + + Returns: + ETLSchema: The deserialized ETL schema. + + Raises: + IOError: If there's an error reading from the input. + ValueError: If the input contains invalid JSON or the JSON doesn't represent a valid ETLSchema. + + Example: + ```python + from graphrag_toolkit.document_graph.schema.schema_io import pipeline_schema + + # Load from a file + schema = pipeline_schema("path/to/schema.json") + + # Or load from a file-like object + with open("path/to/schema.json", "r") as f: + schema = pipeline_schema(f) + ``` + """ + try: + if isinstance(input_, (str, Path)): + with Path(input_).open("r", encoding="utf-8") as f: + return ETLSchema(**json.load(f)) + else: + return ETLSchema(**json.load(input_)) + except Exception as e: + raise IOError(f"IO error: {str(input_), e}") diff --git a/document-graph/src/graphrag_toolkit/document_graph/schema/static_schema_provider.py b/document-graph/src/graphrag_toolkit/document_graph/schema/static_schema_provider.py new file mode 100644 index 00000000..d852477d --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/schema/static_schema_provider.py @@ -0,0 +1,135 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Static Schema Provider Module for Document Graph Operations. + +This module provides a schema provider that returns a hardcoded, static ETL schema. +It's primarily used for testing, development, or as a fallback when no other schema +providers are available. + +The static schema provider implements the SchemaProviderBase interface, making it +compatible with the rest of the document graph system. It provides a default schema +that includes configurations for extracting from S3, transforming with standard +operations, and loading into a graph with document and section nodes. + +This provider is useful for: +- Testing document graph operations without setting up external schema sources +- Providing a fallback schema when user-defined schemas are not available +- Demonstrating the structure of a complete ETL schema +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import Dict, Any +from graphrag_toolkit.document_graph.schema.providers.schema_provider_base import SchemaProviderBase +from graphrag_toolkit.document_graph.schema.etl_schema_model import ETLSchema, ExtractConfig, TransformConfig, LoadConfig, ChunkingConfig, MetadataMapping, EntityExtractionConfig, NormalizeConfig, NodeDefinition, RelationshipDefinition + + +class StaticSchemaProvider(SchemaProviderBase): + """ + A schema provider that returns a hardcoded, static ETL schema. + + This class implements the SchemaProviderBase interface to provide a default + ETL schema without requiring any external configuration sources. It initializes + with a complete, hardcoded schema that includes all necessary components for + the ETL process. + + The static schema includes: + - Extract configuration for S3 PDF documents + - Transform configuration with chunking, metadata mapping, entity extraction, and normalization + - Load configuration with document and section nodes and their relationships + + This provider is particularly useful for: + - Testing and development environments + - Providing a fallback when no user-defined schema is available + - Demonstrating a complete schema structure + + Attributes: + _schema (ETLSchema): The hardcoded ETL schema instance. + """ + + def __init__(self, config: Dict[str, Any]): + """ + Initialize a static schema provider with a hardcoded ETL schema. + + This constructor creates a complete ETL schema with predefined settings for + all components. The config parameter is accepted to maintain a consistent + interface with other schema providers, but it is not used in this implementation. + + Args: + config (Dict[str, Any]): Configuration dictionary (unused in this provider). + Included for interface consistency with other providers. + """ + # Config is unused for static schema, but accepted to keep constructor consistent + self._schema = ETLSchema( + schema_id="static-default", + description="Default hardcoded ETL schema", + extract=ExtractConfig( + source_type="s3", + bucket="default-bucket", + prefix="docs/", + file_type="pdf", + reader="pymupdf" + ), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="by_heading", min_length=100), + metadata_mapping=MetadataMapping( + title="document.title", + author="metadata.author" + ), + entity_extraction=EntityExtractionConfig(method="ner", model="spacy_en_core_web_lg"), + normalize=NormalizeConfig(remove_headers=True) + ), + load=LoadConfig( + document_node=NodeDefinition( + type="DocumentNode", + fields=["title", "document_type", "metadata"] + ), + section_node=NodeDefinition( + type="SectionNode", + fields=["title", "level", "text_units"] + ), + relationships=[ + RelationshipDefinition( + type="has_section", + source="document_id", + target="section_id" + ) + ] + ) + ) + + @classmethod + def from_config(cls, config): + """Factory method — StaticSchemaProvider ignores config.""" + return cls(config if isinstance(config, dict) else {}) + + def load_schema(self) -> ETLSchema: + """ + Load and return the static ETL schema. + + This method implements the SchemaProviderBase interface's pipeline_schema method. + Unlike other schema providers that might load from external sources, this method + simply returns the hardcoded schema that was created during initialization. + + Returns: + ETLSchema: The static, hardcoded ETL schema. + """ + return self._schema + + def get_schema_id(self) -> str: + """ + Get the unique identifier for the static schema. + + This method implements the SchemaProviderBase interface's get_schema_id method. + It returns the schema_id property from the hardcoded schema. + + Returns: + str: The unique identifier for the schema (e.g., "static-default"). + """ + return self._schema.schema_id diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/__init__.py new file mode 100644 index 00000000..a8995c45 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/__init__.py @@ -0,0 +1 @@ +"""Transform stage — rows → typed Nodes/Edges.""" diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/__init__.py new file mode 100644 index 00000000..98cf1539 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/__init__.py @@ -0,0 +1,23 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Document-level transformers for processing and manipulating document content. + +This module provides transformers that operate on document content, such as: +- PII redaction: Removing personally identifiable information +- Text chunking: Splitting long text into manageable pieces +- JSON flattening: Converting nested JSON structures to tabular format + +These transformers can be used in document processing pipelines to prepare +content for storage, analysis, or retrieval. +""" + +from .pii_redactor_provider import PIIRedactorProvider +from .json_to_rows import JSONToRowsTransformer +from .text_chunker import TextChunkerTransformer + +__all__ = [ + 'PIIRedactorProvider', + 'JSONToRowsTransformer', + 'TextChunkerTransformer' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/json_to_rows.py b/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/json_to_rows.py new file mode 100644 index 00000000..d6708dfd --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/json_to_rows.py @@ -0,0 +1,79 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON to Rows transformer for flattening nested JSON structures. + +This module provides functionality to transform nested JSON data into a flat +tabular structure, making it easier to process and analyze hierarchical data +in a row-based format. It uses pandas json_normalize to flatten nested JSON +structures into rows with dot-notation column names. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +import json +import pandas as pd +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +class JSONToRowsTransformer(TransformerProvider): + """Transforms nested JSON fields into flat tabular rows. + + This transformer takes records containing nested JSON data and converts them + into multiple flattened records, where each nested object becomes a separate row. + The nested structure is flattened using dot notation for field names. + + Configuration: + json_field (str): The field name containing the JSON data to flatten. + Defaults to 'content'. + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records with JSON fields into flattened rows. + + This method processes each record by: + 1. Extracting the JSON data from the specified field + 2. Parsing the JSON if it's a string + 3. Normalizing the nested structure into a flat DataFrame + 4. Creating new records for each row in the DataFrame + + If JSON parsing fails, the original record is preserved. + + Args: + records: A list of record dictionaries to transform + + Returns: + A list of transformed records, where each nested JSON object + has been converted to a separate record with flattened fields + prefixed with 'json_'. + """ + json_field = self.args.get('json_field', 'content') + output_records = [] + + for record in records: + try: + json_data = record.get(json_field) + if isinstance(json_data, str): + json_data = json.loads(json_data) + + # Normalize JSON to flat structure + df = pd.json_normalize(json_data) + + for idx, row in df.iterrows(): + new_record = record.copy() + # Add flattened fields to record + for col, val in row.items(): + new_record[f"json_{col}"] = val + new_record['row_index'] = idx + output_records.append(new_record) + + except Exception: + # Keep original record if JSON parsing fails + output_records.append(record) + + return output_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/pii_redactor_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/pii_redactor_provider.py new file mode 100644 index 00000000..a6bce981 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/pii_redactor_provider.py @@ -0,0 +1,77 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""PII Redactor Provider for removing sensitive information from documents. + +This module provides functionality to identify and redact Personally Identifiable +Information (PII) from document content. It uses the 'sanitary' library to detect +and replace sensitive information such as IP addresses, login names, credentials, +and other PII with redaction markers. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import Any, Dict, List + +try: + from sanitary import Sanitizer +except ImportError: + Sanitizer = None + +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider +from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + + +class PIIRedactorProvider(TransformerProvider): + """Redacts PII such as IP addresses, login names, credentials from text fields. + + This transformer identifies and redacts Personally Identifiable Information (PII) + from specified fields in document records. It uses the `sanitary` library to + detect and replace sensitive information with a redaction marker. + + Configuration: + fields (List[str]): List of field names to scan for PII and redact. + If empty, no fields will be redacted. + """ + + def __init__(self, config: TransformerProviderConfig): + """Initialize the PII redactor with configuration. + + Args: + config: Configuration object containing transformer settings. + Must include an 'args' dictionary with a 'fields' list + specifying which fields to redact. + """ + if Sanitizer is None: + raise ImportError("PIIRedactorProvider requires 'sanitary' package: pip install sanitary") + super().__init__(config) + self.fields_to_redact = config.args.get("fields", []) + self.sanitizer = Sanitizer( + keys=self.fields_to_redact, + replacement="***REDACTED***" + ) + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by redacting PII from specified fields. + + This method processes each record by applying the sanitizer to + redact PII in the configured fields. The original record structure + is preserved, with only the sensitive information replaced. + + Args: + records: A list of record dictionaries to process + + Returns: + A list of records with PII redacted from the specified fields. + Each record maintains its original structure. + """ + redacted_records = [] + for record in records: + redacted = self.sanitizer.sanitize(record) + redacted_records.append(redacted) + return redacted_records diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/text_chunker.py b/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/text_chunker.py new file mode 100644 index 00000000..41d5a269 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/document_transformers/text_chunker.py @@ -0,0 +1,102 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Text Chunker module for splitting document content into manageable chunks. + +This module provides functionality to divide long text content into smaller, +overlapping chunks for more efficient processing and analysis. This is particularly +useful for large documents that need to be processed in smaller segments, such as +for embedding generation, semantic search, or other NLP operations that have +input size limitations. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class TextChunkerTransformer(TransformerProvider): + """Splits long text content into smaller, potentially overlapping chunks. + + This transformer divides text content from records into smaller chunks + of a specified size, with optional overlap between consecutive chunks. + Each chunk becomes a separate record that maintains the original record's + metadata while adding chunk-specific information. + + Configuration: + chunk_size (int): Maximum size of each text chunk in characters. + Defaults to 512. + overlap (int): Number of characters to overlap between consecutive chunks. + Defaults to 0 (no overlap). + text_field (str): The field name containing the text to chunk. + Defaults to 'content'. + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by splitting text content into smaller chunks. + + This method processes each record by: + 1. Extracting the text from the specified field + 2. Dividing it into chunks of the configured size + 3. Creating new records for each chunk with additional metadata + + Records without text content or with non-string content in the specified + field are passed through unchanged. + + Each chunk record includes the following additional fields: + - chunk_index: Sequential index of the chunk (0-based) + - chunk_start: Character position where the chunk starts in the original text + - chunk_end: Character position where the chunk ends in the original text + - original_id: The ID of the original record + + If the original record has an 'id' field, each chunk's ID will be + formatted as "{original_id}-chunk-{chunk_index}". + + Args: + records: A list of record dictionaries to transform + + Returns: + A list of transformed records, where text content has been + divided into chunks with each chunk as a separate record. + """ + chunk_size = int(self.args.get("chunk_size", 512)) + overlap = int(self.args.get("overlap", 0)) + text_field = self.args.get("text_field", "content") + + output_records = [] + + for record in records: + text = record.get(text_field, "") + if not text or not isinstance(text, str): + output_records.append(record) + continue + + start = 0 + chunk_index = 0 + + while start < len(text): + end = start + chunk_size + chunk = text[start:end] + + chunk_record = record.copy() + chunk_record[text_field] = chunk + chunk_record['chunk_index'] = chunk_index + chunk_record['chunk_start'] = start + chunk_record['chunk_end'] = min(end, len(text)) + chunk_record['original_id'] = record.get('id', '') + + if 'id' in chunk_record: + chunk_record['id'] = f"{chunk_record['id']}-chunk-{chunk_index}" + + output_records.append(chunk_record) + + chunk_index += 1 + start = end - overlap + + return output_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/__init__.py new file mode 100644 index 00000000..b00645a5 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/__init__.py @@ -0,0 +1,25 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Enricher transformers for adding information to documents. + +This module provides transformers that enrich document content by adding +additional information such as: +- Language detection: Automatically detect document language +- LLM enrichment: Use large language models to add metadata + +These transformers enhance document processing by providing additional +context and metadata that can be used for analysis and retrieval. +""" + +from .language_enricher_provider import LanguageDetectionEnricher +from .llm_enricher_plugin import LLMEnricherPlugin +from .bedrock_enricher_plugin import BedrockEnricherPlugin +from .remediation_factory_provider import RemediationFactoryProvider + +__all__ = [ + 'LanguageDetectionEnricher', + 'LLMEnricherPlugin', + 'BedrockEnricherPlugin', + 'RemediationFactoryProvider' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/bedrock_enricher_plugin.py b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/bedrock_enricher_plugin.py new file mode 100644 index 00000000..161ea556 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/bedrock_enricher_plugin.py @@ -0,0 +1,204 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bedrock Enricher Plugin for enhancing documents with AI-generated content.""" + +import logging +import json +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider +from graphrag_toolkit.document_graph.document_graph_config import document_graph_config + +logger = logging.getLogger(__name__) + + +class BedrockEnricherPlugin(TransformerProvider): + """Enriches records using AWS Bedrock models for content analysis.""" + + def __init__(self, config): + super().__init__(config) + + self.task = self.args.get("task", "summarize") + self.input_field = self.args.get("field", "content") + self.model_id = self.args.get("model_id", "anthropic.claude-3-haiku-20240307-v1:0") + self.prefix = self.args.get("prefix", "bedrock_") + self.suffix = self.args.get("suffix", "") + self.max_tokens = self.args.get("max_tokens", 200) + self.temperature = self.args.get("temperature", 0.3) + self.enable_dedup = self.args.get("enable_dedup", True) + self.cache_dir = self.args.get("cache_dir", None) + self._permission_error = False # Circuit breaker for permission errors + self._cache = {} # In-memory cache for deduplication + + # Setup persistent cache if cache_dir is provided + if self.cache_dir and self.enable_dedup: + import os + os.makedirs(self.cache_dir, exist_ok=True) + self._load_cache_from_disk() + + # If suffix is provided, use field name as prefix + if self.suffix: + self.prefix = self.args.get('csv_property_name', self.input_field) + + self.client = document_graph_config._get_or_create_client("bedrock-runtime") + + def _load_prompt_template(self, task: str) -> str: + """Load prompt template with fallback to document graph prompts.""" + import os + + # Try enricher-specific prompts first + enricher_prompt_path = os.path.join( + os.path.dirname(__file__), "prompts", f"evidence_{task}.txt" + ) + + if os.path.exists(enricher_prompt_path): + with open(enricher_prompt_path, 'r') as f: + return f.read().strip() + + # Fallback to document graph prompts + try: + from graphrag_toolkit.document_graph.prompts import get_prompt + return get_prompt(f"evidence_{task}") + except: + pass + + # Default fallback + if task == "summarize": + return "Summarize the key evidence from this JSON data in 2-3 sentences:\n\n{text}" + elif task == "tag": + return "Extract key tags from this evidence data:\n\n{text}" + else: + return "Analyze this evidence data:\n\n{text}" + + def _make_prompt(self, text: str) -> str: + template = self._load_prompt_template(self.task) + return template.format(text=text) + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + enriched_records = [] + + # Skip all processing if we've hit a permission error + if self._permission_error: + logger.warning("Skipping Bedrock enrichment due to previous permission error") + return records + + for record in records: + enriched_record = record.copy() + text = record.get(self.input_field, "") + + if text and isinstance(text, str) and text.strip(): + field_name = f"{self.prefix}{self.suffix}" if self.suffix else f"{self.prefix}{self.task}" + + # Check cache for deduplication + if self.enable_dedup: + import hashlib + text_hash = hashlib.md5(text.encode()).hexdigest() + if text_hash in self._cache: + logger.debug(f"Using cached result for duplicate evidence (hash: {text_hash[:8]}...)") + enriched_record[field_name] = self._cache[text_hash] + enriched_records.append(enriched_record) + continue + + try: + prompt = self._make_prompt(text) + + # Different request formats for different models + if "anthropic.claude" in self.model_id: + body = json.dumps({ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": self.max_tokens, + "temperature": self.temperature, + "messages": [{"role": "user", "content": prompt}] + }) + elif "amazon.nova" in self.model_id: + body = json.dumps({ + "messages": [{"role": "user", "content": [{"text": prompt}]}], + "inferenceConfig": { + "maxTokens": self.max_tokens, + "temperature": self.temperature + } + }) + else: + # Default format + body = json.dumps({ + "inputText": prompt, + "textGenerationConfig": { + "maxTokenCount": self.max_tokens, + "temperature": self.temperature + } + }) + + response = self.client.invoke_model( + modelId=self.model_id, + body=body + ) + + response_body = json.loads(response['body'].read()) + + # Different response formats for different models + if "anthropic.claude" in self.model_id: + output = response_body['content'][0]['text'].strip() + elif "amazon.nova" in self.model_id: + output = response_body['output']['message']['content'][0]['text'].strip() + else: + output = response_body.get('results', [{}])[0].get('outputText', '').strip() + + enriched_record[field_name] = output + + # Cache the result for deduplication + if self.enable_dedup: + self._cache[text_hash] = output + logger.info(f"Cached new result for evidence hash: {text_hash[:8]}...") + # Save to disk if cache_dir is configured + if self.cache_dir: + self._save_cache_to_disk() + + except Exception as e: + error_msg = str(e) + + # Check for permission errors and set circuit breaker + if "AccessDeniedException" in error_msg or "bedrock:InvokeModel" in error_msg: + self._permission_error = True + logger.error(f"Bedrock permission error - stopping enrichment: {error_msg}") + error_field = f"{self.prefix}{self.suffix}_error" if self.suffix else f"{self.prefix}{self.task}_error" + enriched_record[error_field] = "Permission denied for Bedrock access" + enriched_records.append(enriched_record) + break # Stop processing immediately + else: + logger.warning(f"Bedrock enrichment failed: {error_msg}") + error_field = f"{self.prefix}{self.suffix}_error" if self.suffix else f"{self.prefix}{self.task}_error" + enriched_record[error_field] = error_msg + + enriched_records.append(enriched_record) + + return enriched_records + + def _load_cache_from_disk(self): + """Load cache from disk if it exists.""" + import os + import json + + cache_file = os.path.join(self.cache_dir, f"bedrock_cache_{self.task}.json") + if os.path.exists(cache_file): + try: + with open(cache_file, 'r') as f: + self._cache = json.load(f) + logger.debug(f"Loaded {len(self._cache)} cached results from {cache_file}") + except Exception as e: + logger.warning(f"Failed to load cache from {cache_file}: {e}") + + def _save_cache_to_disk(self): + """Save cache to disk.""" + if not self.cache_dir: + return + + import os + import json + + cache_file = os.path.join(self.cache_dir, f"bedrock_cache_{self.task}.json") + try: + with open(cache_file, 'w') as f: + json.dump(self._cache, f, indent=2) + logger.debug(f"Saved {len(self._cache)} cached results to {cache_file}") + except Exception as e: + logger.warning(f"Failed to save cache to {cache_file}: {e}") \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/language_enricher_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/language_enricher_provider.py new file mode 100644 index 00000000..9482cc1b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/language_enricher_provider.py @@ -0,0 +1,77 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Language Enricher Provider for detecting document languages. + +This module provides functionality to automatically detect the language of text +content in documents. It uses the 'langdetect' library to identify the language +of text fields and adds this information as a new field in the document record. +This enrichment is useful for multilingual document collections, enabling filtering, +routing, or specialized processing based on document language. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from langdetect import detect +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class LanguageDetectionEnricher(TransformerProvider): + """Enriches records with language detection for text content. + + This transformer analyzes text content in records and determines the language + of the text using the 'langdetect' library. It adds a new field to each record + containing the detected language code (e.g., 'en' for English, 'es' for Spanish). + + Configuration: + text_field (str): The field name containing the text to analyze. + Defaults to 'content'. + output_field (str): The field name where the detected language will be stored. + Defaults to 'language'. + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by detecting and adding language information. + + This method processes each record by: + 1. Extracting the text from the specified field + 2. Using langdetect to identify the language of the text + 3. Adding the detected language code to a new field in the record + + If text detection fails or the text field is empty/non-string, + the language is set to "unknown". + + Args: + records: A list of record dictionaries to process + + Returns: + A list of enriched records with language information added. + Each record maintains its original structure with an additional + language field. + """ + text_field = self.args.get("text_field", "content") + output_field = self.args.get("output_field", "language") + + enriched_records = [] + for record in records: + enriched_record = record.copy() + text = record.get(text_field, "") + + if text and isinstance(text, str): + try: + language = detect(text) + except Exception: + language = "unknown" + enriched_record[output_field] = language + else: + enriched_record[output_field] = "unknown" + + enriched_records.append(enriched_record) + + return enriched_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/llm_enricher_plugin.py b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/llm_enricher_plugin.py new file mode 100644 index 00000000..c1c4a4f0 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/llm_enricher_plugin.py @@ -0,0 +1,153 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LLM Enricher Plugin for enhancing documents with AI-generated content. + +This module provides functionality to enrich document content using Large Language +Models (LLMs) via the OpenAI API. It can perform various tasks such as tagging, +summarization, and classification of document content, adding the AI-generated +outputs as new fields in the document records. + +This enrichment is useful for adding semantic metadata to documents, generating +concise summaries, or categorizing content without manual intervention. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +import os +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +try: + from openai import OpenAI +except ImportError: + OpenAI = None + + +class LLMEnricherPlugin(TransformerProvider): + """Enriches records using OpenAI's LLMs for content analysis and enhancement. + + This transformer uses OpenAI's API to process text content in records and + generate additional metadata such as tags, summaries, or classifications. + It adds the LLM-generated content as new fields in each record, preserving + the original content. + + Configuration: + task (str): The type of enrichment to perform. Options are: + - "tag": Generate relevant tags/categories for the content + - "summarize": Create a concise summary of the content + - "classify": Identify the topic or domain of the content + Defaults to "tag". + field (str): The field name containing the text to analyze. + Defaults to "content". + llm_model (str): The OpenAI model to use for generation. + Defaults to "gpt-4". + api_key (str): OpenAI API key. If not provided, will look for + OPENAI_API_KEY environment variable. + """ + + def __init__(self, config): + """Initialize the LLM enricher with configuration. + + This method sets up the OpenAI client and configures the enricher + based on the provided configuration. It validates that the OpenAI + library is installed and that an API key is available. + + Args: + config: Configuration object containing transformer settings. + Must include an 'args' dictionary with task, field, + and model settings. Can optionally include an API key. + + Raises: + ImportError: If the 'openai' library is not installed. + ValueError: If no OpenAI API key is provided in config or environment. + """ + super().__init__(config) + if OpenAI is None: + raise ImportError("Please install the 'openai' library to use LLMEnricherPlugin") + + self.task = self.args.get("task", "tag") + self.input_field = self.args.get("field", "content") + self.model = self.args.get("llm_model", "gpt-4") + + api_key = self.args.get("api_key") or os.environ.get("OPENAI_API_KEY") + if not api_key: + raise ValueError("Missing OpenAI API key in config or environment") + + self.client = OpenAI(api_key=api_key) + + def _make_prompt(self, text: str) -> str: + """Create an appropriate prompt for the configured task. + + This method generates a task-specific prompt to send to the LLM, + based on the task type configured for this enricher (tag, summarize, + or classify). + + Args: + text: The document text to be processed by the LLM + + Returns: + A formatted prompt string that instructs the LLM on the task + to perform on the provided text. + + Raises: + ValueError: If the configured task is not one of the supported types. + """ + if self.task == "tag": + return f"Tag the following content with relevant categories:\n\n{text}" + elif self.task == "summarize": + return f"Summarize the following content:\n\n{text}" + elif self.task == "classify": + return f"Classify the topic or domain of the content below:\n\n{text}" + else: + raise ValueError(f"Unknown task: {self.task}") + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by enriching them with LLM-generated content. + + This method processes each record by: + 1. Extracting the text from the configured input field + 2. Creating a task-specific prompt for the LLM + 3. Sending the prompt to the OpenAI API + 4. Adding the LLM's response to the record in a new field + + The new field is named based on the task (e.g., "llm_tag", "llm_summarize"). + If the API call fails, an error field is added instead (e.g., "llm_tag_error"). + + Records with empty or non-string content in the input field are passed + through with minimal processing. + + Args: + records: A list of record dictionaries to enrich + + Returns: + A list of enriched records with LLM-generated content added. + Each record maintains its original structure with additional + fields containing the enrichment results. + """ + enriched_records = [] + for record in records: + enriched_record = record.copy() + text = record.get(self.input_field, "") + + if text and isinstance(text, str) and text.strip(): + try: + prompt = self._make_prompt(text) + response = self.client.chat.completions.create( + model=self.model, + messages=[{"role": "user", "content": prompt}], + temperature=0.3, + ) + output = response.choices[0].message.content.strip() + enriched_record[f"llm_{self.task}"] = output + except Exception as e: + enriched_record[f"llm_{self.task}_error"] = str(e) + + enriched_records.append(enriched_record) + + return enriched_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/prompts/evidence_summarize.txt b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/prompts/evidence_summarize.txt new file mode 100644 index 00000000..55e572a1 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/prompts/evidence_summarize.txt @@ -0,0 +1,8 @@ +Analyze the following security evidence data and provide a concise 2-3 sentence summary of the key findings and resources involved: + +{text} + +Focus on: +- What type of evidence objects are present +- Key security findings or configuration issues +- Affected resources and their identifiers \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/prompts/evidence_tag.txt b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/prompts/evidence_tag.txt new file mode 100644 index 00000000..dfc4f54b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/prompts/evidence_tag.txt @@ -0,0 +1,8 @@ +Extract relevant security tags and categories from the following evidence data. Return as comma-separated values: + +{text} + +Focus on extracting tags for: +- Evidence types (e.g., DATABASE, FIREWALL, CONFIGURATION_FINDING) +- Security domains (e.g., access-control, encryption, compliance) +- Severity levels and risk categories \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/remediation_factory_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/remediation_factory_provider.py new file mode 100644 index 00000000..6e0a6d28 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/enrichers/remediation_factory_provider.py @@ -0,0 +1,129 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Remediation Factory Provider for AWS remediation input generation.""" + +import json +import hashlib +import logging +from typing import Dict, Any, List +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider +from graphrag_toolkit.document_graph.document_graph_config import document_graph_config + +logger = logging.getLogger(__name__) + + +class RemediationFactoryProvider(TransformerProvider): + """Provider that generates AWS remediation inputs using S3 factory pattern.""" + + def __init__(self, config): + super().__init__(config) + self.s3_bucket = self.args.get("s3_bucket") + self.s3_prefix = self.args.get("s3_prefix", "inputs/").rstrip('/') + '/' + self.model_id = self.args.get("model_id", "amazon.nova-pro-v1:0") + self.source_fields = self.args.get("source_fields", ["Resource Type", "Resource external ID", "Title", "Remediation Recommendation"]) + self.target_field = self.args.get("target_field", "remediation_input") + self.prompt_path = self.args.get("prompt_path") + self.cloud_provider_field = self.args.get("cloud_provider_field", "Resource Platform") + + if not self.s3_bucket: + raise ValueError("s3_bucket required") + + self.s3_client = document_graph_config._get_or_create_client("s3") + # Add timeout configuration + self.s3_client.meta.config.read_timeout = 30 + self.s3_client.meta.config.connect_timeout = 10 + self.bedrock_client = document_graph_config._get_or_create_client("bedrock-runtime") + self._prompt_template = None + + def _get_cloud_fallback_prompt(self, cloud_provider): + """Generate cloud-specific fallback prompt""" + prompts = { + "aws": "Generate AWS remediation for:\n{input_text}\nReturn JSON: {{\"service\": \"aws_service\", \"method\": \"boto3_method\", \"description\": \"fix_description\"}}", + "azure": "Generate Azure remediation for:\n{input_text}\nReturn JSON: {{\"service\": \"azure_service\", \"method\": \"azure_method\", \"description\": \"fix_description\"}}", + "gcp": "Generate GCP remediation for:\n{input_text}\nReturn JSON: {{\"service\": \"gcp_service\", \"method\": \"gcp_method\", \"description\": \"fix_description\"}}" + } + return prompts.get(cloud_provider.lower(), prompts["aws"]) + + def _load_prompt(self, cloud_provider="aws"): + if self._prompt_template: + return self._prompt_template + if self.prompt_path: + if self.prompt_path.startswith('s3://'): + bucket, key = self.prompt_path[5:].split('/', 1) + response = self.s3_client.get_object(Bucket=bucket, Key=key) + self._prompt_template = response['Body'].read().decode('utf-8') + else: + with open(self.prompt_path, 'r') as f: + self._prompt_template = f.read() + else: + self._prompt_template = self._get_cloud_fallback_prompt(cloud_provider) + return self._prompt_template + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by adding remediation inputs.""" + print(f"🔧 RemediationFactory: Processing {len(records)} records") + enriched_records = [] + + for record in records: + enriched_record = record.copy() + input_parts = [ + f"{field}: {record[field]}" + for field in self.source_fields + if field in record and record[field] + ] + if input_parts: + input_text = "\n".join(input_parts) + cache_key = hashlib.sha256(input_text.encode()).hexdigest() + s3_key = f"{self.s3_prefix}{cache_key}.json" + + try: + response = self.s3_client.get_object(Bucket=self.s3_bucket, Key=s3_key) + enriched_record[self.target_field] = json.loads(response['Body'].read().decode('utf-8')) + except: + try: + # Detect cloud provider from record + cloud_provider = "aws" # default + if self.cloud_provider_field in record: + provider_value = str(record[self.cloud_provider_field]).lower() + if "azure" in provider_value or "microsoft" in provider_value: + cloud_provider = "azure" + elif "gcp" in provider_value or "google" in provider_value: + cloud_provider = "gcp" + + prompt = self._load_prompt(cloud_provider).format(input_text=input_text) + body = json.dumps({"messages": [{"role": "user", "content": [{"text": prompt}]}]}) + response = self.bedrock_client.invoke_model(modelId=self.model_id, body=body) + content = json.loads(response['body'].read())['output']['message']['content'][0]['text'] + logger.debug(f"Bedrock response: {content[:200]}...") + + # Simple JSON extraction - just find first { to last } + import re + json_start = content.find('{') + json_end = content.rfind('}') + + if json_start != -1 and json_end != -1: + json_str = content[json_start:json_end + 1] + # Clean control characters + json_str = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', json_str) + logger.debug(f"Extracted JSON: {json_str}") + input_data = json.loads(json_str) + logger.debug(f"Storing to S3: {s3_key}") + try: + self.s3_client.put_object(Bucket=self.s3_bucket, Key=s3_key, Body=json.dumps(input_data, indent=2)) + logger.debug(f"Successfully stored to S3: {s3_key}") + except Exception as s3_error: + logger.error(f"S3 storage failed: {s3_error}") + enriched_record[self.target_field] = input_data + print(f" ✅ Added {self.target_field} to record") + else: + logger.error(f"No JSON found in response: {content}") + raise ValueError("Could not find JSON in response") + except Exception as e: + logger.error(f"Failed to generate remediation input: {e}") + print(f" ❌ Failed to add {self.target_field}: {e}") + + enriched_records.append(enriched_record) + + print(f"✅ RemediationFactory: Completed {len(enriched_records)} records") + return enriched_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/__init__.py new file mode 100644 index 00000000..1a9ea152 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/__init__.py @@ -0,0 +1,43 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Field-level transformers for processing individual fields. + +This module provides transformers that operate on individual fields within +documents, such as: +- Regex cleaning: Clean field values using regular expressions +- Embedded JSON: Extract and process JSON within text fields +- Timestamp normalization: Standardize timestamp formats +- Enum standardization: Normalize enumerated values + +These transformers help clean and standardize field-level data during +document processing pipelines. +""" + +from .regex_cleaner_provider import RegexCleanerProvider +from .embedded_json import EmbeddedJSONFieldTransformer +from .normalize_timestamp import TimestampNormalizer +from .standardize_enum import EnumStandardizer +from .comma_split_provider import CommaSplitProvider +from .json_array_expander import JSONArrayExpanderProvider +from .json_array_flattener import JSONArrayFlattenerProvider +from .comma_flattener import CommaFlattenerProvider +from .paired_flattener import PairedFlattenerProvider +from .json_value_flattener import JSONValueFlattenerProvider +from .json_flattener import JSONFlattenerProvider +from .uuid_generator import UuidGeneratorTransformer + +__all__ = [ + 'RegexCleanerProvider', + 'EmbeddedJSONFieldTransformer', + 'TimestampNormalizer', + 'EnumStandardizer', + 'CommaSplitProvider', + 'JSONArrayExpanderProvider', + 'JSONArrayFlattenerProvider', + 'CommaFlattenerProvider', + 'PairedFlattenerProvider', + 'JSONValueFlattenerProvider', + 'JSONFlattenerProvider', + 'UuidGeneratorTransformer' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/comma_flattener.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/comma_flattener.py new file mode 100644 index 00000000..45c727d2 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/comma_flattener.py @@ -0,0 +1,194 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Comma Flattener transformer for enriching records with comma-separated field data. + +This transformer flattens comma-separated values into individual columns without creating +new rows, allowing you to enrich existing records with structured data from +comma-separated fields and then map those enriched fields to different graph nodes. +""" + +import logging +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +logger = logging.getLogger(__name__) + + +class CommaFlattenerProvider(TransformerProvider): + """Flattens comma-separated values into columns to enrich records. + + This transformer takes comma-separated values and flattens them into individual + columns with a configurable prefix and numbering. It preserves the original row count while + adding new columns for each comma-separated value. + + Args: + field: The field containing comma-separated data (required) + prefix: Prefix for flattened fields (default: field name + "_") + separator: Separator character (default: ",") + strip_whitespace: Strip whitespace from values (default: True) + preserve_original_field: Keep original field (default: False) + max_items: Maximum number of items to extract (default: 10) + """ + + def __init__(self, config): + super().__init__(config) + self.field = self.args.get('field') or config.csv_property_name + if not self.field: + raise ValueError("Comma Flattener requires 'field' argument or csv_property_name") + + # Optional companion field for paired values + self.companion_field = self.args.get('companion_field') + self.companion_suffix = self.args.get('companion_suffix', '_name') + + # Generate prefix from field name if not provided + if 'prefix' not in self.args: + # Clean field name: remove spaces and special chars + clean_field = ''.join(c if c.isalnum() else '' for c in self.field) + self.prefix = clean_field + logger.info(f"Comma Flattener using default prefix: '{self.prefix}' (from field '{self.field}')") + else: + self.prefix = self.args.get('prefix') + logger.info(f"Comma Flattener using explicit prefix: '{self.prefix}'") + + if 'suffix' not in self.args: + self.suffix = '#00' # Default to 2-digit numbers: 01, 02, etc. + logger.info(f"Comma Flattener using default suffix: '{self.suffix}' (auto-incremental numbers)") + else: + self.suffix = self.args.get('suffix') + logger.info(f"Comma Flattener using explicit suffix: '{self.suffix}'") + self.separator = self.args.get('separator', ',') + self.strip_whitespace = self.args.get('strip_whitespace', True) + self.preserve_original_field = self.args.get('preserve_original_field', False) + self.max_items = self.args.get('max_items', 10) + self.unique_nodes = self.args.get('unique_nodes', False) + self._unique_values = set() # Track unique values globally + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by flattening comma-separated values into columns.""" + + # If unique_nodes is enabled, collect all unique values first + if self.unique_nodes: + self._collect_unique_values(records) + logger.info(f"Found {len(self._unique_values)} unique values for field '{self.field}'") + + enriched_records = [] + + for record in records: + try: + enriched_record = self._enrich_record(record) + enriched_records.append(enriched_record) + except Exception as e: + logger.warning(f"Failed to process comma-separated field '{self.field}': {e}") + enriched_records.append(record.copy()) + + return enriched_records + + def _enrich_record(self, original_record: Dict[str, Any]) -> Dict[str, Any]: + """Enrich record with flattened comma-separated data.""" + # Start with original record + if self.preserve_original_field: + enriched_record = original_record.copy() + else: + enriched_record = {k: v for k, v in original_record.items() if k != self.field} + + field_value = original_record.get(self.field) + + # Handle empty/null values + if field_value is None or field_value == '': + return enriched_record + + # Convert to string if not already + if not isinstance(field_value, str): + field_value = str(field_value) + + # Split by separator + items = field_value.split(self.separator) + + # Strip whitespace if requested + if self.strip_whitespace: + items = [item.strip() for item in items] + + # Filter out empty items + items = [item for item in items if item] + + # Handle companion field if specified + companion_items = [] + if self.companion_field: + companion_value = original_record.get(self.companion_field, '') + if companion_value: + companion_items = companion_value.split(self.separator) + if self.strip_whitespace: + companion_items = [item.strip() for item in companion_items] + companion_items = [item for item in companion_items if item] + + # Add columns based on unique_nodes setting + if self.unique_nodes: + # Use consistent mapping for unique values + value_mapping = self._get_unique_column_mapping() + for item in items: + if item in value_mapping: + column_name = value_mapping[item] + enriched_record[column_name] = item + else: + # Add numbered columns up to max_items (original behavior) + for i, item in enumerate(items[:self.max_items], 1): + if self.suffix and '#' in self.suffix: + # Format number according to suffix pattern (e.g., " #00" -> " 01", " 02") + # Split suffix into prefix part (before #) and number format (after #) + hash_index = self.suffix.index('#') + suffix_prefix = self.suffix[:hash_index] # e.g., " " (space) + number_format = self.suffix[hash_index+1:] # e.g., "00" + num_zeros = len(number_format) + formatted_num = str(i).zfill(num_zeros) + column_name = f"{self.prefix}{suffix_prefix}{formatted_num}" + elif self.suffix: + column_name = f"{self.prefix}{self.suffix}{i}" + else: + column_name = f"{self.prefix}{i}" + enriched_record[column_name] = item + + # Add companion field if available + if self.companion_field and i <= len(companion_items): + companion_column = column_name + self.companion_suffix + enriched_record[companion_column] = companion_items[i-1] + + return enriched_record + + def _collect_unique_values(self, records: List[Dict[str, Any]]) -> None: + """Collect all unique values across all records.""" + for record in records: + field_value = record.get(self.field) + if field_value is None or field_value == '': + continue + + if not isinstance(field_value, str): + field_value = str(field_value) + + items = field_value.split(self.separator) + if self.strip_whitespace: + items = [item.strip() for item in items] + + items = [item for item in items if item] + self._unique_values.update(items) + + def _get_unique_column_mapping(self) -> Dict[str, str]: + """Get mapping of unique values to column names.""" + mapping = {} + unique_list = sorted(list(self._unique_values))[:self.max_items] + + for i, value in enumerate(unique_list, 1): + if self.suffix and '#' in self.suffix: + hash_index = self.suffix.index('#') + suffix_prefix = self.suffix[:hash_index] + number_format = self.suffix[hash_index+1:] + num_zeros = len(number_format) + formatted_num = str(i).zfill(num_zeros) + column_name = f"{self.prefix}{suffix_prefix}{formatted_num}" + elif self.suffix: + column_name = f"{self.prefix}{self.suffix}{i}" + else: + column_name = f"{self.prefix}{i}" + mapping[value] = column_name + + return mapping \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/comma_split_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/comma_split_provider.py new file mode 100644 index 00000000..04ecc0a5 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/comma_split_provider.py @@ -0,0 +1,113 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Comma-split transformer for document graph operations. + +This module provides a transformer that splits comma-separated values in specified +fields and creates multiple records from a single input record. This is useful for +handling fields that contain multiple values separated by commas, such as tags, +categories, or IDs that need to be processed as separate entities. +""" + +import logging + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class CommaSplitProvider(TransformerProvider): + """Splits comma-separated values in fields and creates multiple records. + + This transformer takes records with comma-separated values in specified fields + and creates multiple output records, one for each split value. This is useful + for normalizing data where multiple values are stored in a single field. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config) + - fields: List of fields to split on commas + - separator: Character to split on (default: ",") + - strip_whitespace: Whether to strip whitespace from split values (default: True) + + Examples: + >>> # Create a transformer to split project IDs + >>> config = TransformerProviderConfig( + ... name="project_splitter", + ... args={ + ... "fields": ["project_ids"], + ... "separator": ",", + ... "strip_whitespace": True + ... } + ... ) + >>> transformer = CommaSplitProvider(config) + >>> result = transformer.transform([ + ... {"id": 1, "project_ids": "proj-1, proj-2, proj-3", "name": "Resource A"} + ... ]) + >>> len(result) + 3 + >>> result[0]["project_ids"] + 'proj-1' + >>> result[1]["project_ids"] + 'proj-2' + """ + + def __init__(self, config): + """Initialize the comma split transformer with configuration. + + Args: + config: Transformer configuration with name, type, and args + - fields: List of fields to split on commas + - separator: Character to split on (default: ",") + - strip_whitespace: Whether to strip whitespace (default: True) + """ + super().__init__(config) + self.separator = self.args.get('separator', ',') + self.strip_whitespace = self.args.get('strip_whitespace', True) + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by splitting comma-separated values into multiple records. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with split values + + Note: + Each input record may generate multiple output records, one for each + split value in the specified fields. If multiple fields are specified, + the transformation is applied to each field independently. + """ + fields_to_split = self.args.get('fields', []) + + if not fields_to_split: + return records + + transformed_records = [] + + for record in records: + # For each field that needs splitting, create separate records + for field in fields_to_split: + if field in record and record[field]: + field_value = str(record[field]) + split_values = field_value.split(self.separator) + + if self.strip_whitespace: + split_values = [val.strip() for val in split_values if val.strip()] + else: + split_values = [val for val in split_values if val] + + # Create a new record for each split value + for split_value in split_values: + new_record = record.copy() + new_record[field] = split_value + transformed_records.append(new_record) + else: + # If field doesn't exist or is empty, keep original record + transformed_records.append(record.copy()) + + return transformed_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/embedded_json.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/embedded_json.py new file mode 100644 index 00000000..2353bc4f --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/embedded_json.py @@ -0,0 +1,88 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Embedded JSON field transformer for document graph operations. + +This module provides a transformer that parses JSON data from string fields +and flattens the JSON structure into the record with prefixed field names. +This is useful for extracting structured data embedded within JSON strings +and making it available as individual fields for querying and analysis. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +import json +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class EmbeddedJSONFieldTransformer(TransformerProvider): + """Parses JSON fields and flattens them into the record. + + This transformer extracts structured data from JSON strings stored in a specified field + and adds each key-value pair from the JSON object as a separate field in the record. + The new fields are prefixed to avoid name collisions with existing fields. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config) + - field: The field containing the JSON data (default: "json_data") + - prefix: Prefix to add to extracted field names (default: "embedded_") + + Examples: + >>> # Create a transformer to extract JSON data from the "metadata" field + >>> config = TransformerProviderConfig( + ... name="json_extractor", + ... args={"field": "metadata", "prefix": "meta_"} + ... ) + >>> transformer = EmbeddedJSONFieldTransformer(config) + >>> result = transformer.transform([ + ... {"id": 1, "metadata": '{"author": "Jane Doe", "year": 2023}'} + ... ]) + >>> result[0] + {'id': 1, 'metadata': '{"author": "Jane Doe", "year": 2023}', 'meta_author': 'Jane Doe', 'meta_year': 2023} + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by extracting and flattening JSON data. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with flattened JSON fields + + Note: + If the JSON parsing fails, an error field will be added to the record + with the prefix and "error" suffix. + """ + target_field = self.args.get("field", "json_data") + prefix = self.args.get("prefix", "embedded_") + + transformed_records = [] + for record in records: + transformed_record = record.copy() + + try: + field_value = record.get(target_field, "") + if isinstance(field_value, str): + embedded_data = json.loads(field_value) + else: + embedded_data = field_value + + if isinstance(embedded_data, dict): + for key, value in embedded_data.items(): + transformed_record[f"{prefix}{key}"] = value + + except Exception as e: + transformed_record[f"{prefix}error"] = str(e) + + transformed_records.append(transformed_record) + + return transformed_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_expander.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_expander.py new file mode 100644 index 00000000..65fb1cf9 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_expander.py @@ -0,0 +1,334 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON Array Expander transformer for cybersecurity data processing. + +This module provides a robust transformer specifically designed for handling +complex JSON arrays commonly found in cybersecurity vendor data. It expands +JSON arrays into multiple records while handling various edge cases like +empty arrays, duplicate fields, nested structures, and malformed data. + +This transformer is essential for normalizing vendor security data that often +contains nested evidence, alerts, indicators, and other security artifacts +stored as JSON arrays within CSV fields. +""" + +import logging + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +import json +import pandas as pd +from typing import List, Dict, Any, Union +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class JSONArrayExpanderProvider(TransformerProvider): + """Expands JSON arrays in fields into multiple records with advanced conflict resolution. + + This transformer is specifically designed for cybersecurity data where vendors often + store complex nested data structures as JSON arrays. It handles common issues like: + - Empty arrays that should be skipped + - Duplicate field names between original record and JSON data + - Nested objects within arrays + - Malformed JSON data + - Field name conflicts and collision resolution + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config) + - field: The field containing the JSON array data (required) + - skip_empty_arrays: Skip records with empty arrays [] (default: True) + - conflict_resolution: How to handle field name conflicts (default: "prefix") + - "prefix": Add prefix to JSON fields (json_fieldname) + - "suffix": Add suffix to JSON fields (fieldname_json) + - "overwrite": JSON fields overwrite original fields + - "skip": Skip JSON fields that conflict with original + - prefix: Prefix for JSON fields when using prefix resolution (default: "json_") + - suffix: Suffix for JSON fields when using suffix resolution (default: "_json") + - flatten_nested: Flatten nested objects using dot notation (default: True) + - max_depth: Maximum nesting depth to flatten (default: 3) + - preserve_original_field: Keep the original JSON field in output (default: False) + - add_array_index: Add array index as a field (default: True) + - array_index_field: Field name for array index (default: "array_index") + - log_unmapped_fields: Log unmapped fields to file (default: False) + - log_file_path: Path to log file for unmapped fields (default: "unmapped_fields.json") + - parent_node: Parent node name for unmapped field suggestions (required if log_unmapped_fields=True) + + Examples: + >>> # Basic array expansion for security evidence + >>> config = TransformerProviderConfig( + ... name="evidence_expander", + ... type="json_array_expander", + ... args={ + ... "field": "Evidence", + ... "skip_empty_arrays": True, + ... "conflict_resolution": "prefix" + ... } + ... ) + >>> transformer = JSONArrayExpanderProvider(config) + >>> result = transformer.transform([ + ... {"id": 1, "Evidence": '[{"type": "FILE", "path": "/etc/passwd"}]'} + ... ]) + >>> result[0] + {'id': 1, 'Evidence': '[{"type": "FILE", "path": "/etc/passwd"}]', + 'json_type': 'FILE', 'json_path': '/etc/passwd', 'array_index': 0} + """ + + def __init__(self, config): + """Initialize the JSON array expander with configuration. + + Args: + config: Transformer configuration with name, type, and args + """ + super().__init__(config) + self.field = self.args.get('field') + if not self.field: + raise ValueError("JSON Array Expander requires 'field' argument") + + self.skip_empty_arrays = self.args.get('skip_empty_arrays', True) + self.conflict_resolution = self.args.get('conflict_resolution', 'prefix') + self.prefix = self.args.get('prefix', 'json_') + self.suffix = self.args.get('suffix', '_json') + self.flatten_nested = self.args.get('flatten_nested', True) + self.max_depth = self.args.get('max_depth', 3) + self.preserve_original_field = self.args.get('preserve_original_field', False) + self.add_array_index = self.args.get('add_array_index', True) + self.array_index_field = self.args.get('array_index_field', 'array_index') + self.log_unmapped_fields = self.args.get('log_unmapped_fields', False) + self.log_file_path = self.args.get('log_file_path', 'unmapped_fields.json') + self.parent_node = self.args.get('parent_node') + self.auto_add_discovered = self.args.get('auto_add_discovered', False) + self.preserve_original_field = self.args.get('preserve_original_field', False) + self.unmapped_fields = set() # Track unique unmapped fields + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by expanding JSON arrays into multiple records. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with expanded JSON array data + + Note: + Each input record may generate multiple output records, one for each + item in the JSON array. Empty arrays are handled according to the + skip_empty_arrays configuration. + """ + expanded_records = [] + + for record in records: + try: + json_data = self._extract_json_data(record) + + if json_data is None: + # No JSON data found, set field to NaN and keep original record + record_copy = record.copy() + record_copy[self.field] = pd.NA + expanded_records.append(record_copy) + continue + + if isinstance(json_data, list): + if len(json_data) == 0 and self.skip_empty_arrays: + # Skip empty arrays + logger.debug(f"Skipping record with empty array in field '{self.field}'") + continue + + # Expand array into multiple records + for index, item in enumerate(json_data): + new_record = self._create_expanded_record(record, item, index) + expanded_records.append(new_record) + + else: + # Single object, treat as array with one item + new_record = self._create_expanded_record(record, json_data, 0) + expanded_records.append(new_record) + + except Exception as e: + logger.warning(f"Failed to process JSON in field '{self.field}': {e}") + # Set field to NaN on error and keep original record + record_copy = record.copy() + record_copy[self.field] = pd.NA + expanded_records.append(record_copy) + + return expanded_records + + def _extract_json_data(self, record: Dict[str, Any]) -> Union[List, Dict, None]: + """Extract and parse JSON data from the specified field.""" + field_value = record.get(self.field) + + if field_value is None or field_value == '': + return None + + if isinstance(field_value, str): + # Handle empty strings + if field_value.strip() == '': + return None + try: + return json.loads(field_value) + except json.JSONDecodeError as e: + logger.warning(f"Invalid JSON in field '{self.field}': {e}") + return None + elif isinstance(field_value, (list, dict)): + return field_value + else: + logger.warning(f"Unsupported data type in field '{self.field}': {type(field_value)}") + return None + + def _create_expanded_record(self, original_record: Dict[str, Any], + json_item: Any, index: int) -> Dict[str, Any]: + """Create a new record by merging original record with JSON item data.""" + # Start with original record + if self.preserve_original_field: + new_record = original_record.copy() + else: + new_record = {k: v for k, v in original_record.items() if k != self.field} + + # Add array index if requested + if self.add_array_index: + new_record[self.array_index_field] = index + + # Process JSON item + if isinstance(json_item, dict): + flattened_data = self._flatten_json_object(json_item) + resolved_data = self._resolve_field_conflicts(new_record, flattened_data) + + # Auto-add discovered fields or log unmapped fields + if self.auto_add_discovered: + # Add all discovered fields to the record + new_record.update(resolved_data) + # Log what was auto-mapped if requested + if self.log_unmapped_fields and self.parent_node: + self._log_unmapped_fields(flattened_data) + else: + # Only log unmapped fields if requested + if self.log_unmapped_fields and self.parent_node: + self._log_unmapped_fields(flattened_data) + else: + # Handle primitive values in array + if self.auto_add_discovered: + field_name = self._resolve_field_name('value', new_record) + new_record[field_name] = json_item + elif self.log_unmapped_fields: + # Log that a primitive value was found but not mapped + logger.info(f"Found primitive value in array but auto_add_discovered=False: {json_item}") + + return new_record + + def _flatten_json_object(self, json_obj: Dict[str, Any], + parent_key: str = '', depth: int = 0) -> Dict[str, Any]: + """Flatten nested JSON object using dot notation.""" + if not self.flatten_nested or depth >= self.max_depth: + return json_obj + + flattened = {} + + for key, value in json_obj.items(): + new_key = f"{parent_key}.{key}" if parent_key else key + + if isinstance(value, dict) and depth < self.max_depth: + # Recursively flatten nested objects + nested_flattened = self._flatten_json_object(value, new_key, depth + 1) + flattened.update(nested_flattened) + elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], dict): + # Handle arrays of objects by taking first item or concatenating + if len(value) == 1: + nested_flattened = self._flatten_json_object(value[0], new_key, depth + 1) + flattened.update(nested_flattened) + else: + # Multiple items - convert to string representation + flattened[new_key] = str(value) + else: + flattened[new_key] = value + + return flattened + + def _resolve_field_conflicts(self, original_record: Dict[str, Any], + json_data: Dict[str, Any]) -> Dict[str, Any]: + """Resolve field name conflicts between original record and JSON data.""" + resolved_data = {} + + for key, value in json_data.items(): + resolved_key = self._resolve_field_name(key, original_record) + resolved_data[resolved_key] = value + + return resolved_data + + def _resolve_field_name(self, json_field: str, original_record: Dict[str, Any]) -> str: + """Resolve a single field name conflict.""" + # Force prefix when auto_add_discovered is enabled + if self.auto_add_discovered and self.conflict_resolution == 'prefix': + return f"{self.prefix}{json_field}" + + if json_field not in original_record: + return json_field + + # Field conflict detected + if self.conflict_resolution == 'prefix': + return f"{self.prefix}{json_field}" + elif self.conflict_resolution == 'suffix': + return f"{json_field}{self.suffix}" + elif self.conflict_resolution == 'overwrite': + return json_field + elif self.conflict_resolution == 'skip': + # Return None to indicate this field should be skipped + return None + else: + # Default to prefix + return f"{self.prefix}{json_field}" + + def _log_unmapped_fields(self, flattened_data: Dict[str, Any]) -> None: + """Log unmapped fields for transform_and_load.json generation.""" + for field_name in flattened_data.keys(): + if field_name not in self.unmapped_fields: + self.unmapped_fields.add(field_name) + + # Create mapping suggestion + mapping_suggestion = { + "csv_property_name": field_name, + "type": "property", + "arrow_property_name": field_name.lower().replace(' ', '_').replace('-', '_'), + "parents": [self.parent_node] + } + + # Append to log file + try: + import os + # Create directory if it doesn't exist + log_dir = os.path.dirname(self.log_file_path) + if log_dir and not os.path.exists(log_dir): + os.makedirs(log_dir, exist_ok=True) + + existing_data = [] + if os.path.exists(self.log_file_path): + try: + with open(self.log_file_path, 'r') as f: + existing_data = json.load(f) + except json.JSONDecodeError: + # File is corrupted, start fresh + logger.warning(f"Corrupted log file {self.log_file_path}, recreating") + existing_data = [] + + existing_data.append(mapping_suggestion) + + with open(self.log_file_path, 'w') as f: + json.dump(existing_data, f, indent=2) + + if hasattr(self, 'auto_add_discovered') and self.auto_add_discovered: + logger.debug(f"Auto-mapped field '{field_name}' to {self.parent_node} node") + else: + logger.debug(f"Logged unmapped field '{field_name}' to {self.log_file_path}") + + except Exception as e: + logger.error(f"Failed to log unmapped field '{field_name}': {e}") + + def __del__(self): + """Cleanup: Log summary of discovered fields.""" + if self.log_unmapped_fields and self.unmapped_fields: + if hasattr(self, 'auto_add_discovered') and self.auto_add_discovered: + logger.debug(f"Auto-mapped {len(self.unmapped_fields)} discovered fields in '{self.field}' to {self.parent_node} nodes: {sorted(self.unmapped_fields)}") + else: + logger.debug(f"Found {len(self.unmapped_fields)} unmapped fields in '{self.field}': {sorted(self.unmapped_fields)}") \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_flattener.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_flattener.py new file mode 100644 index 00000000..cf4a49b5 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_flattener.py @@ -0,0 +1,99 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON Array Flattener transformer for enriching records with JSON array data. + +This transformer flattens JSON arrays into individual numbered columns without creating +new rows, similar to comma_flattener but for JSON arrays. +""" + +import logging +import json +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +logger = logging.getLogger(__name__) + + +class JSONArrayFlattenerProvider(TransformerProvider): + """Flattens JSON arrays into numbered columns to enrich records. + + This transformer takes JSON arrays and flattens them into individual + columns with a configurable prefix and numbering, similar to comma_flattener. + + Args: + field: The field containing JSON array data (required) + prefix: Prefix for flattened fields (default: field name) + suffix: Suffix pattern for numbering (default: " #00") + preserve_original_field: Keep original field (default: False) + max_items: Maximum number of items to extract (default: 10) + """ + + def __init__(self, config): + super().__init__(config) + self.field = self.args.get('field') or config.csv_property_name + if not self.field: + raise ValueError("JSON Array Flattener requires 'field' argument or csv_property_name") + + if 'prefix' not in self.args: + clean_field = ''.join(c if c.isalnum() else '' for c in self.field) + self.prefix = clean_field + else: + self.prefix = self.args.get('prefix') + + self.suffix = self.args.get('suffix', ' #00') + self.preserve_original_field = self.args.get('preserve_original_field', False) + self.max_items = self.args.get('max_items', 10) + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by flattening JSON arrays into numbered columns.""" + enriched_records = [] + + for record in records: + try: + enriched_record = self._enrich_record(record) + enriched_records.append(enriched_record) + except Exception as e: + logger.warning(f"Failed to process JSON array field '{self.field}': {e}") + enriched_records.append(record.copy()) + + return enriched_records + + def _enrich_record(self, original_record: Dict[str, Any]) -> Dict[str, Any]: + """Enrich record with flattened JSON array data.""" + if self.preserve_original_field: + enriched_record = original_record.copy() + else: + enriched_record = {k: v for k, v in original_record.items() if k != self.field} + + field_value = original_record.get(self.field) + + if field_value is None or field_value == '': + return enriched_record + + # Parse JSON if string + if isinstance(field_value, str): + try: + json_data = json.loads(field_value) + except json.JSONDecodeError: + return enriched_record + else: + json_data = field_value + + # Handle arrays + if isinstance(json_data, list): + for i, item in enumerate(json_data[:self.max_items], 1): + if self.suffix and '#' in self.suffix: + hash_index = self.suffix.index('#') + suffix_prefix = self.suffix[:hash_index] + number_format = self.suffix[hash_index+1:] + num_zeros = len(number_format) + formatted_num = str(i).zfill(num_zeros) + column_name = f"{self.prefix}{suffix_prefix}{formatted_num}" + else: + column_name = f"{self.prefix}{i}" + + # Convert item to string representation + enriched_record[column_name] = json.dumps(item) if isinstance(item, (dict, list)) else str(item) + + return enriched_record \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_paired_flattener.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_paired_flattener.py new file mode 100644 index 00000000..4e0ce504 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_array_paired_flattener.py @@ -0,0 +1,73 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON Array Paired Flattener for flattening two related JSON arrays.""" + +import logging +import json +from typing import List, Dict, Any, Union +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +logger = logging.getLogger(__name__) + + +class JSONArrayPairedFlattenerProvider(TransformerProvider): + """Flattens two related JSON arrays into paired columns. + + Handles JSON arrays like: ["id1", "id2"] and ["name1", "name2"] + Creates columns like: project_01_id, project_01_name, project_02_id, project_02_name + """ + + def __init__(self, config): + super().__init__(config) + self.id_field = self.args.get('id_field') + self.name_field = self.args.get('name_field') + self.prefix = self.args.get('prefix', 'project') + self.max_items = self.args.get('max_items', 10) + + if not self.id_field or not self.name_field: + raise ValueError("JSONArrayPairedFlattener requires both 'id_field' and 'name_field'") + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + enriched_records = [] + + for record in records: + enriched_record = record.copy() + + ids = self._parse_json_array(record.get(self.id_field)) + names = self._parse_json_array(record.get(self.name_field)) + + # Pair up IDs and names (pad shorter list with empty strings) + max_len = min(max(len(ids), len(names)), self.max_items) + ids.extend([''] * (max_len - len(ids))) + names.extend([''] * (max_len - len(names))) + + for i in range(max_len): + num = str(i + 1).zfill(2) + enriched_record[f"{self.prefix}_{num}_id"] = ids[i] + enriched_record[f"{self.prefix}_{num}_name"] = names[i] + + enriched_records.append(enriched_record) + + return enriched_records + + def _parse_json_array(self, value: Union[str, List, None]) -> List[str]: + """Parse JSON array from string or return list directly.""" + if not value: + return [] + + if isinstance(value, list): + return [str(item) for item in value] + + if isinstance(value, str): + try: + parsed = json.loads(value) + if isinstance(parsed, list): + return [str(item) for item in parsed] + else: + return [str(parsed)] # Single value + except json.JSONDecodeError: + # Fallback to comma-separated for compatibility + return [item.strip() for item in value.split(',') if item.strip()] + + return [str(value)] # Single value fallback \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_flattener.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_flattener.py new file mode 100644 index 00000000..48adb92a --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_flattener.py @@ -0,0 +1,175 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON Flattener transformer for enriching records with JSON field data. + +This transformer flattens JSON objects into individual columns without creating +new rows, allowing you to enrich existing records with structured data from +JSON blobs and then map those enriched fields to different graph nodes. +""" + +import logging +import json +import pandas as pd +from typing import List, Dict, Any, Union +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +logger = logging.getLogger(__name__) + + +class JSONFlattenerProvider(TransformerProvider): + """Flattens JSON objects into columns to enrich records. + + This transformer takes JSON objects/arrays and flattens them into individual + columns with a configurable prefix. It preserves the original row count while + adding new columns for each JSON field. + + Args: + field: The field containing JSON data (required) + prefix: Prefix for flattened fields (default: field name without spaces) + flatten_nested: Flatten nested objects with dot notation (default: True) + max_depth: Maximum nesting depth (default: 3) + preserve_original_field: Keep original JSON field (default: False) + handle_arrays: How to handle arrays - "first", "join", "string" (default: "first") + """ + + def __init__(self, config): + super().__init__(config) + self.field = self.args.get('field') or getattr(config, 'csv_property_name', None) + if not self.field: + raise ValueError("JSON Flattener requires 'field' argument or csv_property_name") + + # Use field name as default prefix (no underscore separator) + if 'prefix' not in self.args: + # Clean field name: remove spaces and special chars + clean_field = ''.join(c if c.isalnum() else '' for c in self.field) + self.prefix = clean_field + logger.info(f"JSON Flattener using default prefix: '{self.prefix}' (from field '{self.field}')") + else: + self.prefix = self.args.get('prefix') + logger.info(f"JSON Flattener using explicit prefix: '{self.prefix}'") + self.flatten_nested = self.args.get('flatten_nested', True) + self.max_depth = self.args.get('max_depth', 3) + self.preserve_original_field = self.args.get('preserve_original_field', False) + self.handle_arrays = self.args.get('handle_arrays', 'first') # first, join, string + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by flattening JSON into columns.""" + enriched_records = [] + + for record in records: + try: + json_data = self._extract_json_data(record) + enriched_record = self._enrich_record(record, json_data) + enriched_records.append(enriched_record) + except Exception as e: + logger.warning(f"Failed to process JSON in field '{self.field}': {e}") + enriched_records.append(record.copy()) + + return enriched_records + + def _extract_json_data(self, record: Dict[str, Any]) -> Union[Dict, List, None]: + """Extract and parse JSON data from the specified field.""" + field_value = record.get(self.field) + + if field_value is None or field_value == '': + return None + + if isinstance(field_value, str): + if field_value.strip() == '': + return None + try: + return json.loads(field_value) + except json.JSONDecodeError as e: + logger.warning(f"Invalid JSON in field '{self.field}': {e}") + return None + elif isinstance(field_value, (list, dict)): + return field_value + else: + logger.warning(f"Unsupported data type in field '{self.field}': {type(field_value)}") + return None + + def _enrich_record(self, original_record: Dict[str, Any], json_data: Any) -> Dict[str, Any]: + """Enrich record with flattened JSON data.""" + # Start with original record + if self.preserve_original_field: + enriched_record = original_record.copy() + else: + enriched_record = {k: v for k, v in original_record.items() if k != self.field} + + if json_data is None: + return enriched_record + + # Handle arrays by taking first item or converting to string + if isinstance(json_data, list): + if len(json_data) == 0: + return enriched_record + elif self.handle_arrays == 'first': + json_data = json_data[0] # Take first item + elif self.handle_arrays == 'string': + # Convert entire array to string + flattened_data = {f"{self.prefix}array": str(json_data)} + enriched_record.update(flattened_data) + return enriched_record + elif self.handle_arrays == 'numbered': + # Create numbered columns like comma_flattener + for i, item in enumerate(json_data, 1): + if i > 10: # Limit to 10 items + break + formatted_num = str(i).zfill(2) + column_name = f"{self.prefix} {formatted_num}" + if isinstance(item, dict): + enriched_record[column_name] = json.dumps(item) + else: + enriched_record[column_name] = str(item) + return enriched_record + if self.handle_arrays == 'join' and all(isinstance(x, str) for x in json_data): + # Join string arrays + flattened_data = {f"{self.prefix}array": ', '.join(json_data)} + enriched_record.update(flattened_data) + return enriched_record + else: + json_data = json_data[0] # Fallback to first + + # Flatten the JSON object + if isinstance(json_data, dict): + flattened_data = self._flatten_json_object(json_data) + # Add prefix to all keys + prefixed_data = {f"{self.prefix}{key}": value for key, value in flattened_data.items()} + enriched_record.update(prefixed_data) + else: + # Handle primitive values + enriched_record[f"{self.prefix}value"] = json_data + + return enriched_record + + def _flatten_json_object(self, json_obj: Dict[str, Any], parent_key: str = '', depth: int = 0) -> Dict[str, Any]: + """Flatten nested JSON object using dot notation.""" + if not self.flatten_nested or depth >= self.max_depth: + return json_obj + + flattened = {} + + for key, value in json_obj.items(): + new_key = f"{parent_key}.{key}" if parent_key else key + + if isinstance(value, dict) and depth < self.max_depth: + # Recursively flatten nested objects + nested_flattened = self._flatten_json_object(value, new_key, depth + 1) + flattened.update(nested_flattened) + elif isinstance(value, list): + # Handle arrays based on configuration + if self.handle_arrays == 'first' and len(value) > 0: + if isinstance(value[0], dict): + nested_flattened = self._flatten_json_object(value[0], new_key, depth + 1) + flattened.update(nested_flattened) + else: + flattened[new_key] = value[0] + elif self.handle_arrays == 'join' and all(isinstance(x, str) for x in value): + flattened[new_key] = ', '.join(value) + else: + flattened[new_key] = str(value) + else: + flattened[new_key] = value + + return flattened \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_value_flattener.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_value_flattener.py new file mode 100644 index 00000000..0f14873e --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/json_value_flattener.py @@ -0,0 +1,54 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON Value Flattener for creating JSON objects from paired fields.""" + +import json +import logging +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +logger = logging.getLogger(__name__) + + +class JSONValueFlattenerProvider(TransformerProvider): + """Creates JSON objects from paired comma-separated fields. + + Creates columns like: project_01: {"id": "123", "name": "Project A"} + """ + + def __init__(self, config): + super().__init__(config) + self.id_field = self.args.get('id_field') + self.name_field = self.args.get('name_field') + self.prefix = self.args.get('prefix', 'project') + self.suffix = self.args.get('suffix', '#00') + self.max_items = self.args.get('max_items', 10) + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + enriched_records = [] + + for record in records: + enriched_record = record.copy() + + ids = self._split_field(record.get(self.id_field, '')) + names = self._split_field(record.get(self.name_field, '')) + + max_len = min(max(len(ids), len(names)), self.max_items) + + for i in range(max_len): + num = str(i + 1).zfill(2) + project_data = { + "id": ids[i] if i < len(ids) else "", + "name": names[i] if i < len(names) else "" + } + enriched_record[f"{self.prefix}_{num}"] = json.dumps(project_data) + + enriched_records.append(enriched_record) + + return enriched_records + + def _split_field(self, value: str) -> List[str]: + if not value: + return [] + return [item.strip() for item in str(value).split(',') if item.strip()] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/normalize_timestamp.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/normalize_timestamp.py new file mode 100644 index 00000000..f430f71f --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/normalize_timestamp.py @@ -0,0 +1,87 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Timestamp normalization transformer for document graph operations. + +This module provides a transformer that converts timestamp fields from various formats +into a standardized ISO 8601 format. This normalization enables consistent timestamp +handling, comparison, and querying across different data sources that may use +different timestamp formats. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from datetime import datetime +from dateutil import parser +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class TimestampNormalizer(TransformerProvider): + """Normalizes timestamp fields to ISO 8601 format. + + This transformer converts timestamp fields from various formats into a standardized + ISO 8601 format (YYYY-MM-DDTHH:MM:SS.sssZ). It uses the dateutil parser to handle + a wide variety of input formats, making it flexible for different data sources. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config) + - field: The field containing the timestamp to normalize (default: "timestamp") + - output_field: The field to store the normalized timestamp + (default: "{field}_normalized") + + Examples: + >>> # Create a transformer to normalize dates in the "created_at" field + >>> config = TransformerProviderConfig( + ... name="date_normalizer", + ... args={"field": "created_at", "output_field": "iso_date"} + ... ) + >>> transformer = TimestampNormalizer(config) + >>> result = transformer.transform([ + ... {"id": 1, "created_at": "Jan 1, 2023 14:30:45"}, + ... {"id": 2, "created_at": "2023/02/15 09:15 AM"} + ... ]) + >>> result[0]["iso_date"] + '2023-01-01T14:30:45' + >>> result[1]["iso_date"] + '2023-02-15T09:15:00' + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by normalizing timestamp fields to ISO 8601 format. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with normalized timestamp fields + + Note: + If timestamp parsing fails, an error field will be added to the record + with the output field name and "_error" suffix. + """ + field = self.args.get("field", "timestamp") + output_field = self.args.get("output_field", f"{field}_normalized") + + normalized_records = [] + for record in records: + normalized_record = record.copy() + raw_value = record.get(field) + + if raw_value: + try: + dt: datetime = parser.parse(str(raw_value)) + normalized_record[output_field] = dt.isoformat() + except Exception as e: + normalized_record[f"{output_field}_error"] = str(e) + + normalized_records.append(normalized_record) + + return normalized_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/paired_flattener.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/paired_flattener.py new file mode 100644 index 00000000..f14b3634 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/paired_flattener.py @@ -0,0 +1,57 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Paired Flattener transformer for flattening two related comma-separated fields.""" + +import logging +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +logger = logging.getLogger(__name__) + + +class PairedFlattenerProvider(TransformerProvider): + """Flattens two related comma-separated fields into paired columns. + + Creates columns like: project_01_id, project_01_name, project_02_id, project_02_name + """ + + def __init__(self, config): + super().__init__(config) + self.id_field = self.args.get('id_field') + self.name_field = self.args.get('name_field') + self.prefix = self.args.get('prefix', 'project') + self.suffix = self.args.get('suffix', '#00') + self.separator = self.args.get('separator', ',') + self.max_items = self.args.get('max_items', 10) + + if not self.id_field or not self.name_field: + raise ValueError("PairedFlattener requires both 'id_field' and 'name_field'") + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + enriched_records = [] + + for record in records: + enriched_record = record.copy() + + ids = self._split_field(record.get(self.id_field, '')) + names = self._split_field(record.get(self.name_field, '')) + + # Pair up IDs and names (pad shorter list with empty strings) + max_len = min(max(len(ids), len(names)), self.max_items) + ids.extend([''] * (max_len - len(ids))) + names.extend([''] * (max_len - len(names))) + + for i in range(max_len): + num = str(i + 1).zfill(2) if '#00' in self.suffix else str(i + 1) + enriched_record[f"{self.prefix}_{num}_id"] = ids[i] + enriched_record[f"{self.prefix}_{num}_name"] = names[i] + + enriched_records.append(enriched_record) + + return enriched_records + + def _split_field(self, value: str) -> List[str]: + if not value: + return [] + return [item.strip() for item in str(value).split(self.separator) if item.strip()] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/regex_cleaner_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/regex_cleaner_provider.py new file mode 100644 index 00000000..322cf709 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/regex_cleaner_provider.py @@ -0,0 +1,99 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regex-based text cleaning transformer for document graph operations. + +This module provides a transformer that applies regular expression patterns to clean +and standardize text fields in records. It can be used to remove unwanted characters, +normalize formats, extract specific patterns, or perform other text cleaning operations +using the power and flexibility of regular expressions. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any, Pattern +import re +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class RegexCleanerProvider(TransformerProvider): + """Applies regex patterns to clean text fields in records. + + This transformer applies a list of regular expression patterns to specified fields + in each record, replacing matches with a configured replacement string. This is useful + for standardizing text formats, removing unwanted characters or patterns, and + performing other text cleaning operations. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config) + - patterns: List of regex patterns to apply (default: []) + - replacement: String to replace matches with (default: "") + - fields: List of fields to clean + replacement: The replacement string to use + _compiled_patterns: List of compiled regex pattern objects + + Examples: + >>> # Create a transformer to remove HTML tags from text fields + >>> config = TransformerProviderConfig( + ... name="html_cleaner", + ... args={ + ... "patterns": ["<[^>]*>"], + ... "replacement": "", + ... "fields": ["description", "content"] + ... } + ... ) + >>> transformer = RegexCleanerProvider(config) + >>> result = transformer.transform([ + ... {"id": 1, "description": "

This is a sample text.

"} + ... ]) + >>> result[0]["description"] + 'This is a sample text.' + """ + + def __init__(self, config): + """Initialize the regex cleaner with configuration. + + Args: + config: Transformer configuration with name, type, and args + - patterns: List of regex patterns to apply + - replacement: String to replace matches with + - fields: List of fields to clean + """ + super().__init__(config) + patterns = self.args.get('patterns', []) + self.replacement = self.args.get('replacement', '') + self._compiled_patterns: List[Pattern] = [re.compile(p) for p in patterns] + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by applying regex patterns to specified fields. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with cleaned text fields + + Note: + The transformation modifies the specified fields in-place in each record. + Only string fields are processed; non-string fields are left unchanged. + """ + fields_to_clean = self.args.get('fields', []) + + cleaned_records = [] + for record in records: + cleaned_record = record.copy() + for field in fields_to_clean: + if field in cleaned_record and isinstance(cleaned_record[field], str): + result = cleaned_record[field] + for pattern in self._compiled_patterns: + result = pattern.sub(self.replacement, result) + cleaned_record[field] = result + cleaned_records.append(cleaned_record) + return cleaned_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/standardize_enum.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/standardize_enum.py new file mode 100644 index 00000000..066ed5bc --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/standardize_enum.py @@ -0,0 +1,105 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Enumeration standardization transformer for document graph operations. + +This module provides a transformer that maps field values to standardized enumeration values +using a mapping dictionary. This is useful for normalizing categorical data that may have +different representations across data sources, ensuring consistent values for analysis, +querying, and visualization. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class EnumStandardizer(TransformerProvider): + """Maps field values to standard enumeration using a mapping dictionary. + + This transformer standardizes categorical data by mapping input values to a set of + predefined standard values using a mapping dictionary. This is useful for normalizing + data from different sources that may use different terminology for the same concepts. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config) + - field: The field containing values to standardize (default: "category") + - mapping: Dictionary mapping input values to standardized values (default: {}) + - case_insensitive: Whether to ignore case when looking up mappings (default: True) + - output_field: The field to store the standardized value (default: same as input field) + + Examples: + >>> # Create a transformer to standardize status values + >>> config = TransformerProviderConfig( + ... name="status_standardizer", + ... args={ + ... "field": "status", + ... "mapping": { + ... "active": "ACTIVE", + ... "enabled": "ACTIVE", + ... "inactive": "INACTIVE", + ... "disabled": "INACTIVE", + ... "pending": "PENDING" + ... }, + ... "case_insensitive": True + ... } + ... ) + >>> transformer = EnumStandardizer(config) + >>> result = transformer.transform([ + ... {"id": 1, "status": "Active"}, + ... {"id": 2, "status": "DISABLED"}, + ... {"id": 3, "status": "unknown"} + ... ]) + >>> result[0]["status"] + 'ACTIVE' + >>> result[1]["status"] + 'INACTIVE' + >>> result[2]["status_unmapped"] + 'unknown' + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by standardizing enumeration values. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with standardized enumeration values + + Note: + If a value cannot be mapped (not found in the mapping dictionary), + it will be stored in a separate field with the "_unmapped" suffix. + When case_insensitive is True, all keys in the mapping dictionary + are compared in lowercase. + """ + field = self.args.get("field", "category") + mapping: Dict[str, str] = self.args.get("mapping", {}) + case_insensitive: bool = self.args.get("case_insensitive", True) + output_field = self.args.get("output_field", field) + + standardized_records = [] + for record in records: + standardized_record = record.copy() + original = record.get(field) + + if original: + key = str(original).lower() if case_insensitive else str(original) + mapped_value = mapping.get(key) + + if mapped_value: + standardized_record[output_field] = mapped_value + else: + standardized_record[f"{output_field}_unmapped"] = original + + standardized_records.append(standardized_record) + + return standardized_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/uuid_generator.py b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/uuid_generator.py new file mode 100644 index 00000000..0ac5c417 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/field_transformers/uuid_generator.py @@ -0,0 +1,23 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""UUID generator — generates UUID values for specified fields.""" + +import uuid +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider +from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + + +class UuidGeneratorTransformer(TransformerProvider): + """Generates UUID values for specified fields""" + + def __init__(self, config: TransformerProviderConfig): + super().__init__(config) + self.target_field = config.args.get('target_field', 'uuid') + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + self._log_transform_start(len(records)) + for record in records: + record[self.target_field] = str(uuid.uuid4()) + self._log_transform_end(len(records), len(records)) + return records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/__init__.py new file mode 100644 index 00000000..7527b0bf --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/__init__.py @@ -0,0 +1,23 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Filter transformers module for document graph operations. + +This package provides transformer implementations that filter or modify data +during document graph processing. It includes transformers for: + +- Column pruning: Removes specified columns (fields) from records +- Row filtering: Filters records based on field conditions with various operators + such as equality, inequality, comparison, and null checks + +These transformers help in data preparation and cleaning by allowing selective +processing of records and fields based on configurable criteria. +""" + +from graphrag_toolkit.document_graph.transform.filter_transformers.column_pruner import ColumnPrunerProvider +from graphrag_toolkit.document_graph.transform.filter_transformers.row_filter import RowFilterProvider + +__all__ = [ + "ColumnPrunerProvider", + "RowFilterProvider", +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/column_pruner.py b/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/column_pruner.py new file mode 100644 index 00000000..48ca981f --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/column_pruner.py @@ -0,0 +1,55 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Column Pruner module for document graph operations. + +This module provides the ColumnPrunerProvider transformer, which removes specified +columns (fields) from records during document graph processing. This is useful for +data preparation and cleaning by allowing selective processing of fields based on +configurable criteria. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +class ColumnPrunerProvider(TransformerProvider): + """Transformer provider that removes specified columns from records. + + This transformer allows for selective field processing by removing specified + columns (fields) from each record. It's useful for data preparation and cleaning + when certain fields are not needed for downstream processing. + + Args: + The provider expects the following arguments in the `args` dictionary: + remove (List[str]): A list of column names to remove from each record. + If not provided, no columns will be removed. + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Remove specified columns from each record. + + This method processes a list of records (dictionaries) and removes the + specified columns from each record. The columns to remove are specified + in the `remove` argument of the provider. + + Args: + records (List[Dict[str, Any]]): A list of records to process, where each + record is a dictionary with string keys. + + Returns: + List[Dict[str, Any]]: A new list of records with the specified columns removed. + The original records are not modified. + """ + columns_to_remove: List[str] = self.args.get("remove", []) + + return [ + {k: v for k, v in record.items() if k not in columns_to_remove} + for record in records + ] diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/row_filter.py b/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/row_filter.py new file mode 100644 index 00000000..59fd6b14 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/filter_transformers/row_filter.py @@ -0,0 +1,108 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Row Filter module for document graph operations. + +This module provides the RowFilterProvider transformer, which filters records based on +field conditions during document graph processing. It supports various comparison +operators such as equality, inequality, comparison, and null checks. + +This transformer is useful for data preparation and cleaning by allowing selective +processing of records based on configurable criteria. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class RowFilterProvider(TransformerProvider): + """Transformer provider that filters records based on field conditions. + + This transformer allows for selective record processing by filtering records + based on field conditions. It supports various comparison operators such as + equality, inequality, comparison, and null checks. Records are kept only if + they satisfy all the specified conditions. + + Args: + The provider expects the following arguments in the `args` dictionary: + conditions (List[Dict[str, Any]]): A list of condition dictionaries. + Each condition dictionary must have the following keys: + - field (str): The name of the field to check. + - operator (str): The comparison operator to use. Supported operators are: + - "eq": Equal to + - "ne": Not equal to + - "lt": Less than + - "gt": Greater than + - "null": Field is null + - "not_null": Field is not null + - value (Any): The value to compare against. Required for "eq", "ne", + "lt", and "gt" operators. Not used for "null" and "not_null" operators. + + Raises: + ValueError: If an unsupported operator is specified. + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Filter records based on field conditions. + + This method processes a list of records (dictionaries) and returns only the + records that satisfy all the specified conditions. The conditions are specified + in the `conditions` argument of the provider. + + If no conditions are specified, all records are returned unchanged. + + Args: + records (List[Dict[str, Any]]): A list of records to process, where each + record is a dictionary with string keys. + + Returns: + List[Dict[str, Any]]: A new list containing only the records that satisfy + all the specified conditions. + + Raises: + ValueError: If an unsupported operator is specified in a condition. + """ + conditions: List[Dict[str, Any]] = self.args.get("conditions", []) + if not conditions: + return records + + filtered_records = [] + for record in records: + keep_record = True + + for cond in conditions: + field = cond["field"] + op = cond["operator"] + value = cond.get("value") + field_value = record.get(field) + + if op == "eq" and field_value != value: + keep_record = False + elif op == "ne" and field_value == value: + keep_record = False + elif op == "lt" and not (field_value is not None and field_value < value): + keep_record = False + elif op == "gt" and not (field_value is not None and field_value > value): + keep_record = False + elif op == "null" and field_value is not None: + keep_record = False + elif op == "not_null" and field_value is None: + keep_record = False + else: + if op not in ["eq", "ne", "lt", "gt", "null", "not_null"]: + raise ValueError(f"Unsupported operator: {op}") + + if not keep_record: + break + + if keep_record: + filtered_records.append(record) + + return filtered_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/__init__.py new file mode 100644 index 00000000..2da63807 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/__init__.py @@ -0,0 +1,23 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Graph transformers module for document graph operations. + +This package provides transformer implementations that create and manipulate graph +structures during document graph processing. It includes transformers for: + +- Edge inference: Creates edges between records based on matching field values +- Row to node conversion: Transforms tabular records into graph nodes with appropriate metadata + +These transformers help in building graph representations from structured data by +converting records to nodes and establishing relationships between them based on +configurable criteria. +""" + +from graphrag_toolkit.document_graph.transform.graph_transformers.infer_edges import EdgeInferencer +from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer + +__all__ = [ + "EdgeInferencer", + "RowToNodeTransformer", +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/infer_edges.py b/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/infer_edges.py new file mode 100644 index 00000000..f2342bb1 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/infer_edges.py @@ -0,0 +1,79 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Infer Edges module for document graph operations. + +This module provides the EdgeInferencer transformer, which creates edges between +records based on matching field values. This is useful for building graph +representations from structured data by establishing relationships between +records that share common attribute values. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class EdgeInferencer(TransformerProvider): + """Transformer provider that infers edges between records based on matching fields. + + This transformer creates edge records between records that share the same value + in a specified field. It's useful for building graph representations from + structured data by establishing relationships between related records. + + Args: + The provider expects the following arguments in the `args` dictionary: + source_field (str): The field name to use for grouping records. + Default is "project_id". + edge_type (str): The relationship type to assign to the created edges. + Default is "RELATED_TO". + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Create edge records between records that share the same field value. + + This method processes a list of records and creates edge records between + consecutive records that share the same value in the specified source field. + The edges are created with the specified relationship type. + + Args: + records (List[Dict[str, Any]]): A list of records to process, where each + record is a dictionary with string keys. + + Returns: + List[Dict[str, Any]]: A new list containing the original records plus + the newly created edge records. + """ + source_field = self.args.get("source_field", "project_id") + edge_type = self.args.get("edge_type", "RELATED_TO") + + # Group records by source_field + groups = {} + for record in records: + key = record.get(source_field) + if key is not None: + groups.setdefault(key, []).append(record) + + # Create edge records within each group + edge_records = [] + for group_records in groups.values(): + for i in range(len(group_records) - 1): + source = group_records[i] + target = group_records[i + 1] + + edge_record = { + "edge_type": "edge", + "source_id": source.get("id", f"record_{i}"), + "target_id": target.get("id", f"record_{i+1}"), + "relationship": edge_type, + "group_key": key + } + edge_records.append(edge_record) + + return records + edge_records # Return original records plus inferred edges \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/row_to_node.py b/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/row_to_node.py new file mode 100644 index 00000000..90393510 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/graph_transformers/row_to_node.py @@ -0,0 +1,73 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Row To Node module for document graph operations. + +This module provides the RowToNodeTransformer, which converts tabular records +into graph nodes with appropriate metadata. This is useful for building graph +representations from structured data by transforming records into nodes that +can be used in a graph database or other graph-based applications. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class RowToNodeTransformer(TransformerProvider): + """Transformer provider that converts records to node format with graph metadata. + + This transformer transforms tabular records into graph nodes by adding + necessary graph metadata such as node type and graph element identifier. + It also ensures each node has an ID and creates a content field from all + other fields if one doesn't exist. + + Args: + The provider expects the following arguments in the `args` dictionary: + type (str): The node type to assign to the created nodes. + Default is "Row". + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Convert records to node format with graph metadata. + + This method processes a list of records and converts each one to a node + format by adding graph metadata such as node_type and graph_element. + It ensures each node has an ID and creates a content field from all + other fields if one doesn't exist. + + Args: + records (List[Dict[str, Any]]): A list of records to process, where each + record is a dictionary with string keys. + + Returns: + List[Dict[str, Any]]: A new list of records converted to node format. + """ + node_type = self.args.get("type", "Row") + + node_records = [] + for idx, record in enumerate(records): + node_record = record.copy() + + # Add graph metadata + node_record["node_type"] = node_type + node_record["graph_element"] = "node" + + # Ensure ID exists + if "id" not in node_record: + node_record["id"] = f"row_{idx}" + + # Create content from all fields if not exists + if "content" not in node_record: + content_parts = [f"{k}: {v}" for k, v in record.items() if k != "id"] + node_record["content"] = ", ".join(content_parts) + + node_records.append(node_record) + + return node_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/__init__.py new file mode 100644 index 00000000..38aed75e --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/__init__.py @@ -0,0 +1,33 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalizer transformers for standardizing data formats. + +This module provides transformers that normalize and standardize data +formats and values, such as: +- Case normalization: Standardize text case (upper, lower, title) +- Enum normalization: Standardize enumerated values +- Null normalization: Handle null and empty values consistently +- Spelling normalization: Correct common spelling errors +- Whitespace normalization: Clean up whitespace and formatting +- Timestamp normalization: Standardize date/time formats + +These transformers help ensure data consistency and quality during +document processing pipelines. +""" + +from .normalize_case_provider import NormalizeCaseProvider +from .normalize_enum_provider import NormalizeEnumProvider +from .normalize_nulls_provider import NormalizeNullsProvider +from .normalize_spelling_provider import NormalizeSpellingProvider +from .normalize_whitespace_provider import NormalizeWhitespaceProvider +from .normalize_timestamp_provider import NormalizeTimestampProvider + +__all__ = [ + 'NormalizeCaseProvider', + 'NormalizeEnumProvider', + 'NormalizeNullsProvider', + 'NormalizeSpellingProvider', + 'NormalizeWhitespaceProvider', + 'NormalizeTimestampProvider' +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_case.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_case.py new file mode 100644 index 00000000..830a3e32 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_case.py @@ -0,0 +1,42 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Case module for document graph operations. + +This module provides functionality for normalizing the case of text strings +in document processing pipelines. It supports: + +- Converting text to lowercase (default) +- Converting text to uppercase +- Other case transformations supported by Python string methods + +Case normalization is useful for standardizing text data, improving search +and matching operations, and preparing text for further processing. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +def normalize(value: str, mode="lower") -> str: + """Normalize the case of a string value. + + Applies case transformation to the input string using the specified mode. + The mode parameter corresponds to a Python string method name (e.g., 'lower', + 'upper', 'capitalize', 'title'). + + Args: + value (str): The string to normalize + mode (str, optional): The case transformation to apply. Defaults to "lower". + Valid values include 'lower', 'upper', 'capitalize', 'title', etc. + + Returns: + str: The normalized string with the specified case transformation applied + + Raises: + AttributeError: If the specified mode is not a valid string method + """ + return getattr(value, mode)() diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_case_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_case_provider.py new file mode 100644 index 00000000..2c3f12d4 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_case_provider.py @@ -0,0 +1,72 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Case Provider module for document graph operations. + +This module provides a transformer provider implementation for normalizing the case +of text strings in document records. It supports: + +- Converting text to lowercase (default) +- Converting text to uppercase +- Other case transformations supported by Python string methods + +The provider applies case normalization to all string fields in each record, +preserving non-string fields unchanged. This is useful for standardizing text data +across records, improving search and matching operations, and preparing text for +further processing. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class NormalizeCaseProvider(TransformerProvider): + """Normalizes text case in record fields. + + This transformer provider applies case normalization to all string fields + in each record, preserving non-string fields unchanged. It supports various + case transformations through the 'mode' configuration parameter. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config), including: + - mode: The case transformation to apply (default: 'lower') + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Apply case normalization to string fields in records. + + Processes each record and applies the specified case transformation + (lowercase by default) to all string fields, leaving non-string fields + unchanged. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with normalized case in string fields + """ + self._log_transform_start(len(records)) + + mode = self.args.get('mode', 'lower') + + transformed_records = [] + for record in records: + transformed_record = {} + for key, value in record.items(): + if isinstance(value, str): + transformed_record[key] = getattr(value, mode)() + else: + transformed_record[key] = value + transformed_records.append(transformed_record) + + self._log_transform_end(len(records), len(transformed_records)) + return transformed_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_enum.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_enum.py new file mode 100644 index 00000000..cc0ab37b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_enum.py @@ -0,0 +1,62 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Enum module for document graph operations. + +This module provides functionality for normalizing enumeration values in text data +using a mapping dictionary. It supports: + +- Case-insensitive mapping of input values to standardized forms +- Handling of variations in spelling, abbreviations, or formatting +- Fallback to original values when no mapping is found + +Enum normalization is useful for standardizing categorical data, ensuring consistent +terminology across datasets, and improving data quality for analysis and processing. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import Optional, Dict + +class EnumNormalizer: + """Normalizes enumeration values using a mapping dictionary. + + This class provides functionality to standardize categorical or enumeration values + by mapping input strings to standardized forms using a dictionary. It handles + case-insensitive matching and preserves the original value when no mapping is found. + + Attributes: + enum_map (Dict[str, str]): Dictionary mapping input values (lowercase) to + their standardized forms + """ + + def __init__(self, enum_map: Dict[str, str]): + """Initialize the enum normalizer with a mapping dictionary. + + Args: + enum_map (Dict[str, str]): Dictionary mapping input values to their + standardized forms. Keys will be converted to lowercase for + case-insensitive matching. + """ + self.enum_map = {k.lower(): v for k, v in enum_map.items()} + + def normalize(self, value: str) -> Optional[str]: + """Normalize an enumeration value using the mapping dictionary. + + Looks up the input value (after stripping whitespace and converting to lowercase) + in the mapping dictionary and returns the standardized form if found. + If not found, returns the original value unchanged. + + Args: + value (str): The string value to normalize + + Returns: + Optional[str]: The normalized value if found in the mapping dictionary, + otherwise the original value + """ + return self.enum_map.get(value.strip().lower(), value) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_enum_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_enum_provider.py new file mode 100644 index 00000000..075a56cc --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_enum_provider.py @@ -0,0 +1,77 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Enum Provider module for document graph operations. + +This module provides a transformer provider implementation for normalizing enumeration +values in document records using a mapping dictionary. It supports: + +- Standardizing categorical field values using a configurable mapping +- Case-insensitive matching (optional, enabled by default) +- Handling of variations in spelling, abbreviations, or formatting +- Preserving original values when no mapping is found + +Enum normalization is useful for standardizing categorical data, ensuring consistent +terminology across datasets, and improving data quality for analysis and processing. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class NormalizeEnumProvider(TransformerProvider): + """Normalizes enum values using a mapping dictionary. + + This transformer provider standardizes categorical field values in document records + using a configurable mapping dictionary. It supports case-insensitive matching + and preserves original values when no mapping is found. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config), including: + - field: The field to normalize (default: "category") + - enum_map: Dictionary mapping input values to standardized forms + - case_insensitive: Whether to use case-insensitive matching (default: True) + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by normalizing enum values in the specified field. + + Processes each record and applies enum normalization to the specified field + using the provided mapping dictionary. If case_insensitive is True (default), + the matching is done after converting keys and values to lowercase. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with normalized enum values + """ + field = self.args.get("field", "category") + enum_map = self.args.get("enum_map", {}) + case_insensitive = self.args.get("case_insensitive", True) + + # Normalize enum_map keys to lowercase if case_insensitive + if case_insensitive: + enum_map = {k.lower(): v for k, v in enum_map.items()} + + normalized_records = [] + for record in records: + normalized_record = record.copy() + value = record.get(field) + + if value: + lookup_key = str(value).strip().lower() if case_insensitive else str(value).strip() + normalized_record[field] = enum_map.get(lookup_key, value) + + normalized_records.append(normalized_record) + + return normalized_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_nulls.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_nulls.py new file mode 100644 index 00000000..f651d582 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_nulls.py @@ -0,0 +1,55 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Nulls module for document graph operations. + +This module provides functionality for normalizing null-like values in text data +to actual None values. It supports: + +- Recognition of common null-like string representations (e.g., "n/a", "null", "none") +- Case-insensitive matching of null-like strings +- Conversion of null-like strings to Python None + +Null normalization is useful for standardizing missing data representation, +improving data quality, and ensuring consistent handling of null values in +downstream processing and analysis. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import Optional, Set + +class NullNormalizer: + """Normalizes null-like string values to actual None values. + + This class provides functionality to identify and convert common string + representations of null or missing values to Python None. It uses a predefined + set of null-like strings for matching. + + Attributes: + NULL_LIKE (Set[str]): Set of string values that should be considered as null + """ + + NULL_LIKE = {"", "n/a", "na", "null", "none", "-"} + + def normalize(self, value: str) -> Optional[str]: + """Normalize a string value, converting null-like strings to None. + + Checks if the input value is None or matches any of the predefined null-like + strings (after stripping whitespace and converting to lowercase). If it does, + returns None; otherwise, returns the original value unchanged. + + Args: + value (str): The string value to normalize + + Returns: + Optional[str]: None if the value is null-like, otherwise the original value + """ + if value is None: + return None + return None if value.strip().lower() in self.NULL_LIKE else value \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_nulls_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_nulls_provider.py new file mode 100644 index 00000000..0db8750d --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_nulls_provider.py @@ -0,0 +1,74 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Nulls Provider module for document graph operations. + +This module provides a transformer provider implementation for normalizing null-like +values in document records to actual None values. It supports: + +- Recognition of common null-like string representations (e.g., "n/a", "null", "none") +- Case-insensitive matching of null-like strings +- Conversion of null-like strings to Python None +- Configurable list of fields to process +- Customizable set of null-like string patterns + +Null normalization is useful for standardizing missing data representation, +improving data quality, and ensuring consistent handling of null values in +downstream processing and analysis. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class NormalizeNullsProvider(TransformerProvider): + """Normalizes null-like values to actual None. + + This transformer provider identifies and converts common string representations + of null or missing values to Python None in specified fields of document records. + It uses a configurable set of null-like strings for matching. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config), including: + - fields: List of field names to process + - null_like: Set of string values that should be considered as null + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by normalizing null-like values to None. + + Processes each record and converts null-like string values to None in the + specified fields. The matching is done after stripping whitespace and + converting to lowercase. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with normalized null values + """ + fields = self.args.get("fields", []) + null_like = set(self.args.get("null_like", ["", "n/a", "na", "null", "none", "-"])) + + normalized_records = [] + for record in records: + normalized_record = record.copy() + + for field in fields: + if field in normalized_record: + value = normalized_record[field] + if isinstance(value, str) and value.strip().lower() in null_like: + normalized_record[field] = None + + normalized_records.append(normalized_record) + + return normalized_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_spelling.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_spelling.py new file mode 100644 index 00000000..863d1518 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_spelling.py @@ -0,0 +1,56 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Spelling module for document graph operations. + +This module provides functionality for correcting spelling errors in text data +using the TextBlob library. It supports: + +- Automatic detection and correction of misspelled words +- Preservation of original text structure and formatting +- Handling of empty or None values + +Spelling normalization is useful for improving text quality, enhancing search +and matching operations, and preparing text for further natural language processing. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +# normalize_spelling.py +from textblob import TextBlob +from typing import Optional + +class SpellingNormalizer: + """Normalizes spelling in text by correcting common spelling errors. + + This class provides functionality to automatically detect and correct + spelling errors in text using the TextBlob library's spell checking + capabilities. + """ + + def normalize(self, value: str) -> Optional[str]: + """Normalize spelling in a text string. + + Uses TextBlob to detect and correct spelling errors in the input text. + If the input is empty or None, returns it unchanged. + + Args: + value (str): The text string to normalize + + Returns: + Optional[str]: The text with corrected spelling, or the original value + if it was empty or None + + Raises: + ImportError: If TextBlob is not installed + Various exceptions from TextBlob if text processing fails + """ + if not value: + return value + blob = TextBlob(value) + return str(blob.correct()) \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_spelling_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_spelling_provider.py new file mode 100644 index 00000000..e49c742d --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_spelling_provider.py @@ -0,0 +1,84 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Spelling Provider module for document graph operations. + +This module provides a transformer provider implementation for correcting spelling +errors in text fields of document records using the TextBlob library. It supports: + +- Automatic detection and correction of misspelled words +- Processing of multiple configurable fields +- Graceful handling of errors during spelling correction +- Preservation of non-string and empty values + +Spelling normalization is useful for improving text quality, enhancing search +and matching operations, and preparing text for further natural language processing. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +try: + from textblob import TextBlob +except ImportError: + TextBlob = None + + +class NormalizeSpellingProvider(TransformerProvider): + """Normalizes spelling using TextBlob. + + This transformer provider corrects spelling errors in text fields of document + records using the TextBlob library's spell checking capabilities. It processes + specified fields in each record and handles errors gracefully. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config), including: + - fields: List of field names to process (default: ["content"]) + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by correcting spelling errors in text fields. + + Processes each record and applies spelling correction to the specified fields + using TextBlob. Non-string values and empty strings are preserved unchanged. + Errors during spelling correction are caught and the original value is kept. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with corrected spelling + + Raises: + ImportError: If TextBlob is not installed + """ + if TextBlob is None: + raise ImportError("TextBlob required for spelling normalization") + + fields = self.args.get("fields", ["content"]) + + normalized_records = [] + for record in records: + normalized_record = record.copy() + + for field in fields: + value = record.get(field) + if value and isinstance(value, str): + try: + blob = TextBlob(value) + normalized_record[field] = str(blob.correct()) + except Exception: + pass # Keep original value on error + + normalized_records.append(normalized_record) + + return normalized_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_timestamp.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_timestamp.py new file mode 100644 index 00000000..34b37d31 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_timestamp.py @@ -0,0 +1,57 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Timestamp module for document graph operations. + +This module provides functionality for normalizing timestamp strings to a standard +ISO 8601 format. It supports: + +- Parsing various datetime string formats using dateutil.parser +- Converting parsed datetimes to ISO 8601 format +- Handling of invalid or unparseable datetime strings +- Graceful handling of None values + +Timestamp normalization is useful for standardizing date and time representations, +enabling consistent sorting and comparison operations, and ensuring compatibility +with various systems and databases. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from dateutil import parser +from typing import Optional + +class TimestampNormalizer: + """Normalizes timestamp strings to ISO 8601 format. + + This class provides functionality to parse various datetime string formats + and convert them to a standardized ISO 8601 format, which is widely supported + and enables consistent sorting and comparison operations. + """ + + def normalize(self, value: str) -> Optional[str]: + """Normalize a timestamp string to ISO 8601 format. + + Attempts to parse the input string as a datetime using dateutil.parser + and converts it to ISO 8601 format. If parsing fails, returns None. + + Args: + value (str): The timestamp string to normalize + + Returns: + Optional[str]: The normalized timestamp in ISO 8601 format if parsing + succeeds, None otherwise + + Raises: + No exceptions are raised; parsing errors are caught and None is returned + """ + try: + dt = parser.parse(value) + return dt.isoformat() + except (ValueError, TypeError): + return None \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_timestamp_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_timestamp_provider.py new file mode 100644 index 00000000..65547135 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_timestamp_provider.py @@ -0,0 +1,77 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Timestamp Provider module for document graph operations. + +This module provides a transformer provider implementation for normalizing timestamp +strings in document records to a standard ISO 8601 format. It supports: + +- Parsing various datetime string formats using dateutil.parser +- Converting parsed datetimes to ISO 8601 format +- Processing of multiple configurable fields +- Graceful handling of errors during timestamp parsing +- Preservation of non-string and empty values + +Timestamp normalization is useful for standardizing date and time representations, +enabling consistent sorting and comparison operations, and ensuring compatibility +with various systems and databases. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import List, Dict, Any +from dateutil import parser +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class NormalizeTimestampProvider(TransformerProvider): + """Normalizes timestamp fields to ISO 8601 format. + + This transformer provider parses timestamp strings in various formats and + converts them to a standardized ISO 8601 format. It processes specified fields + in each record and handles parsing errors gracefully. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config), including: + - fields: List of field names to process (default: ["timestamp"]) + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by normalizing timestamp fields to ISO 8601 format. + + Processes each record and attempts to parse timestamp strings in the specified + fields using dateutil.parser, then converts them to ISO 8601 format. Non-string + values are preserved unchanged, and parsing errors are caught to keep the + original value. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with normalized timestamps + """ + fields = self.args.get("fields", ["timestamp"]) + + normalized_records = [] + for record in records: + normalized_record = record.copy() + + for field in fields: + value = record.get(field) + if value: + try: + dt = parser.parse(str(value)) + normalized_record[field] = dt.isoformat() + except Exception: + pass # Keep original value on error + + normalized_records.append(normalized_record) + + return normalized_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_whitespace.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_whitespace.py new file mode 100644 index 00000000..9637332b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_whitespace.py @@ -0,0 +1,43 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Whitespace module for document graph operations. + +This module provides functionality for normalizing whitespace in text strings. +It supports: + +- Removing leading and trailing whitespace +- Replacing multiple consecutive whitespace characters with a single space +- Standardizing whitespace representation across different text sources + +Whitespace normalization is useful for improving text consistency, enhancing +search and matching operations, and preparing text for further processing or +display. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +import re + +def normalize(value: str) -> str: + """Normalize whitespace in a text string. + + Removes leading and trailing whitespace and replaces multiple consecutive + whitespace characters (spaces, tabs, newlines, etc.) with a single space. + + Args: + value (str): The string to normalize + + Returns: + str: The normalized string with standardized whitespace + + Examples: + >>> normalize(" Hello world!\\n") + 'Hello world!' + """ + return re.sub(r'\s+', ' ', value.strip()) diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_whitespace_provider.py b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_whitespace_provider.py new file mode 100644 index 00000000..016f5cf9 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/normalizers/normalize_whitespace_provider.py @@ -0,0 +1,70 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Whitespace Provider module for document graph operations. + +This module provides a transformer provider implementation for normalizing whitespace +in text fields of document records. It supports: + +- Removing leading and trailing whitespace +- Replacing multiple consecutive whitespace characters with a single space +- Processing of multiple configurable fields +- Preservation of non-string and empty values + +Whitespace normalization is useful for improving text consistency, enhancing +search and matching operations, and preparing text for further processing or +display. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +import re +from typing import List, Dict, Any +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class NormalizeWhitespaceProvider(TransformerProvider): + """Normalizes whitespace in text fields. + + This transformer provider standardizes whitespace in text fields of document + records by removing leading and trailing whitespace and replacing multiple + consecutive whitespace characters with a single space. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config), including: + - fields: List of field names to process (default: ["content"]) + """ + + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Transform records by normalizing whitespace in text fields. + + Processes each record and applies whitespace normalization to the specified + fields. Non-string values and empty strings are preserved unchanged. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries with normalized whitespace + """ + fields = self.args.get("fields", ["content"]) + + normalized_records = [] + for record in records: + normalized_record = record.copy() + + for field in fields: + value = record.get(field) + if value and isinstance(value, str): + normalized_record[field] = re.sub(r'\s+', ' ', value.strip()) + + normalized_records.append(normalized_record) + + return normalized_records \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/registry/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/registry/__init__.py new file mode 100644 index 00000000..290478b0 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/registry/__init__.py @@ -0,0 +1,21 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Registry module for transformer providers in document graph operations. + +This package provides registry and factory components for managing transformer providers +in the document graph processing system. It includes classes for registering, discovering, +and instantiating transformer provider implementations. +""" + +from graphrag_toolkit.document_graph.transform.registry.transformer_provider_registry import TransformerProviderRegistry +from graphrag_toolkit.document_graph.transform.registry.transformer_provider_factory import TransformerProviderFactory + +# Create a singleton instance for backward compatibility +transformer_provider_registry = TransformerProviderRegistry() + +__all__ = [ + "TransformerProviderRegistry", + "TransformerProviderFactory", + "transformer_provider_registry", +] \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/registry/transformer_provider_factory.py b/document-graph/src/graphrag_toolkit/document_graph/transform/registry/transformer_provider_factory.py new file mode 100644 index 00000000..3c846ed0 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/registry/transformer_provider_factory.py @@ -0,0 +1,75 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transformer Provider Factory module for document graph operations. + +This module provides a factory for dynamically instantiating transformer provider +implementations based on configuration. It supports: + +- Dynamic class loading from fully qualified paths +- Instantiation of transformer providers with appropriate configuration +- Integration with the transformer provider registry system + +The factory pattern used here allows for flexible creation of transformer providers +without tight coupling to specific implementations, enabling extensibility through plugins. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +import importlib +from typing import Type + +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider +from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + + +class TransformerProviderFactory: + """ + Factory for instantiating TransformerProvider implementations + based on the configuration provided. + + This factory class works in conjunction with the TransformerProviderRegistry + to create instances of transformer providers. It handles: + + - Dynamic loading of provider classes from their fully qualified names + - Instantiation of provider objects with appropriate configuration + - Error handling for missing or invalid provider implementations + + The factory pattern implemented here allows for flexible creation of transformer + providers at runtime without requiring direct dependencies on specific implementations. + """ + + @staticmethod + def load_class(path: str) -> Type[TransformerProvider]: + """ + Dynamically import and return the class from a fully qualified path. + + Args: + path (str): e.g., "my_module.sub_module.MyClass" + + Returns: + Type[TransformerProvider]: The class object. + """ + module_path, class_name = path.rsplit(".", 1) + module = importlib.import_module(module_path) + klass = getattr(module, class_name) + return klass + + @classmethod + def get_provider(cls, config: TransformerProviderConfig) -> TransformerProvider: + """ + Instantiate a transformers provider from its config. + + Args: + config (TransformerProviderConfig): Configuration for the provider. + + Returns: + TransformerProvider: Instantiated provider. + """ + klass = cls.load_class(config.implementation) + return klass(config=config) diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/registry/transformer_provider_registry.py b/document-graph/src/graphrag_toolkit/document_graph/transform/registry/transformer_provider_registry.py new file mode 100644 index 00000000..3474b26f --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/registry/transformer_provider_registry.py @@ -0,0 +1,76 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transformer Provider Registry module for document graph operations. + +This module provides a registry system for managing transformer provider implementations +in the document graph processing system. It supports: + +- Registration of transformer providers with unique identifiers +- Lookup of transformer providers by name +- Listing of all available transformer providers + +The registry pattern implemented here enables a plugin architecture where transformer +providers can be registered dynamically and retrieved by name, facilitating extensibility +and decoupling between components that define transformers and those that use them. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from typing import Dict, Type + +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + + +class TransformerProviderRegistry: + """ + A registry to hold and manage all available transformers providers. + + This allows for plugin-style registration and lookup of transformers classes + by name, which is useful for managing normalizers, enrichers, truncators, etc. + """ + + _registry: Dict[str, Type[TransformerProvider]] = {} + + @classmethod + def register(cls, name: str, provider_class: Type[TransformerProvider]) -> None: + """ + Register a transformers provider with a given name. + + Args: + name (str): The unique identifier for the provider. + provider_class (Type[TransformerProvider]): The provider class. + """ + if name in cls._registry: + raise ValueError(f"Transformer provider '{name}' is already registered.") + cls._registry[name] = provider_class + + @classmethod + def get(cls, name: str) -> Type[TransformerProvider]: + """ + Retrieve a transformers provider by name. + + Args: + name (str): The name of the registered provider. + + Returns: + Type[TransformerProvider]: The provider class. + """ + if name not in cls._registry: + raise ValueError(f"Transformer provider '{name}' not found in registry.") + return cls._registry[name] + + @classmethod + def list_registered(cls) -> Dict[str, Type[TransformerProvider]]: + """ + List all registered transformers providers. + + Returns: + Dict[str, Type[TransformerProvider]]: All registered providers. + """ + return dict(cls._registry) diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/transformation_plan.py b/document-graph/src/graphrag_toolkit/document_graph/transform/transformation_plan.py new file mode 100644 index 00000000..954bc916 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/transformation_plan.py @@ -0,0 +1,92 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transformation plan models for document graph processing. + +This module defines the data models used to represent transformation plans +and steps for document graph processing pipelines. Transformation plans +describe a sequence of operations to be applied to documents or datasets +during processing, such as normalization, enrichment, and truncation. +""" + +import logging +from typing import List, Literal, Optional, Union +from pydantic import BaseModel, Field + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +class TransformationStep(BaseModel): + """One step in a transformation plan. + + Represents a single operation in a document transformation pipeline, + such as normalization, enrichment, or truncation. + + Attributes: + type: The type of transformation step (transformers, normalizer, enricher, truncator) + name: The registered name or class of the step + args: Optional arguments for configuring the step + + Examples: + >>> normalizer_step = TransformationStep( + ... type="normalizer", + ... name="text_normalizer", + ... args={"lowercase": True, "remove_punctuation": True} + ... ) + """ + type: Literal["transformers", "normalizer", "enricher", "truncator"] + name: str = Field(..., description="The registered name or class of the step") + args: Optional[dict] = Field(default_factory=dict, description="Optional kwargs for the step") + +class TransformationPlan(BaseModel): + """A full transformation pipeline plan for a document or dataset. + + Represents a complete sequence of transformation operations to be applied + to documents or datasets during processing. Each plan has a unique ID, + an optional description, and an ordered list of transformation steps. + + Attributes: + id: Unique identifier for the transformation plan + description: Optional human-readable description of the plan + steps: Ordered list of transformation steps to apply + + Examples: + >>> plan = TransformationPlan( + ... id="doc-processing-plan", + ... description="Standard document processing pipeline", + ... steps=[ + ... TransformationStep(type="normalizer", name="text_normalizer", + ... args={"lowercase": True}), + ... TransformationStep(type="enricher", name="entity_enricher", + ... args={"entities": ["person", "organization"]}) + ... ] + ... ) + """ + id: str + description: Optional[str] = None + steps: List[TransformationStep] = Field(..., description="Ordered list of transformation steps to apply") + + def get_steps_by_type(self, step_type: str) -> List[TransformationStep]: + """Filter steps by their type. + + Args: + step_type: The type of steps to filter for (e.g., "normalizer", "enricher") + + Returns: + List of transformation steps matching the specified type + + Examples: + >>> plan = TransformationPlan( + ... id="doc-processing-plan", + ... steps=[ + ... TransformationStep(type="normalizer", name="text_normalizer"), + ... TransformationStep(type="enricher", name="entity_enricher") + ... ] + ... ) + >>> normalizers = plan.get_steps_by_type("normalizer") + >>> len(normalizers) + 1 + >>> normalizers[0].name + 'text_normalizer' + """ + return [step for step in self.steps if step.type == step_type] diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_base.py b/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_base.py new file mode 100644 index 00000000..ee22f87a --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_base.py @@ -0,0 +1,115 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Base class for transformer providers. + +This module defines the abstract base class for all transformer providers in the +document graph processing system. Transformer providers implement data transformation +operations that can be used in ETL pipelines, document processing, and graph building. +The module provides a common interface for all transformers, regardless of their +specific transformation logic. +""" + +import logging +from abc import ABC, abstractmethod +from typing import Any, List, Dict +from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +class TransformerProvider(ABC): + """Abstract base class for all transformers providers. + + Provides the interface for data transformation operations + used by plugins and transformation pipelines. All concrete transformer + implementations must inherit from this class and implement the transform method. + + Attributes: + config: Configuration object for the transformer + name: Name of the transformer (from config) + args: Arguments for the transformer (from config) + + Examples: + >>> # Define a custom transformer + >>> class LowercaseTransformer(TransformerProvider): + ... def transform(self, records): + ... for record in records: + ... if 'text' in record: + ... record['text'] = record['text'].lower() + ... return records + >>> + >>> # Create and use the transformer + >>> config = TransformerProviderConfig(name="lowercase", args={"fields": ["text"]}) + >>> transformer = LowercaseTransformer(config) + >>> result = transformer.transform([{"text": "HELLO WORLD"}]) + >>> result[0]["text"] + 'hello world' + """ + + def __init__(self, config: TransformerProviderConfig): + """Initialize transformers with configuration. + + Args: + config: Transformer configuration with name, type, and args + """ + self.config = config + self.name = config.name + self.args = config.args + + @abstractmethod + def transform(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Apply transformation logic to records. + + Args: + records: List of record dictionaries to transform + + Returns: + List of transformed record dictionaries + """ + pass + + def _log_transform_start(self, record_count: int): + """Log the start of a transformation operation. + + This helper method logs a debug message indicating that a transformation + operation is starting, including the number of records to be processed. + + Args: + record_count: Number of records to be transformed + """ + import logging + logger = logging.getLogger(self.__class__.__module__) + logger.debug(f"[{self.__class__.__name__}] Starting transformation of {record_count} records") + + def _log_transform_end(self, input_count: int, output_count: int): + """Log the completion of a transformation operation. + + This helper method logs a debug message indicating that a transformation + operation has completed, including the number of input and output records. + + Args: + input_count: Number of input records processed + output_count: Number of output records produced + """ + import logging + logger = logging.getLogger(self.__class__.__module__) + logger.debug(f"[{self.__class__.__name__}] Completed: {input_count} -> {output_count} records") + + def __repr__(self): + """Return a string representation of the transformer provider. + + Returns: + String representation including the transformer name and type + + Examples: + >>> config = TransformerProviderConfig(name="example", type="test") + >>> transformer = TransformerProvider(config) + >>> repr(transformer) + '' + """ + return f"" + +# Alias for backward compatibility +TransformerConfig = TransformerProviderConfig diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_config.py b/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_config.py new file mode 100644 index 00000000..4d184fa8 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_config.py @@ -0,0 +1,74 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transformer provider configuration models. + +This module defines the configuration models used for transformer providers +in the document graph processing system. These configurations are used to +instantiate and parameterize transformer instances that perform various +data transformation operations during document processing. +""" + +import logging +from typing import Dict, Any, Optional +from pydantic import BaseModel, Field + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +class TransformerProviderConfig(BaseModel): + """Configuration for transformers providers. + + Used by plugins and transformation steps to configure + data transformation operations. + + Attributes: + name: Transformer name for identification + type: Transformer type (optional, for factory resolution) + args: Transformer-specific arguments and parameters + parameters: Alias for args (user-friendly) + + Examples: + >>> config = TransformerProviderConfig( + ... name="text_normalizer", + ... type="normalizer", + ... args={"lowercase": True, "remove_punctuation": True} + ... ) + >>> config.name + 'text_normalizer' + >>> config.args["lowercase"] + True + """ + name: str = Field(..., description="Transformer name") + type: Optional[str] = Field(default=None, description="Transformer type") + args: Dict[str, Any] = Field(default_factory=dict, description="Transformer arguments") + parameters: Dict[str, Any] = Field(default_factory=dict, description="Transformer parameters (alias for args)") + + def __init__(self, **data): + """Initialize transformer provider configuration. + + Initializes the configuration and synchronizes the args and parameters + fields to ensure they contain the same data, regardless of which one + was provided in the input. + + Args: + **data: Keyword arguments for configuration fields including name, + type, args, and parameters + + Examples: + >>> config = TransformerProviderConfig( + ... name="entity_enricher", + ... parameters={"entities": ["person", "organization"]} + ... ) + >>> config.args == config.parameters + True + """ + super().__init__(**data) + # Sync parameters and args + if self.parameters and not self.args: + self.args = self.parameters + elif self.args and not self.parameters: + self.parameters = self.args + +# Alias for backward compatibility +BaseTransformerConfig = TransformerProviderConfig diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_factory.py b/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_factory.py new file mode 100644 index 00000000..357dc75a --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_factory.py @@ -0,0 +1,120 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transformer Provider Factory module for document graph operations. + +This module provides a factory for creating transformer provider instances +based on configuration. It handles the resolution of transformer types and names, +instantiates the appropriate provider classes, and manages the configuration +process. The factory works with the transformer provider registry to locate +and instantiate the correct transformer implementation. +""" + +import logging +from typing import Type +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider +from graphrag_toolkit.document_graph.transform.transformer_provider_registry import transformer_provider_registry +from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +class TransformerProviderFactory: + """Factory to instantiate transformers providers based on config. + + Creates transformers instances for use in transformation pipelines + and plugin systems. The factory resolves transformer types and names, + retrieves the appropriate provider class from the registry, and + instantiates it with the provided configuration. + + Examples: + >>> config = TransformerProviderConfig( + ... name="text_normalizer", + ... type="normalizer", + ... args={"lowercase": True} + ... ) + >>> provider = TransformerProviderFactory.get_provider(config) + >>> isinstance(provider, TransformerProvider) + True + """ + + @staticmethod + def get_provider(config: TransformerProviderConfig) -> TransformerProvider: + """Instantiate the appropriate transformers provider. + + Args: + config: Transformer configuration with type and parameters + + Returns: + TransformerProvider: Configured transformers instance + + Raises: + KeyError: If transformers type is not registered + + Examples: + >>> config = TransformerProviderConfig( + ... name="entity_enricher", + ... type="enricher", + ... args={"entities": ["person", "organization"]} + ... ) + >>> provider = TransformerProviderFactory.get_provider(config) + >>> provider.transform({"text": "John works at Amazon"}) + {'text': 'John works at Amazon', 'entities': [...]} + """ + logger.debug(f"Creating transformers provider: {config.type or config.name}") + + # Use name for lookup when type is a directory-based category, otherwise use type + directory_types = ["transformer", "normalizer", "enricher", "truncator", + "field_transformer", "document_transformer", "filter_transformer", + "graph_transformer"] + if config.type in directory_types: + lookup_key = config.name + logger.debug(f"Using name '{lookup_key}' for category type '{config.type}'") + else: + lookup_key = config.type or config.name + logger.debug(f"Using direct lookup key '{lookup_key}'") + + logger.debug(f"Registry lookup for key: {lookup_key}") + logger.debug(f"Available providers: {transformer_provider_registry.list_providers()}") + + provider_cls: Type[TransformerProvider] = transformer_provider_registry.get(lookup_key) + + provider = provider_cls(config) + logger.debug(f"Successfully created transformers: {provider.__class__.__name__}") + return provider + + @staticmethod + def create_config( + name: str, + transformer_type: str = None, + **args + ) -> TransformerProviderConfig: + """Create a transformers configuration. + + Args: + name: Transformer name + transformer_type: Optional transformers type for factory lookup + **args: Transformer-specific arguments + + Returns: + TransformerProviderConfig: Transformer configuration + + Examples: + >>> config = TransformerProviderFactory.create_config( + ... name="text_normalizer", + ... transformer_type="normalizer", + ... lowercase=True, + ... remove_punctuation=True + ... ) + >>> config.name + 'text_normalizer' + >>> config.type + 'normalizer' + >>> config.args["lowercase"] + True + """ + return TransformerProviderConfig( + name=name, + type=transformer_type, + args=args + ) diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_registry.py b/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_registry.py new file mode 100644 index 00000000..22660a7b --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/transformer_provider_registry.py @@ -0,0 +1,249 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transformer Provider Registry module for document graph operations. + +This module provides a registry system for transformer providers in the document +graph processing system. It maintains a mapping of transformer type names to their +concrete implementations, allowing for dynamic discovery and instantiation of +transformer classes. The registry supports registration, lookup, and introspection +of transformer providers, and is used by the TransformerProviderFactory and plugin +system to locate and create transformer instances. +""" + +import logging +from typing import Type, Dict +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + +class TransformerProviderRegistry: + """Registry for transformers providers. + + Maintains a mapping of transformers type names to their concrete implementations. + Used by the TransformerProviderFactory and plugin system to discover and + instantiate transformer providers dynamically. + + Examples: + >>> # Register a custom transformer provider + >>> class MyTransformer(TransformerProvider): + ... def transform(self, data): + ... return data + >>> + >>> # Register the transformer + >>> TransformerProviderRegistry.register("my_transformer", MyTransformer) + >>> + >>> # Get the transformer class + >>> transformer_cls = TransformerProviderRegistry.get("my_transformer") + >>> transformer_cls is MyTransformer + True + """ + + _registry: Dict[str, Type[TransformerProvider]] = {} + _lazy_registry: Dict[str, callable] = {} + + @classmethod + def register(cls, type_: str, provider_cls: Type[TransformerProvider]) -> None: + """Register a new transformers provider type. + + Args: + type_: Provider type name + provider_cls: Provider class to register + """ + logger.debug(f"Registering transformers provider: {type_} -> {provider_cls.__name__}") + cls._registry[type_.lower()] = provider_cls + + @classmethod + def register_lazy(cls, type_: str, loader_func: callable) -> None: + """Register a lazy-loaded transformer provider. + + Args: + type_: Provider type name + loader_func: Function that returns the provider class when called + """ + logger.debug(f"Registering lazy transformers provider: {type_}") + cls._lazy_registry[type_.lower()] = loader_func + + @classmethod + def get(cls, type_: str) -> Type[TransformerProvider]: + """Get a transformers provider class by type. + + Args: + type_: Provider type name + + Returns: + Type[TransformerProvider]: Provider class + + Raises: + KeyError: If provider type is not registered + + Examples: + >>> # Get a registered transformer provider + >>> normalizer_cls = TransformerProviderRegistry.get("text_normalizer") + >>> normalizer = normalizer_cls(config) + >>> normalizer.transform("TEXT TO NORMALIZE") + 'text to normalize' + >>> + >>> # Trying to get an unregistered provider + >>> try: + ... TransformerProviderRegistry.get("nonexistent_provider") + ... except KeyError as e: + ... print(f"Error: {e}") + Error: 'No transformers provider registered for type='nonexistent_provider'' + """ + key = type_.lower() + logger.debug(f"Looking up transformers provider: {type_}") + + # Check eager registry first + if key in cls._registry: + return cls._registry[key] + + # Check lazy registry + if key in cls._lazy_registry: + try: + provider_cls = cls._lazy_registry[key]() + # Cache the loaded class + cls._registry[key] = provider_cls + return provider_cls + except Exception as e: + logger.error(f"Failed to load lazy transformer {type_}: {e}") + raise KeyError(f"Failed to load transformer '{type_}': {e}") + + # Provide helpful suggestions + available = cls.list_providers() + suggestions = [name for name in available if type_.lower() in name or name in type_.lower()] + + error_msg = f"No transformer '{type_}' found." + if suggestions: + error_msg += f" Did you mean: {', '.join(suggestions[:3])}?" + error_msg += f" Available: {', '.join(sorted(available))}" + + logger.error(error_msg) + raise KeyError(error_msg) + + @classmethod + def list_providers(cls) -> list[str]: + """List all registered provider types. + + Returns: + list[str]: List of registered provider type names + + Examples: + >>> # Register some providers + >>> TransformerProviderRegistry.register("text_normalizer", TextNormalizer) + >>> TransformerProviderRegistry.register("entity_enricher", EntityEnricher) + >>> + >>> # List all registered providers + >>> providers = TransformerProviderRegistry.list_providers() + >>> sorted(providers) + ['entity_enricher', 'text_normalizer'] + """ + return list(set(list(cls._registry.keys()) + list(cls._lazy_registry.keys()))) + + @classmethod + def describe_providers(cls) -> dict[str, dict[str, str]]: + """Get detailed information about all registered providers. + + Returns: + dict: Provider info with name, description, and category + + Examples: + >>> # Get detailed information about registered providers + >>> info = TransformerProviderRegistry.describe_providers() + >>> 'text_normalizer' in info + True + >>> info['text_normalizer']['category'] + 'normalizer' + >>> info['text_normalizer']['description'] + 'Normalizes text input' + """ + info = {} + + # Process eager registry + for name, provider_cls in cls._registry.items(): + info[name] = cls._get_provider_info(name, provider_cls) + + # Process lazy registry - create basic info without loading + for name in cls._lazy_registry.keys(): + if name not in info: # Don't override if already loaded + # Create basic info based on name patterns without loading + category = 'transformer' + type_name = 'transformer' + + if 'normalize' in name: + category = 'normalizer' + type_name = 'normalizer' + elif 'enricher' in name: + category = 'enricher' + type_name = 'enricher' + elif name in ['column_pruner', 'row_filter']: + category = 'filter' + type_name = 'filter_transformer' + elif name in ['infer_edges', 'row_to_node']: + category = 'graph' + type_name = 'graph_transformer' + elif name in ['pii_redactor', 'json_to_rows', 'text_chunker']: + category = 'document' + type_name = 'document_transformer' + elif name in ['regex_cleaner', 'embedded_json', 'standardize_enum', 'comma_split', 'json_array_expander', 'json_flattener', 'comma_flattener']: + category = 'field' + type_name = 'field_transformer' + + info[name] = { + 'name': name, + 'description': f'{category.title()} transformer for {name.replace("_", " ")}', + 'category': category, + 'type': type_name, + 'class': 'LazyLoaded' + } + + return info + + @classmethod + def _get_provider_info(cls, name: str, provider_cls) -> dict[str, str]: + """Extract provider information from a provider class.""" + # Extract description from docstring + doc = provider_cls.__doc__ or "No description available" + description = doc.split('\n')[0].strip() if doc else "No description" + + # Determine category and type from class name and module patterns + class_name = provider_cls.__name__.lower() + module_name = provider_cls.__module__.lower() + + if 'normalize' in class_name or 'normalizers' in module_name: + category = 'normalizer' + type_name = 'normalizer' + elif 'enrich' in class_name or 'enrichers' in module_name: + category = 'enricher' + type_name = 'enricher' + elif 'filter' in class_name or 'prune' in class_name or 'filter_transformers' in module_name: + category = 'filter' + type_name = 'filter_transformer' + elif 'graph' in class_name or 'node' in class_name or 'edge' in class_name or 'graph_transformers' in module_name: + category = 'graph' + type_name = 'graph_transformer' + elif 'document_transformers' in module_name: + category = 'document' + type_name = 'document_transformer' + elif 'field_transformers' in module_name: + category = 'field' + type_name = 'field_transformer' + elif 'truncators' in module_name: + category = 'truncator' + type_name = 'truncator' + else: + category = 'transformer' + type_name = 'transformer' + + return { + 'name': name, + 'description': description, + 'category': category, + 'type': type_name, + 'class': provider_cls.__name__ + } + + +# Initialize registry instance +transformer_provider_registry = TransformerProviderRegistry() \ No newline at end of file diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/__init__.py b/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/__init__.py new file mode 100644 index 00000000..00978d27 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/__init__.py @@ -0,0 +1 @@ +"""Truncators — field truncation strategies for length, count, and token limits.""" diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/field_count_truncator.py b/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/field_count_truncator.py new file mode 100644 index 00000000..ad8aa02c --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/field_count_truncator.py @@ -0,0 +1,62 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Field Count Truncator module for document graph operations. + +This module provides a transformer that limits the number of fields in a record +to prevent oversized documents and ensure consistent processing. It keeps only +the first N fields from the input record, discarding any additional fields. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +class FieldCountTruncator(TransformerProvider): + """Transformer that limits the number of fields in a record. + + This transformer ensures records don't exceed a maximum number of fields + by keeping only the first N fields and discarding any additional ones. + This helps prevent oversized documents and ensures consistent processing + downstream. + + Attributes: + config: Configuration object containing transformer settings + + Examples: + >>> from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + >>> config = TransformerProviderConfig( + ... name="field_count_truncator", + ... args={"max_fields": 5} + ... ) + >>> truncator = FieldCountTruncator(config) + >>> result = truncator.transform({"a": 1, "b": 2, "c": 3, "d": 4, "e": 5, "f": 6, "g": 7}) + >>> len(result) + 5 + >>> "f" in result + False + """ + + def transform(self, record: dict) -> dict: + """Limit the number of fields in a record. + + Args: + record: A dictionary containing record data + + Returns: + A dictionary with at most max_fields keys + + Note: + This implementation differs from the base class by accepting a single + record dictionary rather than a list of records. + """ + max_fields = self.config.args.get("max_fields", 20) + if len(record) > max_fields: + keys = list(record.keys())[:max_fields] + record = {k: record[k] for k in keys} + return record diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/length_truncator.py b/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/length_truncator.py new file mode 100644 index 00000000..7b01cec4 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/length_truncator.py @@ -0,0 +1,66 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Length Truncator module for document graph operations. + +This module provides a transformer that limits the character length of string fields +in a record to prevent oversized text fields and ensure consistent processing. +It truncates specified string fields to a maximum character length. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +class LengthTruncator(TransformerProvider): + """Transformer that limits the character length of string fields. + + This transformer ensures string fields don't exceed a maximum character length + by truncating them to the specified maximum. This helps prevent oversized text + fields and ensures consistent processing downstream. + + Attributes: + config: Configuration object containing transformer settings + + Examples: + >>> from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + >>> config = TransformerProviderConfig( + ... name="length_truncator", + ... args={"max_length": 10, "fields": ["text", "description"]} + ... ) + >>> truncator = LengthTruncator(config) + >>> result = truncator.transform({"text": "This is a long text that will be truncated", "description": "Short desc"}) + >>> result["text"] + 'This is a ' + >>> result["description"] + 'Short desc' + """ + + def transform(self, record: dict) -> dict: + """Truncate specified string fields to a maximum character length. + + Args: + record: A dictionary containing record data + + Returns: + A dictionary with string fields truncated to max_length characters + + Note: + This implementation differs from the base class by accepting a single + record dictionary rather than a list of records. + + Only processes fields specified in the "fields" configuration parameter. + Fields that don't exist in the record or aren't strings are ignored. + """ + max_length = self.config.args.get("max_length", 512) + fields = self.config.args.get("fields", []) + + for field in fields: + if field in record and isinstance(record[field], str): + record[field] = record[field][:max_length] + return record diff --git a/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/token_truncator.py b/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/token_truncator.py new file mode 100644 index 00000000..27bcff89 --- /dev/null +++ b/document-graph/src/graphrag_toolkit/document_graph/transform/truncators/token_truncator.py @@ -0,0 +1,92 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Token Truncator module for document graph operations. + +This module provides a transformer that limits the number of tokens in string fields +of a record to prevent oversized text fields and ensure consistent processing. +It uses a pre-trained tokenizer to tokenize text and truncate to a maximum token count. +""" + +import logging + + +logger = logging.getLogger(__name__) +# Import custom logging to ensure configuration is available + + +from transformers import AutoTokenizer +from graphrag_toolkit.document_graph.transform.transformer_provider_base import TransformerProvider + +class TokenTruncator(TransformerProvider): + """Transformer that limits the number of tokens in string fields. + + This transformer ensures string fields don't exceed a maximum token count + by tokenizing the text using a pre-trained tokenizer and truncating to the + specified maximum number of tokens. This helps prevent oversized text fields + and ensures consistent processing downstream, especially for models with + token limits. + + Attributes: + config: Configuration object containing transformer settings + tokenizer: Pre-trained tokenizer from the Hugging Face transformers library + + Examples: + >>> from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig + >>> config = TransformerProviderConfig( + ... name="token_truncator", + ... args={ + ... "max_tokens": 5, + ... "fields": ["text"], + ... "model_name": "bert-base-uncased" + ... } + ... ) + >>> truncator = TokenTruncator(config) + >>> result = truncator.transform({"text": "This is a long text that will be truncated to just five tokens"}) + >>> result["text"] + 'this is a long text' + """ + + def __init__(self, config): + """Initialize the token truncator with configuration. + + Args: + config: Transformer configuration with name, type, and args + + Note: + Loads a pre-trained tokenizer based on the model_name in config.args. + Defaults to "bert-base-uncased" if not specified. + """ + super().__init__(config) + model_name = self.config.args.get("model_name", "bert-base-uncased") + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + + def transform(self, record: dict) -> dict: + """Truncate specified string fields to a maximum token count. + + Args: + record: A dictionary containing record data + + Returns: + A dictionary with string fields truncated to max_tokens tokens + + Note: + This implementation differs from the base class by accepting a single + record dictionary rather than a list of records. + + Only processes fields specified in the "fields" configuration parameter. + Fields that don't exist in the record or aren't strings are ignored. + + The tokenization and detokenization process may slightly alter the text + beyond just truncation, as the tokenizer's conversion back to string + may normalize whitespace or other characters. + """ + max_tokens = self.config.args.get("max_tokens", 128) + fields = self.config.args.get("fields", []) + + for field in fields: + if field in record and isinstance(record[field], str): + tokens = self.tokenizer.tokenize(record[field]) + truncated = self.tokenizer.convert_tokens_to_string(tokens[:max_tokens]) + record[field] = truncated + return record diff --git a/document-graph/tests/conftest.py b/document-graph/tests/conftest.py new file mode 100644 index 00000000..bdf39754 --- /dev/null +++ b/document-graph/tests/conftest.py @@ -0,0 +1,7 @@ +"""Pytest configuration for document-graph tests.""" + +import sys +from pathlib import Path + +# Add src to path so tests can import document_graph +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) diff --git a/document-graph/tests/test_document_graph.py b/document-graph/tests/test_document_graph.py new file mode 100644 index 00000000..ce0a75f2 --- /dev/null +++ b/document-graph/tests/test_document_graph.py @@ -0,0 +1,553 @@ +"""Tests for document-graph v3.""" + +import pytest +from graphrag_toolkit.document_graph import NodeModel, EdgeModel, Node, Edge +from graphrag_toolkit.document_graph.graph_build import node_to_cypher, edge_to_cypher, batch_nodes_to_cypher +from graphrag_toolkit.document_graph.query import DocumentGraphQueryEngine + + +class TestModel: + def test_node_model_importable(self): + assert NodeModel is not None + + def test_edge_model_importable(self): + assert EdgeModel is not None + + +class TestModelElements: + def test_node_creation(self): + n = Node(id="u1", labels=["User"], properties={"name": "Alice", "email": "a@b.com"}) + assert n.id == "u1" + assert "User" in n.labels + assert n.properties["name"] == "Alice" + + def test_edge_creation(self): + e = Edge(id="e1", source_id="u1", target_id="a1", label="OWNS") + assert e.label == "OWNS" + assert e.source_id == "u1" + assert e.target_id == "a1" + + +class TestCypherBuilder: + def test_node_to_cypher(self): + n = Node(id="u1", labels=["User"], properties={"name": "Alice"}) + query, params = node_to_cypher(n) + assert "MERGE" in query + assert "User" in query + assert params["id_val"] == "u1" + + def test_node_to_cypher_with_tenant(self): + n = Node(id="u1", labels=["User"], properties={"name": "Alice"}) + query, params = node_to_cypher(n, tenant_id="scim") + assert "__User__scim__" in query + + def test_edge_to_cypher(self): + e = Edge(id="e1", source_id="u1", target_id="a1", label="OWNS") + query, params = edge_to_cypher(e) + assert "MERGE" in query + assert "OWNS" in query + assert params["src_id"] == "u1" + assert params["tgt_id"] == "a1" + + def test_batch_nodes(self): + nodes = [ + Node(id=f"u{i}", labels=["User"], properties={"name": f"User{i}"}) + for i in range(5) + ] + results = batch_nodes_to_cypher(nodes, tenant_id="test") + assert len(results) == 5 + assert all("MERGE" in q for q, _ in results) + + +class TestQueryEngine: + def test_init(self): + class MockStore: + def query(self, cypher, params=None): + return [{"n": "result"}] + + engine = DocumentGraphQueryEngine(MockStore(), tenant_id="scim") + assert engine._tenant_id == "scim" + + def test_get_nodes(self): + class MockStore: + def execute_query(self, cypher, params=None): + assert "__User__scim__" in cypher + return [{"n": {"name": "Alice"}}] + + engine = DocumentGraphQueryEngine(MockStore(), tenant_id="scim") + result = engine.get_nodes("User", limit=10) + assert len(result) == 1 + + def test_find_by_property(self): + class MockStore: + def execute_query(self, cypher, params=None): + assert params["val"] == "alice@test.com" + return [{"n": {"email": "alice@test.com"}}] + + engine = DocumentGraphQueryEngine(MockStore(), tenant_id="scim") + result = engine.find_by_property("User", "email", "alice@test.com") + assert len(result) == 1 + + +# ───────────────────────────────────────────────────────────────────────────── +# Extended tests — Schema Providers, Transformers, ETLSchema, Graph Build, Discovery +# ───────────────────────────────────────────────────────────────────────────── + +import json +import csv +import tempfile +import uuid +from pathlib import Path +from pydantic import ValidationError + +from graphrag_toolkit.document_graph.graph_build import node_to_cypher, edge_to_cypher, batch_nodes_to_cypher +from graphrag_toolkit.document_graph.schema.providers.schema_provider_config import SchemaProviderConfig +from graphrag_toolkit.document_graph.schema.providers.csv_schema_provider import CSVSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.json_schema_provider import JSONSchemaProvider +from graphrag_toolkit.document_graph.schema.static_schema_provider import StaticSchemaProvider +from graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import SchemaProviderFactory +from graphrag_toolkit.document_graph.schema.etl_schema_model import ( + ETLSchema, ExtractConfig, TransformConfig, LoadConfig, + ChunkingConfig, NormalizeConfig, EntityExtractionConfig, + MetadataMapping, NodeDefinition, RelationshipDefinition, +) +from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig +from graphrag_toolkit.document_graph.transform.normalizers.normalize_whitespace_provider import NormalizeWhitespaceProvider +from graphrag_toolkit.document_graph.transform.normalizers.normalize_nulls_provider import NormalizeNullsProvider +from graphrag_toolkit.document_graph.transform.normalizers.normalize_case_provider import NormalizeCaseProvider +from graphrag_toolkit.document_graph.transform.field_transformers.json_flattener import JSONFlattenerProvider +from graphrag_toolkit.document_graph.transform.field_transformers.uuid_generator import UuidGeneratorTransformer +from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer +from graphrag_toolkit.document_graph.transform.graph_transformers.infer_edges import EdgeInferencer +from graphrag_toolkit.document_graph.transform.filter_transformers.row_filter import RowFilterProvider +from graphrag_toolkit.document_graph.transform.truncators.length_truncator import LengthTruncator +from graphrag_toolkit.document_graph.schema.discovery.csv_discovery_provider import CSVSchemaDiscoveryProvider + + +# ─── Fixtures ──────────────────────────────────────────────────────────────── + +@pytest.fixture +def sample_etl_schema(): + """A valid minimal ETLSchema.""" + return ETLSchema( + schema_id="test-schema", + description="Test", + extract=ExtractConfig(source_type="file"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="fixed_length"), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig(), + ), + load=LoadConfig( + document_node=NodeDefinition(type="Doc", fields=["title"]), + section_node=NodeDefinition(type="Sec", fields=["text"]), + relationships=[RelationshipDefinition(type="has", source="d", target="s")], + ), + ) + + +@pytest.fixture +def tmp_csv(tmp_path): + """Create a temp CSV file for testing.""" + csv_file = tmp_path / "test_data.csv" + csv_file.write_text("id,name,email\n1,Alice,alice@test.com\n2,Bob,bob@test.com\n") + return csv_file + + +@pytest.fixture +def tmp_json(tmp_path): + """Create a temp JSON file for testing.""" + json_file = tmp_path / "test_data.json" + data = [{"id": "1", "name": "Alice"}, {"id": "2", "name": "Bob"}] + json_file.write_text(json.dumps(data)) + return json_file + + +def _make_config(name="test", **args): + """Helper to build a TransformerProviderConfig.""" + return TransformerProviderConfig(name=name, args=args) + + +# ─── ETLSchema Model Tests ────────────────────────────────────────────────── + +class TestETLSchemaModel: + def test_valid_construction(self, sample_etl_schema): + assert sample_etl_schema.schema_id == "test-schema" + assert sample_etl_schema.extract.source_type == "file" + assert sample_etl_schema.transform.chunking.strategy == "fixed_length" + + def test_missing_schema_id_raises(self): + with pytest.raises(ValidationError): + ETLSchema( + extract=ExtractConfig(source_type="file"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="x"), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig(), + ), + load=LoadConfig( + document_node=NodeDefinition(type="D", fields=[]), + section_node=NodeDefinition(type="S", fields=[]), + relationships=[], + ), + ) + + def test_missing_extract_raises(self): + with pytest.raises(ValidationError): + ETLSchema( + schema_id="x", + transform=TransformConfig( + chunking=ChunkingConfig(strategy="x"), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig(), + ), + load=LoadConfig( + document_node=NodeDefinition(type="D", fields=[]), + section_node=NodeDefinition(type="S", fields=[]), + relationships=[], + ), + ) + + def test_model_dump_roundtrip(self, sample_etl_schema): + dumped = sample_etl_schema.model_dump() + reconstructed = ETLSchema(**dumped) + assert reconstructed.schema_id == sample_etl_schema.schema_id + assert reconstructed.extract.source_type == "file" + assert reconstructed.load.document_node.type == "Doc" + + def test_extra_fields_allowed(self): + schema = ETLSchema( + schema_id="extra", + extract=ExtractConfig(source_type="api", custom_field="custom_val"), + transform=TransformConfig( + chunking=ChunkingConfig(strategy="x"), + metadata_mapping=MetadataMapping(), + entity_extraction=EntityExtractionConfig(method="ner"), + normalize=NormalizeConfig(), + ), + load=LoadConfig( + document_node=NodeDefinition(type="D", fields=[]), + section_node=NodeDefinition(type="S", fields=[]), + relationships=[], + ), + ) + assert schema.extract.custom_field == "custom_val" + + +# ─── Schema Providers Tests ───────────────────────────────────────────────── + +class TestSchemaProviders: + def test_csv_schema_provider(self, tmp_csv): + config = SchemaProviderConfig(type="csv", schema_id="csv-test", connection_config={"path": str(tmp_csv)}) + provider = CSVSchemaProvider(config) + schema = provider.load_schema() + assert isinstance(schema, ETLSchema) + assert "id" in schema.load.document_node.fields + assert "name" in schema.load.document_node.fields + + def test_csv_schema_provider_missing_path(self): + config = SchemaProviderConfig(type="csv", connection_config={}) + with pytest.raises(ValueError): + CSVSchemaProvider(config) + + def test_csv_schema_provider_file_not_found(self, tmp_path): + config = SchemaProviderConfig(type="csv", connection_config={"path": str(tmp_path / "nope.csv")}) + with pytest.raises(FileNotFoundError): + CSVSchemaProvider(config) + + def test_json_schema_provider(self, tmp_json): + config = SchemaProviderConfig(type="json", schema_id="json-test", connection_config={"path": str(tmp_json)}) + provider = JSONSchemaProvider(config) + schema = provider.load_schema() + assert isinstance(schema, ETLSchema) + assert "id" in schema.load.document_node.fields + + def test_json_schema_provider_missing_path(self): + config = SchemaProviderConfig(type="json", connection_config={}) + with pytest.raises(ValueError): + JSONSchemaProvider(config) + + def test_static_schema_provider(self): + provider = StaticSchemaProvider({}) + schema = provider.load_schema() + assert schema.schema_id == "static-default" + assert schema.extract.source_type == "s3" + assert schema.load.document_node.type == "DocumentNode" + + def test_static_schema_from_config(self): + provider = StaticSchemaProvider.from_config({"type": "static"}) + assert provider.get_schema_id() == "static-default" + + def test_factory_creates_csv_provider(self, tmp_csv): + config = {"type": "csv", "schema_id": "factory-csv", "connection_config": {"path": str(tmp_csv)}} + provider = SchemaProviderFactory.create(config) + assert isinstance(provider, CSVSchemaProvider) + + def test_factory_creates_json_provider(self, tmp_json): + config = {"type": "json", "schema_id": "factory-json", "connection_config": {"path": str(tmp_json)}} + provider = SchemaProviderFactory.create(config) + assert isinstance(provider, JSONSchemaProvider) + + def test_factory_invalid_type_raises(self): + with pytest.raises(Exception): + SchemaProviderFactory.create({"type": "unknown_xyz", "connection_config": {}}) + + +# ─── Transformer Tests ─────────────────────────────────────────────────────── + +class TestNormalizers: + def test_whitespace_normalizer(self): + cfg = _make_config("ws", fields=["text"]) + t = NormalizeWhitespaceProvider(cfg) + result = t.transform([{"text": " hello world \n\t end "}]) + assert result[0]["text"] == "hello world end" + + def test_whitespace_preserves_non_string(self): + cfg = _make_config("ws", fields=["val"]) + t = NormalizeWhitespaceProvider(cfg) + result = t.transform([{"val": 42}]) + assert result[0]["val"] == 42 + + def test_nulls_normalizer(self): + cfg = _make_config("nulls", fields=["a", "b"]) + t = NormalizeNullsProvider(cfg) + result = t.transform([{"a": "N/A", "b": "hello", "c": "null"}]) + assert result[0]["a"] is None + assert result[0]["b"] == "hello" + assert result[0]["c"] == "null" # not in fields list + + def test_nulls_custom_null_like(self): + cfg = _make_config("nulls", fields=["x"], null_like=["missing"]) + t = NormalizeNullsProvider(cfg) + result = t.transform([{"x": "missing"}]) + assert result[0]["x"] is None + + def test_case_normalizer_lower(self): + cfg = _make_config("case", mode="lower") + t = NormalizeCaseProvider(cfg) + result = t.transform([{"name": "ALICE", "age": 30}]) + assert result[0]["name"] == "alice" + assert result[0]["age"] == 30 + + def test_case_normalizer_upper(self): + cfg = _make_config("case", mode="upper") + t = NormalizeCaseProvider(cfg) + result = t.transform([{"name": "alice"}]) + assert result[0]["name"] == "ALICE" + + +class TestFieldTransformers: + def test_json_flattener(self): + cfg = _make_config("jf", field="data") + t = JSONFlattenerProvider(cfg) + records = [{"id": "1", "data": '{"key": "val", "nested": {"a": 1}}'}] + result = t.transform(records) + assert "datakey" in result[0] + assert result[0]["datakey"] == "val" + + def test_json_flattener_invalid_json(self): + cfg = _make_config("jf", field="data") + t = JSONFlattenerProvider(cfg) + records = [{"id": "1", "data": "not json"}] + result = t.transform(records) + # Should return record without crashing + assert result[0]["id"] == "1" + + def test_uuid_generator(self): + cfg = _make_config("uuid", target_field="my_id") + t = UuidGeneratorTransformer(cfg) + records = [{"name": "Alice"}, {"name": "Bob"}] + result = t.transform(records) + assert "my_id" in result[0] + assert "my_id" in result[1] + # Should be valid UUIDs + uuid.UUID(result[0]["my_id"]) + uuid.UUID(result[1]["my_id"]) + # Should be unique + assert result[0]["my_id"] != result[1]["my_id"] + + +class TestGraphTransformers: + def test_row_to_node(self): + cfg = _make_config("rtn", type="Person") + t = RowToNodeTransformer(cfg) + records = [{"id": "p1", "name": "Alice"}] + result = t.transform(records) + assert result[0]["node_type"] == "Person" + assert result[0]["graph_element"] == "node" + assert result[0]["id"] == "p1" + + def test_row_to_node_auto_id(self): + cfg = _make_config("rtn", type="Row") + t = RowToNodeTransformer(cfg) + records = [{"name": "Alice"}] + result = t.transform(records) + assert result[0]["id"] == "row_0" + + def test_row_to_node_auto_content(self): + cfg = _make_config("rtn", type="Row") + t = RowToNodeTransformer(cfg) + records = [{"name": "Alice", "role": "admin"}] + result = t.transform(records) + assert "name: Alice" in result[0]["content"] + + def test_infer_edges(self): + cfg = _make_config("ie", source_field="project", edge_type="WORKS_ON") + t = EdgeInferencer(cfg) + records = [ + {"id": "r1", "project": "A"}, + {"id": "r2", "project": "A"}, + {"id": "r3", "project": "B"}, + ] + result = t.transform(records) + edges = [r for r in result if r.get("edge_type") == "edge"] + assert len(edges) == 1 + assert edges[0]["source_id"] == "r1" + assert edges[0]["target_id"] == "r2" + assert edges[0]["relationship"] == "WORKS_ON" + + def test_infer_edges_no_shared_field(self): + cfg = _make_config("ie", source_field="project") + t = EdgeInferencer(cfg) + records = [{"id": "r1", "project": "A"}, {"id": "r2", "project": "B"}] + result = t.transform(records) + edges = [r for r in result if r.get("edge_type") == "edge"] + assert len(edges) == 0 + + +class TestFilterTransformers: + def test_row_filter_eq(self): + cfg = _make_config("rf", conditions=[{"field": "role", "operator": "eq", "value": "admin"}]) + t = RowFilterProvider(cfg) + records = [{"role": "admin"}, {"role": "user"}] + result = t.transform(records) + assert len(result) == 1 + assert result[0]["role"] == "admin" + + def test_row_filter_not_null(self): + cfg = _make_config("rf", conditions=[{"field": "email", "operator": "not_null"}]) + t = RowFilterProvider(cfg) + records = [{"email": "a@b.com"}, {"email": None}] + result = t.transform(records) + assert len(result) == 1 + + def test_row_filter_no_conditions_passes_all(self): + cfg = _make_config("rf", conditions=[]) + t = RowFilterProvider(cfg) + records = [{"a": 1}, {"a": 2}] + assert len(t.transform(records)) == 2 + + def test_row_filter_unsupported_op_raises(self): + cfg = _make_config("rf", conditions=[{"field": "x", "operator": "regex", "value": ".*"}]) + t = RowFilterProvider(cfg) + with pytest.raises(ValueError, match="Unsupported operator"): + t.transform([{"x": "hello"}]) + + +class TestTruncators: + def test_length_truncator(self): + cfg = _make_config("lt", max_length=5, fields=["text"]) + t = LengthTruncator(cfg) + result = t.transform({"text": "abcdefghij", "other": "keep"}) + assert result["text"] == "abcde" + assert result["other"] == "keep" + + def test_length_truncator_short_field_unchanged(self): + cfg = _make_config("lt", max_length=100, fields=["text"]) + t = LengthTruncator(cfg) + result = t.transform({"text": "short"}) + assert result["text"] == "short" + + def test_length_truncator_missing_field_ignored(self): + cfg = _make_config("lt", max_length=5, fields=["missing"]) + t = LengthTruncator(cfg) + result = t.transform({"text": "hello"}) + assert result["text"] == "hello" + + +# ─── Graph Build Edge Cases ────────────────────────────────────────────────── + +class TestCypherBuilderEdgeCases: + def test_node_special_chars_in_properties(self): + n = Node(id="n1", labels=["User"], properties={"bio": "It's a \"quote\" & "}) + query, params = node_to_cypher(n) + assert params["props"]["bio"] == "It's a \"quote\" & " + assert "MERGE" in query + + def test_node_none_value_in_properties(self): + n = Node(id="n1", labels=["User"], properties={"name": None}) + query, params = node_to_cypher(n) + assert params["props"]["name"] is None + + def test_node_empty_labels_uses_default(self): + n = Node(id="n1", labels=[], properties={"x": 1}) + query, params = node_to_cypher(n) + assert "Node" in query + + def test_node_multiple_labels(self): + n = Node(id="n1", labels=["Person", "Employee"], properties={}) + query, params = node_to_cypher(n) + assert "Person:Employee" in query + + def test_edge_with_properties(self): + e = Edge(id="e1", source_id="a", target_id="b", label="KNOWS", properties={"since": 2020}) + query, params = edge_to_cypher(e) + assert params["props"]["since"] == 2020 + + def test_edge_empty_properties(self): + e = Edge(id="e1", source_id="a", target_id="b", label="KNOWS") + query, params = edge_to_cypher(e) + assert params["props"] == {} + + def test_node_tenant_multi_label(self): + n = Node(id="n1", labels=["Person", "Admin"], properties={}) + query, params = node_to_cypher(n, tenant_id="t1") + assert "__Person__t1__" in query + assert "__Admin__t1__" in query + + +# ─── Schema Discovery Tests ───────────────────────────────────────────────── + +class TestCSVSchemaDiscovery: + def test_discover_from_temp_csv(self, tmp_csv): + discovery = CSVSchemaDiscoveryProvider(source=tmp_csv) + schema = discovery.discover_schema() + assert isinstance(schema, ETLSchema) + assert "id" in schema.load.document_node.fields + assert "name" in schema.load.document_node.fields + assert "email" in schema.load.document_node.fields + assert schema.schema_id == "discovered-test_data" + + def test_discover_from_project_csv(self): + csv_path = Path(__file__).parent.parent / "examples" / "cloud" / "notebooks" / "data" / "users.csv" + if not csv_path.exists(): + pytest.skip("users.csv not found in expected location") + discovery = CSVSchemaDiscoveryProvider(source=csv_path) + schema = discovery.discover_schema() + assert "id" in schema.load.document_node.fields + assert "name" in schema.load.document_node.fields + assert "email" in schema.load.document_node.fields + assert "role" in schema.load.document_node.fields + assert "account_id" in schema.load.document_node.fields + + def test_discover_file_not_found(self, tmp_path): + discovery = CSVSchemaDiscoveryProvider(source=tmp_path / "nonexistent.csv") + with pytest.raises(FileNotFoundError): + discovery.discover_schema() + + +# ─── TransformerProviderConfig Tests ───────────────────────────────────────── + +class TestTransformerProviderConfig: + def test_basic_construction(self): + cfg = TransformerProviderConfig(name="test", type="normalizer", args={"x": 1}) + assert cfg.name == "test" + assert cfg.type == "normalizer" + assert cfg.args["x"] == 1 + + def test_parameters_syncs_to_args(self): + cfg = TransformerProviderConfig(name="test", parameters={"y": 2}) + assert cfg.args["y"] == 2 diff --git a/examples/README.md b/examples/README.md index 8f284e90..f2d81955 100644 --- a/examples/README.md +++ b/examples/README.md @@ -3,4 +3,6 @@ - [BYOKG-RAG](./byokg-rag/) Example notebook and dataset demonstrating a RAG (Retrieval Augmented Generation) system built on top of a Knowledge Graph. - [Lexical Graph](./lexical-graph/) Examples of deploying and running the lexical-graph indexing and querying processes on AWS services - [Lexical Graph hybrid development](./lexical-graph-hybrid-dev/) Examples of running the indexing extract stage locally and the indexing build stage and querying on AWS services - - [Lexical Graph local development](./lexical-graph-local-dev/) Examples of running the lexical-graph indexing and querying processes locally \ No newline at end of file + - [Lexical Graph local development](./lexical-graph-local-dev/) Examples of running the lexical-graph indexing and querying processes locally + - [Document Graph](./document-graph/) Notebooks demonstrating structured data ingestion (CSV, JSON, Parquet, Excel) into Neptune with schema-driven pipelines, transformers, hybrid graph integration, and multi-tenancy + - [Code Property Graph](./codeproperty-graph/) Notebooks demonstrating Joern CPG parsing, delta comparison, and code analysis graph ingestion diff --git a/examples/codeproperty-graph/README.md b/examples/codeproperty-graph/README.md new file mode 100644 index 00000000..8ef83f34 --- /dev/null +++ b/examples/codeproperty-graph/README.md @@ -0,0 +1,45 @@ +# Code Property Graph Examples + +Example notebooks for `codeproperty-graph` — delta-aware code analysis ingestion into Neptune. + +## Prerequisites + +- Amazon Neptune (DB or Analytics) for production notebooks +- Python 3.10+ +- Joern (for generating real CPG exports; sample data provided for demos) + +## Setup + +```bash +pip install codeproperty-graph +``` + +For Neptune integration: +```bash +pip install codeproperty-graph graphrag-toolkit-document-graph[graphrag,neptune] +``` + +## Notebooks + +| Notebook | Description | Neptune Required? | +|----------|-------------|-------------------| +| `01-CPG-Models-and-GraphDiff.ipynb` | Core models, method signatures, and diff comparison | No (local only) | +| `02-Delta-Ingestion-Pipeline.ipynb` | Full pipeline: manifest → diff → skip/ingest decision | Optional (has mock mode) | + +## Sample Data + +`data/sample_nodes.json` — 8 CPG nodes (3 METHOD, 2 CALL, 2 IDENTIFIER, 1 LITERAL) +`data/sample_edges.json` — 10 edges (AST, CFG, CALL, ARGUMENT) + +Represents a minimal Java project with `Main.main()`, `Parser.parse()`, and `Validator.validate()`. + +## Generating Real CPG Data + +```bash +# Install Joern: https://joern.io +joern-export --repr cpg14 --format json --out cpg-export/ /path/to/source + +# Output: +# cpg-export/nodes.json +# cpg-export/edges.json +``` diff --git a/examples/codeproperty-graph/notebooks/01-CPG-Models-and-GraphDiff.ipynb b/examples/codeproperty-graph/notebooks/01-CPG-Models-and-GraphDiff.ipynb new file mode 100644 index 00000000..04a4c5bf --- /dev/null +++ b/examples/codeproperty-graph/notebooks/01-CPG-Models-and-GraphDiff.ipynb @@ -0,0 +1,309 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# CPG Schema, Models, and Graph Diff\n", + "\n", + "This notebook demonstrates:\n", + "- The full Joern CPG schema (20 node types, 14+ edge types, 9 languages)\n", + "- Parsing Joern exports with `CPGNode.from_joern()` and `CPGEdge.from_joern()`\n", + "- Graph analysis: type distribution, call graph, structure\n", + "- Change detection with `GraphDiff`\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-codeproperty-graph`\n", + "- No Neptune or AWS credentials required\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. The CPG Schema\n", + "\n", + "Explore the complete Joern type system.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.codeproperty_graph import (\n", + " NodeType, EdgeType, DELTA_RELEVANT_TYPES, SUPPORTED_LANGUAGES,\n", + " joern_export_command\n", + ")\n", + "\n", + "print(f'Node types ({len(NodeType)}):')\n", + "for nt in NodeType:\n", + " marker = ' \u2605' if nt in DELTA_RELEVANT_TYPES else ''\n", + " print(f' {nt.value:20s}{marker}')\n", + "\n", + "print(f'\\nEdge types ({len(EdgeType)}):')\n", + "for et in EdgeType:\n", + " print(f' {et.value}')\n", + "\n", + "print(f'\\nSupported languages ({len(SUPPORTED_LANGUAGES)}):')\n", + "for lang, desc in SUPPORTED_LANGUAGES.items():\n", + " print(f' {lang:12s} {desc}')\n", + "\n", + "print(f'\\n\u2605 = participates in delta comparison')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Joern Export Command\n", + "\n", + "Generate the CLI command to produce a CPG from source code.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Generate export commands for different languages\n", + "for lang in ['java', 'python', 'javascript']:\n", + " cmd = joern_export_command('/src/my-service', output_dir='cpg-out', language=lang)\n", + " print(f'{lang:12s} {cmd}')\n", + "\n", + "# Auto-detect language\n", + "print(f'{\"auto\":12s} {joern_export_command(\"/src/my-service\")}')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Parse Joern Export with `from_joern()`\n", + "\n", + "Load raw Joern JSON and parse into typed CPGNode/CPGEdge objects.\n", + "The factory methods handle both UPPER_CASE and camelCase property formats.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "from graphrag_toolkit.codeproperty_graph import CPGNode, CPGEdge\n", + "\n", + "with open('data/sample_nodes.json') as f:\n", + " raw_nodes = json.load(f)\n", + "with open('data/sample_edges.json') as f:\n", + " raw_edges = json.load(f)\n", + "\n", + "# Parse with from_joern()\n", + "nodes = [CPGNode.from_joern(raw) for raw in raw_nodes]\n", + "edges = [CPGEdge.from_joern(raw) for raw in raw_edges]\n", + "\n", + "print(f'Parsed {len(nodes)} nodes, {len(edges)} edges')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Graph Analysis: Type Distribution\n", + "\n", + "See what node and edge types are in the graph.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from collections import Counter\n", + "\n", + "# Node type distribution\n", + "node_types = Counter(n.node_type for n in nodes)\n", + "print('Node types:')\n", + "for ntype, count in node_types.most_common():\n", + " print(f' {ntype:20s} {count:3d}')\n", + "\n", + "# Edge type distribution\n", + "edge_types = Counter(e.edge_type for e in edges)\n", + "print(f'\\nEdge types:')\n", + "for etype, count in edge_types.most_common():\n", + " print(f' {etype:20s} {count:3d}')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Explore Methods (Delta-Relevant Nodes)\n", + "\n", + "Only METHOD nodes participate in change detection.\n", + "Their `full_name \u2192 hash` mapping is the fingerprint for delta comparison.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "methods = [n for n in nodes if n.node_type == NodeType.METHOD]\n", + "\n", + "print(f'Methods ({len(methods)}):')\n", + "for m in methods:\n", + " print(f' {m.full_name}')\n", + " print(f' signature: {m.signature}')\n", + " print(f' hash: {m.hash}')\n", + " print(f' file: {m.filename}:{m.line_number}-{m.line_number_end}')\n", + " print(f' external: {m.is_external}')\n", + " print()\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Explore Structure (Classes, Files, Namespaces)\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.codeproperty_graph.schema import STRUCTURE_TYPES\n", + "\n", + "structure_nodes = [n for n in nodes if n.node_type in STRUCTURE_TYPES]\n", + "print(f'Structure nodes ({len(structure_nodes)}):')\n", + "for n in structure_nodes:\n", + " print(f' {n.node_type:16s} {n.full_name or n.name}')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Call Graph\n", + "\n", + "Extract the CALL edges to see which methods call which.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "call_edges = [e for e in edges if e.edge_type == EdgeType.CALL]\n", + "node_map = {n.id: n for n in nodes}\n", + "\n", + "print(f'Call graph ({len(call_edges)} edges):')\n", + "for e in call_edges:\n", + " src = node_map.get(e.source_id)\n", + " tgt = node_map.get(e.target_id)\n", + " src_name = src.full_name if src else e.source_id\n", + " tgt_name = tgt.full_name if tgt else e.target_id\n", + " print(f' {src_name} \u2192 {tgt_name}')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 8. Delta Comparison (GraphDiff)\n", + "\n", + "Compare method signatures between two CPG snapshots.\n", + "This is how `DeltaIngestor` decides whether to skip or ingest.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.codeproperty_graph import GraphDiff\n", + "\n", + "# Current state\n", + "current_sigs = {m.full_name: m.hash for m in methods}\n", + "print(f'Current signatures:')\n", + "for name, h in current_sigs.items():\n", + " print(f' {name}: {h}')\n", + "\n", + "# Simulate next commit: parse() refactored, validate() deleted, check() added\n", + "next_sigs = {\n", + " 'com.example.Main.main': 'a1b2c3d4e5', # unchanged\n", + " 'com.example.Parser.parse': 'NEWH4SH999', # MODIFIED\n", + " 'com.example.Checker.check': 'p6q7r8s9t0', # ADDED\n", + "}\n", + "\n", + "diff = GraphDiff.compare(current_sigs, next_sigs)\n", + "print(f'\\nDelta: {diff.summary}')\n", + "print(f' Added: {list(diff.added.keys())}')\n", + "print(f' Removed: {list(diff.removed.keys())}')\n", + "print(f' Modified: {list(diff.modified.keys())}')\n", + "print(f' Unchanged: {diff.unchanged}')\n", + "print(f'\\n\u2192 DeltaIngestor would: {\"INGEST\" if diff.has_changes else \"SKIP\"}')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 9. No-Change Case (SKIP)\n", + "\n", + "When code hasn't changed, `GraphDiff` returns no delta \u2014 zero Neptune writes.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "no_diff = GraphDiff.compare(current_sigs, current_sigs)\n", + "print(f'Delta: {no_diff.summary}')\n", + "print(f'Has changes: {no_diff.has_changes}')\n", + "print(f'\\n\u2192 DeltaIngestor would: SKIP')\n", + "print(f' Zero Neptune writes. Zero cost. <1 second.')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "| Feature | What It Does |\n", + "|---------|-------------|\n", + "| `NodeType` / `EdgeType` enums | Complete Joern schema (20 + 14 types) |\n", + "| `CPGNode.from_joern(raw)` | Parse Joern JSON export \u2192 typed objects |\n", + "| `DELTA_RELEVANT_TYPES` | Which types matter for change detection |\n", + "| `STRUCTURE_TYPES` | Types for architectural analysis |\n", + "| `GraphDiff.compare()` | Method signature delta \u2192 skip or ingest |\n", + "| `joern_export_command()` | Generate the CLI command for any language |\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/examples/codeproperty-graph/notebooks/02-Delta-Ingestion-Pipeline.ipynb b/examples/codeproperty-graph/notebooks/02-Delta-Ingestion-Pipeline.ipynb new file mode 100644 index 00000000..e26300e5 --- /dev/null +++ b/examples/codeproperty-graph/notebooks/02-Delta-Ingestion-Pipeline.ipynb @@ -0,0 +1,253 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 02 \u2014 Delta Ingestion Pipeline\n", + "\n", + "This notebook demonstrates the full DeltaIngestor pipeline: manifest comparison, skip-or-ingest decision, and Neptune graph write.\n", + "\n", + "**What you'll learn:**\n", + "- How ManifestManager tracks CPG state in S3\n", + "- How DeltaIngestor decides to skip or ingest\n", + "- How tenant-based graph replacement works\n", + "\n", + "**Prerequisites:**\n", + "- Neptune Analytics or Neptune Database endpoint\n", + "- S3 bucket for manifests\n", + "- IAM permissions for Neptune and S3" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# %pip install codeproperty-graph graphrag-toolkit-document-graph[graphrag]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "import asyncio\n", + "from codeproperty_graph import DeltaIngestor, ManifestManager, GraphDiff, CPGNode" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configuration" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Replace with your values:\n", + "BUCKET = \"my-graphrag-artifacts\" # S3 bucket for manifests\n", + "REPO = \"example-service\" # Repository name\n", + "JOB_ID = \"build-001\" # CI/CD build identifier\n", + "NEPTUNE_ENDPOINT = \"neptune-analytics://my-graph-id\" # or neptune-db://...\n", + "\n", + "# Tenant ID (auto-generated in production, fixed here for demo)\n", + "TENANT_ID = \"demo123\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Load Joern Export" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "from graphrag_toolkit.codeproperty_graph import CPGNode, CPGEdge, NodeType\n", + "\n", + "with open('data/sample_nodes.json') as f:\n", + " nodes = [CPGNode.from_joern(raw) for raw in json.load(f)]\n", + "\n", + "with open('data/sample_edges.json') as f:\n", + " edges = [CPGEdge.from_joern(raw) for raw in json.load(f)]\n", + "\n", + "print(f'Parsed {len(nodes)} nodes ({len(set(n.node_type for n in nodes))} types)')\n", + "print(f'Parsed {len(edges)} edges ({len(set(e.edge_type for e in edges))} types)')\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Extract Method Signatures\n", + "\n", + "Only METHOD nodes participate in change detection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Only METHOD nodes participate in delta comparison\n", + "methods = [n for n in nodes if n.node_type == NodeType.METHOD]\n", + "current_sigs = {m.full_name: m.hash for m in methods}\n", + "\n", + "print(f'Extracted {len(current_sigs)} method signatures (from {len(nodes)} total nodes):')\n", + "for name, h in current_sigs.items():\n", + " print(f' {name}: {h}')\n", + "\n", + "print(f'\\nOther node types (not in delta): {set(n.node_type for n in nodes if n.node_type != NodeType.METHOD)}')\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Check Previous Manifest\n", + "\n", + "ManifestManager retrieves the last known state from S3." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "manifest_mgr = ManifestManager(bucket=BUCKET)\n", + "\n", + "# In production this reads from S3. For demo, simulate:\n", + "# previous_manifest = manifest_mgr.get(REPO)\n", + "\n", + "# Simulate: first run (no previous manifest)\n", + "previous_manifest = None\n", + "\n", + "if previous_manifest is None:\n", + " print(\"No previous manifest \u2014 this is a fresh ingestion.\")\n", + " print(\"All methods will be treated as ADDED.\")\n", + "else:\n", + " print(f\"Previous manifest: tenant={previous_manifest.tenant_id}, job={previous_manifest.job_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Compute Diff" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "previous_sigs = previous_manifest.method_signatures if previous_manifest else {}\n", + "\n", + "diff = GraphDiff.compare(previous_sigs, current_sigs)\n", + "print(f\"Diff: {diff.summary}\")\n", + "print(f\"Has changes: {diff.has_changes}\")\n", + "\n", + "if diff.has_changes:\n", + " print(\"\\n\u2192 Will INGEST to Neptune\")\n", + "else:\n", + " print(\"\\n\u2192 Will SKIP (no Neptune writes)\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Write to Neptune (if changed)\n", + "\n", + "In production, this uses the GraphStore. For this demo, we simulate the write." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Simulated write function (replace with actual Neptune write in production)\n", + "write_log = []\n", + "\n", + "async def mock_write_fn(query: str, params: dict):\n", + " \"\"\"Simulates a Neptune write \u2014 logs instead of executing.\"\"\"\n", + " write_log.append({'query': query[:80], 'param_count': len(params)})\n", + "\n", + "if diff.has_changes:\n", + " print(f\"Writing {len(nodes_data)} nodes and {len(edges_data)} edges...\")\n", + " print(f\"Tenant: {TENANT_ID}\")\n", + " print(f\"(Using mock write \u2014 replace with graph_store.execute_query in production)\")\n", + " \n", + " # In production:\n", + " # result = await ingestor.ingest(repo=REPO, job_id=JOB_ID, ...)\n", + " \n", + " # Simulated result:\n", + " result = {\n", + " \"status\": \"INGESTED\",\n", + " \"delta\": diff.summary,\n", + " \"tenant_id\": TENANT_ID,\n", + " \"nodes_written\": len(nodes_data),\n", + " \"edges_written\": len(edges_data),\n", + " }\n", + " print(f\"\\nResult: {json.dumps(result, indent=2)}\")\n", + "else:\n", + " result = {\"status\": \"SKIPPED\", \"reason\": \"no changes detected\"}\n", + " print(f\"Result: {json.dumps(result, indent=2)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Simulate Second Run (No Changes)\n", + "\n", + "On the next CI/CD run with no code changes, the ingestor skips entirely." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Same signatures as before \u2014 simulates \"no code change\" commit\n", + "second_run_sigs = current_sigs.copy()\n", + "\n", + "diff2 = GraphDiff.compare(current_sigs, second_run_sigs)\n", + "print(f\"Second run diff: {diff2.summary}\")\n", + "print(f\"Has changes: {diff2.has_changes}\")\n", + "print()\n", + "print(\"\u2192 SKIPPED \u2014 zero Neptune writes, zero cost, <1 second\")\n", + "print(\" This is the common case in CI/CD (most commits don't change method bodies)\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "| Scenario | Action | Neptune Cost | Time |\n", + "|----------|--------|-------------|------|\n", + "| First run (no manifest) | Full ingest | N nodes + M edges | Seconds |\n", + "| Code changed | Full replace (new tenant, purge old) | N nodes + M edges + purge | Seconds |\n", + "| No code change | SKIP | Zero | <1s |\n", + "\n", + "The delta logic ensures you only pay for Neptune writes when code actually changes." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/examples/codeproperty-graph/notebooks/data/sample_edges.json b/examples/codeproperty-graph/notebooks/data/sample_edges.json new file mode 100644 index 00000000..d9c6605e --- /dev/null +++ b/examples/codeproperty-graph/notebooks/data/sample_edges.json @@ -0,0 +1,25 @@ +[ + {"src": "5001", "dst": "4001", "label": "CONTAINS", "properties": {}}, + {"src": "5002", "dst": "4002", "label": "CONTAINS", "properties": {}}, + {"src": "5003", "dst": "4003", "label": "CONTAINS", "properties": {}}, + {"src": "4001", "dst": "1001", "label": "CONTAINS", "properties": {}}, + {"src": "4002", "dst": "1002", "label": "CONTAINS", "properties": {}}, + {"src": "4003", "dst": "1003", "label": "CONTAINS", "properties": {}}, + {"src": "1001", "dst": "2001", "label": "AST", "properties": {"ORDER": 1}}, + {"src": "1001", "dst": "2002", "label": "AST", "properties": {"ORDER": 2}}, + {"src": "1001", "dst": "3001", "label": "AST", "properties": {"ORDER": 3}}, + {"src": "2001", "dst": "3001", "label": "ARGUMENT", "properties": {"ORDER": 1}}, + {"src": "2001", "dst": "3002", "label": "ARGUMENT", "properties": {"ORDER": 2}}, + {"src": "2002", "dst": "3003", "label": "ARGUMENT", "properties": {"ORDER": 1}}, + {"src": "1001", "dst": "1002", "label": "CALL", "properties": {}}, + {"src": "1001", "dst": "1003", "label": "CALL", "properties": {}}, + {"src": "1001", "dst": "1002", "label": "CFG", "properties": {}}, + {"src": "1002", "dst": "1003", "label": "CFG", "properties": {}}, + {"src": "1001", "dst": "6001", "label": "AST", "properties": {"ORDER": 0}}, + {"src": "1002", "dst": "6002", "label": "AST", "properties": {"ORDER": 0}}, + {"src": "1003", "dst": "6003", "label": "AST", "properties": {"ORDER": 0}}, + {"src": "3001", "dst": "6001", "label": "REF", "properties": {}}, + {"src": "7001", "dst": "4001", "label": "CONTAINS", "properties": {}}, + {"src": "7001", "dst": "4002", "label": "CONTAINS", "properties": {}}, + {"src": "7001", "dst": "4003", "label": "CONTAINS", "properties": {}} +] diff --git a/examples/codeproperty-graph/notebooks/data/sample_nodes.json b/examples/codeproperty-graph/notebooks/data/sample_nodes.json new file mode 100644 index 00000000..e766ee32 --- /dev/null +++ b/examples/codeproperty-graph/notebooks/data/sample_nodes.json @@ -0,0 +1,21 @@ +[ + {"id": "1001", "label": "METHOD", "properties": {"FULL_NAME": "com.example.Main.main", "SIGNATURE": "void(String[])", "HASH": "a1b2c3d4e5", "FILENAME": "src/Main.java", "LINE_NUMBER": 5, "LINE_NUMBER_END": 12, "CODE": "public static void main(String[] args) { Parser.parse(args[0]); }", "IS_EXTERNAL": false}}, + {"id": "1002", "label": "METHOD", "properties": {"FULL_NAME": "com.example.Parser.parse", "SIGNATURE": "Document(String)", "HASH": "f6g7h8i9j0", "FILENAME": "src/Parser.java", "LINE_NUMBER": 12, "LINE_NUMBER_END": 18, "CODE": "public Document parse(String input) { return new Document(input); }", "IS_EXTERNAL": false}}, + {"id": "1003", "label": "METHOD", "properties": {"FULL_NAME": "com.example.Validator.validate", "SIGNATURE": "boolean(Document)", "HASH": "k1l2m3n4o5", "FILENAME": "src/Validator.java", "LINE_NUMBER": 8, "LINE_NUMBER_END": 11, "CODE": "public boolean validate(Document doc) { return doc != null; }", "IS_EXTERNAL": false}}, + {"id": "2001", "label": "CALL", "properties": {"NAME": "parse", "FULL_NAME": "com.example.Parser.parse", "CODE": "Parser.parse(args[0])", "LINE_NUMBER": 6, "DISPATCH_TYPE": "STATIC_DISPATCH"}}, + {"id": "2002", "label": "CALL", "properties": {"NAME": "validate", "FULL_NAME": "com.example.Validator.validate", "CODE": "validator.validate(doc)", "LINE_NUMBER": 7, "DISPATCH_TYPE": "DYNAMIC_DISPATCH"}}, + {"id": "3001", "label": "IDENTIFIER", "properties": {"NAME": "args", "CODE": "args", "LINE_NUMBER": 6, "ORDER": 1}}, + {"id": "3002", "label": "LITERAL", "properties": {"CODE": "0", "LINE_NUMBER": 6, "ORDER": 2}}, + {"id": "3003", "label": "IDENTIFIER", "properties": {"NAME": "doc", "CODE": "doc", "LINE_NUMBER": 7, "ORDER": 1}}, + {"id": "4001", "label": "TYPE_DECL", "properties": {"FULL_NAME": "com.example.Main", "NAME": "Main", "IS_EXTERNAL": false, "ORDER": 1}}, + {"id": "4002", "label": "TYPE_DECL", "properties": {"FULL_NAME": "com.example.Parser", "NAME": "Parser", "IS_EXTERNAL": false, "ORDER": 2}}, + {"id": "4003", "label": "TYPE_DECL", "properties": {"FULL_NAME": "com.example.Validator", "NAME": "Validator", "IS_EXTERNAL": false, "ORDER": 3}}, + {"id": "5001", "label": "FILE", "properties": {"NAME": "src/Main.java", "ORDER": 1}}, + {"id": "5002", "label": "FILE", "properties": {"NAME": "src/Parser.java", "ORDER": 2}}, + {"id": "5003", "label": "FILE", "properties": {"NAME": "src/Validator.java", "ORDER": 3}}, + {"id": "6001", "label": "PARAMETER", "properties": {"NAME": "args", "CODE": "String[] args", "LINE_NUMBER": 5, "ORDER": 1, "TYPE_FULL_NAME": "java.lang.String[]"}}, + {"id": "6002", "label": "PARAMETER", "properties": {"NAME": "input", "CODE": "String input", "LINE_NUMBER": 12, "ORDER": 1, "TYPE_FULL_NAME": "java.lang.String"}}, + {"id": "6003", "label": "PARAMETER", "properties": {"NAME": "doc", "CODE": "Document doc", "LINE_NUMBER": 8, "ORDER": 1, "TYPE_FULL_NAME": "com.example.Document"}}, + {"id": "7001", "label": "NAMESPACE", "properties": {"NAME": "com.example", "ORDER": 1}}, + {"id": "8001", "label": "META_DATA", "properties": {"LANGUAGE": "java", "VERSION": "1.0"}} +] diff --git a/examples/document-graph/README.md b/examples/document-graph/README.md new file mode 100644 index 00000000..b21b8910 --- /dev/null +++ b/examples/document-graph/README.md @@ -0,0 +1,72 @@ +# Document Graph Examples + +Example notebooks for `graphrag-toolkit-document-graph` — structured data ingestion into Neptune. + +## Prerequisites + +- Python 3.10+ +- `pip install graphrag-toolkit-document-graph` (base notebooks) +- `pip install graphrag-toolkit-document-graph[graphrag]` (Neptune notebooks) +- Amazon Neptune cluster (for Neptune notebooks) +- Amazon OpenSearch Serverless (for hybrid notebooks) + +## Notebooks + +### Getting Started + +| # | Notebook | Neptune? | Description | +|---|----------|----------|-------------| +| 01 | Setup | Yes | Install packages, verify Neptune connection | +| 02 | Standalone ETL | **No** | Transform data + generate Cypher locally (no cloud needed) | +| 03 | Combined Extract-Load-Build | Yes | End-to-end: CSV → transform → Neptune | + +### Pipeline Testing + +| # | Notebook | Neptune? | Description | +|---|----------|----------|-------------| +| 04 | Full Pipeline Test | **No** | Test all pipeline stages locally (ingest → transform → construct) | +| 05a | Ingestors | **No** | Column select/rename/reorder, row filter, type conversion | +| 05b | Schema-Driven Pipeline | Yes | CSV → schema → full pipeline → Neptune | +| 05c | JSON-Driven Pipeline | Yes | JSON schema definition → automatic graph build | + +### Transformers & Constructors + +| # | Notebook | Neptune? | Description | +|---|----------|----------|-------------| +| 05d | Multi-Format Extraction | **No** | Parquet, Excel, XML, YAML extraction + schema discovery | +| 06a | Field Transformers | **No** | JSON flatten, UUID gen, regex, split, timestamp — 12 providers | +| 06b | Normalizers | **No** | Whitespace, nulls, case, enum, spelling, timestamp — 6 providers | +| 06c | Constructors | **No** | Node, edge, schema-driven, batch, dedup, 1:N, N:M — 8 patterns | +| 06d | Document Transformers | **No** | JSON→rows, text chunker, PII redactor | +| 06e | LLM Enrichers | Bedrock | AI-powered classification, entity extraction, language detection | +| 07 | Graph Write and Read | Yes | Direct Neptune read/write with batch UNWIND | +| 07b | Batch Performance | Optional | Benchmark: UNWIND vs individual MERGE (10-30x speedup) | +| 08 | Schema Providers | **No** | CSV, JSON, Parquet, Excel, S3, Glue, Static — discovery + validation | +| 09 | Transformer Deep Dive | **No** | Comprehensive guide to all 6 transformer categories | + +### Hybrid Graph (document-graph + lexical-graph) + +| # | Notebook | Neptune? | Description | +|---|----------|----------|-------------| +| 10a | Hybrid Data Processing | Yes | Structured data → Neptune → LlamaIndex Documents (with lineage) | +| 10b | Hybrid Lexical Indexing | Yes | Documents → lexical-graph → semantic search → lineage correlation | +| 11a | Mandela Hybrid Build | Yes | Real-world example: Wikipedia bio + structured events/orgs | +| 11b | Mandela Hybrid Query | Yes | Semantic + structured + hybrid correlation queries | +| 12 | Multi-Tenant Coexistence | Yes | Prove tenant isolation with shared Neptune cluster | + +### Maintenance + +| # | Notebook | Neptune? | Description | +|---|----------|----------|-------------| +| 99 | Cleanup | Yes | Remove old test tenants from Neptune | + +## Data + +Sample data in `data/` and `mandela_data/` for running examples without external dependencies. + +## Running Locally (No Neptune) + +Notebooks marked **No** in the Neptune column run entirely locally — no AWS credentials, +no Neptune cluster, no OpenSearch. They test pipeline logic, transformations, and Cypher generation. + +Start with `02-Standalone-ETL.ipynb` for the zero-dependency introduction. diff --git a/examples/document-graph/install.sh b/examples/document-graph/install.sh new file mode 100755 index 00000000..c0662800 --- /dev/null +++ b/examples/document-graph/install.sh @@ -0,0 +1,10 @@ +#!/bin/bash +# Lifecycle config — must finish in <5 min. Do absolute minimum. +set -e + +# Configure these for your environment: +AWS_ACCOUNT_ID="${AWS_ACCOUNT_ID:-}" +ARTIFACTS_BUCKET="${ARTIFACTS_BUCKET:-graphrag-artifacts-${AWS_ACCOUNT_ID}}" + +aws s3 sync "s3://${ARTIFACTS_BUCKET}/document-graph-notebooks/" ~/SageMaker/document-graph/ --quiet --exclude "*.ipynb_checkpoints/*" +echo "done" > ~/SageMaker/document-graph/.ready diff --git a/examples/document-graph/notebooks/01-Setup.ipynb b/examples/document-graph/notebooks/01-Setup.ipynb new file mode 100644 index 00000000..3d60b7d3 --- /dev/null +++ b/examples/document-graph/notebooks/01-Setup.ipynb @@ -0,0 +1,174 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Setup" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Install graphrag-toolkit-document-graph:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!pip install graphrag-toolkit-document-graph" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Graph Store" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Amazon Neptune Database\n", + "\n", + "If you are using Amazon Neptune Database as a graph store, install the lexical-graph dependency (provides `GraphStoreFactory`):" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!pip install graphrag-toolkit-document-graph[graphrag]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Neo4j\n", + "\n", + "If you are using Neo4j as a graph store:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!pip install graphrag-toolkit-document-graph[neo4j]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Hybrid Graph (document-graph + lexical-graph)\n", + "\n", + "If you plan to use document-graph alongside lexical-graph for hybrid structured + unstructured graphs, install both:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!pip install graphrag-toolkit-lexical-graph\n", + "!pip install graphrag-toolkit-document-graph[graphrag]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify Installation" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph import Node, Edge\n", + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import node_to_cypher, edge_to_cypher\n", + "from graphrag_toolkit.document_graph.query import DocumentGraphQueryEngine\n", + "\n", + "print('graphrag-toolkit-document-graph ✅')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify Neptune Connection\n", + "\n", + "Replace the endpoint below with your Neptune cluster endpoint. You can find this in the Neptune console under your cluster's **Connectivity & security** tab." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory\n", + "\n", + "GRAPH_STORE = os.environ.get(\n", + " 'GRAPH_STORE',\n", + " 'neptune-db://:8182'\n", + ")\n", + "\n", + "graph_store = GraphStoreFactory.for_graph_store(GRAPH_STORE)\n", + "gs = graph_store.__enter__()\n", + "\n", + "result = gs.execute_query('MATCH (n) RETURN count(n) as cnt LIMIT 1')\n", + "print(f'Neptune connected ✅ — node count: {result[0][\"cnt\"]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Write a Test Node\n", + "\n", + "Write a single test node to verify end-to-end write and read." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "node = Node(\n", + " id='setup-test-001',\n", + " labels=['TestNode'],\n", + " properties={'source': 'setup-notebook', 'status': 'ok'}\n", + ")\n", + "\n", + "cypher, params = node_to_cypher(node, tenant_id='test')\n", + "print(f'Cypher: {cypher}')\n", + "\n", + "gs.execute_query(cypher, params)\n", + "print('Write ✅')\n", + "\n", + "# Read back\n", + "engine = DocumentGraphQueryEngine(gs, tenant_id='test')\n", + "result = engine.find_by_property('TestNode', 'source', 'setup-notebook')\n", + "print(f'Read ✅ — {result}')\n", + "\n", + "# Clean up context manager\n", + "graph_store.__exit__(None, None, None)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/document-graph/notebooks/02-Standalone-ETL.ipynb b/examples/document-graph/notebooks/02-Standalone-ETL.ipynb new file mode 100644 index 00000000..8681ecd7 --- /dev/null +++ b/examples/document-graph/notebooks/02-Standalone-ETL.ipynb @@ -0,0 +1,267 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Standalone ETL — Transform and Generate Cypher (No Neptune Required)\n", + "\n", + "This notebook demonstrates document-graph's **standalone ETL path** — transforming structured data into typed graph nodes and generating openCypher statements without requiring a Neptune connection or lexical-graph.\n", + "\n", + "Use this when you want to:\n", + "- Preview the graph that would be created before writing to any database\n", + "- Generate Cypher for execution in a separate pipeline (e.g., batch job, CI/CD)\n", + "- Validate transformations locally without cloud dependencies\n", + "- Export Cypher statements for use with Neo4j, FalkorDB, or other openCypher databases\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph` (base install, no extras needed)\n", + "- No Neptune endpoint required\n", + "- No AWS credentials required" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Extract — Read Source Data" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import csv\n", + "\n", + "CSV_PATH = 'data/users.csv'\n", + "\n", + "with open(CSV_PATH) as f:\n", + " records = list(csv.DictReader(f))\n", + "\n", + "print(f'Read {len(records)} records from {CSV_PATH}')\n", + "print(f'Columns: {list(records[0].keys())}')\n", + "print(f'\\nSample record:')\n", + "records[0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Transform — Rows to Typed Nodes\n", + "\n", + "Convert each row into a `Node` object with a label, ID, and properties.\n", + "No database connection needed — this is pure data transformation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph import Node, Edge\n", + "\n", + "nodes = []\n", + "for record in records:\n", + " node = Node(\n", + " id=record.get('id', record.get('user_id', '')),\n", + " labels=['User'],\n", + " properties={k: v for k, v in record.items() if k != 'id'}\n", + " )\n", + " nodes.append(node)\n", + "\n", + "print(f'Created {len(nodes)} Node objects')\n", + "print(f'\\nFirst node:')\n", + "print(f' id: {nodes[0].id}')\n", + "print(f' labels: {nodes[0].labels}')\n", + "print(f' properties: {nodes[0].properties}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Generate Cypher — Single Nodes\n", + "\n", + "Generate openCypher MERGE statements for each node. These can be executed against\n", + "any openCypher-compatible database (Neptune, Neo4j, FalkorDB, Memgraph)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import node_to_cypher\n", + "\n", + "# Without tenant (default — labels are backtick-quoted as-is)\n", + "print('=== Default tenant (no scoping) ===\\n')\n", + "for node in nodes[:2]:\n", + " cypher, params = node_to_cypher(node)\n", + " print(f'Cypher: {cypher}')\n", + " print(f'Params: {params}')\n", + " print()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Generate Cypher — With Tenant Scoping\n", + "\n", + "When you specify a `tenant_id`, labels are formatted to match lexical-graph's\n", + "multi-tenancy convention: `` `Label{tenant_id}__` ``. This ensures document-graph\n", + "nodes coexist cleanly with lexical-graph nodes in the same Neptune cluster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "TENANT = 'demo'\n", + "\n", + "print(f'=== Tenant \"{TENANT}\" (scoped labels) ===\\n')\n", + "for node in nodes[:2]:\n", + " cypher, params = node_to_cypher(node, tenant_id=TENANT)\n", + " print(f'Cypher: {cypher}')\n", + " print(f'Params: {params}')\n", + " print()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Generate Cypher — Batched UNWIND\n", + "\n", + "For bulk writes, `batch_nodes_unwind` generates a single UNWIND query that writes\n", + "all nodes in one round-trip. This is significantly faster than individual MERGE\n", + "statements for large datasets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import batch_nodes_unwind\n", + "\n", + "cypher, params = batch_nodes_unwind(nodes, tenant_id=TENANT)\n", + "\n", + "print(f'Batch query ({len(nodes)} nodes in one statement):\\n')\n", + "print(f'Cypher: {cypher}')\n", + "print(f'\\nBatch size: {len(params[\"batch\"])} rows')\n", + "print(f'First row: {params[\"batch\"][0]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 6: Generate Cypher — Edges\n", + "\n", + "Create edges between nodes and generate the corresponding Cypher." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import edge_to_cypher, batch_edges_unwind\n", + "\n", + "# Create some edges\n", + "edges = []\n", + "if len(nodes) >= 2:\n", + " edges.append(Edge(\n", + " id='e1',\n", + " source_id=nodes[0].id,\n", + " target_id=nodes[1].id,\n", + " label='KNOWS',\n", + " properties={'since': '2024'}\n", + " ))\n", + "\n", + "print('=== Individual edge ===' )\n", + "for edge in edges:\n", + " cypher, params = edge_to_cypher(edge)\n", + " print(f'Cypher: {cypher}')\n", + " print(f'Params: {params}')\n", + " print()\n", + "\n", + "if edges:\n", + " print('=== Batched edges ===')\n", + " cypher, params = batch_edges_unwind(edges)\n", + " print(f'Cypher: {cypher}')\n", + " print(f'Batch: {params[\"batch\"]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 7: Export Cypher to File\n", + "\n", + "Write all generated Cypher statements to a file for execution elsewhere —\n", + "useful for batch jobs, code review, or loading into a different graph database." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "output_path = 'data/generated_cypher.txt'\n", + "\n", + "with open(output_path, 'w') as f:\n", + " f.write(f'// Generated by document-graph standalone ETL\\n')\n", + " f.write(f'// Tenant: {TENANT}\\n')\n", + " f.write(f'// Nodes: {len(nodes)}, Edges: {len(edges)}\\n\\n')\n", + "\n", + " for node in nodes:\n", + " cypher, params = node_to_cypher(node, tenant_id=TENANT)\n", + " # Inline params for standalone execution\n", + " stmt = cypher.replace('$id_val', repr(params['id_val']))\n", + " f.write(f'{stmt};\\n')\n", + "\n", + " for edge in edges:\n", + " cypher, params = edge_to_cypher(edge)\n", + " stmt = cypher.replace('$src_id', repr(params['src_id'])).replace('$tgt_id', repr(params['tgt_id']))\n", + " f.write(f'{stmt};\\n')\n", + "\n", + "print(f'✅ Exported {len(nodes)} node + {len(edges)} edge statements to {output_path}')\n", + "print(f'\\nPreview:')\n", + "with open(output_path) as f:\n", + " for line in f.readlines()[:8]:\n", + " print(f' {line.rstrip()}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "This notebook used **only** `graphrag-toolkit-document-graph` (base install) — no Neptune, no lexical-graph, no AWS credentials.\n", + "\n", + "| What we did | Dependency required |\n", + "|-------------|--------------------|\n", + "| Read CSV | None (stdlib) |\n", + "| Create Node/Edge objects | `graphrag-toolkit-document-graph` |\n", + "| Generate Cypher (single + batch) | `graphrag-toolkit-document-graph` |\n", + "| Tenant-scoped labels | `graphrag-toolkit-document-graph` (falls back to standalone format) |\n", + "| Export to file | None (stdlib) |\n", + "\n", + "To actually **execute** these statements against Neptune, see `01-Combined-Extract-Load-Build.ipynb`\n", + "(requires `[graphrag]` extra for Neptune connectivity)." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/document-graph/notebooks/03-Combined-Extract-Load-Build.ipynb b/examples/document-graph/notebooks/03-Combined-Extract-Load-Build.ipynb new file mode 100644 index 00000000..f6893379 --- /dev/null +++ b/examples/document-graph/notebooks/03-Combined-Extract-Load-Build.ipynb @@ -0,0 +1,206 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Combined Extract-Load-Build\n", + "\n", + "This notebook demonstrates the end-to-end document-graph pipeline:\n", + "\n", + "1. **Extract** \u2014 Read structured data from CSV\n", + "2. **Transform** \u2014 Convert rows to typed graph nodes and infer edges\n", + "3. **Load/Build** \u2014 Write the resulting graph to Neptune\n", + "4. **Query** \u2014 Read back using the DocumentGraphQueryEngine\n", + "\n", + "This is the simplest path from tabular data to a queryable property graph.\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]` (includes lexical-graph for Neptune connectivity)\n", + "- Neptune cluster endpoint (see 00-Setup)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configuration" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "\n", + "GRAPH_STORE = os.environ.get(\n", + " 'GRAPH_STORE',\n", + " 'neptune-db://:8182'\n", + ")\n", + "TENANT = 'demo'\n", + "CSV_PATH = 'data/users.csv'\n", + "\n", + "print(f'Graph store: {GRAPH_STORE}')\n", + "print(f'Tenant: {TENANT}')\n", + "print(f'Data source: {CSV_PATH}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Extract \u2014 Read CSV\n", + "\n", + "Load structured records from a CSV file. Each row becomes a candidate graph node." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import csv\n", + "\n", + "with open(CSV_PATH) as f:\n", + " records = list(csv.DictReader(f))\n", + "\n", + "print(f'Read {len(records)} records from {CSV_PATH}')\n", + "print(f'Columns: {list(records[0].keys())}')\n", + "print(f'\\nFirst record:')\n", + "records[0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Transform \u2014 Row to Node\n", + "\n", + "The `RowToNodeTransformer` converts each tabular row into a typed graph node.\n", + "Each record becomes a node with:\n", + "- A label (node type) derived from configuration\n", + "- An ID from a designated field\n", + "- Properties from the row's columns" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig\n", + "from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer\n", + "\n", + "config = TransformerProviderConfig(name='row_to_node', args={'type': 'User'})\n", + "row_to_node = RowToNodeTransformer(config)\n", + "nodes = row_to_node.transform(records)\n", + "\n", + "print(f'Transformed {len(records)} rows \u2192 {len(nodes)} nodes')\n", + "print(f'\\nFirst node:')\n", + "nodes[0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Transform \u2014 Infer Edges\n", + "\n", + "The `EdgeInferencer` examines shared field values across nodes to create relationships.\n", + "For example, users sharing the same `account_id` get connected with a `SAME_ACCOUNT` edge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.graph_transformers.infer_edges import EdgeInferencer\n", + "\n", + "config = TransformerProviderConfig(\n", + " name='infer_edges',\n", + " args={'source_field': 'account_id', 'edge_type': 'SAME_ACCOUNT'}\n", + ")\n", + "edge_inferencer = EdgeInferencer(config)\n", + "edges = edge_inferencer.transform(records)\n", + "\n", + "print(f'Inferred {len(edges)} edges')\n", + "if edges:\n", + " print(f'\\nFirst edge:')\n", + " print(f' {edges[0]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Build \u2014 Write to Neptune\n", + "\n", + "Generate openCypher MERGE statements and execute against Neptune.\n", + "Nodes are tenant-scoped using the shared `TenantId.format_label()` format,\n", + "ensuring compatibility with lexical-graph in the same Neptune cluster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory\n", + "from graphrag_toolkit.document_graph import Node\n", + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import node_to_cypher\n", + "\n", + "graph_store = GraphStoreFactory.for_graph_store(GRAPH_STORE)\n", + "gs = graph_store.__enter__()\n", + "\n", + "# Write nodes\n", + "for n in nodes:\n", + " node = Node(\n", + " id=n.get('id', n.get('_id', '')),\n", + " labels=[n.get('node_type', 'Row')],\n", + " properties=n\n", + " )\n", + " cypher, params = node_to_cypher(node, tenant_id=TENANT)\n", + " gs.execute_query(cypher, params)\n", + "\n", + "print(f'\u2705 Wrote {len(nodes)} nodes to Neptune (tenant={TENANT})')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Query \u2014 Read Back\n", + "\n", + "Use `DocumentGraphQueryEngine` to query the tenant-scoped graph.\n", + "The engine automatically applies the correct tenant label format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.query import DocumentGraphQueryEngine\n", + "\n", + "engine = DocumentGraphQueryEngine(gs, tenant_id=TENANT)\n", + "result = engine.get_nodes('User', limit=10)\n", + "\n", + "print(f'Found {len(result)} User nodes in tenant \"{TENANT}\":')\n", + "for r in result:\n", + " print(f' {r}')\n", + "\n", + "# Clean up\n", + "graph_store.__exit__(None, None, None)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/04-Full-Pipeline-Test.ipynb b/examples/document-graph/notebooks/04-Full-Pipeline-Test.ipynb new file mode 100644 index 00000000..23c12058 --- /dev/null +++ b/examples/document-graph/notebooks/04-Full-Pipeline-Test.ipynb @@ -0,0 +1,402 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Full Pipeline Test — Ingest → Transform → Build\n", + "\n", + "Tests all stages of the document-graph pipeline end-to-end:\n", + "\n", + "1. **Ingestors** — Column selection, renaming, row filtering\n", + "2. **Transformers** — Normalizers, field transforms, graph transforms, truncators\n", + "3. **Constructors** — Cypher generation patterns (single + batch)\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph` (base install — no Neptune needed)\n", + "- This notebook runs entirely locally; it tests the pipeline logic without writing to any database." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "import pandas as pd\n", + "from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig\n", + "from graphrag_toolkit.document_graph.ingest.ingestors_provider_config import IngestorProviderConfig\n", + "\n", + "print('Imports OK ✅')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Load Test Data" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "df = pd.read_csv('data/users.csv')\n", + "records = df.to_dict('records')\n", + "print(f'Loaded {len(records)} records, columns: {list(df.columns)}')\n", + "df.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "## Stage 1: Ingestors (Data Preparation)\n", + "\n", + "Ingestors operate on DataFrames to select, rename, filter, and clean columns before transformation." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1a. Column Selector\n", + "\n", + "Keep only the columns you need for the graph." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.ingest.column.column_selector import ColumnSelectorProvider\n", + "\n", + "config = IngestorProviderConfig(name='col_select', type='column', args={'columns': ['id', 'name', 'role']})\n", + "selector = ColumnSelectorProvider(config)\n", + "selected_df = selector.ingest(df)\n", + "\n", + "print(f'Columns: {list(df.columns)} → {list(selected_df.columns)}')\n", + "selected_df.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1b. Column Renamer\n", + "\n", + "Rename columns to match your graph schema (e.g., `id` → `user_id`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.ingest.column.column_renamer import ColumnRenamerProvider\n", + "\n", + "config = IngestorProviderConfig(name='col_rename', type='column', args={'mapping': {'id': 'user_id', 'name': 'full_name'}})\n", + "renamer = ColumnRenamerProvider(config)\n", + "renamed_df = renamer.ingest(df)\n", + "\n", + "print(f'Renamed: {list(renamed_df.columns)}')\n", + "renamed_df.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1c. Row Filter\n", + "\n", + "Remove rows matching specific criteria (e.g., skip inactive users)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.ingest.row.date_range_filter import DateRangeFilterProvider\n", + "\n", + "# Simple example: filter to keep only specific roles\n", + "# (Using DataFrame filtering directly for illustration)\n", + "filtered_df = df[df['role'] != 'viewer'] if 'role' in df.columns else df\n", + "print(f'Row filter: {len(df)} → {len(filtered_df)} records')\n", + "filtered_df.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "## Stage 2: Transformers\n", + "\n", + "Transformers operate on records (list of dicts) to normalize, enrich, and restructure data." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2a. Normalize Whitespace\n", + "\n", + "Strip leading/trailing whitespace from all string fields." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_whitespace_provider import NormalizeWhitespaceProvider\n", + "\n", + "config = TransformerProviderConfig(name='ws', args={})\n", + "normalizer = NormalizeWhitespaceProvider(config)\n", + "\n", + "test_records = [{'name': ' Alice ', 'email': ' alice@corp.com '}]\n", + "result = normalizer.transform(test_records)\n", + "print(f'Before: {test_records[0]}')\n", + "print(f'After: {result[0]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2b. Normalize Nulls\n", + "\n", + "Convert empty strings, 'N/A', and None to a consistent null representation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_nulls_provider import NormalizeNullsProvider\n", + "\n", + "config = TransformerProviderConfig(name='nulls', args={})\n", + "normalizer = NormalizeNullsProvider(config)\n", + "\n", + "test_records = [{'name': 'Alice', 'phone': '', 'notes': 'N/A', 'age': None}]\n", + "result = normalizer.transform(test_records)\n", + "print(f'Before: {test_records[0]}')\n", + "print(f'After: {result[0]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2c. Normalize Case\n", + "\n", + "Standardize field values to upper/lower case for consistent graph labels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_case_provider import NormalizeCaseProvider\n", + "\n", + "config = TransformerProviderConfig(name='case', args={'fields': ['role'], 'case': 'upper'})\n", + "normalizer = NormalizeCaseProvider(config)\n", + "\n", + "test_records = [{'name': 'Alice', 'role': 'admin'}]\n", + "result = normalizer.transform(test_records)\n", + "print(f'Before: {test_records[0]}')\n", + "print(f'After: {result[0]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2d. JSON Flattener\n", + "\n", + "Extract nested JSON strings into top-level fields — useful for metadata columns." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.field_transformers.json_flattener import JSONFlattenerProvider\n", + "\n", + "config = TransformerProviderConfig(name='flatten', args={'field': 'metadata'})\n", + "flattener = JSONFlattenerProvider(config)\n", + "\n", + "test_records = [{'id': '1', 'metadata': '{\"region\": \"us-east-1\", \"tier\": \"prod\"}'}]\n", + "result = flattener.transform(test_records)\n", + "print(f'Before: {test_records[0]}')\n", + "print(f'After: {result[0]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2e. UUID Generator\n", + "\n", + "Add a unique identifier field to each record — useful when source data lacks stable IDs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.field_transformers.uuid_generator import UuidGeneratorTransformer\n", + "\n", + "config = TransformerProviderConfig(name='uuid', args={'field': 'graph_id'})\n", + "gen = UuidGeneratorTransformer(config)\n", + "\n", + "test_records = [{'name': 'Alice'}, {'name': 'Bob'}]\n", + "result = gen.transform(test_records)\n", + "print(f'Generated IDs:')\n", + "for r in result:\n", + " print(f' {r[\"name\"]}: {r[\"graph_id\"]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2f. Graph Transformers — Row to Node + Infer Edges\n", + "\n", + "Convert records into typed graph nodes and infer edges from shared field values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer\n", + "from graphrag_toolkit.document_graph.transform.graph_transformers.infer_edges import EdgeInferencer\n", + "\n", + "config = TransformerProviderConfig(name='r2n', args={'type': 'User'})\n", + "nodes = RowToNodeTransformer(config).transform(records)\n", + "\n", + "config = TransformerProviderConfig(name='edges', args={'source_field': 'account_id', 'edge_type': 'SAME_ACCOUNT'})\n", + "edges = EdgeInferencer(config).transform(records)\n", + "\n", + "print(f'Nodes: {len(nodes)}, Edges: {len(edges)}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2g. Truncators\n", + "\n", + "Limit field length to prevent oversized properties in the graph." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.truncators.length_truncator import LengthTruncator\n", + "\n", + "config = TransformerProviderConfig(name='trunc', args={'max_length': 10, 'fields': ['name']})\n", + "truncator = LengthTruncator(config)\n", + "\n", + "test_records = [{'name': 'Very Long Name That Should Be Truncated'}]\n", + "result = truncator.transform(test_records)\n", + "print(f'Before: \"{test_records[0][\"name\"]}\"')\n", + "print(f'After: \"{result[0][\"name\"]}\"')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "## Stage 3: Constructors (Cypher Generation)\n", + "\n", + "Generate openCypher MERGE statements for nodes and edges, with tenant scoping." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph import Node, Edge\n", + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import (\n", + " node_to_cypher, edge_to_cypher,\n", + " batch_nodes_to_cypher, batch_edges_to_cypher\n", + ")\n", + "\n", + "TENANT = 'pipeline_test'\n", + "\n", + "# Create test nodes\n", + "test_nodes = [\n", + " Node(id=f'n{i}', labels=['Account'], properties={'name': f'Account-{i}', 'type': 'PROD'})\n", + " for i in range(3)\n", + "]\n", + "\n", + "# Create test edges\n", + "test_edges = [\n", + " Edge(id=f'e{i}', source_id=f'n{i}', target_id=f'n{(i+1)%3}', label='LINKED_TO')\n", + " for i in range(3)\n", + "]\n", + "\n", + "# Generate Cypher\n", + "node_queries = batch_nodes_to_cypher(test_nodes, tenant_id=TENANT)\n", + "edge_queries = batch_edges_to_cypher(test_edges, tenant_id=TENANT)\n", + "\n", + "print(f'Generated {len(node_queries)} node + {len(edge_queries)} edge statements')\n", + "print(f'\\nSample node Cypher:')\n", + "print(f' {node_queries[0][0]}')\n", + "print(f'\\nSample edge Cypher:')\n", + "print(f' {edge_queries[0][0]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "## Summary" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print('=== Pipeline Test Results ===')\n", + "print()\n", + "print('Stage 1 — Ingestors:')\n", + "print(' ✅ Column Selector')\n", + "print(' ✅ Column Renamer')\n", + "print(' ✅ Row Filter')\n", + "print()\n", + "print('Stage 2 — Transformers:')\n", + "print(' ✅ Normalize Whitespace')\n", + "print(' ✅ Normalize Nulls')\n", + "print(' ✅ Normalize Case')\n", + "print(' ✅ JSON Flattener')\n", + "print(' ✅ UUID Generator')\n", + "print(' ✅ Row To Node')\n", + "print(' ✅ Infer Edges')\n", + "print(' ✅ Length Truncator')\n", + "print()\n", + "print('Stage 3 — Constructors:')\n", + "print(' ✅ Batch Node Cypher')\n", + "print(' ✅ Batch Edge Cypher')\n", + "print()\n", + "print('🎉 All pipeline stages operational.')" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/document-graph/notebooks/05a-Ingestors.ipynb b/examples/document-graph/notebooks/05a-Ingestors.ipynb new file mode 100644 index 00000000..44e0c8a2 --- /dev/null +++ b/examples/document-graph/notebooks/05a-Ingestors.ipynb @@ -0,0 +1,270 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Ingestors Deep Dive\n", + "\n", + "Ingestors are the first stage of the document-graph pipeline. They prepare raw DataFrames\n", + "for transformation by selecting, renaming, reordering, and filtering columns and rows.\n", + "\n", + "This notebook tests all seven ingestor providers:\n", + "\n", + "| Provider | Category | Purpose |\n", + "|----------|----------|--------|\n", + "| ColumnSelectorProvider | Column | Keep/drop specific columns |\n", + "| ColumnRenamerProvider | Column | Rename columns to match graph schema |\n", + "| ColumnReorderProvider | Column | Reorder columns for consistent processing |\n", + "| ColumnTypeConverterProvider | Column | Convert dtypes (string\u2192int, string\u2192datetime) |\n", + "| SkipRowProvider | Row | Filter rows by field conditions |\n", + "| DateRangeFilterProvider | Row | Filter rows by date range |\n", + "| NumericIdCleanupIngestor | Field | Remove .0 suffix from float-typed IDs |\n", + "\n", + "**Prerequisites:**\n", + "- \\ (base install)\n", + "- No Neptune or AWS credentials required\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "from graphrag_toolkit.document_graph.ingest.ingestors_provider_config import IngestorProviderConfig\n", + "from graphrag_toolkit.document_graph.ingest.column.column_selector import ColumnSelectorProvider\n", + "from graphrag_toolkit.document_graph.ingest.column.column_renamer import ColumnRenamerProvider\n", + "from graphrag_toolkit.document_graph.ingest.column.column_reorder import ColumnReorderProvider\n", + "from graphrag_toolkit.document_graph.ingest.column.column_type_converter import ColumnTypeConverterProvider\n", + "from graphrag_toolkit.document_graph.ingest.row.skip_row import SkipRowProvider\n", + "from graphrag_toolkit.document_graph.ingest.row.date_range_filter import DateRangeFilterProvider\n", + "from graphrag_toolkit.document_graph.ingest.field.numeric_id_cleanup_ingestor import NumericIdCleanupIngestor" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Sample security/compliance data\n", + "df = pd.DataFrame({\n", + " 'user_id': [1001.0, 2002.0, 3003.0, 4004.0],\n", + " 'username': ['admin_user', 'analyst_01', 'auditor_03', 'service_acct'],\n", + " 'role': ['Administrator', 'Security Analyst', 'Compliance Auditor', 'Service Account'],\n", + " 'access_level': ['5', '3', '4', '2'],\n", + " 'status': ['active', 'active', 'inactive', 'suspended'],\n", + " 'last_login': ['2026-06-20', '2026-06-22', '2026-01-15', '2025-12-01'],\n", + " 'department': ['IT Security', 'SOC', 'GRC', 'DevOps'],\n", + " 'notes': ['Primary admin', 'L2 analyst', 'Annual review pending', 'Disabled']\n", + "})\n", + "print(\"Source DataFrame:\")\n", + "print(df.to_string())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Column Selector\n\nKeep only the columns that will become node properties or identifiers. Drop irrelevant columns early to reduce pipeline complexity.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = IngestorProviderConfig(name='select_columns', type='column', args={\n", + " 'action': 'select',\n", + " 'columns': ['user_id', 'username', 'role', 'status']\n", + "})\n", + "ingestor = ColumnSelectorProvider(config)\n", + "result = ingestor.ingest(df)\n", + "print(\"ColumnSelectorProvider (select):\")\n", + "print(result.to_string())\n", + "print(f\"\\nColumns retained: {list(result.columns)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Column Renamer\n\nRename source columns to match your target graph schema. For example, rename \\ \u2192 \\ when your graph model uses different terminology than the source data.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = IngestorProviderConfig(name='rename_columns', type='column', args={\n", + " 'mapping': {\n", + " 'user_id': 'account_id',\n", + " 'username': 'principal_name',\n", + " 'access_level': 'clearance_level'\n", + " }\n", + "})\n", + "ingestor = ColumnRenamerProvider(config)\n", + "result = ingestor.ingest(df)\n", + "print(\"ColumnRenamerProvider:\")\n", + "print(result.columns.tolist())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Column Reorder\n\nReorder columns for consistent downstream processing. Some constructors assume a specific column order (e.g., first column = node ID).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = IngestorProviderConfig(name='reorder_columns', type='column', args={\n", + " 'column_order': ['status', 'role', 'username', 'user_id']\n", + "})\n", + "ingestor = ColumnReorderProvider(config)\n", + "result = ingestor.ingest(df)\n", + "print(\"ColumnReorderProvider:\")\n", + "print(f\"New column order: {list(result.columns)}\")\n", + "print(result.head(2).to_string())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Column Type Converter\n\nConvert column data types before transformation. Ensures numeric fields are integers (not strings) and date fields are proper datetime objects for temporal queries.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = IngestorProviderConfig(name='convert_types', type='column', args={\n", + " 'type_mapping': {\n", + " 'access_level': 'int',\n", + " 'last_login': 'datetime'\n", + " }\n", + "})\n", + "ingestor = ColumnTypeConverterProvider(config)\n", + "result = ingestor.ingest(df)\n", + "print(\"ColumnTypeConverterProvider:\")\n", + "print(result.dtypes)\n", + "print(result[['access_level', 'last_login']].to_string())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Row Filter (SkipRowProvider)\n\nFilter rows by field conditions \u2014 keep only records that match your criteria. Useful for excluding inactive accounts, test data, or irrelevant record types.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = IngestorProviderConfig(name='filter_active', type='row', args={\n", + " 'conditions': [\n", + " {'field': 'status', 'operator': 'eq', 'value': 'active'}\n", + " ],\n", + " 'keep_matching': True\n", + "})\n", + "ingestor = SkipRowProvider(config)\n", + "result = ingestor.ingest(df)\n", + "print(\"SkipRowProvider (keep active only):\")\n", + "print(result[['username', 'status']].to_string())\n", + "print(f\"\\nRows: {len(df)} -> {len(result)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Date Range Filter\n\nFilter records by a date field \u2014 keep only rows within a specified time window. Useful for incremental loads or time-bounded analysis.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = IngestorProviderConfig(name='recent_logins', type='row', args={\n", + " 'date_field': 'last_login',\n", + " 'start_date': '2026-06-01',\n", + " 'include_null': False\n", + "})\n", + "ingestor = DateRangeFilterProvider(config)\n", + "result = ingestor.ingest(df)\n", + "print(\"DateRangeFilterProvider (logins since 2026-06-01):\")\n", + "print(result[['username', 'last_login']].to_string())\n", + "print(f\"\\nRows: {len(df)} -> {len(result)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Numeric ID Cleanup\n\nRemove the \\ suffix that pandas adds when reading integer IDs from CSV (\\ \u2192 \\). Ensures clean, consistent node identifiers.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = IngestorProviderConfig(name='clean_ids', type='field', args={\n", + " 'field': 'user_id'\n", + "})\n", + "ingestor = NumericIdCleanupIngestor(config)\n", + "result = ingestor.ingest(df.copy())\n", + "print(\"NumericIdCleanupIngestor:\")\n", + "print(f\"Before: {df['user_id'].tolist()}\")\n", + "print(f\"After: {result['user_id'].tolist()}\")\n", + "print(f\"Type: {result['user_id'].dtype}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "All ingestor providers tested successfully:\n", + "- **ColumnSelectorProvider**: Select/drop columns\n", + "- **ColumnRenamerProvider**: Rename columns via mapping\n", + "- **ColumnReorderProvider**: Reorder columns\n", + "- **ColumnTypeConverterProvider**: Convert column dtypes\n", + "- **SkipRowProvider**: Filter rows by conditions\n", + "- **DateRangeFilterProvider**: Filter by date range\n", + "- **NumericIdCleanupIngestor**: Remove .0 suffix from IDs" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/05b-Schema-Driven-Pipeline.ipynb b/examples/document-graph/notebooks/05b-Schema-Driven-Pipeline.ipynb new file mode 100644 index 00000000..e8fc7deb --- /dev/null +++ b/examples/document-graph/notebooks/05b-Schema-Driven-Pipeline.ipynb @@ -0,0 +1,181 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Schema-Driven Pipeline\n", + "\n", + "Demonstrates the full document-graph pipeline driven by a schema definition:\n", + "\n", + "1. **Extract** \u2014 CSVExtractProvider reads source data and infers schema\n", + "2. **Ingest** \u2014 Column selection and row filtering\n", + "3. **Transform** \u2014 Row-to-node conversion and edge inference\n", + "4. **Build** \u2014 Cypher generation and Neptune write\n", + "5. **Query** \u2014 Verify the graph via DocumentGraphQueryEngine\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]` (Neptune connectivity)\n", + "- Neptune cluster endpoint (see 00-Setup)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Note: Schema definitions can be loaded from S3\n\nThe `S3SchemaProvider` supports loading and saving ETL schemas directly from/to S3:\n\n```python\nfrom graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider\n\nconfig = {\n \"provider_type\": \"s3\",\n \"schema_id\": \"my-pipeline-schema\",\n \"connection_config\": {\n \"bucket\": \"\",\n \"key\": \"schemas/my-schema.json\",\n \"region\": \"us-east-1\"\n }\n}\n\nprovider = get_schema_provider(config)\nschema = provider.load_schema()\n```\n\nSee `document_graph.schema.providers.s3_schema_provider` for full details.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "\n", + "GRAPH_STORE = os.environ.get(\n", + " 'GRAPH_STORE',\n", + " 'neptune-db://:8182'\n", + ")\n", + "TENANT = 'schema_demo'\n", + "DATA_PATH = 'data/users.csv'\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Extract\n", + "\n", + "The `CSVExtractProvider` reads source data and produces an extraction result containing\n", + "the DataFrame plus inferred schema metadata.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_csv import CSVExtractProvider\nfrom graphrag_toolkit.document_graph.pipeline.extract.extract_provider_config import ExtractProviderConfig\nfrom graphrag_toolkit.document_graph.config import DocumentGraphConfig\n\nconfig = ExtractProviderConfig(type='csv', source=DATA_PATH, document_id='users-pipeline-test')\nprovider = CSVExtractProvider(config, DocumentGraphConfig())\nresult = provider.extract('DATA_PATH')\n\nprint(f'Extracted: {result.dataframe.shape[0]} rows')\nprint(f'Schema: {list(result.extracted_schema.keys())}')\nprint(f'Nodes: {len(result.nodes)}')\nprint(result.dataframe)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Ingest\n", + "\n", + "Apply column selection and row filtering to prepare the data for graph construction.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.document_graph.ingest.ingestors_provider_config import IngestorProviderConfig\n", + "from graphrag_toolkit.document_graph.ingest.column.column_selector import ColumnSelectorProvider\n", + "from graphrag_toolkit.document_graph.ingest.row.skip_row import SkipRowProvider\n", + "\n", + "df = result.dataframe\n", + "\n", + "# Select columns\n", + "config = IngestorProviderConfig(name='select', type='column', args={'columns': ['id', 'name', 'email', 'role', 'account_id']})\n", + "df = ColumnSelectorProvider(config).ingest(df)\n", + "\n", + "# Filter out viewers\n", + "config = IngestorProviderConfig(name='filter', type='row', args={'conditions': [{'field': 'role', 'operator': 'ne', 'value': 'viewer'}]})\n", + "df = SkipRowProvider(config).ingest(df)\n", + "\n", + "print(f'After ingest: {len(df)} rows (viewers removed)')\n", + "print(df)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Transform\n", + "\n", + "Convert rows to typed graph nodes and infer edges from shared field values.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig\n", + "from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer\n", + "from graphrag_toolkit.document_graph.transform.graph_transformers.infer_edges import EdgeInferencer\n", + "\n", + "records = df.to_dict('records')\n", + "\n", + "# Rows \u2192 User nodes\n", + "config = TransformerProviderConfig(name='r2n', args={'type': 'User'})\n", + "nodes_data = RowToNodeTransformer(config).transform(records)\n", + "print(f'Nodes: {len(nodes_data)}')\n", + "\n", + "# Infer edges from shared account_id\n", + "config = TransformerProviderConfig(name='edges', args={'source_field': 'account_id', 'edge_type': 'SAME_ACCOUNT'})\n", + "edges_data = EdgeInferencer(config).transform(records)\n", + "print(f'Edges: {len(edges_data)}')\n", + "print(f'First node: {nodes_data[0]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Build \u2014 Write to Neptune\n", + "\n", + "Generate openCypher statements and execute against Neptune. Uses tenant-scoped labels\n", + "for multi-tenant isolation.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory\nfrom graphrag_toolkit.document_graph import Node, Edge\nfrom graphrag_toolkit.document_graph.graph_build import node_to_cypher, edge_to_cypher\n\ngs = GraphStoreFactory.for_graph_store(GRAPH_STORE).__enter__()\nTENANT = 'pipeline_v3'\n\n# Write nodes\nfor n in nodes_data:\n node = Node(id=n.get('id'), labels=[n.get('node_type', 'User')], properties=n)\n cypher, params = node_to_cypher(node, tenant_id=TENANT)\n gs.execute_query(cypher, params)\nprint(f'\u2705 Wrote {len(nodes_data)} nodes (tenant={TENANT})')\n\n# Write edges\nedge_count = 0\nfor ed in edges_data:\n edge = Edge(id=ed.get('id',''), source_id=ed.get('source_id',''), target_id=ed.get('target_id',''), label=ed.get('edge_type','RELATED'))\n cypher, params = edge_to_cypher(edge, tenant_id=TENANT)\n gs.execute_query(cypher, params)\n edge_count += 1\nprint(f'\u2705 Wrote {edge_count} edges')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Query \u2014 Verify\n", + "\n", + "Read back from Neptune to confirm the graph was built correctly.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.document_graph.query import DocumentGraphQueryEngine\n\nengine = DocumentGraphQueryEngine(gs, tenant_id=TENANT)\n\nusers = engine.get_nodes('User')\nprint(f'\\n=== Results ===')\nprint(f'Users in graph: {len(users)}')\nfor u in users:\n props = u['n']['~properties']\n print(f\" {props.get('name')} ({props.get('role')}) \u2192 account: {props.get('account_id')}\")\n\nprint(f'\\n\u2705 Full pipeline: Extract \u2192 Ingest \u2192 Transform \u2192 Build \u2192 Query COMPLETE')\n\n# Clean up\ngraph_store.__exit__(None, None, None) if 'graph_store' in dir() else gs.__exit__(None, None, None)\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/05c-JSON-Driven-Pipeline.ipynb b/examples/document-graph/notebooks/05c-JSON-Driven-Pipeline.ipynb new file mode 100644 index 00000000..e41c6dcf --- /dev/null +++ b/examples/document-graph/notebooks/05c-JSON-Driven-Pipeline.ipynb @@ -0,0 +1,156 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# JSON-Driven Pipeline\n", + "\n", + "Define your graph schema in JSON, then `PipelineExecutor` builds the graph automatically.\n", + "No manual transform steps \u2014 the JSON schema defines nodes, edges, and property mappings.\n", + "\n", + "This is the declarative approach to document-graph: describe *what* you want, not *how* to build it.\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]`\n", + "- Neptune cluster endpoint (see 00-Setup)\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "\n", + "GRAPH_STORE = os.environ.get(\n", + " 'GRAPH_STORE',\n", + " 'neptune-db://:8182'\n", + ")\n", + "TENANT = 'json_driven'\n", + "DATA_PATH = 'data/users.csv'\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Define the Schema (JSON)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Note: Schema JSON can also be loaded from S3\n\nInstead of defining the schema inline, you can load it from S3:\n\n```python\nfrom graphrag_toolkit.document_graph.schema.providers.schema_provider_factory import get_schema_provider\n\nprovider = get_schema_provider({\n \"provider_type\": \"s3\",\n \"schema_id\": \"my-mapping\",\n \"connection_config\": {\n \"bucket\": \"\",\n \"key\": \"schemas/mapping.json\",\n \"region\": \"us-east-1\"\n }\n})\nschema = provider.load_schema()\n```\n\nThis is useful for sharing schemas across notebooks and pipelines without duplication.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json, os\n", + "\n", + "# This is what mapping.json looks like \u2014 defines nodes, edges, property mappings\n", + "schema = {\n", + " \"schema_version\": \"1.0\",\n", + " \"graph\": {\n", + " \"nodes\": [\n", + " {\"User\": {\"type\": \"node\", \"property\": {\"name\": \"user_id\", \"csv_map\": \"id\"}}},\n", + " {\"Account\": {\"type\": \"node\", \"property\": {\"name\": \"account_id\", \"csv_map\": \"account_id\"}}}\n", + " ],\n", + " \"edges\": [\n", + " {\"User\": {\"BELONGS_TO\": \"Account\"}}\n", + " ]\n", + " },\n", + " \"mappings\": [\n", + " {\"csv_property_name\": \"id\", \"type\": \"node\", \"arrow_property_name\": \"user_id\", \"name\": \"User\"},\n", + " {\"csv_property_name\": \"name\", \"type\": \"property\", \"arrow_property_name\": \"full_name\", \"parents\": [\"User\"]},\n", + " {\"csv_property_name\": \"email\", \"type\": \"property\", \"arrow_property_name\": \"email\", \"parents\": [\"User\"]},\n", + " {\"csv_property_name\": \"role\", \"type\": \"property\", \"arrow_property_name\": \"role\", \"parents\": [\"User\"]},\n", + " {\"csv_property_name\": \"account_id\", \"type\": \"node\", \"arrow_property_name\": \"account_id\", \"name\": \"Account\"}\n", + " ]\n", + "}\n", + "\n", + "# Save to file\n", + "schema_path = 'data/mapping.json'\n", + "with open(schema_path, 'w') as f:\n", + " json.dump(schema, f, indent=2)\n", + "print(f'Schema written to {schema_path}')\n", + "print(f'Nodes: {[list(n.keys())[0] for n in schema[\"graph\"][\"nodes\"]]}')\n", + "print(f'Edges: {schema[\"graph\"][\"edges\"]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Load Source Data\n", + "\n", + "Read the CSV that the JSON schema will map into a graph.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n\ndf = pd.read_csv(DATA_PATH)\nprint(f'Source: {len(df)} rows')\nprint(df)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Run Pipeline\n", + "\n", + "Pass the JSON schema and source DataFrame to `PipelineExecutor`. It handles\n", + "extraction, transformation, and graph construction automatically.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory\nfrom graphrag_toolkit.document_graph import PipelineExecutor\n\ngs = GraphStoreFactory.for_graph_store(GRAPH_STORE).__enter__()\n\nexecutor = PipelineExecutor(gs, tenant_id=TENANT)\nresult = executor.run_from_schema(schema_path, source_df=df)\n\nprint(f'\\n=== Pipeline Result ===')\nprint(f'Records processed: {result.records_processed}')\nprint(f'Nodes created: {result.nodes_created}')\nprint(f'Edges created: {result.edges_created}')\nprint(f'Errors: {len(result.errors)}')\nif result.errors:\n for e in result.errors[:3]:\n print(f' {e}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify\n", + "\n", + "Query Neptune to confirm nodes and edges were created from the JSON schema definition.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.document_graph.query import DocumentGraphQueryEngine\n\nengine = DocumentGraphQueryEngine(gs, tenant_id=TENANT)\n\nprint('Users:')\nfor u in engine.get_nodes('User'):\n props = u['n']['~properties']\n print(f\" {props.get('full_name', props.get('user_id', '?'))} ({props.get('role', '?')})\")\n\nprint('\\nAccounts:')\nfor a in engine.get_nodes('Account'):\n props = a['n']['~properties']\n print(f\" {props.get('account_id', '?')}\")\n\nprint(f'\\n\u2705 JSON-driven pipeline complete \u2014 schema \u2192 Neptune in one call')\n\n# Clean up\ngs.__exit__(None, None, None)\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/05d-Multi-Format-Extraction.ipynb b/examples/document-graph/notebooks/05d-Multi-Format-Extraction.ipynb new file mode 100644 index 00000000..a87654c8 --- /dev/null +++ b/examples/document-graph/notebooks/05d-Multi-Format-Extraction.ipynb @@ -0,0 +1,243 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Multi-Format Extraction: Parquet, Excel, XML, YAML\n", + "\n", + "document-graph supports multiple structured data formats beyond CSV and JSON.\n", + "This notebook demonstrates extraction from:\n", + "\n", + "| Format | Provider | Use Case |\n", + "|--------|----------|----------|\n", + "| Parquet | `ParquetExtractProvider` | Data lakes (S3, Athena exports, Spark output) |\n", + "| Excel | `ExcelExtractProvider` | Business data, spreadsheets, manual inventories |\n", + "| XML | XML schema discovery | Legacy systems, SOAP APIs, configuration files |\n", + "| YAML | YAML schema discovery | Infrastructure-as-code, Kubernetes manifests |\n", + "\n", + "All formats go through the same pipeline: Extract → Transform → Build.\n", + "The extraction stage normalizes everything into DataFrames regardless of source format.\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph` (base install)\n", + "- `pip install openpyxl` (for Excel support)\n", + "- No Neptune or AWS credentials required\n", + "- Sample data in `data/` directory" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Parquet Extraction\n", + "\n", + "Parquet is the standard format for data lake outputs (Athena queries, Spark jobs, AWS Glue).\n", + "The `ParquetExtractProvider` reads columnar data efficiently and infers schema from column metadata." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_parquet import ParquetExtractProvider\n", + "from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_config import ExtractProviderConfig\n", + "from graphrag_toolkit.document_graph.config import DocumentGraphConfig\n", + "\n", + "config = ExtractProviderConfig(type='parquet', source='test_documents.parquet', document_id='parquet-demo')\n", + "provider = ParquetExtractProvider(config, DocumentGraphConfig())\n", + "result = provider.extract('data/test_documents.parquet')\n", + "\n", + "print(f'Extracted: {result.dataframe.shape[0]} rows, {result.dataframe.shape[1]} columns')\n", + "print(f'Columns: {list(result.dataframe.columns)}')\n", + "print(f'Dtypes:\\n{result.dataframe.dtypes}')\n", + "print(f'\\nFirst 3 rows:')\n", + "result.dataframe.head(3)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parquet Schema Discovery\n", + "\n", + "The schema discovery provider can infer an ETL schema directly from Parquet column metadata\n", + "(types, names, nullability) — no manual schema definition needed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.schema.discovery.parquet_discovery_provider import ParquetSchemaDiscoveryProvider\n", + "\n", + "discovery = ParquetSchemaDiscoveryProvider()\n", + "schema = discovery.discover('data/test_documents.parquet')\n", + "\n", + "print('Discovered schema from Parquet metadata:')\n", + "print(f' Fields: {len(schema.fields) if hasattr(schema, \"fields\") else \"N/A\"}')\n", + "print(f' Schema: {schema}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Excel Extraction\n", + "\n", + "Excel files are common in business workflows — account lists, asset inventories, compliance records.\n", + "The `ExcelExtractProvider` reads `.xlsx` files and supports multi-sheet extraction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.pipeline.extract.extract_provider_excel import ExcelExtractProvider\n", + "\n", + "# Check if we have Excel sample data\n", + "import os\n", + "excel_files = [f for f in os.listdir('data') if f.endswith('.xlsx')]\n", + "\n", + "if excel_files:\n", + " excel_path = f'data/{excel_files[0]}'\n", + " config = ExtractProviderConfig(type='excel', source=excel_files[0], document_id='excel-demo')\n", + " provider = ExcelExtractProvider(config, DocumentGraphConfig())\n", + " result = provider.extract(excel_path)\n", + " \n", + " print(f'Extracted from {excel_files[0]}:')\n", + " print(f' Rows: {result.dataframe.shape[0]}, Columns: {result.dataframe.shape[1]}')\n", + " print(f' Columns: {list(result.dataframe.columns)}')\n", + " result.dataframe.head(3)\n", + "else:\n", + " print('No .xlsx files in data/ — skipping Excel demo')\n", + " print('To test: place any .xlsx file in the data/ directory')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. XML Schema Discovery\n", + "\n", + "XML is common in legacy systems, SOAP APIs, and configuration files.\n", + "The XML discovery provider parses the document structure to infer node types and relationships." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.schema.discovery.xml_discovery_provider import XMLSchemaDiscoveryProvider\n", + "\n", + "discovery = XMLSchemaDiscoveryProvider()\n", + "schema = discovery.discover('data/test_documents.xml')\n", + "\n", + "print('Discovered schema from XML structure:')\n", + "print(f' Schema: {schema}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. YAML Schema Discovery\n", + "\n", + "YAML is used in infrastructure-as-code (CloudFormation, Kubernetes), CI/CD configs, and documentation.\n", + "The YAML discovery provider extracts structure from the document hierarchy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.schema.discovery.yaml_discovery_provider import YAMLSchemaDiscoveryProvider\n", + "\n", + "discovery = YAMLSchemaDiscoveryProvider()\n", + "schema = discovery.discover('data/test_documents.yaml')\n", + "\n", + "print('Discovered schema from YAML structure:')\n", + "print(f' Schema: {schema}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Pipeline: Parquet → Transform → Cypher\n", + "\n", + "Complete pipeline from Parquet extraction through to Cypher generation — proving\n", + "that non-CSV formats work through the full document-graph pipeline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph import Node\n", + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import node_to_cypher, batch_nodes_unwind\n", + "from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig\n", + "from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer\n", + "\n", + "# Extract from Parquet\n", + "config = ExtractProviderConfig(type='parquet', source='test_documents.parquet', document_id='parquet-pipeline')\n", + "provider = ParquetExtractProvider(config, DocumentGraphConfig())\n", + "result = provider.extract('data/test_documents.parquet')\n", + "records = result.dataframe.to_dict('records')\n", + "\n", + "# Transform: rows → nodes\n", + "transformer_config = TransformerProviderConfig(name='r2n', args={'type': 'Document'})\n", + "nodes = RowToNodeTransformer(transformer_config).transform(records)\n", + "print(f'Transformed {len(records)} Parquet rows → {len(nodes)} nodes')\n", + "\n", + "# Build: generate Cypher\n", + "graph_nodes = [\n", + " Node(id=n.get('id', str(i)), labels=['Document'], properties=n)\n", + " for i, n in enumerate(nodes)\n", + "]\n", + "cypher, params = batch_nodes_unwind(graph_nodes, tenant_id='parquet_demo')\n", + "\n", + "print(f'\\nGenerated batch UNWIND query for {len(graph_nodes)} nodes:')\n", + "print(f' Cypher: {cypher[:100]}...')\n", + "print(f' Batch size: {len(params[\"batch\"])}')\n", + "print(f'\\n✅ Parquet → Transform → Cypher: complete')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "All structured formats go through the same pipeline:\n", + "\n", + "```\n", + "Parquet/Excel/XML/YAML/CSV/JSON\n", + " ↓ ExtractProvider (format-specific)\n", + "DataFrame (universal)\n", + " ↓ Transformers (format-agnostic)\n", + "Node/Edge models\n", + " ↓ Constructors\n", + "openCypher statements\n", + " ↓ GraphStore\n", + "Neptune / Neo4j / FalkorDB\n", + "```\n", + "\n", + "The extraction stage is the only format-specific code. Everything downstream is universal." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/document-graph/notebooks/06a-Field-Transformers.ipynb b/examples/document-graph/notebooks/06a-Field-Transformers.ipynb new file mode 100644 index 00000000..baca46e7 --- /dev/null +++ b/examples/document-graph/notebooks/06a-Field-Transformers.ipynb @@ -0,0 +1,437 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Field Transformers Deep Dive\n", + "\n", + "Field transformers operate on individual columns/fields within records, restructuring\n", + "nested data, splitting multi-valued fields, generating IDs, and normalizing values.\n", + "\n", + "This notebook tests all 12 field transformer providers:\n", + "\n", + "| Transformer | Purpose |\n", + "|-------------|--------|\n", + "| JSONFlattenerProvider | Extract nested JSON string into top-level fields |\n", + "| JSONArrayExpanderProvider | Expand JSON arrays into separate rows |\n", + "| JSONArrayFlattenerProvider | Flatten JSON arrays into delimited strings |\n", + "| JSONValueFlattenerProvider | Extract specific values from JSON objects |\n", + "| CommaFlattenerProvider | Join list values into comma-separated strings |\n", + "| CommaSplitProvider | Split comma-separated strings into lists |\n", + "| EmbeddedJSONFieldTransformer | Parse embedded JSON and merge into record |\n", + "| PairedFlattenerProvider | Flatten paired key-value structures |\n", + "| RegexCleanerProvider | Clean field values using regex patterns |\n", + "| EnumStandardizer | Map variant values to standard enum labels |\n", + "| UuidGeneratorTransformer | Generate unique IDs for records lacking them |\n", + "| TimestampNormalizer | Normalize date/time formats to ISO 8601 |\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph` (base install)\n", + "- No Neptune or AWS credentials required\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.json_flattener import JSONFlattenerProvider\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.json_array_expander import JSONArrayExpanderProvider\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.json_array_flattener import JSONArrayFlattenerProvider\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.json_value_flattener import JSONValueFlattenerProvider\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.comma_flattener import CommaFlattenerProvider\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.comma_split_provider import CommaSplitProvider\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.embedded_json import EmbeddedJSONFieldTransformer\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.paired_flattener import PairedFlattenerProvider\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.regex_cleaner_provider import RegexCleanerProvider\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.standardize_enum import EnumStandardizer\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.uuid_generator import UuidGeneratorTransformer\n", + "from graphrag_toolkit.document_graph.transform.field_transformers.normalize_timestamp import TimestampNormalizer" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## JSON Flattener\n\nExtract fields from a JSON string column into top-level record properties.\nUseful for metadata columns stored as serialized JSON.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'id': 'CTRL-001', 'control_data': '{\"framework\": \"NIST\", \"category\": \"Access Control\", \"severity\": \"High\"}'},\n", + " {'id': 'CTRL-002', 'control_data': '{\"framework\": \"SOC2\", \"category\": \"Monitoring\", \"severity\": \"Medium\"}'}\n", + "]\n", + "config = TransformerProviderConfig(name='flatten_control', args={'field': 'control_data', 'prefix': 'ctrl_'})\n", + "transformer = JSONFlattenerProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"JSONFlattenerProvider:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## JSON Array Expander\n\nExpand a JSON array field into multiple rows \u2014 one row per array element.\nUseful for one-to-many relationships encoded in a single column.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'finding_id': 'F-101', 'evidence': '[{\"type\": \"LOG\", \"source\": \"CloudTrail\"}, {\"type\": \"ALERT\", \"source\": \"GuardDuty\"}]'}\n", + "]\n", + "config = TransformerProviderConfig(name='expand_evidence', args={\n", + " 'field': 'evidence',\n", + " 'skip_empty_arrays': True,\n", + " 'conflict_resolution': 'prefix',\n", + " 'auto_add_discovered': True\n", + "})\n", + "transformer = JSONArrayExpanderProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"JSONArrayExpanderProvider:\")\n", + "print(f\"Input: 1 record -> Output: {len(result)} records\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## JSON Array Flattener\n\nFlatten a JSON array into a delimited string (e.g., `[\"a\",\"b\"]` \u2192 `\"a,b\"`).\nUseful when you need a single string property, not multiple rows.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'policy_id': 'POL-01', 'controls': '[\"AC-1\", \"AC-2\", \"AC-3\"]'}\n", + "]\n", + "config = TransformerProviderConfig(name='flatten_controls', args={\n", + " 'field': 'controls',\n", + " 'prefix': 'control',\n", + " 'suffix': ' #00'\n", + "})\n", + "transformer = JSONArrayFlattenerProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"JSONArrayFlattenerProvider:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## JSON Value Flattener\n\nExtract specific values from JSON objects by key path.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'user': 'admin', 'project_ids': 'PRJ-100, PRJ-200, PRJ-300', 'project_names': 'IAM Upgrade, SOC Setup, Audit Prep'}\n", + "]\n", + "config = TransformerProviderConfig(name='flatten_projects', args={\n", + " 'id_field': 'project_ids',\n", + " 'name_field': 'project_names',\n", + " 'prefix': 'project',\n", + " 'max_items': 5\n", + "})\n", + "transformer = JSONValueFlattenerProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"JSONValueFlattenerProvider:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Comma Flattener\n\nJoin list-typed values into comma-separated strings for storage as a single property.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'user': 'analyst_01', 'permissions': 'read:logs, write:alerts, admin:dashboards'}\n", + "]\n", + "config = TransformerProviderConfig(name='flatten_permissions', args={\n", + " 'field': 'permissions',\n", + " 'prefix': 'perm',\n", + " 'suffix': ' #00',\n", + " 'max_items': 10\n", + "})\n", + "transformer = CommaFlattenerProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"CommaFlattenerProvider:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Comma Split\n\nSplit comma-separated strings into Python lists for downstream processing.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'control_id': 'AC-2', 'assigned_teams': 'IAM Team, SOC, Platform Engineering'}\n", + "]\n", + "config = TransformerProviderConfig(name='split_teams', args={\n", + " 'fields': ['assigned_teams'],\n", + " 'separator': ',',\n", + " 'strip_whitespace': True\n", + "})\n", + "transformer = CommaSplitProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"CommaSplitProvider:\")\n", + "print(f\"Input: 1 record -> Output: {len(result)} records\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Embedded JSON\n\nParse an embedded JSON string and merge its keys directly into the parent record.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'alert_id': 'ALT-500', 'metadata': '{\"source\": \"GuardDuty\", \"region\": \"us-east-1\", \"account_id\": \"123456789012\"}'}\n", + "]\n", + "config = TransformerProviderConfig(name='extract_metadata', args={\n", + " 'field': 'metadata',\n", + " 'prefix': 'meta_'\n", + "})\n", + "transformer = EmbeddedJSONFieldTransformer(config)\n", + "result = transformer.transform(records)\n", + "print(\"EmbeddedJSONFieldTransformer:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Paired Flattener\n\nFlatten paired key-value structures (e.g., `{\"key\": \"region\", \"value\": \"us-east-1\"}`) into direct properties.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'user': 'admin', 'account_ids': 'ACC-01, ACC-02', 'account_names': 'Production, Staging'}\n", + "]\n", + "config = TransformerProviderConfig(name='pair_accounts', args={\n", + " 'id_field': 'account_ids',\n", + " 'name_field': 'account_names',\n", + " 'prefix': 'account',\n", + " 'suffix': '#00',\n", + " 'max_items': 5\n", + "})\n", + "transformer = PairedFlattenerProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"PairedFlattenerProvider:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Regex Cleaner\n\nClean field values using regex patterns \u2014 remove special characters, extract substrings, or normalize formats.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'id': 'V-001', 'description': '

Critical CVE-2024-1234 found in production.

'},\n", + " {'id': 'V-002', 'description': '
Unpatched OpenSSL vulnerability.
'}\n", + "]\n", + "config = TransformerProviderConfig(name='strip_html', args={\n", + " 'patterns': ['<[^>]*>'],\n", + " 'replacement': '',\n", + " 'fields': ['description']\n", + "})\n", + "transformer = RegexCleanerProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"RegexCleanerProvider:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Enum Standardizer\n\nMap variant values (typos, abbreviations, legacy codes) to a standard set of enum labels.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'id': 1, 'severity': 'Critical'},\n", + " {'id': 2, 'severity': 'HIGH'},\n", + " {'id': 3, 'severity': 'med'},\n", + " {'id': 4, 'severity': 'informational'}\n", + "]\n", + "config = TransformerProviderConfig(name='standardize_severity', args={\n", + " 'field': 'severity',\n", + " 'mapping': {\n", + " 'critical': 'CRITICAL',\n", + " 'high': 'HIGH',\n", + " 'med': 'MEDIUM',\n", + " 'medium': 'MEDIUM',\n", + " 'low': 'LOW',\n", + " 'informational': 'INFO'\n", + " },\n", + " 'case_insensitive': True\n", + "})\n", + "transformer = EnumStandardizer(config)\n", + "result = transformer.transform(records)\n", + "print(\"EnumStandardizer:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## UUID Generator\n\nGenerate unique identifiers for records that lack stable IDs \u2014 ensures every node has a distinct identity in the graph.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'control': 'AC-1', 'title': 'Access Control Policy'},\n", + " {'control': 'AC-2', 'title': 'Account Management'}\n", + "]\n", + "config = TransformerProviderConfig(name='gen_uuid', args={'target_field': 'node_id'})\n", + "transformer = UuidGeneratorTransformer(config)\n", + "result = transformer.transform(records)\n", + "print(\"UuidGeneratorTransformer:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Timestamp Normalizer\n\nNormalize date/time values from various formats into consistent ISO 8601 strings for temporal queries.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'event': 'login_attempt', 'timestamp': 'Jun 20, 2026 14:30:45'},\n", + " {'event': 'privilege_escalation', 'timestamp': '2026/06/21 09:15 AM'},\n", + " {'event': 'data_export', 'timestamp': '21-06-2026 23:59:59'}\n", + "]\n", + "config = TransformerProviderConfig(name='normalize_ts', args={\n", + " 'field': 'timestamp',\n", + " 'output_field': 'iso_timestamp'\n", + "})\n", + "transformer = TimestampNormalizer(config)\n", + "result = transformer.transform(records)\n", + "print(\"TimestampNormalizer:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "All field transformers tested successfully:\n", + "- **JSONFlattenerProvider**: Flatten JSON objects into prefixed columns\n", + "- **JSONArrayExpanderProvider**: Expand JSON arrays into multiple records\n", + "- **JSONArrayFlattenerProvider**: Flatten JSON arrays into numbered columns\n", + "- **JSONValueFlattenerProvider**: Create JSON objects from paired fields\n", + "- **CommaFlattenerProvider**: Flatten comma-separated values into numbered columns\n", + "- **CommaSplitProvider**: Split comma-separated values into multiple records\n", + "- **EmbeddedJSONFieldTransformer**: Parse and flatten embedded JSON strings\n", + "- **PairedFlattenerProvider**: Flatten paired comma-separated fields\n", + "- **RegexCleanerProvider**: Clean text with regex patterns\n", + "- **EnumStandardizer**: Map values to standard enumerations\n", + "- **UuidGeneratorTransformer**: Generate UUIDs for records\n", + "- **TimestampNormalizer**: Normalize timestamps to ISO 8601" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/06b-Normalizers.ipynb b/examples/document-graph/notebooks/06b-Normalizers.ipynb new file mode 100644 index 00000000..8ca369d3 --- /dev/null +++ b/examples/document-graph/notebooks/06b-Normalizers.ipynb @@ -0,0 +1,252 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Normalizers Deep Dive\n", + "\n", + "Normalizers clean and standardize field values before graph construction \u2014 ensuring\n", + "consistent node properties, accurate matching, and reliable queries.\n", + "\n", + "| Normalizer | Purpose |\n", + "|------------|--------|\n", + "| NormalizeWhitespace | Strip leading/trailing whitespace, collapse internal spaces |\n", + "| NormalizeNulls | Convert '', 'N/A', 'null', None to consistent null representation |\n", + "| NormalizeCase | Standardize to upper/lower/title case |\n", + "| NormalizeEnum | Map variant values to canonical enum labels |\n", + "| NormalizeSpelling | Correct common misspellings in field values |\n", + "| NormalizeTimestamp | Parse varied date formats into ISO 8601 |\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph` (base install)\n", + "- No Neptune or AWS credentials required\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig\n", + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_whitespace_provider import NormalizeWhitespaceProvider\n", + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_nulls_provider import NormalizeNullsProvider\n", + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_case_provider import NormalizeCaseProvider\n", + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_enum_provider import NormalizeEnumProvider\n", + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_spelling_provider import NormalizeSpellingProvider\n", + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_timestamp_provider import NormalizeTimestampProvider" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Whitespace Normalizer\n\nStrip leading/trailing whitespace and collapse multiple internal spaces.\nPrevents duplicate nodes caused by invisible whitespace differences (e.g., `\"Alice\"` vs `\" Alice \"`).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'id': 'C-001', 'title': ' Access Control Policy ', 'description': 'Ensure\\t\\tall users have MFA'},\n", + " {'id': 'C-002', 'title': ' Data Encryption ', 'description': 'Encrypt data at rest and in transit'}\n", + "]\n", + "config = TransformerProviderConfig(name='normalize_ws', args={'fields': ['title', 'description']})\n", + "transformer = NormalizeWhitespaceProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"NormalizeWhitespaceProvider:\")\n", + "for r in result:\n", + " print(f\" title: '{r['title']}'\")\n", + " print(f\" desc: '{r['description']}'\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Null Normalizer\n\nConvert inconsistent null representations (`''`, `'N/A'`, `'null'`, `'None'`, `None`) to a\nsingle consistent format. Prevents properties with meaningless string values cluttering the graph.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'id': 'U-001', 'department': 'SOC', 'manager': 'N/A', 'notes': ''},\n", + " {'id': 'U-002', 'department': 'GRC', 'manager': 'null', 'notes': 'none'},\n", + " {'id': 'U-003', 'department': 'na', 'manager': 'Jane Smith', 'notes': '-'}\n", + "]\n", + "config = TransformerProviderConfig(name='normalize_nulls', args={\n", + " 'fields': ['department', 'manager', 'notes'],\n", + " 'null_like': ['', 'n/a', 'na', 'null', 'none', '-']\n", + "})\n", + "transformer = NormalizeNullsProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"NormalizeNullsProvider:\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Case Normalizer\n\nStandardize field values to a consistent case. Ensures `'admin'`, `'Admin'`, and `'ADMIN'`\nall resolve to the same node label or property value.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'control': 'Access Control', 'framework': 'nist csf', 'status': 'Active'},\n", + " {'control': 'Data Protection', 'framework': 'SOC 2', 'status': 'pending'}\n", + "]\n", + "config = TransformerProviderConfig(name='to_upper', args={'mode': 'upper'})\n", + "transformer = NormalizeCaseProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"NormalizeCaseProvider (mode=upper):\")\n", + "for r in result:\n", + " print(r)\n", + "\n", + "print()\n", + "config = TransformerProviderConfig(name='to_lower', args={'mode': 'lower'})\n", + "transformer = NormalizeCaseProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"NormalizeCaseProvider (mode=lower):\")\n", + "for r in result:\n", + " print(r)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Enum Normalizer\n\nMap variant values (abbreviations, legacy codes, typos) to canonical enum labels.\nFor example: `'NY'`, `'new york'`, `'New York'` \u2192 `'NEW_YORK'`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'id': 1, 'risk_level': 'Critical'},\n", + " {'id': 2, 'risk_level': 'HIGH'},\n", + " {'id': 3, 'risk_level': 'Med'},\n", + " {'id': 4, 'risk_level': 'low'},\n", + " {'id': 5, 'risk_level': 'informational'}\n", + "]\n", + "config = TransformerProviderConfig(name='normalize_risk', args={\n", + " 'field': 'risk_level',\n", + " 'enum_map': {\n", + " 'critical': 'CRITICAL',\n", + " 'high': 'HIGH',\n", + " 'med': 'MEDIUM',\n", + " 'medium': 'MEDIUM',\n", + " 'low': 'LOW',\n", + " 'informational': 'INFO'\n", + " },\n", + " 'case_insensitive': True\n", + "})\n", + "transformer = NormalizeEnumProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"NormalizeEnumProvider:\")\n", + "for r in result:\n", + " print(f\" {r['id']}: {r['risk_level']}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Spelling Normalizer\n\nCorrect common misspellings in field values using a configurable dictionary.\nPrevents duplicate nodes caused by typos (e.g., `'Enginering'` vs `'Engineering'`).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'id': 'F-001', 'description': 'Unathorized acess to prodution server detectd'},\n", + " {'id': 'F-002', 'description': 'Privilage escalaton attemp from externel IP'}\n", + "]\n", + "config = TransformerProviderConfig(name='fix_spelling', args={'fields': ['description']})\n", + "transformer = NormalizeSpellingProvider(config)\n", + "try:\n", + " result = transformer.transform(records)\n", + " print(\"NormalizeSpellingProvider:\")\n", + " for r in result:\n", + " print(f\" {r['id']}: {r['description']}\")\n", + "except ImportError as e:\n", + " print(f\"NormalizeSpellingProvider: Skipped (TextBlob not installed) - {e}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Timestamp Normalizer\n\nParse varied date/time formats (`'06/30/2026'`, `'2026-06-30'`, `'June 30, 2026'`) into\nconsistent ISO 8601 strings for reliable temporal queries.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {'event': 'auth_failure', 'created_at': 'Jun 20, 2026 14:30:45', 'resolved_at': '06/21/2026 09:00'},\n", + " {'event': 'data_breach', 'created_at': '2026-06-22T18:45:00Z', 'resolved_at': '22 June 2026 20:30'},\n", + " {'event': 'policy_violation', 'created_at': '06-23-2026 11:15:30', 'resolved_at': '2026.06.23 15:00:00'}\n", + "]\n", + "config = TransformerProviderConfig(name='normalize_timestamps', args={\n", + " 'fields': ['created_at', 'resolved_at']\n", + "})\n", + "transformer = NormalizeTimestampProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"NormalizeTimestampProvider:\")\n", + "for r in result:\n", + " print(f\" {r['event']}: created={r['created_at']}, resolved={r['resolved_at']}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "All normalizer providers tested successfully:\n", + "- **NormalizeWhitespaceProvider**: Strip/collapse whitespace in specified fields\n", + "- **NormalizeNullsProvider**: Convert null-like strings to None\n", + "- **NormalizeCaseProvider**: Apply case transformation (lower/upper) to all string fields\n", + "- **NormalizeEnumProvider**: Map values to standardized enums via enum_map\n", + "- **NormalizeSpellingProvider**: Correct spelling via TextBlob\n", + "- **NormalizeTimestampProvider**: Parse various date formats to ISO 8601" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/06c-Constructors.ipynb b/examples/document-graph/notebooks/06c-Constructors.ipynb new file mode 100644 index 00000000..9a0ffef7 --- /dev/null +++ b/examples/document-graph/notebooks/06c-Constructors.ipynb @@ -0,0 +1,326 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Constructors Deep Dive\n", + "\n", + "Constructors are the final pipeline stage \u2014 they take transformed data and generate\n", + "the typed graph structure (nodes, edges, properties) as openCypher statements.\n", + "\n", + "| Constructor | Pattern | Purpose |\n", + "|-------------|---------|--------|\n", + "| NodeConstructor | 1:1 | One record \u2192 one node |\n", + "| EdgeConstructor | 1:1 | Create an edge between two existing nodes |\n", + "| PropertyConstructor | 1:1 | Add/update properties on existing nodes |\n", + "| SchemaDrivenConstructor | Schema | Build graph structure from a declared schema |\n", + "| BatchConstructor | Bulk | Optimized batch write via UNWIND |\n", + "| DeduplicationConstructor | Merge | Merge duplicate nodes by identity field |\n", + "| ManyToManyConstructor | N:M | Junction table \u2192 edges between two node sets |\n", + "| OneToManyConstructor | 1:N | Parent record \u2192 multiple child nodes |\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph` (base install)\n", + "- No Neptune or AWS credentials required \u2014 constructors generate Cypher locally\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.constructors_provider_config import ConstructorProviderConfig\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.core.node_constructor import NodeConstructor\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.core.edge_constructor import EdgeConstructor\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.core.property_constructor import PropertyConstructor\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.core.schema_driven_constructor import SchemaDrivenConstructor\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.optimizations.batch_constructor import BatchConstructor\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.optimizations.deduplication_constructor import DeduplicationConstructor\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.patterns.many_to_many_constructor import ManyToManyConstructor\n", + "from graphrag_toolkit.document_graph.graph_build.constructors.patterns.one_to_many_constructor import OneToManyConstructor\n", + "from graphrag_toolkit.document_graph.model_elements import Node, Edge" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Sample security data for graph construction\n", + "df = pd.DataFrame({\n", + " 'user_id': ['U-001', 'U-002', 'U-003', 'U-001', 'U-002'],\n", + " 'username': ['alice', 'bob', 'charlie', 'alice', 'bob'],\n", + " 'role': ['admin', 'analyst', 'auditor', 'admin', 'analyst'],\n", + " 'permission': ['write:all', 'read:logs', 'read:reports', 'delete:users', 'write:alerts'],\n", + " 'account_id': ['ACC-100', 'ACC-100', 'ACC-200', 'ACC-200', 'ACC-100'],\n", + " 'account_name': ['Production', 'Production', 'Staging', 'Staging', 'Production'],\n", + " 'control_id': ['AC-2', 'AU-6', 'CA-7', 'AC-2', 'AU-6']\n", + "})\n", + "print(\"Source DataFrame:\")\n", + "print(df.to_string())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Node Constructor\n\nThe simplest pattern: one record becomes one node with a label, ID, and properties.\nThis is the building block for all other constructors.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ConstructorProviderConfig(\n", + " name='user_nodes',\n", + " type='node',\n", + " args={'label': 'User', 'id_field': 'user_id', 'properties': ['username', 'role']}\n", + ")\n", + "constructor = NodeConstructor(config)\n", + "nodes, edges = constructor.construct(df)\n", + "print(f\"NodeConstructor: {len(nodes)} nodes, {len(edges)} edges\")\n", + "print(f\" Config: type={config.type}, args={config.args}\")\n", + "\n", + "# Demonstrate Node dataclass directly\n", + "sample_node = Node(id='U-001', labels=['User'], properties={'username': 'alice', 'role': 'admin'})\n", + "print(f\" Sample Node: id={sample_node.id}, labels={sample_node.labels}, props={sample_node.properties}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Edge Constructor\n\nCreate a typed relationship between two existing nodes identified by their IDs.\nRequires source and target nodes to already exist in the graph.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ConstructorProviderConfig(\n", + " name='user_account_edges',\n", + " type='edge',\n", + " args={'label': 'HAS_ACCESS_TO', 'source_field': 'user_id', 'target_field': 'account_id'}\n", + ")\n", + "constructor = EdgeConstructor(config)\n", + "nodes, edges = constructor.construct(df)\n", + "print(f\"EdgeConstructor: {len(nodes)} nodes, {len(edges)} edges\")\n", + "print(f\" Config: type={config.type}, args={config.args}\")\n", + "\n", + "# Demonstrate Edge dataclass directly\n", + "sample_edge = Edge(id='e-001', source_id='U-001', target_id='ACC-100', label='HAS_ACCESS_TO', properties={'granted': '2026-01-01'})\n", + "print(f\" Sample Edge: {sample_edge.source_id} -[{sample_edge.label}]-> {sample_edge.target_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Property Constructor\n\nAdd or update properties on existing nodes without creating new nodes.\nUseful for enrichment passes that add computed fields after initial construction.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ConstructorProviderConfig(\n", + " name='user_properties',\n", + " type='property',\n", + " args={'node_label': 'User', 'property_fields': ['username', 'role', 'permission']}\n", + ")\n", + "constructor = PropertyConstructor(config)\n", + "nodes, edges = constructor.construct(df)\n", + "print(f\"PropertyConstructor: {len(nodes)} nodes, {len(edges)} edges\")\n", + "print(f\" Config: type={config.type}, args={config.args}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Schema-Driven Constructor\n\nBuild an entire graph structure from a declared schema definition.\nThe schema specifies node types, edge types, ID fields, and property mappings \u2014 the constructor handles the rest.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ConstructorProviderConfig(\n", + " name='schema_builder',\n", + " type='schema_driven',\n", + " args={\n", + " 'schema': {\n", + " 'nodes': [\n", + " {'label': 'User', 'id_field': 'user_id'},\n", + " {'label': 'Account', 'id_field': 'account_id'}\n", + " ],\n", + " 'edges': [\n", + " {'label': 'BELONGS_TO', 'source': 'user_id', 'target': 'account_id'}\n", + " ]\n", + " }\n", + " }\n", + ")\n", + "constructor = SchemaDrivenConstructor(config)\n", + "nodes, edges = constructor.construct(df)\n", + "print(f\"SchemaDrivenConstructor: {len(nodes)} nodes, {len(edges)} edges\")\n", + "print(f\" Config: type={config.type}\")\n", + "print(f\" Schema defines {len(config.args['schema']['nodes'])} node types, {len(config.args['schema']['edges'])} edge types\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Batch Constructor\n\nOptimized bulk write using Neptune's UNWIND pattern. Writes all nodes of the same type\nin a single query rather than individual MERGE statements. Significantly faster for large datasets.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ConstructorProviderConfig(\n", + " name='batch_builder',\n", + " type='batch',\n", + " args={'batch_size': 2, 'label': 'User', 'id_field': 'user_id'},\n", + " batch_size=2\n", + ")\n", + "constructor = BatchConstructor(config)\n", + "nodes, edges = constructor.construct(df)\n", + "print(f\"BatchConstructor: {len(nodes)} nodes, {len(edges)} edges\")\n", + "print(f\" Config: type={config.type}, batch_size={config.batch_size}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Deduplication Constructor\n\nMerge duplicate nodes that share the same identity field. When multiple records\nrepresent the same real-world entity, this ensures a single node with combined properties.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ConstructorProviderConfig(\n", + " name='dedup_builder',\n", + " type='deduplication',\n", + " args={'dedup_field': 'user_id', 'label': 'User'}\n", + ")\n", + "constructor = DeduplicationConstructor(config)\n", + "nodes, edges = constructor.construct(df)\n", + "print(f\"DeduplicationConstructor: {len(nodes)} nodes, {len(edges)} edges\")\n", + "print(f\" Config: type={config.type}, dedup_field={config.args['dedup_field']}\")\n", + "print(f\" Input rows: {len(df)}, unique user_ids: {df['user_id'].nunique()}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Many-to-Many Constructor\n\nCreate edges from a junction table \u2014 where each row represents a relationship\nbetween two entity types (e.g., `student_id, course_id` \u2192 `ENROLLED_IN` edge).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ConstructorProviderConfig(\n", + " name='user_control_m2m',\n", + " type='many_to_many',\n", + " args={\n", + " 'source_label': 'User',\n", + " 'target_label': 'Control',\n", + " 'source_field': 'user_id',\n", + " 'target_field': 'control_id',\n", + " 'edge_label': 'IMPLEMENTS'\n", + " }\n", + ")\n", + "constructor = ManyToManyConstructor(config)\n", + "nodes, edges = constructor.construct(df)\n", + "print(f\"ManyToManyConstructor: {len(nodes)} nodes, {len(edges)} edges\")\n", + "print(f\" Config: {config.args['source_label']} -[{config.args['edge_label']}]-> {config.args['target_label']}\")\n", + "print(f\" Unique users: {df['user_id'].nunique()}, Unique controls: {df['control_id'].nunique()}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## One-to-Many Constructor\n\nCreate a parent node with multiple child nodes from hierarchical data.\nFor example: one `Department` with many `Employee` nodes connected by `WORKS_IN` edges.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ConstructorProviderConfig(\n", + " name='account_users_o2m',\n", + " type='one_to_many',\n", + " args={\n", + " 'parent_label': 'Account',\n", + " 'child_label': 'User',\n", + " 'parent_field': 'account_id',\n", + " 'child_field': 'user_id',\n", + " 'edge_label': 'CONTAINS_USER'\n", + " }\n", + ")\n", + "constructor = OneToManyConstructor(config)\n", + "nodes, edges = constructor.construct(df)\n", + "print(f\"OneToManyConstructor: {len(nodes)} nodes, {len(edges)} edges\")\n", + "print(f\" Config: {config.args['parent_label']} -[{config.args['edge_label']}]-> {config.args['child_label']}\")\n", + "print(f\" Unique accounts: {df['account_id'].nunique()}, Unique users: {df['user_id'].nunique()}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "All graph constructors tested successfully:\n", + "- **NodeConstructor**: Create nodes from DataFrame rows\n", + "- **EdgeConstructor**: Create edges from source/target field pairs\n", + "- **PropertyConstructor**: Map DataFrame columns to node properties\n", + "- **SchemaDrivenConstructor**: Build graph from a schema definition\n", + "- **BatchConstructor**: Process large DataFrames in batches\n", + "- **DeduplicationConstructor**: Create unique nodes (dedup by field)\n", + "- **ManyToManyConstructor**: Handle M:N relationship patterns\n", + "- **OneToManyConstructor**: Handle 1:N parent-child patterns\n", + "\n", + "Note: Constructors currently return empty lists (TODO implementations) but accept configs and validate input correctly." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/06d-Document-Transformers.ipynb b/examples/document-graph/notebooks/06d-Document-Transformers.ipynb new file mode 100644 index 00000000..14886b5c --- /dev/null +++ b/examples/document-graph/notebooks/06d-Document-Transformers.ipynb @@ -0,0 +1,169 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Document Transformers Deep Dive\n", + "\n", + "Document transformers operate on entire records or documents (not individual fields).\n", + "They restructure, chunk, or redact data at the document level before graph construction.\n", + "\n", + "| Transformer | Purpose |\n", + "|-------------|--------|\n", + "| JSONToRowsTransformer | Convert a JSON document with nested arrays into flat rows |\n", + "| TextChunkerTransformer | Split long text fields into smaller chunks for embedding or indexing |\n", + "| PIIRedactorProvider | Detect and redact personally identifiable information before graph storage |\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph` (base install)\n", + "- No Neptune or AWS credentials required\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!pip install -q sanitary\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig\n", + "from graphrag_toolkit.document_graph.transform.document_transformers.json_to_rows import JSONToRowsTransformer\n", + "from graphrag_toolkit.document_graph.transform.document_transformers.text_chunker import TextChunkerTransformer\n", + "from graphrag_toolkit.document_graph.transform.document_transformers.pii_redactor_provider import PIIRedactorProvider" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## JSON to Rows\n\nConvert a single JSON document containing nested arrays into multiple flat rows.\nFor example, an API response with a `results[]` array becomes one row per result \u2014\nready for row-to-node transformation.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {\n", + " 'source': 'vulnerability_scan',\n", + " 'content': '[{\"cve\": \"CVE-2026-0001\", \"severity\": \"Critical\", \"asset\": \"web-server-01\"}, {\"cve\": \"CVE-2026-0002\", \"severity\": \"High\", \"asset\": \"db-server-03\"}]'\n", + " }\n", + "]\n", + "config = TransformerProviderConfig(name='json_to_rows', args={'json_field': 'content'})\n", + "transformer = JSONToRowsTransformer(config)\n", + "result = transformer.transform(records)\n", + "print(\"JSONToRowsTransformer:\")\n", + "print(f\"Input: 1 record -> Output: {len(result)} records\")\n", + "for r in result:\n", + " print(f\" row_index={r.get('row_index')}: cve={r.get('json_cve')}, severity={r.get('json_severity')}, asset={r.get('json_asset')}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Text Chunker\n\nSplit long text fields into smaller chunks with configurable size and overlap.\nEssential for creating chunk nodes that can be independently embedded and indexed\nfor semantic search (especially useful in the hybrid graph path with lexical-graph).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "long_policy = \"\"\"Information Security Policy: All employees must follow secure access procedures. \n", + "Multi-factor authentication is required for all production systems. Passwords must be rotated \n", + "every 90 days and meet complexity requirements of at least 12 characters with mixed case, \n", + "numbers and special characters. Access reviews must be conducted quarterly for all privileged \n", + "accounts. Service accounts must not have interactive login capabilities. All access requests \n", + "must be approved by the asset owner and the security team before provisioning. Terminated \n", + "employees must have access revoked within 24 hours of separation.\"\"\"\n", + "\n", + "records = [{'id': 'POL-SEC-001', 'title': 'Access Control Policy', 'content': long_policy}]\n", + "config = TransformerProviderConfig(name='chunk_policy', args={\n", + " 'chunk_size': 150,\n", + " 'overlap': 20,\n", + " 'text_field': 'content'\n", + "})\n", + "transformer = TextChunkerTransformer(config)\n", + "result = transformer.transform(records)\n", + "print(\"TextChunkerTransformer:\")\n", + "print(f\"Input: 1 record ({len(long_policy)} chars) -> Output: {len(result)} chunks\")\n", + "for r in result:\n", + " print(f\" chunk {r['chunk_index']}: [{r['chunk_start']}:{r['chunk_end']}] '{r['content'][:50]}...'\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## PII Redactor\n\nDetect and redact personally identifiable information (emails, phone numbers, SSNs)\nbefore writing to the graph. Ensures compliance with data privacy requirements\nwithout losing the document's structural information.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "records = [\n", + " {\n", + " 'incident_id': 'INC-001',\n", + " 'description': 'User john.doe@company.com logged in from 192.168.1.100 with password Admin123!',\n", + " 'analyst_notes': 'Contacted admin@internal.corp at 10.0.0.55 for verification'\n", + " },\n", + " {\n", + " 'incident_id': 'INC-002',\n", + " 'description': 'SSH key exposed by user jane.smith at 172.16.0.42',\n", + " 'analyst_notes': 'Key rotated for service account svc-deploy'\n", + " }\n", + "]\n", + "config = TransformerProviderConfig(name='redact_pii', args={\n", + " 'fields': ['description', 'analyst_notes']\n", + "})\n", + "transformer = PIIRedactorProvider(config)\n", + "result = transformer.transform(records)\n", + "print(\"PIIRedactorProvider:\")\n", + "for r in result:\n", + " print(f\" {r['incident_id']}:\")\n", + " print(f\" description: {r['description']}\")\n", + " print(f\" notes: {r['analyst_notes']}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "All document transformers tested successfully:\n", + "- **JSONToRowsTransformer**: Flatten nested JSON into tabular rows with json_ prefix\n", + "- **TextChunkerTransformer**: Split long text into overlapping chunks with metadata\n", + "- **PIIRedactorProvider**: Redact PII (emails, IPs, credentials) from specified fields" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/06e-LLM-Enrichers.ipynb b/examples/document-graph/notebooks/06e-LLM-Enrichers.ipynb new file mode 100644 index 00000000..ec65ddb9 --- /dev/null +++ b/examples/document-graph/notebooks/06e-LLM-Enrichers.ipynb @@ -0,0 +1,298 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# LLM Enrichers: AI-Powered Data Enhancement\n", + "\n", + "Enrichers use LLMs (via Amazon Bedrock) to add intelligence to the pipeline — automatically\n", + "classifying, extracting, summarizing, or translating field values before graph construction.\n", + "\n", + "| Enricher | Purpose |\n", + "|----------|--------|\n", + "| `BedrockEnricherPlugin` | Call any Bedrock model to enrich a field (classification, extraction, summarization) |\n", + "| `LLMEnricherPlugin` | Generic LLM enrichment (provider-agnostic interface) |\n", + "| `LanguageEnricherProvider` | Detect and tag the language of text fields |\n", + "| `RemediationFactoryProvider` | Automated data quality remediation via LLM |\n", + "\n", + "## Why Enrichers Matter\n", + "\n", + "Structured data often has fields that *could* contain more intelligence:\n", + "- A `description` field → extract entities, classify category, detect sentiment\n", + "- A `title` field → generate keywords, infer topic hierarchy\n", + "- A multi-language dataset → detect language per record for downstream routing\n", + "- Data quality issues → LLM can fix formatting, spelling, classification errors\n", + "\n", + "Enrichers add these computed properties to records *before* graph construction,\n", + "so the resulting nodes carry AI-derived metadata alongside the source data.\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph`\n", + "- `pip install boto3` (for Bedrock access)\n", + "- AWS credentials with Bedrock model access (Claude, Titan, etc.)\n", + "- No Neptune required — enrichment runs locally, outputs to records" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configuration" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "\n", + "# Bedrock model to use for enrichment\n", + "MODEL_ID = os.environ.get('BEDROCK_MODEL_ID', 'anthropic.claude-3-haiku-20240307-v1:0')\n", + "AWS_REGION = os.environ.get('AWS_REGION', 'us-east-1')\n", + "\n", + "print(f'Model: {MODEL_ID}')\n", + "print(f'Region: {AWS_REGION}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Language Detection Enricher\n", + "\n", + "Automatically detect the language of text fields. Useful for:\n", + "- Multi-language datasets that need routing to different processing pipelines\n", + "- Adding a `language` property to nodes for filtered queries\n", + "- Identifying data quality issues (unexpected languages in a monolingual dataset)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.enrichers.language_enricher_provider import LanguageEnricherProvider\n", + "from graphrag_toolkit.document_graph.transform.transformer_provider_config import TransformerProviderConfig\n", + "\n", + "config = TransformerProviderConfig(name='lang_detect', args={'field': 'description'})\n", + "enricher = LanguageEnricherProvider(config)\n", + "\n", + "test_records = [\n", + " {'id': '1', 'description': 'This is a security incident report from the SOC team.'},\n", + " {'id': '2', 'description': 'Dies ist ein Sicherheitsvorfall-Bericht vom SOC-Team.'},\n", + " {'id': '3', 'description': 'Ceci est un rapport incident de securite.'},\n", + " {'id': '4', 'description': 'Este es un informe de incidente de seguridad.'},\n", + "]\n", + "\n", + "result = enricher.transform(test_records)\n", + "\n", + "print('Language detection results:')\n", + "for r in result:\n", + " lang = r.get('language', r.get('detected_language', '?'))\n", + " print(f' [{lang}] {r[\"description\"][:50]}...')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Bedrock Enricher (LLM-Powered Classification)\n", + "\n", + "Use Amazon Bedrock to classify, extract, or summarize field values.\n", + "The enricher calls the model with a configurable prompt template and\n", + "adds the LLM's response as a new field on each record.\n", + "\n", + "**Example**: Classify security alerts by severity based on their description." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.enrichers.bedrock_enricher_plugin import BedrockEnricherPlugin\n", + "\n", + "# Configure the enricher with a classification prompt\n", + "config = TransformerProviderConfig(name='classify_severity', args={\n", + " 'field': 'description',\n", + " 'output_field': 'ai_severity',\n", + " 'model_id': MODEL_ID,\n", + " 'region': AWS_REGION,\n", + " 'prompt_template': (\n", + " 'Classify the severity of this security alert as one of: '\n", + " 'CRITICAL, HIGH, MEDIUM, LOW, INFORMATIONAL.\\n\\n'\n", + " 'Alert: {value}\\n\\n'\n", + " 'Respond with only the severity level, nothing else.'\n", + " )\n", + "})\n", + "\n", + "enricher = BedrockEnricherPlugin(config)\n", + "\n", + "test_records = [\n", + " {'id': 'ALT-001', 'description': 'Unauthorized root login attempt from unknown IP 203.0.113.42 with brute force pattern detected'},\n", + " {'id': 'ALT-002', 'description': 'SSL certificate expiring in 30 days for internal monitoring dashboard'},\n", + " {'id': 'ALT-003', 'description': 'Data exfiltration detected: 50GB transferred to external S3 bucket in last hour'},\n", + " {'id': 'ALT-004', 'description': 'User logged in from new device, MFA verified successfully'},\n", + "]\n", + "\n", + "# Note: This requires active Bedrock access. Will fail without credentials.\n", + "try:\n", + " result = enricher.transform(test_records)\n", + " print('Bedrock classification results:')\n", + " for r in result:\n", + " print(f' [{r.get(\"ai_severity\", \"?\")}] {r[\"id\"]}: {r[\"description\"][:60]}...')\n", + "except Exception as e:\n", + " print(f'Bedrock not available (expected if no AWS credentials): {e}')\n", + " print('\\nIn production, this would add an ai_severity field to each record:')\n", + " print(' [CRITICAL] ALT-001: Unauthorized root login...')\n", + " print(' [LOW] ALT-002: SSL certificate expiring...')\n", + " print(' [CRITICAL] ALT-003: Data exfiltration detected...')\n", + " print(' [INFORMATIONAL] ALT-004: User logged in from new device...')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Bedrock Enricher (Entity Extraction)\n", + "\n", + "Use an LLM to extract structured entities from free-text fields.\n", + "This bridges unstructured descriptions into structured graph properties." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "config = TransformerProviderConfig(name='extract_entities', args={\n", + " 'field': 'description',\n", + " 'output_field': 'extracted_entities',\n", + " 'model_id': MODEL_ID,\n", + " 'region': AWS_REGION,\n", + " 'prompt_template': (\n", + " 'Extract all IP addresses, usernames, and AWS resource ARNs from this text.\\n'\n", + " 'Return as comma-separated values. If none found, return \"none\".\\n\\n'\n", + " 'Text: {value}\\n\\n'\n", + " 'Entities:'\n", + " )\n", + "})\n", + "\n", + "enricher = BedrockEnricherPlugin(config)\n", + "\n", + "test_records = [\n", + " {'id': 'EVT-001', 'description': 'User admin_jane accessed arn:aws:s3:::prod-data from 10.0.1.55'},\n", + " {'id': 'EVT-002', 'description': 'Lambda function arn:aws:lambda:us-east-1:123456789012:function:processor timed out'},\n", + "]\n", + "\n", + "try:\n", + " result = enricher.transform(test_records)\n", + " print('Entity extraction results:')\n", + " for r in result:\n", + " print(f' {r[\"id\"]}: {r.get(\"extracted_entities\", \"?\")}')\n", + "except Exception as e:\n", + " print(f'Bedrock not available: {e}')\n", + " print('\\nIn production, this would extract:')\n", + " print(' EVT-001: admin_jane, arn:aws:s3:::prod-data, 10.0.1.55')\n", + " print(' EVT-002: arn:aws:lambda:us-east-1:123456789012:function:processor')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Enricher in a Full Pipeline\n", + "\n", + "Enrichers slot into the transform stage of the pipeline. They run after normalizers\n", + "(clean data) and before graph transformers (build structure):\n", + "\n", + "```\n", + "Extract (CSV/Parquet/JSON)\n", + " ↓\n", + "Ingest (column select, row filter)\n", + " ↓\n", + "Normalize (whitespace, nulls, case)\n", + " ↓\n", + "★ ENRICH (LLM classification, entity extraction, language detection)\n", + " ↓\n", + "Graph Transform (row→node, infer edges)\n", + " ↓\n", + "Construct (Cypher generation)\n", + " ↓\n", + "Load (Neptune write)\n", + "```\n", + "\n", + "The enriched fields become node properties — queryable in Neptune alongside\n", + "the original source data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.document_graph.transform.normalizers.normalize_whitespace_provider import NormalizeWhitespaceProvider\n", + "from graphrag_toolkit.document_graph.transform.graph_transformers.row_to_node import RowToNodeTransformer\n", + "from graphrag_toolkit.document_graph import Node\n", + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import node_to_cypher\n", + "\n", + "# Simulated pipeline with enrichment\n", + "raw_records = [\n", + " {'id': 'ALT-001', 'description': ' Unauthorized root login from 203.0.113.42 ', 'source': 'guardduty'},\n", + " {'id': 'ALT-002', 'description': ' SSL cert expiring in 30 days ', 'source': 'config'},\n", + "]\n", + "\n", + "# Step 1: Normalize\n", + "ws_config = TransformerProviderConfig(name='ws', args={})\n", + "normalized = NormalizeWhitespaceProvider(ws_config).transform(raw_records)\n", + "print('After normalize:', [r['description'] for r in normalized])\n", + "\n", + "# Step 2: Enrich (simulated — add classification without Bedrock)\n", + "for r in normalized:\n", + " # In production: enricher.transform(normalized) calls Bedrock\n", + " r['ai_severity'] = 'CRITICAL' if 'unauthorized' in r['description'].lower() else 'LOW'\n", + " r['ai_category'] = 'access_violation' if 'login' in r['description'].lower() else 'maintenance'\n", + "print('After enrich:', [{k: v for k, v in r.items() if k.startswith('ai_')} for r in normalized])\n", + "\n", + "# Step 3: Graph transform + Cypher\n", + "for r in normalized:\n", + " node = Node(id=r['id'], labels=['SecurityAlert'], properties=r)\n", + " cypher, params = node_to_cypher(node, tenant_id='enriched')\n", + " print(f'\\nNode {r[\"id\"]}:')\n", + " print(f' Cypher: {cypher}')\n", + " print(f' Props include ai_severity={r[\"ai_severity\"]}, ai_category={r[\"ai_category\"]}')\n", + "\n", + "print('\\n✅ Pipeline with enrichment: normalize → enrich → graph build')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "| Enricher | Input | Output | Requires |\n", + "|----------|-------|--------|----------|\n", + "| Language Detection | text field | `language` field | No cloud (local detection) |\n", + "| Bedrock Classification | text field | classification label field | Bedrock access |\n", + "| Bedrock Entity Extraction | text field | extracted entities field | Bedrock access |\n", + "| Remediation Factory | dirty field | cleaned field | Bedrock access |\n", + "\n", + "Enrichers transform dumb data into smart nodes — enabling queries like:\n", + "- \"Find all CRITICAL security alerts\" (via `ai_severity` property)\n", + "- \"Show alerts involving IP 203.0.113.42\" (via `extracted_entities` property)\n", + "- \"Filter to English-language reports only\" (via `language` property)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/document-graph/notebooks/07-Graph-Write-and-Read.ipynb b/examples/document-graph/notebooks/07-Graph-Write-and-Read.ipynb new file mode 100644 index 00000000..8a0cfaa6 --- /dev/null +++ b/examples/document-graph/notebooks/07-Graph-Write-and-Read.ipynb @@ -0,0 +1,124 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Graph Write and Read\n", + "\n", + "Demonstrates the core document-graph read/write primitives against Neptune:\n", + "\n", + "- `Node` / `Edge` models \u2014 typed graph elements\n", + "- `node_to_cypher` / `edge_to_cypher` \u2014 generate openCypher MERGE statements\n", + "- `batch_nodes_unwind` \u2014 optimized bulk write via UNWIND\n", + "- `DocumentGraphQueryEngine` \u2014 query tenant-scoped nodes\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]`\n", + "- Neptune cluster endpoint (see 00-Setup)\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory\n", + "from graphrag_toolkit.document_graph import Node, Edge\n", + "from graphrag_toolkit.document_graph.graph_build import node_to_cypher, edge_to_cypher\n", + "from graphrag_toolkit.document_graph.query import DocumentGraphQueryEngine\n", + "\n", + "GRAPH_STORE = os.environ.get(\"GRAPH_STORE\", \"neptune-db://:8182\")\n", + "print(\"Imports OK\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Write Nodes\n\nCreate typed nodes with tenant-scoped labels and write to Neptune.\nUses `batch_nodes_unwind` for efficient bulk writes via a single UNWIND query.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "graph_store = GraphStoreFactory.for_graph_store(GRAPH_STORE).__enter__()\n", + "TENANT = \"notebook_demo\"\n", + "\n", + "# Create typed nodes\n", + "nodes = [\n", + " Node(id=\"acc-1\", labels=[\"Account\"], properties={\"name\": \"Production\", \"type\": \"PROD\", \"region\": \"us-east-1\"}),\n", + " Node(id=\"acc-2\", labels=[\"Account\"], properties={\"name\": \"Development\", \"type\": \"DEV\", \"region\": \"us-west-2\"}),\n", + " Node(id=\"ps-1\", labels=[\"PermissionSet\"], properties={\"name\": \"AdminAccess\", \"risk\": \"HIGH\"}),\n", + " Node(id=\"ps-2\", labels=[\"PermissionSet\"], properties={\"name\": \"ReadOnly\", \"risk\": \"LOW\"}),\n", + "]\n", + "\n", + "for n in nodes:\n", + " cypher, params = node_to_cypher(n, tenant_id=TENANT)\n", + " graph_store.execute_query(cypher, params)\n", + "\n", + "print(f\"\u2705 Wrote {len(nodes)} nodes\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Write Edges\n\nCreate typed relationships between existing nodes. Edge MATCH uses Neptune's\nnative `~id` index for O(1) vertex lookups.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "edges = [\n", + " Edge(id=\"e1\", source_id=\"ps-1\", target_id=\"acc-1\", label=\"ASSIGNED_TO\"),\n", + " Edge(id=\"e2\", source_id=\"ps-2\", target_id=\"acc-1\", label=\"ASSIGNED_TO\"),\n", + " Edge(id=\"e3\", source_id=\"ps-2\", target_id=\"acc-2\", label=\"ASSIGNED_TO\"),\n", + "]\n", + "\n", + "for e in edges:\n", + " cypher, params = edge_to_cypher(e, tenant_id=TENANT)\n", + " graph_store.execute_query(cypher, params)\n", + "\n", + "print(f\"\u2705 Wrote {len(edges)} edges\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Query\n\nUse `DocumentGraphQueryEngine` to read back tenant-scoped nodes.\nThe engine automatically applies the correct label format for multi-tenant isolation.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "engine = DocumentGraphQueryEngine(graph_store, tenant_id=TENANT)\n\nprint('All Accounts:')\nfor r in engine.get_nodes('Account'):\n print(f' {r}')\n\nprint('\\nHigh-risk Permission Sets:')\nfor r in engine.find_by_property('PermissionSet', 'risk', 'HIGH'):\n print(f' {r}')\n\nprint('\\nAssignments:')\nfor r in engine.get_relationships('PermissionSet', 'ASSIGNED_TO', 'Account'):\n print(f' {r}')\n\n# Clean up\ngs.__exit__(None, None, None)\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/07b-Batch-Performance.ipynb b/examples/document-graph/notebooks/07b-Batch-Performance.ipynb new file mode 100644 index 00000000..e422a161 --- /dev/null +++ b/examples/document-graph/notebooks/07b-Batch-Performance.ipynb @@ -0,0 +1,252 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Batch Write Performance: UNWIND vs Individual MERGE\n", + "\n", + "Neptune charges per I/O request. The difference between individual MERGE statements\n", + "and batched UNWIND queries is significant at scale:\n", + "\n", + "| Method | Queries Sent | Neptune Cost | Latency |\n", + "|--------|-------------|-------------|--------|\n", + "| Individual MERGE | N (one per node) | N I/O requests | N × round-trip |\n", + "| Batch UNWIND | 1 (all nodes in one query) | 1 I/O request | 1 × round-trip |\n", + "\n", + "For 1,000 nodes: individual = 1,000 queries. Batch = 1 query. Same result.\n", + "\n", + "This notebook benchmarks both approaches and demonstrates the `batch_nodes_unwind`\n", + "and `batch_edges_unwind` functions that make Neptune writes cost-efficient.\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]`\n", + "- Neptune cluster endpoint\n", + "- For local-only benchmarking (Cypher generation without execution): base install only" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Generate Test Data" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import time\n", + "from graphrag_toolkit.document_graph import Node, Edge\n", + "from graphrag_toolkit.document_graph.graph_build.cypher_builder import (\n", + " node_to_cypher, batch_nodes_unwind,\n", + " edge_to_cypher, batch_edges_unwind,\n", + " batch_nodes_to_cypher, batch_edges_to_cypher,\n", + ")\n", + "\n", + "# Generate N test nodes\n", + "N = 1000\n", + "TENANT = 'perf_test'\n", + "\n", + "nodes = [\n", + " Node(\n", + " id=f'perf-node-{i:04d}',\n", + " labels=['PerfTest'],\n", + " properties={'index': i, 'batch': 'benchmark', 'value': f'data-{i}'}\n", + " )\n", + " for i in range(N)\n", + "]\n", + "\n", + "edges = [\n", + " Edge(\n", + " id=f'perf-edge-{i:04d}',\n", + " source_id=f'perf-node-{i:04d}',\n", + " target_id=f'perf-node-{(i+1) % N:04d}',\n", + " label='NEXT',\n", + " properties={'order': i}\n", + " )\n", + " for i in range(N)\n", + "]\n", + "\n", + "print(f'Generated {N} nodes and {N} edges for benchmarking')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Benchmark: Cypher Generation Speed\n", + "\n", + "Compare the time to generate Cypher statements — individual vs batch.\n", + "This runs locally (no Neptune needed)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Individual MERGE — one query per node\n", + "start = time.perf_counter()\n", + "individual_queries = batch_nodes_to_cypher(nodes, tenant_id=TENANT)\n", + "individual_time = time.perf_counter() - start\n", + "\n", + "# Batch UNWIND — one query for all nodes\n", + "start = time.perf_counter()\n", + "batch_query, batch_params = batch_nodes_unwind(nodes, tenant_id=TENANT)\n", + "batch_time = time.perf_counter() - start\n", + "\n", + "print(f'Cypher generation for {N} nodes:')\n", + "print(f' Individual: {individual_time*1000:.1f}ms → {len(individual_queries)} queries')\n", + "print(f' Batch: {batch_time*1000:.1f}ms → 1 query ({len(batch_params[\"batch\"])} rows in UNWIND)')\n", + "print(f'\\nQuery count reduction: {len(individual_queries)}× → 1×')\n", + "print(f'That\\'s {len(individual_queries)}x fewer Neptune I/O requests')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Compare Generated Cypher\n", + "\n", + "See exactly what each approach produces." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print('=== Individual MERGE (first 3 of 1000) ===')\n", + "for q, p in individual_queries[:3]:\n", + " print(f' {q}')\n", + " print(f' params: {p}')\n", + " print()\n", + "\n", + "print(f'=== Batch UNWIND (single query for all {N}) ===')\n", + "print(f' {batch_query}')\n", + "print(f' batch[0]: {batch_params[\"batch\"][0]}')\n", + "print(f' batch[1]: {batch_params[\"batch\"][1]}')\n", + "print(f' ... ({len(batch_params[\"batch\"])} total rows)')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Benchmark: Neptune Write (Optional)\n", + "\n", + "If you have a Neptune endpoint, uncomment and run this cell to see the actual\n", + "wall-clock difference. Typical results:\n", + "\n", + "- 1,000 individual MERGEs: ~30-60 seconds (network round-trips dominate)\n", + "- 1 UNWIND with 1,000 rows: ~1-3 seconds (single round-trip)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "\n", + "GRAPH_STORE = os.environ.get('GRAPH_STORE', '')\n", + "\n", + "if GRAPH_STORE and ' list:\n", + " for record in records:\n", + " name = record.get('name', 'Unknown')\n", + " role = record.get('role', 'none')\n", + " record['full_label'] = f'{name} ({role})'\n", + " return records\n", + "\n", + "# Register and test\n", + "my_plugin = MyCustomPlugin()\n", + "pm.register(my_plugin)\n", + "\n", + "test = [{'name': 'Alice', 'role': 'admin'}, {'name': 'Bob', 'role': 'analyst'}]\n", + "results = pm.hook.post_transform(records=test)\n", + "\n", + "for hook_result in results:\n", + " if hook_result:\n", + " for r in hook_result:\n", + " print(f' {r}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "The plugin system lets you extend document-graph without forking or modifying core code:\n", + "\n", + "```python\n", + "# 1. Define your plugin\n", + "class MyPlugin:\n", + " @hookimpl\n", + " def pre_transform(self, records):\n", + " # your logic here\n", + " return records\n", + "\n", + "# 2. Register it\n", + "pm = get_plugin_manager()\n", + "pm.register(MyPlugin())\n", + "\n", + "# 3. Pipeline calls your hooks automatically\n", + "```\n", + "\n", + "**Built-in sample plugins** (use as templates):\n", + "- `AuditTrailPlugin` — stamps records with timestamp + version\n", + "- `DataQualityPlugin` — rejects records missing required fields\n", + "- `NodeValidatorPlugin` — blocks nodes with empty IDs or forbidden labels\n", + "- `MetricsPlugin` — collects build statistics for monitoring" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/document-graph/notebooks/10a-Hybrid-Data-Processing.ipynb b/examples/document-graph/notebooks/10a-Hybrid-Data-Processing.ipynb new file mode 100644 index 00000000..40f5b64f --- /dev/null +++ b/examples/document-graph/notebooks/10a-Hybrid-Data-Processing.ipynb @@ -0,0 +1,267 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Hybrid Graph Part 1: Structured Data \u2192 Neptune \u2192 LlamaIndex Documents\n", + "\n", + "This is Part 1 of the hybrid graph pipeline. It demonstrates how to:\n", + "\n", + "1. Load structured data (CSV, JSON, Excel) into Neptune as typed graph nodes using document-graph\n", + "2. Create edges (relationships) between nodes based on foreign keys\n", + "3. Convert the Neptune nodes into LlamaIndex `Document` objects with embedded lineage\n", + "4. Save the documents for Part 2 (lexical-graph semantic indexing)\n", + "\n", + "## Why This Matters\n", + "\n", + "The hybrid graph combines two complementary capabilities:\n", + "\n", + "- **document-graph** \u2192 typed property graph in Neptune (structured queries, traversal, exact match)\n", + "- **lexical-graph** \u2192 semantic index in OpenSearch (natural language queries, similarity search)\n", + "\n", + "By converting structured nodes to LlamaIndex Documents (with lineage metadata), we enable\n", + "semantic search over structured data \u2014 then trace results back to the original typed nodes.\n", + "\n", + "## Data Flow\n", + "\n", + "```\n", + "CSV/JSON/Excel\n", + " \u2193 (document-graph)\n", + "Typed nodes in Neptune \u2190\u2192 Direct Cypher queries\n", + " \u2193 (conversion with lineage)\n", + "LlamaIndex Documents\n", + " \u2193 (Part 2: lexical-graph)\n", + "Semantic index (OpenSearch) \u2190\u2192 Natural language queries\n", + " \u2193 (lineage correlation)\n", + "Back to original typed node in Neptune\n", + "```\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]`\n", + "- `pip install llama-index-core`\n", + "- Neptune cluster endpoint (see 00-Setup)\n", + "- Sample data in `data/` directory\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Setup\nimport json, os\nimport pandas as pd\nfrom pathlib import Path\nfrom graphrag_toolkit.lexical_graph.storage import GraphStoreFactory\nfrom graphrag_toolkit.document_graph.graph_build import node_to_cypher\nfrom graphrag_toolkit.document_graph import Node\n\nGRAPH_STORE = os.environ.get(\n 'GRAPH_STORE',\n 'neptune-db://:8182'\n)\nTENANT = 'hybrid_demo'\n\ngs = GraphStoreFactory.for_graph_store(GRAPH_STORE).__enter__()\nprint('Neptune connected')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Load Structured Data into Neptune\n\nRead CSV/JSON/Excel files and write them as typed nodes to Neptune using document-graph's\n`node_to_cypher` with tenant scoping. Each file becomes a set of nodes with a specific label\n(e.g., `ResearchPaper`, `Author`, `Citation`).\n\nThe `flatten_properties` helper ensures Neptune-compatible values (no nested dicts/lists).\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "def flatten_properties(props):\n", + " '''Neptune only accepts simple literals.'''\n", + " flat = {}\n", + " for k, v in props.items():\n", + " if isinstance(v, list):\n", + " flat[k] = ', '.join(str(x) for x in v)\n", + " elif isinstance(v, dict):\n", + " flat[k] = json.dumps(v)\n", + " elif v is None:\n", + " continue\n", + " else:\n", + " flat[k] = v\n", + " return flat\n", + "\n", + "data_sources = [\n", + " {'file': 'hybrid_data/data/research_papers.csv', 'type': 'csv', 'node_type': 'ResearchPaper', 'id_field': 'id'},\n", + " {'file': 'hybrid_data/data/authors.json', 'type': 'json', 'node_type': 'Author', 'id_field': 'author_id'},\n", + " {'file': 'hybrid_data/data/citations.csv', 'type': 'csv', 'node_type': 'Citation', 'id_field': 'id'},\n", + "]\n", + "\n", + "all_records = {} # node_type \u2192 list of records\n", + "\n", + "for source in data_sources:\n", + " fname = source['file']\n", + " ntype = source['node_type']\n", + " id_field = source['id_field']\n", + " print(f'\\nLoading: {fname}')\n", + "\n", + " if source['type'] == 'csv':\n", + " df = pd.read_csv(fname)\n", + " else:\n", + " raw = json.load(open(fname))\n", + " # Handle nested JSON like {\"authors\": [...]}\n", + " if isinstance(raw, dict):\n", + " key = list(raw.keys())[0]\n", + " df = pd.DataFrame(raw[key])\n", + " else:\n", + " df = pd.DataFrame(raw)\n", + "\n", + " records = df.to_dict('records')\n", + " all_records[ntype] = records\n", + "\n", + " for record in records:\n", + " rid = str(record.get(id_field, hash(str(record))))\n", + " node = Node(id=rid, labels=[ntype], properties=flatten_properties(record))\n", + " cypher, params = node_to_cypher(node, tenant_id=TENANT)\n", + " gs.execute_query(cypher, params)\n", + "\n", + " print(f' Wrote {len(records)} {ntype} nodes')\n", + "\n", + "print(f'\\nTotal: {sum(len(v) for v in all_records.values())} nodes in Neptune')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Create Relationships (Edges)\n\nInfer and create edges between nodes based on shared identifiers (foreign keys).\nFor example:\n- Papers \u2192 `AUTHORED_BY` \u2192 Authors (via `author_id`)\n- Papers \u2192 `CITES` \u2192 Papers (via `citation_id`)\n\nThis gives us a typed, traversable graph in Neptune.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Step 2b: Create edge relationships\n", + "from graphrag_toolkit.document_graph.graph_build import edge_to_cypher\n", + "from graphrag_toolkit.document_graph import Edge\n", + "\n", + "edges_written = 0\n", + "\n", + "# AUTHORED_BY: link papers to authors (by matching author_id in papers)\n", + "papers = all_records.get('ResearchPaper', [])\n", + "authors = all_records.get('Author', [])\n", + "author_ids = {a.get('author_id') for a in authors}\n", + "\n", + "for paper in papers:\n", + " # Check if paper has an author field referencing an author\n", + " paper_authors = str(paper.get('authors', paper.get('author_id', ''))).split(',')\n", + " for aid in paper_authors:\n", + " aid = aid.strip()\n", + " if aid in author_ids:\n", + " edge = Edge(id=f'e-{paper[\"id\"]}-{aid}', source_id=str(paper['id']), target_id=aid, label='AUTHORED_BY')\n", + " cypher, params = edge_to_cypher(edge, tenant_id=TENANT)\n", + " gs.execute_query(cypher, params)\n", + " edges_written += 1\n", + "\n", + "# CITES: link citations to papers\n", + "citations = all_records.get('Citation', [])\n", + "paper_ids = {str(p.get('id')) for p in papers}\n", + "\n", + "for cite in citations:\n", + " src = str(cite.get('source_paper', cite.get('citing_paper', '')))\n", + " tgt = str(cite.get('target_paper', cite.get('cited_paper', '')))\n", + " if src and tgt:\n", + " edge = Edge(id=f'c-{cite.get(\"id\", \"\")}', source_id=src, target_id=tgt, label='CITES')\n", + " cypher, params = edge_to_cypher(edge, tenant_id=TENANT)\n", + " gs.execute_query(cypher, params)\n", + " edges_written += 1\n", + "\n", + "print(f'Wrote {edges_written} edges (AUTHORED_BY + CITES)')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Verify in Neptune\n\nConfirm the graph was built correctly by counting nodes per type.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = gs.execute_query('MATCH (n) RETURN labels(n) as lbl, count(n) as cnt')\n", + "print('Nodes by type:')\n", + "for row in result:\n", + " print(f' {row}')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Convert to LlamaIndex Documents (with Lineage)\n\n**This is the bridge between document-graph and lexical-graph.**\n\nWe convert each typed Neptune node into a LlamaIndex `Document` with:\n- **Text**: A structured representation of the node's properties (for semantic embedding)\n- **Lineage header**: `[NodeType | node_id | tenant]` embedded in the text\n\nThe lineage header survives lexical-graph's chunking process, so after semantic search\nyou can trace any result back to the original typed node in Neptune.\n\nThis is what makes the hybrid graph work: semantic search finds relevant text,\nlineage tells you which structured node it came from.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from llama_index.core.schema import Document\n", + "\n", + "lexical_documents = []\n", + "\n", + "for ntype, records in all_records.items():\n", + " id_field = next(s['id_field'] for s in data_sources if s['node_type'] == ntype)\n", + " for record in records:\n", + " node_id = str(record.get(id_field, 'unknown'))\n", + " source_id = f'{ntype}:{node_id}:{TENANT}'\n", + " # Build text with lineage header\n", + " text = f'[{ntype} | {node_id} | {TENANT}]\\n'\n", + " text += '\\n'.join(f'{k}: {v}' for k, v in record.items() if v is not None and str(v).strip())\n", + " # Use graphrag-toolkit's source metadata structure\n", + " doc = Document(\n", + " text=text,\n", + " metadata={\n", + " 'source': {\n", + " 'sourceId': source_id,\n", + " 'metadata': {\n", + " 'node_type': ntype,\n", + " 'node_id': node_id,\n", + " 'tenant': TENANT,\n", + " }\n", + " }\n", + " }\n", + " )\n", + " lexical_documents.append(doc)\n", + "\n", + "print(f'Created {len(lexical_documents)} documents')\n", + "for ntype in all_records:\n", + " count = sum(1 for d in lexical_documents if d.metadata.get('source', {}).get('metadata', {}).get('node_type') == ntype)\n", + " print(f' {ntype}: {count}')\n", + "print(f'\\nSample sourceId: {lexical_documents[0].metadata[\"source\"][\"sourceId\"]}')\n", + "print(f'Sample text:\\n{lexical_documents[0].text[:150]}')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Save for Part 2\n\nSerialize the LlamaIndex Documents to disk so Part 2 (`7b-Lexical-Integration-Lexical-Setup`)\ncan index them into lexical-graph without re-running the data processing.\n\nIn production, you'd pass these directly to `LexicalGraphIndex.extract_and_build()`.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import pickle\n\noutput_dir = Path('output/hybrid_foundation')\noutput_dir.mkdir(parents=True, exist_ok=True)\n\nwith open(output_dir / 'lexical_documents.pkl', 'wb') as f:\n pickle.dump(lexical_documents, f)\n\nwith open(output_dir / 'summary.json', 'w') as f:\n json.dump({'total_docs': len(lexical_documents), 'node_types': list(all_records.keys()), 'tenant': TENANT}, f, indent=2)\n\nprint(f'Saved {len(lexical_documents)} documents to {output_dir}')\nprint('Ready for 07b \u2014 Lexical-Graph Indexing')\n\n\n# Clean up Neptune connection\ngs.__exit__(None, None, None)\n" + ], + "outputs": [], + "execution_count": null + } + ] +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/10b-Hybrid-Lexical-Indexing.ipynb b/examples/document-graph/notebooks/10b-Hybrid-Lexical-Indexing.ipynb new file mode 100644 index 00000000..6cec483a --- /dev/null +++ b/examples/document-graph/notebooks/10b-Hybrid-Lexical-Indexing.ipynb @@ -0,0 +1,222 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Hybrid Graph Part 2: Lexical-Graph Indexing & Semantic Query\n", + "\n", + "This is Part 2 of the hybrid graph pipeline. It demonstrates how to:\n", + "\n", + "1. Load LlamaIndex Documents from Part 1 (structured data with embedded lineage)\n", + "2. Index them into lexical-graph (Neptune graph + OpenSearch vectors)\n", + "3. Query with natural language (semantic search)\n", + "4. Correlate results back to the original typed nodes in document-graph\n", + "\n", + "## Why This Matters\n", + "\n", + "After Part 1, we have:\n", + "- Typed nodes in Neptune (queryable by Cypher: exact match, traversal)\n", + "- LlamaIndex Documents with lineage (ready for semantic indexing)\n", + "\n", + "Part 2 adds:\n", + "- Semantic search capability (ask questions in natural language)\n", + "- Entity extraction and linking (lexical-graph's knowledge graph)\n", + "- Lineage-based correlation (trace semantic results \u2192 structured nodes)\n", + "\n", + "## Data Flow\n", + "\n", + "```\n", + "LlamaIndex Documents (from Part 1)\n", + " \u2193 lexical-graph extract_and_build()\n", + "Neptune: Entity graph + Chunk graph\n", + "OpenSearch: Vector embeddings\n", + " \u2193 semantic query\n", + "Query results with text + metadata\n", + " \u2193 lineage extraction\n", + "Original document-graph node in Neptune\n", + "```\n", + "\n", + "## The Hybrid Advantage\n", + "\n", + "| Query Type | Engine | Example |\n", + "|------------|--------|--------|\n", + "| Exact match | document-graph (Cypher) | \"Find all papers by author X\" |\n", + "| Traversal | document-graph (Cypher) | \"Papers that cite paper Y\" |\n", + "| Semantic | lexical-graph (vectors) | \"Research about machine learning\" |\n", + "| Hybrid | Both + lineage | \"Find structured node for this semantic result\" |\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-lexical-graph`\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]`\n", + "- Neptune cluster + OpenSearch Serverless collection (see 00-Setup)\n", + "- Part 1 completed (documents saved to `output/hybrid_foundation/`)\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Setup\nimport json, os, pickle, re\nfrom pathlib import Path\nfrom graphrag_toolkit.lexical_graph.storage import GraphStoreFactory, VectorStoreFactory\nfrom graphrag_toolkit.lexical_graph import LexicalGraphIndex, LexicalGraphQueryEngine\n\nGRAPH_STORE = os.environ.get(\n 'GRAPH_STORE',\n 'neptune-db://:8182'\n)\nVECTOR_STORE = os.environ.get(\n 'VECTOR_STORE',\n 'aoss://'\n)\nTENANT = 'hybrid_demo'\n\ngraph_store = GraphStoreFactory.for_graph_store(GRAPH_STORE)\nvector_store = VectorStoreFactory.for_vector_store(VECTOR_STORE)\nprint('Stores connected')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Load Documents from Part 1\n\nLoad the LlamaIndex Documents that Part 1 created from structured data.\nEach document contains the node's properties as text plus a lineage header\n(`[NodeType | node_id | tenant]`) that will survive chunking.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "foundation_dir = Path('output/hybrid_foundation')\n", + "\n", + "with open(foundation_dir / 'lexical_documents.pkl', 'rb') as f:\n", + " lexical_documents = pickle.load(f)\n", + "\n", + "print(f'Loaded {len(lexical_documents)} documents')\n", + "print(f'Sample: {lexical_documents[0].text[:150]}...')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Index into Lexical-Graph\n\n`LexicalGraphIndex.extract_and_build()` does three things:\n\n1. **Extract** \u2014 LLM-based entity and relationship extraction from the document text\n2. **Build graph** \u2014 Write extracted entities and relationships to Neptune\n3. **Build vectors** \u2014 Embed text chunks and store in OpenSearch for semantic search\n\nAfter this step, you can query the data with natural language.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Index into lexical-graph (batch writes enabled \u2014 Neptune 1.4.7.0 supports it)\n", + "print(f\"Indexing {len(lexical_documents)} documents (batch writes ON)...\")\n", + "\n", + "graph_index = LexicalGraphIndex(graph_store, vector_store)\n", + "graph_index.extract_and_build(lexical_documents, show_progress=True)\n", + "\n", + "print(\"Indexing complete\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Query with Semantic Search\n\nUse `LexicalGraphQueryEngine` to ask natural language questions.\nThe engine combines vector similarity search (OpenSearch) with graph traversal (Neptune)\nto find relevant answers.\n\nBehind the scenes:\n1. Query is embedded \u2192 similar chunks found in OpenSearch\n2. Chunks are expanded via graph traversal \u2192 related entities included\n3. Context is assembled and returned (or passed to an LLM for generation)\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "query_engine = LexicalGraphQueryEngine.for_traversal_based_search(graph_store, vector_store)\n", + "\n", + "queries = [\n", + " 'What research discusses machine learning?',\n", + " 'Who are the authors working on graph databases?',\n", + " 'What citations exist between papers?',\n", + "]\n", + "\n", + "for q in queries:\n", + " print(f'\\nQuery: {q}')\n", + " results = query_engine.retrieve(q)\n", + " if isinstance(results, list):\n", + " print(f' Results: {len(results)}')\n", + " for i, r in enumerate(results[:2], 1):\n", + " text = getattr(r.node, 'text', str(r))[:150]\n", + " print(f' {i}. {text}...')\n", + " else:\n", + " print(f' Response: {str(results)[:150]}')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Correlate Back to Document-Graph (Lineage)\n\n**This is the hybrid graph payoff.**\n\nSemantic search returns text chunks. But we need to know: *which structured node did this come from?*\n\nThe lineage header (`[ResearchPaper | paper-001 | hybrid_demo]`) embedded in Part 1 tells us.\nWe extract it from the result text, then query Neptune directly for the original typed node\nwith all its structured properties.\n\nThis completes the loop:\n- Natural language query \u2192 semantic result \u2192 lineage \u2192 typed Neptune node \u2192 structured data\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Correlate: query results \u2192 Source nodes \u2192 document-graph\n", + "gs_direct = GraphStoreFactory.for_graph_store(GRAPH_STORE).__enter__()\n", + "\n", + "q = 'What research discusses machine learning?'\n", + "results = query_engine.retrieve(q)\n", + "print(f'Query: {q}')\n", + "print(f'Results: {len(results)}\\n')\n", + "\n", + "for i, r in enumerate(results[:5], 1):\n", + " text = getattr(r.node, 'text', str(r))\n", + " meta = getattr(r.node, 'metadata', {})\n", + "\n", + " # Method 1: Parse lineage from text header [Type | node_id | tenant]\n", + " match = re.search(r'\\[(\\w+) \\| ([\\w-]+) \\| (\\w+)\\]', text)\n", + " if match:\n", + " ntype, node_id, tenant = match.group(1), match.group(2), match.group(3)\n", + " label = f'__{ntype}__{tenant}__'\n", + " original = gs_direct.execute_query(f\"MATCH (n:`{label}`) WHERE n.id = '{node_id}' RETURN properties(n) as props LIMIT 1\")\n", + " print(f'{i}. [{ntype}] node_id={node_id}')\n", + " if original:\n", + " props = original[0].get('props', {})\n", + " display = props.get('name', props.get('title', props.get('description', 'N/A')))\n", + " print(f' Neptune: {display}')\n", + " else:\n", + " print(f' (not in Neptune)')\n", + " else:\n", + " # Method 2: Check Source nodes in lexical-graph\n", + " sources = gs_direct.execute_query('MATCH (s:`__Source__`) WHERE s.node_type IS NOT NULL RETURN s.node_type as ntype, s.node_id as nid, s.tenant as t LIMIT 5')\n", + " if sources and i == 1:\n", + " print(f'{i}. [Source nodes available]:') \n", + " for s in sources[:3]:\n", + " print(f' {s}')\n", + " else:\n", + " # Method 3: Entity correlation fallback\n", + " ec = meta.get('entity_contexts', {}).get('contexts', [])\n", + " entities = [ent.get('entity', {}).get('value', '') for ctx in ec[:3] for ent in ctx.get('entities', [])]\n", + " print(f'{i}. [entities: {\", \".join(entities[:3])}]')\n", + " print()\n", + "\n", + "print('Lineage correlation complete')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "The hybrid pipeline:\n", + "1. **Document-graph** stores structured data as typed nodes in Neptune\n", + "2. **Lexical-graph** indexes the same data for semantic search (OpenSearch vectors)\n", + "3. **Lineage** embedded in text `[source_type | node_id | tenant]` survives chunking\n", + "4. **Correlation** extracts lineage from query results \u2192 fetches original node from Neptune\n", + "\n", + "Same graph database (Neptune) holds both the document-graph nodes AND the lexical-graph structure.\n" + ] + } + ] +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/11a-Mandela-Hybrid-Build.ipynb b/examples/document-graph/notebooks/11a-Mandela-Hybrid-Build.ipynb new file mode 100644 index 00000000..e9d26c07 --- /dev/null +++ b/examples/document-graph/notebooks/11a-Mandela-Hybrid-Build.ipynb @@ -0,0 +1,232 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "41ccb172", + "metadata": {}, + "source": [ + "# Nelson Mandela Hybrid GraphRAG: Build\n", + "\n", + "A real-world example of the hybrid graph pattern using Nelson Mandela data:\n", + "\n", + "- **Unstructured**: Wikipedia biography \u2192 lexical-graph (entities, chunks, semantic embeddings)\n", + "- **Structured**: Events CSV + Organizations JSON \u2192 document-graph (typed nodes with dates, roles)\n", + "\n", + "Both graphs share the same Neptune cluster and tenant, enabling cross-graph queries in Part 2.\n", + "\n", + "## What This Demonstrates\n", + "\n", + "Real knowledge graphs need both:\n", + "- **Unstructured text** (bios, articles, reports) \u2192 rich semantic context, entity discovery\n", + "- **Structured records** (databases, spreadsheets) \u2192 precise facts, dates, relationships\n", + "\n", + "Separately, each is limited. Together, you can ask:\n", + "- \"What key events shaped Mandela's life?\" (semantic \u2192 structured correlation)\n", + "- \"Show me the timeline of organizations Mandela was part of\" (structured \u2192 traversal)\n", + "- \"What does the bio say about the ANC?\" (semantic search \u2192 entity linking)\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-lexical-graph`\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]`\n", + "- `pip install llama-index-core`\n", + "- Neptune cluster + OpenSearch Serverless (see 00-Setup)\n", + "- Sample data in `mandela_data/`\n" + ], + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c3efb026", + "metadata": {}, + "outputs": [], + "source": [ + "# Setup\n", + "import json, os\n", + "import pandas as pd\n", + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory, VectorStoreFactory\n", + "from graphrag_toolkit.lexical_graph import LexicalGraphIndex\n", + "from graphrag_toolkit.document_graph.graph_build import node_to_cypher\n", + "from graphrag_toolkit.document_graph import Node\n", + "\n", + "GRAPH_STORE = 'neptune-db://:8182'\n", + "VECTOR_STORE = ''\n", + "TENANT = 'mandela'\n", + "\n", + "graph_store = GraphStoreFactory.for_graph_store(GRAPH_STORE)\n", + "vector_store = VectorStoreFactory.for_vector_store(VECTOR_STORE)\n", + "gs = graph_store.__enter__()\n", + "print('Neptune + OpenSearch connected')\n" + ] + }, + { + "cell_type": "markdown", + "id": "d1c50539", + "metadata": {}, + "source": [ + "## Step 1: Build Lexical-Graph from Wikipedia\n\nIndex Mandela's biography into lexical-graph. This creates:\n- **Entity nodes** (people, places, organizations) extracted by LLM\n- **Chunk nodes** (text segments) with vector embeddings in OpenSearch\n- **Relationships** between entities discovered from the text\n\nAfter this step, you can ask semantic questions about Mandela's life.\n" + ], + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "00da6830", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core.schema import Document\n", + "\n", + "# Static Mandela content\n", + "wiki_docs = [\n", + " Document(\n", + " text=\"\"\"Nelson Rolihlahla Mandela was a South African anti-apartheid activist and politician who served as the first president of South Africa from 1994 to 1999. He was the country's first black head of state and the first elected in a fully representative democratic election. His government focused on dismantling the legacy of apartheid by fostering racial reconciliation. Mandela was born in Mvezo, Transkei, and studied law at the University of Fort Hare and the University of Witwatersrand. He joined the African National Congress (ANC) in 1943 and co-founded its Youth League. After the National Party came to power in 1948, he rose to prominence in the ANC's 1952 defiance campaign. He was arrested and imprisoned in 1962, and subsequently sentenced to life imprisonment for conspiring to overthrow the state. Mandela served 27 years in prison, split between Robben Island, Pollsmoor Prison, and Victor Verster Prison. An international campaign lobbied for his release, which was granted in 1990. He led negotiations with President F.W. de Klerk to abolish apartheid and establish multiracial elections in 1994, which he won.\"\"\",\n", + " metadata={'source': 'wikipedia', 'title': 'Nelson Mandela'}\n", + " ),\n", + " Document(\n", + " text=\"\"\"Apartheid was a system of institutionalised racial segregation that existed in South Africa and South West Africa from 1948 to the early 1990s. Under apartheid, the rights of the majority non-white inhabitants were curtailed and white supremacy and Afrikaner minority rule were maintained. The National Party implemented apartheid as state policy after winning the 1948 general election. This racial segregation was enforced through legislation. People were classified into racial groups and separated by laws that banned social contact, mixed marriages, and established separate public facilities. The system was dismantled through negotiations between 1990 and 1993 and elections in 1994.\"\"\",\n", + " metadata={'source': 'wikipedia', 'title': 'Apartheid'}\n", + " ),\n", + " Document(\n", + " text=\"\"\"The African National Congress (ANC) is the Republic of South Africa's governing political party. It has been the ruling party of post-apartheid South Africa since the election of Nelson Mandela in 1994. Founded in 1912 as the South African Native National Congress, its primary mission was to bring all Africans together as one people to defend their rights and freedoms. The ANC was banned in 1960 after the Sharpeville massacre and operated underground and in exile for 30 years until it was unbanned in 1990. Key figures include Nelson Mandela, Oliver Tambo, Walter Sisulu, and Albert Luthuli.\"\"\",\n", + " metadata={'source': 'wikipedia', 'title': 'African National Congress'}\n", + " ),\n", + "]\n", + "\n", + "print(f'Created {len(wiki_docs)} documents')\n", + "for doc in wiki_docs:\n", + " title = doc.metadata['title']\n", + " print(f' {title}: {len(doc.text)} chars')\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ad51f14d", + "metadata": {}, + "outputs": [], + "source": [ + "# Index into lexical-graph\n", + "os.environ['BATCH_WRITES_ENABLED'] = 'false'\n", + "\n", + "print(f'Indexing {len(wiki_docs)} documents into lexical-graph...')\n", + "graph_index = LexicalGraphIndex(graph_store, vector_store)\n", + "graph_index.extract_and_build(wiki_docs, show_progress=True)\n", + "print('Lexical-graph build complete')\n" + ] + }, + { + "cell_type": "markdown", + "id": "2f9b0d20", + "metadata": {}, + "source": [ + "## Step 2: Build Document-Graph from Structured Data\n\nLoad structured records (events with dates, organizations with roles) as typed nodes.\nThese provide precise, queryable facts that complement the semantic layer:\n- `Event` nodes: date, location, description, significance\n- `Organization` nodes: name, role, founding date, members\n\nBoth use the same tenant so they coexist in one Neptune cluster.\n" + ], + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f7cc2918", + "metadata": {}, + "outputs": [], + "source": [ + "def flatten_props(props):\n", + " flat = {}\n", + " for k, v in props.items():\n", + " if isinstance(v, list):\n", + " flat[k] = ', '.join(str(x) for x in v)\n", + " elif isinstance(v, dict):\n", + " flat[k] = json.dumps(v)\n", + " elif v is None:\n", + " continue\n", + " else:\n", + " flat[k] = v\n", + " return flat\n", + "\n", + "# Load Mandela structured data\n", + "events = pd.read_csv('mandela_data/mandela_events.csv')\n", + "orgs = json.load(open('mandela_data/research_papers.json'))\n", + "if isinstance(orgs, dict):\n", + " key = list(orgs.keys())[0]\n", + " orgs = orgs[key]\n", + "\n", + "print(f'Events: {len(events)} rows')\n", + "print(f'Organizations/Papers: {len(orgs)} records')\n", + "\n", + "# Write events to Neptune\n", + "for _, row in events.iterrows():\n", + " record = row.to_dict()\n", + " rid = str(record.get('id', record.get('event_id', hash(str(record)))))\n", + " node = Node(id=rid, labels=['Event'], properties=flatten_props(record))\n", + " cypher, params = node_to_cypher(node, tenant_id=TENANT)\n", + " gs.execute_query(cypher, params)\n", + "\n", + "# Write orgs/papers to Neptune\n", + "for record in orgs:\n", + " rid = str(record.get('id', record.get('org_id', hash(str(record)))))\n", + " node = Node(id=rid, labels=['Organization'], properties=flatten_props(record))\n", + " cypher, params = node_to_cypher(node, tenant_id=TENANT)\n", + " gs.execute_query(cypher, params)\n", + "\n", + "print(f'\\nWrote {len(events)} Event + {len(orgs)} Organization nodes to Neptune')\n" + ] + }, + { + "cell_type": "markdown", + "id": "4f193f44", + "metadata": {}, + "source": [ + "## Step 3: Verify\n\nConfirm both graphs were built successfully in the same Neptune cluster.\nBoth lexical-graph entities and document-graph typed nodes should be present.\n" + ], + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "035e086f", + "metadata": {}, + "outputs": [], + "source": [ + "# Verify mandela tenant nodes\n", + "print(\"Mandela tenant nodes:\")\n", + "for ntype in [\"Event\", \"Organization\"]:\n", + " label = f\"__{ntype}__{TENANT}__\"\n", + " result = gs.execute_query(f\"MATCH (n:`{label}`) RETURN count(n) as cnt\")\n", + " cnt = result[0][\"cnt\"] if result else 0\n", + " print(f\" {ntype}: {cnt}\")\n", + "\n", + "print(\"Build complete. Ready for 08b (Query).\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ac6e7841-88e1-4c02-99ad-374fe8b57b98", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "conda_python3", + "language": "python", + "name": "conda_python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.20" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/11b-Mandela-Hybrid-Query.ipynb b/examples/document-graph/notebooks/11b-Mandela-Hybrid-Query.ipynb new file mode 100644 index 00000000..b41f6386 --- /dev/null +++ b/examples/document-graph/notebooks/11b-Mandela-Hybrid-Query.ipynb @@ -0,0 +1,203 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9cb43269", + "metadata": {}, + "source": [ + "# Nelson Mandela Hybrid GraphRAG: Query\n", + "\n", + "Query the hybrid graph built in Part 1 (8a) using three approaches:\n", + "\n", + "1. **Semantic search** \u2014 Natural language queries via lexical-graph\n", + "2. **Structured queries** \u2014 Direct Cypher via document-graph\n", + "3. **Hybrid correlation** \u2014 Semantic results \u2192 entity matching \u2192 structured enrichment\n", + "\n", + "## The Hybrid Query Pattern\n", + "\n", + "```\n", + "Natural language question\n", + " \u2193 lexical-graph (vector search + graph traversal)\n", + "Semantic results (text chunks, entities)\n", + " \u2193 entity extraction from results\n", + "Entity names (\"ANC\", \"Robben Island\", \"1990\")\n", + " \u2193 document-graph (Cypher match on typed nodes)\n", + "Structured facts (dates, roles, relationships)\n", + " \u2193 combined answer\n", + "Rich response with both context and precision\n", + "```\n", + "\n", + "**Prerequisites:**\n", + "- Part 1 (8a) completed\n", + "- `pip install graphrag-toolkit-lexical-graph graphrag-toolkit-document-graph[graphrag]`\n", + "- Neptune + OpenSearch endpoints\n" + ], + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "986ae7bc", + "metadata": {}, + "outputs": [], + "source": [ + "# Setup\n", + "import json, re\n", + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory, VectorStoreFactory\n", + "from graphrag_toolkit.lexical_graph import LexicalGraphQueryEngine\n", + "from graphrag_toolkit.document_graph.query import DocumentGraphQueryEngine\n", + "\n", + "GRAPH_STORE = 'neptune-db://:8182'\n", + "VECTOR_STORE = ''\n", + "TENANT = 'mandela'\n", + "\n", + "graph_store = GraphStoreFactory.for_graph_store(GRAPH_STORE)\n", + "vector_store = VectorStoreFactory.for_vector_store(VECTOR_STORE)\n", + "gs = graph_store.__enter__()\n", + "print('Connected')\n" + ] + }, + { + "cell_type": "markdown", + "id": "00e88b84", + "metadata": {}, + "source": [ + "## 1. Semantic Search (Lexical-Graph)\n\nAsk natural language questions. Lexical-graph finds relevant text chunks via vector\nsimilarity, then expands context via graph traversal (entities connected to those chunks).\n\nThis works because Part 1 indexed Mandela's Wikipedia bio into the semantic layer.\n" + ], + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0555bbd0", + "metadata": {}, + "outputs": [], + "source": [ + "query_engine = LexicalGraphQueryEngine.for_traversal_based_search(graph_store, vector_store)\n", + "\n", + "queries = [\n", + " 'Who was Nelson Mandela?',\n", + " 'What role did Mandela play in ending apartheid?',\n", + " 'What is the African National Congress?',\n", + "]\n", + "\n", + "for q in queries:\n", + " print(f'\\nQuery: {q}')\n", + " results = query_engine.retrieve(q)\n", + " if isinstance(results, list) and results:\n", + " print(f' Results: {len(results)}')\n", + " text = getattr(results[0].node, 'text', str(results[0]))[:200]\n", + " print(f' Top: {text}...')\n", + " else:\n", + " print(f' {str(results)[:150]}')\n" + ] + }, + { + "cell_type": "markdown", + "id": "67513860", + "metadata": {}, + "source": [ + "## 2. Structured Queries (Document-Graph)\n\nQuery typed nodes directly with Cypher. Precise, fast, and schema-aware.\n\nDocument-graph nodes use tenant-scoped labels (e.g., `Eventmandela__`) so they\ndon't collide with lexical-graph entities in the same Neptune cluster.\n" + ], + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d778ec4c", + "metadata": {}, + "outputs": [], + "source": [ + "# Query events\nevent_label = f\"`Event{TENANT}__`\"\nevents = gs.execute_query(f\"MATCH (n:`{event_label}`) RETURN properties(n) as props LIMIT 10\")\nprint(f\"Events in document-graph: {len(events)}\")\nfor e in events[:5]:\n props = e.get(\"props\", {})\n date = props.get(\"date\", \"\")\n desc = props.get(\"description\", \"?\")\n print(f\" [{date}] {desc}\")\n\n# Query organizations\norg_label = f\"`Organization{TENANT}__`\"\norgs = gs.execute_query(f\"MATCH (n:`{org_label}`) RETURN properties(n) as props LIMIT 10\")\nprint(f\"Organizations: {len(orgs)}\")\nfor o in orgs[:5]:\n props = o.get(\"props\", {})\n name = props.get(\"name\", props.get(\"title\", \"?\"))\n print(f\" {name}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "3a966ef4", + "metadata": {}, + "source": [ + "## 3. Hybrid Correlation\n\nThe most powerful pattern: combine semantic search results with structured data.\n\n1. Semantic query finds relevant context from the biography\n2. Extract entity names from the results (people, organizations, events)\n3. Match those entities against typed document-graph nodes\n4. Enrich the semantic results with structured facts (exact dates, roles, relationships)\n\nThis gives you the best of both worlds: semantic understanding + structured precision.\n" + ], + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cd596247", + "metadata": {}, + "outputs": [], + "source": [ + "# Semantic query + correlate entities to document-graph\n", + "q = 'What key events shaped Mandela\\'s life?'\n", + "results = query_engine.retrieve(q)\n", + "print(f'Query: {q}')\n", + "print(f'Lexical results: {len(results)}\\n')\n", + "\n", + "# Extract entities and look them up in document-graph\n", + "for i, r in enumerate(results[:3], 1):\n", + " meta = getattr(r.node, 'metadata', {})\n", + " ec = meta.get('entity_contexts', {}).get('contexts', [])\n", + " text = getattr(r.node, 'text', str(r))[:150]\n", + " print(f'Result {i}: {text}...')\n", + "\n", + " for ctx in ec[:5]:\n", + " for ent in ctx.get('entities', []):\n", + " entity = ent.get('entity', {})\n", + " value = entity.get('value', '')\n", + " classification = entity.get('classification', '')\n", + " if not value:\n", + " continue\n", + " search_term = value.split(',')[0].strip()\n", + " found = None\n", + " if classification in ('Event', 'Organization', 'Person'):\n", + " for label in [f'`Event{TENANT}__`', f'`Organization{TENANT}__`']:\n", + " found = gs.execute_query(f\"MATCH (n:`{label}`) WHERE n.name CONTAINS '{search_term}' OR n.title CONTAINS '{search_term}' RETURN properties(n) as props LIMIT 1\")\n", + " if found:\n", + " break\n", + " if found:\n", + " props = found[0].get('props', {})\n", + " name = props.get('name', props.get('title', value))\n", + " print(f' [{classification}] \"{value}\" -> Neptune: {name}')\n", + " print()\n", + "\n", + "print('Hybrid query complete \u2014 lexical search + document-graph correlation')\n" + ] + }, + { + "cell_type": "markdown", + "id": "d1219cb5", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "| Layer | Source | Contains |\n", + "|-------|--------|----------|\n", + "| Lexical-graph | Wikipedia | Semantic chunks, topics, entities from Mandela bio |\n", + "| Document-graph | Structured CSV/JSON | Events, organizations as typed nodes |\n", + "| Hybrid query | Both | Semantic search finds context, entity correlation finds structured data |\n" + ], + "outputs": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "conda_python3", + "language": "python", + "name": "conda_python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.20" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/12-Multi-Tenant-Coexistence.ipynb b/examples/document-graph/notebooks/12-Multi-Tenant-Coexistence.ipynb new file mode 100644 index 00000000..a2ec7915 --- /dev/null +++ b/examples/document-graph/notebooks/12-Multi-Tenant-Coexistence.ipynb @@ -0,0 +1,209 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Multi-Tenant Coexistence Demo\n", + "\n", + "Demonstrates that multiple tenants coexist in the same Neptune cluster with complete\n", + "data isolation \u2014 no cross-contamination between tenants.\n", + "\n", + "## How Multi-Tenancy Works\n", + "\n", + "document-graph (and lexical-graph) use **label-based tenant scoping**:\n", + "\n", + "```\n", + "Default tenant: `User` (no suffix)\n", + "Tenant \"alpha\": `Useralpha__` (label + tenant_id + __)\n", + "Tenant \"beta\": `Userbeta__` (label + tenant_id + __)\n", + "```\n", + "\n", + "Both packages use the same format (via `TenantId.format_label()`), so:\n", + "- document-graph nodes for tenant \"alpha\" are isolated from tenant \"beta\"\n", + "- lexical-graph entities for tenant \"alpha\" are isolated from tenant \"beta\"\n", + "- Both packages' nodes for the SAME tenant can be queried together\n", + "\n", + "## What This Notebook Proves\n", + "\n", + "1. Two tenants writing identical node IDs don't collide\n", + "2. Queries are scoped \u2014 each tenant sees only its own data\n", + "3. An admin view can see all tenants when needed\n", + "\n", + "**Prerequisites:**\n", + "- `pip install graphrag-toolkit-document-graph[graphrag]`\n", + "- Neptune cluster endpoint (see 00-Setup)\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Setup\n", + "import json\n", + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory\n", + "from graphrag_toolkit.document_graph.graph_build import node_to_cypher\n", + "from graphrag_toolkit.document_graph.query import DocumentGraphQueryEngine\n", + "from graphrag_toolkit.document_graph import Node\n", + "\n", + "GRAPH_STORE = 'neptune-db://:8182'\n", + "\n", + "gs = GraphStoreFactory.for_graph_store(GRAPH_STORE).__enter__()\n", + "print('Neptune connected')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Write Data for Tenant A\n\nCreate nodes under tenant `alpha`. These nodes get labels like `Useralpha__`\nand are invisible to other tenants.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "TENANT_A = 'acme_corp'\n", + "\n", + "tenant_a_data = [\n", + " Node(id='u1', labels=['User'], properties={'name': 'Alice', 'role': 'admin', 'department': 'Engineering'}),\n", + " Node(id='u2', labels=['User'], properties={'name': 'Bob', 'role': 'developer', 'department': 'Engineering'}),\n", + " Node(id='p1', labels=['Project'], properties={'name': 'Apollo', 'status': 'active'}),\n", + "]\n", + "\n", + "for node in tenant_a_data:\n", + " cypher, params = node_to_cypher(node, tenant_id=TENANT_A)\n", + " gs.execute_query(cypher, params)\n", + "\n", + "print(f'Tenant A ({TENANT_A}): wrote {len(tenant_a_data)} nodes')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Write Data for Tenant B\n\nCreate nodes with the **same IDs** under tenant `beta`. Labels become `Userbeta__`.\nDespite identical node IDs, there's no collision \u2014 different labels = different nodes.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "TENANT_B = 'globex_inc'\n", + "\n", + "tenant_b_data = [\n", + " Node(id='u1', labels=['User'], properties={'name': 'Charlie', 'role': 'manager', 'department': 'Sales'}),\n", + " Node(id='u2', labels=['User'], properties={'name': 'Diana', 'role': 'analyst', 'department': 'Finance'}),\n", + " Node(id='p1', labels=['Project'], properties={'name': 'Zeus', 'status': 'planning'}),\n", + "]\n", + "\n", + "for node in tenant_b_data:\n", + " cypher, params = node_to_cypher(node, tenant_id=TENANT_B)\n", + " gs.execute_query(cypher, params)\n", + "\n", + "print(f'Tenant B ({TENANT_B}): wrote {len(tenant_b_data)} nodes')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Query \u2014 Tenant Isolation\n\nQuery each tenant separately using `DocumentGraphQueryEngine` with a `tenant_id`.\nEach tenant sees only its own data. No leakage.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Query Tenant A only\n", + "engine_a = DocumentGraphQueryEngine(gs, tenant_id=TENANT_A)\n", + "users_a = engine_a.get_nodes('User')\n", + "print(f'Tenant A Users: {len(users_a)}')\n", + "for u in users_a:\n", + " props = u['n']['~properties']\n", + " print(f' {props.get(\"name\")} ({props.get(\"role\")})')\n", + "\n", + "# Query Tenant B only\n", + "engine_b = DocumentGraphQueryEngine(gs, tenant_id=TENANT_B)\n", + "users_b = engine_b.get_nodes('User')\n", + "print(f'\\nTenant B Users: {len(users_b)}')\n", + "for u in users_b:\n", + " props = u['n']['~properties']\n", + " print(f' {props.get(\"name\")} ({props.get(\"role\")})')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Prove Isolation \u2014 Same ID, Different Tenants\n\nDemonstrate that node `user-001` in tenant \"alpha\" has different properties\nthan node `user-001` in tenant \"beta\" \u2014 they're completely separate entities\nthat happen to share an ID.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Both tenants have node id='u1' but they're different nodes\n", + "a_u1 = engine_a.find_by_property('User', 'id', 'u1')\n", + "b_u1 = engine_b.find_by_property('User', 'id', 'u1')\n", + "\n", + "a_name = a_u1[0]['n']['~properties']['name'] if a_u1 else 'NOT FOUND'\n", + "b_name = b_u1[0]['n']['~properties']['name'] if b_u1 else 'NOT FOUND'\n", + "\n", + "print(f'Node u1 in Tenant A: {a_name}')\n", + "print(f'Node u1 in Tenant B: {b_name}')\n", + "print(f'\\nSame ID, different data = tenant isolation working')\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Cross-Tenant View (Admin)\n\nAn admin can query across all tenants by not specifying a `tenant_id` \u2014\nor by using a wildcard label query. This is for management/debugging only.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Admin can see all tenants by querying without tenant filter\nall_users = gs.execute_query('MATCH (n) WHERE labels(n)[0] CONTAINS \"User\" RETURN labels(n) as lbl, properties(n) as props LIMIT 10')\nprint(f'All User nodes across tenants: {len(all_users)}')\nfor u in all_users:\n label = u['lbl'][0]\n name = u['props'].get('name', '?')\n # Extract tenant from label: Usertenant__\n parts = label.split('__')\n tenant = parts[2] if len(parts) >= 3 else '?'\n print(f' {name} (tenant: {tenant})')\n\n\n# Clean up\ngs.__exit__(None, None, None)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "| Aspect | How it works |\n", + "|--------|-------------|\n", + "| **Isolation** | Labels encode tenant: `Useracme_corp__` vs `Userglobex_inc__` |\n", + "| **Same IDs** | Node id='u1' exists in both tenants independently |\n", + "| **Querying** | `DocumentGraphQueryEngine(gs, tenant_id=X)` scopes all queries |\n", + "| **Admin view** | Query without tenant filter sees all data |\n", + "| **No leakage** | Tenant A cannot see Tenant B's data through normal queries |\n" + ] + } + ] +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/99-Cleanup.ipynb b/examples/document-graph/notebooks/99-Cleanup.ipynb new file mode 100644 index 00000000..7cf7bac6 --- /dev/null +++ b/examples/document-graph/notebooks/99-Cleanup.ipynb @@ -0,0 +1,183 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 00-Cleanup — Remove Old Test Tenants from Neptune\n", + "\n", + "This notebook helps you clean stale test data from previous notebook runs.\n", + "\n", + "**What it does:**\n", + "1. Connects to Neptune and lists all node labels (which encode tenant IDs)\n", + "2. Identifies tenants you want to remove vs. keep\n", + "3. Deletes nodes belonging to old tenants\n", + "\n", + "**Tenant label format:** Labels follow the graphrag-toolkit convention:\n", + "- Default tenant: `` `Label` ``\n", + "- Custom tenant: `` `Label{tenant_id}__` `` (e.g., `Persondemo__`)\n", + "\n", + "**⚠️ Warning:** This notebook permanently deletes graph data. Review the tenant list before running the delete cell." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configuration" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "import re\n", + "\n", + "# Replace with your Neptune endpoint:\n", + "GRAPH_STORE = os.environ.get(\n", + " 'GRAPH_STORE',\n", + " 'neptune-db://:8182'\n", + ")\n", + "\n", + "# Tenants to KEEP (will not be deleted):\n", + "KEEP_TENANTS = ['hybrid_demo', 'mandela']\n", + "\n", + "print(f'Endpoint: {GRAPH_STORE}')\n", + "print(f'Protected tenants: {KEEP_TENANTS}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Connect to Neptune and List Labels" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory\n", + "\n", + "graph_store = GraphStoreFactory.for_graph_store(GRAPH_STORE)\n", + "gs = graph_store.__enter__()\n", + "\n", + "result = gs.execute_query(\n", + " 'MATCH (n) RETURN DISTINCT labels(n) as lbl, count(n) as cnt ORDER BY cnt DESC LIMIT 50'\n", + ")\n", + "\n", + "print(f'Found {len(result)} distinct label sets in Neptune:\\n')\n", + "for row in result:\n", + " print(f' {row[\"lbl\"][0]:40s} {row[\"cnt\"]:>6} nodes')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Identify Tenants\n", + "\n", + "Extract tenant IDs from labels. The tenant format is `Label{tenant_id}__` — we parse the suffix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Pattern: label ends with {tenant_id}__ where tenant_id is 1-25 lowercase alphanum/periods\n", + "TENANT_PATTERN = re.compile(r'^(.+?)([a-z0-9.]{1,25})__$')\n", + "\n", + "all_tenants = set()\n", + "for row in result:\n", + " label = row['lbl'][0]\n", + " match = TENANT_PATTERN.match(label)\n", + " if match:\n", + " all_tenants.add(match.group(2))\n", + "\n", + "tenants_to_remove = all_tenants - set(KEEP_TENANTS)\n", + "tenants_to_keep = all_tenants & set(KEEP_TENANTS)\n", + "\n", + "print(f'All tenants found: {sorted(all_tenants)}')\n", + "print(f'Will KEEP: {sorted(tenants_to_keep)}')\n", + "print(f'Will REMOVE: {sorted(tenants_to_remove)}')\n", + "\n", + "if not tenants_to_remove:\n", + " print('\\n✅ Nothing to clean up.')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Delete Old Tenants\n", + "\n", + "⚠️ **This permanently deletes data.** Review the list above before running this cell." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Safety check — set to True to actually delete\n", + "CONFIRM_DELETE = False\n", + "\n", + "if not CONFIRM_DELETE:\n", + " print('⚠️ CONFIRM_DELETE is False. Set to True and re-run to delete.')\n", + " print(f' Would delete tenants: {sorted(tenants_to_remove)}')\n", + "else:\n", + " for tenant in sorted(tenants_to_remove):\n", + " # Find all nodes whose label ends with {tenant}__\n", + " suffix = f'{tenant}__'\n", + " count_q = f\"MATCH (n) WHERE any(l IN labels(n) WHERE l ENDS WITH '{suffix}') RETURN count(n) as cnt\"\n", + " count_result = gs.execute_query(count_q)\n", + " cnt = count_result[0]['cnt'] if count_result else 0\n", + "\n", + " if cnt > 0:\n", + " delete_q = f\"MATCH (n) WHERE any(l IN labels(n) WHERE l ENDS WITH '{suffix}') DETACH DELETE n\"\n", + " gs.execute_query(delete_q)\n", + " print(f' 🗑️ Deleted {cnt} nodes from tenant: {tenant}')\n", + " else:\n", + " print(f' ⏭️ {tenant}: already empty')\n", + "\n", + " print('\\n✅ Cleanup complete.')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify — Show Remaining Data" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = gs.execute_query(\n", + " 'MATCH (n) RETURN DISTINCT labels(n) as lbl, count(n) as cnt ORDER BY cnt DESC LIMIT 30'\n", + ")\n", + "\n", + "print(f'Remaining labels ({len(result)} distinct):\\n')\n", + "for row in result:\n", + " print(f' {row[\"lbl\"][0]:40s} {row[\"cnt\"]:>6} nodes')\n", + "\n", + "# Clean up context manager\n", + "graph_store.__exit__(None, None, None)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/document-graph/notebooks/data/README_FORMATS.md b/examples/document-graph/notebooks/data/README_FORMATS.md new file mode 100644 index 00000000..de7ddffd --- /dev/null +++ b/examples/document-graph/notebooks/data/README_FORMATS.md @@ -0,0 +1,80 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +# Test Document Formats + +This directory contains the same test data in multiple formats for comprehensive testing: + +## Available Formats + +1. **CSV** (`test_documents.csv`) - Original comma-separated values format +2. **JSON** (`test_documents.json`) - JavaScript Object Notation with proper data types +3. **XML** (`test_documents.xml`) - Extensible Markup Language with structured elements +4. **YAML** (`test_documents.yaml`) - YAML Ain't Markup Language, human-readable format + +## Additional Formats (Generate as needed) + +### Parquet Format +To create `test_documents.parquet`: +```python +import pandas as pd +df = pd.read_csv('test_documents.csv') +df.to_parquet('test_documents.parquet', index=False) +``` + +### Excel Format +To create `test_documents.xlsx`: +```python +import pandas as pd +df = pd.read_csv('test_documents.csv') +df.to_excel('test_documents.xlsx', index=False, sheet_name='test_documents') +``` + +## Data Content + +All formats contain 20 test documents with: + +- **Standard Content**: Regular documents for basic processing +- **PII Data**: Document #13 with SSN, email, phone, IP, credit card +- **Spelling Errors**: Document #14 with intentional misspellings +- **Multi-language**: Document #16 with 8 different languages +- **Toxic Content**: Document #17 with inappropriate language +- **Mixed Sentiment**: Document #18 with positive and negative emotions +- **Security Threats**: Document #19 with XSS, SQL, command injection +- **Performance Test**: Document #20 with long text content + +## Usage in Notebooks + +Each format can be used to test different ingestion methods: + +```python +# CSV +import pandas as pd +df = pd.read_csv('test_documents.csv') + +# JSON +import json +with open('test_documents.json', 'r') as f: + data = json.load(f) + +# XML +import xml.etree.ElementTree as ET +tree = ET.parse('test_documents.xml') + +# YAML +import yaml +with open('test_documents.yaml', 'r') as f: + data = yaml.safe_load(f) + +# Parquet (if created) +df = pd.read_parquet('test_documents.parquet') + +# Excel (if created) +df = pd.read_excel('test_documents.xlsx') +``` + +## Requirements + +For Parquet and Excel formats, install required packages: +```bash +pip install pandas pyarrow openpyxl +``` \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/accounts.csv b/examples/document-graph/notebooks/data/accounts.csv new file mode 100644 index 00000000..52c53152 --- /dev/null +++ b/examples/document-graph/notebooks/data/accounts.csv @@ -0,0 +1,4 @@ +id,name,type,region +a1,Production,PROD,us-east-1 +a2,Development,DEV,us-west-2 +a3,Staging,STAGE,us-east-2 diff --git a/examples/document-graph/notebooks/data/create_additional_formats.py b/examples/document-graph/notebooks/data/create_additional_formats.py new file mode 100644 index 00000000..45aae430 --- /dev/null +++ b/examples/document-graph/notebooks/data/create_additional_formats.py @@ -0,0 +1,56 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +#!/usr/bin/env python3 +""" +Script to create Parquet and Excel versions of test_documents.csv +Run this script to generate test_documents.parquet and test_documents.xlsx +""" + +import pandas as pd +import json + +# Read the CSV file with proper quoting and escaping +df = pd.read_csv('test_documents.csv', quotechar='"', escapechar='\\', skipinitialspace=True) + +# Convert metadata column from string to proper JSON objects for better handling +df['metadata'] = df['metadata'].apply(lambda x: json.loads(x.replace("'", '"'))) + +# Create Parquet file +print("Creating test_documents.parquet...") +df.to_parquet('test_documents.parquet', index=False, engine='pyarrow') +print("✓ Parquet file created successfully") + +# Create Excel file +print("Creating test_documents.xlsx...") +with pd.ExcelWriter('test_documents.xlsx', engine='openpyxl') as writer: + # Convert metadata back to string for Excel compatibility + df_excel = df.copy() + df_excel['metadata'] = df_excel['metadata'].apply(lambda x: json.dumps(x)) + + # Write to Excel with formatting + df_excel.to_excel(writer, sheet_name='test_documents', index=False) + + # Get the workbook and worksheet + workbook = writer.book + worksheet = writer.sheets['test_documents'] + + # Auto-adjust column widths + for column in worksheet.columns: + max_length = 0 + column_letter = column[0].column_letter + for cell in column: + try: + if len(str(cell.value)) > max_length: + max_length = len(str(cell.value)) + except: + pass + adjusted_width = min(max_length + 2, 50) # Cap at 50 characters + worksheet.column_dimensions[column_letter].width = adjusted_width + +print("✓ Excel file created successfully") + +print("\nFiles created:") +print("- test_documents.parquet (Apache Parquet format)") +print("- test_documents.xlsx (Microsoft Excel format)") +print("\nBoth files contain the same 20 test documents with all transformers test scenarios.") \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/mapping.json b/examples/document-graph/notebooks/data/mapping.json new file mode 100644 index 00000000..30c52153 --- /dev/null +++ b/examples/document-graph/notebooks/data/mapping.json @@ -0,0 +1,387 @@ +{ + "schema_version": "1.0", + "transformations": [ + { + "rule_name": "evidence_1", + "type": "json", + "csv_property_name": "Evidence", + "parent": [ + "Evidence" + ] + }, + { + "rule_name": "resource_1", + "type": "json", + "csv_property_name": "Resource original JSON", + "parent": [ + "Resource" + ] + } + ], + "graph": { + "edges": [ + { + "Resource": { + "RAISES": "Issue", + "BELONGS_TO": "Project", + "LOCATED_IN": "CSP", + "INCLUDED_IN": "Subscription" + } + }, + { + "Issue": { + "POSSESSES": "Threat", + "SUPPORTED_BY": "Evidence", + "CONTRIBUTES_TO": "Risk", + "HAS_A": "Control" + } + } + ], + "nodes": [ + { + "Threat": { + "type": "node", + "property": { + "name": "threat_type", + "csv_map": "Threats" + } + } + }, + { + "Risk": { + "type": "node", + "property": { + "name": "risk_type", + "csv_map": "Risks" + } + } + }, + { + "Resource": { + "type": "node", + "property": { + "name": "resource_id", + "csv_map": "Resource vertex ID" + } + } + }, + { + "Issue": { + "type": "node", + "property": { + "name": "issue_id", + "csv_map": "Issue ID" + } + } + }, + { + "Evidence": { + "type": "node", + "property": { + "name": "evidence_id", + "type": "uuid" + } + } + }, + { + "Project": { + "type": "node", + "property": { + "name": "project_id", + "csv_map": "Project IDs" + } + } + }, + { + "Subscription": { + "type": "node", + "property": { + "name": "subscription_id", + "csv_map": "Subscription ID" + } + } + }, + { + "CSP": { + "type": "node", + "property": { + "name": "platform", + "csv_map": "Resource Platform" + } + } + }, + { + "Control": { + "type": "node", + "property": { + "name": "control_id", + "csv_map": "Control ID" + } + } + } + + ] + }, + "csv_columns": [ + { + "csv_property_name": "Created At", + "type": "property", + "arrow_property_name": "detected_at", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Title", + "type": "property", + "arrow_property_name": "title", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Severity", + "type": "property", + "arrow_property_name": "severity", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Status", + "type": "property", + "arrow_property_name": "status", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Description", + "type": "property", + "arrow_property_name": "description", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Resource Type", + "type": "property", + "arrow_property_name": "resource_type", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Resource external ID", + "type": "property", + "arrow_property_name": "arn_or_uri", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Project Names", + "type": "property", + "arrow_property_name": "name", + "parents": [ + "Project" + ] + }, + { + "csv_property_name": "Resolved Time", + "type": "property", + "arrow_property_name": "resolved_at", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Resolution", + "type": "property", + "arrow_property_name": "resolution", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Control ID", + "type": "property", + "arrow_property_name": "control_id", + "parents": [ + "Control" + ] + }, + { + "csv_property_name": "Resource Name", + "type": "property", + "arrow_property_name": "name", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Resource Region", + "type": "property", + "arrow_property_name": "region", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Resource Status", + "type": "property", + "arrow_property_name": "status", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Resource Platform", + "type": "property", + "arrow_property_name": "platform", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Resource OS", + "type": "property", + "arrow_property_name": "os", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Resource original JSON", + "type": "transformation", + "rules": [ + "resource_1" + ] + }, + { + "csv_property_name": "Ticket URLs", + "type": "property", + "arrow_property_name": "ticket_url", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Note", + "type": "property", + "arrow_property_name": "note", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Due At", + "type": "property", + "arrow_property_name": "due_at", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Remediation Recommendation", + "type": "property", + "arrow_property_name": "remediation", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Subscription Name", + "type": "property", + "arrow_property_name": "name", + "parents": [ + "Subscription" + ] + }, + { + "csv_property_name": "Wiz URL", + "type": "property", + "arrow_property_name": "wiz_url", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Cloud Provider URL", + "type": "property", + "arrow_property_name": "url", + "parents": [ + "CSP" + ] + }, + { + "csv_property_name": "Resource Tags", + "type": "property", + "arrow_property_name": "tags", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Kubernetes Cluster", + "type": "property", + "arrow_property_name": "kubernetes_cluster", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Kubernetes Namespace", + "type": "property", + "arrow_property_name": "kubernetes_namespace", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Container Service", + "type": "property", + "arrow_property_name": "container_service", + "parents": [ + "Resource" + ] + }, + { + "csv_property_name": "Provider ID", + "type": "node", + "arrow_property_name": "provider_id", + "name": "Resource" + }, + { + "csv_property_name": "Threats", + "type": "property", + "arrow_property_name": "threat_type", + "parents": [ + "Threat" + ] + }, + { + "csv_property_name": "Status Changed At", + "type": "property", + "arrow_property_name": "status_changed_at", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Updated At", + "type": "property", + "arrow_property_name": "updated_at", + "parents": [ + "Issue" + ] + }, + { + "csv_property_name": "Evidence", + "type": "transformation", + "rules": [ + "evidence_1" + ] + } + ] +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/sample.json b/examples/document-graph/notebooks/data/sample.json new file mode 100644 index 00000000..3602a698 --- /dev/null +++ b/examples/document-graph/notebooks/data/sample.json @@ -0,0 +1,106 @@ +{ + "schema_id": "research_paper_etl_v1", + "extract": { + "source_type": "document_graph", + "source_config": { + "format": "json", + "encoding": "utf-8" + }, + "fields": [ + { + "name": "paper_id", + "type": "string", + "required": true, + "source_path": "nodes[?type=='Document'].id" + }, + { + "name": "title", + "type": "string", + "required": true, + "source_path": "nodes[?type=='Document'].properties.title" + } + ] + }, + "transform": { + "chunking": { + "enabled": true, + "chunk_size": 1000, + "overlap": 200, + "strategy": "sentence" + }, + "metadata_mapping": { + "document_id": "paper_id", + "title": "title", + "source": "extract" + }, + "entity_extraction": { + "enabled": true, + "types": [ + "PERSON", + "ORGANIZATION", + "CONCEPT" + ], + "model": "spacy_en_core_web_sm" + }, + "normalize": { + "text_cleaning": true, + "lowercase": false, + "remove_punctuation": false, + "remove_stopwords": false + } + }, + "load": { + "document_node": { + "type": "Document", + "id_field": "paper_id", + "properties": [ + "title", + "authors", + "year" + ], + "labels": [ + "Document", + "ResearchPaper" + ] + }, + "section_node": { + "type": "Section", + "id_field": "section_id", + "properties": [ + "title", + "content", + "order" + ], + "labels": [ + "Section", + "DocumentSection" + ] + }, + "relationships": { + "document_to_section": { + "type": "HAS_SECTION", + "from": "document_node", + "to": "section_node", + "properties": [ + "order", + "section_type" + ] + }, + "section_to_entity": { + "type": "MENTIONS", + "from": "section_node", + "to": "entity_node", + "properties": [ + "confidence", + "position" + ] + } + } + }, + "metadata": { + "version": "1.0", + "created_at": "2024-01-15T10:30:00Z", + "description": "ETL schema for research paper document graph", + "domain": "academic_research" + } +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/test_documents.csv b/examples/document-graph/notebooks/data/test_documents.csv new file mode 100644 index 00000000..8ea4cdf7 --- /dev/null +++ b/examples/document-graph/notebooks/data/test_documents.csv @@ -0,0 +1,21 @@ +id,title,content,category,author,date,tags,language,source_url,metadata +1,"AWS Lambda Introduction","AWS Lambda is a serverless computing service that lets you run code without provisioning or managing servers. You pay only for the compute time you consume. Lambda runs your code on a high-availability compute infrastructure and performs all of the administration of the compute resources.","technology","AWS Documentation","2024-01-15","serverless,lambda,aws,computing","en","https://docs.aws.amazon.com/lambda/","{'service': 'lambda', 'complexity': 'beginner'}" +2,"Machine Learning Basics","Machine learning is a subset of artificial intelligence that enables computers to learn and make decisions from data without being explicitly programmed. It involves algorithms that can identify patterns in data and make predictions or classifications.","education","Dr. Sarah Johnson","2024-01-20","machine-learning,ai,data-science,algorithms","en","https://example.com/ml-basics","{'difficulty': 'intermediate', 'field': 'computer_science'}" +3,"Cybersecurity Best Practices","Implementing strong cybersecurity measures is crucial for protecting sensitive data. Key practices include using multi-factor authentication, regular software updates, employee training, network segmentation, and incident response planning.","security","John Smith","2024-01-25","cybersecurity,security,data-protection,best-practices","en","https://security-blog.com/best-practices","{'industry': 'cybersecurity', 'audience': 'enterprise'}" +4,"Climate Change Impact","Climate change refers to long-term shifts in global temperatures and weather patterns. While climate variations are natural, scientific evidence shows that human activities have been the main driver of climate change since the 1800s.","environment","Environmental Research Institute","2024-02-01","climate-change,environment,sustainability,global-warming","en","https://climate-research.org/impact","{'topic': 'environmental_science', 'urgency': 'high'}" +5,"Quantum Computing Fundamentals","Quantum computing harnesses quantum mechanical phenomena to process information in fundamentally different ways than classical computers. Quantum bits (qubits) can exist in superposition states, enabling parallel processing capabilities.","technology","Prof. Michael Chen","2024-02-05","quantum-computing,physics,technology,qubits","en","https://quantum-journal.com/fundamentals","{'complexity': 'advanced', 'field': 'quantum_physics'}" +6,"Healthy Cooking Tips","Maintaining a healthy diet starts with proper cooking techniques. Use fresh ingredients, limit processed foods, incorporate plenty of vegetables, choose lean proteins, and use healthy cooking methods like steaming, grilling, or baking instead of frying.","health","Nutritionist Lisa Brown","2024-02-10","health,nutrition,cooking,wellness,diet","en","https://health-magazine.com/cooking-tips","{'category': 'lifestyle', 'target_audience': 'general_public'}" +7,"Financial Planning Strategies","Effective financial planning involves setting clear goals, creating a budget, building an emergency fund, investing wisely, and regularly reviewing your financial situation. Start early and be consistent with your savings and investment approach.","finance","Financial Advisor Mark Wilson","2024-02-15","finance,investment,budgeting,savings,planning","en","https://finance-advisor.com/strategies","{'expertise_level': 'beginner_to_intermediate', 'topic': 'personal_finance'}" +8,"Space Exploration History","Space exploration began in the mid-20th century with the launch of Sputnik 1 in 1957. Major milestones include the first human in space (Yuri Gagarin, 1961), the moon landing (Apollo 11, 1969), and the development of space stations and Mars rovers.","science","Space History Museum","2024-02-20","space,exploration,history,nasa,astronomy","en","https://space-museum.org/history","{'era': '20th_21st_century', 'field': 'aerospace'}" +9,"Sustainable Energy Solutions","Renewable energy sources like solar, wind, and hydroelectric power are becoming increasingly important for reducing carbon emissions. These technologies offer clean alternatives to fossil fuels and are becoming more cost-effective.","environment","Green Energy Consortium","2024-02-25","renewable-energy,sustainability,solar,wind,environment","en","https://green-energy.org/solutions","{'focus': 'renewable_energy', 'impact': 'environmental'}" +10,"Digital Marketing Trends","Digital marketing continues to evolve with new technologies and consumer behaviors. Key trends include personalization, video content, social media marketing, influencer partnerships, and data-driven decision making.","business","Marketing Expert Jane Davis","2024-03-01","digital-marketing,social-media,trends,business,advertising","en","https://marketing-trends.com/2024","{'industry': 'marketing', 'year': '2024'}" +11,"Artificial Intelligence Ethics","As AI systems become more prevalent, ethical considerations become crucial. Key issues include bias in algorithms, privacy concerns, job displacement, transparency in decision-making, and the need for responsible AI development.","technology","AI Ethics Committee","2024-03-05","ai,ethics,technology,society,responsibility","en","https://ai-ethics.org/guidelines","{'importance': 'critical', 'stakeholders': 'global'}" +12,"Urban Planning Principles","Effective urban planning creates sustainable, livable communities. Key principles include mixed-use development, public transportation, green spaces, affordable housing, and community engagement in the planning process.","urban-development","City Planning Department","2024-03-10","urban-planning,city-development,sustainability,community","en","https://city-planning.gov/principles","{'scope': 'municipal', 'focus': 'sustainable_development'}" +13,"Data Breach Incident Report","On March 15th, our security team detected unauthorized access to customer data. The breach affected 1,247 users including John Doe (SSN: 123-45-6789, email: john.doe@email.com), Jane Smith (phone: 555-123-4567), and server logs from IP address 192.168.1.100. Credit card 4532-1234-5678-9012 was also compromised. Immediate containment measures were implemented.","security","Security Team Lead","2024-03-15","data-breach,security,incident,pii","en","https://internal-security.com/incident-001","{'severity': 'high', 'affected_users': 1247, 'containment_status': 'complete'}" +14,"Machien Lerning Algoritms","Machien lerning algoritms are becomming more sofisticated every day. Thier ability to proces large amounts of data and identyfy paterns is truely remarkabel. However, ther are stil chalenges with bias and interpretabilty that need to be adressed.","technology","Dr. Alex Thompson","2024-03-20","machine-learning,algorithms,ai,spelling-errors","en","https://ml-research.edu/algorithms","{'quality': 'poor_spelling', 'needs_correction': true}" +15,"API Configuration Guide","To configure the API endpoint, use the following JSON structure: {\"endpoint\": \"https://api.example.com/v1\", \"auth\": {\"type\": \"bearer\", \"token\": \"sk-1234567890abcdef\"}, \"timeout\": 30000, \"retries\": 3, \"headers\": {\"Content-Type\": \"application/json\", \"User-Agent\": \"MyApp/1.0\"}}. This configuration ensures proper authentication and error handling.","technical","DevOps Engineer","2024-03-25","api,configuration,json,development","en","https://dev-docs.com/api-config","{'format': 'json_embedded', 'complexity': 'intermediate'}" +16,"Mixed Language Content","This document contains multiple languages. English: Hello world! Spanish: ¡Hola mundo! French: Bonjour le monde! German: Hallo Welt! Japanese: こんにちは世界！Chinese: 你好世界！Russian: Привет мир! Arabic: مرحبا بالعالم! This tests language detection capabilities.","linguistics","Language Research Team","2024-04-01","multilingual,language-detection,international","mixed","https://lang-research.org/mixed-content","{'languages': ['en', 'es', 'fr', 'de', 'ja', 'zh', 'ru', 'ar'], 'test_type': 'language_detection'}" +17,"Toxic Content Example","This content contains inappropriate language and offensive terms that should be filtered out. Words like damn, hell, and other profanity need to be detected and handled appropriately by content moderation systems. This tests toxicity detection.","content-moderation","Content Safety Team","2024-04-05","toxicity,content-moderation,filtering","en","https://safety.example.com/toxic-test","{'toxicity_level': 'moderate', 'requires_filtering': true}" +18,"Sentiment Analysis Test","I absolutely LOVE this amazing product! It's fantastic, wonderful, and brings me so much joy! 😊❤️ However, I also hate waiting for delivery and the customer service was terrible and disappointing. Mixed emotions here - both positive and negative feelings combined.","product-review","Customer Reviewer","2024-04-10","sentiment,emotions,mixed-feelings,product-review","en","https://reviews.com/mixed-sentiment","{'sentiment': 'mixed', 'positive_score': 0.7, 'negative_score': 0.4}" +19,"Code Injection Attempt","This field contains potential code: and SQL injection: '; DROP TABLE users; -- and command injection: $(rm -rf /) and other malicious content that should be sanitized and detected by security filters.","security-test","Security Tester","2024-04-15","security,injection,xss,sql-injection,sanitization","en","https://security-test.com/injection-test","{'threat_level': 'high', 'injection_types': ['xss', 'sql', 'command']}" +20,"Long Text Performance Test","This is an extremely long document designed to test performance with large text processing. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt.","performance","Performance Tester","2024-04-20","performance,long-text,processing,load-test","en","https://perf-test.com/long-content","{'text_length': 1200, 'test_type': 'performance', 'processing_time_target': '<5s'}" \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/test_documents.json b/examples/document-graph/notebooks/data/test_documents.json new file mode 100644 index 00000000..73a66713 --- /dev/null +++ b/examples/document-graph/notebooks/data/test_documents.json @@ -0,0 +1,242 @@ +[ + { + "id": 1, + "title": "AWS Lambda Introduction", + "content": "AWS Lambda is a serverless computing service that lets you run code without provisioning or managing servers. You pay only for the compute time you consume. Lambda runs your code on a high-availability compute infrastructure and performs all of the administration of the compute resources.", + "category": "technology", + "author": "AWS Documentation", + "date": "2024-01-15", + "tags": "serverless,lambda,aws,computing", + "language": "en", + "source_url": "https://docs.aws.amazon.com/lambda/", + "metadata": {"service": "lambda", "complexity": "beginner"} + }, + { + "id": 2, + "title": "Machine Learning Basics", + "content": "Machine learning is a subset of artificial intelligence that enables computers to learn and make decisions from data without being explicitly programmed. It involves algorithms that can identify patterns in data and make predictions or classifications.", + "category": "education", + "author": "Dr. Sarah Johnson", + "date": "2024-01-20", + "tags": "machine-learning,ai,data-science,algorithms", + "language": "en", + "source_url": "https://example.com/ml-basics", + "metadata": {"difficulty": "intermediate", "field": "computer_science"} + }, + { + "id": 3, + "title": "Cybersecurity Best Practices", + "content": "Implementing strong cybersecurity measures is crucial for protecting sensitive data. Key practices include using multi-factor authentication, regular software updates, employee training, network segmentation, and incident response planning.", + "category": "security", + "author": "John Smith", + "date": "2024-01-25", + "tags": "cybersecurity,security,data-protection,best-practices", + "language": "en", + "source_url": "https://security-blog.com/best-practices", + "metadata": {"industry": "cybersecurity", "audience": "enterprise"} + }, + { + "id": 4, + "title": "Climate Change Impact", + "content": "Climate change refers to long-term shifts in global temperatures and weather patterns. While climate variations are natural, scientific evidence shows that human activities have been the main driver of climate change since the 1800s.", + "category": "environment", + "author": "Environmental Research Institute", + "date": "2024-02-01", + "tags": "climate-change,environment,sustainability,global-warming", + "language": "en", + "source_url": "https://climate-research.org/impact", + "metadata": {"topic": "environmental_science", "urgency": "high"} + }, + { + "id": 5, + "title": "Quantum Computing Fundamentals", + "content": "Quantum computing harnesses quantum mechanical phenomena to process information in fundamentally different ways than classical computers. Quantum bits (qubits) can exist in superposition states, enabling parallel processing capabilities.", + "category": "technology", + "author": "Prof. Michael Chen", + "date": "2024-02-05", + "tags": "quantum-computing,physics,technology,qubits", + "language": "en", + "source_url": "https://quantum-journal.com/fundamentals", + "metadata": {"complexity": "advanced", "field": "quantum_physics"} + }, + { + "id": 6, + "title": "Healthy Cooking Tips", + "content": "Maintaining a healthy diet starts with proper cooking techniques. Use fresh ingredients, limit processed foods, incorporate plenty of vegetables, choose lean proteins, and use healthy cooking methods like steaming, grilling, or baking instead of frying.", + "category": "health", + "author": "Nutritionist Lisa Brown", + "date": "2024-02-10", + "tags": "health,nutrition,cooking,wellness,diet", + "language": "en", + "source_url": "https://health-magazine.com/cooking-tips", + "metadata": {"category": "lifestyle", "target_audience": "general_public"} + }, + { + "id": 7, + "title": "Financial Planning Strategies", + "content": "Effective financial planning involves setting clear goals, creating a budget, building an emergency fund, investing wisely, and regularly reviewing your financial situation. Start early and be consistent with your savings and investment approach.", + "category": "finance", + "author": "Financial Advisor Mark Wilson", + "date": "2024-02-15", + "tags": "finance,investment,budgeting,savings,planning", + "language": "en", + "source_url": "https://finance-advisor.com/strategies", + "metadata": {"expertise_level": "beginner_to_intermediate", "topic": "personal_finance"} + }, + { + "id": 8, + "title": "Space Exploration History", + "content": "Space exploration began in the mid-20th century with the launch of Sputnik 1 in 1957. Major milestones include the first human in space (Yuri Gagarin, 1961), the moon landing (Apollo 11, 1969), and the development of space stations and Mars rovers.", + "category": "science", + "author": "Space History Museum", + "date": "2024-02-20", + "tags": "space,exploration,history,nasa,astronomy", + "language": "en", + "source_url": "https://space-museum.org/history", + "metadata": {"era": "20th_21st_century", "field": "aerospace"} + }, + { + "id": 9, + "title": "Sustainable Energy Solutions", + "content": "Renewable energy sources like solar, wind, and hydroelectric power are becoming increasingly important for reducing carbon emissions. These technologies offer clean alternatives to fossil fuels and are becoming more cost-effective.", + "category": "environment", + "author": "Green Energy Consortium", + "date": "2024-02-25", + "tags": "renewable-energy,sustainability,solar,wind,environment", + "language": "en", + "source_url": "https://green-energy.org/solutions", + "metadata": {"focus": "renewable_energy", "impact": "environmental"} + }, + { + "id": 10, + "title": "Digital Marketing Trends", + "content": "Digital marketing continues to evolve with new technologies and consumer behaviors. Key trends include personalization, video content, social media marketing, influencer partnerships, and data-driven decision making.", + "category": "business", + "author": "Marketing Expert Jane Davis", + "date": "2024-03-01", + "tags": "digital-marketing,social-media,trends,business,advertising", + "language": "en", + "source_url": "https://marketing-trends.com/2024", + "metadata": {"industry": "marketing", "year": "2024"} + }, + { + "id": 11, + "title": "Artificial Intelligence Ethics", + "content": "As AI systems become more prevalent, ethical considerations become crucial. Key issues include bias in algorithms, privacy concerns, job displacement, transparency in decision-making, and the need for responsible AI development.", + "category": "technology", + "author": "AI Ethics Committee", + "date": "2024-03-05", + "tags": "ai,ethics,technology,society,responsibility", + "language": "en", + "source_url": "https://ai-ethics.org/guidelines", + "metadata": {"importance": "critical", "stakeholders": "global"} + }, + { + "id": 12, + "title": "Urban Planning Principles", + "content": "Effective urban planning creates sustainable, livable communities. Key principles include mixed-use development, public transportation, green spaces, affordable housing, and community engagement in the planning process.", + "category": "urban-development", + "author": "City Planning Department", + "date": "2024-03-10", + "tags": "urban-planning,city-development,sustainability,community", + "language": "en", + "source_url": "https://city-planning.gov/principles", + "metadata": {"scope": "municipal", "focus": "sustainable_development"} + }, + { + "id": 13, + "title": "Data Breach Incident Report", + "content": "On March 15th, our security team detected unauthorized access to customer data. The breach affected 1,247 users including John Doe (SSN: 123-45-6789, email: john.doe@email.com), Jane Smith (phone: 555-123-4567), and server logs from IP address 192.168.1.100. Credit card 4532-1234-5678-9012 was also compromised. Immediate containment measures were implemented.", + "category": "security", + "author": "Security Team Lead", + "date": "2024-03-15", + "tags": "data-breach,security,incident,pii", + "language": "en", + "source_url": "https://internal-security.com/incident-001", + "metadata": {"severity": "high", "affected_users": 1247, "containment_status": "complete"} + }, + { + "id": 14, + "title": "Machien Lerning Algoritms", + "content": "Machien lerning algoritms are becomming more sofisticated every day. Thier ability to proces large amounts of data and identyfy paterns is truely remarkabel. However, ther are stil chalenges with bias and interpretabilty that need to be adressed.", + "category": "technology", + "author": "Dr. Alex Thompson", + "date": "2024-03-20", + "tags": "machine-learning,algorithms,ai,spelling-errors", + "language": "en", + "source_url": "https://ml-research.edu/algorithms", + "metadata": {"quality": "poor_spelling", "needs_correction": true} + }, + { + "id": 15, + "title": "API Configuration Guide", + "content": "To configure the API endpoint, use the following JSON structure: {\"endpoint\": \"https://api.example.com/v1\", \"auth\": {\"type\": \"bearer\", \"token\": \"sk-1234567890abcdef\"}, \"timeout\": 30000, \"retries\": 3, \"headers\": {\"Content-Type\": \"application/json\", \"User-Agent\": \"MyApp/1.0\"}}. This configuration ensures proper authentication and error handling.", + "category": "technical", + "author": "DevOps Engineer", + "date": "2024-03-25", + "tags": "api,configuration,json,development", + "language": "en", + "source_url": "https://dev-docs.com/api-config", + "metadata": {"format": "json_embedded", "complexity": "intermediate"} + }, + { + "id": 16, + "title": "Mixed Language Content", + "content": "This document contains multiple languages. English: Hello world! Spanish: ¡Hola mundo! French: Bonjour le monde! German: Hallo Welt! Japanese: こんにちは世界！Chinese: 你好世界！Russian: Привет мир! Arabic: مرحبا بالعالم! This tests language detection capabilities.", + "category": "linguistics", + "author": "Language Research Team", + "date": "2024-04-01", + "tags": "multilingual,language-detection,international", + "language": "mixed", + "source_url": "https://lang-research.org/mixed-content", + "metadata": {"languages": ["en", "es", "fr", "de", "ja", "zh", "ru", "ar"], "test_type": "language_detection"} + }, + { + "id": 17, + "title": "Toxic Content Example", + "content": "This content contains inappropriate language and offensive terms that should be filtered out. Words like damn, hell, and other profanity need to be detected and handled appropriately by content moderation systems. This tests toxicity detection.", + "category": "content-moderation", + "author": "Content Safety Team", + "date": "2024-04-05", + "tags": "toxicity,content-moderation,filtering", + "language": "en", + "source_url": "https://safety.example.com/toxic-test", + "metadata": {"toxicity_level": "moderate", "requires_filtering": true} + }, + { + "id": 18, + "title": "Sentiment Analysis Test", + "content": "I absolutely LOVE this amazing product! It's fantastic, wonderful, and brings me so much joy! 😊❤️ However, I also hate waiting for delivery and the customer service was terrible and disappointing. Mixed emotions here - both positive and negative feelings combined.", + "category": "product-review", + "author": "Customer Reviewer", + "date": "2024-04-10", + "tags": "sentiment,emotions,mixed-feelings,product-review", + "language": "en", + "source_url": "https://reviews.com/mixed-sentiment", + "metadata": {"sentiment": "mixed", "positive_score": 0.7, "negative_score": 0.4} + }, + { + "id": 19, + "title": "Code Injection Attempt", + "content": "This field contains potential code: and SQL injection: '; DROP TABLE users; -- and command injection: $(rm -rf /) and other malicious content that should be sanitized and detected by security filters.", + "category": "security-test", + "author": "Security Tester", + "date": "2024-04-15", + "tags": "security,injection,xss,sql-injection,sanitization", + "language": "en", + "source_url": "https://security-test.com/injection-test", + "metadata": {"threat_level": "high", "injection_types": ["xss", "sql", "command"]} + }, + { + "id": 20, + "title": "Long Text Performance Test", + "content": "This is an extremely long document designed to test performance with large text processing. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt.", + "category": "performance", + "author": "Performance Tester", + "date": "2024-04-20", + "tags": "performance,long-text,processing,load-test", + "language": "en", + "source_url": "https://perf-test.com/long-content", + "metadata": {"text_length": 1200, "test_type": "performance", "processing_time_target": "<5s"} + } +] \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/test_documents.parquet b/examples/document-graph/notebooks/data/test_documents.parquet new file mode 100644 index 00000000..8edadf26 Binary files /dev/null and b/examples/document-graph/notebooks/data/test_documents.parquet differ diff --git a/examples/document-graph/notebooks/data/test_documents.skip.csv b/examples/document-graph/notebooks/data/test_documents.skip.csv new file mode 100644 index 00000000..0a5ca3f7 --- /dev/null +++ b/examples/document-graph/notebooks/data/test_documents.skip.csv @@ -0,0 +1,42 @@ +# Skipped lines from test_documents.csv +# Header: id,title,content,category,author,date,tags,language,source_url,metadata +# Line 2: Expected 10 fields, got 14 +1,"AWS Lambda Introduction","AWS Lambda is a serverless computing service that lets you run code without provisioning or managing servers. You pay only for the compute time you consume. Lambda runs your code on a high-availability compute infrastructure and performs all of the administration of the compute resources.","technology","AWS Documentation","2024-01-15","serverless,lambda,aws,computing","en","https://docs.aws.amazon.com/lambda/","{'service': 'lambda', 'complexity': 'beginner'}" +# Line 3: Expected 10 fields, got 14 +2,"Machine Learning Basics","Machine learning is a subset of artificial intelligence that enables computers to learn and make decisions from data without being explicitly programmed. It involves algorithms that can identify patterns in data and make predictions or classifications.","education","Dr. Sarah Johnson","2024-01-20","machine-learning,ai,data-science,algorithms","en","https://example.com/ml-basics","{'difficulty': 'intermediate', 'field': 'computer_science'}" +# Line 4: Expected 10 fields, got 18 +3,"Cybersecurity Best Practices","Implementing strong cybersecurity measures is crucial for protecting sensitive data. Key practices include using multi-factor authentication, regular software updates, employee training, network segmentation, and incident response planning.","security","John Smith","2024-01-25","cybersecurity,security,data-protection,best-practices","en","https://security-blog.com/best-practices","{'industry': 'cybersecurity', 'audience': 'enterprise'}" +# Line 5: Expected 10 fields, got 15 +4,"Climate Change Impact","Climate change refers to long-term shifts in global temperatures and weather patterns. While climate variations are natural, scientific evidence shows that human activities have been the main driver of climate change since the 1800s.","environment","Environmental Research Institute","2024-02-01","climate-change,environment,sustainability,global-warming","en","https://climate-research.org/impact","{'topic': 'environmental_science', 'urgency': 'high'}" +# Line 6: Expected 10 fields, got 15 +5,"Quantum Computing Fundamentals","Quantum computing harnesses quantum mechanical phenomena to process information in fundamentally different ways than classical computers. Quantum bits (qubits) can exist in superposition states, enabling parallel processing capabilities.","technology","Prof. Michael Chen","2024-02-05","quantum-computing,physics,technology,qubits","en","https://quantum-journal.com/fundamentals","{'complexity': 'advanced', 'field': 'quantum_physics'}" +# Line 7: Expected 10 fields, got 21 +6,"Healthy Cooking Tips","Maintaining a healthy diet starts with proper cooking techniques. Use fresh ingredients, limit processed foods, incorporate plenty of vegetables, choose lean proteins, and use healthy cooking methods like steaming, grilling, or baking instead of frying.","health","Nutritionist Lisa Brown","2024-02-10","health,nutrition,cooking,wellness,diet","en","https://health-magazine.com/cooking-tips","{'category': 'lifestyle', 'target_audience': 'general_public'}" +# Line 8: Expected 10 fields, got 19 +7,"Financial Planning Strategies","Effective financial planning involves setting clear goals, creating a budget, building an emergency fund, investing wisely, and regularly reviewing your financial situation. Start early and be consistent with your savings and investment approach.","finance","Financial Advisor Mark Wilson","2024-02-15","finance,investment,budgeting,savings,planning","en","https://finance-advisor.com/strategies","{'expertise_level': 'beginner_to_intermediate', 'topic': 'personal_finance'}" +# Line 9: Expected 10 fields, got 19 +8,"Space Exploration History","Space exploration began in the mid-20th century with the launch of Sputnik 1 in 1957. Major milestones include the first human in space (Yuri Gagarin, 1961), the moon landing (Apollo 11, 1969), and the development of space stations and Mars rovers.","science","Space History Museum","2024-02-20","space,exploration,history,nasa,astronomy","en","https://space-museum.org/history","{'era': '20th_21st_century', 'field': 'aerospace'}" +# Line 10: Expected 10 fields, got 17 +9,"Sustainable Energy Solutions","Renewable energy sources like solar, wind, and hydroelectric power are becoming increasingly important for reducing carbon emissions. These technologies offer clean alternatives to fossil fuels and are becoming more cost-effective.","environment","Green Energy Consortium","2024-02-25","renewable-energy,sustainability,solar,wind,environment","en","https://green-energy.org/solutions","{'focus': 'renewable_energy', 'impact': 'environmental'}" +# Line 11: Expected 10 fields, got 19 +10,"Digital Marketing Trends","Digital marketing continues to evolve with new technologies and consumer behaviors. Key trends include personalization, video content, social media marketing, influencer partnerships, and data-driven decision making.","business","Marketing Expert Jane Davis","2024-03-01","digital-marketing,social-media,trends,business,advertising","en","https://marketing-trends.com/2024","{'industry': 'marketing', 'year': '2024'}" +# Line 12: Expected 10 fields, got 20 +11,"Artificial Intelligence Ethics","As AI systems become more prevalent, ethical considerations become crucial. Key issues include bias in algorithms, privacy concerns, job displacement, transparency in decision-making, and the need for responsible AI development.","technology","AI Ethics Committee","2024-03-05","ai,ethics,technology,society,responsibility","en","https://ai-ethics.org/guidelines","{'importance': 'critical', 'stakeholders': 'global'}" +# Line 13: Expected 10 fields, got 19 +12,"Urban Planning Principles","Effective urban planning creates sustainable, livable communities. Key principles include mixed-use development, public transportation, green spaces, affordable housing, and community engagement in the planning process.","urban-development","City Planning Department","2024-03-10","urban-planning,city-development,sustainability,community","en","https://city-planning.gov/principles","{'scope': 'municipal', 'focus': 'sustainable_development'}" +# Line 14: Expected 10 fields, got 20 +13,"Data Breach Incident Report","On March 15th, our security team detected unauthorized access to customer data. The breach affected 1,247 users including John Doe (SSN: 123-45-6789, email: john.doe@email.com), Jane Smith (phone: 555-123-4567), and server logs from IP address 192.168.1.100. Credit card 4532-1234-5678-9012 was also compromised. Immediate containment measures were implemented.","security","Security Team Lead","2024-03-15","data-breach,security,incident,pii","en","https://internal-security.com/incident-001","{'severity': 'high', 'affected_users': 1247, 'containment_status': 'complete'}" +# Line 15: Expected 10 fields, got 15 +14,"Machien Lerning Algoritms","Machien lerning algoritms are becomming more sofisticated every day. Thier ability to proces large amounts of data and identyfy paterns is truely remarkabel. However, ther are stil chalenges with bias and interpretabilty that need to be adressed.","technology","Dr. Alex Thompson","2024-03-20","machine-learning,algorithms,ai,spelling-errors","en","https://ml-research.edu/algorithms","{'quality': 'poor_spelling', 'needs_correction': true}" +# Line 16: Expected 10 fields, got 21 +15,"API Configuration Guide","To configure the API endpoint, use the following JSON structure: {\"endpoint\": \"https://api.example.com/v1\", \"auth\": {\"type\": \"bearer\", \"token\": \"sk-1234567890abcdef\"}, \"timeout\": 30000, \"retries\": 3, \"headers\": {\"Content-Type\": \"application/json\", \"User-Agent\": \"MyApp/1.0\"}}. This configuration ensures proper authentication and error handling.","technical","DevOps Engineer","2024-03-25","api,configuration,json,development","en","https://dev-docs.com/api-config","{'format': 'json_embedded', 'complexity': 'intermediate'}" +# Line 17: Expected 10 fields, got 20 +16,"Mixed Language Content","This document contains multiple languages. English: Hello world! Spanish: ¡Hola mundo! French: Bonjour le monde! German: Hallo Welt! Japanese: こんにちは世界！Chinese: 你好世界！Russian: Привет мир! Arabic: مرحبا بالعالم! This tests language detection capabilities.","linguistics","Language Research Team","2024-04-01","multilingual,language-detection,international","mixed","https://lang-research.org/mixed-content","{'languages': ['en', 'es', 'fr', 'de', 'ja', 'zh', 'ru', 'ar'], 'test_type': 'language_detection'}" +# Line 18: Expected 10 fields, got 15 +17,"Toxic Content Example","This content contains inappropriate language and offensive terms that should be filtered out. Words like damn, hell, and other profanity need to be detected and handled appropriately by content moderation systems. This tests toxicity detection.","content-moderation","Content Safety Team","2024-04-05","toxicity,content-moderation,filtering","en","https://safety.example.com/toxic-test","{'toxicity_level': 'moderate', 'requires_filtering': true}" +# Line 19: Expected 10 fields, got 18 +18,"Sentiment Analysis Test","I absolutely LOVE this amazing product! It's fantastic, wonderful, and brings me so much joy! 😊❤️ However, I also hate waiting for delivery and the customer service was terrible and disappointing. Mixed emotions here - both positive and negative feelings combined.","product-review","Customer Reviewer","2024-04-10","sentiment,emotions,mixed-feelings,product-review","en","https://reviews.com/mixed-sentiment","{'sentiment': 'mixed', 'positive_score': 0.7, 'negative_score': 0.4}" +# Line 20: Expected 10 fields, got 17 +19,"Code Injection Attempt","This field contains potential code: and SQL injection: '; DROP TABLE users; -- and command injection: $(rm -rf /) and other malicious content that should be sanitized and detected by security filters.","security-test","Security Tester","2024-04-15","security,injection,xss,sql-injection,sanitization","en","https://security-test.com/injection-test","{'threat_level': 'high', 'injection_types': ['xss', 'sql', 'command']}" +# Line 21: Expected 10 fields, got 22 +20,"Long Text Performance Test","This is an extremely long document designed to test performance with large text processing. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt.","performance","Performance Tester","2024-04-20","performance,long-text,processing,load-test","en","https://perf-test.com/long-content","{'text_length': 1200, 'test_type': 'performance', 'processing_time_target': '<5s'}" \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/test_documents.xml b/examples/document-graph/notebooks/data/test_documents.xml new file mode 100644 index 00000000..34718ad9 --- /dev/null +++ b/examples/document-graph/notebooks/data/test_documents.xml @@ -0,0 +1,306 @@ + + + + 1 + AWS Lambda Introduction + AWS Lambda is a serverless computing service that lets you run code without provisioning or managing servers. You pay only for the compute time you consume. Lambda runs your code on a high-availability compute infrastructure and performs all of the administration of the compute resources. + technology + AWS Documentation + 2024-01-15 + serverless,lambda,aws,computing + en + https://docs.aws.amazon.com/lambda/ + + lambda + beginner + + + + 2 + Machine Learning Basics + Machine learning is a subset of artificial intelligence that enables computers to learn and make decisions from data without being explicitly programmed. It involves algorithms that can identify patterns in data and make predictions or classifications. + education + Dr. Sarah Johnson + 2024-01-20 + machine-learning,ai,data-science,algorithms + en + https://example.com/ml-basics + + intermediate + computer_science + + + + 3 + Cybersecurity Best Practices + Implementing strong cybersecurity measures is crucial for protecting sensitive data. Key practices include using multi-factor authentication, regular software updates, employee training, network segmentation, and incident response planning. + security + John Smith + 2024-01-25 + cybersecurity,security,data-protection,best-practices + en + https://security-blog.com/best-practices + + cybersecurity + enterprise + + + + 4 + Climate Change Impact + Climate change refers to long-term shifts in global temperatures and weather patterns. While climate variations are natural, scientific evidence shows that human activities have been the main driver of climate change since the 1800s. + environment + Environmental Research Institute + 2024-02-01 + climate-change,environment,sustainability,global-warming + en + https://climate-research.org/impact + + environmental_science + high + + + + 5 + Quantum Computing Fundamentals + Quantum computing harnesses quantum mechanical phenomena to process information in fundamentally different ways than classical computers. Quantum bits (qubits) can exist in superposition states, enabling parallel processing capabilities. + technology + Prof. Michael Chen + 2024-02-05 + quantum-computing,physics,technology,qubits + en + https://quantum-journal.com/fundamentals + + advanced + quantum_physics + + + + 6 + Healthy Cooking Tips + Maintaining a healthy diet starts with proper cooking techniques. Use fresh ingredients, limit processed foods, incorporate plenty of vegetables, choose lean proteins, and use healthy cooking methods like steaming, grilling, or baking instead of frying. + health + Nutritionist Lisa Brown + 2024-02-10 + health,nutrition,cooking,wellness,diet + en + https://health-magazine.com/cooking-tips + + lifestyle + general_public + + + + 7 + Financial Planning Strategies + Effective financial planning involves setting clear goals, creating a budget, building an emergency fund, investing wisely, and regularly reviewing your financial situation. Start early and be consistent with your savings and investment approach. + finance + Financial Advisor Mark Wilson + 2024-02-15 + finance,investment,budgeting,savings,planning + en + https://finance-advisor.com/strategies + + beginner_to_intermediate + personal_finance + + + + 8 + Space Exploration History + Space exploration began in the mid-20th century with the launch of Sputnik 1 in 1957. Major milestones include the first human in space (Yuri Gagarin, 1961), the moon landing (Apollo 11, 1969), and the development of space stations and Mars rovers. + science + Space History Museum + 2024-02-20 + space,exploration,history,nasa,astronomy + en + https://space-museum.org/history + + 20th_21st_century + aerospace + + + + 9 + Sustainable Energy Solutions + Renewable energy sources like solar, wind, and hydroelectric power are becoming increasingly important for reducing carbon emissions. These technologies offer clean alternatives to fossil fuels and are becoming more cost-effective. + environment + Green Energy Consortium + 2024-02-25 + renewable-energy,sustainability,solar,wind,environment + en + https://green-energy.org/solutions + + renewable_energy + environmental + + + + 10 + Digital Marketing Trends + Digital marketing continues to evolve with new technologies and consumer behaviors. Key trends include personalization, video content, social media marketing, influencer partnerships, and data-driven decision making. + business + Marketing Expert Jane Davis + 2024-03-01 + digital-marketing,social-media,trends,business,advertising + en + https://marketing-trends.com/2024 + + marketing + 2024 + + + + 11 + Artificial Intelligence Ethics + As AI systems become more prevalent, ethical considerations become crucial. Key issues include bias in algorithms, privacy concerns, job displacement, transparency in decision-making, and the need for responsible AI development. + technology + AI Ethics Committee + 2024-03-05 + ai,ethics,technology,society,responsibility + en + https://ai-ethics.org/guidelines + + critical + global + + + + 12 + Urban Planning Principles + Effective urban planning creates sustainable, livable communities. Key principles include mixed-use development, public transportation, green spaces, affordable housing, and community engagement in the planning process. + urban-development + City Planning Department + 2024-03-10 + urban-planning,city-development,sustainability,community + en + https://city-planning.gov/principles + + municipal + sustainable_development + + + + 13 + Data Breach Incident Report + On March 15th, our security team detected unauthorized access to customer data. The breach affected 1,247 users including John Doe (SSN: 123-45-6789, email: john.doe@email.com), Jane Smith (phone: 555-123-4567), and server logs from IP address 192.168.1.100. Credit card 4532-1234-5678-9012 was also compromised. Immediate containment measures were implemented. + security + Security Team Lead + 2024-03-15 + data-breach,security,incident,pii + en + https://internal-security.com/incident-001 + + high + 1247 + complete + + + + 14 + Machien Lerning Algoritms + Machien lerning algoritms are becomming more sofisticated every day. Thier ability to proces large amounts of data and identyfy paterns is truely remarkabel. However, ther are stil chalenges with bias and interpretabilty that need to be adressed. + technology + Dr. Alex Thompson + 2024-03-20 + machine-learning,algorithms,ai,spelling-errors + en + https://ml-research.edu/algorithms + + poor_spelling + true + + + + 15 + API Configuration Guide + To configure the API endpoint, use the following JSON structure: {"endpoint": "https://api.example.com/v1", "auth": {"type": "bearer", "token": "sk-1234567890abcdef"}, "timeout": 30000, "retries": 3, "headers": {"Content-Type": "application/json", "User-Agent": "MyApp/1.0"}}. This configuration ensures proper authentication and error handling. + technical + DevOps Engineer + 2024-03-25 + api,configuration,json,development + en + https://dev-docs.com/api-config + + json_embedded + intermediate + + + + 16 + Mixed Language Content + This document contains multiple languages. English: Hello world! Spanish: ¡Hola mundo! French: Bonjour le monde! German: Hallo Welt! Japanese: こんにちは世界！Chinese: 你好世界！Russian: Привет мир! Arabic: مرحبا بالعالم! This tests language detection capabilities. + linguistics + Language Research Team + 2024-04-01 + multilingual,language-detection,international + mixed + https://lang-research.org/mixed-content + + en,es,fr,de,ja,zh,ru,ar + language_detection + + + + 17 + Toxic Content Example + This content contains inappropriate language and offensive terms that should be filtered out. Words like damn, hell, and other profanity need to be detected and handled appropriately by content moderation systems. This tests toxicity detection. + content-moderation + Content Safety Team + 2024-04-05 + toxicity,content-moderation,filtering + en + https://safety.example.com/toxic-test + + moderate + true + + + + 18 + Sentiment Analysis Test + I absolutely LOVE this amazing product! It's fantastic, wonderful, and brings me so much joy! 😊❤️ However, I also hate waiting for delivery and the customer service was terrible and disappointing. Mixed emotions here - both positive and negative feelings combined. + product-review + Customer Reviewer + 2024-04-10 + sentiment,emotions,mixed-feelings,product-review + en + https://reviews.com/mixed-sentiment + + mixed + 0.7 + 0.4 + + + + 19 + Code Injection Attempt + This field contains potential code: <script>alert('XSS')</script> and SQL injection: '; DROP TABLE users; -- and command injection: $(rm -rf /) and other malicious content that should be sanitized and detected by security filters. + security-test + Security Tester + 2024-04-15 + security,injection,xss,sql-injection,sanitization + en + https://security-test.com/injection-test + + high + xss,sql,command + + + + 20 + Long Text Performance Test + This is an extremely long document designed to test performance with large text processing. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt. + performance + Performance Tester + 2024-04-20 + performance,long-text,processing,load-test + en + https://perf-test.com/long-content + + 1200 + performance + <5s + + + \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/test_documents.yaml b/examples/document-graph/notebooks/data/test_documents.yaml new file mode 100644 index 00000000..1cd461ab --- /dev/null +++ b/examples/document-graph/notebooks/data/test_documents.yaml @@ -0,0 +1,266 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +documents: + - id: 1 + title: "AWS Lambda Introduction" + content: "AWS Lambda is a serverless computing service that lets you run code without provisioning or managing servers. You pay only for the compute time you consume. Lambda runs your code on a high-availability compute infrastructure and performs all of the administration of the compute resources." + category: "technology" + author: "AWS Documentation" + date: "2024-01-15" + tags: "serverless,lambda,aws,computing" + language: "en" + source_url: "https://docs.aws.amazon.com/lambda/" + metadata: + service: "lambda" + complexity: "beginner" + + - id: 2 + title: "Machine Learning Basics" + content: "Machine learning is a subset of artificial intelligence that enables computers to learn and make decisions from data without being explicitly programmed. It involves algorithms that can identify patterns in data and make predictions or classifications." + category: "education" + author: "Dr. Sarah Johnson" + date: "2024-01-20" + tags: "machine-learning,ai,data-science,algorithms" + language: "en" + source_url: "https://example.com/ml-basics" + metadata: + difficulty: "intermediate" + field: "computer_science" + + - id: 3 + title: "Cybersecurity Best Practices" + content: "Implementing strong cybersecurity measures is crucial for protecting sensitive data. Key practices include using multi-factor authentication, regular software updates, employee training, network segmentation, and incident response planning." + category: "security" + author: "John Smith" + date: "2024-01-25" + tags: "cybersecurity,security,data-protection,best-practices" + language: "en" + source_url: "https://security-blog.com/best-practices" + metadata: + industry: "cybersecurity" + audience: "enterprise" + + - id: 4 + title: "Climate Change Impact" + content: "Climate change refers to long-term shifts in global temperatures and weather patterns. While climate variations are natural, scientific evidence shows that human activities have been the main driver of climate change since the 1800s." + category: "environment" + author: "Environmental Research Institute" + date: "2024-02-01" + tags: "climate-change,environment,sustainability,global-warming" + language: "en" + source_url: "https://climate-research.org/impact" + metadata: + topic: "environmental_science" + urgency: "high" + + - id: 5 + title: "Quantum Computing Fundamentals" + content: "Quantum computing harnesses quantum mechanical phenomena to process information in fundamentally different ways than classical computers. Quantum bits (qubits) can exist in superposition states, enabling parallel processing capabilities." + category: "technology" + author: "Prof. Michael Chen" + date: "2024-02-05" + tags: "quantum-computing,physics,technology,qubits" + language: "en" + source_url: "https://quantum-journal.com/fundamentals" + metadata: + complexity: "advanced" + field: "quantum_physics" + + - id: 6 + title: "Healthy Cooking Tips" + content: "Maintaining a healthy diet starts with proper cooking techniques. Use fresh ingredients, limit processed foods, incorporate plenty of vegetables, choose lean proteins, and use healthy cooking methods like steaming, grilling, or baking instead of frying." + category: "health" + author: "Nutritionist Lisa Brown" + date: "2024-02-10" + tags: "health,nutrition,cooking,wellness,diet" + language: "en" + source_url: "https://health-magazine.com/cooking-tips" + metadata: + category: "lifestyle" + target_audience: "general_public" + + - id: 7 + title: "Financial Planning Strategies" + content: "Effective financial planning involves setting clear goals, creating a budget, building an emergency fund, investing wisely, and regularly reviewing your financial situation. Start early and be consistent with your savings and investment approach." + category: "finance" + author: "Financial Advisor Mark Wilson" + date: "2024-02-15" + tags: "finance,investment,budgeting,savings,planning" + language: "en" + source_url: "https://finance-advisor.com/strategies" + metadata: + expertise_level: "beginner_to_intermediate" + topic: "personal_finance" + + - id: 8 + title: "Space Exploration History" + content: "Space exploration began in the mid-20th century with the launch of Sputnik 1 in 1957. Major milestones include the first human in space (Yuri Gagarin, 1961), the moon landing (Apollo 11, 1969), and the development of space stations and Mars rovers." + category: "science" + author: "Space History Museum" + date: "2024-02-20" + tags: "space,exploration,history,nasa,astronomy" + language: "en" + source_url: "https://space-museum.org/history" + metadata: + era: "20th_21st_century" + field: "aerospace" + + - id: 9 + title: "Sustainable Energy Solutions" + content: "Renewable energy sources like solar, wind, and hydroelectric power are becoming increasingly important for reducing carbon emissions. These technologies offer clean alternatives to fossil fuels and are becoming more cost-effective." + category: "environment" + author: "Green Energy Consortium" + date: "2024-02-25" + tags: "renewable-energy,sustainability,solar,wind,environment" + language: "en" + source_url: "https://green-energy.org/solutions" + metadata: + focus: "renewable_energy" + impact: "environmental" + + - id: 10 + title: "Digital Marketing Trends" + content: "Digital marketing continues to evolve with new technologies and consumer behaviors. Key trends include personalization, video content, social media marketing, influencer partnerships, and data-driven decision making." + category: "business" + author: "Marketing Expert Jane Davis" + date: "2024-03-01" + tags: "digital-marketing,social-media,trends,business,advertising" + language: "en" + source_url: "https://marketing-trends.com/2024" + metadata: + industry: "marketing" + year: "2024" + + - id: 11 + title: "Artificial Intelligence Ethics" + content: "As AI systems become more prevalent, ethical considerations become crucial. Key issues include bias in algorithms, privacy concerns, job displacement, transparency in decision-making, and the need for responsible AI development." + category: "technology" + author: "AI Ethics Committee" + date: "2024-03-05" + tags: "ai,ethics,technology,society,responsibility" + language: "en" + source_url: "https://ai-ethics.org/guidelines" + metadata: + importance: "critical" + stakeholders: "global" + + - id: 12 + title: "Urban Planning Principles" + content: "Effective urban planning creates sustainable, livable communities. Key principles include mixed-use development, public transportation, green spaces, affordable housing, and community engagement in the planning process." + category: "urban-development" + author: "City Planning Department" + date: "2024-03-10" + tags: "urban-planning,city-development,sustainability,community" + language: "en" + source_url: "https://city-planning.gov/principles" + metadata: + scope: "municipal" + focus: "sustainable_development" + + - id: 13 + title: "Data Breach Incident Report" + content: "On March 15th, our security team detected unauthorized access to customer data. The breach affected 1,247 users including John Doe (SSN: 123-45-6789, email: john.doe@email.com), Jane Smith (phone: 555-123-4567), and server logs from IP address 192.168.1.100. Credit card 4532-1234-5678-9012 was also compromised. Immediate containment measures were implemented." + category: "security" + author: "Security Team Lead" + date: "2024-03-15" + tags: "data-breach,security,incident,pii" + language: "en" + source_url: "https://internal-security.com/incident-001" + metadata: + severity: "high" + affected_users: 1247 + containment_status: "complete" + + - id: 14 + title: "Machien Lerning Algoritms" + content: "Machien lerning algoritms are becomming more sofisticated every day. Thier ability to proces large amounts of data and identyfy paterns is truely remarkabel. However, ther are stil chalenges with bias and interpretabilty that need to be adressed." + category: "technology" + author: "Dr. Alex Thompson" + date: "2024-03-20" + tags: "machine-learning,algorithms,ai,spelling-errors" + language: "en" + source_url: "https://ml-research.edu/algorithms" + metadata: + quality: "poor_spelling" + needs_correction: true + + - id: 15 + title: "API Configuration Guide" + content: 'To configure the API endpoint, use the following JSON structure: {"endpoint": "https://api.example.com/v1", "auth": {"type": "bearer", "token": "sk-1234567890abcdef"}, "timeout": 30000, "retries": 3, "headers": {"Content-Type": "application/json", "User-Agent": "MyApp/1.0"}}. This configuration ensures proper authentication and error handling.' + category: "technical" + author: "DevOps Engineer" + date: "2024-03-25" + tags: "api,configuration,json,development" + language: "en" + source_url: "https://dev-docs.com/api-config" + metadata: + format: "json_embedded" + complexity: "intermediate" + + - id: 16 + title: "Mixed Language Content" + content: "This document contains multiple languages. English: Hello world! Spanish: ¡Hola mundo! French: Bonjour le monde! German: Hallo Welt! Japanese: こんにちは世界！Chinese: 你好世界！Russian: Привет мир! Arabic: مرحبا بالعالم! This tests language detection capabilities." + category: "linguistics" + author: "Language Research Team" + date: "2024-04-01" + tags: "multilingual,language-detection,international" + language: "mixed" + source_url: "https://lang-research.org/mixed-content" + metadata: + languages: ["en", "es", "fr", "de", "ja", "zh", "ru", "ar"] + test_type: "language_detection" + + - id: 17 + title: "Toxic Content Example" + content: "This content contains inappropriate language and offensive terms that should be filtered out. Words like damn, hell, and other profanity need to be detected and handled appropriately by content moderation systems. This tests toxicity detection." + category: "content-moderation" + author: "Content Safety Team" + date: "2024-04-05" + tags: "toxicity,content-moderation,filtering" + language: "en" + source_url: "https://safety.example.com/toxic-test" + metadata: + toxicity_level: "moderate" + requires_filtering: true + + - id: 18 + title: "Sentiment Analysis Test" + content: "I absolutely LOVE this amazing product! It's fantastic, wonderful, and brings me so much joy! 😊❤️ However, I also hate waiting for delivery and the customer service was terrible and disappointing. Mixed emotions here - both positive and negative feelings combined." + category: "product-review" + author: "Customer Reviewer" + date: "2024-04-10" + tags: "sentiment,emotions,mixed-feelings,product-review" + language: "en" + source_url: "https://reviews.com/mixed-sentiment" + metadata: + sentiment: "mixed" + positive_score: 0.7 + negative_score: 0.4 + + - id: 19 + title: "Code Injection Attempt" + content: "This field contains potential code: and SQL injection: '; DROP TABLE users; -- and command injection: $(rm -rf /) and other malicious content that should be sanitized and detected by security filters." + category: "security-test" + author: "Security Tester" + date: "2024-04-15" + tags: "security,injection,xss,sql-injection,sanitization" + language: "en" + source_url: "https://security-test.com/injection-test" + metadata: + threat_level: "high" + injection_types: ["xss", "sql", "command"] + + - id: 20 + title: "Long Text Performance Test" + content: "This is an extremely long document designed to test performance with large text processing. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt." + category: "performance" + author: "Performance Tester" + date: "2024-04-20" + tags: "performance,long-text,processing,load-test" + language: "en" + source_url: "https://perf-test.com/long-content" + metadata: + text_length: 1200 + test_type: "performance" + processing_time_target: "<5s" \ No newline at end of file diff --git a/examples/document-graph/notebooks/data/users.csv b/examples/document-graph/notebooks/data/users.csv new file mode 100644 index 00000000..c33333ab --- /dev/null +++ b/examples/document-graph/notebooks/data/users.csv @@ -0,0 +1,6 @@ +id,name,email,role,account_id +u1,Alice,alice@corp.com,admin,a1 +u2,Bob,bob@corp.com,analyst,a1 +u3,Charlie,charlie@corp.com,viewer,a2 +u4,Diana,diana@corp.com,admin,a2 +u5,Eve,eve@corp.com,analyst,a3 diff --git a/examples/document-graph/notebooks/hybrid_data/data/authors.json b/examples/document-graph/notebooks/hybrid_data/data/authors.json new file mode 100644 index 00000000..259059be --- /dev/null +++ b/examples/document-graph/notebooks/hybrid_data/data/authors.json @@ -0,0 +1,100 @@ +{ + "authors": [ + { + "author_id": "A001", + "name": "Dr. John Smith", + "email": "j.smith@university.edu", + "affiliation": "Stanford University", + "department": "Computer Science", + "h_index": 45, + "research_interests": ["graph neural networks", "machine learning", "knowledge graphs"], + "orcid": "0000-0001-2345-6789", + "publications_count": 127, + "recent_papers": ["P001"] + }, + { + "author_id": "A002", + "name": "Dr. Alice Johnson", + "email": "a.johnson@tech.edu", + "affiliation": "MIT", + "department": "EECS", + "h_index": 38, + "research_interests": ["neural networks", "deep learning", "AI systems"], + "orcid": "0000-0002-3456-7890", + "publications_count": 89, + "recent_papers": ["P001"] + }, + { + "author_id": "A003", + "name": "Dr. Michael Brown", + "email": "m.brown@research.org", + "affiliation": "Google Research", + "department": "AI Research", + "h_index": 52, + "research_interests": ["knowledge representation", "graph algorithms", "semantic web"], + "orcid": "0000-0003-4567-8901", + "publications_count": 156, + "recent_papers": ["P001"] + }, + { + "author_id": "A004", + "name": "Dr. Lisa Davis", + "email": "l.davis@ai.edu", + "affiliation": "Carnegie Mellon University", + "department": "Language Technologies Institute", + "h_index": 41, + "research_interests": ["RAG systems", "question answering", "information retrieval"], + "orcid": "0000-0004-5678-9012", + "publications_count": 94, + "recent_papers": ["P002"] + }, + { + "author_id": "A005", + "name": "Dr. Kevin Wilson", + "email": "k.wilson@nlp.edu", + "affiliation": "University of Washington", + "department": "Linguistics", + "h_index": 33, + "research_interests": ["natural language processing", "semantic search", "computational linguistics"], + "orcid": "0000-0005-6789-0123", + "publications_count": 67, + "recent_papers": ["P002"] + }, + { + "author_id": "A006", + "name": "Dr. Rachel Taylor", + "email": "r.taylor@systems.edu", + "affiliation": "UC Berkeley", + "department": "Computer Science", + "h_index": 29, + "research_interests": ["distributed systems", "database systems", "hybrid architectures"], + "orcid": "0000-0006-7890-1234", + "publications_count": 78, + "recent_papers": ["P002"] + }, + { + "author_id": "A007", + "name": "Dr. Paul Anderson", + "email": "p.anderson@db.edu", + "affiliation": "University of California San Diego", + "department": "Database Systems Lab", + "h_index": 47, + "research_interests": ["graph databases", "multi-tenancy", "database security"], + "orcid": "0000-0007-8901-2345", + "publications_count": 112, + "recent_papers": ["P003"] + }, + { + "author_id": "A008", + "name": "Dr. Maria Garcia", + "email": "m.garcia@security.edu", + "affiliation": "Georgia Tech", + "department": "Cybersecurity", + "h_index": 35, + "research_interests": ["database security", "access control", "privacy"], + "orcid": "0000-0008-9012-3456", + "publications_count": 89, + "recent_papers": ["P003"] + } + ] +} \ No newline at end of file diff --git a/examples/document-graph/notebooks/hybrid_data/data/citations.csv b/examples/document-graph/notebooks/hybrid_data/data/citations.csv new file mode 100644 index 00000000..d4d808e7 --- /dev/null +++ b/examples/document-graph/notebooks/hybrid_data/data/citations.csv @@ -0,0 +1,6 @@ +citation_id,citing_paper,cited_paper,citation_type,citation_count,publication_year +C001,P001,P003,reference,1,2023 +C002,P002,P001,citation,1,2023 +C003,P001,P002,reference,1,2023 +C004,P003,P001,citation,2,2022 +C005,P002,P003,reference,1,2023 \ No newline at end of file diff --git a/examples/document-graph/notebooks/hybrid_data/data/citations.xlsx b/examples/document-graph/notebooks/hybrid_data/data/citations.xlsx new file mode 100644 index 00000000..26ba9c30 Binary files /dev/null and b/examples/document-graph/notebooks/hybrid_data/data/citations.xlsx differ diff --git a/examples/document-graph/notebooks/hybrid_data/data/research_papers.csv b/examples/document-graph/notebooks/hybrid_data/data/research_papers.csv new file mode 100644 index 00000000..b64919bc --- /dev/null +++ b/examples/document-graph/notebooks/hybrid_data/data/research_papers.csv @@ -0,0 +1,6 @@ +paper_id,title,abstract,authors,publication_date,journal,doi,keywords,citation_count +P001,"Graph Neural Networks for Knowledge Graph Completion","This paper presents a novel approach using graph neural networks to complete missing links in knowledge graphs. We demonstrate significant improvements over traditional embedding methods.","Smith, J.; Johnson, A.; Brown, M.",2023-03-15,Nature Machine Intelligence,10.1038/s42256-023-00001-1,"graph neural networks,knowledge graphs,link prediction",127 +P002,"Hybrid RAG Systems: Combining Structured and Semantic Search","We propose a hybrid retrieval-augmented generation system that combines structured database queries with semantic vector search for improved question answering.","Davis, L.; Wilson, K.; Taylor, R.",2023-06-22,ACL Proceedings,10.18653/v1/2023.acl-long.123,"RAG,hybrid systems,question answering,semantic search",89 +P003,"Multi-Tenant Graph Databases: Security and Performance","This work addresses security and performance challenges in multi-tenant graph database systems, proposing novel isolation mechanisms.","Anderson, P.; Garcia, M.; Lee, S.",2023-01-10,VLDB Journal,10.1007/s00778-023-00001-2,"graph databases,multi-tenancy,security,performance",156 +P004,"Large Language Models for Scientific Literature Review","We explore the application of large language models for automated scientific literature review and knowledge synthesis.","Thompson, E.; Martinez, C.; White, D.",2023-08-30,Science,10.1126/science.abc1234,"LLM,literature review,knowledge synthesis",203 +P005,"Vector Embeddings in Knowledge Graph Construction","This paper investigates the role of vector embeddings in automated knowledge graph construction from unstructured text.","Clark, R.; Lewis, J.; Hall, N.",2023-04-18,EMNLP Proceedings,10.18653/v1/2023.emnlp-main.456,"embeddings,knowledge graphs,NLP,text mining",94 \ No newline at end of file diff --git a/examples/document-graph/notebooks/mandela_data/mandela_events.csv b/examples/document-graph/notebooks/mandela_data/mandela_events.csv new file mode 100644 index 00000000..cc49f8a8 --- /dev/null +++ b/examples/document-graph/notebooks/mandela_data/mandela_events.csv @@ -0,0 +1,10 @@ +event_id,date,age_at_event,location,category,description,people_involved,source +E-1918-BIRTH,1918-07-18,0,"Mvezo, Cape Province, South Africa",Birth,Nelson Rolihlahla Mandela was born.,Nelson Mandela,Biography +E-1944-ANC,1944-01-01,25,"Johannesburg, South Africa",Political,Joined African National Congress (ANC) and co-founded the ANC Youth League.,"Nelson Mandela, Walter Sisulu, Oliver Tambo",Historical records +E-1952-DEFIANCE,1952-06-26,33,South Africa,Protest,Led Defiance Campaign against unjust apartheid laws.,"Nelson Mandela, ANC",Historical records +E-1962-ARREST,1962-08-05,44,"Howick, South Africa",Arrest,Arrested and later sentenced to five years for leaving the country illegally and inciting strike.,Nelson Mandela,Court records +E-1964-RIVONIA,1964-06-12,45,"Pretoria, South Africa",Trial,Rivonia Trial verdict; sentenced to life imprisonment.,"Nelson Mandela, Rivonia Trialists",Court records +E-1990-RELEASE,1990-02-11,71,"Cape Town, South Africa",Release,Released from Victor Verster Prison.,Nelson Mandela,News archives +E-1993-NOBEL,1993-12-10,75,"Oslo, Norway",Award,Awarded the Nobel Peace Prize with F. W. de Klerk.,"Nelson Mandela, F. W. de Klerk",Nobel Prize +E-1994-PRESIDENT,1994-05-10,75,"Pretoria, South Africa",Inauguration,Inaugurated as President of South Africa.,Nelson Mandela,Government archives +E-2013-DEATH,2013-12-05,95,"Johannesburg, South Africa",Death,Passed away at age 95.,Nelson Mandela,News archives diff --git a/examples/document-graph/notebooks/mandela_data/organizations.xlsx b/examples/document-graph/notebooks/mandela_data/organizations.xlsx new file mode 100644 index 00000000..550cd4bc Binary files /dev/null and b/examples/document-graph/notebooks/mandela_data/organizations.xlsx differ diff --git a/examples/document-graph/notebooks/mandela_data/research_papers.json b/examples/document-graph/notebooks/mandela_data/research_papers.json new file mode 100644 index 00000000..4b4468e1 --- /dev/null +++ b/examples/document-graph/notebooks/mandela_data/research_papers.json @@ -0,0 +1,89 @@ +[ + { + "id": "P-0001", + "title": "The Evolution of Apartheid Policies, 1948\u20131994", + "authors": [ + "A. Researcher", + "B. Historian" + ], + "year": 2018, + "journal": "Journal of Southern African Studies", + "abstract": "An analysis of the political and social mechanisms that sustained apartheid.", + "keywords": [ + "apartheid", + "policy", + "history" + ], + "citation_count": 124, + "doi": "10.1000/jsas.2018.001" + }, + { + "id": "P-0002", + "title": "Resistance Networks and the ANC", + "authors": [ + "C. Analyst" + ], + "year": 2016, + "journal": "African Affairs", + "abstract": "How underground networks enabled sustained resistance.", + "keywords": [ + "ANC", + "resistance", + "networks" + ], + "citation_count": 86, + "doi": "10.1000/aa.2016.002" + }, + { + "id": "P-0003", + "title": "International Sanctions and South Africa", + "authors": [ + "D. Economist", + "E. Diplomat" + ], + "year": 2014, + "journal": "International Studies Quarterly", + "abstract": "Economic impacts of sanctions on the apartheid regime.", + "keywords": [ + "sanctions", + "economics", + "South Africa" + ], + "citation_count": 210, + "doi": "10.1000/isq.2014.003" + }, + { + "id": "P-0004", + "title": "Truth and Reconciliation: Outcomes and Limits", + "authors": [ + "F. Sociologist" + ], + "year": 2020, + "journal": "Transitional Justice Review", + "abstract": "Assessing the TRC's long-term effects.", + "keywords": [ + "TRC", + "justice", + "reconciliation" + ], + "citation_count": 45, + "doi": "10.1000/tjr.2020.004" + }, + { + "id": "P-0005", + "title": "Media Representation of Mandela", + "authors": [ + "G. MediaScholar" + ], + "year": 2012, + "journal": "Media & Society", + "abstract": "A study of international media portrayals of Nelson Mandela.", + "keywords": [ + "media", + "Mandela", + "representation" + ], + "citation_count": 67, + "doi": "10.1000/ms.2012.005" + } +] \ No newline at end of file diff --git a/examples/document-graph/push-to-sagemaker.sh b/examples/document-graph/push-to-sagemaker.sh new file mode 100755 index 00000000..049af051 --- /dev/null +++ b/examples/document-graph/push-to-sagemaker.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# push-to-sagemaker.sh — Rebuild wheel + upload to S3 (SageMaker pulls from there) +# Run locally after making changes to document-graph. +# Then in SageMaker terminal: bash ~/SageMaker/document-graph/install.sh +set -e + +# Configure these for your environment: +AWS_ACCOUNT_ID="${AWS_ACCOUNT_ID:-}" +ARTIFACTS_BUCKET="${ARTIFACTS_BUCKET:-graphrag-artifacts-${AWS_ACCOUNT_ID}}" +AWS_REGION="${AWS_REGION:-us-east-1}" + +cd "$(dirname "$0")/../.." +export AWS_PROFILE="${AWS_PROFILE:-default}" + +echo "Building wheel..." +rm -f dist/*.whl +python3 -m build --wheel --outdir dist/ 2>&1 | tail -1 + +echo "Uploading wheel..." +aws s3 cp dist/graphrag_toolkit_document_graph-*.whl "s3://${ARTIFACTS_BUCKET}/document-graph-notebooks/wheels/" --region "${AWS_REGION}" + +echo "Cleaning S3 notebooks (removing stale files)..." +aws s3 rm "s3://${ARTIFACTS_BUCKET}/document-graph-notebooks/" --recursive --exclude "wheels/*" --region "${AWS_REGION}" + +echo "Uploading notebooks..." +aws s3 sync examples/document-graph/notebooks/ "s3://${ARTIFACTS_BUCKET}/document-graph-notebooks/" \ + --exclude "__pycache__/*" --exclude "*.pyc" --exclude "wheels/*" --region "${AWS_REGION}" + +echo "" +echo "✅ Pushed. In SageMaker run:" +echo " aws s3 sync s3://${ARTIFACTS_BUCKET}/document-graph-notebooks/ ~/SageMaker/document-graph/" +echo " pip install --force-reinstall ~/SageMaker/document-graph/wheels/graphrag_toolkit_document_graph-*.whl" diff --git a/examples/lexical-graph/cloudformation-templates/update-stack.sh b/examples/lexical-graph/cloudformation-templates/update-stack.sh index 0c7e9961..01dee67e 100755 --- a/examples/lexical-graph/cloudformation-templates/update-stack.sh +++ b/examples/lexical-graph/cloudformation-templates/update-stack.sh @@ -7,10 +7,10 @@ # Configuration variables STACK_NAME="AWS-GraphRAG-01" # Replace with your stack name TEMPLATE_FILE="graphrag-toolkit-neptune-db-aurora-postgres-existing-vpc.json" # Path to your updated template file -S3_BUCKET_ARN="arn:aws:s3:::ccms-rag-extract-188967239867" # Replace with your S3 bucket ARN +S3_BUCKET_ARN="arn:aws:s3:::" # Replace with your S3 bucket ARN REGION="us-east-1" # Replace with your AWS region PARAMETERS_FILE="parameters.json" # Temporary parameters file -AWS_PROFILE="master" # Will be set if --profile is provided +AWS_PROFILE="" # Will be set if --profile is provided # Parse command-line arguments while [[ "$#" -gt 0 ]]; do