From 76c4763b9b8a07d5656821d83d4ea35010eb0a9e Mon Sep 17 00:00:00 2001 From: Sunny Yang Date: Fri, 3 Jul 2026 08:16:39 -0600 Subject: [PATCH 1/4] feat: add tg-export-workspace / tg-import-workspace bundle commands (#877) (#1019) Phase 1 (config only): export a workspace's full configuration as a portable .tgx bundle (gzipped tar with manifest.json and one pretty-printed, self-describing JSON file per config key under config//), and import it into another deployment or workspace. Import defaults to WorkspaceInit's re-run semantics (existing keys kept, missing keys added; --overwrite replaces), supports --workspace rename, --dry-run, and --config-only, and refuses to silently drop knowledge data from future Phase-2 bundles it cannot import yet. --- .../test_workspace_bundle_commands.py | 224 ++++++++++++++++++ trustgraph-cli/pyproject.toml | 2 + .../trustgraph/cli/export_workspace.py | 141 +++++++++++ .../trustgraph/cli/import_workspace.py | 194 +++++++++++++++ 4 files changed, 561 insertions(+) create mode 100644 tests/unit/test_cli/test_workspace_bundle_commands.py create mode 100644 trustgraph-cli/trustgraph/cli/export_workspace.py create mode 100644 trustgraph-cli/trustgraph/cli/import_workspace.py diff --git a/tests/unit/test_cli/test_workspace_bundle_commands.py b/tests/unit/test_cli/test_workspace_bundle_commands.py new file mode 100644 index 00000000..e45c0ef4 --- /dev/null +++ b/tests/unit/test_cli/test_workspace_bundle_commands.py @@ -0,0 +1,224 @@ +""" +Tests for tg-export-workspace / tg-import-workspace (.tgx bundle commands). + +The Api class is mocked in each command module's namespace (same pattern as +test_config_commands.py); bundles are written to and read from tmp_path so +the archive format itself is exercised end-to-end. +""" + +import json +import tarfile +from unittest.mock import Mock, patch + +import pytest + +from trustgraph.api.types import ConfigValue +from trustgraph.cli.export_workspace import export_workspace +from trustgraph.cli.import_workspace import import_workspace + +SAMPLE_CONFIG = { + "prompt": { + "extract-concepts": json.dumps({"template": "Extract {{q}}"}), + "answer": json.dumps({"template": "Answer {{q}}"}), + }, + "tool": { + "web-search": json.dumps({"name": "web-search", "kind": "http"}), + }, +} + + +def make_mock_api(): + mock_config = Mock() + mock_api = Mock() + mock_api.config.return_value = mock_config + return mock_api, mock_config + + +@pytest.fixture +def bundle(tmp_path): + """Export SAMPLE_CONFIG to a real .tgx and yield its path.""" + path = tmp_path / "ws.tgx" + with patch("trustgraph.cli.export_workspace.Api") as api_cls: + mock_api, mock_config = make_mock_api() + api_cls.return_value = mock_api + mock_config.all.return_value = (SAMPLE_CONFIG, "v42") + export_workspace( + url="http://api/", workspace="source-ws", output=str(path), + ) + return path + + +class TestExportWorkspace: + + def test_bundle_contains_manifest_and_per_key_entries(self, bundle): + with tarfile.open(bundle, "r:gz") as tar: + names = tar.getnames() + manifest = json.load(tar.extractfile("manifest.json")) + + assert manifest["format"] == "tgx" + assert manifest["workspace"] == "source-ws" + assert manifest["config_version"] == "v42" + assert manifest["contents"] == {"config": True, "knowledge": False} + + assert "config/prompt/extract-concepts.json" in names + assert "config/prompt/answer.json" in names + assert "config/tool/web-search.json" in names + + def test_entries_are_parsed_and_self_describing(self, bundle): + with tarfile.open(bundle, "r:gz") as tar: + entry = json.load( + tar.extractfile("config/prompt/extract-concepts.json") + ) + # Values are pretty-printed objects, not double-encoded strings, + # and each entry embeds its own type/key (filenames are cosmetic). + assert entry == { + "type": "prompt", + "key": "extract-concepts", + "value": {"template": "Extract {{q}}"}, + } + + def test_path_unsafe_keys_are_quoted_in_filenames(self, tmp_path): + path = tmp_path / "ws.tgx" + with patch("trustgraph.cli.export_workspace.Api") as api_cls: + mock_api, mock_config = make_mock_api() + api_cls.return_value = mock_api + mock_config.all.return_value = ( + {"prompt": {"a/b": json.dumps({"x": 1})}}, "v1", + ) + export_workspace( + url="http://api/", workspace="ws", output=str(path), + ) + with tarfile.open(path, "r:gz") as tar: + names = tar.getnames() + entry = json.load(tar.extractfile("config/prompt/a%2Fb.json")) + assert "config/prompt/a%2Fb.json" in names + assert entry["key"] == "a/b" + + +class TestImportWorkspace: + + def test_roundtrip_puts_all_values_with_overwrite(self, bundle): + with patch("trustgraph.cli.import_workspace.Api") as api_cls: + mock_api, mock_config = make_mock_api() + api_cls.return_value = mock_api + import_workspace( + url="http://api/", input=str(bundle), overwrite=True, + ) + + # Target workspace defaults to the manifest's workspace. + api_cls.assert_called_once_with( + "http://api/", token=None, workspace="source-ws", + ) + values = mock_config.put.call_args.args[0] + assert sorted((v.type, v.key) for v in values) == [ + ("prompt", "answer"), + ("prompt", "extract-concepts"), + ("tool", "web-search"), + ] + # Values are re-serialized to JSON strings, as config-svc stores. + by_key = {(v.type, v.key): v for v in values} + assert json.loads(by_key[("prompt", "answer")].value) == { + "template": "Answer {{q}}", + } + assert all(isinstance(v, ConfigValue) for v in values) + + def test_workspace_flag_renames_target(self, bundle): + with patch("trustgraph.cli.import_workspace.Api") as api_cls: + mock_api, mock_config = make_mock_api() + api_cls.return_value = mock_api + import_workspace( + url="http://api/", input=str(bundle), workspace="staging", + overwrite=True, + ) + api_cls.assert_called_once_with( + "http://api/", token=None, workspace="staging", + ) + + def test_default_skips_existing_keys(self, bundle): + """WorkspaceInit re-run semantics: only missing keys are written.""" + with patch("trustgraph.cli.import_workspace.Api") as api_cls: + mock_api, mock_config = make_mock_api() + api_cls.return_value = mock_api + mock_config.list.side_effect = lambda t: { + "prompt": ["extract-concepts"], + "tool": [], + }[t] + import_workspace(url="http://api/", input=str(bundle)) + + values = mock_config.put.call_args.args[0] + assert sorted((v.type, v.key) for v in values) == [ + ("prompt", "answer"), + ("tool", "web-search"), + ] + + def test_dry_run_writes_nothing(self, bundle, capsys): + with patch("trustgraph.cli.import_workspace.Api") as api_cls: + mock_api, mock_config = make_mock_api() + api_cls.return_value = mock_api + import_workspace( + url="http://api/", input=str(bundle), overwrite=True, + dry_run=True, + ) + mock_config.put.assert_not_called() + out = capsys.readouterr().out + assert "would import prompt/extract-concepts" in out + assert "3 item(s) would be imported" in out + + def test_rejects_bundle_without_manifest(self, tmp_path): + path = tmp_path / "bad.tgx" + with tarfile.open(path, "w:gz"): + pass + with patch("trustgraph.cli.import_workspace.Api"): + with pytest.raises(RuntimeError, match="manifest.json missing"): + import_workspace(url="http://api/", input=str(path)) + + def test_rejects_newer_format_version(self, tmp_path): + import io + path = tmp_path / "future.tgx" + manifest = json.dumps({ + "format": "tgx", "format_version": 99, "workspace": "w", + "contents": {"config": True, "knowledge": False}, + }).encode() + with tarfile.open(path, "w:gz") as tar: + info = tarfile.TarInfo("manifest.json") + info.size = len(manifest) + tar.addfile(info, io.BytesIO(manifest)) + with patch("trustgraph.cli.import_workspace.Api"): + with pytest.raises(RuntimeError, match="newer than this tool"): + import_workspace(url="http://api/", input=str(path)) + + def test_refuses_knowledge_bundle_without_config_only(self, tmp_path): + import io + path = tmp_path / "knowledge.tgx" + manifest = json.dumps({ + "format": "tgx", "format_version": 1, "workspace": "w", + "contents": {"config": True, "knowledge": True}, + }).encode() + with tarfile.open(path, "w:gz") as tar: + info = tarfile.TarInfo("manifest.json") + info.size = len(manifest) + tar.addfile(info, io.BytesIO(manifest)) + with patch("trustgraph.cli.import_workspace.Api"): + with pytest.raises(RuntimeError, match="--config-only"): + import_workspace(url="http://api/", input=str(path)) + + def test_config_only_flag_allows_knowledge_bundle(self, tmp_path): + import io + path = tmp_path / "knowledge.tgx" + manifest = json.dumps({ + "format": "tgx", "format_version": 1, "workspace": "w", + "contents": {"config": True, "knowledge": True}, + }).encode() + with tarfile.open(path, "w:gz") as tar: + info = tarfile.TarInfo("manifest.json") + info.size = len(manifest) + tar.addfile(info, io.BytesIO(manifest)) + with patch("trustgraph.cli.import_workspace.Api") as api_cls: + mock_api, mock_config = make_mock_api() + api_cls.return_value = mock_api + import_workspace( + url="http://api/", input=str(path), config_only=True, + overwrite=True, + ) + # No config entries in this bundle; nothing written, no error. + mock_config.put.assert_not_called() diff --git a/trustgraph-cli/pyproject.toml b/trustgraph-cli/pyproject.toml index 193ee1cd..a89be8ca 100644 --- a/trustgraph-cli/pyproject.toml +++ b/trustgraph-cli/pyproject.toml @@ -116,6 +116,8 @@ tg-put-config-item = "trustgraph.cli.put_config_item:main" tg-delete-config-item = "trustgraph.cli.delete_config_item:main" tg-export-workspace-config = "trustgraph.cli.export_workspace_config:main" tg-import-workspace-config = "trustgraph.cli.import_workspace_config:main" +tg-export-workspace = "trustgraph.cli.export_workspace:main" +tg-import-workspace = "trustgraph.cli.import_workspace:main" tg-list-collections = "trustgraph.cli.list_collections:main" tg-set-collection = "trustgraph.cli.set_collection:main" tg-delete-collection = "trustgraph.cli.delete_collection:main" diff --git a/trustgraph-cli/trustgraph/cli/export_workspace.py b/trustgraph-cli/trustgraph/cli/export_workspace.py new file mode 100644 index 00000000..41d3403d --- /dev/null +++ b/trustgraph-cli/trustgraph/cli/export_workspace.py @@ -0,0 +1,141 @@ +""" +Exports a workspace's full configuration state as a portable .tgx bundle +(a gzipped tar archive) for backup, migration between deployments, or +sharing a pre-configured workspace. + +The bundle is human-readable: a manifest.json plus one pretty-printed JSON +file per config key under config//, so it can be inspected and +hand-edited before import. Each entry file embeds its own type and key, so +filenames are cosmetic. Knowledge export (triples, documents, embeddings) +is not yet included; the manifest records that so future importers can +distinguish config-only bundles. +""" + +import argparse +import io +import json +import os +import sys +import tarfile +import time +from urllib.parse import quote + +from trustgraph.api import Api + +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) +default_workspace = os.getenv("TRUSTGRAPH_WORKSPACE", "default") + +MANIFEST_FORMAT = "tgx" +MANIFEST_FORMAT_VERSION = 1 + + +def _add_bytes(tar, name, data): + info = tarfile.TarInfo(name=name) + info.size = len(data) + info.mtime = int(time.time()) + tar.addfile(info, io.BytesIO(data)) + + +def export_workspace(url, workspace, output, token=None): + + api = Api(url, token=token, workspace=workspace).config() + + config, version = api.all() + + manifest = { + "format": MANIFEST_FORMAT, + "format_version": MANIFEST_FORMAT_VERSION, + "workspace": workspace, + "config_version": version, + "exported_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "contents": {"config": True, "knowledge": False}, + } + + count = 0 + + with tarfile.open(output, "w:gz") as tar: + + _add_bytes( + tar, "manifest.json", + json.dumps(manifest, indent=2).encode("utf-8"), + ) + + for type_, entries in sorted(config.items()): + for key, raw in sorted(entries.items()): + + # Config values are stored as JSON strings; parse so the + # bundle is pretty-printed and hand-editable. A value that + # isn't valid JSON is preserved verbatim. + try: + value = json.loads(raw) + except (TypeError, json.JSONDecodeError): + value = raw + + entry = {"type": type_, "key": key, "value": value} + + # Keys may contain path-unsafe characters; the entry embeds + # the real key, so the quoted filename is cosmetic only. + name = f"config/{quote(type_, safe='')}/{quote(key, safe='')}.json" + + _add_bytes( + tar, name, + json.dumps(entry, indent=2).encode("utf-8"), + ) + + count += 1 + + print(f"Exported {count} config item(s) from workspace " + f"'{workspace}' to {output}", flush=True) + + +def main(): + + parser = argparse.ArgumentParser( + prog='tg-export-workspace', + description=__doc__, + ) + + parser.add_argument( + '-u', '--api-url', + default=default_url, + help=f'API URL (default: {default_url})', + ) + + parser.add_argument( + '-t', '--token', + default=default_token, + help='API token (default: TRUSTGRAPH_TOKEN environment variable)', + ) + + parser.add_argument( + '-w', '--workspace', + default=default_workspace, + help=f'Workspace to export (default: {default_workspace})', + ) + + parser.add_argument( + '-o', '--output', + required=True, + help='Output bundle path, e.g. workspace-default.tgx', + ) + + args = parser.parse_args() + + try: + + export_workspace( + url=args.api_url, + workspace=args.workspace, + output=args.output, + token=args.token, + ) + + except Exception as e: + + print("Exception:", e, flush=True) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/trustgraph-cli/trustgraph/cli/import_workspace.py b/trustgraph-cli/trustgraph/cli/import_workspace.py new file mode 100644 index 00000000..71885bfb --- /dev/null +++ b/trustgraph-cli/trustgraph/cli/import_workspace.py @@ -0,0 +1,194 @@ +""" +Imports a workspace bundle (.tgx, produced by tg-export-workspace) into a +TrustGraph deployment. The target workspace defaults to the name recorded +in the bundle's manifest and can be renamed with --workspace. + +By default existing (type, key) entries in the target workspace are left +untouched and only missing keys are added, matching WorkspaceInit's +re-run behaviour; pass --overwrite to replace every imported key. Use +--dry-run to show what would be written without changing anything. +""" + +import argparse +import json +import os +import sys +import tarfile + +from trustgraph.api import Api +from trustgraph.api.types import ConfigValue + +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) + +SUPPORTED_FORMAT = "tgx" +SUPPORTED_FORMAT_VERSION = 1 + + +def _read_bundle(path): + """Read manifest and config entries from a .tgx bundle.""" + + manifest = None + entries = [] + + with tarfile.open(path, "r:gz") as tar: + for member in tar.getmembers(): + if not member.isfile(): + continue + f = tar.extractfile(member) + if f is None: + continue + data = f.read() + if member.name == "manifest.json": + manifest = json.loads(data) + elif member.name.startswith("config/") and \ + member.name.endswith(".json"): + entries.append(json.loads(data)) + + if manifest is None: + raise RuntimeError("not a workspace bundle: manifest.json missing") + + if manifest.get("format") != SUPPORTED_FORMAT: + raise RuntimeError( + f"unsupported bundle format: {manifest.get('format')!r}" + ) + + if manifest.get("format_version", 0) > SUPPORTED_FORMAT_VERSION: + raise RuntimeError( + f"bundle format version {manifest.get('format_version')} is " + f"newer than this tool supports ({SUPPORTED_FORMAT_VERSION}); " + "upgrade trustgraph-cli" + ) + + return manifest, entries + + +def import_workspace( + url, input, workspace=None, overwrite=False, config_only=False, + dry_run=False, token=None, +): + + manifest, entries = _read_bundle(input) + + # Knowledge import (triples, documents, embeddings) is not implemented + # yet; refuse to silently drop it from a bundle that carries it. + if manifest.get("contents", {}).get("knowledge") and not config_only: + raise RuntimeError( + "bundle contains knowledge data, which this tool cannot import " + "yet; re-run with --config-only to import just the configuration" + ) + + target = workspace or manifest.get("workspace") or "default" + + api = Api(url, token=token, workspace=target).config() + + # Mirror WorkspaceInit's re-run behaviour: without --overwrite, keys + # already present in the target workspace are skipped (per key, not per + # type). The config API's put is a blanket upsert, so filter client-side. + existing = {} + if not overwrite: + for type_ in sorted({e["type"] for e in entries}): + existing[type_] = set(api.list(type_)) + + values = [] + skipped = 0 + + for e in entries: + type_, key, value = e["type"], e["key"], e["value"] + if not overwrite and key in existing.get(type_, set()): + skipped += 1 + continue + # Config values are stored as JSON strings (see WorkspaceInit). + values.append( + ConfigValue(type=type_, key=key, value=json.dumps(value)) + ) + + if dry_run: + for v in values: + print(f"would import {v.type}/{v.key}", flush=True) + print(f"Dry run: {len(values)} item(s) would be imported into " + f"workspace '{target}', {skipped} skipped as existing", + flush=True) + return + + if values: + api.put(values) + + print(f"Imported {len(values)} config item(s) into workspace " + f"'{target}', {skipped} skipped as existing", flush=True) + + +def main(): + + parser = argparse.ArgumentParser( + prog='tg-import-workspace', + description=__doc__, + ) + + parser.add_argument( + '-u', '--api-url', + default=default_url, + help=f'API URL (default: {default_url})', + ) + + parser.add_argument( + '-t', '--token', + default=default_token, + help='API token (default: TRUSTGRAPH_TOKEN environment variable)', + ) + + parser.add_argument( + '-i', '--input', + required=True, + help='Input bundle path, e.g. workspace-default.tgx', + ) + + parser.add_argument( + '-w', '--workspace', + default=None, + help='Target workspace (default: the workspace recorded in the ' + 'bundle manifest)', + ) + + parser.add_argument( + '--overwrite', + action='store_true', + help='Replace existing keys in the target workspace (default: ' + 'keep existing keys and only add missing ones)', + ) + + parser.add_argument( + '--config-only', + action='store_true', + help='Import only the configuration, skipping any knowledge data ' + 'in the bundle', + ) + + parser.add_argument( + '--dry-run', + action='store_true', + help='Show what would be imported without writing anything', + ) + + args = parser.parse_args() + + try: + + import_workspace( + url=args.api_url, + input=args.input, + workspace=args.workspace, + overwrite=args.overwrite, + config_only=args.config_only, + dry_run=args.dry_run, + token=args.token, + ) + + except Exception as e: + + print("Exception:", e, flush=True) + sys.exit(1) + + +if __name__ == "__main__": + main() From 68e816e65ce27541b38d2aa0300b54d8106c5c4c Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Fri, 3 Jul 2026 15:51:04 +0100 Subject: [PATCH 2/4] feat: filter and cap GraphRAG reranker input across full stack (#1021) - Filter out RDF/RDFS/OWL schema predicates (rdfs:domain, owl:inverseOf, etc.) from hop traversal, keeping rdf:type for data signal - Skip edges where reranker-visible components are unlabeled IRIs, since the cross-encoder cannot meaningfully score raw URIs - Add max-reranker-input safety cap (default 350) to prevent overloading the reranker, applied after filtering for maximum useful candidates - Expose max-reranker-input as per-request parameter through schema, translator, REST API, socket client, CLI, and OpenAPI spec - Update tests - Update tech spec --- docs/tech-specs/graph-rag-semantic-filter.md | 31 ++++++- .../schemas/rag/GraphRagRequest.yaml | 7 ++ .../test_graph_rag_direction_aware_text.py | 92 ++++++++++++++----- trustgraph-base/trustgraph/api/flow.py | 3 + .../trustgraph/api/socket_client.py | 4 + .../messaging/translators/retrieval.py | 2 + .../trustgraph/schema/services/retrieval.py | 1 + .../trustgraph/cli/invoke_graph_rag.py | 18 +++- .../retrieval/graph_rag/graph_rag.py | 64 +++++++++++-- .../trustgraph/retrieval/graph_rag/rag.py | 19 +++- 10 files changed, 198 insertions(+), 43 deletions(-) diff --git a/docs/tech-specs/graph-rag-semantic-filter.md b/docs/tech-specs/graph-rag-semantic-filter.md index 58497d10..cb2fd24a 100644 --- a/docs/tech-specs/graph-rag-semantic-filter.md +++ b/docs/tech-specs/graph-rag-semantic-filter.md @@ -404,10 +404,33 @@ no LLM call. These fields are dropped from the Focus entity. a. Retrieve all edges one hop from the current frontier nodes. - b. Represent each edge using direction-aware text: from a - subject node use `"{predicate} {object}"`, from an object - node use `"{subject} {predicate}"`, from a predicate node - use `"{subject} {object}"`. + b. Filter and represent edges for scoring: + + - **Schema predicate filter.** Edges with RDF/RDFS/OWL + schema predicates (`rdfs:domain`, `owl:inverseOf`, etc.) + are removed. `rdf:type` is kept as it carries useful + data signal. + + - **IRI filter.** Edges where the reranker-visible text + components (after label resolution) are still raw IRIs + are removed — the cross-encoder cannot meaningfully score + unresolved URIs. Only the components that would appear + in the reranker text are checked, based on traversal + direction. + + - **Direction-aware text.** Each surviving edge is + represented using direction-aware text: from a subject + node use `"{predicate} {object}"`, from an object node + use `"{subject} {predicate}"`, from a predicate node + use `"{subject} {object}"`. + + - **Reranker input cap.** The candidate set is truncated + to `max_reranker_input` (default 350) edges. This is a + safety measure, not an accuracy optimisation — there is + no point in producing a perfectly ranked edge set if the + reranker crashes or times out because it was handed + thousands of candidates. The cap is applied after + filtering so that the most useful edges fill the budget. c. Score edges against the extracted concepts using the cross-encoder service. diff --git a/specs/api/components/schemas/rag/GraphRagRequest.yaml b/specs/api/components/schemas/rag/GraphRagRequest.yaml index 754dcc92..f1899fc4 100644 --- a/specs/api/components/schemas/rag/GraphRagRequest.yaml +++ b/specs/api/components/schemas/rag/GraphRagRequest.yaml @@ -42,6 +42,13 @@ properties: minimum: 1 maximum: 5 example: 3 + max-reranker-input: + type: integer + description: Maximum candidate edges sent to the reranker per hop + default: 350 + minimum: 1 + maximum: 1000 + example: 350 streaming: type: boolean description: Enable streaming response delivery diff --git a/tests/unit/test_retrieval/test_graph_rag_direction_aware_text.py b/tests/unit/test_retrieval/test_graph_rag_direction_aware_text.py index cc95228a..c58ac3b6 100644 --- a/tests/unit/test_retrieval/test_graph_rag_direction_aware_text.py +++ b/tests/unit/test_retrieval/test_graph_rag_direction_aware_text.py @@ -18,15 +18,30 @@ from trustgraph.schema import Term, IRI, LITERAL # Helpers # --------------------------------------------------------------------------- -def _make_rag(reranker_results=None): - """Create a mock GraphRag with all clients stubbed.""" +LABEL = "http://www.w3.org/2000/01/rdf-schema#label" + + +def _make_rag(reranker_results=None, labels=None): + """Create a mock GraphRag with all clients stubbed. + + labels is an optional dict mapping URI -> label string. When provided, + the mock triples_client.query will return matching label triples so + that hop_and_filter resolves labels instead of falling back to raw URIs + (which are now filtered out by the IRI filter). + """ rag = MagicMock() rag.label_cache = LRUCacheWithTTL() rag.triples_client = AsyncMock() rag.reranker_client = AsyncMock() - # Label lookups return empty (fall back to URI) - rag.triples_client.query.return_value = [] + if labels: + async def label_query(s=None, p=None, o=None, limit=1, **kwargs): + if p == LABEL and s in labels: + return [MagicMock(o=labels[s])] + return [] + rag.triples_client.query.side_effect = label_query + else: + rag.triples_client.query.return_value = [] if reranker_results is not None: rag.reranker_client.rerank.return_value = reranker_results @@ -147,8 +162,13 @@ class TestDirectionAwareRerankerText: "http://ex/likes", "http://ex/entity-B", ) + labels = { + "http://ex/entity-A": "Alice", + "http://ex/likes": "likes", + "http://ex/entity-B": "Bob", + } reranker_result = _reranker_result(0) - rag = _make_rag(reranker_results=[reranker_result]) + rag = _make_rag(reranker_results=[reranker_result], labels=labels) async def query_stream(s=None, p=None, o=None, **kwargs): if s is not None: @@ -166,9 +186,8 @@ class TestDirectionAwareRerankerText: call_args = rag.reranker_client.rerank.call_args documents = call_args.kwargs["documents"] - # Text should be "{p} {o}" — the URIs since no labels found assert len(documents) == 1 - assert documents[0]["text"] == "http://ex/likes http://ex/entity-B" + assert documents[0]["text"] == "likes Bob" @pytest.mark.asyncio async def test_from_o_uses_subject_predicate(self): @@ -178,8 +197,13 @@ class TestDirectionAwareRerankerText: "http://ex/likes", "http://ex/entity-B", ) + labels = { + "http://ex/entity-A": "Alice", + "http://ex/likes": "likes", + "http://ex/entity-B": "Bob", + } reranker_result = _reranker_result(0) - rag = _make_rag(reranker_results=[reranker_result]) + rag = _make_rag(reranker_results=[reranker_result], labels=labels) async def query_stream(s=None, p=None, o=None, **kwargs): if o is not None: @@ -198,7 +222,7 @@ class TestDirectionAwareRerankerText: call_args = rag.reranker_client.rerank.call_args documents = call_args.kwargs["documents"] assert len(documents) == 1 - assert documents[0]["text"] == "http://ex/entity-A http://ex/likes" + assert documents[0]["text"] == "Alice likes" @pytest.mark.asyncio async def test_from_p_uses_subject_object(self): @@ -208,8 +232,13 @@ class TestDirectionAwareRerankerText: "http://ex/likes", "http://ex/entity-B", ) + labels = { + "http://ex/entity-A": "Alice", + "http://ex/likes": "likes", + "http://ex/entity-B": "Bob", + } reranker_result = _reranker_result(0) - rag = _make_rag(reranker_results=[reranker_result]) + rag = _make_rag(reranker_results=[reranker_result], labels=labels) async def query_stream(s=None, p=None, o=None, **kwargs): if p is not None: @@ -228,7 +257,7 @@ class TestDirectionAwareRerankerText: call_args = rag.reranker_client.rerank.call_args documents = call_args.kwargs["documents"] assert len(documents) == 1 - assert documents[0]["text"] == "http://ex/entity-A http://ex/entity-B" + assert documents[0]["text"] == "Alice Bob" @pytest.mark.asyncio async def test_mixed_directions_produce_different_text(self): @@ -239,10 +268,18 @@ class TestDirectionAwareRerankerText: triple_from_o = _make_schema_triple( "http://ex/other", "http://ex/ref", "http://ex/seed", ) + labels = { + "http://ex/seed": "Seed", + "http://ex/rel": "relates to", + "http://ex/target": "Target", + "http://ex/other": "Other", + "http://ex/ref": "references", + } - rag = _make_rag(reranker_results=[ - _reranker_result(0), _reranker_result(1), - ]) + rag = _make_rag( + reranker_results=[_reranker_result(0), _reranker_result(1)], + labels=labels, + ) async def query_stream(s=None, p=None, o=None, **kwargs): if s == "http://ex/seed": @@ -264,10 +301,10 @@ class TestDirectionAwareRerankerText: documents = call_args.kwargs["documents"] texts = {d["text"] for d in documents} - # From S: "{p} {o}" = "http://ex/rel http://ex/target" - assert "http://ex/rel http://ex/target" in texts - # From O: "{s} {p}" = "http://ex/other http://ex/ref" - assert "http://ex/other http://ex/ref" in texts + # From S: "{p} {o}" = "relates to Target" + assert "relates to Target" in texts + # From O: "{s} {p}" = "Other references" + assert "Other references" in texts @pytest.mark.asyncio async def test_labels_applied_to_direction_text(self): @@ -280,8 +317,6 @@ class TestDirectionAwareRerankerText: reranker_result = _reranker_result(0) rag = _make_rag(reranker_results=[reranker_result]) - LABEL = "http://www.w3.org/2000/01/rdf-schema#label" - async def query_stream(s=None, p=None, o=None, **kwargs): if s is not None and p is None: return [triple] @@ -323,10 +358,17 @@ class TestDirectionAwareRerankerText: triple_b = _make_schema_triple( "http://ex/cpu-B", "http://ex/hasCategory", "http://ex/Processors", ) + labels = { + "http://ex/cpu-A": "CPU Alpha", + "http://ex/cpu-B": "CPU Beta", + "http://ex/hasCategory": "has category", + "http://ex/Processors": "Processors", + } - rag = _make_rag(reranker_results=[ - _reranker_result(0), _reranker_result(1), - ]) + rag = _make_rag( + reranker_results=[_reranker_result(0), _reranker_result(1)], + labels=labels, + ) async def query_stream(s=None, p=None, o=None, **kwargs): if o == "http://ex/Processors": @@ -349,5 +391,5 @@ class TestDirectionAwareRerankerText: assert len(texts) == 2 # From O: "{s} {p}" — subjects differ, so texts differ assert texts[0] != texts[1] - assert "http://ex/cpu-A" in texts[0] - assert "http://ex/cpu-B" in texts[1] + assert "CPU Alpha" in texts[0] + assert "CPU Beta" in texts[1] diff --git a/trustgraph-base/trustgraph/api/flow.py b/trustgraph-base/trustgraph/api/flow.py index b9e9487b..95fc009b 100644 --- a/trustgraph-base/trustgraph/api/flow.py +++ b/trustgraph-base/trustgraph/api/flow.py @@ -357,6 +357,7 @@ class FlowInstance: self, query,collection="default", entity_limit=50, triple_limit=30, max_subgraph_size=150, max_path_length=2, edge_score_limit=30, edge_limit=25, + max_reranker_input=350, ): """ Execute graph-based Retrieval-Augmented Generation (RAG) query. @@ -373,6 +374,7 @@ class FlowInstance: max_path_length: Maximum traversal depth (default: 2) edge_score_limit: Max edges for semantic pre-filter (default: 50) edge_limit: Max edges after LLM scoring (default: 25) + max_reranker_input: Max candidate edges sent to reranker per hop (default: 350) Returns: str: Generated response incorporating graph context @@ -399,6 +401,7 @@ class FlowInstance: "max-path-length": max_path_length, "edge-score-limit": edge_score_limit, "edge-limit": edge_limit, + "max-reranker-input": max_reranker_input, } result = self.request( diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index efa887a1..6c51210f 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -682,6 +682,7 @@ class SocketFlowInstance: max_path_length: int = 2, edge_score_limit: int = 30, edge_limit: int = 25, + max_reranker_input: int = 350, streaming: bool = False, **kwargs: Any ) -> Union[TextCompletionResult, Iterator[RAGChunk]]: @@ -699,6 +700,7 @@ class SocketFlowInstance: "max-path-length": max_path_length, "edge-score-limit": edge_score_limit, "edge-limit": edge_limit, + "max-reranker-input": max_reranker_input, "streaming": streaming } request.update(kwargs) @@ -725,6 +727,7 @@ class SocketFlowInstance: max_path_length: int = 2, edge_score_limit: int = 30, edge_limit: int = 25, + max_reranker_input: int = 350, **kwargs: Any ) -> Iterator[Union[RAGChunk, ProvenanceEvent]]: """Execute graph-based RAG query with explainability support.""" @@ -737,6 +740,7 @@ class SocketFlowInstance: "max-path-length": max_path_length, "edge-score-limit": edge_score_limit, "edge-limit": edge_limit, + "max-reranker-input": max_reranker_input, "streaming": True, "explainable": True, } diff --git a/trustgraph-base/trustgraph/messaging/translators/retrieval.py b/trustgraph-base/trustgraph/messaging/translators/retrieval.py index f2a0b29a..556ad758 100644 --- a/trustgraph-base/trustgraph/messaging/translators/retrieval.py +++ b/trustgraph-base/trustgraph/messaging/translators/retrieval.py @@ -103,6 +103,7 @@ class GraphRagRequestTranslator(MessageTranslator): max_path_length=int(data.get("max-path-length", 2)), edge_score_limit=int(data.get("edge-score-limit", 30)), edge_limit=int(data.get("edge-limit", 25)), + max_reranker_input=int(data.get("max-reranker-input", 350)), streaming=data.get("streaming", False) ) @@ -116,6 +117,7 @@ class GraphRagRequestTranslator(MessageTranslator): "max-path-length": obj.max_path_length, "edge-score-limit": obj.edge_score_limit, "edge-limit": obj.edge_limit, + "max-reranker-input": obj.max_reranker_input, "streaming": getattr(obj, "streaming", False) } diff --git a/trustgraph-base/trustgraph/schema/services/retrieval.py b/trustgraph-base/trustgraph/schema/services/retrieval.py index 2d4e01e1..47ced73d 100644 --- a/trustgraph-base/trustgraph/schema/services/retrieval.py +++ b/trustgraph-base/trustgraph/schema/services/retrieval.py @@ -15,6 +15,7 @@ class GraphRagQuery: max_path_length: int = 0 edge_score_limit: int = 0 edge_limit: int = 0 + max_reranker_input: int = 0 streaming: bool = False parent_uri: str = "" diff --git a/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py b/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py index 892d2d35..97bb8db9 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py +++ b/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py @@ -27,11 +27,13 @@ default_max_subgraph_size = 150 default_max_path_length = 2 default_edge_score_limit = 30 default_edge_limit = 25 +default_max_reranker_input = 350 def _question_explainable_api( url, flow_id, question_text, collection, entity_limit, triple_limit, max_subgraph_size, max_path_length, edge_score_limit=30, - edge_limit=25, token=None, debug=False, workspace="default", + edge_limit=25, max_reranker_input=350, token=None, debug=False, + workspace="default", ): """Execute graph RAG with explainability using the new API classes.""" api = Api(url=url, token=token, workspace=workspace) @@ -50,6 +52,7 @@ def _question_explainable_api( max_path_length=max_path_length, edge_score_limit=edge_score_limit, edge_limit=edge_limit, + max_reranker_input=max_reranker_input, ): if isinstance(item, RAGChunk): # Print response content @@ -138,7 +141,7 @@ def _question_explainable_api( def question( url, flow_id, question, collection, entity_limit, triple_limit, max_subgraph_size, max_path_length, edge_score_limit=50, - edge_limit=25, streaming=True, token=None, + edge_limit=25, max_reranker_input=350, streaming=True, token=None, explainable=False, debug=False, show_usage=False, workspace="default", ): @@ -156,6 +159,7 @@ def question( max_path_length=max_path_length, edge_score_limit=edge_score_limit, edge_limit=edge_limit, + max_reranker_input=max_reranker_input, token=token, debug=debug, workspace=workspace, @@ -180,6 +184,7 @@ def question( max_path_length=max_path_length, edge_score_limit=edge_score_limit, edge_limit=edge_limit, + max_reranker_input=max_reranker_input, streaming=True ) @@ -212,6 +217,7 @@ def question( max_path_length=max_path_length, edge_score_limit=edge_score_limit, edge_limit=edge_limit, + max_reranker_input=max_reranker_input, ) print(result.text) @@ -308,6 +314,13 @@ def main(): help=f'Max edges after LLM scoring (default: {default_edge_limit})' ) + parser.add_argument( + '--max-reranker-input', + type=int, + default=default_max_reranker_input, + help=f'Max candidate edges sent to reranker per hop (default: {default_max_reranker_input})' + ) + parser.add_argument( '--no-streaming', action='store_true', @@ -347,6 +360,7 @@ def main(): max_path_length=args.max_path_length, edge_score_limit=args.edge_score_limit, edge_limit=args.edge_limit, + max_reranker_input=args.max_reranker_input, streaming=not args.no_streaming, token=args.token, explainable=args.explainable, diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py index 2054cb0f..c094e395 100644 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py @@ -34,6 +34,22 @@ logger = logging.getLogger(__name__) LABEL="http://www.w3.org/2000/01/rdf-schema#label" +RDF_NS = "http://www.w3.org/1999/02/22-rdf-syntax-ns#" +RDFS_NS = "http://www.w3.org/2000/01/rdf-schema#" +OWL_NS = "http://www.w3.org/2002/07/owl#" +RDF_TYPE = RDF_NS + "type" +SCHEMA_NAMESPACES = (RDF_NS, RDFS_NS, OWL_NS) + + +def is_schema_predicate(predicate): + """Return True if the predicate is an RDF/RDFS/OWL schema predicate. + + rdf:type is excluded from filtering as it carries useful data signal. + """ + if predicate == RDF_TYPE: + return False + return predicate.startswith(SCHEMA_NAMESPACES) + def term_to_string(term): """Extract string value from a Term object.""" @@ -120,7 +136,8 @@ class Query: def __init__( self, rag, collection, verbose, entity_limit=50, triple_limit=30, max_subgraph_size=1000, - max_path_length=2, edge_limit=25, track_usage=None, + max_path_length=2, edge_limit=25, max_reranker_input=350, + track_usage=None, ): self.rag = rag self.collection = collection @@ -130,6 +147,7 @@ class Query: self.max_subgraph_size = max_subgraph_size self.max_path_length = max_path_length self.edge_limit = edge_limit + self.max_reranker_input = max_reranker_input self.track_usage = track_usage async def extract_concepts(self, query): @@ -346,7 +364,7 @@ class Query: hop_directions = {} for triple, direction in triples: triple_tuple = (str(triple.s), str(triple.p), str(triple.o)) - if triple_tuple[1] == LABEL: + if is_schema_predicate(triple_tuple[1]): continue if triple_tuple in seen_edges: continue @@ -385,25 +403,50 @@ class Query: # The reranker text highlights the NEW information relative # to the traversal direction: arriving from S means p,o are # new; from O means s,p are new; from P means s,o are new. + # Edges where the reranker-visible components are unlabeled + # IRIs are skipped — the cross-encoder can't score them. + def is_iri(val): + return val.startswith(("http://", "https://", "urn:")) + + filtered_triples = [] labeled_hop = [] + documents = [] for s, p, o in hop_triples: ls = label_map.get(s, s) lp = label_map.get(p, p) lo = label_map.get(o, o) - labeled_hop.append((ls, lp, lo)) - documents = [] - for i, (triple_tuple, (ls, lp, lo)) in enumerate( - zip(hop_triples, labeled_hop) - ): - direction = hop_directions[triple_tuple] + direction = hop_directions[(s, p, o)] if direction == self.FROM_S: + if is_iri(lp) or is_iri(lo): + continue text = f"{lp} {lo}" elif direction == self.FROM_O: + if is_iri(ls) or is_iri(lp): + continue text = f"{ls} {lp}" else: + if is_iri(ls) or is_iri(lo): + continue text = f"{ls} {lo}" - documents.append({"id": str(i), "text": text}) + + idx = len(filtered_triples) + filtered_triples.append((s, p, o)) + labeled_hop.append((ls, lp, lo)) + documents.append({"id": str(idx), "text": text}) + + hop_triples = filtered_triples + + # Cap the number of candidates sent to the reranker + if len(hop_triples) > self.max_reranker_input: + if self.verbose: + logger.debug( + f"Hop {hop + 1}: truncating {len(hop_triples)} " + f"candidates to {self.max_reranker_input}" + ) + hop_triples = hop_triples[:self.max_reranker_input] + labeled_hop = labeled_hop[:self.max_reranker_input] + documents = documents[:self.max_reranker_input] queries = [ {"id": str(i), "text": c} @@ -588,7 +631,7 @@ class GraphRag: async def query( self, query, collection = "default", entity_limit = 50, triple_limit = 30, max_subgraph_size = 1000, - max_path_length = 2, edge_limit = 25, + max_path_length = 2, edge_limit = 25, max_reranker_input = 350, streaming = False, chunk_callback = None, explain_callback = None, save_answer_callback = None, @@ -642,6 +685,7 @@ class GraphRag: max_subgraph_size = max_subgraph_size, max_path_length = max_path_length, edge_limit = edge_limit, + max_reranker_input = max_reranker_input, track_usage = track_usage, ) diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index 27ec4937..9ed802f4 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py @@ -34,6 +34,7 @@ class Processor(FlowProcessor): max_subgraph_size = params.get("max_subgraph_size", 150) max_path_length = params.get("max_path_length", 2) edge_limit = params.get("edge_limit", 25) + max_reranker_input = params.get("max_reranker_input", 350) super(Processor, self).__init__( **params | { @@ -44,6 +45,7 @@ class Processor(FlowProcessor): "max_subgraph_size": max_subgraph_size, "max_path_length": max_path_length, "edge_limit": edge_limit, + "max_reranker_input": max_reranker_input, } ) @@ -52,6 +54,7 @@ class Processor(FlowProcessor): self.default_max_subgraph_size = max_subgraph_size self.default_max_path_length = max_path_length self.default_edge_limit = edge_limit + self.default_max_reranker_input = max_reranker_input # Workspace isolation is enforced by the flow layer (flow.workspace). # Per-request caching (see GraphRag) keeps within-request state @@ -197,6 +200,11 @@ class Processor(FlowProcessor): else: edge_limit = self.default_edge_limit + if v.max_reranker_input: + max_reranker_input = v.max_reranker_input + else: + max_reranker_input = self.default_max_reranker_input + async def save_answer(doc_id, answer_text): await flow.librarian.save_document( doc_id=doc_id, @@ -226,8 +234,8 @@ class Processor(FlowProcessor): entity_limit = entity_limit, triple_limit = triple_limit, max_subgraph_size = max_subgraph_size, max_path_length = max_path_length, - edge_limit = edge_limit, + max_reranker_input = max_reranker_input, streaming = True, chunk_callback = send_chunk, explain_callback = send_explainability, @@ -242,8 +250,8 @@ class Processor(FlowProcessor): entity_limit = entity_limit, triple_limit = triple_limit, max_subgraph_size = max_subgraph_size, max_path_length = max_path_length, - edge_limit = edge_limit, + max_reranker_input = max_reranker_input, explain_callback = send_explainability, save_answer_callback = save_answer, parent_uri = v.parent_uri, @@ -346,6 +354,13 @@ class Processor(FlowProcessor): help=f'Max edges selected per hop by cross-encoder (default: 25)' ) + parser.add_argument( + '--max-reranker-input', + type=int, + default=350, + help=f'Max candidate edges sent to the reranker per hop (default: 350)' + ) + # Note: Explainability triples are now stored in the request's collection # with the named graph urn:graph:retrieval (no separate collection needed) From cfbd5b90791d44bcb5d2bd587b484bed6b962900 Mon Sep 17 00:00:00 2001 From: Sunny Yang Date: Mon, 6 Jul 2026 03:45:23 -0600 Subject: [PATCH 3/4] feat: export/import workspace knowledge in .tgx bundles (#877 Phase 2) (#1024) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: streaming N-Quads serializer for wire-format triples Groundwork for Phase 2 of #877 (knowledge export). Hand-rolled N-Triples term encoding: rdflib's term.n3() emits Turtle-style forms (numeric shorthand, unescaped newlines) that are invalid in line-oriented N-Quads, so literals are escaped per the ECHAR grammar and IRIs validated for representability. Round-trip tests parse the output back with rdflib's nquads parser and compare term-for-term. * feat: export/import workspace knowledge in .tgx bundles (#877) Phase 2 of the workspace bundle commands: tg-export-workspace now includes the workspace's knowledge by default — per-collection knowledge-graph triples as N-Quads (the collection names the graph, streamed through a tempfile so memory stays flat regardless of knowledge-base size) and the document library (metadata plus content, fetched one document at a time). --config-only skips knowledge on both sides; --triples-limit bounds very large graphs; -f/--flow-id selects the flow the triples services run through. tg-import-workspace streams triples back through the bulk import per collection and recreates library documents (children after parents). Knowledge import is additive, unlike config's skip-existing semantics. Embedding vectors are not carried in bundles: --process re-runs imported documents through the flow, which regenerates extraction output and embeddings; --process-collection targets it. Round-trip covered by unit tests over real archives: export with a mocked Api, re-import, and assert the bulk triples stream and library add calls reproduce the original values (including datatyped literals via the N-Quads path). --- tests/unit/test_cli/conftest.py | 16 +- tests/unit/test_cli/test_nquads.py | 97 +++++ .../test_workspace_bundle_commands.py | 405 ++++++++++++++---- .../trustgraph/cli/export_workspace.py | 260 ++++++++--- .../trustgraph/cli/import_workspace.py | 306 +++++++++++-- trustgraph-cli/trustgraph/cli/nquads.py | 137 ++++++ 6 files changed, 1040 insertions(+), 181 deletions(-) create mode 100644 tests/unit/test_cli/test_nquads.py create mode 100644 trustgraph-cli/trustgraph/cli/nquads.py diff --git a/tests/unit/test_cli/conftest.py b/tests/unit/test_cli/conftest.py index b085345f..ff293844 100644 --- a/tests/unit/test_cli/conftest.py +++ b/tests/unit/test_cli/conftest.py @@ -45,4 +45,18 @@ def sample_metadata(): "metadata": [], "user": "test-user", "collection": "test-collection" - } \ No newline at end of file + } + +def iri(v): + """Wire-format IRI term dict, as triples_query_stream yields them.""" + return {"t": "i", "i": v} + + +def lit(v, d=None, lang=None): + """Wire-format literal term dict (optional datatype / language).""" + t = {"t": "l", "v": v} + if d: + t["d"] = d + if lang: + t["l"] = lang + return t diff --git a/tests/unit/test_cli/test_nquads.py b/tests/unit/test_cli/test_nquads.py new file mode 100644 index 00000000..d395b324 --- /dev/null +++ b/tests/unit/test_cli/test_nquads.py @@ -0,0 +1,97 @@ +""" +Round-trip tests for the streaming N-Quads serializer: wire-format triples +are serialized line-by-line, then parsed back with rdflib's nquads parser +and compared term-for-term — proving the output is valid N-Quads and the +encoding (escaping, datatypes, language tags, unicode) is lossless. +""" + +import io + +import rdflib + +from trustgraph.cli.nquads import serialize_nquads, triple_to_nquad + +from tests.unit.test_cli.conftest import iri, lit + +GRAPH = "urn:trustgraph:collection:default" + + +def roundtrip(batches): + """Serialize then parse back; return (parsed_dataset, written, skipped).""" + out = io.StringIO() + written, skipped = serialize_nquads(batches, GRAPH, out) + ds = rdflib.Dataset() + ds.parse(data=out.getvalue(), format="nquads") + return ds, written, skipped + + +class TestNquadsRoundTrip: + + def test_iri_and_literal_flavours_survive_roundtrip(self): + batches = [[ + {"s": iri("http://example.com/s"), "p": iri("http://example.com/p"), + "o": iri("http://example.com/o")}, + {"s": iri("http://example.com/s"), "p": iri("http://example.com/label"), + "o": lit("plain value")}, + {"s": iri("http://example.com/s"), "p": iri("http://example.com/label"), + "o": lit("bonjour", lang="fr")}, + {"s": iri("http://example.com/s"), "p": iri("http://example.com/count"), + "o": lit("42", d="http://www.w3.org/2001/XMLSchema#integer")}, + ]] + ds, written, skipped = roundtrip(batches) + + assert (written, skipped) == (4, 0) + quads = list(ds.quads((None, None, None, None))) + assert len(quads) == 4 + g = rdflib.URIRef(GRAPH) + assert all(q[3] == g for q in quads) + + objects = {q[2] for q in quads} + assert rdflib.URIRef("http://example.com/o") in objects + assert rdflib.Literal("plain value") in objects + assert rdflib.Literal("bonjour", lang="fr") in objects + assert rdflib.Literal( + "42", datatype=rdflib.URIRef("http://www.w3.org/2001/XMLSchema#integer") + ) in objects + + def test_hostile_literal_content_is_escaped_losslessly(self): + nasty = 'line1\nline2\t"quoted" back\\slash 中文 emoji\U0001f680' + batches = [[{ + "s": iri("http://example.com/s"), + "p": iri("http://example.com/note"), + "o": lit(nasty), + }]] + ds, written, skipped = roundtrip(batches) + + assert (written, skipped) == (1, 0) + obj = next(iter(ds.quads((None, None, None, None))))[2] + assert str(obj) == nasty + + def test_malformed_and_unrepresentable_terms_are_skipped_not_emitted(self): + batches = [[ + # IRI with a space (matches graph_to_turtle's malformed skip) + {"s": iri("http://example.com/bad iri"), "p": iri("http://example.com/p"), + "o": lit("x")}, + # RDF-star quoted triple: no N-Quads encoding + {"s": iri("http://example.com/s"), "p": iri("http://example.com/p"), + "o": {"t": "r", "r": {}}}, + # literal in predicate position: invalid RDF + {"s": iri("http://example.com/s"), "p": lit("not-a-predicate"), + "o": lit("x")}, + # one good triple to prove the stream continues past skips + {"s": iri("http://example.com/s"), "p": iri("http://example.com/p"), + "o": lit("good")}, + ]] + ds, written, skipped = roundtrip(batches) + + assert (written, skipped) == (1, 3) + assert len(list(ds.quads((None, None, None, None)))) == 1 + + def test_streaming_shape_one_line_per_triple(self): + line = triple_to_nquad( + {"s": iri("http://example.com/s"), "p": iri("http://example.com/p"), + "o": lit("v")}, + f"<{GRAPH}>", + ) + assert line.endswith(" .\n") + assert line.count("\n") == 1 diff --git a/tests/unit/test_cli/test_workspace_bundle_commands.py b/tests/unit/test_cli/test_workspace_bundle_commands.py index e45c0ef4..915e0dee 100644 --- a/tests/unit/test_cli/test_workspace_bundle_commands.py +++ b/tests/unit/test_cli/test_workspace_bundle_commands.py @@ -3,19 +3,25 @@ Tests for tg-export-workspace / tg-import-workspace (.tgx bundle commands). The Api class is mocked in each command module's namespace (same pattern as test_config_commands.py); bundles are written to and read from tmp_path so -the archive format itself is exercised end-to-end. +the archive format itself is exercised end-to-end, including the Phase-2 +knowledge tree (per-collection N-Quads + library documents). """ +import datetime +import io import json import tarfile +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest -from trustgraph.api.types import ConfigValue +from trustgraph.api.types import ConfigValue, Triple from trustgraph.cli.export_workspace import export_workspace from trustgraph.cli.import_workspace import import_workspace +from tests.unit.test_cli.conftest import iri, lit + SAMPLE_CONFIG = { "prompt": { "extract-concepts": json.dumps({"template": "Extract {{q}}"}), @@ -26,25 +32,110 @@ SAMPLE_CONFIG = { }, } +# Wire-format triples for one collection, incl. a datatyped literal. +WIRE_BATCHES = [[ + {"s": iri("http://ex.com/s"), "p": iri("http://ex.com/p"), + "o": iri("http://ex.com/o")}, + {"s": iri("http://ex.com/s"), "p": iri("http://ex.com/count"), + "o": lit("42", d="http://www.w3.org/2001/XMLSchema#integer")}, +]] -def make_mock_api(): - mock_config = Mock() +DOC = SimpleNamespace( + id="doc-1", + time=datetime.datetime(2026, 7, 1, 12, 0, 0), + kind="text/plain", + title="Policy", + comments="returns policy", + metadata=[Triple(s="http://ex.com/doc-1", p="http://ex.com/about", + o="returns")], + tags=["policy"], + parent_id="", + document_type="source", +) + + +def make_mock_api(collections=(), batches=(), docs=(), contents=None): + """Full-surface Api mock; returns (mock_api, mock_config).""" mock_api = Mock() + + mock_config = Mock() mock_api.config.return_value = mock_config + mock_config.all.return_value = (SAMPLE_CONFIG, "v42") + + mock_api.collection.return_value.list_collections.return_value = \ + list(collections) + + flow = mock_api.socket.return_value.flow.return_value + flow.triples_query_stream.return_value = iter(batches) + + library = mock_api.library.return_value + library.get_documents.return_value = list(docs) + library.get_document_content.side_effect = \ + lambda id: (contents or {}).get(id, b"") + return mock_api, mock_config +def export_bundle(path, collections=(), batches=(), docs=(), contents=None, + **kwargs): + """Export SAMPLE_CONFIG (+ optional knowledge mocks) to path.""" + mock_api, mock_config = make_mock_api( + collections=collections, batches=batches, docs=docs, + contents=contents, + ) + with patch("trustgraph.cli.export_workspace.Api") as api_cls: + api_cls.return_value = mock_api + export_workspace( + url="http://api/", workspace="source-ws", output=str(path), + **kwargs, + ) + return mock_api, mock_config + + +DEFAULT_MANIFEST = { + "format": "tgx", "format_version": 1, "workspace": "w", + "contents": {"config": True, "knowledge": True}, +} + + +def write_bundle(path, members, manifest=DEFAULT_MANIFEST): + """Write a raw .tgx from name -> bytes (manifest added first).""" + entries = {"manifest.json": json.dumps(manifest).encode(), **members} + with tarfile.open(path, "w:gz") as tar: + for name, data in entries.items(): + info = tarfile.TarInfo(name) + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + return path + + +def run_import(mock_api, input, **kwargs): + """Run import_workspace against a mocked Api.""" + with patch("trustgraph.cli.import_workspace.Api") as api_cls: + api_cls.return_value = mock_api + import_workspace(url="http://api/", input=str(input), **kwargs) + return api_cls + + @pytest.fixture def bundle(tmp_path): - """Export SAMPLE_CONFIG to a real .tgx and yield its path.""" + """Config-only-shaped bundle (no collections/docs mocked).""" path = tmp_path / "ws.tgx" - with patch("trustgraph.cli.export_workspace.Api") as api_cls: - mock_api, mock_config = make_mock_api() - api_cls.return_value = mock_api - mock_config.all.return_value = (SAMPLE_CONFIG, "v42") - export_workspace( - url="http://api/", workspace="source-ws", output=str(path), - ) + export_bundle(path) + return path + + +@pytest.fixture +def knowledge_bundle(tmp_path): + """Bundle with one collection (2 triples) and one document.""" + path = tmp_path / "kws.tgx" + export_bundle( + path, + collections=[SimpleNamespace(collection="research")], + batches=WIRE_BATCHES, + docs=[DOC], + contents={"doc-1": b"Customers may return items."}, + ) return path @@ -58,7 +149,10 @@ class TestExportWorkspace: assert manifest["format"] == "tgx" assert manifest["workspace"] == "source-ws" assert manifest["config_version"] == "v42" - assert manifest["contents"] == {"config": True, "knowledge": False} + assert manifest["contents"] == {"config": True, "knowledge": True} + assert manifest["knowledge"] == { + "collections": [], "documents": 0, "triples": {}, + } assert "config/prompt/extract-concepts.json" in names assert "config/prompt/answer.json" in names @@ -79,12 +173,12 @@ class TestExportWorkspace: def test_path_unsafe_keys_are_quoted_in_filenames(self, tmp_path): path = tmp_path / "ws.tgx" + mock_api, mock_config = make_mock_api() + mock_config.all.return_value = ( + {"prompt": {"a/b": json.dumps({"x": 1})}}, "v1", + ) with patch("trustgraph.cli.export_workspace.Api") as api_cls: - mock_api, mock_config = make_mock_api() api_cls.return_value = mock_api - mock_config.all.return_value = ( - {"prompt": {"a/b": json.dumps({"x": 1})}}, "v1", - ) export_workspace( url="http://api/", workspace="ws", output=str(path), ) @@ -94,16 +188,63 @@ class TestExportWorkspace: assert "config/prompt/a%2Fb.json" in names assert entry["key"] == "a/b" + def test_knowledge_tree_written_per_collection_and_document( + self, knowledge_bundle): + with tarfile.open(knowledge_bundle, "r:gz") as tar: + names = tar.getnames() + manifest = json.load(tar.extractfile("manifest.json")) + nq = tar.extractfile( + "knowledge/research/triples.nq").read().decode() + meta = json.load( + tar.extractfile("knowledge/library/doc-1.meta.json")) + content = tar.extractfile( + "knowledge/library/doc-1.content").read() + + assert manifest["contents"]["knowledge"] is True + assert manifest["knowledge"] == { + "collections": ["research"], "documents": 1, + "triples": {"research": 2}, + } + assert "knowledge/research/triples.nq" in names + # N-Quads: one line per triple, graph = the collection IRI, and + # the datatyped literal keeps its full quoted form. + lines = [ln for ln in nq.splitlines() if ln] + assert len(lines) == 2 + assert all("" in ln + for ln in lines) + assert '"42"^^' in nq + + assert meta["id"] == "doc-1" + assert meta["title"] == "Policy" + assert meta["metadata"] == [{ + "s": "http://ex.com/doc-1", "p": "http://ex.com/about", + "o": "returns", + }] + assert content == b"Customers may return items." + + def test_config_only_skips_knowledge(self, tmp_path): + path = tmp_path / "co.tgx" + mock_api, _ = export_bundle( + path, + collections=[SimpleNamespace(collection="research")], + config_only=True, + ) + with tarfile.open(path, "r:gz") as tar: + names = tar.getnames() + manifest = json.load(tar.extractfile("manifest.json")) + assert manifest["contents"] == {"config": True, "knowledge": False} + assert "knowledge" not in manifest + assert not any(n.startswith("knowledge/") for n in names) + mock_api.collection.return_value.list_collections.assert_not_called() + mock_api.socket.assert_not_called() + mock_api.library.assert_not_called() + class TestImportWorkspace: def test_roundtrip_puts_all_values_with_overwrite(self, bundle): - with patch("trustgraph.cli.import_workspace.Api") as api_cls: - mock_api, mock_config = make_mock_api() - api_cls.return_value = mock_api - import_workspace( - url="http://api/", input=str(bundle), overwrite=True, - ) + mock_api, mock_config = make_mock_api() + api_cls = run_import(mock_api, bundle, overwrite=True) # Target workspace defaults to the manifest's workspace. api_cls.assert_called_once_with( @@ -123,27 +264,22 @@ class TestImportWorkspace: assert all(isinstance(v, ConfigValue) for v in values) def test_workspace_flag_renames_target(self, bundle): - with patch("trustgraph.cli.import_workspace.Api") as api_cls: - mock_api, mock_config = make_mock_api() - api_cls.return_value = mock_api - import_workspace( - url="http://api/", input=str(bundle), workspace="staging", - overwrite=True, - ) + mock_api, mock_config = make_mock_api() + api_cls = run_import( + mock_api, bundle, workspace="staging", overwrite=True, + ) api_cls.assert_called_once_with( "http://api/", token=None, workspace="staging", ) def test_default_skips_existing_keys(self, bundle): """WorkspaceInit re-run semantics: only missing keys are written.""" - with patch("trustgraph.cli.import_workspace.Api") as api_cls: - mock_api, mock_config = make_mock_api() - api_cls.return_value = mock_api - mock_config.list.side_effect = lambda t: { - "prompt": ["extract-concepts"], - "tool": [], - }[t] - import_workspace(url="http://api/", input=str(bundle)) + mock_api, mock_config = make_mock_api() + mock_config.list.side_effect = lambda t: { + "prompt": ["extract-concepts"], + "tool": [], + }[t] + run_import(mock_api, bundle) values = mock_config.put.call_args.args[0] assert sorted((v.type, v.key) for v in values) == [ @@ -152,17 +288,11 @@ class TestImportWorkspace: ] def test_dry_run_writes_nothing(self, bundle, capsys): - with patch("trustgraph.cli.import_workspace.Api") as api_cls: - mock_api, mock_config = make_mock_api() - api_cls.return_value = mock_api - import_workspace( - url="http://api/", input=str(bundle), overwrite=True, - dry_run=True, - ) + mock_api, mock_config = make_mock_api() + run_import(mock_api, bundle, overwrite=True, dry_run=True) mock_config.put.assert_not_called() out = capsys.readouterr().out assert "would import prompt/extract-concepts" in out - assert "3 item(s) would be imported" in out def test_rejects_bundle_without_manifest(self, tmp_path): path = tmp_path / "bad.tgx" @@ -173,52 +303,147 @@ class TestImportWorkspace: import_workspace(url="http://api/", input=str(path)) def test_rejects_newer_format_version(self, tmp_path): - import io - path = tmp_path / "future.tgx" - manifest = json.dumps({ - "format": "tgx", "format_version": 99, "workspace": "w", - "contents": {"config": True, "knowledge": False}, - }).encode() - with tarfile.open(path, "w:gz") as tar: - info = tarfile.TarInfo("manifest.json") - info.size = len(manifest) - tar.addfile(info, io.BytesIO(manifest)) + path = write_bundle( + tmp_path / "future.tgx", {}, + manifest={**DEFAULT_MANIFEST, "format_version": 99}, + ) with patch("trustgraph.cli.import_workspace.Api"): with pytest.raises(RuntimeError, match="newer than this tool"): import_workspace(url="http://api/", input=str(path)) - def test_refuses_knowledge_bundle_without_config_only(self, tmp_path): - import io - path = tmp_path / "knowledge.tgx" - manifest = json.dumps({ - "format": "tgx", "format_version": 1, "workspace": "w", - "contents": {"config": True, "knowledge": True}, - }).encode() - with tarfile.open(path, "w:gz") as tar: - info = tarfile.TarInfo("manifest.json") - info.size = len(manifest) - tar.addfile(info, io.BytesIO(manifest)) - with patch("trustgraph.cli.import_workspace.Api"): - with pytest.raises(RuntimeError, match="--config-only"): - import_workspace(url="http://api/", input=str(path)) - def test_config_only_flag_allows_knowledge_bundle(self, tmp_path): - import io - path = tmp_path / "knowledge.tgx" - manifest = json.dumps({ - "format": "tgx", "format_version": 1, "workspace": "w", - "contents": {"config": True, "knowledge": True}, - }).encode() - with tarfile.open(path, "w:gz") as tar: - info = tarfile.TarInfo("manifest.json") - info.size = len(manifest) - tar.addfile(info, io.BytesIO(manifest)) - with patch("trustgraph.cli.import_workspace.Api") as api_cls: - mock_api, mock_config = make_mock_api() - api_cls.return_value = mock_api - import_workspace( - url="http://api/", input=str(path), config_only=True, - overwrite=True, - ) - # No config entries in this bundle; nothing written, no error. - mock_config.put.assert_not_called() +class TestImportKnowledge: + + def test_roundtrip_imports_triples_and_documents(self, knowledge_bundle): + mock_api, mock_config = make_mock_api() + run_import(mock_api, knowledge_bundle, overwrite=True) + + # Triples land in the bulk import stream for the right collection. + bulk = mock_api.bulk.return_value + call = bulk.import_triples.call_args + assert call.args[0] == "default" # flow id + triples = sorted(list(call.args[1]), key=lambda t: t.p) + assert triples == [ + Triple(s="http://ex.com/s", p="http://ex.com/count", o="42"), + Triple(s="http://ex.com/s", p="http://ex.com/p", + o="http://ex.com/o"), + ] + assert call.kwargs["metadata"]["collection"] == "research" + + # The document is recreated with its metadata and content. + add = mock_api.library.return_value.add_document.call_args + assert add.kwargs["id"] == "doc-1" + assert add.kwargs["document"] == b"Customers may return items." + assert add.kwargs["title"] == "Policy" + assert add.kwargs["kind"] == "text/plain" + assert add.kwargs["tags"] == ["policy"] + assert add.kwargs["metadata"] == [ + Triple(s="http://ex.com/doc-1", p="http://ex.com/about", + o="returns"), + ] + # No processing unless asked: embeddings re-derivation is opt-in. + mock_api.library.return_value.start_processing.assert_not_called() + + def test_config_only_skips_knowledge_on_import(self, knowledge_bundle): + mock_api, mock_config = make_mock_api() + run_import(mock_api, knowledge_bundle, overwrite=True, + config_only=True) + mock_api.bulk.assert_not_called() + mock_api.library.return_value.add_document.assert_not_called() + mock_config.put.assert_called_once() + + def test_dry_run_covers_knowledge(self, knowledge_bundle, capsys): + mock_api, mock_config = make_mock_api() + run_import(mock_api, knowledge_bundle, overwrite=True, dry_run=True) + mock_api.bulk.assert_not_called() + mock_api.library.return_value.add_document.assert_not_called() + out = capsys.readouterr().out + assert "would import 2 triple(s) into collection 'research'" in out + assert "would import document doc-1" in out + + def test_process_flag_reprocesses_documents(self, knowledge_bundle): + mock_api, mock_config = make_mock_api() + run_import(mock_api, knowledge_bundle, overwrite=True, process=True, + process_collection="research") + proc = mock_api.library.return_value.start_processing.call_args + assert proc.kwargs["document_id"] == "doc-1" + assert proc.kwargs["flow"] == "default" + assert proc.kwargs["collection"] == "research" + + +class TestImportDocumentSkipOverwrite: + """Live-verified semantics: existing documents skip by default, replace + with --overwrite (remove + add; the library API has no content update).""" + + def _bundle_with_doc(self, tmp_path): + return write_bundle(tmp_path / "doc.tgx", { + "knowledge/library/doc-1.meta.json": json.dumps( + {"id": "doc-1", "title": "T", "metadata": []}).encode(), + "knowledge/library/doc-1.content": b"hello", + }) + + def test_existing_document_is_skipped_not_fatal(self, tmp_path): + path = self._bundle_with_doc(tmp_path) + mock_api, mock_config = make_mock_api( + docs=[SimpleNamespace(id="doc-1")], + ) + run_import(mock_api, path, overwrite=False) + lib = mock_api.library.return_value + lib.add_document.assert_not_called() + lib.remove_document.assert_not_called() + + def test_overwrite_replaces_existing_document(self, tmp_path): + path = self._bundle_with_doc(tmp_path) + mock_api, mock_config = make_mock_api( + docs=[SimpleNamespace(id="doc-1")], + ) + run_import(mock_api, path, overwrite=True) + lib = mock_api.library.return_value + lib.remove_document.assert_called_once_with("doc-1") + lib.add_document.assert_called_once() + + +class TestCollectionDiscovery: + """Live-verified: implicitly-created collections (raw triple loads) are + queryable but unlisted, so export merges --collection extras and import + registers what it restores.""" + + def test_export_includes_extra_unregistered_collections(self, tmp_path): + path = tmp_path / "ws.tgx" + mock_api, _ = export_bundle( + path, + collections=[SimpleNamespace(collection="default")], + batches=[[]], + extra_collections=["research"], + ) + with tarfile.open(path, "r:gz") as tar: + manifest = json.load(tar.extractfile("manifest.json")) + assert manifest["knowledge"]["collections"] == ["default", "research"] + + def test_import_registers_each_restored_collection(self, tmp_path): + path = write_bundle(tmp_path / "kb.tgx", { + "knowledge/research/triples.nq": ( + b' "v" ' + b' .\n' + ), + }) + mock_api, mock_config = make_mock_api() + run_import(mock_api, path) + mock_api.collection.return_value.update_collection.assert_called_once_with( + "research", name="research", + ) + + def test_registered_collections_are_not_reregistered(self, tmp_path): + """update_collection is an upsert that clears omitted fields, so a + collection already in the registry must be left untouched.""" + path = write_bundle(tmp_path / "kb.tgx", { + "knowledge/research/triples.nq": ( + b' "v" ' + b' .\n' + ), + }) + mock_api, mock_config = make_mock_api( + collections=[SimpleNamespace(collection="research")], + ) + run_import(mock_api, path) + mock_api.collection.return_value.update_collection.assert_not_called() diff --git a/trustgraph-cli/trustgraph/cli/export_workspace.py b/trustgraph-cli/trustgraph/cli/export_workspace.py index 41d3403d..2474ce36 100644 --- a/trustgraph-cli/trustgraph/cli/export_workspace.py +++ b/trustgraph-cli/trustgraph/cli/export_workspace.py @@ -1,14 +1,14 @@ """ -Exports a workspace's full configuration state as a portable .tgx bundle -(a gzipped tar archive) for backup, migration between deployments, or -sharing a pre-configured workspace. +Exports a workspace's full state as a portable .tgx bundle (a gzipped tar +archive) for backup, migration between deployments, or sharing a +pre-configured workspace. -The bundle is human-readable: a manifest.json plus one pretty-printed JSON -file per config key under config//, so it can be inspected and -hand-edited before import. Each entry file embeds its own type and key, so -filenames are cosmetic. Knowledge export (triples, documents, embeddings) -is not yet included; the manifest records that so future importers can -distinguish config-only bundles. +The bundle carries the workspace configuration (one pretty-printed JSON +file per config key) and, by default, its knowledge: per-collection +knowledge-graph triples as N-Quads (the collection names the graph) and +the document library (metadata plus content). Pass --config-only to +export just the configuration. Embedding vectors are not exported — +re-processing imported documents through a flow regenerates them. """ import argparse @@ -17,11 +17,14 @@ import json import os import sys import tarfile +import tempfile import time from urllib.parse import quote from trustgraph.api import Api +from . nquads import serialize_nquads + default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_token = os.getenv("TRUSTGRAPH_TOKEN", None) default_workspace = os.getenv("TRUSTGRAPH_WORKSPACE", "default") @@ -29,6 +32,10 @@ default_workspace = os.getenv("TRUSTGRAPH_WORKSPACE", "default") MANIFEST_FORMAT = "tgx" MANIFEST_FORMAT_VERSION = 1 +# triples_query_stream is bounded by a limit; exports want "everything", so +# default high and let --triples-limit override for truly huge graphs. +DEFAULT_TRIPLES_LIMIT = 1_000_000 + def _add_bytes(tar, name, data): info = tarfile.TarInfo(name=name) @@ -37,56 +44,176 @@ def _add_bytes(tar, name, data): tar.addfile(info, io.BytesIO(data)) -def export_workspace(url, workspace, output, token=None): - - api = Api(url, token=token, workspace=workspace).config() - - config, version = api.all() - - manifest = { - "format": MANIFEST_FORMAT, - "format_version": MANIFEST_FORMAT_VERSION, - "workspace": workspace, - "config_version": version, - "exported_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), - "contents": {"config": True, "knowledge": False}, - } - +def _export_config(tar, config): + """Write one self-describing JSON file per config key; return count.""" count = 0 + for type_, entries in sorted(config.items()): + for key, raw in sorted(entries.items()): + + # Config values are stored as JSON strings; parse so the + # bundle is pretty-printed and hand-editable. A value that + # isn't valid JSON is preserved verbatim. + try: + value = json.loads(raw) + except (TypeError, json.JSONDecodeError): + value = raw + + entry = {"type": type_, "key": key, "value": value} + + # Keys may contain path-unsafe characters; the entry embeds + # the real key, so the quoted filename is cosmetic only. + name = f"config/{quote(type_, safe='')}/{quote(key, safe='')}.json" + + _add_bytes( + tar, name, + json.dumps(entry, indent=2).encode("utf-8"), + ) + + count += 1 + return count + + +def _export_triples(tar, api, flow_id, collections, triples_limit): + """Stream each collection's triples into knowledge//triples.nq. + + N-Quads are written to a tempfile first (tar members need their size + upfront), so memory stays flat regardless of knowledge-base size. + Returns {collection: written} and a total skipped count. + """ + counts = {} + skipped_total = 0 + socket = api.socket() + try: + flow = socket.flow(flow_id) + for c in collections: + graph_iri = f"urn:trustgraph:collection:{quote(c, safe='')}" + tmp = tempfile.NamedTemporaryFile( + "w", encoding="utf-8", suffix=".nq", delete=False, + ) + try: + with tmp: + written, skipped = serialize_nquads( + flow.triples_query_stream( + s=None, p=None, o=None, + collection=c, + limit=triples_limit, + batch_size=100, + ), + graph_iri, + tmp, + ) + if written: + tar.add( + tmp.name, + arcname=( + f"knowledge/{quote(c, safe='')}/triples.nq" + ), + ) + counts[c] = written + skipped_total += skipped + finally: + os.unlink(tmp.name) + finally: + socket.close() + return counts, skipped_total + + +def _export_library(tar, api): + """Write each library document's metadata + content; return count.""" + library = api.library() + count = 0 + for doc in library.get_documents(include_children=True): + # Content is fetched one document at a time so memory is bounded + # by the largest single document, not the whole library. + content = library.get_document_content(doc.id) + meta = { + "id": doc.id, + "time": doc.time.isoformat() if doc.time else None, + "kind": doc.kind, + "title": doc.title, + "comments": doc.comments, + "metadata": [ + {"s": t.s, "p": t.p, "o": t.o} for t in (doc.metadata or []) + ], + "tags": list(doc.tags or []), + "parent_id": doc.parent_id or "", + "document_type": getattr(doc, "document_type", "") or "", + } + base = f"knowledge/library/{quote(doc.id, safe='')}" + _add_bytes( + tar, f"{base}.meta.json", + json.dumps(meta, indent=2).encode("utf-8"), + ) + _add_bytes(tar, f"{base}.content", content or b"") + count += 1 + return count + + +def export_workspace( + url, workspace, output, token=None, config_only=False, + flow_id="default", triples_limit=DEFAULT_TRIPLES_LIMIT, + extra_collections=(), +): + + api = Api(url, token=token, workspace=workspace) + + config, version = api.config().all() + + # Collection discovery is registry-based: collections created implicitly + # by raw triple loads (e.g. tg-load-knowledge) are queryable but not + # listed, so they would silently drop out of the bundle. --collection + # names them explicitly; the enumeration is printed so what's included + # is never a guess. + collections = [] + if not config_only: + registered = [c.collection for c in api.collection().list_collections()] + collections = sorted(set(registered) | set(extra_collections)) + print(f"Exporting collections: {', '.join(collections)}", flush=True) with tarfile.open(output, "w:gz") as tar: + config_count = _export_config(tar, config) + + triple_counts = {} + skipped = 0 + doc_count = 0 + if not config_only: + triple_counts, skipped = _export_triples( + tar, api, flow_id, collections, triples_limit, + ) + doc_count = _export_library(tar, api) + + manifest = { + "format": MANIFEST_FORMAT, + "format_version": MANIFEST_FORMAT_VERSION, + "workspace": workspace, + "config_version": version, + "exported_at": time.strftime( + "%Y-%m-%dT%H:%M:%SZ", time.gmtime(), + ), + "contents": {"config": True, "knowledge": not config_only}, + } + if not config_only: + manifest["knowledge"] = { + "collections": collections, + "documents": doc_count, + "triples": triple_counts, + } + _add_bytes( tar, "manifest.json", json.dumps(manifest, indent=2).encode("utf-8"), ) - for type_, entries in sorted(config.items()): - for key, raw in sorted(entries.items()): - - # Config values are stored as JSON strings; parse so the - # bundle is pretty-printed and hand-editable. A value that - # isn't valid JSON is preserved verbatim. - try: - value = json.loads(raw) - except (TypeError, json.JSONDecodeError): - value = raw - - entry = {"type": type_, "key": key, "value": value} - - # Keys may contain path-unsafe characters; the entry embeds - # the real key, so the quoted filename is cosmetic only. - name = f"config/{quote(type_, safe='')}/{quote(key, safe='')}.json" - - _add_bytes( - tar, name, - json.dumps(entry, indent=2).encode("utf-8"), - ) - - count += 1 - - print(f"Exported {count} config item(s) from workspace " - f"'{workspace}' to {output}", flush=True) + summary = f"Exported {config_count} config item(s)" + if not config_only: + summary += ( + f", {sum(triple_counts.values())} triple(s) across " + f"{len(triple_counts)} collection(s), {doc_count} document(s)" + ) + if skipped: + summary += f" ({skipped} triple(s) not representable, skipped)" + print(f"{summary} from workspace '{workspace}' to {output}", flush=True) def main(): @@ -114,12 +241,41 @@ def main(): help=f'Workspace to export (default: {default_workspace})', ) + parser.add_argument( + '-c', '--collection', + action='append', + default=[], + help='Additionally export this collection even if it is not ' + 'registered in collection management (repeatable)', + ) + parser.add_argument( '-o', '--output', required=True, help='Output bundle path, e.g. workspace-default.tgx', ) + parser.add_argument( + '-f', '--flow-id', + default="default", + help='Flow to query triples through (default: default)', + ) + + parser.add_argument( + '--config-only', + action='store_true', + help='Export only the configuration, skipping knowledge ' + '(triples and library documents)', + ) + + parser.add_argument( + '--triples-limit', + type=int, + default=DEFAULT_TRIPLES_LIMIT, + help='Maximum triples to export per collection ' + f'(default: {DEFAULT_TRIPLES_LIMIT})', + ) + args = parser.parse_args() try: @@ -129,6 +285,10 @@ def main(): workspace=args.workspace, output=args.output, token=args.token, + config_only=args.config_only, + flow_id=args.flow_id, + triples_limit=args.triples_limit, + extra_collections=args.collection, ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/import_workspace.py b/trustgraph-cli/trustgraph/cli/import_workspace.py index 71885bfb..13fabdd1 100644 --- a/trustgraph-cli/trustgraph/cli/import_workspace.py +++ b/trustgraph-cli/trustgraph/cli/import_workspace.py @@ -3,10 +3,16 @@ Imports a workspace bundle (.tgx, produced by tg-export-workspace) into a TrustGraph deployment. The target workspace defaults to the name recorded in the bundle's manifest and can be renamed with --workspace. -By default existing (type, key) entries in the target workspace are left -untouched and only missing keys are added, matching WorkspaceInit's -re-run behaviour; pass --overwrite to replace every imported key. Use ---dry-run to show what would be written without changing anything. +Configuration import follows WorkspaceInit's re-run behaviour: existing +(type, key) entries are left untouched and only missing keys are added; +pass --overwrite to replace every imported key. Knowledge import (triples +and library documents) is additive — triples are streamed into the target +collection and documents are added to the library; re-importing the same +bundle twice will duplicate knowledge, not merge it. Use --dry-run to +show what would be written without changing anything, --config-only to +skip a bundle's knowledge, and --process to re-run imported documents +through a flow (which regenerates embeddings, so bundles don't carry +vectors). """ import argparse @@ -14,9 +20,12 @@ import json import os import sys import tarfile +import uuid +from urllib.parse import unquote from trustgraph.api import Api -from trustgraph.api.types import ConfigValue +from trustgraph.api.types import ConfigValue, Triple +from trustgraph.cli.nquads import parse_nquads default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_token = os.getenv("TRUSTGRAPH_TOKEN", None) @@ -26,10 +35,21 @@ SUPPORTED_FORMAT_VERSION = 1 def _read_bundle(path): - """Read manifest and config entries from a .tgx bundle.""" + """Read manifest, config entries, triples and documents from a .tgx. + + Returns (manifest, config_entries, triples_by_collection, documents) + where triples_by_collection maps collection -> list[Triple] and + documents is a list of {meta: dict, content: bytes}. + """ manifest = None - entries = [] + config_entries = [] + triples = {} + doc_meta = {} + doc_content = {} + + def member_id(name, prefix, suffix): + return unquote(name[len(prefix):-len(suffix)]) with tarfile.open(path, "r:gz") as tar: for member in tar.getmembers(): @@ -39,11 +59,28 @@ def _read_bundle(path): if f is None: continue data = f.read() - if member.name == "manifest.json": + name = member.name + + if name == "manifest.json": manifest = json.loads(data) - elif member.name.startswith("config/") and \ - member.name.endswith(".json"): - entries.append(json.loads(data)) + + elif name.startswith("config/") and name.endswith(".json"): + config_entries.append(json.loads(data)) + + elif name.startswith("knowledge/library/") and \ + name.endswith(".meta.json"): + doc_id = member_id(name, "knowledge/library/", ".meta.json") + doc_meta[doc_id] = json.loads(data) + + elif name.startswith("knowledge/library/") and \ + name.endswith(".content"): + doc_id = member_id(name, "knowledge/library/", ".content") + doc_content[doc_id] = data + + elif name.startswith("knowledge/") and \ + name.endswith("/triples.nq"): + collection = member_id(name, "knowledge/", "/triples.nq") + triples[collection] = parse_nquads(data) if manifest is None: raise RuntimeError("not a workspace bundle: manifest.json missing") @@ -60,27 +97,18 @@ def _read_bundle(path): "upgrade trustgraph-cli" ) - return manifest, entries + documents = [ + {"meta": meta, "content": doc_content.get(doc_id, b"")} + for doc_id, meta in doc_meta.items() + ] + + return manifest, config_entries, triples, documents -def import_workspace( - url, input, workspace=None, overwrite=False, config_only=False, - dry_run=False, token=None, -): +def _import_config(api, entries, overwrite, dry_run): + """Import config entries; returns (imported, skipped) counts.""" - manifest, entries = _read_bundle(input) - - # Knowledge import (triples, documents, embeddings) is not implemented - # yet; refuse to silently drop it from a bundle that carries it. - if manifest.get("contents", {}).get("knowledge") and not config_only: - raise RuntimeError( - "bundle contains knowledge data, which this tool cannot import " - "yet; re-run with --config-only to import just the configuration" - ) - - target = workspace or manifest.get("workspace") or "default" - - api = Api(url, token=token, workspace=target).config() + config = api.config() # Mirror WorkspaceInit's re-run behaviour: without --overwrite, keys # already present in the target workspace are skipped (per key, not per @@ -88,7 +116,7 @@ def import_workspace( existing = {} if not overwrite: for type_ in sorted({e["type"] for e in entries}): - existing[type_] = set(api.list(type_)) + existing[type_] = set(config.list(type_)) values = [] skipped = 0 @@ -106,16 +134,191 @@ def import_workspace( if dry_run: for v in values: print(f"would import {v.type}/{v.key}", flush=True) - print(f"Dry run: {len(values)} item(s) would be imported into " - f"workspace '{target}', {skipped} skipped as existing", - flush=True) - return + elif values: + config.put(values) - if values: - api.put(values) + return len(values), skipped - print(f"Imported {len(values)} config item(s) into workspace " - f"'{target}', {skipped} skipped as existing", flush=True) + +def _import_triples(api, flow_id, triples_by_collection, dry_run): + """Stream each collection's triples into the flow; returns count.""" + + # Collections restored by the bulk triples path (unlike document + # processing) are not auto-registered, and an unregistered collection + # would silently drop out of a future export. Register only the missing + # ones: update_collection is an upsert whose omitted fields clear the + # description/tags that _import_config may have just restored. + # Best-effort — a registry hiccup shouldn't fail a completed triple + # import (tg-set-collection is the manual remedy). + registered = None + if triples_by_collection and not dry_run: + try: + registered = { + c.collection for c in api.collection().list_collections() + } + except Exception as e: + print( + f"warning: could not list collections for registration: " + f"{e} — register manually with tg-set-collection", + flush=True, + ) + + total = 0 + for collection, triples in sorted(triples_by_collection.items()): + if dry_run: + print( + f"would import {len(triples)} triple(s) into " + f"collection '{collection}'", + flush=True, + ) + else: + api.bulk().import_triples( + flow_id, + triples, + metadata={ + "id": f"workspace-import-{uuid.uuid4()}", + "metadata": [], + "collection": collection, + }, + ) + if registered is not None and collection not in registered: + try: + api.collection().update_collection( + collection, name=collection, + ) + except Exception as e: + print( + f"warning: could not register collection " + f"'{collection}': {e} — register manually with " + f"tg-set-collection", + flush=True, + ) + total += len(triples) + return total + + +def _import_documents( + api, documents, flow_id, process, process_collection, overwrite, + dry_run, +): + """Add library documents back (children after parents). + + Mirrors the config semantics: documents already present in the target + workspace are skipped unless --overwrite, which replaces them (the + library API has no in-place content update, so replace = remove + add). + Returns (imported, skipped). + """ + + library = api.library() + + existing = set() + if documents: + existing = {d.id for d in library.get_documents(include_children=True)} + + # Parents must exist before their children. + ordered = sorted( + documents, key=lambda d: bool(d["meta"].get("parent_id")), + ) + + count = 0 + doc_skipped = 0 + for doc in ordered: + meta = doc["meta"] + + if meta["id"] in existing and not overwrite: + doc_skipped += 1 + if dry_run: + print(f"would skip existing document {meta['id']}", flush=True) + continue + + if dry_run: + print(f"would import document {meta['id']}", flush=True) + count += 1 + continue + + if meta["id"] in existing: + library.remove_document(meta["id"]) + + metadata = [ + Triple(s=t["s"], p=t["p"], o=t["o"]) + for t in meta.get("metadata", []) + ] + if meta.get("parent_id"): + library.add_child_document( + document=doc["content"], + id=meta["id"], + parent_id=meta["parent_id"], + title=meta.get("title", ""), + comments=meta.get("comments", ""), + kind=meta.get("kind", "text/plain"), + tags=meta.get("tags", []), + metadata=metadata, + ) + else: + library.add_document( + document=doc["content"], + id=meta["id"], + metadata=metadata, + title=meta.get("title", ""), + comments=meta.get("comments", ""), + kind=meta.get("kind", "text/plain"), + tags=meta.get("tags", []), + ) + + # Re-processing regenerates extraction output and embeddings for + # the imported content (bundles carry no vectors). + if process: + library.start_processing( + id=f"proc-{uuid.uuid4()}", + document_id=meta["id"], + flow=flow_id, + collection=process_collection, + tags=meta.get("tags", []), + ) + + count += 1 + return count, doc_skipped + + +def import_workspace( + url, input, workspace=None, overwrite=False, config_only=False, + flow_id="default", process=False, process_collection="default", + dry_run=False, token=None, +): + + manifest, config_entries, triples, documents = _read_bundle(input) + + target = workspace or manifest.get("workspace") or "default" + + api = Api(url, token=token, workspace=target) + + imported, skipped = _import_config( + api, config_entries, overwrite, dry_run, + ) + + triple_count = 0 + doc_count = 0 + has_knowledge = bool(triples or documents) + if has_knowledge and not config_only: + triple_count = _import_triples(api, flow_id, triples, dry_run) + doc_count, doc_skipped = _import_documents( + api, documents, flow_id, process, process_collection, overwrite, + dry_run, + ) + + verb = "Dry run:" if dry_run else "Imported" + summary = ( + f"{verb} {imported} config item(s) into workspace '{target}', " + f"{skipped} skipped as existing" + ) + if has_knowledge and not config_only: + summary += ( + f"; {triple_count} triple(s), {doc_count} document(s)" + f" ({doc_skipped} skipped as existing)" + ) + elif has_knowledge: + summary += "; knowledge skipped (--config-only)" + print(summary, flush=True) def main(): @@ -150,11 +353,18 @@ def main(): 'bundle manifest)', ) + parser.add_argument( + '-f', '--flow-id', + default="default", + help='Flow to import triples through and process documents with ' + '(default: default)', + ) + parser.add_argument( '--overwrite', action='store_true', - help='Replace existing keys in the target workspace (default: ' - 'keep existing keys and only add missing ones)', + help='Replace existing config keys in the target workspace ' + '(default: keep existing keys and only add missing ones)', ) parser.add_argument( @@ -164,6 +374,19 @@ def main(): 'in the bundle', ) + parser.add_argument( + '--process', + action='store_true', + help='Re-process imported documents through the flow after import ' + '(regenerates extraction output and embeddings)', + ) + + parser.add_argument( + '--process-collection', + default="default", + help='Collection that --process targets (default: default)', + ) + parser.add_argument( '--dry-run', action='store_true', @@ -180,6 +403,9 @@ def main(): workspace=args.workspace, overwrite=args.overwrite, config_only=args.config_only, + flow_id=args.flow_id, + process=args.process, + process_collection=args.process_collection, dry_run=args.dry_run, token=args.token, ) diff --git a/trustgraph-cli/trustgraph/cli/nquads.py b/trustgraph-cli/trustgraph/cli/nquads.py new file mode 100644 index 00000000..f38a7747 --- /dev/null +++ b/trustgraph-cli/trustgraph/cli/nquads.py @@ -0,0 +1,137 @@ +""" +N-Quads serialization and parsing for workspace knowledge bundles: the +wire-format triples yielded by triples_query_stream go out one line per +triple (so an export never holds a whole graph in memory), and bundle +members come back as api Triple values. Term encoding is hand-rolled to +the N-Triples grammar: rdflib's term.n3() emits Turtle-style forms +(numeric shorthand, unescaped newlines) that are not valid in +line-oriented N-Quads. +""" + +import re + +import rdflib + +from trustgraph.schema import IRI, LITERAL +from trustgraph.api.types import Triple + +# RDF-star quoted triples have no standard N-Quads encoding; they are +# skipped with a count so callers can surface the omission. + +# N-Triples string-literal escapes (ECHAR): backslash first, then the rest. +_ESCAPES = [ + ("\\", "\\\\"), + ('"', '\\"'), + ("\n", "\\n"), + ("\r", "\\r"), + ("\t", "\\t"), +] + +# Characters the IRIREF production cannot carry: controls/space plus the +# explicitly forbidden set. One compiled scan keeps this off the profile +# for large exports (it runs per term). +_BAD_IRI = re.compile(r'[\x00-\x20<>"{}|^\x60]') + + +def _escape_literal(value): + for raw, esc in _ESCAPES: + value = value.replace(raw, esc) + return value + + +def _encode_iri(iri): + """, or None for values the grammar cannot carry.""" + if not iri or _BAD_IRI.search(iri): + return None + return f"<{iri}>" + + +def encode_term(term, is_object=False): + """Encode one wire-format term dict for an N-Quads line. + + :param is_object: literals are only valid in object position; + subjects and predicates must be IRIs (bnodes never appear on + the wire). + :returns: encoded string, or None when the term can't be represented. + """ + if term is None: + return None + + t = term.get("t", "") + + if t == IRI: + return _encode_iri(term.get("i", "")) + + if t == LITERAL and is_object: + value = _escape_literal(term.get("v", "")) + language = term.get("l") + datatype = term.get("d") + if language: + return f'"{value}"@{language}' + if datatype: + dt = _encode_iri(datatype) + if dt is None: + return None + return f'"{value}"^^{dt}' + return f'"{value}"' + + # literals outside object position, RDF-star, unknown types + return None + + +def triple_to_nquad(triple, graph_encoded): + """One wire-format triple dict -> an N-Quads line, or None to skip. + + :param triple: {"s": term, "p": term, "o": term} wire dict + :param graph_encoded: pre-encoded string + """ + s = encode_term(triple.get("s")) + p = encode_term(triple.get("p")) + o = encode_term(triple.get("o"), is_object=True) + if s is None or p is None or o is None: + return None + return f"{s} {p} {o} {graph_encoded} .\n" + + +def serialize_nquads(batches, graph_iri, out): + """Write wire-format triple batches to a text file-like as N-Quads. + + :param batches: iterable of lists of wire triple dicts + (e.g. triples_query_stream output) + :param graph_iri: graph name for every quad (str) + :param out: text-mode file-like with .write() + :returns: (written, skipped) counts + """ + g = _encode_iri(graph_iri) + if g is None: + raise ValueError(f"graph IRI not representable in N-Quads: {graph_iri!r}") + written = 0 + skipped = 0 + for batch in batches: + for t in batch: + line = triple_to_nquad(t, g) + if line is None: + skipped += 1 + else: + out.write(line) + written += 1 + return written, skipped + + +def parse_nquads(data): + """Parse N-Quads bytes back into api Triple values. + + Terms are stringified with str(), the same convention tg-load-knowledge + uses, so values survive the store round trip unchanged. The whole + member is materialized in memory (bundles are bounded by + --triples-limit at export); line-streaming is a possible follow-up. + + :param data: N-Quads bytes (one bundle member) + :returns: list of Triple + """ + ds = rdflib.Dataset() + ds.parse(data=data.decode("utf-8"), format="nquads") + return [ + Triple(s=str(s), p=str(p), o=str(o)) + for s, p, o, _g in ds.quads((None, None, None, None)) + ] From 08f69007c91046e9a704c1fde734a0baa6ad2ab7 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 6 Jul 2026 10:47:24 +0100 Subject: [PATCH 4/4] fix: handle missing time field in library API document/processing responses (#1028) Documents and processing records without a stored time field caused a KeyError crash in get_documents and list_children. Use optional access so time resolves to None instead of failing. --- trustgraph-base/trustgraph/api/library.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/trustgraph-base/trustgraph/api/library.py b/trustgraph-base/trustgraph/api/library.py index b3506bb7..1e87ed6e 100644 --- a/trustgraph-base/trustgraph/api/library.py +++ b/trustgraph-base/trustgraph/api/library.py @@ -363,7 +363,7 @@ class Library: return [ DocumentMetadata( id = v["id"], - time = datetime.datetime.fromtimestamp(v["time"]), + time = datetime.datetime.fromtimestamp(v["time"]) if "time" in v else None, kind = v["kind"], title = v.get("title", ""), comments = v.get("comments", ""), @@ -678,7 +678,7 @@ class Library: ProcessingMetadata( id = v["id"], document_id = v["document-id"], - time = datetime.datetime.fromtimestamp(v["time"]), + time = datetime.datetime.fromtimestamp(v["time"]) if "time" in v else None, flow = v["flow"], collection = v["collection"], tags = v["tags"], @@ -983,7 +983,7 @@ class Library: return [ DocumentMetadata( id=v["id"], - time=datetime.datetime.fromtimestamp(v["time"]), + time=datetime.datetime.fromtimestamp(v["time"]) if "time" in v else None, kind=v["kind"], title=v["title"], comments=v.get("comments", ""),