Compare commits

..

63 commits

Author SHA1 Message Date
cybermaggedon
0374098ee9
fix: lazily retry signing key fetch on JWT authentication (#1032)
If the gateway started before IAM was ready and exhausted its startup
retries, _signing_public_pem stayed None and all JWT authentication
permanently 401'd with no recovery path short of a restart. Now
_verify_jwt retries the fetch once on demand, so the gateway self-heals
as soon as IAM becomes available.
2026-07-07 16:00:34 +01:00
Keshav
f76f2abe08
feat: make fetch timeout for chunks configurable (#1031)
* tests: processor accepts default and overrides of fetch_chunk_timeout

* feat: configure fetch timeout for chunks from librarian
2026-07-07 15:55:33 +01:00
Sunny Yang
2bdc930b2a
feat: hybrid retrieval (BM25 + vector RRF fusion) for document-RAG (#875) (#1030)
Adds a sparse keyword retrieval path beside the existing vector path in
document-RAG, fused by weighted Reciprocal Rank Fusion on chunk_id, behind
--retrieval-mode (vector | keyword | hybrid, default vector).

The keyword index is a new pluggable service (KeywordIndexService /
KeywordIndexClientSpec); the first backend is SQLite FTS5, consuming Chunk
messages off the ingestion stream and answering BM25 queries from one
process, since the index is a single local file. Query text is sanitized
into per-term quoted phrases (raw text is not valid FTS5 syntax), which
also makes dotted clause numbers and error codes exact-match without a
trigram index. Indexes are scoped per (workspace, collection) and dropped
on collection deletion.

The keyword-index client spec is only registered when the sparse path is
enabled, so existing flow definitions without keyword-index queues are
untouched; with retrieval_mode=vector the retrieval path is unchanged. In
hybrid mode a keyword-path failure degrades to vector-only.
2026-07-07 12:54:02 +01:00
Cyber MacGeddon
e5206bddd0 Merge branch 'master' into release/v2.7 2026-07-06 10:54:01 +01:00
Cyber MacGeddon
5a1b42ab33 Merge branch 'release/v2.6' 2026-07-06 10:53:47 +01:00
cybermaggedon
a84cffa485 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.
2026-07-06 10:48:53 +01:00
Sunny Yang
0b6bae02f5 feat: export/import workspace knowledge in .tgx bundles (#877 Phase 2) (#1024)
* 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).
2026-07-06 10:48:41 +01:00
cybermaggedon
d5f3b6d9f6
feat: add structured audit event system (#1027)
Add a complete audit event pipeline that emits structured, machine-
parseable events for every gateway request and IAM decision.

Schema and publisher:
- AuditEvent dataclass and notify-class queue (audit_events_queue)
- AuditPublisher utility: fire-and-forget emission with envelope
  (schema_version, event_id, event_type, timestamp, producer)
- New request_id and client_ip fields on IamRequest for correlation

Gateway (gateway.request events):
- aiohttp middleware assigns request_id, captures timing/status/sizes
  and emits an event after every HTTP request completes
- IamAuth.authenticate annotates the request with identity
- Main endpoint handlers annotate capability and workspace
- request_id and client_ip forwarded to IAM on authenticate/authorise

IAM service (iam.authenticate, iam.authorise, iam.management events):
- Emits iam.authenticate for resolve-api-key, login, anonymous auth
- Emits iam.authorise for authorise and authorise-many decisions
- Emits iam.management for user/workspace/key mutations
- All events include request_id for correlation with gateway events

Design: events land on a pub/sub notify topic — non-persistent,
per-subscriber delivery. If no audit consumer is deployed, events
are silently discarded. Storage, retention, and alerting are
consumer-side concerns outside this boundary.

Added unit tests for the publisher and gateway middleware, unit
tests for IAM audit emission, and a contract test for the AuditEvent
schema.

Tech spec: docs/tech-specs/audit-events.md
2026-07-06 10:47:49 +01:00
cybermaggedon
08f69007c9
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.
2026-07-06 10:47:24 +01:00
Sunny Yang
cfbd5b9079
feat: export/import workspace knowledge in .tgx bundles (#877 Phase 2) (#1024)
* 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).
2026-07-06 10:45:23 +01:00
cybermaggedon
12ea6d46bc
maint: Open 2.7 branch (#1026)
Update package definitions and CI pipeline to refer to 2.7.x
2026-07-06 09:19:12 +01:00
Jack Colquitt
4ac6bc0986
Update section titles and video links in README (#1025) 2026-07-05 07:39:29 -07:00
Jack Colquitt
8d630ef204
Update README with access request link (#1023) 2026-07-03 11:43:56 -07:00
Jack Colquitt
d2a4abb7c5
Revise README with new TrustGraph context description (#1022)
Updated the README to enhance the description of TrustGraph and its features.
2026-07-03 11:40:52 -07:00
cybermaggedon
68e816e65c
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
2026-07-03 15:51:04 +01:00
Sunny Yang
76c4763b9b
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/<type>/), 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.
2026-07-03 15:16:39 +01:00
Cyber MacGeddon
508d0bb5c1 Merge branch 'release/v2.6' 2026-07-03 13:45:02 +01:00
cybermaggedon
c05296376e
fix: remove test import (#1017)
Convention in the tests is to just import the libraries as production
code would. This is fragile, and could possibly be used to inject
malicious code in the CI environment.
2026-07-03 13:44:09 +01:00
YingzuoLiu
f04ae5331d
Add diversity-aware selection after Document-RAG reranking (#1014)
* Add Document-RAG diversity selection helper

* Add optional MMR diversity selection after reranking

* Fix Document-RAG diversity test method signatures
2026-07-03 13:35:42 +01:00
cybermaggedon
db7fdbc652
feat: direction-aware reranker text in GraphRAG hop-and-filter (#1016)
The reranker document text now reflects the traversal direction,
showing only the new information relative to the frontier entity:
- From S (subject is frontier): text = "{predicate} {object}"
- From O (object is frontier): text = "{subject} {predicate}"
- From P (predicate is frontier): text = "{subject} {object}"

This eliminates duplicate reranker texts when traversing inward
from shared object nodes (e.g. 18 CPUs all producing identical
"hasSubcategory Processors" text when the subject was dropped).

execute_batch_triple_queries now returns (triple, direction)
tuples so hop_and_filter can select the appropriate text format.

Updates tech spec to document the direction-aware approach.
Adds unit tests for direction tracking and reranker text
construction.
2026-07-02 21:14:47 +01:00
Cyber MacGeddon
4aaa1ce915 Merge branch 'release/v2.6' into master to keep sync. 2026-07-02 14:58:37 +01:00
cybermaggedon
9cf7dcb578
fix: wire variant into remaining streaming integration test mocks (#1013)
Three more streaming tests were missing _wire_variant after the
async for change in create_completion_stream.
2026-07-02 11:14:54 +01:00
Sunny
6c9a545a06
feat: add cross-encoder reranking to Document-RAG with two-limit control (#878) (#1011)
Wire the FlashRank reranker subsystem from #1005 into Document-RAG: after
vector retrieval, over-fetch a wider candidate pool, rerank with the
cross-encoder, and keep the top doc_limit chunks for synthesis.

Per maintainer review, the fetch and select sizes are two caller-controlled
limits rather than one internal heuristic:

- doc_limit:   chunks selected into the synthesis prompt (unchanged meaning).
- fetch_limit: candidate pool pulled from the vector store before reranking.
  0 = derive (OVERFETCH_FACTOR x doc_limit); values below doc_limit are
  raised to it. Lets the caller control how hard the reranker has to work.

Details:
- schema: DocumentRagQuery.fetch_limit (additive, backward compatible).
- document_rag.py / rag.py: fetch_limit resolved in the processor (mirrors
  doc_limit); the core applies the heuristic default and derives synthesis
  provenance from the chunk-selection focus when reranking ran.
- provenance: tg:ChunkSelection focus stage (mirrors tg:EdgeSelection).
- request translator + client SDKs + CLI: fetch-limit / --fetch-limit,
  threaded exactly like doc_limit and the GraphRAG limits.
- tests: no-op identity, over-fetch/narrow, explicit fetch_limit, heuristic
  default, floor-at-doc_limit, provenance lineage, cross-repo topic wiring.

Reranking is skipped byte-identically when no reranker role is wired.
Requires the companion trustgraph-templates change wiring the reranker
topics into the document-rag flow (mirrors #279 for GraphRAG).
2026-07-02 09:50:13 +01:00
cybermaggedon
f18d48dc39
fix: simplify dashscope variant and route API calls through variants (#1012)
Replace the client.post()/httpx bypass with standard SDK extra_body,
confirmed working against DashScope. Make DashScope the base variant
with Qwen as a subclass alias. Route all API calls through variant
create_completion/create_completion_stream methods.
2026-07-02 09:12:55 +01:00
cybermaggedon
6887076ce0
feat: add dashscope variant for Alibaba Cloud DashScope API (#1010)
DashScope uses enable_thinking as a top-level parameter rather than
inside extra_body as the Qwen docs suggest.
2026-07-01 16:50:47 +01:00
cybermaggedon
55e2a2a3ce
feat: add guided macOS installer and developer install guide (#1003)
Interactive bash installer (install_trustgraph.sh) that detects hardware,
recommends an LLM mode (OpenAI or Ollama), installs missing prerequisites
via Homebrew, sets up a Python venv, runs the test suite, generates a
deployment via npx @trustgraph/config, starts the Docker Compose stack,
health-checks the API gateway, and opens the Workbench UI.

Includes README.dev-install.md with usage documentation covering CLI
options, environment variables, LLM mode selection, non-interactive/CI
usage, uninstall, and troubleshooting. Currently macOS only.
2026-07-01 16:50:14 +01:00
cybermaggedon
11ca7c89c4
feat: add GLM (Zhipu AI) variant for OpenAI processor (#1009) 2026-07-01 16:20:43 +01:00
cybermaggedon
656ca430b9
fix: wire variant into text-completion integration test mocks (#1008)
Tests using MagicMock processors need the variant, thinking mode,
and _build_kwargs/_extract_content methods bound to work with the
new variant-based API kwargs construction.
2026-07-01 15:40:23 +01:00
cybermaggedon
f20b50cfb2
feat: add API variant profiles and thinking support to OpenAI processor (#1007)
Add a --variant flag (openai, deepseek, qwen, mistral, llama) that
encapsulates provider-specific API differences: output token parameter
names, thinking/reasoning toggles, temperature rules, and thinking
output extraction. Add --thinking flag (off, low, medium, high) to
control reasoning effort.
2026-07-01 14:48:32 +01:00
Jack Colquitt
04c5921687
Fix Discord link in README (#1006)
Updated Discord link in README.md
2026-06-30 19:39:33 -07:00
cybermaggedon
01cc8dbc64
feat: replace LLM edge scoring with cross-encoder reranker in GraphRAG (#1005)
Replace the three-prompt LLM scoring pipeline (kg-edge-scoring,
kg-edge-reasoning, kg-edge-selection) with a cross-encoder reranker
service backed by FlashRank. The new hop_and_filter() method performs
iterative graph traversal with semantic scoring at each hop, replacing
the previous follow_edges/get_subgraph approach.

- Add reranker service (trustgraph-base client/service, FlashRank processor)
- Add gateway dispatch for reranker via API and WebSocket
- Rewrite GraphRAG pipeline: hop_and_filter() with per-hop cross-encoder scoring
- Remove kg_prompt() and edge_score_limit from prompt client
- Update provenance: add tg:EdgeSelection type, tg:concept, tg:score predicates
- Update CLIs (tg-invoke-graph-rag, tg-show-explain-trace) for new metadata
- Add tg-invoke-reranker CLI tool
- Add tech spec and UX developer guidance
- Update all unit and integration tests
2026-06-30 14:36:37 +01:00
corvus-0x
1aa9549912
feat: make bootstrapper initialiser timeouts configurable (#999)
* feat: make bootstrapper initialiser timeouts configurable

DefaultFlowStart and WorkspaceInit hardcoded the request timeouts for
their flow-svc and IAM calls, leaving operators no way to tune them for
high-latency environments (#874).

Expose them as constructor parameters threaded through the existing
initialiser `params:` mechanism, defaulting to the current values so
behaviour is unchanged unless explicitly overridden:

- DefaultFlowStart: list_timeout=10 (list-flows), start_timeout=30 (start-flow)
- WorkspaceInit: iam_timeout=10 (create-workspace)

Add unit tests for the defaults, override storage, and that configured
values reach the underlying request calls.

* test: mark async bootstrap test with @pytest.mark.asyncio

Addresses review feedback on PR #999: add the explicit
@pytest.mark.asyncio decorator to test_run_forwards_configured_timeouts
so it does not rely on asyncio_mode=auto and stays consistent with the
rest of the suite.
2026-06-30 09:37:22 +01:00
cybermaggedon
5cb4f83afa
fix: list-my-workspaces permissions were broken (#1002)
list-my-workspaces has AUTHENTICATED scope, so anyone is permitted
to run the operation.  No specific permission grant is needed.
2026-06-29 09:13:05 +01:00
cybermaggedon
0a828379be
feat: global usernames and rename workspace to default_workspace (#1001)
Users are global entities, not scoped to workspaces. This change:

Track A — Global usernames:
- Change iam_users_by_username to PRIMARY KEY (username), removing
  workspace from the lookup key
- Login looks up username globally, no workspace required
- Username uniqueness is enforced globally, not per-workspace
- Login -w now overrides the JWT workspace (session workspace)
  rather than selecting which user registry to search

Track B — Rename workspace to default_workspace:
- UserRecord.workspace → UserRecord.default_workspace
- Identity.workspace → Identity.default_workspace
- JWT claim "workspace" → "default_workspace"
- IamResponse.resolved_workspace → resolved_default_workspace
- WebSocket auth-ok frame field → default_workspace
- Socket clients read default_workspace from auth-ok
- _user_record_to_dict wire key → default_workspace
- CLI help text and output updated throughout
- Test files updated for renamed fields
2026-06-25 16:34:31 +01:00
cybermaggedon
16f8cfd972
fix: use envelope workspace for mux authorisation, not inner request body (#1000)
The mux was extracting the authorisation resource workspace from the
inner request body via registry extractors. But workspace-scoped
services (config, flow, librarian, etc.) receive workspace from the
queue identity, not the message body — the inner workspace field is
a dead field that no service handler reads.

This caused access-denied errors when the inner body's workspace
(e.g. CLI default "default") disagreed with the caller's assigned
workspace, even though the envelope workspace was correct.

Fix: resolve workspace from the envelope only. Split the non-flow
authorisation path by resource level — WORKSPACE ops use the envelope
workspace directly; SYSTEM ops (IAM) still use registry extractors
since they legitimately read operation-specific body fields.
2026-06-25 13:44:57 +01:00
Cyber MacGeddon
a3df4f62bb Merge branch 'master' into release/v2.6 2026-06-22 21:20:29 +01:00
cybermaggedon
09b8a1d347
feat: fine-grained capabilities and enterprise IAM schema extensions (#996)
Split coarse gateway capabilities into fine-grained variants to
support per-operation access control in the enterprise IAM regime.
Add additive schema fields for enterprise group and grant management.

Capability split (gateway registry):
- graph:read -> triples:read, sparql:read, graph-rag:read,
  graph-embeddings:read
- graph:write -> triples:write, graph-embeddings:write,
  entity-contexts:write
- documents:read -> documents:read, document-rag:read,
  document-embeddings:read, entity-contexts:read
- documents:write -> documents:write, document-embeddings:write
- rows:read -> rows:read, nlp-query:read, structured-query:read,
  row-embeddings:read

OSS role definitions expanded to include all new fine-grained
capability names — no behavioral change for OSS deployments.

Schema additions (IamRequest):
- group_id, member_type, member_id for group membership operations
- group (GroupInput), grant (GrantInput) for create/update payloads
- Decoder now handles capability, resource_json, parameters_json,
  authorise_checks (previously missing from translator)

Schema additions (IamResponse):
- group_json, groups_json, members_json, grants_json,
  effective_permissions_json for enterprise operation responses
- Encoder now emits authorise decision fields

Gateway registry:
- 16 enterprise IAM operations registered (create-group,
  add-group-member, add-user-grant, etc.) under iam:admin capability
2026-06-22 20:23:34 +01:00
Jack Colquitt
fa264ded46
Update section titles for Holonic Context Graph (#995) 2026-06-18 19:56:44 -07:00
Jack Colquitt
cae931409a
Update TrustGraph description in README (#994)
Clarified the description of TrustGraph's capabilities and API integrations.
2026-06-18 19:46:25 -07:00
Jack Colquitt
6b0475e315
Revise README for clarity on TrustGraph features (#993)
Updated the README to clarify the concept of holons and the functionality of TrustGraph. Improved the structure and flow of information regarding context management and agent explainability.
2026-06-18 19:42:16 -07:00
Jack Colquitt
cb0ad1a450
Change video link in README (#992)
Updated video source link in README.md.
2026-06-17 17:52:30 -07:00
Jack Colquitt
fc0ecc770a
Format terms as code in README.md (#991) 2026-06-17 16:53:22 -07:00
Jack Colquitt
345da375b1
Document Workspaces, Collections, and Flows in README (#990)
Added section on Workspaces, Collections, and Flows to explain the organizational structure of TrustGraph.
2026-06-17 16:48:22 -07:00
Jack Colquitt
0ba1eeeda0
Enhance README with token consumption details (#989)
Added a note about reducing token consumption in context management.
2026-06-17 16:27:16 -07:00
Jack Colquitt
eb1e38d7d0
Add hyperlink to 'holon' in README.md (#988) 2026-06-17 16:13:41 -07:00
Jack Colquitt
b8770a6005
Update README with new context and features (#987) 2026-06-17 16:08:22 -07:00
Jack Colquitt
28802a644a
Update license badge in README.md (#986) 2026-06-11 20:45:39 -07:00
cybermaggedon
8797d9d9ff feat: per-caller Bearer token auth and new query tools for MCP server (#984)
Replace the broken GATEWAY_SECRET auth (token was sent as a query
parameter, silently ignored by the gateway) with end-to-end Bearer
token forwarding.  Each MCP caller gets a dedicated WebSocket
authenticated via the gateway's in-band first-frame protocol, with
whoami verification on first connect.

Also fix and extend the tool surface:
- embeddings: accept list of texts (was single string)
- triples_query: use Term wire format with compact keys (was legacy
  Value format), add collection and graph parameters
- sparql_query: new tool for SPARQL SELECT/ASK/CONSTRUCT/DESCRIBE
- graphql_query: new tool for structured data (rows) GraphQL queries
- all tools: add optional workspace parameter
2026-06-10 14:11:49 +01:00
cybermaggedon
627c669097
feat: per-caller Bearer token auth and new query tools for MCP server (#984)
Replace the broken GATEWAY_SECRET auth (token was sent as a query
parameter, silently ignored by the gateway) with end-to-end Bearer
token forwarding.  Each MCP caller gets a dedicated WebSocket
authenticated via the gateway's in-band first-frame protocol, with
whoami verification on first connect.

Also fix and extend the tool surface:
- embeddings: accept list of texts (was single string)
- triples_query: use Term wire format with compact keys (was legacy
  Value format), add collection and graph parameters
- sparql_query: new tool for SPARQL SELECT/ASK/CONSTRUCT/DESCRIBE
- graphql_query: new tool for structured data (rows) GraphQL queries
- all tools: add optional workspace parameter
2026-06-10 14:10:43 +01:00
cybermaggedon
8b0619e5d8
Bump version numbers to 2.6 (#983) 2026-06-09 20:03:14 +01:00
cybermaggedon
e3f9f8c357
Merge pull request #982 from trustgraph-ai/master
master -> release/v2.6
2026-06-09 19:46:50 +01:00
Cyber MacGeddon
81d57826c8 Merge branch 'release/v2.5' 2026-06-09 19:43:31 +01:00
Jacob Molz
79d7ef6a90 fix: reject invalid PDF decoder input (#977) 2026-06-09 16:37:39 +01:00
Jacob Molz
28a51c244f
fix: reject invalid PDF decoder input (#977) 2026-06-09 16:37:10 +01:00
Cyber MacGeddon
fa5ebe2393 Merge branch 'release/v2.5' 2026-06-09 16:34:20 +01:00
cybermaggedon
e1c9351454
fix: update row query tests to mock async_execute_paged and async_scan (#979)
The query service now uses async_execute_paged (indexed path) and
async_scan (scan path) instead of async_execute. Tests were mocking
the old function, causing them to hang indefinitely.
2026-06-09 16:29:32 +01:00
cybermaggedon
dbc21c0bb9
fix: structured data query and auth fixes (#978)
- Pass auth token to schema discovery and descriptor generation in
  tg-load-structured-data, fixing 401 errors with IAM enabled
- Fix row query pagination: replace single-page async_execute with
  async_scan that streams pages and applies filters without
  materialising the full result set (OOM on large datasets)
- Add missing filter operators (not, startsWith, endsWith, not_in)
  to row query post-filter matching
- Fall back to scan path when an indexed field is queried with an
  empty string value, since empty index values are not stored
- Revert top-level indexes array support — the current table schema
  overwrites rows with duplicate index values, so only primary_key
  fields are safe to index until the schema is redesigned
2026-06-08 15:22:11 +01:00
cybermaggedon
08bfec1539
fix: wire replication params through YAML/params path for Cassandra and Qdrant (#976)
resolve_cassandra_config did not accept replication_factor as a kwarg,
so cassandra_replication_factor from YAML params was silently ignored
by all 6 callers. Add the kwarg and pass it from every caller.

Same fix for Qdrant: 3 writers now pass qdrant_replication_factor and
qdrant_shard_number from params.

Add tests covering the params path for both helpers.
2026-06-04 12:36:36 +01:00
cybermaggedon
4913f8c2eb
feat: data store replication configuration and TLS upgrade (#975)
- Add centralised qdrant_config.py helper with env-var fallback for
  QDRANT_URL, QDRANT_API_KEY, QDRANT_REPLICATION_FACTOR, QDRANT_SHARD_NUMBER
- Update all 6 Qdrant processors to use the helper; writers pass
  replication_factor and shard_number to create_collection
- Fix hardcoded Cassandra replication_factor=1 in cassandra_kg.py,
  write.py, and sparql_cassandra.py to respect CASSANDRA_REPLICATION_FACTOR
- Upgrade Cassandra TLS from deprecated PROTOCOL_TLSv1_2 to
  ssl.create_default_context() across all connectors
2026-06-04 11:49:29 +01:00
Jack Colquitt
97453d9b83
Change project title to 'The semantic deployment platform' (#968)
Updated the project title in the README.
2026-06-01 14:08:30 -07:00
Jack Colquitt
6dfa47aac8
Revise README for semantic infrastructure terminology (#962)
Updated the README to reflect changes in terminology and improve clarity regarding the platform's features.
2026-05-30 17:07:19 -07:00
Cyber MacGeddon
dcee842455 Merge branch 'release/v2.5' 2026-05-28 11:26:43 +01:00
cybermaggedon
36eadbda3a
Merge pull request #953 from trustgraph-ai/release/v2.5
release/v2.5 -> master
2026-05-26 15:01:44 +01:00
153 changed files with 12718 additions and 2507 deletions

View file

@ -22,7 +22,7 @@ jobs:
uses: actions/checkout@v3
- name: Setup packages
run: make update-package-versions VERSION=2.5.999
run: make update-package-versions VERSION=2.7.999
- name: Setup environment
run: python3 -m venv env

218
README.dev-install.md Normal file
View file

@ -0,0 +1,218 @@
# TrustGraph Developer Install Guide
A guided installer that gets TrustGraph running locally in a single
command. It detects your hardware, recommends an LLM backend, installs
missing prerequisites, runs the test suite, generates a compose deployment,
starts the stack, and opens the Workbench UI.
> **macOS only.** This installer has only been tested on macOS. If you are
> on Linux or Windows, use the standard docker-compose / podman-compose
> installation instructions instead.
## Quick start
```bash
./install_trustgraph.sh
```
The installer walks you through each step interactively. When it finishes,
the Workbench UI opens at `http://localhost:8888` and the API gateway is
available at `http://localhost:8088/`.
## Prerequisites
The installer checks for these and offers to install any that are missing
(via Homebrew):
- **Python 3** with venv support
- **Node.js / npx** (drives the `@trustgraph/config` deployment generator)
- **Docker** (with Compose) or **Podman** (with podman-compose)
- **curl** and **unzip**
- **Ollama** (only if you choose local LLMs)
The installer can also launch Docker Desktop or the Ollama app for you if
they are installed but not running.
## What the installer does
1. **Detects hardware** -- OS, architecture, CPU cores, memory, and GPU.
2. **Recommends an LLM mode** -- `ollama` for machines with >= 16 GB RAM and
a GPU or >= 8 cores; `openai` otherwise.
3. **Collects configuration** -- API key, LLM provider, model choices,
install directory. Answers are saved to
`<install-dir>/trustgraph-installer.env` and reused on subsequent runs.
4. **Checks and installs prerequisites** -- Python, Node/npx, Docker or
Podman, Ollama (if selected).
5. **Downloads Ollama models** (if using Ollama) -- chat model
(`granite4:350m` by default) and embeddings model (`mxbai-embed-large`).
6. **Creates a Python venv** and installs the local TrustGraph packages into
it, along with NLTK data and tiktoken caches.
7. **Runs the full pytest suite** against the local source tree.
8. **Runs `npx @trustgraph/config`** -- the existing interactive config
wizard that produces a `deploy.zip` with a compose file.
9. **Starts the compose stack** and waits for the API gateway to respond.
10. **Bootstraps IAM** and verifies the API key authenticates.
11. **Opens the Workbench UI** in your default browser.
## Command-line options
| Option | Description |
|---|---|
| `--install-dir PATH` | Directory for deployment files (default: `./trustgraph-deploy`) |
| `--api-url URL` | API gateway URL for health checks (default: `http://localhost:8088/`) |
| `--ui-url URL` | Workbench UI URL to open (default: `http://localhost:8888`) |
| `--use-existing-compose FILE` | Skip config generation and start this compose file directly |
| `--skip-tests` | Do not run the pytest suite |
| `--no-launch` | Do not open the Workbench UI at the end |
| `--non-interactive` | Accept all defaults without prompting |
| `--yes` | Auto-accept confirmation prompts |
| `--fresh` | Remove installer-managed files before generating a new deployment |
| `--remove-all` | Uninstall: stop containers, remove compose volumes, delete installer files |
| `--dry-run` | Print detected hardware and planned defaults, then exit |
| `-h`, `--help` | Show the built-in help text |
## Environment variables
These override the interactive prompts when set:
| Variable | Purpose |
|---|---|
| `TRUSTGRAPH_TOKEN` | Admin/bootstrap API key (must start with `tg_`) |
| `TRUSTGRAPH_URL` | API gateway URL |
| `TRUSTGRAPH_UI_URL` | Workbench UI URL |
| `OPENAI_TOKEN` | OpenAI-compatible API key |
| `OPENAI_BASE_URL` | OpenAI-compatible base URL |
| `OLLAMA_HOST` / `OLLAMA_BASE_URL` | Ollama service URL |
| `OLLAMA_MODEL` | Ollama chat model (default: `granite4:350m`) |
| `OLLAMA_EMBEDDINGS_MODEL` | Ollama embeddings model (default: `mxbai-embed-large`) |
| `TG_INSTALL_DIR` | Override the install directory |
| `TG_VENV_DIR` | Override the Python venv location |
| `TG_NLTK_DATA_DIR` | Override the NLTK data directory |
| `TIKTOKEN_CACHE_DIR` | Override the tiktoken cache directory |
| `TG_HEALTH_TIMEOUT` | Seconds to wait for the API gateway (default: 240) |
## Choosing an LLM mode
### OpenAI (or any OpenAI-compatible provider)
Best when you already have an API key or are running against a remote
endpoint. The installer asks for a base URL and an API key.
```bash
OPENAI_TOKEN=sk-... ./install_trustgraph.sh
```
### Ollama (local models)
Best on machines with enough RAM to run a small model. The installer detects
locally installed Ollama models and offers to pull missing ones. It uses
`host.docker.internal` so the Docker containers can reach the host-side
Ollama service.
```bash
./install_trustgraph.sh # choose "ollama" when prompted
```
### None
Start the platform without an LLM. Agent and RAG features will not work
until you configure one later through the Workbench.
## Saved answers and re-running
The installer saves your answers to
`<install-dir>/trustgraph-installer.env`. On the next run it loads those
answers as defaults, so you can re-run with a single Enter through each
prompt.
To start completely fresh:
```bash
./install_trustgraph.sh --fresh
```
This stops any running containers (keeping Docker volumes), removes
installer-managed files, and re-runs the full flow.
## Using an existing compose file
If you already have a compose file from the config tool or another source:
```bash
./install_trustgraph.sh --use-existing-compose path/to/docker-compose.yaml
```
This skips the config wizard and `npx` prerequisite check, and goes straight
to starting the stack.
## Non-interactive / CI usage
```bash
TRUSTGRAPH_TOKEN=tg_my-token \
OPENAI_TOKEN=sk-... \
./install_trustgraph.sh --non-interactive --yes --skip-tests
```
In non-interactive mode the installer uses defaults for every prompt. Pair
with `--yes` to auto-accept confirmation prompts and `--skip-tests` if you
want a faster run.
## Dry run
Preview what the installer would do without making any changes:
```bash
./install_trustgraph.sh --dry-run
```
This prints the detected hardware, recommended LLM mode, and planned
install paths, then exits.
## Uninstalling
```bash
./install_trustgraph.sh --remove-all
```
This stops containers, removes compose-managed volumes, and deletes
installer-managed files (venv, deploy output, logs, saved answers). It does
**not** remove Docker/Podman itself, container images, Ollama, or Ollama
models.
## Troubleshooting
### Logs
All long-running operations write logs to `<install-dir>/logs/`. Key files:
- `pytest.log` -- test suite output
- `compose-up.log` -- docker compose output
- `iam-bootstrap.log` -- IAM bootstrap output
- `ollama-pull-*.log` -- Ollama model downloads
- `pip-*.log` -- Python package installs
- `brew-install-*.log` -- Homebrew installs
### API key rejected after reinstall
If the API gateway returns 401/403 with your saved key, the compose volumes
likely contain IAM data from a previous install with a different key. Run:
```bash
./install_trustgraph.sh --remove-all
./install_trustgraph.sh
```
This clears the old volumes and starts fresh.
### Ollama not reachable from containers
The Ollama base URL should use `host.docker.internal` instead of
`localhost` so that containers running in Docker Desktop can reach the
host-side Ollama service. The installer sets this automatically; if you
override `OLLAMA_HOST`, make sure the URL is reachable from inside the
container network.
### Docker daemon not running
The installer detects Docker Desktop and offers to start it. If that
doesn't work, start Docker Desktop manually and re-run the installer.

293
README.md
View file

@ -3,52 +3,104 @@
<img src="TG-fullname-logo.svg" width=100% />
[![PyPI version](https://img.shields.io/pypi/v/trustgraph.svg)](https://pypi.org/project/trustgraph/) [![License](https://img.shields.io/github/license/trustgraph-ai/trustgraph?color=blue)](LICENSE) ![E2E Tests](https://github.com/trustgraph-ai/trustgraph/actions/workflows/release.yaml/badge.svg)
[![PyPI version](https://img.shields.io/pypi/v/trustgraph.svg)](https://pypi.org/project/trustgraph/) ![License](https://img.shields.io/badge/license-Apache%202.0-blue) ![E2E Tests](https://github.com/trustgraph-ai/trustgraph/actions/workflows/release.yaml/badge.svg)
[![Discord](https://img.shields.io/discord/1251652173201149994
)](https://discord.gg/sQMwkRz5GX) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/trustgraph-ai/trustgraph)
[**Website**](https://trustgraph.ai) | [**Docs**](https://docs.trustgraph.ai) | [**YouTube**](https://www.youtube.com/@TrustGraphAI?sub_confirmation=1) | [**Configuration Terminal**](https://config-ui.demo.trustgraph.ai/) | [**Discord**](https://discord.gg/sQMwkRz5GX) | [**Blog**](https://blog.trustgraph.ai/subscribe)
[**Website**](https://trustgraph.ai) | [**Docs**](https://docs.trustgraph.ai) | [**YouTube**](https://www.youtube.com/@TrustGraphAI?sub_confirmation=1) | [**Configuration Terminal**](https://config-ui.demo.trustgraph.ai/) | [**Discord**](https://discord.gg/yUWRkfbD) | [**Blog**](https://blog.trustgraph.ai/subscribe)
### The Anti-Palantir
<a href="https://trendshift.io/repositories/17291" target="_blank"><img src="https://trendshift.io/api/badge/repositories/17291" alt="trustgraph-ai%2Ftrustgraph | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
# The agent runtime platform
**Open Source · Open Standards · Total Transparency**
[**Request Access to the Playground Preview**](https://docs.google.com/forms/d/e/1FAIpQLSeTnF22ZjUP20FWV--VvS5606x-5cOvnKty6AqcPdtlnPuqbQ/viewform)
</div>
TrustGraph is an agent runtime platform built around context graphs — structured, queryable representations of your domain knowledge that ground every agent query in verified, explainable facts in private deployments with sovereign control. The platform is the full stack for agentic systems: context graphs, memory, retrieval, orchestration, and inference for precision-critical agent workloads.
---
The platform:
- [x] Multi-model and multimodal database system
- [x] Tabular/relational, key-value
- [x] Document, graph, and vectors
- [x] Images, video, and audio
- [x] Context Graph engine
- [x] Automated entity and relationship extraction
- [x] Ontology-driven graph construction
- [x] Graph-grounded retrieval for explainable outputs
- [x] Automated data ingest and loading
- [x] Quick ingest with semantic similarity retrieval
- [x] Ontology structuring for precision retrieval
- [x] Out-of-the-box RAG pipelines
- [x] DocumentRAG
- [x] GraphRAG
- [x] OntologyRAG
- [x] 3D GraphViz for exploring context
- [x] Fully Agentic System
- [x] Single or Multi Agent
- [x] ReAct, Plan-then-Execute, and Supervisor patterns
- [x] MCP integration
- [x] Run anywhere
- [x] Deploy locally with Docker
- [x] Deploy in cloud with Kubernetes
- [x] Support for all major LLMs
- [x] API support for Anthropic, Cohere, Gemini, Mistral, OpenAI, and others
- [x] Model inferencing with vLLM, Ollama, TGI, LM Studio, and Llamafiles
- [x] Developer friendly
- [x] REST API [Docs](https://docs.trustgraph.ai/reference/apis/rest.html)
- [x] Websocket API [Docs](https://docs.trustgraph.ai/reference/apis/websocket.html)
- [x] Python API [Docs](https://docs.trustgraph.ai/reference/apis/python)
- [x] CLI [Docs](https://docs.trustgraph.ai/reference/cli/)
Write context once. Run agents anywhere. Own your data and the models.
Stop rebuilding context from scratch. TrustGraph treats context as a holon — a modular, independent whole that naturally snaps into a larger domain-wide intelligence layer. By deploying context as holonic context graphs, TrustGraph powers multi-tenant agent workflows, dramatically reduces token consumption, and aligns with semantic web standards (RDF, OWL, SKOS, SHACL). Version your context, share it across teams, and scale with full provenance.
## What TrustGraph Does
TrustGraph is a complete holonic context harness for all LLMs. It provides the full infrastructure layer underneath your agents: knowledge ingestion, structured storage, graph-grounded retrieval, agent orchestration, and a full LLM inferencing stack.
TrustGraph relies on absolutely no 3rd party services aside from optional API integrations to cloud-hosted LLMs. Whether you are using Anthropic's or OpenAI's API, or self-hosting Qwen3.7 via vLLM, TrustGraph handles it all with pre-built API connectors and a full LLM inferencing stack to enrich the models with a sovereign, private holonic system that grounds your agents in reality.
## The Problem: Why Agents Break
When you build an AI agent today, you spend most of your time fighting context:
- **RAG retrieves fragments, not meaning**. Chunks of text have no structure. Relationships between facts are invisible. Your agent guesses at the connections.
- **Context is disposable**. What the agent learned in one session is gone in the next. There is no persistent, structured knowledge layer underneath.
- **Answers aren't traceable**. You can't explain why the agent said what it said, which means you can't trust it in production.
- **Knowledge can't be reused**. You rebuild the same context pipelines for every new project, every new agent, every new environment.
These aren't retrieval problems. They are structural problems. Context needs to be organized, versioned, and composable — exactly the way software infrastructure is.
## The Solution: A Holonic Context System
The philosopher Arthur Koestler coined the word [holon](https://en.wikipedia.org/wiki/Holon_(philosophy)) to describe something that is simultaneously a whole in itself and a part of something larger. A fact is whole. It is also part of a domain. A domain is whole. It is also part of an organization's knowledge.
AI agents break down because this holonic structure is never built. Context gets shoved into flat text windows, scattered across vector stores, or hardwired into one-off prompts. Facts lose their relationships.
TrustGraph solves this by organizing your domain into holonic context graphs. Entities, relationships, and evidence are treated as first-class objects. Every agent query is grounded against these holons—marrying symbolic graph structures with vector embeddings. Every answer carries provenance. Every fact is traceable.
## Context Cores: Knowledge as a First-Class Citizen
A Context Core is the deployable unit of knowledge in TrustGraph. It packages everything an agent needs to reason reliably over a domain into a single, portable artifact.
### What's inside a Context Core
- **Ontology** — your domain schema and entity mappings
- **Holon** — entities, relationships, and supporting evidence
- **Embeddings** — vector indexes for fast semantic entry-point lookup
- **Provenance** — where every fact came from, when, and how it was derived
- **Retrieval policies** — traversal rules, freshness controls, authority ranking
Context Cores decouple what agents know from how agents are deployed. Build once. Run in Docker locally, Kubernetes in production, or on any cloud. Pin a version. Roll back. Promote across environments. This is context engineering — and it works because knowledge is finally treated like the infrastructure it is.
## Explainability: Trust Your Agents in Production
LLMs are black boxes, and traditional RAG makes it worse. When an agent pulls flat text chunks from a vector store, you have no idea how it connected those fragments to form an answer. You cannot ship agents to production if you can't explain why they said what they said.
### How TrustGraph makes agents explainable:
- **Traceable Reasoning Paths**: Instead of guessing at connections between text chunks, TrustGraph traverses explicit relationship paths in the holonic context graph. You can inspect exactly which entities, relationships, and sub-graphs were pulled into the LLM's context window to generate a given response.
- **Fact-Level Provenance**: Every node and edge in the graph carries strict provenance. When an agent makes a claim, you can trace it back to the exact source document, the time it was ingested, and the extraction method used to derive it.
- **No Black-Box Guesses**: By grounding the LLM in a structured, symbolic graph, you eliminate the hallucinations that occur when models are forced to infer relationships from unstructured text. If a fact isn't in the graph, the agent doesn't use it.
TrustGraph doesn't just give you answers - it gives you the receipt. Every fact is traceable, every connection is visible, and every output is verifiable.
## Workspaces, Collections, and Flows
TrustGraph has a [three-level system](https://docs.trustgraph.ai/overview/workspaces) for organizing and isolating knowledge.
A `Workspace` is the outermost boundary — a fully isolated tenancy scope where all data, users, configuration, and pipelines live independently from every other workspace. Isolation is structural: enforced at the pub/sub queue, storage, and API gateway layers, not by trusting a field in a message body.
Within a workspace, a `Collection` groups related holons, graph structures, embeddings, and documents together — think of it as a dedicated shelf in a library, scoped to a specific domain, project, or customer.
A `Flow` is a running data processing pipeline that defines how raw data moves through ingestion, extraction, structuring, and storage — the assembly line that turns documents into queryable knowledge. Together, the three layers let you run multiple isolated tenants on a single deployment, separate knowledge by domain within each tenant, and process that knowledge through fully configurable pipelines — all without restarting the system or rebuilding your infrastructure.
## The Full Stack
TrustGraph is not a wrapper around a graph database. It is the complete backend for production agentic systems.
- **Holonic context graph engine**: automated entity and relationship extraction, ontology-driven graph construction, graph-grounded retrieval for explainable outputs
- **Multi-model database**: tabular/relational, key-value, document, graph, vectors, images, video, and audio — all managed in Cassandra and S3-compatible Garage
- **Out-of-the-box RAG pipelines**: DocumentRAG, GraphRAG, and OntologyRAG ready to deploy
- **Fully agentic orchestration**: single or multi-agent, ReAct, Plan-then-Execute, Supervisor patterns, and MCP integration
- **3D Knowledge Explorer**: interactive graph visualization with BFS neighborhood extraction and edge pulse animation
- **Automated data ingest**: quick ingest with semantic similarity or ontology-structured precision retrieval
- **Run anywhere**: Docker/Podman locally, Kubernetes in the cloud
All major LLMs — Anthropic, Cohere, Gemini, Mistral, OpenAI, and more via API.
vLLM, Ollama, TGI, LM Studio, and Llamafiles for fully local inferencing.
Verified cloud deployments for Alibaba Cloud, AWS, Azure, GCP, OVHcloud, and Scaleway.
## No API Keys Required
@ -62,12 +114,12 @@ Everything else is included.
- [x] Managed Multi-model storage in [Cassandra](https://cassandra.apache.org/_/index.html)
- [x] Managed Vector embedding storage in [Qdrant](https://github.com/qdrant/qdrant)
- [x] Managed File and Object storage in [Garage](https://github.com/deuxfleurs-org/garage) (S3 compatible)
- [x] Managed High-speed Pub/Sub messaging fabric with [Pulsar](https://github.com/apache/pulsar)
- [x] Managed High-speed Pub/Sub messaging fabric with [Pulsar](https://github.com/apache/pulsar) or [RabbitMQ](https://www.rabbitmq.com/)
- [x] Complete LLM inferencing stack for open LLMs with [vLLM](https://github.com/vllm-project/vllm), [TGI](https://github.com/huggingface/text-generation-inference), [Ollama](https://github.com/ollama/ollama), [LM Studio](https://github.com/lmstudio-ai), and [Llamafiles](https://github.com/mozilla-ai/llamafile)
## Quickstart
There's no need to clone this repo, unless you want to build from source. TrustGraph is a fully containerized app that deploys as a set of Docker containers. To configure TrustGraph on the command line:
No need to clone the repo unless you are building from source. TrustGraph deploys as a set of Docker containers. Configure it on the command line in one step:
```
npx @trustgraph/config
@ -78,44 +130,39 @@ The config process will generate an app config that can be run locally with Dock
- Deployment instructions as `INSTALLATION.md`
<p align="center">
<video src="https://github.com/user-attachments/assets/2978a6aa-4c9c-4d7c-ad02-8f3d01a1c602"
<video src="https://github.com/user-attachments/assets/33434c3c-f586-4610-8bb2-d7b7b586a672"
width="80%" controls></video>
</p>
For a browser based configuration, try the [Configuration Terminal](https://config-ui.demo.trustgraph.ai/).
## Watch What is a Context Graph?
## Watch What is a Holonic Context Graph?
[![What is a Context Graph?](https://img.youtube.com/vi/gZjlt5WcWB4/maxresdefault.jpg)](https://www.youtube.com/watch?v=gZjlt5WcWB4)
## Watch Context Graphs in Action
## Watch Building Real Agents from a Context Graph
[![Context Graphs in Action with TrustGraph](https://img.youtube.com/vi/sWc7mkhITIo/maxresdefault.jpg)](https://www.youtube.com/watch?v=sWc7mkhITIo)
[![Real Agents from context graphs with TrustGraph](https://img.youtube.com/vi/lmhmrJ7zRE0/maxresdefault.jpg)](https://www.youtube.com/watch?v=lmhmrJ7zRE0)
## Getting Started with TrustGraph
- [**Getting Started Guides**](https://docs.trustgraph.ai/getting-started)
- [**Using the Workbench**](#workbench)
- [**Developer APIs and CLI**](https://docs.trustgraph.ai/reference)
- [**Deployment Guides**](https://docs.trustgraph.ai/deployment)
## Workbench
## TrustGraph UI
The **Workbench** provides tools for all major features of TrustGraph. The **Workbench** is on port `8888` by default.
<img width="1389" height="961" alt="Image" src="https://github.com/user-attachments/assets/35c9250d-0f01-40cb-9294-1ee8fd9a1b56" />
- **Vector Search**: Search the installed knowledge bases
- **Agentic, GraphRAG and LLM Chat**: Chat interface for agents, GraphRAG queries, or direct to LLMs
- **Relationships**: Analyze deep relationships in the installed knowledge bases
- **Graph Visualizer**: 3D GraphViz of the installed knowledge bases
- **Library**: Staging area for installing knowledge bases
- **Flow Classes**: Workflow preset configurations
- **Flows**: Create custom workflows and adjust LLM parameters during runtime
- **Knowledge Cores**: Manage resuable knowledge bases
- **Prompts**: Manage and adjust prompts during runtime
- **Schemas**: Define custom schemas for structured data knowledge bases
- **Ontologies**: Define custom ontologies for unstructured data knowledge bases
- **Agent Tools**: Define tools with collections, knowledge cores, MCP connections, and tool groups
- **MCP Tools**: Connect to MCP servers
The UI provides tools for all major features of TrustGraph. The UI deploys on port `8888` by default.
- **Agent Console** — Query your agents directly with streaming responses and live explainability event tracking, so you can watch reasoning unfold in real time
- **GraphRAG View** — Interactive graph RAG queries with a visual explainability DAG and inline provenance display, making it easy to see exactly where answers came from
- **Context Explorer** — An interactive 3D context graph explorer with dynamic graph loading, BFS neighborhood extraction, edge pulse animation, and multiple navigation views
- **Document Ingestion** — A complete upload and submission workflow with page and chunk inspection and document structure browsing
- **Ontology Workbench** — A full ontology editor with class and property trees, OWL/XML and Turtle import/export with round-trip fidelity, circular dependency detection, and safe-delete confirmation dialogs
- **Schema Workbench** — Interactive schema management with list, create, edit, and delete operations including field and index management
- **Prompt Editor** — A dedicated prompt editing workflow
## TypeScript Library for UIs
@ -125,134 +172,6 @@ There are 3 libraries for quick UI integration of TrustGraph services.
- [@trustgraph/react-state](https://www.npmjs.com/package/@trustgraph/react-state)
- [@trustgraph/react-provider](https://www.npmjs.com/package/@trustgraph/react-provider)
## Context Cores
Context Cores are how TrustGraph treats context like code. A Context Core is a **portable, versioned bundle of context** that you can ship between projects and environments, pin in production, and reuse across agents. It packages the “stuff agents need to know” (structured knowledge + embeddings + evidence + policies) into a single artifact, so you can treat context like code: build it, test it, version it, promote it, and roll it back. TrustGraph is built to support this kind of end-to-end context engineering and orchestration workflow.
### Whats inside a Context Core
A Context Core typically includes:
- Ontology (your domain schema) and mappings
- Context Graph (entities, relationships, supporting evidence)
- Embeddings / vector indexes for fast semantic entry-point lookup
- Source manifests + provenance (where facts came from, when, and how they were derived)
- Retrieval policies (traversal rules, freshness, authority ranking)
## Tech Stack
TrustGraph provides component flexibility to optimize agent workflows.
<details>
<summary>LLM APIs</summary>
<br>
- Anthropic<br>
- AWS Bedrock<br>
- AzureAI<br>
- AzureOpenAI<br>
- Cohere<br>
- Google AI Studio<br>
- Google VertexAI<br>
- Mistral<br>
- OpenAI<br>
</details>
<details>
<summary>LLM Orchestration</summary>
<br>
- LM Studio<br>
- Llamafiles<br>
- Ollama<br>
- TGI<br>
- vLLM<br>
</details>
<details>
<summary>Multi-model storage</summary>
<br>
- Apache Cassandra<br>
</details>
<details>
<summary>VectorDB</summary>
<br>
- Qdrant<br>
</details>
<details>
<summary>File and Object Storage</summary>
<br>
- Garage<br>
</details>
<details>
<summary>Observability</summary>
<br>
- Prometheus<br>
- Grafana<br>
- Loki<br>
</details>
<details>
<summary>Data Streaming</summary>
<br>
- Apache Pulsar<br>
- RabbitMQ<br>
- Apache Kafka<br>
</details>
<details>
<summary>Clouds</summary>
<br>
- AWS<br>
- Azure<br>
- Google Cloud<br>
- OVHcloud<br>
- Scaleway<br>
</details>
## Observability & Telemetry
Once the platform is running, access the Grafana dashboard at:
```
http://localhost:3000
```
Default credentials are:
```
user: admin
password: admin
```
The default Grafana dashboard tracks the following:
<details>
<summary>Telemetry</summary>
<br>
- LLM Latency<br>
- Error Rate<br>
- Service Request Rates<br>
- Queue Backlogs<br>
- Chunking Histogram<br>
- Error Source by Service<br>
- Rate Limit Events<br>
- CPU usage by Service<br>
- Memory usage by Service<br>
- Models Deployed<br>
- Token Throughput (Tokens/second)<br>
- Cost Throughput (Cost/second)<br>
</details>
## Contributing
[Developer's Guide](https://docs.trustgraph.ai/guides/building/introduction.html)
@ -261,7 +180,7 @@ The default Grafana dashboard tracks the following:
**TrustGraph** is licensed under [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0).
Copyright 2024-2025 TrustGraph
Copyright 2024-2026 TrustGraph
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

View file

@ -32,6 +32,14 @@ processors:
id: graph-embeddings-write
store_uri: http://localhost:6333
# Keyword (BM25) index: ingest-write and query in one processor, since
# the FTS5 index is a single local file.
- class: trustgraph.storage.kw_index.fts5.Processor
params:
<<: *defaults
id: kw-index
index_path: /tmp/tg-kw-index.db
- class: trustgraph.query.row_embeddings.qdrant.Processor
params:
<<: *defaults

View file

@ -0,0 +1,402 @@
---
layout: default
title: "Audit Events Technical Specification"
parent: "Tech Specs"
---
# Audit Events Technical Specification
## Overview
This specification defines the audit event system for TrustGraph.
Audit events provide a structured, complete record of security-
relevant operations: API gateway invocations and IAM decisions.
The design principle is: **emit everything, let consumers decide.**
Audit events are cheap to produce (a pub/sub message per operation)
and rich enough to support any downstream consumer — compliance
dashboards, SIEM integration, anomaly detection, billing metering,
or simple grep-based debugging. This spec covers event production
only. Storage, retention, alerting, and presentation are
deployment-specific concerns handled by consumers outside this
boundary.
## Motivation
TrustGraph currently has operational logging (Python `logging` to
stdout/Loki) but no structured audit trail. Operational logs are
unstructured, filtered by level, and designed for debugging — not
for answering "who did what, when, and was it allowed?"
Enterprise deployments need:
- **Compliance evidence** — demonstrable record of access for
auditors.
- **Incident investigation** — reconstruct what happened around a
security event.
- **Anomaly detection** — feed structured events into monitoring
systems.
- **Accountability** — attribute actions to identities across
workspaces.
The current logging infrastructure cannot serve these needs because
it is unstructured, inconsistently formatted, and interleaves
debug noise with security-relevant signals.
## Design Principles
- **Complete.** Every gateway request and every IAM decision emits
an event. No sampling, no level-gating. The pub/sub cost is
negligible; consumers filter what they need.
- **Structured.** Events are typed, versioned, machine-parseable
JSON objects with a fixed envelope and operation-specific payloads.
No free-text messages.
- **Cheap to produce.** Events land on a pub/sub topic. No
synchronous writes, no blocking on consumer availability. If no
consumer is subscribed, events are discarded by the broker — that
is acceptable.
- **Rich.** Events carry enough context to reconstruct the full
security narrative without correlating against operational logs.
Identity, workspace, capability, resource, outcome, timing,
client metadata.
- **Immutable.** Once emitted, an event is a fact. Consumers may
filter, aggregate, or discard events, but never mutate them.
- **Decoupled.** Producers (gateway, IAM service) have no knowledge
of consumers. The topic is fire-and-forget. This keeps the
critical path fast and allows diverse consumer deployments.
## Architecture
### Event transport
Audit events are published to a dedicated pub/sub topic, declared
in the schema layer following the project's queue naming convention:
```python
audit_events_queue = queue('audit-events', cls='notify')
```
This produces the queue identifier `notify:tg:audit-events`, which
each backend maps to its native topic format (e.g. Pulsar maps
`notify` to `non-persistent://tg/notify/audit-events`).
The `notify` class is the right fit: non-persistent, per-subscriber
delivery, no competing-consumer semantics. Audit event production
must never block the gateway or IAM service. Consumers that need
durability persist events themselves on receipt.
A single topic carries all event types, distinguished by the
`event_type` field in the envelope. This simplifies producer
logic and allows consumers to subscribe once and filter client-side.
### Producers
Two components emit audit events:
1. **API Gateway** — emits a `gateway.request` event for every
inbound HTTP/WebSocket request after the request completes
(or fails).
2. **IAM Service** — emits `iam.authenticate` and `iam.authorise`
events for every authentication and authorisation decision.
Both producers emit asynchronously — the event is published after
the response is sent (gateway) or after the decision is returned
(IAM). Audit emission is never on the critical path.
### Consumers
Not defined by this spec. Example consumers that deployments
may wire up:
- Append to an immutable log store (S3, Cassandra, ClickHouse).
- Forward to a SIEM (Splunk, Elastic, Sentinel).
- Aggregate for billing/metering.
- Feed an anomaly detection model.
- Write to stdout for development debugging.
## Event Envelope
Every audit event shares a common envelope:
```json
{
"schema_version": 1,
"event_id": "uuid-v4",
"event_type": "gateway.request",
"timestamp": "2026-07-05T14:23:01.123Z",
"producer": "api-gateway",
"payload": { ... }
}
```
| Field | Type | Description |
|---|---|---|
| `schema_version` | int | Envelope schema version. Consumers must ignore events with versions they don't understand. |
| `event_id` | string | Globally unique event identifier (UUID v4). |
| `event_type` | string | Dot-separated event type from the vocabulary below. |
| `timestamp` | string | ISO 8601 UTC timestamp at event emission. |
| `producer` | string | Component identity that emitted the event. |
| `payload` | object | Event-type-specific structured data. |
## Event Types
### `gateway.request`
Emitted by the API gateway for every completed request.
```json
{
"request_id": "uuid-v4",
"method": "POST",
"path": "/api/v1/flow/default/graph-rag",
"capability": "graph-rag:query",
"workspace": "production",
"identity": "user:mark",
"client_ip": "192.168.1.42",
"user_agent": "trustgraph-cli/2.6.11",
"status_code": 200,
"outcome": "success",
"duration_ms": 1423,
"request_size_bytes": 256,
"response_size_bytes": 4096,
"parameters": {
"collection": "default",
"entity_limit": 50
}
}
```
| Field | Type | Description |
|---|---|---|
| `request_id` | string | Unique ID for this request, propagated to IAM events for correlation. |
| `method` | string | HTTP method. |
| `path` | string | Request path (no query string). |
| `capability` | string | The capability required for this endpoint (from the capability vocabulary). |
| `workspace` | string | Resolved workspace for this request. |
| `identity` | string | Authenticated identity handle, or `"anonymous"` if unauthenticated. |
| `client_ip` | string | Client IP address (may be from X-Forwarded-For). |
| `user_agent` | string | Client User-Agent header. |
| `status_code` | int | HTTP response status code. |
| `outcome` | string | One of `success`, `denied`, `error`, `unauthenticated`. |
| `duration_ms` | int | Request duration in milliseconds. |
| `request_size_bytes` | int | Request body size. |
| `response_size_bytes` | int | Response body size. |
| `error` | string | Error category. Present only when outcome is not `success`. |
| `parameters` | object | Operation-specific parameters extracted from the request (not the full body — only semantically relevant fields). |
### `iam.authenticate`
Emitted by the IAM service for every authentication attempt.
```json
{
"request_id": "uuid-v4",
"credential_type": "api-key",
"identity": "user:mark",
"outcome": "success",
"client_ip": "192.168.1.42",
"key_id": "key-abc123"
}
```
| Field | Type | Description |
|---|---|---|
| `request_id` | string | Correlates with the gateway request that triggered this authentication. |
| `credential_type` | string | One of `api-key`, `jwt`, `login-password`. |
| `identity` | string | Resolved identity on success, or `"unknown"` on failure. |
| `outcome` | string | One of `success`, `failure`. |
| `failure_reason` | string | Internal failure category (not exposed to clients): `invalid-key`, `expired-jwt`, `bad-signature`, `user-disabled`, `unknown-user`. Present only on failure. |
| `client_ip` | string | Forwarded from the gateway request. |
| `key_id` | string | API key identifier (not the secret). Present only on key-based auth. |
**Note:** `failure_reason` is for the audit log only. The client
response is always the same masked error per the IAM contract's
security rule. The audit consumer sees the real reason; the
attacker does not.
### `iam.authorise`
Emitted by the IAM service for every authorisation decision.
```json
{
"request_id": "uuid-v4",
"identity": "user:mark",
"capability": "graph-rag:query",
"workspace": "production",
"resource": "flow:default",
"outcome": "allow",
"evaluated_roles": ["workspace-user"],
"evaluation_time_us": 42
}
```
| Field | Type | Description |
|---|---|---|
| `request_id` | string | Correlates with the gateway request. |
| `identity` | string | Identity being authorised. |
| `capability` | string | Capability being checked. |
| `workspace` | string | Workspace scope of the resource. |
| `resource` | string | Structured resource identifier. |
| `outcome` | string | One of `allow`, `deny`. |
| `denial_reason` | string | Why denied: `no-matching-role`, `capability-not-in-role`, `workspace-not-accessible`, `user-disabled`. Present only on denial. |
| `evaluated_roles` | list of string | Roles evaluated during the decision (OSS regime specific — other regimes may populate differently). |
| `evaluation_time_us` | int | Time to evaluate the decision in microseconds. |
### `iam.management`
Emitted by the IAM service for administrative mutations.
```json
{
"request_id": "uuid-v4",
"actor": "user:admin",
"operation": "create-user",
"target_identity": "user:new-hire",
"target_workspace": "engineering",
"outcome": "success",
"details": {
"roles_assigned": ["workspace-user"]
}
}
```
| Field | Type | Description |
|---|---|---|
| `request_id` | string | Correlates with the gateway request. |
| `actor` | string | Identity performing the action. |
| `operation` | string | IAM operation name (`create-user`, `delete-api-key`, `assign-role`, `create-workspace`, etc.). |
| `target_identity` | string | Identity being acted upon. Present only when applicable. |
| `target_workspace` | string | Workspace being acted upon. Present only when applicable. |
| `outcome` | string | One of `success`, `error`. |
| `details` | object | Operation-specific details (roles assigned, key created, etc.). |
## Correlation
All events from a single gateway request share the same
`request_id`. A typical request produces:
1. One `gateway.request` event (after completion).
2. One `iam.authenticate` event (credential validation).
3. One or more `iam.authorise` events (capability checks).
Consumers can reconstruct the full request lifecycle by grouping
on `request_id`.
## Implementation
### Gateway changes
The gateway emits `gateway.request` events. Implementation:
- Assign a UUID `request_id` at request entry.
- Pass `request_id` and `client_ip` to the IAM service in the
`IamRequest` (new fields on the dataclass).
- After the response is sent, publish the audit event to the
audit topic. This is a non-blocking fire-and-forget publish.
### IAM service changes
The IAM service emits `iam.authenticate`, `iam.authorise`, and
`iam.management` events. Implementation:
- Accept `request_id` and `client_ip` from the gateway on each
`IamRequest`.
- After each decision or mutation, publish the corresponding audit
event. Non-blocking.
### Schema additions
New queue declaration in `trustgraph-base/trustgraph/schema/`:
```python
from trustgraph.schema.core.topic import queue
audit_events_queue = queue('audit-events', cls='notify')
```
New fields on `IamRequest`:
```python
@dataclass
class IamRequest:
...
request_id: str = ""
client_ip: str = ""
```
These are informational — the IAM service does not act on them
beyond echoing them into audit events.
### Pub/sub producer
A lightweight audit publisher utility in `trustgraph-base`:
```python
class AuditPublisher:
def __init__(self, producer):
self.producer = producer
async def emit(self, event_type, payload):
event = {
"schema_version": 1,
"event_id": str(uuid4()),
"event_type": event_type,
"timestamp": datetime.utcnow().isoformat() + "Z",
"producer": self.component_name,
"payload": payload,
}
await self.producer.send(json.dumps(event).encode())
```
The publisher is instantiated once per component and shared across
request handlers.
## What This Spec Does Not Cover
- **Storage.** Where audit events are persisted, for how long,
and in what format. Deployment-specific.
- **Alerting.** What conditions trigger alerts. Consumer logic.
- **Retention policy.** How long events are kept. Compliance-
dependent.
- **UI.** Audit log viewers, dashboards, search interfaces.
- **Filtering/routing.** Topic partitioning, consumer-side
filtering, event routing to different backends.
- **Redaction.** PII handling in audit events (may be needed for
GDPR — a future concern for enterprise consumers).
These are all consumer-side concerns. The value of this boundary
is that producers remain simple and fast while consumers can be
as sophisticated as the deployment requires.
## Open Questions
- **Should WebSocket upgrade events emit separately from per-message
events?** Current proposal: one `gateway.request` per WebSocket
session (on close), with `duration_ms` covering the full session.
Per-message audit for long-lived sockets (e.g. streaming RAG) may
be needed for metering but adds volume.
- **Should `parameters` in `gateway.request` be standardised per
endpoint, or free-form?** Standardised is more useful for
consumers but requires maintenance as endpoints evolve.
- **Event ordering guarantees.** Pub/sub does not guarantee
ordering across partitions. Consumers that need strict ordering
must sort by `timestamp` or `request_id` sequence.
## References
- [IAM Contract](iam-contract.md) — the authentication/authorisation
abstraction.
- [IAM Protocol](iam-protocol.md) — the OSS regime wire protocol.
- [Capability Vocabulary](capabilities.md) — the capability strings
used in authorisation and audit events.
- [Logging Strategy](logging-strategy.md) — operational logging
(complementary, not overlapping).

View file

@ -100,7 +100,6 @@ multi-word subsystems.
| `users:admin` | Assign / remove roles on users within the workspace |
| `keys:self` | Create / revoke / list **own** API keys |
| `keys:admin` | Create / revoke / list **any user's** API keys within the workspace |
| `workspaces:list-own` | List workspaces the caller has access to |
| `workspaces:admin` | Create / delete / disable workspaces (system-level) |
| `iam:admin` | JWT signing-key rotation, IAM-level operations |
| `metrics:read` | Prometheus metrics proxy |
@ -111,7 +110,7 @@ The open-source edition ships three roles:
| Role | Capabilities |
|---|---|
| `reader` | `agent`, `graph:read`, `documents:read`, `rows:read`, `llm`, `embeddings`, `mcp`, `collections:read`, `knowledge:read`, `flows:read`, `config:read`, `keys:self`, `workspaces:list-own` |
| `reader` | `agent`, `graph:read`, `documents:read`, `rows:read`, `llm`, `embeddings`, `mcp`, `collections:read`, `knowledge:read`, `flows:read`, `config:read`, `keys:self` |
| `writer` | everything in `reader` **+** `graph:write`, `documents:write`, `rows:write`, `collections:write`, `knowledge:write` |
| `admin` | everything in `writer` **+** `config:write`, `flows:write`, `users:read`, `users:write`, `users:admin`, `keys:admin`, `workspaces:admin`, `iam:admin`, `metrics:read` |

View file

@ -0,0 +1,564 @@
# GraphRAG Semantic Filter Improvement
## Problem Statement
The GraphRAG semantic filter is observed to be ineffective with certain
LLM models. Smaller models in particular produce poor-quality edge
relevance scores, and there is a suspicion that models trained or
evaluated heavily on non-Roman-script datasets offer lower performance
on the semantic ranking operation.
The root cause is that the current implementation delegates edge
relevance scoring to the LLM via a prompt that asks the model to
assign a 110 relevance score to each knowledge-graph edge. This
task — ranking structured triples for relevance to a natural-language
query — is not well covered in standard LLM evaluation suites, so
model benchmark scores are not predictive of performance on this
operation. The result is that GraphRAG quality varies unpredictably
across model choices, undermining confidence in the pipeline.
Beyond model variability, the LLM scoring step has further problems:
- **Cost and latency.** The LLM call consumes tokens and adds
latency to every query, yet its output is unreliable. Even when
the model performs well, the cost is disproportionate for what is
fundamentally a ranking operation.
- **Subjective scoring scale.** The 110 relevance scale gives the
model no objective criteria for what constitutes a 5 versus a 7.
Different models interpret the scale differently, and even the same
model can produce inconsistent scores across runs.
- **Redundancy with the embedding pre-filter.** The pipeline already
contains a cosine-similarity stage that ranks edges by semantic
relevance using embeddings. The LLM scoring step is a second
filter applied on top of this, and it is not clear that it adds
enough value to justify the additional cost and risk of
degradation.
### Industry context
Semantic ranking is rigorously evaluated on dedicated benchmarks such
as MTEB (Massive Text Embedding Benchmark) and BEIR (Benchmarking
Information Retrieval), which test retrieval and reranking across
diverse domains. The current TrustGraph approach — prompting a
general-purpose LLM to score and rank documents (the "listwise"
approach) — is known to be poorly optimized for this task. It
suffers from positional bias, formatting failures, and
inconsistency at scale.
The industry standard for semantic ranking has moved to
cross-encoder models: lightweight, purpose-built models that take a
querydocument pair as input and produce a single relevance score.
These models are fine-tuned on millions of relevance-labelled pairs
and dominate retrieval benchmarks. They are fast, deterministic,
and do not require an LLM inference call.
## Architecture
### Cross-encoder service
A new request/response service that exposes a generic semantic
ranking API. The service is not specific to GraphRAG — it is a
reusable building block for any component that needs to rank text
by relevance.
The service interface is pluggable. Alternative implementations
can be swapped in behind the same API.
**Packaging options considered:**
- *`sentence-transformers`.* Full-featured, widely used.
However, it pulls in PyTorch (~2 GB), making containers
very large. Tested at ~1.8 seconds for 2200 edges.
- *`optimum.onnxruntime`.* ONNX-based inference. Still
depends on PyTorch at import time despite using ONNX for
inference. Tested at ~4.2 seconds for 2200 edges.
- *`flashrank`.* Lightweight wrapper around ONNX Runtime
with a clean API (`Ranker`, `RerankRequest`). No PyTorch
dependency. Tested at ~4.4 seconds for 2200 edges.
- *Pure `onnxruntime` + `tokenizers`.* Leanest option
(~200 MB total). Requires manual tokenisation, padding,
and numpy array management — more boilerplate to maintain.
- *External API (e.g. Cohere Rerank).* No local model at
all. Adds network latency and an external dependency.
**Decision:** `flashrank` for the initial implementation.
No PyTorch dependency, clean API, comparable performance.
The pluggable interface allows swapping to another backend
later.
**Request:**
- `queries` — list of `{id, text}` objects. In the GraphRAG use
case these are the concepts extracted from the user's question.
- `documents` — list of `{id, text}` objects. In the GraphRAG
use case these are the candidate knowledge-graph edges
represented as text.
- `limit` — integer. Maximum number of results to return.
**Scoring:**
The service produces the cartesian product of all querydocument
pairs and scores each pair through the cross-encoder model. For
each document, the maximum score across all queries is taken as the
document's relevance score. Documents are then ranked by this
score and the top `limit` results are returned.
**Response:**
A list of the top `limit` results, each containing:
- `document_id` — the ID of the matched document.
- `query_id` — the ID of the query (concept) that produced the
highest score for this document.
- `score` — the relevance score.
Including `query_id` in the response supports the explainability
interface: it records that an edge was selected because it is
related to a specific concept.
### Integration
The cross-encoder service follows the standard TrustGraph service
integration pattern:
- **Base package (trustgraph-base).** Schema definitions for the
cross-encoder request/response messages. A client class that
other components (e.g. GraphRAG) can use to call the
cross-encoder service. Message translator registration so the
pub/sub layer can serialise/deserialise the messages.
- **Flow package (trustgraph-flow).** The cross-encoder service
implementation itself — loads the model, listens for requests,
scores pairs, returns results. Flow definition support so the
cross-encoder can be introduced into a processing flow via the
standard flow configuration. `flashrank` is added as a
dependency of `trustgraph-flow`. The service runs in its own
container.
- **API gateway.** A gateway endpoint that routes cross-encoder
requests from the HTTP API to the service over pub/sub and
returns the response.
- **CLI tool.** A command-line utility
(e.g. `tg-invoke-cross-encoder`) that calls the gateway
endpoint for manual testing and debugging.
### Current GraphRAG pipeline
The current pipeline follows these steps:
1. **Concept extraction.** An LLM prompt extracts key concepts
from the user's query.
2. **Graph exploration.** Seed entities are found via embedding
similarity. A subgraph is built by multi-hop traversal from
the seed entities (up to `max_path_length` hops, capped at
`max_subgraph_size` edges).
3. **Embedding pre-filter.** Each edge is embedded as
`"subject, predicate, object"` and scored by cosine similarity
against the concept embeddings. The top `edge_score_limit`
(default 30) edges are kept.
4. **LLM edge scoring.** The `kg-edge-scoring` prompt asks the
LLM to assign a 110 relevance score to each remaining edge.
The top `edge_limit` (default 25) edges are kept.
5. **LLM edge reasoning.** The `kg-edge-reasoning` prompt asks
the LLM to explain why each selected edge is relevant to the
query. Used for the explainability interface.
6. **Document tracing.** Selected edges are traced back to their
source documents in the librarian. Runs concurrently with
step 5.
7. **Synthesis.** The `kg-synthesis` prompt generates the final
answer from the selected edges and source document metadata.
### Potential improvements
#### Replace LLM edge scoring with cross-encoder (step 4)
The LLM edge scoring step is replaced by a call to the
cross-encoder service. The candidate edges are the documents and
`edge_limit` is the limit. This is a direct substitution: faster,
cheaper, deterministic, and more reliable across model choices.
The LLM `kg-edge-scoring` prompt is retired.
**Cross-encoder query input: concepts vs. raw query.** There are
two options for what to use as the cross-encoder queries:
- *Option A: Raw user query.* Pass the original question as a
single query string. Simpler, no dependency on concept
extraction. However, raw queries contain noise words and
conversational phrasing that do not match well against the
structured vocabulary of knowledge-graph edges. A single query
also means every edge competes against the full question — a
partial match on one aspect is diluted.
- *Option B: Extracted concepts.* Pass the concepts from step 1
as separate queries. The concepts are distilled, focused terms
that are closer to the language of the edges. With multiple
concepts as independent queries, the cross-encoder scores each
edge against each concept separately, giving better coverage —
an edge only needs to match one concept well to be selected.
The trade-off is a dependency on the LLM concept extraction
step, but this is already in the pipeline and is a lightweight,
reliable LLM call.
**Decision:** Option B — use extracted concepts. The concept
extraction is fast, and the resulting terms produce better
cross-encoder matches against structured triples.
#### Edge text representation
The current embedding pre-filter represents each edge as
`"subject, predicate, object"`. Two changes:
- **Drop commas.** Commas add tokenisation noise without semantic
value.
- **Direction-aware text.** The reranker text should highlight
the *new* information relative to the traversal direction.
The frontier entity is already known context — repeating it
adds noise and, when traversing from an object node, causes
many edges to produce identical reranker text (e.g. 18
products sharing the same `hasSubcategory Processors` triple
all collapse to the same string when the subject is dropped).
The text is constructed based on which position the frontier
entity occupied in the triple:
- **From subject** (s=entity): `"{predicate} {object}"`
the subject is known, predicate and object are new.
- **From object** (o=entity): `"{subject} {predicate}"`
the object is known, subject and predicate are new.
- **From predicate** (p=entity): `"{subject} {object}"`
the predicate is known, subject and object are new.
This eliminates the duplicate-text problem that arises when
traversing inward from a shared object node, and gives the
cross-encoder a more informative signal at every hop.
#### Remove the embedding pre-filter (step 3)
The embedding pre-filter was introduced to reduce the number of
edges before the expensive LLM scoring call. With the
cross-encoder replacing the LLM call, this cost equation changes.
**Arguments for removal:**
- The cross-encoder is fast enough to score the full subgraph
directly. In testing, 2200 edges scored in ~1.8 seconds; at
the default `max_subgraph_size` of 150 edges, scoring takes
a fraction of a second.
- The pre-filter is a weaker version of what the cross-encoder
does. Bi-encoder cosine similarity embeds the query and
document independently and compares vectors; the cross-encoder
processes both texts together through the full transformer,
giving it much better relevance judgement. Running a weaker
filter before a stronger one adds latency without improving
quality.
- Removing it eliminates an embedding service call (two batches:
concepts + edges) and the associated latency.
**Arguments for keeping it:**
- If the subgraph is very large (thousands of edges), the
cross-encoder's linear scaling could become a bottleneck.
The pre-filter would act as a safety valve.
- The embedding call is cheap compared to an LLM call, so the
overhead is modest.
**Decision:** Remove the pre-filter. The `max_subgraph_size`
parameter (default 150) already caps the number of edges entering
this stage, so the cross-encoder will not face an unbounded
workload. If very large subgraphs become a concern in future,
the pre-filter can be reintroduced or `max_subgraph_size` can be
tuned.
#### Iterative graph traversal with cross-encoder filtering
The current pipeline performs graph exploration and edge filtering
as separate phases: first build the full subgraph (up to
`max_path_length` hops), then score and filter edges. An
alternative is to interleave traversal and filtering — at each
hop, use the cross-encoder to select relevant edges before
expanding further.
**Option A: Big-bang traversal then filter.** Traverse the full
subgraph up to `max_path_length` hops from the seed entities,
collecting all edges up to `max_subgraph_size`. Then
cross-encode the entire result to select the top edges.
- Simple to implement — the current traversal logic is largely
unchanged.
- Produces large, unfocused subgraphs. Irrelevant branches are
explored and scored even though they will be discarded.
- Poorly suited to multi-hop reasoning. For a query about
Voyager 1, the subgraph includes Voyager 2's edges because
they are within hop distance, and the filter must then
separate them.
**Option B: Iterative hop-and-filter.** At each hop:
1. Retrieve all edges one hop from the current frontier nodes.
2. Cross-encode these edges against the query concepts.
3. Select the top relevant edges.
4. The target nodes of the selected edges become the frontier
for the next hop.
5. Repeat up to `max_path_length` hops.
The final set of selected edges across all hops is the input to
synthesis.
- **Guided exploration.** Each hop focuses the search by
pruning irrelevant branches before expanding further. The
working set stays small and relevant at every step.
- **Multi-hop reasoning works naturally.** Following
"Voyager 1 → has-event → crossed the heliopause" succeeds
because each hop is individually relevant and leads to the
next.
- **Smaller total workload.** Fewer edges are scored overall
because irrelevant branches are never expanded.
- **Trade-off: greedy pruning.** An edge discarded at hop 1
cannot lead to relevant edges at hop 2. This is inherent in
any bounded traversal, and the cross-encoder is better
equipped to make this relevance judgement than a blind hop
limit.
- **Trade-off: sequential latency.** Hops cannot be
parallelised since each depends on the previous. However,
each cross-encoder call on a small edge set is very fast
(sub-second for typical working sets).
**Decision:** Option B — iterative hop-and-filter. The guided
traversal produces more focused subgraphs and supports multi-hop
reasoning, which is a significant quality improvement over the
current approach.
#### Replace LLM edge reasoning with cross-encoder metadata (step 5)
The current `kg-edge-reasoning` prompt asks the LLM to explain why
each edge is relevant. With the cross-encoder now making the
selection, this explanation would be a post-hoc fabrication — the
LLM was not involved in the decision.
- *Option A: Keep LLM reasoning.* Generates natural-language
explanations but they are not grounded in the actual selection
process. Adds an LLM call per query.
- *Option B: Record cross-encoder metadata.* The cross-encoder
already returns the matched concept and score for each selected
edge. Use this directly as the explanation.
**Decision:** Option B. The cross-encoder metadata is the true
reason the edge was selected. The `kg-edge-reasoning` prompt is
retired.
#### Explainability interface update
The explainability interface uses a `Focus` entity containing
`EdgeSelection` sub-entities. Each `EdgeSelection` currently
carries an `edge` (the quoted triple) and a `reasoning` field
(free-text LLM prose), stored as `tg:reasoning` in the
provenance graph.
With the cross-encoder replacing LLM reasoning, the
`EdgeSelection` type gains two new predicates and drops one:
- **Remove** `tg:reasoning` — no longer produced.
- **Add** `tg:concept` — the concept text that produced the
highest cross-encoder score for this edge.
- **Add** `tg:score` — the cross-encoder relevance score.
This is an evolution of the existing `EdgeSelection` type, not a
new entity type. The edge selection sub-entities currently have
no `rdf:type` declared; a new `tg:EdgeSelection` type should be
added so that consumers can identify them in the provenance
graph. The `Focus` entity and its relationship to `Exploration`
are unchanged.
The `Focus` entity's token-usage metadata (`tg:inToken`,
`tg:outToken`, `tg:llmModel`) no longer applies since there is
no LLM call. These fields are dropped from the Focus entity.
### Proposed pipeline
1. **Concept extraction.** Unchanged — LLM extracts key concepts
from the user's query.
2. **Seed entity lookup.** Find seed entities via embedding
similarity against the extracted concepts.
3. **Iterative hop-and-filter.** For each hop up to
`max_path_length`:
a. Retrieve all edges one hop from the current frontier nodes.
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.
d. Select the top relevant edges. The target nodes of the
selected edges become the frontier for the next hop.
4. **Document tracing.** Selected edges are traced back to source
documents.
5. **Synthesis.** The `kg-synthesis` prompt generates the final
answer from the selected edges and source document metadata.
### Implementation order
1. Cross-encoder service with full integration (base schema,
flow service, gateway endpoint, CLI tool).
2. GraphRAG pipeline changes (iterative hop-and-filter,
edge representation, remove pre-filter).
3. Explainability update (`tg:EdgeSelection` type, concept
and score predicates, retire `tg:reasoning`).
4. Retire `kg-edge-scoring` and `kg-edge-reasoning` prompts.
5. Update `tg-invoke-graph-rag` and `tg-show-explain-trace`
to display the new metadata. Use these as the main
end-to-end test.
6. Fix any failing unit tests, then add new tests as needed.
7. Write guidance for UX devs to update the UI for the new
explainability predicates.
## UX developer guidance
This section describes the changes to the explainability interface
that affect frontend rendering of GraphRAG Focus events.
### What changed
Edge selection in GraphRAG previously used LLM-based scoring and
reasoning. Each selected edge carried a `tg:reasoning` predicate
with free-text explanation from the LLM. This has been replaced
by a cross-encoder reranker that scores edges against query
concepts. The explainability data now carries structured metadata
instead of free text.
### Removed
- **`tg:reasoning`** is no longer emitted on edge selection
entities in GraphRAG Focus events. UX code that reads
`edge_sel.reasoning` will get an empty string. Remove any
rendering that displays a "Reasoning" or "Reason" field for
Focus edges.
- The **`kg-edge-scoring`**, **`kg-edge-reasoning`**, and
**`kg-edge-selection`** prompts are retired. Any UX that
references these prompt names should be cleaned up.
### Added
Each edge selection entity within a Focus event now has three
new properties:
| RDF predicate | API field | Type | Description |
|---|---|---|---|
| `rdf:type tg:EdgeSelection` | (type check) | — | Each edge selection entity is now explicitly typed |
| `tg:concept` | `edge_sel.concept` | `str` | The query concept that matched this edge |
| `tg:score` | `edge_sel.score` | `float` or `None` | Cross-encoder relevance score (0.01.0) |
The `tg:edge` predicate (RDF-star quoted triple) is unchanged.
### How to render
The recommended rendering for each selected edge in a Focus event:
```
Edge: (subject_label, predicate_label, object_label)
Concept: <concept> Score: <score formatted to 4 decimal places>
```
Scores near 1.0 indicate high relevance; scores near 0.0 indicate
low relevance. UX could use the score to drive visual indicators
such as colour intensity or a relevance bar.
Edges are not returned in score order — they arrive in traversal
order across hops. If the UX wants to display edges ranked by
relevance, sort by `edge_sel.score` descending.
### API classes (Python)
The `EdgeSelection` dataclass in `trustgraph.api.explainability`
has these fields:
```python
@dataclass
class EdgeSelection:
uri: str
edge: Optional[Dict[str, str]] # {"s": ..., "p": ..., "o": ...}
reasoning: str = "" # Legacy, always empty for new traces
concept: str = "" # Query concept that matched
score: Optional[float] = None # Cross-encoder relevance score
```
These are populated when calling
`ExplainabilityClient.fetch_focus_with_edges()` or when parsing
inline provenance triples from the streaming response.
### WebSocket response format
For inline explainability via the streaming WebSocket, Focus events
arrive as `message_type: "explain"` responses. The `explain_triples`
array contains the edge selection triples. The relevant predicates
in wire format are:
```json
{"s": {"t": "i", "i": "<edge_sel_uri>"},
"p": {"t": "i", "i": "https://trustgraph.ai/ns/concept"},
"o": {"t": "l", "v": "flyby event"}}
{"s": {"t": "i", "i": "<edge_sel_uri>"},
"p": {"t": "i", "i": "https://trustgraph.ai/ns/score"},
"o": {"t": "l", "v": "0.9962"}}
```
Note that `tg:score` is transmitted as a string literal and must
be parsed to a float on the client side.
### Exploration event
The Exploration event's `edge_count` field now reports the number
of edges selected by the cross-encoder across all hops (previously
it reported the total number of edges retrieved before filtering).
The `entities` list continues to report the seed entities found
by vector search.

2603
install_trustgraph.sh Normal file

File diff suppressed because it is too large Load diff

View file

@ -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

View file

@ -24,6 +24,8 @@ from trustgraph.schema import (
EntityContext,
EntityEmbeddings,
ChunkEmbeddings,
AuditEvent,
IamRequest,
)
@ -72,3 +74,21 @@ class TestSchemaFieldContracts:
"context",
"chunk_id",
}
def test_audit_event_fields(self):
assert _field_names(AuditEvent) == {
"schema_version",
"event_id",
"event_type",
"timestamp",
"producer",
"payload_json",
}
def test_iam_request_has_audit_fields(self):
"""IamRequest must carry request_id and client_ip for audit
correlation. Removing these breaks the gatewayIAM audit
chain."""
names = _field_names(IamRequest)
assert "request_id" in names
assert "client_ip" in names

View file

@ -95,10 +95,6 @@ class TestGraphRagIntegration:
async def mock_prompt(prompt_name, variables=None, streaming=False, chunk_callback=None):
if prompt_name == "extract-concepts":
return PromptResult(response_type="text", text="")
elif prompt_name == "kg-edge-scoring":
return PromptResult(response_type="text", text="")
elif prompt_name == "kg-edge-reasoning":
return PromptResult(response_type="text", text="")
elif prompt_name == "kg-synthesis":
return PromptResult(
response_type="text",
@ -113,14 +109,22 @@ class TestGraphRagIntegration:
client.prompt.side_effect = mock_prompt
return client
@pytest.fixture
def mock_reranker_client(self):
"""Mock reranker client for cross-encoder edge filtering"""
client = AsyncMock()
client.rerank.return_value = []
return client
@pytest.fixture
def graph_rag(self, mock_embeddings_client, mock_graph_embeddings_client,
mock_triples_client, mock_prompt_client):
mock_triples_client, mock_reranker_client, mock_prompt_client):
"""Create GraphRag instance with mocked dependencies"""
return GraphRag(
embeddings_client=mock_embeddings_client,
graph_embeddings_client=mock_graph_embeddings_client,
triples_client=mock_triples_client,
reranker_client=mock_reranker_client,
prompt_client=mock_prompt_client,
verbose=True
)
@ -167,8 +171,8 @@ class TestGraphRagIntegration:
# 3. Should query triples to build knowledge subgraph
assert mock_triples_client.query_stream.call_count > 0
# 4. Should call prompt four times (extract-concepts + edge-scoring + edge-reasoning + synthesis)
assert mock_prompt_client.prompt.call_count == 4
# 4. Should call prompt twice (extract-concepts + synthesis)
assert mock_prompt_client.prompt.call_count == 2
# Verify final response
response, usage = response

View file

@ -63,11 +63,6 @@ class TestGraphRagStreaming:
async def prompt_side_effect(prompt_id, variables, streaming=False, chunk_callback=None, **kwargs):
if prompt_id == "extract-concepts":
return PromptResult(response_type="text", text="")
elif prompt_id == "kg-edge-scoring":
# Edge scoring returns JSONL with IDs and scores
return PromptResult(response_type="text", text='{"id": "abc12345", "score": 0.9}\n')
elif prompt_id == "kg-edge-reasoning":
return PromptResult(response_type="text", text='{"id": "abc12345", "reasoning": "Relevant to query"}\n')
elif prompt_id == "kg-synthesis":
if streaming and chunk_callback:
# Simulate streaming chunks with end_of_stream flags
@ -88,14 +83,23 @@ class TestGraphRagStreaming:
client.prompt.side_effect = prompt_side_effect
return client
@pytest.fixture
def mock_reranker_client(self):
"""Mock reranker client for cross-encoder edge filtering"""
client = AsyncMock()
client.rerank.return_value = []
return client
@pytest.fixture
def graph_rag_streaming(self, mock_embeddings_client, mock_graph_embeddings_client,
mock_triples_client, mock_streaming_prompt_client):
mock_triples_client, mock_reranker_client,
mock_streaming_prompt_client):
"""Create GraphRag instance with streaming support"""
return GraphRag(
embeddings_client=mock_embeddings_client,
graph_embeddings_client=mock_graph_embeddings_client,
triples_client=mock_triples_client,
reranker_client=mock_reranker_client,
prompt_client=mock_streaming_prompt_client,
verbose=True
)

View file

@ -46,7 +46,7 @@ class TestGraphRagStreamingProtocol:
client = AsyncMock()
async def prompt_side_effect(prompt_name, variables=None, streaming=False, chunk_callback=None):
if prompt_name == "kg-edge-selection":
if prompt_name == "extract-concepts":
return PromptResult(response_type="text", text="")
elif prompt_name == "kg-synthesis":
if streaming and chunk_callback:
@ -63,14 +63,23 @@ class TestGraphRagStreamingProtocol:
client.prompt.side_effect = prompt_side_effect
return client
@pytest.fixture
def mock_reranker_client(self):
"""Mock reranker client for cross-encoder edge filtering"""
client = AsyncMock()
client.rerank.return_value = []
return client
@pytest.fixture
def graph_rag(self, mock_embeddings_client, mock_graph_embeddings_client,
mock_triples_client, mock_streaming_prompt_client):
mock_triples_client, mock_reranker_client,
mock_streaming_prompt_client):
"""Create GraphRag instance with mocked dependencies"""
return GraphRag(
embeddings_client=mock_embeddings_client,
graph_embeddings_client=mock_graph_embeddings_client,
triples_client=mock_triples_client,
reranker_client=mock_reranker_client,
prompt_client=mock_streaming_prompt_client,
verbose=False
)
@ -327,7 +336,7 @@ class TestStreamingProtocolEdgeCases:
client = AsyncMock()
async def prompt_with_empties(prompt_name, variables=None, streaming=False, chunk_callback=None):
if prompt_name == "kg-edge-selection":
if prompt_name == "extract-concepts":
return PromptResult(response_type="text", text="")
elif prompt_name == "kg-synthesis":
if streaming and chunk_callback:
@ -342,10 +351,14 @@ class TestStreamingProtocolEdgeCases:
client.prompt.side_effect = prompt_with_empties
mock_reranker = AsyncMock()
mock_reranker.rerank.return_value = []
rag = GraphRag(
embeddings_client=AsyncMock(embed=AsyncMock(return_value=[[[0.1]]])),
graph_embeddings_client=AsyncMock(query=AsyncMock(return_value=[])),
triples_client=AsyncMock(query=AsyncMock(return_value=[])),
reranker_client=mock_reranker,
prompt_client=client,
verbose=False
)

View file

@ -15,11 +15,20 @@ from openai.types.chat.chat_completion import Choice
from openai.types.completion_usage import CompletionUsage
from trustgraph.model.text_completion.openai.llm import Processor
from trustgraph.model.text_completion.openai.variants import get_variant
from trustgraph.exceptions import TooManyRequests
from trustgraph.base import LlmResult
from trustgraph.schema import TextCompletionRequest, TextCompletionResponse, Error
def _wire_variant(processor):
"""Attach variant methods to a MagicMock processor."""
processor.variant = get_variant("openai")
processor.thinking = "off"
processor._build_kwargs = Processor._build_kwargs.__get__(processor, Processor)
processor._extract_content = Processor._extract_content.__get__(processor, Processor)
@pytest.mark.integration
class TestTextCompletionIntegration:
"""Integration tests for OpenAI text completion service coordination"""
@ -66,6 +75,7 @@ class TestTextCompletionIntegration:
# Add the actual generate_content method from Processor class
processor.generate_content = Processor.generate_content.__get__(processor, Processor)
_wire_variant(processor)
return processor
@ -119,6 +129,7 @@ class TestTextCompletionIntegration:
# Add the actual generate_content method
processor.generate_content = Processor.generate_content.__get__(processor, Processor)
_wire_variant(processor)
# Act
result = await processor.generate_content("System prompt", "User prompt")
@ -247,6 +258,7 @@ class TestTextCompletionIntegration:
processor.max_output = processor_config["max_output"]
processor.openai = mock_openai_client
processor.generate_content = Processor.generate_content.__get__(processor, Processor)
_wire_variant(processor)
processors.append(processor)
# Simulate multiple concurrent requests
@ -354,6 +366,7 @@ class TestTextCompletionIntegration:
processor.max_output = 2048
processor.openai = mock_openai_client
processor.generate_content = Processor.generate_content.__get__(processor, Processor)
_wire_variant(processor)
# Act
await processor.generate_content("System prompt", "User prompt")

View file

@ -11,6 +11,7 @@ from openai.types.chat import ChatCompletionChunk
from openai.types.chat.chat_completion_chunk import Choice as StreamChoice, ChoiceDelta
from trustgraph.model.text_completion.openai.llm import Processor
from trustgraph.model.text_completion.openai.variants import get_variant
from trustgraph.base import LlmChunk
from tests.utils.streaming_assertions import (
assert_streaming_chunks_valid,
@ -18,6 +19,14 @@ from tests.utils.streaming_assertions import (
)
def _wire_variant(processor):
"""Attach variant methods to a MagicMock processor."""
processor.variant = get_variant("openai")
processor.thinking = "off"
processor._build_kwargs = Processor._build_kwargs.__get__(processor, Processor)
processor._extract_content = Processor._extract_content.__get__(processor, Processor)
@pytest.mark.integration
class TestTextCompletionStreaming:
"""Integration tests for Text Completion streaming"""
@ -69,6 +78,7 @@ class TestTextCompletionStreaming:
processor.generate_content_stream = Processor.generate_content_stream.__get__(
processor, Processor
)
_wire_variant(processor)
return processor
@ -190,6 +200,7 @@ class TestTextCompletionStreaming:
processor.generate_content_stream = Processor.generate_content_stream.__get__(
processor, Processor
)
_wire_variant(processor)
# Act
chunks = []
@ -223,6 +234,7 @@ class TestTextCompletionStreaming:
processor.generate_content_stream = Processor.generate_content_stream.__get__(
processor, Processor
)
_wire_variant(processor)
# Act
chunks = []
@ -258,6 +270,7 @@ class TestTextCompletionStreaming:
processor.generate_content_stream = Processor.generate_content_stream.__get__(
processor, Processor
)
_wire_variant(processor)
# Act & Assert
with pytest.raises(Exception) as exc_info:
@ -295,6 +308,7 @@ class TestTextCompletionStreaming:
processor.generate_content_stream = Processor.generate_content_stream.__get__(
processor, Processor
)
_wire_variant(processor)
# Act
chunks = []
@ -318,6 +332,7 @@ class TestTextCompletionStreaming:
processor.generate_content_stream = Processor.generate_content_stream.__get__(
processor, Processor
)
_wire_variant(processor)
system_prompt = "You are an expert."
user_prompt = "Explain quantum physics."

View file

@ -0,0 +1,90 @@
"""
Tests for the AuditPublisher utility.
Verifies envelope construction, fire-and-forget semantics, and
failure suppression.
"""
import json
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from trustgraph.base.audit_publisher import AuditPublisher
from trustgraph.schema import AuditEvent, audit_events_queue
class TestAuditPublisherInit:
def test_queue_is_notify_class(self):
assert audit_events_queue == "notify:tg:audit-events"
def test_creates_producer_with_audit_queue(self):
backend = MagicMock()
pub = AuditPublisher(
backend=backend,
component_name="test-component",
)
assert pub.producer.topic == audit_events_queue
assert pub.producer.schema == AuditEvent
assert pub.component_name == "test-component"
class TestAuditPublisherEmit:
@pytest.fixture
def publisher(self):
backend = MagicMock()
pub = AuditPublisher(
backend=backend,
component_name="test-svc",
processor_id="proc-1",
)
pub.producer = AsyncMock()
return pub
@pytest.mark.asyncio
async def test_emit_sends_structured_envelope(self, publisher):
await publisher.emit("gateway.request", {"path": "/test"})
publisher.producer.send.assert_called_once()
event = publisher.producer.send.call_args[0][0]
assert isinstance(event, AuditEvent)
assert event.schema_version == 1
assert event.event_type == "gateway.request"
assert event.producer == "test-svc"
assert event.event_id != ""
assert event.timestamp != ""
@pytest.mark.asyncio
async def test_emit_serializes_payload_as_json(self, publisher):
payload = {"method": "POST", "status_code": 200}
await publisher.emit("gateway.request", payload)
event = publisher.producer.send.call_args[0][0]
decoded = json.loads(event.payload_json)
assert decoded == payload
@pytest.mark.asyncio
async def test_emit_generates_unique_event_ids(self, publisher):
await publisher.emit("test.a", {})
await publisher.emit("test.b", {})
ids = [
call[0][0].event_id
for call in publisher.producer.send.call_args_list
]
assert ids[0] != ids[1]
@pytest.mark.asyncio
async def test_emit_swallows_send_failure(self, publisher):
publisher.producer.send.side_effect = RuntimeError("pub/sub down")
await publisher.emit("test.event", {"key": "value"})
@pytest.mark.asyncio
async def test_emit_timestamp_is_utc_iso(self, publisher):
await publisher.emit("test.event", {})
event = publisher.producer.send.call_args[0][0]
assert "T" in event.timestamp
assert "+" in event.timestamp or "Z" in event.timestamp

View file

@ -410,3 +410,56 @@ class TestEdgeCases:
assert hosts == ['mixed-host']
assert username is None # Stays None
assert password == 'mixed-pass'
class TestReplicationFactorParamPath:
def test_explicit_kwarg(self):
with patch.dict(os.environ, {}, clear=True):
_, _, _, _, rf = resolve_cassandra_config(
replication_factor=3,
)
assert rf == 3
def test_kwarg_overrides_env(self):
with patch.dict(os.environ, {'CASSANDRA_REPLICATION_FACTOR': '5'}, clear=True):
_, _, _, _, rf = resolve_cassandra_config(
replication_factor=3,
)
assert rf == 3
def test_env_fallback_when_kwarg_none(self):
with patch.dict(os.environ, {'CASSANDRA_REPLICATION_FACTOR': '5'}, clear=True):
_, _, _, _, rf = resolve_cassandra_config(
replication_factor=None,
)
assert rf == 5
def test_default_when_no_kwarg_no_env(self):
with patch.dict(os.environ, {}, clear=True):
_, _, _, _, rf = resolve_cassandra_config()
assert rf == 1
def test_params_dict_path(self):
with patch.dict(os.environ, {}, clear=True):
params = {'cassandra_replication_factor': 3}
_, _, _, _, rf = resolve_cassandra_config(
replication_factor=params.get('cassandra_replication_factor'),
)
assert rf == 3
def test_params_dict_overrides_env(self):
with patch.dict(os.environ, {'CASSANDRA_REPLICATION_FACTOR': '5'}, clear=True):
params = {'cassandra_replication_factor': 3}
_, _, _, _, rf = resolve_cassandra_config(
replication_factor=params.get('cassandra_replication_factor'),
)
assert rf == 3
def test_params_dict_missing_falls_to_env(self):
with patch.dict(os.environ, {'CASSANDRA_REPLICATION_FACTOR': '5'}, clear=True):
params = {}
_, _, _, _, rf = resolve_cassandra_config(
replication_factor=params.get('cassandra_replication_factor'),
)
assert rf == 5

View file

@ -0,0 +1,81 @@
"""
Tests for RequestResponseSpec's optional flag: an optional client spec
binds only when the flow definition declares its topics, so a definition
predating the topics skips the binding (flow(name) then returns None)
instead of raising KeyError during Flow construction which would wedge
the processor's start-flow retry loop.
"""
import pytest
from unittest.mock import MagicMock
from trustgraph.base.request_response_spec import RequestResponseSpec
class StubImpl:
"""Captures constructor kwargs; stands in for RequestResponse."""
def __init__(self, **kwargs):
self.kwargs = kwargs
def make_spec(optional):
return RequestResponseSpec(
request_name="keyword-index-request",
request_schema=object,
response_name="keyword-index-response",
response_schema=object,
impl=StubImpl,
optional=optional,
)
def make_flow():
flow = MagicMock()
flow.id = "f-id"
flow.name = "f-name"
flow.workspace = "ws"
flow.consumer = {}
return flow
FULL_TOPICS = {
"topics": {
"keyword-index-request": "request:tg:keyword-index:ws:f",
"keyword-index-response": "response:tg:keyword-index:ws:f",
}
}
class TestOptionalRequestResponseSpec:
def test_optional_spec_skips_binding_when_topics_absent(self):
flow = make_flow()
make_spec(optional=True).add(flow, MagicMock(), {"topics": {}})
assert flow.consumer == {}
def test_optional_spec_skips_when_only_one_topic_present(self):
flow = make_flow()
definition = {
"topics": {
"keyword-index-request": "request:tg:keyword-index:ws:f",
}
}
make_spec(optional=True).add(flow, MagicMock(), definition)
assert flow.consumer == {}
def test_optional_spec_binds_when_topics_present(self):
flow = make_flow()
make_spec(optional=True).add(flow, MagicMock(), FULL_TOPICS)
client = flow.consumer["keyword-index-request"]
assert isinstance(client, StubImpl)
assert client.kwargs["request_topic"] == \
"request:tg:keyword-index:ws:f"
def test_default_spec_still_requires_topics(self):
# Non-optional specs keep the existing contract: a missing topic
# is a definition error, surfaced immediately.
with pytest.raises(KeyError):
make_spec(optional=False).add(
make_flow(), MagicMock(), {"topics": {}},
)

View file

@ -195,38 +195,6 @@ class TestPromptClientStreamingCallback:
assert callback.call_args_list[0] == call("test", False)
assert callback.call_args_list[1] == call("", True)
@pytest.mark.asyncio
async def test_kg_prompt_passes_parameters_to_callback(self, prompt_client):
"""Test that kg_prompt correctly passes streaming parameters"""
# Arrange
async def mock_request(request, recipient=None, timeout=600):
if recipient:
responses = [
PromptResponse(text="Answer", object=None, error=None, end_of_stream=False),
PromptResponse(text="", object=None, error=None, end_of_stream=True),
]
for resp in responses:
should_stop = await recipient(resp)
if should_stop:
break
prompt_client.request = mock_request
callback = AsyncMock()
# Act
await prompt_client.kg_prompt(
query="What is machine learning?",
kg=[("subject", "predicate", "object")],
streaming=True,
chunk_callback=callback
)
# Assert
assert callback.call_count == 2
assert callback.call_args_list[0] == call("Answer", False)
assert callback.call_args_list[1] == call("", True)
@pytest.mark.asyncio
async def test_document_prompt_passes_parameters_to_callback(self, prompt_client):
"""Test that document_prompt correctly passes streaming parameters"""

View file

@ -0,0 +1,136 @@
import os
import pytest
from unittest.mock import patch
from trustgraph.base.qdrant_config import (
get_qdrant_defaults,
resolve_qdrant_config,
)
class TestGetQdrantDefaults:
def test_defaults_with_no_env_vars(self):
with patch.dict(os.environ, {}, clear=True):
defaults = get_qdrant_defaults()
assert defaults['url'] == 'http://localhost:6333'
assert defaults['api_key'] is None
assert defaults['replication_factor'] == 1
assert defaults['shard_number'] == 1
def test_defaults_from_env(self):
env = {
'QDRANT_URL': 'http://qdrant:6333',
'QDRANT_API_KEY': 'secret',
'QDRANT_REPLICATION_FACTOR': '3',
'QDRANT_SHARD_NUMBER': '5',
}
with patch.dict(os.environ, env, clear=True):
defaults = get_qdrant_defaults()
assert defaults['url'] == 'http://qdrant:6333'
assert defaults['api_key'] == 'secret'
assert defaults['replication_factor'] == 3
assert defaults['shard_number'] == 5
class TestResolveQdrantConfig:
def test_defaults(self):
with patch.dict(os.environ, {}, clear=True):
url, api_key, rf, sn = resolve_qdrant_config()
assert url == 'http://localhost:6333'
assert api_key is None
assert rf == 1
assert sn == 1
def test_explicit_kwargs(self):
with patch.dict(os.environ, {}, clear=True):
url, api_key, rf, sn = resolve_qdrant_config(
url='http://custom:6333',
api_key='key',
replication_factor=3,
shard_number=5,
)
assert url == 'http://custom:6333'
assert api_key == 'key'
assert rf == 3
assert sn == 5
def test_kwargs_override_env(self):
env = {
'QDRANT_URL': 'http://env:6333',
'QDRANT_REPLICATION_FACTOR': '10',
'QDRANT_SHARD_NUMBER': '10',
}
with patch.dict(os.environ, env, clear=True):
url, _, rf, sn = resolve_qdrant_config(
url='http://explicit:6333',
replication_factor=3,
shard_number=5,
)
assert url == 'http://explicit:6333'
assert rf == 3
assert sn == 5
def test_env_fallback_when_kwargs_none(self):
env = {
'QDRANT_URL': 'http://env:6333',
'QDRANT_REPLICATION_FACTOR': '3',
'QDRANT_SHARD_NUMBER': '5',
}
with patch.dict(os.environ, env, clear=True):
url, _, rf, sn = resolve_qdrant_config()
assert url == 'http://env:6333'
assert rf == 3
assert sn == 5
def test_params_dict_path(self):
with patch.dict(os.environ, {}, clear=True):
params = {
'store_uri': 'http://params:6333',
'api_key': 'pkey',
'qdrant_replication_factor': 3,
'qdrant_shard_number': 5,
}
url, api_key, rf, sn = resolve_qdrant_config(
url=params.get('store_uri'),
api_key=params.get('api_key'),
replication_factor=params.get('qdrant_replication_factor'),
shard_number=params.get('qdrant_shard_number'),
)
assert url == 'http://params:6333'
assert api_key == 'pkey'
assert rf == 3
assert sn == 5
def test_params_dict_overrides_env(self):
env = {
'QDRANT_REPLICATION_FACTOR': '10',
'QDRANT_SHARD_NUMBER': '10',
}
with patch.dict(os.environ, env, clear=True):
params = {
'qdrant_replication_factor': 3,
'qdrant_shard_number': 5,
}
_, _, rf, sn = resolve_qdrant_config(
replication_factor=params.get('qdrant_replication_factor'),
shard_number=params.get('qdrant_shard_number'),
)
assert rf == 3
assert sn == 5
def test_params_dict_missing_falls_to_env(self):
env = {
'QDRANT_REPLICATION_FACTOR': '3',
'QDRANT_SHARD_NUMBER': '5',
}
with patch.dict(os.environ, env, clear=True):
params = {}
_, _, rf, sn = resolve_qdrant_config(
replication_factor=params.get('qdrant_replication_factor'),
shard_number=params.get('qdrant_shard_number'),
)
assert rf == 3
assert sn == 5

View file

@ -0,0 +1,54 @@
"""
Unit tests for trustgraph.bootstrap.initialisers.DefaultFlowStart
Verifies the list/start timeouts are configurable and that the
configured values actually reach the flow-client request calls.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from trustgraph.bootstrap.initialisers.default_flow_start import (
DefaultFlowStart,
)
def test_default_timeouts():
init = DefaultFlowStart(blueprint="bp")
assert init.list_timeout == 10
assert init.start_timeout == 30
def test_timeout_overrides_are_stored():
init = DefaultFlowStart(blueprint="bp", list_timeout=5, start_timeout=99)
assert init.list_timeout == 5
assert init.start_timeout == 99
@pytest.mark.asyncio
async def test_run_forwards_configured_timeouts():
init = DefaultFlowStart(blueprint="bp", list_timeout=5, start_timeout=99)
# Flow client: list-flows returns no error + empty flow list,
# start-flow returns no error.
flow = MagicMock()
flow.start = AsyncMock()
flow.stop = AsyncMock()
flow.request = AsyncMock(side_effect=[
MagicMock(error=None, flow_ids=[]), # list-flows response
MagicMock(error=None), # start-flow response
])
# Context: workspace "default" exists, hands back our mock flow client.
ctx = MagicMock()
ctx.logger = MagicMock()
ctx.config.keys = AsyncMock(return_value=["default"])
ctx.make_flow_client = MagicMock(return_value=flow)
await init.run(ctx, None, "v1")
calls = flow.request.call_args_list
assert len(calls) == 2
assert calls[0].kwargs["timeout"] == 5
assert calls[1].kwargs["timeout"] == 99

View file

@ -0,0 +1,13 @@
"""Unit tests for trustgraph.bootstrap.initialisers.WorkspaceInit."""
from trustgraph.bootstrap.initialisers.workspace_init import WorkspaceInit
def test_default_iam_timeout():
init = WorkspaceInit()
assert init.iam_timeout == 10
def test_iam_timeout_override_is_stored():
init = WorkspaceInit(iam_timeout=42)
assert init.iam_timeout == 42

View file

@ -46,3 +46,17 @@ def sample_metadata():
"user": "test-user",
"collection": "test-collection"
}
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

View file

@ -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

View file

@ -0,0 +1,449 @@
"""
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, 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, 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}}"}),
"answer": json.dumps({"template": "Answer {{q}}"}),
},
"tool": {
"web-search": json.dumps({"name": "web-search", "kind": "http"}),
},
}
# 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")},
]]
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):
"""Config-only-shaped bundle (no collections/docs mocked)."""
path = tmp_path / "ws.tgx"
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
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": True}
assert manifest["knowledge"] == {
"collections": [], "documents": 0, "triples": {},
}
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"
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:
api_cls.return_value = mock_api
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"
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("<urn:trustgraph:collection:research>" in ln
for ln in lines)
assert '"42"^^<http://www.w3.org/2001/XMLSchema#integer>' 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):
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(
"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):
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."""
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) == [
("prompt", "answer"),
("tool", "web-search"),
]
def test_dry_run_writes_nothing(self, bundle, capsys):
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
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):
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))
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'<http://ex.com/s> <http://ex.com/p> "v" '
b'<urn:trustgraph:collection:research> .\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'<http://ex.com/s> <http://ex.com/p> "v" '
b'<urn:trustgraph:collection:research> .\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()

View file

@ -129,6 +129,9 @@ class TestBatchTripleQueries:
# 3 queries, alternating results
assert len(result) == 3
# Each result is a (triple, direction) tuple
for triple, direction in result:
assert direction in (Query.FROM_S, Query.FROM_P, Query.FROM_O)
@pytest.mark.asyncio
async def test_exception_in_one_query_does_not_block_others(self):
@ -153,6 +156,8 @@ class TestBatchTripleQueries:
# 3 queries: 2 succeed, 1 fails → 2 triples
assert len(result) == 2
for triple, direction in result:
assert direction in (Query.FROM_S, Query.FROM_P, Query.FROM_O)
@pytest.mark.asyncio
async def test_none_results_filtered(self):
@ -176,6 +181,8 @@ class TestBatchTripleQueries:
# 3 queries: 1 returns None, 2 return triples
assert len(result) == 2
for triple, direction in result:
assert direction in (Query.FROM_S, Query.FROM_P, Query.FROM_O)
@pytest.mark.asyncio
async def test_empty_entities_no_queries(self):
@ -220,6 +227,80 @@ class TestBatchTripleQueries:
assert calls[2].kwargs["p"] is None
assert calls[2].kwargs["o"] == "ent-1"
@pytest.mark.asyncio
async def test_directions_assigned_correctly(self):
"""Each query position should produce the correct direction tag."""
triple = _make_triple("s", "p", "o")
call_count = 0
async def one_triple_each(**kwargs):
nonlocal call_count
call_count += 1
return [triple]
client = AsyncMock()
client.query_stream = one_triple_each
query = _make_query(triples_client=client)
result = await query.execute_batch_triple_queries(
["e1"], limit_per_entity=10
)
assert len(result) == 3
# Order matches query order: s-position, p-position, o-position
assert result[0][1] == Query.FROM_S
assert result[1][1] == Query.FROM_P
assert result[2][1] == Query.FROM_O
@pytest.mark.asyncio
async def test_directions_correct_for_multiple_entities(self):
"""Direction tags cycle correctly across multiple entities."""
triple = _make_triple("s", "p", "o")
client = AsyncMock()
client.query_stream = AsyncMock(return_value=[triple])
query = _make_query(triples_client=client)
result = await query.execute_batch_triple_queries(
["e1", "e2"], limit_per_entity=10
)
assert len(result) == 6
expected_directions = [
Query.FROM_S, Query.FROM_P, Query.FROM_O,
Query.FROM_S, Query.FROM_P, Query.FROM_O,
]
for (_, direction), expected in zip(result, expected_directions):
assert direction == expected
@pytest.mark.asyncio
async def test_direction_preserved_with_multiple_triples(self):
"""All triples from one query share the same direction."""
t1 = _make_triple("a", "p1", "b")
t2 = _make_triple("a", "p2", "c")
call_count = 0
async def multi_results(**kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
return [t1, t2]
return []
client = AsyncMock()
client.query_stream = multi_results
query = _make_query(triples_client=client)
result = await query.execute_batch_triple_queries(
["e1"], limit_per_entity=10
)
# First query (FROM_S) returns 2 triples, both should be FROM_S
assert len(result) == 2
assert result[0] == (t1, Query.FROM_S)
assert result[1] == (t2, Query.FROM_S)
class TestLRUCacheWithTTL:

View file

@ -49,7 +49,7 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase):
async def test_on_message_success(self, mock_pdf_loader_class, mock_producer, mock_consumer):
"""Test successful PDF processing"""
# Mock PDF content
pdf_content = b"fake pdf content"
pdf_content = b"%PDF-1.7\nfake pdf content"
pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
# Mock PyPDFLoader
@ -88,13 +88,55 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase):
# Verify triples were sent for each page (provenance)
assert mock_triples_flow.send.call_count == 2
@patch('trustgraph.base.librarian_client.Consumer')
@patch('trustgraph.base.librarian_client.Producer')
@patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader')
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
async def test_on_message_rejects_librarian_content_that_is_not_pdf(self, mock_pdf_loader_class, mock_producer, mock_consumer):
"""Test rejecting non-PDF content before invoking the PDF loader"""
html_content = b"<html><body>Not found</body></html>"
html_base64 = base64.b64encode(html_content)
mock_metadata = Metadata(id="test-doc")
mock_document = Document(metadata=mock_metadata, document_id="doc-123")
mock_msg = MagicMock()
mock_msg.value.return_value = mock_document
mock_output_flow = AsyncMock()
mock_triples_flow = AsyncMock()
mock_flow = MagicMock(side_effect=lambda name: {
"output": mock_output_flow,
"triples": mock_triples_flow,
}.get(name))
mock_flow.librarian.fetch_document_metadata = AsyncMock(
return_value=MagicMock(kind="application/pdf")
)
mock_flow.librarian.fetch_document_content = AsyncMock(
return_value=html_base64
)
mock_flow.librarian.save_child_document = AsyncMock()
config = {
'id': 'test-pdf-decoder',
'taskgroup': AsyncMock()
}
processor = Processor(**config)
await processor.on_message(mock_msg, None, mock_flow)
mock_pdf_loader_class.assert_not_called()
mock_output_flow.send.assert_not_called()
mock_triples_flow.send.assert_not_called()
mock_flow.librarian.save_child_document.assert_not_called()
@patch('trustgraph.base.librarian_client.Consumer')
@patch('trustgraph.base.librarian_client.Producer')
@patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader')
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
async def test_on_message_empty_pdf(self, mock_pdf_loader_class, mock_producer, mock_consumer):
"""Test handling of empty PDF"""
pdf_content = b"fake pdf content"
pdf_content = b"%PDF-1.7\nfake pdf content"
pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
mock_loader = MagicMock()
@ -126,7 +168,7 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase):
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
async def test_on_message_unicode_content(self, mock_pdf_loader_class, mock_producer, mock_consumer):
"""Test handling of unicode content in PDF"""
pdf_content = b"fake pdf content"
pdf_content = b"%PDF-1.7\nfake pdf content"
pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
mock_loader = MagicMock()

View file

@ -0,0 +1,198 @@
"""
Tests for the gateway audit middleware.
Verifies that gateway.request events are emitted with correct
metadata for success, error, and auth failure paths.
"""
import pytest
from unittest.mock import AsyncMock, MagicMock
from aiohttp import web
from aiohttp.test_utils import make_mocked_request
from trustgraph.gateway.audit import make_audit_middleware, _client_ip, _outcome_from_status
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_request(method="POST", path="/api/v1/config", headers=None,
remote="10.0.0.1"):
hdrs = headers or {}
req = make_mocked_request(method, path, headers=hdrs)
req._transport_peername = (remote, 0)
return req
# ---------------------------------------------------------------------------
# Pure helpers
# ---------------------------------------------------------------------------
class TestClientIp:
def test_uses_x_forwarded_for(self):
req = _make_request(headers={"X-Forwarded-For": "1.2.3.4, 10.0.0.1"})
assert _client_ip(req) == "1.2.3.4"
def test_falls_back_to_remote(self):
req = _make_request(remote="192.168.1.1")
assert _client_ip(req) == "192.168.1.1"
class TestOutcomeFromStatus:
def test_success_range(self):
assert _outcome_from_status(200) == "success"
assert _outcome_from_status(201) == "success"
assert _outcome_from_status(204) == "success"
assert _outcome_from_status(301) == "success"
def test_unauthenticated(self):
assert _outcome_from_status(401) == "unauthenticated"
def test_denied(self):
assert _outcome_from_status(403) == "denied"
def test_error(self):
assert _outcome_from_status(400) == "error"
assert _outcome_from_status(404) == "error"
assert _outcome_from_status(500) == "error"
# ---------------------------------------------------------------------------
# Middleware integration
# ---------------------------------------------------------------------------
class TestAuditMiddleware:
@pytest.fixture
def audit_publisher(self):
pub = AsyncMock()
pub.emit = AsyncMock()
return pub
@pytest.fixture
def middleware(self, audit_publisher):
return make_audit_middleware(audit_publisher)
@pytest.mark.asyncio
async def test_emits_event_on_success(self, middleware, audit_publisher):
async def handler(request):
return web.json_response({"ok": True})
req = _make_request()
await middleware(req, handler)
audit_publisher.emit.assert_called_once()
event_type, payload = audit_publisher.emit.call_args[0]
assert event_type == "gateway.request"
assert payload["method"] == "POST"
assert payload["path"] == "/api/v1/config"
assert payload["status_code"] == 200
assert payload["outcome"] == "success"
assert "request_id" in payload
assert "duration_ms" in payload
assert "error" not in payload
@pytest.mark.asyncio
async def test_emits_event_on_http_exception(self, middleware,
audit_publisher):
async def handler(request):
raise web.HTTPForbidden(reason="access denied")
req = _make_request()
with pytest.raises(web.HTTPForbidden):
await middleware(req, handler)
audit_publisher.emit.assert_called_once()
_, payload = audit_publisher.emit.call_args[0]
assert payload["status_code"] == 403
assert payload["outcome"] == "denied"
assert payload["error"] == "access denied"
@pytest.mark.asyncio
async def test_emits_event_on_unhandled_exception(self, middleware,
audit_publisher):
async def handler(request):
raise ValueError("boom")
req = _make_request()
with pytest.raises(ValueError):
await middleware(req, handler)
audit_publisher.emit.assert_called_once()
_, payload = audit_publisher.emit.call_args[0]
assert payload["status_code"] == 500
assert payload["outcome"] == "error"
assert payload["error"] == "ValueError"
@pytest.mark.asyncio
async def test_includes_identity_when_annotated(self, middleware,
audit_publisher):
async def handler(request):
request['audit_identity'] = "user:mark"
return web.json_response({"ok": True})
req = _make_request()
await middleware(req, handler)
_, payload = audit_publisher.emit.call_args[0]
assert payload["identity"] == "user:mark"
@pytest.mark.asyncio
async def test_includes_capability_when_annotated(self, middleware,
audit_publisher):
async def handler(request):
request['audit_capability'] = "config:read"
request['audit_workspace'] = "production"
return web.json_response({"ok": True})
req = _make_request()
await middleware(req, handler)
_, payload = audit_publisher.emit.call_args[0]
assert payload["capability"] == "config:read"
assert payload["workspace"] == "production"
@pytest.mark.asyncio
async def test_omits_unannotated_fields(self, middleware,
audit_publisher):
async def handler(request):
return web.json_response({"ok": True})
req = _make_request()
await middleware(req, handler)
_, payload = audit_publisher.emit.call_args[0]
assert "identity" not in payload
assert "capability" not in payload
assert "workspace" not in payload
@pytest.mark.asyncio
async def test_assigns_request_id(self, middleware, audit_publisher):
captured_id = None
async def handler(request):
nonlocal captured_id
captured_id = request.get('audit_request_id')
return web.json_response({"ok": True})
req = _make_request()
await middleware(req, handler)
assert captured_id is not None
_, payload = audit_publisher.emit.call_args[0]
assert payload["request_id"] == captured_id
@pytest.mark.asyncio
async def test_still_returns_response_if_emit_fails(self, middleware,
audit_publisher):
audit_publisher.emit.side_effect = RuntimeError("pub/sub down")
async def handler(request):
return web.json_response({"ok": True})
req = _make_request()
resp = await middleware(req, handler)
assert resp.status == 200

View file

@ -86,19 +86,19 @@ class TestVerifyJwtEddsa:
def test_valid_jwt_passes(self):
priv, pub = make_keypair()
claims = {
"sub": "user-1", "workspace": "default",
"sub": "user-1", "default_workspace": "default",
"iat": int(time.time()),
"exp": int(time.time()) + 60,
}
token = sign_jwt(priv, claims)
got = _verify_jwt_eddsa(token, pub)
assert got["sub"] == "user-1"
assert got["workspace"] == "default"
assert got["default_workspace"] == "default"
def test_expired_jwt_rejected(self):
priv, pub = make_keypair()
claims = {
"sub": "user-1", "workspace": "default",
"sub": "user-1", "default_workspace": "default",
"iat": int(time.time()) - 3600,
"exp": int(time.time()) - 1,
}
@ -110,7 +110,7 @@ class TestVerifyJwtEddsa:
priv_a, _ = make_keypair()
_, pub_b = make_keypair()
claims = {
"sub": "user-1", "workspace": "default",
"sub": "user-1", "default_workspace": "default",
"iat": int(time.time()),
"exp": int(time.time()) + 60,
}
@ -130,7 +130,7 @@ class TestVerifyJwtEddsa:
# since we expect it to bail before verifying.
header = {"alg": "HS256", "typ": "JWT", "kid": "x"}
payload = {
"sub": "user-1", "workspace": "default",
"sub": "user-1", "default_workspace": "default",
"iat": int(time.time()), "exp": int(time.time()) + 60,
}
h = _b64url(json.dumps(header, separators=(",", ":")).encode())
@ -148,11 +148,11 @@ class TestIdentity:
def test_fields(self):
i = Identity(
handle="u", workspace="w",
handle="u", default_workspace="w",
principal_id="u", source="api-key",
)
assert i.handle == "u"
assert i.workspace == "w"
assert i.default_workspace == "w"
assert i.principal_id == "u"
assert i.source == "api-key"
@ -208,7 +208,7 @@ class TestIamAuthDispatch:
async def test_valid_jwt_resolves_to_identity(self):
priv, pub = make_keypair()
claims = {
"sub": "user-1", "workspace": "default",
"sub": "user-1", "default_workspace": "default",
"iat": int(time.time()),
"exp": int(time.time()) + 60,
}
@ -221,7 +221,7 @@ class TestIamAuthDispatch:
make_request(f"Bearer {token}")
)
assert ident.handle == "user-1"
assert ident.workspace == "default"
assert ident.default_workspace == "default"
assert ident.principal_id == "user-1"
assert ident.source == "jwt"
@ -231,7 +231,7 @@ class TestIamAuthDispatch:
# must not validate — even ones that would otherwise pass.
priv, _ = make_keypair()
claims = {
"sub": "user-1", "workspace": "default",
"sub": "user-1", "default_workspace": "default",
"iat": int(time.time()), "exp": int(time.time()) + 60,
}
token = sign_jwt(priv, claims)
@ -244,7 +244,7 @@ class TestIamAuthDispatch:
async def test_api_key_path(self):
auth = IamAuth(backend=Mock())
async def fake_resolve(api_key):
async def fake_resolve(api_key, **kwargs):
assert api_key == "tg_testkey"
# Roles are returned by the regime as a hint but the
# gateway ignores them — kept here so the resolve
@ -259,7 +259,7 @@ class TestIamAuthDispatch:
make_request("Bearer tg_testkey")
)
assert ident.handle == "user-xyz"
assert ident.workspace == "default"
assert ident.default_workspace == "default"
assert ident.principal_id == "user-xyz"
assert ident.source == "api-key"
@ -309,7 +309,7 @@ class TestApiKeyCache:
seen = []
async def fake_with_client(op):
async def resolve(plaintext):
async def resolve(plaintext, **kwargs):
seen.append(plaintext)
return ("u-" + plaintext, "default", ["reader"])
return await op(Mock(resolve_api_key=resolve))
@ -338,9 +338,9 @@ class TestAuthorise:
decision for the regime's TTL (clamped above), and raises 403
on deny / 401 on regime error (fail closed)."""
def _make_identity(self, handle="u-1", workspace="default"):
def _make_identity(self, handle="u-1", default_workspace="default"):
return Identity(
handle=handle, workspace=workspace,
handle=handle, default_workspace=default_workspace,
principal_id=handle, source="api-key",
)

View file

@ -25,11 +25,11 @@ from trustgraph.gateway.capabilities import (
class _Identity:
"""Stand-in for auth.Identity — under the IAM contract it has
just ``handle``, ``workspace``, ``principal_id``, ``source``."""
just ``handle``, ``default_workspace``, ``principal_id``, ``source``."""
def __init__(self, handle="user-1", workspace="default"):
def __init__(self, handle="user-1", default_workspace="default"):
self.handle = handle
self.workspace = workspace
self.default_workspace = default_workspace
self.principal_id = handle
self.source = "api-key"
@ -105,14 +105,14 @@ class TestEnforceWorkspace:
async def test_default_fills_from_identity(self):
data = {"operation": "x"}
auth = _allow_auth()
await enforce_workspace(data, _Identity(workspace="default"), auth)
await enforce_workspace(data, _Identity(default_workspace="default"), auth)
assert data["workspace"] == "default"
@pytest.mark.asyncio
async def test_caller_supplied_workspace_kept(self):
data = {"workspace": "acme", "operation": "x"}
auth = _allow_auth()
await enforce_workspace(data, _Identity(workspace="default"), auth)
await enforce_workspace(data, _Identity(default_workspace="default"), auth)
assert data["workspace"] == "acme"
@pytest.mark.asyncio

View file

@ -28,7 +28,7 @@ TEST_CAP = "graph:write"
def _valid_identity():
return Identity(
handle="test-user",
workspace="default",
default_workspace="default",
principal_id="test-user",
source="api-key",
)

View file

@ -0,0 +1,345 @@
"""
Tests for IAM service audit event emission.
Verifies that the IAM Processor emits correctly categorised audit
events (iam.authenticate, iam.authorise, iam.management) after
handling each request type.
"""
import asyncio
import pytest
from unittest.mock import AsyncMock, MagicMock, Mock
from trustgraph.schema import IamRequest, IamResponse, Error
from trustgraph.iam.service.service import Processor
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_msg(request, msg_id="test-msg-1"):
msg = MagicMock()
msg.value.return_value = request
msg.properties.return_value = {"id": msg_id}
return msg
def _make_processor():
"""Create a Processor with stubbed dependencies."""
proc = object.__new__(Processor)
proc.id = "test-iam-svc"
proc.iam_response_producer = AsyncMock()
proc.audit = AsyncMock()
return proc
# ---------------------------------------------------------------------------
# Authentication events
# ---------------------------------------------------------------------------
class TestAuthenticateAuditEvents:
@pytest.mark.asyncio
async def test_resolve_api_key_success(self):
proc = _make_processor()
resp = IamResponse(
resolved_user_id="user-123",
resolved_default_workspace="default",
)
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="resolve-api-key",
api_key="tg_test",
request_id="req-1",
client_ip="10.0.0.1",
)
await proc.on_iam_request(_make_msg(req), None, None)
proc.audit.emit.assert_called_once()
event_type, payload = proc.audit.emit.call_args[0]
assert event_type == "iam.authenticate"
assert payload["credential_type"] == "api-key"
assert payload["identity"] == "user-123"
assert payload["outcome"] == "success"
assert payload["request_id"] == "req-1"
assert payload["client_ip"] == "10.0.0.1"
assert "failure_reason" not in payload
@pytest.mark.asyncio
async def test_resolve_api_key_failure(self):
proc = _make_processor()
resp = IamResponse(error=Error(type="auth-failed", message="unknown"))
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="resolve-api-key",
api_key="tg_bad",
request_id="req-2",
client_ip="10.0.0.2",
)
await proc.on_iam_request(_make_msg(req), None, None)
_, payload = proc.audit.emit.call_args[0]
assert payload["outcome"] == "failure"
assert payload["identity"] == "unknown"
assert payload["failure_reason"] == "auth-failed"
@pytest.mark.asyncio
async def test_login_emits_authenticate(self):
proc = _make_processor()
resp = IamResponse(
resolved_user_id="user-456",
jwt="eyJ...",
jwt_expires="2026-07-06T10:00:00Z",
)
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="login",
username="mark",
password="secret",
request_id="req-3",
)
await proc.on_iam_request(_make_msg(req), None, None)
event_type, payload = proc.audit.emit.call_args[0]
assert event_type == "iam.authenticate"
assert payload["credential_type"] == "login-password"
@pytest.mark.asyncio
async def test_anonymous_emits_authenticate(self):
proc = _make_processor()
resp = IamResponse(error=Error(type="auth-failed", message="no"))
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(operation="authenticate-anonymous")
await proc.on_iam_request(_make_msg(req), None, None)
event_type, payload = proc.audit.emit.call_args[0]
assert event_type == "iam.authenticate"
assert payload["credential_type"] == "anonymous"
assert payload["outcome"] == "failure"
# ---------------------------------------------------------------------------
# Authorise events
# ---------------------------------------------------------------------------
class TestAuthoriseAuditEvents:
@pytest.mark.asyncio
async def test_authorise_allow(self):
proc = _make_processor()
resp = IamResponse(decision_allow=True, decision_ttl_seconds=60)
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="authorise",
user_id="user-1",
capability="config:read",
resource_json='{"workspace": "production"}',
request_id="req-4",
)
await proc.on_iam_request(_make_msg(req), None, None)
event_type, payload = proc.audit.emit.call_args[0]
assert event_type == "iam.authorise"
assert payload["outcome"] == "allow"
assert payload["capability"] == "config:read"
assert payload["identity"] == "user-1"
assert payload["workspace"] == "production"
assert "denial_reason" not in payload
@pytest.mark.asyncio
async def test_authorise_deny(self):
proc = _make_processor()
resp = IamResponse(decision_allow=False, decision_ttl_seconds=5)
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="authorise",
user_id="user-2",
capability="users:admin",
resource_json='{}',
request_id="req-5",
)
await proc.on_iam_request(_make_msg(req), None, None)
_, payload = proc.audit.emit.call_args[0]
assert payload["outcome"] == "deny"
assert payload["denial_reason"] == "capability-not-in-role"
@pytest.mark.asyncio
async def test_authorise_extracts_workspace_from_resource_json(self):
proc = _make_processor()
resp = IamResponse(decision_allow=True, decision_ttl_seconds=60)
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="authorise",
user_id="user-1",
capability="graph-rag:query",
resource_json='{"workspace": "engineering"}',
)
await proc.on_iam_request(_make_msg(req), None, None)
_, payload = proc.audit.emit.call_args[0]
assert payload["workspace"] == "engineering"
@pytest.mark.asyncio
async def test_authorise_omits_empty_workspace(self):
proc = _make_processor()
resp = IamResponse(decision_allow=True, decision_ttl_seconds=60)
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="authorise",
user_id="user-1",
capability="metrics:read",
resource_json='{}',
)
await proc.on_iam_request(_make_msg(req), None, None)
_, payload = proc.audit.emit.call_args[0]
assert "workspace" not in payload
# ---------------------------------------------------------------------------
# Management events
# ---------------------------------------------------------------------------
class TestManagementAuditEvents:
@pytest.mark.asyncio
async def test_create_user_emits_management(self):
proc = _make_processor()
resp = IamResponse()
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="create-user",
actor="admin-1",
user_id="new-user",
workspace="default",
request_id="req-6",
)
await proc.on_iam_request(_make_msg(req), None, None)
event_type, payload = proc.audit.emit.call_args[0]
assert event_type == "iam.management"
assert payload["operation"] == "create-user"
assert payload["actor"] == "admin-1"
assert payload["target_identity"] == "new-user"
assert payload["target_workspace"] == "default"
assert payload["outcome"] == "success"
@pytest.mark.asyncio
async def test_revoke_api_key_emits_management(self):
proc = _make_processor()
resp = IamResponse()
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="revoke-api-key",
actor="admin-1",
key_id="key-abc",
workspace="production",
request_id="req-7",
)
await proc.on_iam_request(_make_msg(req), None, None)
event_type, payload = proc.audit.emit.call_args[0]
assert event_type == "iam.management"
assert payload["operation"] == "revoke-api-key"
assert payload["outcome"] == "success"
@pytest.mark.asyncio
async def test_management_error_outcome(self):
proc = _make_processor()
resp = IamResponse(
error=Error(type="not-found", message="user not found"),
)
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="delete-user",
actor="admin-1",
user_id="ghost",
)
await proc.on_iam_request(_make_msg(req), None, None)
_, payload = proc.audit.emit.call_args[0]
assert payload["outcome"] == "error"
@pytest.mark.asyncio
async def test_management_omits_empty_target_fields(self):
proc = _make_processor()
resp = IamResponse()
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(
operation="rotate-signing-key",
actor="admin-1",
)
await proc.on_iam_request(_make_msg(req), None, None)
_, payload = proc.audit.emit.call_args[0]
assert "target_identity" not in payload
assert "target_workspace" not in payload
# ---------------------------------------------------------------------------
# Non-audited operations
# ---------------------------------------------------------------------------
class TestNonAuditedOperations:
@pytest.mark.asyncio
async def test_whoami_does_not_emit(self):
proc = _make_processor()
resp = IamResponse()
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(operation="whoami", actor="user-1")
await proc.on_iam_request(_make_msg(req), None, None)
proc.audit.emit.assert_not_called()
@pytest.mark.asyncio
async def test_list_users_does_not_emit(self):
proc = _make_processor()
resp = IamResponse()
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(operation="list-users", workspace="default")
await proc.on_iam_request(_make_msg(req), None, None)
proc.audit.emit.assert_not_called()
@pytest.mark.asyncio
async def test_get_signing_key_does_not_emit(self):
proc = _make_processor()
resp = IamResponse(signing_key_public="PEM...")
proc.iam = MagicMock()
proc.iam.handle = AsyncMock(return_value=resp)
req = IamRequest(operation="get-signing-key-public")
await proc.on_iam_request(_make_msg(req), None, None)
proc.audit.emit.assert_not_called()

View file

@ -32,7 +32,7 @@ class TestAuthenticateAnonymous:
)
assert resp.error is None
assert resp.resolved_user_id == "anon"
assert resp.resolved_workspace == "ws"
assert resp.resolved_default_workspace == "ws"
assert "admin" in list(resp.resolved_roles)
@pytest.mark.asyncio
@ -44,7 +44,7 @@ class TestAuthenticateAnonymous:
_make_request(operation="authenticate-anonymous")
)
assert resp.resolved_user_id == "dev-user"
assert resp.resolved_workspace == "dev-ws"
assert resp.resolved_default_workspace == "dev-ws"
class TestResolveApiKey:
@ -57,7 +57,7 @@ class TestResolveApiKey:
)
assert resp.error is None
assert resp.resolved_user_id == "anonymous"
assert resp.resolved_workspace == "default"
assert resp.resolved_default_workspace == "default"
class TestAuthorise:

View file

@ -107,6 +107,7 @@ class TestGraphRagDagStructure:
embeddings_client = AsyncMock()
graph_embeddings_client = AsyncMock()
triples_client = AsyncMock()
reranker_client = AsyncMock()
embeddings_client.embed.return_value = [[0.1, 0.2]]
graph_embeddings_client.query.return_value = [
@ -121,27 +122,22 @@ class TestGraphRagDagStructure:
]
triples_client.query.return_value = []
result = MagicMock()
result.document_id = "0"
result.query_id = "0"
result.score = 0.95
reranker_client.rerank.return_value = [result]
async def mock_prompt(template_id, variables=None, **kwargs):
if template_id == "extract-concepts":
return PromptResult(response_type="text", text="concept")
elif template_id == "kg-edge-scoring":
edges = variables.get("knowledge", [])
return PromptResult(
response_type="jsonl",
objects=[{"id": e["id"], "score": 10} for e in edges],
)
elif template_id == "kg-edge-reasoning":
edges = variables.get("knowledge", [])
return PromptResult(
response_type="jsonl",
objects=[{"id": e["id"], "reasoning": "relevant"} for e in edges],
)
elif template_id == "kg-synthesis":
return PromptResult(response_type="text", text="Answer.")
return PromptResult(response_type="text", text="")
prompt_client.prompt.side_effect = mock_prompt
return prompt_client, embeddings_client, graph_embeddings_client, triples_client
return (prompt_client, embeddings_client, graph_embeddings_client,
triples_client, reranker_client)
@pytest.mark.asyncio
async def test_dag_chain(self, mock_clients):
@ -152,7 +148,7 @@ class TestGraphRagDagStructure:
events.append({"explain_id": explain_id, "triples": triples})
await rag.query(
query="test", explain_callback=explain_cb, edge_score_limit=0,
query="test", explain_callback=explain_cb,
)
dag = _collect_events(events)

View file

@ -333,8 +333,8 @@ class TestUnifiedTableQueries:
"""Test queries against the unified rows table"""
@pytest.mark.asyncio
@patch('trustgraph.query.rows.cassandra.service.async_execute', new_callable=AsyncMock)
async def test_query_with_index_match(self, mock_async_execute):
@patch('trustgraph.query.rows.cassandra.service.async_execute_paged', new_callable=AsyncMock)
async def test_query_with_index_match(self, mock_async_execute_paged):
"""Test query execution with matching index"""
processor = MagicMock()
processor.session = MagicMock()
@ -344,10 +344,10 @@ class TestUnifiedTableQueries:
processor.find_matching_index = Processor.find_matching_index.__get__(processor, Processor)
processor.query_cassandra = Processor.query_cassandra.__get__(processor, Processor)
# Mock async_execute to return test data
# Mock async_execute_paged to return test data (list of pages)
mock_row = MagicMock()
mock_row.data = {"id": "123", "name": "Test Product", "category": "electronics"}
mock_async_execute.return_value = [mock_row]
mock_async_execute_paged.return_value = [[mock_row]]
schema = RowSchema(
name="products",
@ -370,10 +370,10 @@ class TestUnifiedTableQueries:
# Verify Cassandra was connected and queried
processor.connect_cassandra.assert_called_once()
mock_async_execute.assert_called_once()
mock_async_execute_paged.assert_called_once()
# Verify query structure - should query unified rows table
call_args = mock_async_execute.call_args
call_args = mock_async_execute_paged.call_args
query = call_args[0][1]
params = call_args[0][2]
@ -394,8 +394,8 @@ class TestUnifiedTableQueries:
assert results[0]["category"] == "electronics"
@pytest.mark.asyncio
@patch('trustgraph.query.rows.cassandra.service.async_execute', new_callable=AsyncMock)
async def test_query_without_index_match(self, mock_async_execute):
@patch('trustgraph.query.rows.cassandra.service.async_scan', new_callable=AsyncMock)
async def test_query_without_index_match(self, mock_async_scan):
"""Test query execution without matching index (scan mode)"""
processor = MagicMock()
processor.session = MagicMock()
@ -406,12 +406,10 @@ class TestUnifiedTableQueries:
processor._matches_filters = Processor._matches_filters.__get__(processor, Processor)
processor.query_cassandra = Processor.query_cassandra.__get__(processor, Processor)
# Mock async_execute to return test data
# Mock async_scan to return filtered test data
mock_row1 = MagicMock()
mock_row1.data = {"id": "1", "name": "Product A", "price": "100"}
mock_row2 = MagicMock()
mock_row2.data = {"id": "2", "name": "Product B", "price": "200"}
mock_async_execute.return_value = [mock_row1, mock_row2]
mock_async_scan.return_value = [mock_row1]
schema = RowSchema(
name="products",
@ -432,13 +430,16 @@ class TestUnifiedTableQueries:
limit=10
)
# Query should use ALLOW FILTERING for scan
call_args = mock_async_execute.call_args
# Verify async_scan was called
mock_async_scan.assert_called_once()
# Verify query structure
call_args = mock_async_scan.call_args
query = call_args[0][1]
assert "ALLOW FILTERING" in query
# Should post-filter results
# Should return filtered results
assert len(results) == 1
assert results[0]["name"] == "Product A"

View file

@ -259,6 +259,8 @@ class TestGraphEmbeddingsNullProtection:
proc.collection_exists = MagicMock(return_value=True)
proc._cache_lock = asyncio.Lock()
proc._known_collections = set()
proc.replication_factor = 1
proc.shard_number = 1
msg = MagicMock()
msg.metadata.collection = "graphs"

View file

@ -101,27 +101,27 @@ class TestQuery:
assert query.rag == mock_rag
assert query.collection == "test_collection"
assert query.verbose is False
assert query.doc_limit == 20 # Default value
assert query.fetch_limit == 20 # Default value
def test_query_initialization_with_custom_doc_limit(self):
"""Test Query initialization with custom doc_limit"""
def test_query_initialization_with_custom_fetch_limit(self):
"""Test Query initialization with custom fetch_limit"""
# Create mock DocumentRag
mock_rag = MagicMock()
# Initialize Query with custom doc_limit
# Initialize Query with custom fetch_limit
query = Query(
rag=mock_rag,
workspace="test_workspace",
collection="custom_collection",
verbose=True,
doc_limit=50
fetch_limit=50
)
# Verify initialization
assert query.rag == mock_rag
assert query.collection == "custom_collection"
assert query.verbose is True
assert query.doc_limit == 50
assert query.fetch_limit == 50
@pytest.mark.asyncio
async def test_extract_concepts(self):
@ -224,7 +224,7 @@ class TestQuery:
workspace="test_workspace",
collection="test_collection",
verbose=False,
doc_limit=15
fetch_limit=15
)
# Call get_docs with concepts list
@ -377,7 +377,7 @@ class TestQuery:
workspace="test_workspace",
collection="test_collection",
verbose=True,
doc_limit=5
fetch_limit=5
)
# Call get_docs with concepts
@ -615,7 +615,7 @@ class TestQuery:
workspace="test_workspace",
collection="test_collection",
verbose=False,
doc_limit=10
fetch_limit=10
)
docs, chunk_ids = await query.get_docs(["concept A", "concept B"])

View file

@ -0,0 +1,92 @@
from trustgraph.retrieval.document_rag.rerank import (
RerankCandidate, normalize_candidate_scores, mmr_select,
_pair_diversity_penalty
)
def candidate(index, chunk_id, text, score):
return RerankCandidate(
index=index,
chunk_id=chunk_id,
text=text,
reranker_score=score,
)
def test_normalize_candidate_scores_min_max_scales_raw_scores():
candidates = [
candidate(0, "a", "alpha", -2.0),
candidate(1, "b", "beta", 0.0),
candidate(2, "c", "gamma", 4.0),
]
normalized = normalize_candidate_scores(candidates)
assert normalized[0].normalized_score == 0.0
assert normalized[1].normalized_score == 1.0 / 3.0
assert normalized[2].normalized_score == 1.0
def test_normalize_candidate_scores_handles_equal_scores():
candidates = [
candidate(0, "a", "alpha", 3.0),
candidate(1, "b", "beta", 3.0),
candidate(2, "c", "gamma", 3.0),
]
normalized = normalize_candidate_scores(candidates)
assert [c.normalized_score for c in normalized] == [0.5, 0.5, 0.5]
def test_mmr_select_limits_results():
candidates = [
candidate(0, "a", "alpha policy", 0.9),
candidate(1, "b", "beta refund", 0.8),
candidate(2, "c", "gamma shipping", 0.7),
]
selected = mmr_select(candidates, limit=2)
assert len(selected) == 2
def test_mmr_select_prefers_highest_reranker_score_first():
candidates = [
candidate(0, "a", "weakly relevant text", 0.1),
candidate(1, "b", "strongly relevant answer", 10.0),
candidate(2, "c", "medium relevant text", 5.0),
]
selected = mmr_select(candidates, limit=1)
assert selected[0].chunk_id == "b"
def test_mmr_select_penalizes_near_duplicate_chunks():
candidates = [
candidate(0, "a", "apple banana fruit return policy", 1.00),
candidate(1, "b", "apple banana fruit return policy duplicate", 0.95),
candidate(2, "c", "engine motor vehicle warranty", 0.90),
]
selected = mmr_select(
candidates,
limit=2,
lambda_mult=0.2,
token_overlap_weight=1.0,
)
assert [c.chunk_id for c in selected] == ["a", "c"]
def test_pair_diversity_penalty_is_clamped():
left = candidate(0, "a", "same same same", 1.0)
right = candidate(1, "b", "same same same", 0.9)
penalty = _pair_diversity_penalty(
left,
right,
token_overlap_weight=10.0,
)
assert penalty == 1.0

View file

@ -0,0 +1,211 @@
"""
Tests for the retrieval-mode dispatch in DocumentRag (issue: hybrid
BM25 + vector retrieval).
Covered behaviours:
1. Default: retrieval_mode="vector" never touches the keyword client and
produces the same chunks as before the sparse path is strictly opt-in.
2. keyword: only the keyword index is queried (no vector-store query, no
embedding of concepts); chunk order follows the BM25 ranking.
3. hybrid: both paths run and are fused by weighted RRF on chunk_id; a
keyword-path failure degrades to vector-only instead of failing the
query.
4. Constructing with keyword/hybrid but no keyword client is an error.
Pure orchestration tests: all subsidiary clients are stubs.
"""
import pytest
from unittest.mock import AsyncMock
from trustgraph.retrieval.document_rag.document_rag import (
DocumentRag, rrf_fuse, RRF_K,
)
from trustgraph.base import PromptResult
from trustgraph.schema import ChunkMatch
CONTENT = {
"v1": "vector chunk one",
"v2": "vector chunk two",
"k1": "keyword chunk one",
"both": "chunk found by both paths",
}
def build_clients(vector_ids, keyword_ids):
prompt_client = AsyncMock()
embeddings_client = AsyncMock()
doc_embeddings_client = AsyncMock()
kw_index_client = AsyncMock()
fetch_chunk = AsyncMock()
async def mock_prompt(template_id, variables=None, **kwargs):
if template_id == "extract-concepts":
return PromptResult(response_type="text", text="concept")
return PromptResult(response_type="text", text="")
prompt_client.prompt.side_effect = mock_prompt
prompt_client.document_prompt.return_value = PromptResult(
response_type="text", text="answer",
)
embeddings_client.embed.return_value = [[0.1, 0.2]]
doc_embeddings_client.query.return_value = [
ChunkMatch(chunk_id=c) for c in vector_ids
]
kw_index_client.query.return_value = [
ChunkMatch(chunk_id=c, score=1.0) for c in keyword_ids
]
fetch_chunk.side_effect = lambda chunk_id: CONTENT[chunk_id]
return (
prompt_client, embeddings_client, doc_embeddings_client,
kw_index_client, fetch_chunk,
)
def build_rag(vector_ids, keyword_ids, **kwargs):
prompt, embeddings, doc_embeddings, kw, fetch = build_clients(
vector_ids, keyword_ids,
)
rag = DocumentRag(
prompt_client=prompt,
embeddings_client=embeddings,
doc_embeddings_client=doc_embeddings,
fetch_chunk=fetch,
kw_index_client=kw,
**kwargs,
)
return rag, doc_embeddings, kw, embeddings, prompt
# ---------------------------------------------------------------------------
# rrf_fuse
# ---------------------------------------------------------------------------
class TestRrfFuse:
def test_chunk_in_both_lists_outranks_single_list_leaders(self):
a = ChunkMatch("a")
b = ChunkMatch("b")
both = ChunkMatch("both")
fused = rrf_fuse([[a, both], [both, b]], [1.0, 1.0], 10)
assert [m.chunk_id for m in fused][0] == "both"
assert {m.chunk_id for m in fused} == {"a", "b", "both"}
def test_weights_bias_the_fusion(self):
a, b = ChunkMatch("a"), ChunkMatch("b")
fused = rrf_fuse([[a], [b]], [1.0, 10.0], 10)
assert [m.chunk_id for m in fused] == ["b", "a"]
def test_limit_truncates(self):
matches = [ChunkMatch(f"c{i}") for i in range(5)]
assert len(rrf_fuse([matches], [1.0], 2)) == 2
def test_cross_list_accumulation_beats_single_top_rank(self):
# b sums 1/(K+2) + 1/(K+3) across two lists, beating the single
# 1/(K+1) that a gets — the accumulation property that
# distinguishes RRF from a best-rank merge.
a, b, x, y = (ChunkMatch(c) for c in "abxy")
fused = rrf_fuse([[a, b], [x, y, b]], [1.0, 1.0], 10)
assert fused[0].chunk_id == "b"
assert 1 / (RRF_K + 2) + 1 / (RRF_K + 3) > 1 / (RRF_K + 1)
def test_empty_chunk_ids_are_skipped(self):
fused = rrf_fuse([[ChunkMatch(""), ChunkMatch("a")]], [1.0], 10)
assert [m.chunk_id for m in fused] == ["a"]
# ---------------------------------------------------------------------------
# Mode dispatch through DocumentRag.query()
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_vector_mode_never_touches_keyword_client():
rag, doc_embeddings, kw, _, prompt = build_rag(
["v1", "v2"], ["k1"], retrieval_mode="vector",
)
await rag.query("question")
kw.query.assert_not_called()
doc_embeddings.query.assert_called()
docs = prompt.document_prompt.call_args.kwargs["documents"]
assert docs == [CONTENT["v1"], CONTENT["v2"]]
@pytest.mark.asyncio
async def test_default_mode_is_vector_with_no_keyword_client():
prompt, embeddings, doc_embeddings, _, fetch = build_clients(
["v1"], [],
)
rag = DocumentRag(
prompt_client=prompt,
embeddings_client=embeddings,
doc_embeddings_client=doc_embeddings,
fetch_chunk=fetch,
)
await rag.query("question")
docs = prompt.document_prompt.call_args.kwargs["documents"]
assert docs == [CONTENT["v1"]]
@pytest.mark.asyncio
async def test_keyword_mode_skips_vector_store_and_embeddings():
rag, doc_embeddings, kw, embeddings, prompt = build_rag(
["v1", "v2"], ["k1", "both"], retrieval_mode="keyword",
)
await rag.query("what does clause 7.3.2 say")
doc_embeddings.query.assert_not_called()
embeddings.embed.assert_not_called()
# No dense path -> no concept-extraction LLM call either
prompt.prompt.assert_not_called()
# The sparse path searches the raw query text, not extracted concepts
assert kw.query.call_args.kwargs["query"] == "what does clause 7.3.2 say"
docs = prompt.document_prompt.call_args.kwargs["documents"]
assert docs == [CONTENT["k1"], CONTENT["both"]]
@pytest.mark.asyncio
async def test_hybrid_mode_fuses_both_paths():
# both appears in both rankings, so RRF must put it first
rag, doc_embeddings, kw, _, prompt = build_rag(
["v1", "both"], ["both", "k1"], retrieval_mode="hybrid",
)
await rag.query("question")
doc_embeddings.query.assert_called()
kw.query.assert_called()
docs = prompt.document_prompt.call_args.kwargs["documents"]
assert docs[0] == CONTENT["both"]
assert set(docs) == {CONTENT["both"], CONTENT["v1"], CONTENT["k1"]}
@pytest.mark.asyncio
async def test_hybrid_degrades_to_vector_when_keyword_path_fails():
rag, doc_embeddings, kw, _, prompt = build_rag(
["v1", "v2"], [], retrieval_mode="hybrid",
)
kw.query.side_effect = RuntimeError("keyword index down")
await rag.query("question")
docs = prompt.document_prompt.call_args.kwargs["documents"]
assert docs == [CONTENT["v1"], CONTENT["v2"]]
def test_non_vector_mode_without_client_is_an_error():
prompt, embeddings, doc_embeddings, _, fetch = build_clients([], [])
for mode in ("keyword", "hybrid"):
with pytest.raises(ValueError):
DocumentRag(
prompt_client=prompt,
embeddings_client=embeddings,
doc_embeddings_client=doc_embeddings,
fetch_chunk=fetch,
retrieval_mode=mode,
)

View file

@ -0,0 +1,550 @@
"""
Tests for the optional cross-encoder reranking pass in DocumentRag.query().
Two behaviours are covered:
1. No-op: when no reranker_client is wired (the default), query() must feed
the LLM the exact same chunks, in the same order, that retrieval produced
- byte-identical to the pre-reranker behaviour - and must NOT emit a
chunk-selection provenance event.
2. Rerank: when a reranker_client is wired, the retrieved chunks are reordered
and truncated according to the reranker's results, the LLM receives the
reranked top-N, and a tg:ChunkSelection (focus) provenance event is emitted
carrying the per-surviving-chunk scores and chunk references.
These are pure orchestration tests - the reranker is a stub, so there is no
torch / network dependency.
"""
import pytest
from unittest.mock import AsyncMock
from dataclasses import dataclass
from trustgraph.retrieval.document_rag.document_rag import DocumentRag
from trustgraph.base import PromptResult
from trustgraph.schema import RerankerResult
from trustgraph.provenance.namespaces import (
RDF_TYPE, PROV_WAS_DERIVED_FROM,
TG_DOC_RAG_QUESTION, TG_GROUNDING, TG_EXPLORATION,
TG_FOCUS, TG_SYNTHESIS,
TG_CHUNK_SELECTION, TG_SELECTED_CHUNK, TG_SCORE, TG_DOCUMENT,
)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def find_triple(triples, predicate, subject=None):
for t in triples:
if t.p.iri == predicate:
if subject is None or t.s.iri == subject:
return t
return None
def find_triples(triples, predicate, subject=None):
return [
t for t in triples
if t.p.iri == predicate
and (subject is None or t.s.iri == subject)
]
def has_type(triples, subject, rdf_type):
return any(
t.s.iri == subject and t.p.iri == RDF_TYPE and t.o.iri == rdf_type
for t in triples
)
def derived_from(triples, subject):
t = find_triple(triples, PROV_WAS_DERIVED_FROM, subject)
return t.o.iri if t else None
@dataclass
class ChunkMatch:
"""Mimics the result from doc_embeddings_client.query()."""
chunk_id: str
# ---------------------------------------------------------------------------
# Fixtures: three retrievable chunks
# ---------------------------------------------------------------------------
CHUNK_A = "urn:chunk:policy-doc-1:chunk-0"
CHUNK_B = "urn:chunk:policy-doc-1:chunk-1"
CHUNK_C = "urn:chunk:policy-doc-1:chunk-2"
CHUNK_A_CONTENT = "Customers may return items within 30 days of purchase."
CHUNK_B_CONTENT = "Our stores are open from 9am to 5pm on weekdays."
CHUNK_C_CONTENT = "Refunds are processed to the original payment method."
# Retrieval (post-dedupe) order is A, B, C.
ORDERED_CONTENT = [CHUNK_A_CONTENT, CHUNK_B_CONTENT, CHUNK_C_CONTENT]
ORDERED_CHUNK_IDS = [CHUNK_A, CHUNK_B, CHUNK_C]
def build_mock_clients():
"""
Build mock subsidiary clients for a document-rag query returning three
distinct chunks (A, B, C) in that order.
"""
prompt_client = AsyncMock()
embeddings_client = AsyncMock()
doc_embeddings_client = AsyncMock()
fetch_chunk = AsyncMock()
async def mock_prompt(template_id, variables=None, **kwargs):
if template_id == "extract-concepts":
return PromptResult(response_type="text", text="return policy\nrefund")
return PromptResult(response_type="text", text="")
prompt_client.prompt.side_effect = mock_prompt
embeddings_client.embed.return_value = [[0.1, 0.2], [0.3, 0.4]]
# Each concept query returns the same three chunks; dedupe keeps A, B, C.
doc_embeddings_client.query.return_value = [
ChunkMatch(chunk_id=CHUNK_A),
ChunkMatch(chunk_id=CHUNK_B),
ChunkMatch(chunk_id=CHUNK_C),
]
async def mock_fetch(chunk_id):
return {
CHUNK_A: CHUNK_A_CONTENT,
CHUNK_B: CHUNK_B_CONTENT,
CHUNK_C: CHUNK_C_CONTENT,
}[chunk_id]
fetch_chunk.side_effect = mock_fetch
prompt_client.document_prompt.return_value = PromptResult(
response_type="text",
text="Items can be returned within 30 days for a full refund.",
)
return prompt_client, embeddings_client, doc_embeddings_client, fetch_chunk
class StubReranker:
"""
Stub reranker_client mirroring RerankerClient.rerank(): returns a fixed,
pre-sorted, truncated list of RerankerResult - exactly the contract the
flashrank service guarantees (sorted desc by score, truncated to limit).
"""
def __init__(self, results):
self._results = results
self.calls = []
async def rerank(self, queries, documents, limit=10, timeout=300):
self.calls.append(
{"queries": queries, "documents": documents, "limit": limit}
)
return self._results
# ---------------------------------------------------------------------------
# 1. No-op: reranker_client=None must not change anything
# ---------------------------------------------------------------------------
class TestRerankNoOp:
@pytest.mark.asyncio
async def test_documents_passed_to_llm_are_unchanged(self):
"""
With no reranker wired, document_prompt must receive the retrieved
chunks in the original order and length.
"""
clients = build_mock_clients()
rag = DocumentRag(*clients) # reranker_client defaults to None
await rag.query(query="What is the return policy?")
call = rag.prompt_client.document_prompt.call_args
passed_docs = call.kwargs["documents"]
assert passed_docs == ORDERED_CONTENT
@pytest.mark.asyncio
async def test_no_chunk_selection_event_emitted(self):
"""
Without a reranker, the provenance chain is the original 4 stages:
question, grounding, exploration, synthesis - no focus stage.
"""
clients = build_mock_clients()
rag = DocumentRag(*clients)
events = []
async def explain_callback(triples, explain_id):
events.append({"triples": triples, "explain_id": explain_id})
await rag.query(
query="What is the return policy?",
explain_callback=explain_callback,
)
assert len(events) == 4
types = [
TG_DOC_RAG_QUESTION, TG_GROUNDING, TG_EXPLORATION, TG_SYNTHESIS,
]
for i, expected in enumerate(types):
assert has_type(events[i]["triples"], events[i]["explain_id"], expected)
# No chunk-selection entity anywhere.
for e in events:
assert not any(
t.o.iri == TG_CHUNK_SELECTION
for t in e["triples"]
if t.p.iri == RDF_TYPE
)
@pytest.mark.asyncio
async def test_synthesis_derives_from_exploration_when_no_rerank(self):
"""
No-op lineage is unchanged: synthesis derives from exploration
(there is no focus stage). Guards the conditional synthesis parent.
"""
clients = build_mock_clients()
rag = DocumentRag(*clients)
events = []
async def explain_callback(triples, explain_id):
events.append({"triples": triples, "explain_id": explain_id})
await rag.query(
query="What is the return policy?",
explain_callback=explain_callback,
)
# events: question, grounding, exploration, synthesis
exp_uri = events[2]["explain_id"]
syn_event = events[3]
assert derived_from(syn_event["triples"], syn_event["explain_id"]) == exp_uri
# ---------------------------------------------------------------------------
# 2. Rerank: reorder + truncate + provenance
# ---------------------------------------------------------------------------
class TestRerankActive:
def _reranker_keeping_C_then_A(self):
# Reranker says chunk index 2 (C) is best, then index 0 (A); B dropped.
# Pre-sorted desc by score and truncated to limit, per the contract.
return StubReranker([
RerankerResult(document_id="2", query_id="0", score=0.95),
RerankerResult(document_id="0", query_id="0", score=0.42),
])
@pytest.mark.asyncio
async def test_documents_reordered_and_truncated(self):
clients = build_mock_clients()
reranker = self._reranker_keeping_C_then_A()
rag = DocumentRag(*clients, reranker_client=reranker)
await rag.query(query="What is the return policy?")
call = rag.prompt_client.document_prompt.call_args
passed_docs = call.kwargs["documents"]
assert passed_docs == [CHUNK_C_CONTENT, CHUNK_A_CONTENT]
@pytest.mark.asyncio
async def test_reranker_called_with_single_query_and_all_docs(self):
clients = build_mock_clients()
reranker = self._reranker_keeping_C_then_A()
rag = DocumentRag(*clients, reranker_client=reranker)
await rag.query(query="What is the return policy?", doc_limit=2)
assert len(reranker.calls) == 1
c = reranker.calls[0]
assert c["queries"] == [{"id": "0", "text": "What is the return policy?"}]
assert c["documents"] == [
{"id": "0", "text": CHUNK_A_CONTENT},
{"id": "1", "text": CHUNK_B_CONTENT},
{"id": "2", "text": CHUNK_C_CONTENT},
]
# The rerank narrows down to the final doc_limit, NOT fetch_limit
# (fetch_limit is the over-fetched candidate pool size).
assert c["limit"] == 2
@pytest.mark.asyncio
async def test_explicit_fetch_limit_over_fetches_then_narrows(self):
"""
Semantic guard for the value of reranking AND the maintainer's two-limit
contract: an explicit fetch_limit makes retrieval OVER-FETCH a wider
candidate pool so the cross-encoder can surface chunks the bi-encoder
ranked outside the final doc_limit, then the rerank narrows the pool back
down to doc_limit. The fetch_limit is honoured directly (caller controls
how hard the reranker works), not overridden by any heuristic.
"""
clients = build_mock_clients()
prompt_client, embeddings_client, doc_embeddings_client, fetch_chunk = clients
reranker = self._reranker_keeping_C_then_A()
# Candidate pool (fetch_limit=60) >> final doc_limit (6).
rag = DocumentRag(*clients, reranker_client=reranker)
await rag.query(
query="What is the return policy?", doc_limit=6, fetch_limit=60,
)
# Over-fetch: the embeddings store is queried with the fetch_limit
# budget (60 // 2 concept-vectors = 30 per concept), NOT the doc_limit
# budget (6 // 2 = 3). This is the bug guard.
q_limit = doc_embeddings_client.query.call_args.kwargs["limit"]
assert q_limit == 30
# Narrow: the rerank keeps the final doc_limit (6), not fetch_limit.
assert reranker.calls[0]["limit"] == 6
@pytest.mark.asyncio
async def test_default_fetch_limit_derives_overfetch_from_doc_limit(self):
"""
With no fetch_limit passed to query(), the candidate pool falls back to
the OVERFETCH_FACTOR x doc_limit heuristic, so over-fetch scales with
doc_limit and reranking keeps its recall benefit out of the box.
"""
clients = build_mock_clients()
prompt_client, embeddings_client, doc_embeddings_client, fetch_chunk = clients
reranker = self._reranker_keeping_C_then_A()
# No fetch_limit -> heuristic default.
rag = DocumentRag(*clients, reranker_client=reranker)
await rag.query(query="What is the return policy?", doc_limit=20)
# fetch = 3 x 20 = 60 -> 60 // 2 concept-vectors = 30 per concept.
q_limit = doc_embeddings_client.query.call_args.kwargs["limit"]
assert q_limit == 30
# Rerank narrows to the final doc_limit (20).
assert reranker.calls[0]["limit"] == 20
@pytest.mark.asyncio
async def test_fetch_limit_floored_at_doc_limit(self):
"""
A fetch_limit below doc_limit is floored up to doc_limit: retrieval must
never fetch fewer candidates than the rerank is asked to keep, else the
prompt could not be filled.
"""
clients = build_mock_clients()
prompt_client, embeddings_client, doc_embeddings_client, fetch_chunk = clients
reranker = self._reranker_keeping_C_then_A()
rag = DocumentRag(*clients, reranker_client=reranker)
await rag.query(
query="What is the return policy?", doc_limit=10, fetch_limit=4,
)
# fetch = max(4, 10) = 10 -> 10 // 2 concept-vectors = 5 per concept.
q_limit = doc_embeddings_client.query.call_args.kwargs["limit"]
assert q_limit == 5
assert reranker.calls[0]["limit"] == 10
@pytest.mark.asyncio
async def test_chunk_selection_event_emitted(self):
clients = build_mock_clients()
reranker = self._reranker_keeping_C_then_A()
rag = DocumentRag(*clients, reranker_client=reranker)
events = []
async def explain_callback(triples, explain_id):
events.append({"triples": triples, "explain_id": explain_id})
await rag.query(
query="What is the return policy?",
explain_callback=explain_callback,
)
# Now 5 stages: question, grounding, exploration, focus, synthesis.
assert len(events) == 5
ordered_types = [
TG_DOC_RAG_QUESTION, TG_GROUNDING, TG_EXPLORATION,
TG_FOCUS, TG_SYNTHESIS,
]
for i, expected in enumerate(ordered_types):
assert has_type(events[i]["triples"], events[i]["explain_id"], expected)
@pytest.mark.asyncio
async def test_chunk_selection_carries_scores_and_chunk_refs(self):
clients = build_mock_clients()
reranker = self._reranker_keeping_C_then_A()
rag = DocumentRag(*clients, reranker_client=reranker)
events = []
async def explain_callback(triples, explain_id):
events.append({"triples": triples, "explain_id": explain_id})
await rag.query(
query="What is the return policy?",
explain_callback=explain_callback,
)
focus_event = events[3]
foc_uri = focus_event["explain_id"]
triples = focus_event["triples"]
# focus is derived from exploration
exp_uri = events[2]["explain_id"]
assert derived_from(triples, foc_uri) == exp_uri
# Two ChunkSelection sub-entities, linked from focus.
sel_links = find_triples(triples, TG_SELECTED_CHUNK, foc_uri)
assert len(sel_links) == 2
# Each selection has a ChunkSelection type, a chunk document ref and a score.
chunk_refs = set()
scores = set()
for link in sel_links:
sel_uri = link.o.iri
assert has_type(triples, sel_uri, TG_CHUNK_SELECTION)
doc_ref = find_triple(triples, TG_DOCUMENT, sel_uri)
assert doc_ref is not None
chunk_refs.add(doc_ref.o.iri)
score_t = find_triple(triples, TG_SCORE, sel_uri)
assert score_t is not None
scores.add(score_t.o.value)
# Surviving chunks are C and A (B dropped), with the reranker scores.
assert chunk_refs == {CHUNK_C, CHUNK_A}
assert scores == {"0.95", "0.42"}
@pytest.mark.asyncio
async def test_all_focus_triples_in_retrieval_graph(self):
clients = build_mock_clients()
reranker = self._reranker_keeping_C_then_A()
rag = DocumentRag(*clients, reranker_client=reranker)
events = []
async def explain_callback(triples, explain_id):
events.append({"triples": triples, "explain_id": explain_id})
await rag.query(
query="What is the return policy?",
explain_callback=explain_callback,
)
for t in events[3]["triples"]:
assert t.g == "urn:graph:retrieval"
@pytest.mark.asyncio
async def test_synthesis_derives_from_focus_when_reranking(self):
"""
When reranking runs, synthesis must derive from the focus node (the
reranked chunks actually fed to the LLM), mirroring GraphRAG - not from
exploration, which would leave focus as a dangling branch and
misrepresent what fed the answer.
"""
clients = build_mock_clients()
reranker = self._reranker_keeping_C_then_A()
rag = DocumentRag(*clients, reranker_client=reranker)
events = []
async def explain_callback(triples, explain_id):
events.append({"triples": triples, "explain_id": explain_id})
await rag.query(
query="What is the return policy?",
doc_limit=2,
explain_callback=explain_callback,
)
# events: question, grounding, exploration, focus, synthesis
foc_uri = events[3]["explain_id"]
syn_event = events[4]
assert derived_from(syn_event["triples"], syn_event["explain_id"]) == foc_uri
@pytest.mark.asyncio
async def test_empty_docs_skips_reranker(self):
"""If retrieval returns no chunks, the reranker is never called."""
clients = build_mock_clients()
prompt_client, embeddings_client, doc_embeddings_client, fetch_chunk = clients
doc_embeddings_client.query.return_value = [] # no matches
reranker = self._reranker_keeping_C_then_A()
rag = DocumentRag(*clients, reranker_client=reranker)
await rag.query(query="What is the return policy?")
assert reranker.calls == []
# ---------------------------------------------------------------------------
# 3. Diversity selection: optional MMR after cross-encoder scoring
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_diversity_mode_scores_full_candidate_pool_before_selecting(self):
"""
With diversity selection enabled, the cross-encoder should score the full
fetched candidate pool before MMR narrows it down to doc_limit.
"""
clients = build_mock_clients()
reranker = StubReranker([
RerankerResult(document_id="0", query_id="0", score=1.00),
RerankerResult(document_id="1", query_id="0", score=0.95),
RerankerResult(document_id="2", query_id="0", score=0.90),
])
rag = DocumentRag(
*clients,
reranker_client=reranker,
rerank_diversity_mode="mmr",
)
await rag.query(query="What is the return policy?", doc_limit=2)
assert reranker.calls[0]["limit"] == len(ORDERED_CONTENT)
call = rag.prompt_client.document_prompt.call_args
passed_docs = call.kwargs["documents"]
assert len(passed_docs) == 2
@pytest.mark.asyncio
async def test_diversity_mode_selects_less_redundant_context_set(self):
"""
MMR should use cross-encoder scores as relevance while penalizing redundant
chunks, so a slightly lower-scored but less redundant chunk can be selected.
"""
clients = build_mock_clients()
prompt_client, embeddings_client, doc_embeddings_client, fetch_chunk = clients
duplicate_a = "apple banana fruit return policy"
duplicate_b = "apple banana fruit return policy duplicate"
diverse_c = "engine motor vehicle warranty"
async def mock_fetch(chunk_id):
return {
CHUNK_A: duplicate_a,
CHUNK_B: duplicate_b,
CHUNK_C: diverse_c,
}[chunk_id]
fetch_chunk.side_effect = mock_fetch
reranker = StubReranker([
RerankerResult(document_id="0", query_id="0", score=1.00),
RerankerResult(document_id="1", query_id="0", score=0.95),
RerankerResult(document_id="2", query_id="0", score=0.90),
])
rag = DocumentRag(
*clients,
reranker_client=reranker,
rerank_diversity_mode="mmr",
rerank_diversity_lambda=0.2,
)
await rag.query(query="What is the return policy?", doc_limit=2)
call = rag.prompt_client.document_prompt.call_args
passed_docs = call.kwargs["documents"]
assert passed_docs == [duplicate_a, diverse_c]

View file

@ -0,0 +1,89 @@
"""
Cross-layer wiring contract for the Document-RAG reranker (issue #878).
The Document-RAG processor registers a ``RerankerClientSpec`` for the
``reranker-request`` / ``reranker-response`` roles (see
``retrieval/document_rag/rag.py``). At flow construction every spec runs
``spec.add(flow, processor, definition)``, and ``RequestResponseSpec.add``
resolves its topics via ``definition["topics"][name]`` - which raises
``KeyError`` if the flow blueprint does not provide those topics.
This means the monorepo code change is only safe to deploy together with the
companion ``trustgraph-templates`` change that wires ``reranker-request`` /
``reranker-response`` into the Document-RAG flow (mirroring what templates
PR #279 did for GraphRAG via ``graph-store.jsonnet``). These tests pin that
contract from the monorepo side:
* with the reranker topics present (as the updated templates compile them),
the spec binds cleanly and registers the client;
* without them (the pre-companion blueprint), construction fails fast with a
KeyError naming the missing role - documenting exactly why the templates
change is required.
No broker/network: the pub/sub backend is mocked (topics are bound at add()
time, connections happen later at start()).
"""
import pytest
from unittest.mock import MagicMock
from trustgraph.base import RerankerClientSpec
def _flow():
f = MagicMock()
f.workspace = "ws"
f.name = "document-rag"
f.id = "proc1"
f.consumer = {}
return f
def _processor():
p = MagicMock()
p.pubsub = MagicMock()
p.id = "proc1"
p.taskgroup = MagicMock()
return p
def _spec():
return RerankerClientSpec(
request_name="reranker-request",
response_name="reranker-response",
)
# Topics dict as the UPDATED document-store.jsonnet compiles them
# (verified by compiling the template: reranker-request -> request:tg:reranker:{workspace}:{id}).
DEFINITION_WITH_RERANKER = {
"topics": {
"request": "request:tg:document-rag:ws:id",
"response": "response:tg:document-rag:ws:id",
"reranker-request": "request:tg:reranker:ws:id",
"reranker-response": "response:tg:reranker:ws:id",
}
}
# Pre-companion blueprint: no reranker topics (document-rag before the templates change).
DEFINITION_WITHOUT_RERANKER = {
"topics": {
"request": "request:tg:document-rag:ws:id",
"response": "response:tg:document-rag:ws:id",
}
}
def test_reranker_client_binds_when_flow_provides_topics():
flow = _flow()
_spec().add(flow, _processor(), DEFINITION_WITH_RERANKER)
# The client consumer is registered against the reranker role.
assert "reranker-request" in flow.consumer
def test_reranker_client_keyerrors_without_companion_template_topics():
with pytest.raises(KeyError) as exc:
_spec().add(_flow(), _processor(), DEFINITION_WITHOUT_RERANKER)
# Fails fast naming the missing role -> the trustgraph-templates companion
# change (wire reranker-request/response into the document-rag flow) is required.
assert "reranker-request" in str(exc.value)

View file

@ -66,6 +66,7 @@ class TestDocumentRagService:
workspace=ANY, # Workspace comes from flow.workspace (mock)
collection="test_coll_1", # Must be from message, not hardcoded default
doc_limit=5,
fetch_limit=0, # Unset -> core derives the candidate pool
explain_callback=ANY, # Explainability callback is always passed
save_answer_callback=ANY, # Librarian save callback is always passed
)
@ -130,3 +131,72 @@ class TestDocumentRagService:
assert sent_response.end_of_stream is True, "Non-streaming response must have end_of_stream=True"
assert sent_response.end_of_session is True
assert sent_response.error is None
def test_fetch_chunk_timeout_defaults_to_120(self):
"""Processor without fetch_chunk_timeout override uses default of 120."""
processor = Processor(
taskgroup=MagicMock(),
id="test-processor",
doc_limit=10
)
assert processor.fetch_chunk_timeout == 120
def test_fetch_chunk_timeout_uses_overridden_value(self):
"""Processor with fetch_chunk_timeout override stores that value."""
processor = Processor(
taskgroup=MagicMock(),
id="test-processor",
doc_limit=10,
fetch_chunk_timeout=45
)
assert processor.fetch_chunk_timeout == 45
@patch('trustgraph.retrieval.document_rag.rag.DocumentRag')
@pytest.mark.asyncio
async def test_fetch_chunk_uses_configured_timeout(self, mock_document_rag_class):
"""
Test that the fetch_chunk closure built in on_request calls
flow.librarian.fetch_document_text with the configured
fetch_chunk_timeout, not the old hardcoded 120.
"""
processor = Processor(
taskgroup=MagicMock(),
id="test-processor",
doc_limit=10,
fetch_chunk_timeout=45
)
mock_rag_instance = AsyncMock()
mock_document_rag_class.return_value = mock_rag_instance
mock_rag_instance.query.return_value = (
"test response", {"in_token": None, "out_token": None, "model": None})
msg = MagicMock()
msg.value.return_value = DocumentRagQuery(
query="test query",
collection="default",
doc_limit=5
)
msg.properties.return_value = {"id": "test-id"}
consumer = MagicMock()
flow = MagicMock()
flow.librarian.fetch_document_text = AsyncMock(return_value="chunk text")
mock_producer = AsyncMock()
def flow_router(service_name):
if service_name == "response":
return mock_producer
return AsyncMock()
flow.side_effect = flow_router
await processor.on_request(msg, consumer, flow)
# Retrieve the fetch_chunk callable that on_request built and passed into DocumentRag(...)
fetch_chunk = mock_document_rag_class.call_args.kwargs["fetch_chunk"]
await fetch_chunk("some-chunk-id")
flow.librarian.fetch_document_text.assert_called_once_with(
document_id="some-chunk-id", timeout=45,
)

View file

@ -15,54 +15,52 @@ class TestGraphRag:
def test_graph_rag_initialization_with_defaults(self):
"""Test GraphRag initialization with default verbose setting"""
# Create mock clients
mock_prompt_client = MagicMock()
mock_embeddings_client = MagicMock()
mock_graph_embeddings_client = MagicMock()
mock_triples_client = MagicMock()
mock_reranker_client = MagicMock()
# Initialize GraphRag
graph_rag = GraphRag(
prompt_client=mock_prompt_client,
embeddings_client=mock_embeddings_client,
graph_embeddings_client=mock_graph_embeddings_client,
triples_client=mock_triples_client
)
# Verify initialization
assert graph_rag.prompt_client == mock_prompt_client
assert graph_rag.embeddings_client == mock_embeddings_client
assert graph_rag.graph_embeddings_client == mock_graph_embeddings_client
assert graph_rag.triples_client == mock_triples_client
assert graph_rag.verbose is False # Default value
# Verify label_cache is an LRUCacheWithTTL instance
from trustgraph.retrieval.graph_rag.graph_rag import LRUCacheWithTTL
assert isinstance(graph_rag.label_cache, LRUCacheWithTTL)
def test_graph_rag_initialization_with_verbose(self):
"""Test GraphRag initialization with verbose enabled"""
# Create mock clients
mock_prompt_client = MagicMock()
mock_embeddings_client = MagicMock()
mock_graph_embeddings_client = MagicMock()
mock_triples_client = MagicMock()
# Initialize GraphRag with verbose=True
graph_rag = GraphRag(
prompt_client=mock_prompt_client,
embeddings_client=mock_embeddings_client,
graph_embeddings_client=mock_graph_embeddings_client,
triples_client=mock_triples_client,
verbose=True
reranker_client=mock_reranker_client,
)
# Verify initialization
assert graph_rag.prompt_client == mock_prompt_client
assert graph_rag.embeddings_client == mock_embeddings_client
assert graph_rag.graph_embeddings_client == mock_graph_embeddings_client
assert graph_rag.triples_client == mock_triples_client
assert graph_rag.reranker_client == mock_reranker_client
assert graph_rag.verbose is False
from trustgraph.retrieval.graph_rag.graph_rag import LRUCacheWithTTL
assert isinstance(graph_rag.label_cache, LRUCacheWithTTL)
def test_graph_rag_initialization_with_verbose(self):
"""Test GraphRag initialization with verbose enabled"""
mock_prompt_client = MagicMock()
mock_embeddings_client = MagicMock()
mock_graph_embeddings_client = MagicMock()
mock_triples_client = MagicMock()
mock_reranker_client = MagicMock()
graph_rag = GraphRag(
prompt_client=mock_prompt_client,
embeddings_client=mock_embeddings_client,
graph_embeddings_client=mock_graph_embeddings_client,
triples_client=mock_triples_client,
reranker_client=mock_reranker_client,
verbose=True,
)
assert graph_rag.prompt_client == mock_prompt_client
assert graph_rag.embeddings_client == mock_embeddings_client
assert graph_rag.graph_embeddings_client == mock_graph_embeddings_client
assert graph_rag.triples_client == mock_triples_client
assert graph_rag.reranker_client == mock_reranker_client
assert graph_rag.verbose is True
# Verify label_cache is an LRUCacheWithTTL instance
from trustgraph.retrieval.graph_rag.graph_rag import LRUCacheWithTTL
assert isinstance(graph_rag.label_cache, LRUCacheWithTTL)
@ -365,244 +363,162 @@ class TestQuery:
assert "workspace" not in c.kwargs
@pytest.mark.asyncio
async def test_follow_edges_never_passes_workspace(self):
"""Verify follow_edges never passes workspace to query_stream."""
async def test_hop_and_filter_never_passes_workspace(self):
"""Verify hop_and_filter never passes workspace to query_stream."""
mock_rag = MagicMock()
mock_triples_client = AsyncMock()
mock_reranker_client = AsyncMock()
mock_rag.triples_client = mock_triples_client
mock_rag.reranker_client = mock_reranker_client
mock_rag.label_cache = MagicMock()
mock_rag.label_cache.get.return_value = None
mock_triple = MagicMock()
mock_triple.s, mock_triple.p, mock_triple.o = "e1", "p1", "o1"
mock_triple.s = "e1"
mock_triple.p = "p1"
mock_triple.o = "o1"
mock_triples_client.query_stream.return_value = [mock_triple]
mock_triples_client.query.return_value = []
result = MagicMock()
result.document_id = "0"
result.query_id = "0"
result.score = 0.9
mock_reranker_client.rerank.return_value = [result]
query = Query(
rag=mock_rag,
collection="test_collection",
verbose=False,
triple_limit=10
triple_limit=10,
)
subgraph = set()
await query.follow_edges("e1", subgraph, path_length=1)
await query.hop_and_filter(["e1"], ["concept"])
for c in mock_triples_client.query_stream.call_args_list:
assert "workspace" not in c.kwargs
@pytest.mark.asyncio
async def test_follow_edges_basic_functionality(self):
"""Test Query.follow_edges method basic triple discovery"""
async def test_hop_and_filter_basic_functionality(self):
"""Test hop_and_filter retrieves edges and scores them with reranker."""
mock_rag = MagicMock()
mock_triples_client = AsyncMock()
mock_reranker_client = AsyncMock()
mock_rag.triples_client = mock_triples_client
mock_rag.reranker_client = mock_reranker_client
mock_rag.label_cache = MagicMock()
mock_rag.label_cache.get.return_value = None
mock_triple1 = MagicMock()
mock_triple1.s, mock_triple1.p, mock_triple1.o = "entity1", "predicate1", "object1"
mock_triple = MagicMock()
mock_triple.s = "entity1"
mock_triple.p = "predicate1"
mock_triple.o = "object1"
mock_triples_client.query_stream.return_value = [mock_triple]
mock_triples_client.query.return_value = []
mock_triple2 = MagicMock()
mock_triple2.s, mock_triple2.p, mock_triple2.o = "subject2", "entity1", "object2"
mock_triple3 = MagicMock()
mock_triple3.s, mock_triple3.p, mock_triple3.o = "subject3", "predicate3", "entity1"
mock_triples_client.query_stream.side_effect = [
[mock_triple1], # s=ent
[mock_triple2], # p=ent
[mock_triple3], # o=ent
]
result = MagicMock()
result.document_id = "0"
result.query_id = "0"
result.score = 0.95
mock_reranker_client.rerank.return_value = [result]
query = Query(
rag=mock_rag,
collection="test_collection",
verbose=False,
triple_limit=10
triple_limit=10,
edge_limit=25,
)
subgraph = set()
await query.follow_edges("entity1", subgraph, path_length=1)
assert mock_triples_client.query_stream.call_count == 3
mock_triples_client.query_stream.assert_any_call(
s="entity1", p=None, o=None, limit=10,
collection="test_collection", batch_size=20, g=""
)
mock_triples_client.query_stream.assert_any_call(
s=None, p="entity1", o=None, limit=10,
collection="test_collection", batch_size=20, g=""
)
mock_triples_client.query_stream.assert_any_call(
s=None, p=None, o="entity1", limit=10,
collection="test_collection", batch_size=20, g=""
selected, uri_map, edge_meta = await query.hop_and_filter(
["entity1"], ["test concept"],
)
expected_subgraph = {
("entity1", "predicate1", "object1"),
("subject2", "entity1", "object2"),
("subject3", "predicate3", "entity1")
}
assert subgraph == expected_subgraph
assert len(selected) == 1
assert len(uri_map) == 1
assert len(edge_meta) == 1
mock_reranker_client.rerank.assert_called_once()
call_kwargs = mock_reranker_client.rerank.call_args
assert call_kwargs.kwargs["limit"] == 25
@pytest.mark.asyncio
async def test_follow_edges_with_path_length_zero(self):
"""Test Query.follow_edges method with path_length=0"""
async def test_hop_and_filter_with_empty_frontier(self):
"""Test hop_and_filter with no seed entities returns empty."""
mock_rag = MagicMock()
query = Query(
rag=mock_rag,
collection="test_collection",
verbose=False,
)
selected, uri_map, edge_meta = await query.hop_and_filter([], ["concept"])
assert selected == []
assert uri_map == {}
assert edge_meta == {}
@pytest.mark.asyncio
async def test_hop_and_filter_filters_label_triples(self):
"""Test hop_and_filter skips rdfs:label edges."""
mock_rag = MagicMock()
mock_triples_client = AsyncMock()
mock_reranker_client = AsyncMock()
mock_rag.triples_client = mock_triples_client
mock_rag.reranker_client = mock_reranker_client
mock_rag.label_cache = MagicMock()
mock_rag.label_cache.get.return_value = None
query = Query(
rag=mock_rag,
collection="test_collection",
verbose=False
)
label_triple = MagicMock()
label_triple.s = "entity1"
label_triple.p = "http://www.w3.org/2000/01/rdf-schema#label"
label_triple.o = "Entity One"
subgraph = set()
await query.follow_edges("entity1", subgraph, path_length=0)
mock_triples_client.query_stream.assert_not_called()
assert subgraph == set()
@pytest.mark.asyncio
async def test_follow_edges_with_max_subgraph_size_limit(self):
"""Test Query.follow_edges method respects max_subgraph_size"""
mock_rag = MagicMock()
mock_triples_client = AsyncMock()
mock_rag.triples_client = mock_triples_client
mock_triples_client.query_stream.return_value = [label_triple]
mock_triples_client.query.return_value = []
query = Query(
rag=mock_rag,
collection="test_collection",
verbose=False,
max_subgraph_size=2
triple_limit=10,
)
subgraph = {("s1", "p1", "o1"), ("s2", "p2", "o2"), ("s3", "p3", "o3")}
await query.follow_edges("entity1", subgraph, path_length=1)
mock_triples_client.query_stream.assert_not_called()
assert len(subgraph) == 3
@pytest.mark.asyncio
async def test_get_subgraph_method(self):
"""Test Query.get_subgraph returns (subgraph, entities, concepts) tuple"""
mock_rag = MagicMock()
query = Query(
rag=mock_rag,
collection="test_collection",
verbose=False,
max_path_length=1
selected, uri_map, edge_meta = await query.hop_and_filter(
["entity1"], ["concept"],
)
# Mock get_entities to return (entities, concepts) tuple
query.get_entities = AsyncMock(
return_value=(["entity1", "entity2"], ["concept1"])
)
query.follow_edges_batch = AsyncMock(return_value=(
{
("entity1", "predicate1", "object1"),
("entity2", "predicate2", "object2")
},
{}
))
subgraph, term_map, entities, concepts = await query.get_subgraph("test query")
query.get_entities.assert_called_once_with("test query")
query.follow_edges_batch.assert_called_once_with(["entity1", "entity2"], 1)
assert isinstance(subgraph, list)
assert len(subgraph) == 2
assert ("entity1", "predicate1", "object1") in subgraph
assert ("entity2", "predicate2", "object2") in subgraph
assert entities == ["entity1", "entity2"]
assert concepts == ["concept1"]
@pytest.mark.asyncio
async def test_get_labelgraph_method(self):
"""Test Query.get_labelgraph returns (labeled_edges, uri_map, entities, concepts)"""
mock_rag = MagicMock()
query = Query(
rag=mock_rag,
collection="test_collection",
verbose=False,
max_subgraph_size=100
)
test_subgraph = [
("entity1", "predicate1", "object1"),
("subject2", "http://www.w3.org/2000/01/rdf-schema#label", "Label Value"),
("entity3", "predicate3", "object3")
]
test_entities = ["entity1", "entity3"]
test_concepts = ["concept1"]
query.get_subgraph = AsyncMock(
return_value=(test_subgraph, {}, test_entities, test_concepts)
)
async def mock_maybe_label(entity):
label_map = {
"entity1": "Human Entity One",
"predicate1": "Human Predicate One",
"object1": "Human Object One",
"entity3": "Human Entity Three",
"predicate3": "Human Predicate Three",
"object3": "Human Object Three"
}
return label_map.get(entity, entity)
query.maybe_label = AsyncMock(side_effect=mock_maybe_label)
labeled_edges, uri_map, entities, concepts = await query.get_labelgraph("test query")
query.get_subgraph.assert_called_once_with("test query")
# Label triples filtered out
assert len(labeled_edges) == 2
# maybe_label called for non-label triples
assert query.maybe_label.call_count == 6
expected_edges = [
("Human Entity One", "Human Predicate One", "Human Object One"),
("Human Entity Three", "Human Predicate Three", "Human Object Three")
]
assert labeled_edges == expected_edges
assert len(uri_map) == 2
assert entities == test_entities
assert concepts == test_concepts
assert selected == []
mock_reranker_client.rerank.assert_not_called()
@pytest.mark.asyncio
async def test_graph_rag_query_method(self):
"""Test GraphRag.query method orchestrates full RAG pipeline with provenance"""
import json
from trustgraph.retrieval.graph_rag.graph_rag import edge_id
mock_prompt_client = AsyncMock()
mock_embeddings_client = AsyncMock()
mock_graph_embeddings_client = AsyncMock()
mock_triples_client = AsyncMock()
mock_reranker_client = AsyncMock()
expected_response = "This is the RAG response"
test_labelgraph = [("Subject", "Predicate", "Object")]
test_edge_id = edge_id("Subject", "Predicate", "Object")
test_selected_edges = [("Subject", "Predicate", "Object")]
test_eid = edge_id("Subject", "Predicate", "Object")
test_uri_map = {
test_edge_id: ("http://example.org/subject", "http://example.org/predicate", "http://example.org/object")
test_eid: ("http://example.org/subject", "http://example.org/predicate", "http://example.org/object")
}
test_edge_metadata = {
test_eid: {"concept": "test concept", "score": 0.95}
}
test_entities = ["http://example.org/subject"]
test_concepts = ["test concept"]
# Mock prompt responses for the multi-step process
mock_embeddings_client.embed.return_value = [[0.1, 0.2]]
mock_graph_embeddings_client.query.return_value = []
async def mock_prompt(prompt_name, variables=None, streaming=False, chunk_callback=None):
if prompt_name == "extract-concepts":
return PromptResult(response_type="text", text="")
elif prompt_name == "kg-edge-scoring":
return PromptResult(response_type="jsonl", objects=[{"id": test_edge_id, "score": 0.9}])
elif prompt_name == "kg-edge-reasoning":
return PromptResult(response_type="jsonl", objects=[{"id": test_edge_id, "reasoning": "relevant"}])
return PromptResult(response_type="text", text="test concept")
elif prompt_name == "kg-synthesis":
return PromptResult(response_type="text", text=expected_response)
return PromptResult(response_type="text", text="")
@ -614,16 +530,16 @@ class TestQuery:
embeddings_client=mock_embeddings_client,
graph_embeddings_client=mock_graph_embeddings_client,
triples_client=mock_triples_client,
verbose=False
reranker_client=mock_reranker_client,
verbose=False,
)
# Patch Query.get_labelgraph to return test data
original_get_labelgraph = Query.get_labelgraph
original_hop_and_filter = Query.hop_and_filter
async def mock_get_labelgraph(self, query_text):
return test_labelgraph, test_uri_map, test_entities, test_concepts
async def mock_hop_and_filter(self, seed_entities, concepts):
return test_selected_edges, test_uri_map, test_edge_metadata
Query.get_labelgraph = mock_get_labelgraph
Query.hop_and_filter = mock_hop_and_filter
provenance_events = []
@ -636,7 +552,7 @@ class TestQuery:
collection="test_collection",
entity_limit=25,
triple_limit=15,
explain_callback=collect_provenance
explain_callback=collect_provenance,
)
response_text, usage = response
@ -650,7 +566,6 @@ class TestQuery:
assert len(triples) > 0
assert prov_id.startswith("urn:trustgraph:")
# Verify order
assert "question" in provenance_events[0][1]
assert "grounding" in provenance_events[1][1]
assert "exploration" in provenance_events[2][1]
@ -658,4 +573,4 @@ class TestQuery:
assert "synthesis" in provenance_events[4][1]
finally:
Query.get_labelgraph = original_get_labelgraph
Query.hop_and_filter = original_hop_and_filter

View file

@ -0,0 +1,395 @@
"""
Tests for direction-aware reranker text in GraphRAG hop-and-filter.
The reranker document text varies by traversal direction:
- From S (subject is the frontier entity): text = "{p} {o}"
- From O (object is the frontier entity): text = "{s} {p}"
- From P (predicate is the frontier entity): text = "{s} {o}"
"""
import pytest
from unittest.mock import MagicMock, AsyncMock
from trustgraph.retrieval.graph_rag.graph_rag import Query, LRUCacheWithTTL
from trustgraph.schema import Term, IRI, LITERAL
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
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()
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
else:
rag.reranker_client.rerank.return_value = []
return rag
def _make_query(rag, max_path_length=1, edge_limit=25):
return Query(
rag=rag,
collection="test",
verbose=False,
entity_limit=50,
triple_limit=30,
max_subgraph_size=1000,
max_path_length=max_path_length,
edge_limit=edge_limit,
)
def _make_schema_triple(s, p, o):
"""Create a mock triple matching the schema interface."""
t = MagicMock()
t.s = s
t.p = p
t.o = o
return t
def _reranker_result(document_id, query_id="0", score=0.9):
r = MagicMock()
r.document_id = str(document_id)
r.query_id = str(query_id)
r.score = score
return r
# ---------------------------------------------------------------------------
# Tests: execute_batch_triple_queries direction tracking
# ---------------------------------------------------------------------------
class TestDirectionTracking:
@pytest.mark.asyncio
async def test_from_s_direction(self):
"""Triples from s=entity queries are tagged FROM_S."""
triple = _make_schema_triple("ent1", "pred", "obj")
rag = _make_rag()
async def query_stream(s=None, p=None, o=None, **kwargs):
if s is not None:
return [triple]
return []
rag.triples_client.query_stream.side_effect = query_stream
q = _make_query(rag)
result = await q.execute_batch_triple_queries(["ent1"], 10)
from_s = [(t, d) for t, d in result if d == Query.FROM_S]
assert len(from_s) == 1
assert from_s[0][0] is triple
@pytest.mark.asyncio
async def test_from_o_direction(self):
"""Triples from o=entity queries are tagged FROM_O."""
triple = _make_schema_triple("subj", "pred", "ent1")
rag = _make_rag()
async def query_stream(s=None, p=None, o=None, **kwargs):
if o is not None:
return [triple]
return []
rag.triples_client.query_stream.side_effect = query_stream
q = _make_query(rag)
result = await q.execute_batch_triple_queries(["ent1"], 10)
from_o = [(t, d) for t, d in result if d == Query.FROM_O]
assert len(from_o) == 1
assert from_o[0][0] is triple
@pytest.mark.asyncio
async def test_from_p_direction(self):
"""Triples from p=entity queries are tagged FROM_P."""
triple = _make_schema_triple("subj", "ent1", "obj")
rag = _make_rag()
async def query_stream(s=None, p=None, o=None, **kwargs):
if p is not None:
return [triple]
return []
rag.triples_client.query_stream.side_effect = query_stream
q = _make_query(rag)
result = await q.execute_batch_triple_queries(["ent1"], 10)
from_p = [(t, d) for t, d in result if d == Query.FROM_P]
assert len(from_p) == 1
assert from_p[0][0] is triple
# ---------------------------------------------------------------------------
# Tests: hop_and_filter reranker document text
# ---------------------------------------------------------------------------
class TestDirectionAwareRerankerText:
@pytest.mark.asyncio
async def test_from_s_uses_predicate_object(self):
"""From-S traversal: reranker text should be '{p} {o}'."""
triple = _make_schema_triple(
"http://ex/entity-A",
"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], labels=labels)
async def query_stream(s=None, p=None, o=None, **kwargs):
if s is not None:
return [triple]
return []
rag.triples_client.query_stream.side_effect = query_stream
q = _make_query(rag, max_path_length=1, edge_limit=10)
await q.hop_and_filter(
seed_entities=["http://ex/entity-A"],
concepts=["likes"],
)
call_args = rag.reranker_client.rerank.call_args
documents = call_args.kwargs["documents"]
assert len(documents) == 1
assert documents[0]["text"] == "likes Bob"
@pytest.mark.asyncio
async def test_from_o_uses_subject_predicate(self):
"""From-O traversal: reranker text should be '{s} {p}'."""
triple = _make_schema_triple(
"http://ex/entity-A",
"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], labels=labels)
async def query_stream(s=None, p=None, o=None, **kwargs):
if o is not None:
return [triple]
return []
rag.triples_client.query_stream.side_effect = query_stream
q = _make_query(rag, max_path_length=1, edge_limit=10)
await q.hop_and_filter(
seed_entities=["http://ex/entity-B"],
concepts=["likes"],
)
call_args = rag.reranker_client.rerank.call_args
documents = call_args.kwargs["documents"]
assert len(documents) == 1
assert documents[0]["text"] == "Alice likes"
@pytest.mark.asyncio
async def test_from_p_uses_subject_object(self):
"""From-P traversal: reranker text should be '{s} {o}'."""
triple = _make_schema_triple(
"http://ex/entity-A",
"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], labels=labels)
async def query_stream(s=None, p=None, o=None, **kwargs):
if p is not None:
return [triple]
return []
rag.triples_client.query_stream.side_effect = query_stream
q = _make_query(rag, max_path_length=1, edge_limit=10)
await q.hop_and_filter(
seed_entities=["http://ex/likes"],
concepts=["entity"],
)
call_args = rag.reranker_client.rerank.call_args
documents = call_args.kwargs["documents"]
assert len(documents) == 1
assert documents[0]["text"] == "Alice Bob"
@pytest.mark.asyncio
async def test_mixed_directions_produce_different_text(self):
"""Edges from different directions use different text formats."""
triple_from_s = _make_schema_triple(
"http://ex/seed", "http://ex/rel", "http://ex/target",
)
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)],
labels=labels,
)
async def query_stream(s=None, p=None, o=None, **kwargs):
if s == "http://ex/seed":
return [triple_from_s]
if o == "http://ex/seed":
return [triple_from_o]
return []
rag.triples_client.query_stream.side_effect = query_stream
q = _make_query(rag, max_path_length=1, edge_limit=10)
await q.hop_and_filter(
seed_entities=["http://ex/seed"],
concepts=["test"],
)
call_args = rag.reranker_client.rerank.call_args
documents = call_args.kwargs["documents"]
texts = {d["text"] for d in documents}
# 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):
"""Labels should be resolved and used in the direction-aware text."""
triple = _make_schema_triple(
"http://ex/entity-A",
"http://ex/likes",
"http://ex/entity-B",
)
reranker_result = _reranker_result(0)
rag = _make_rag(reranker_results=[reranker_result])
async def query_stream(s=None, p=None, o=None, **kwargs):
if s is not None and p is None:
return [triple]
return []
async def label_query(s=None, p=None, o=None, limit=1, **kwargs):
if p == LABEL:
labels = {
"http://ex/entity-A": "Alice",
"http://ex/likes": "likes",
"http://ex/entity-B": "Bob",
}
if s in labels:
return [MagicMock(o=labels[s])]
return []
rag.triples_client.query_stream.side_effect = query_stream
rag.triples_client.query.side_effect = label_query
q = _make_query(rag, max_path_length=1, edge_limit=10)
await q.hop_and_filter(
seed_entities=["http://ex/entity-A"],
concepts=["friendship"],
)
call_args = rag.reranker_client.rerank.call_args
documents = call_args.kwargs["documents"]
assert len(documents) == 1
# From S with labels: "{p_label} {o_label}"
assert documents[0]["text"] == "likes Bob"
@pytest.mark.asyncio
async def test_no_duplicate_text_from_shared_object(self):
"""Multiple edges sharing an object should produce distinct texts."""
triple_a = _make_schema_triple(
"http://ex/cpu-A", "http://ex/hasCategory", "http://ex/Processors",
)
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)],
labels=labels,
)
async def query_stream(s=None, p=None, o=None, **kwargs):
if o == "http://ex/Processors":
return [triple_a, triple_b]
return []
rag.triples_client.query_stream.side_effect = query_stream
q = _make_query(rag, max_path_length=1, edge_limit=10)
await q.hop_and_filter(
seed_entities=["http://ex/Processors"],
concepts=["CPUs"],
)
call_args = rag.reranker_client.rerank.call_args
documents = call_args.kwargs["documents"]
texts = [d["text"] for d in documents]
assert len(texts) == 2
# From O: "{s} {p}" — subjects differ, so texts differ
assert texts[0] != texts[1]
assert "CPU Alpha" in texts[0]
assert "CPU Beta" in texts[1]

View file

@ -20,7 +20,7 @@ from trustgraph.provenance.namespaces import (
TG_GRAPH_RAG_QUESTION, TG_GROUNDING, TG_EXPLORATION,
TG_FOCUS, TG_SYNTHESIS, TG_ANSWER_TYPE,
TG_QUERY, TG_CONCEPT, TG_ENTITY, TG_EDGE_COUNT,
TG_SELECTED_EDGE, TG_EDGE, TG_REASONING,
TG_SELECTED_EDGE, TG_EDGE, TG_SCORE, TG_EDGE_SELECTION,
)
@ -91,17 +91,17 @@ def build_mock_clients():
1. prompt_client.prompt("extract-concepts", ...) -> concepts
2. embeddings_client.embed(concepts) -> vectors
3. graph_embeddings_client.query(vector, ...) -> entity matches
4. triples_client.query_stream(s/p/o, ...) -> edges (follow_edges_batch)
4. triples_client.query_stream(s/p/o, ...) -> edges (hop_and_filter)
5. triples_client.query(s, LABEL, ...) -> labels (maybe_label)
6. prompt_client.prompt("kg-edge-scoring", ...) -> scored edges
7. prompt_client.prompt("kg-edge-reasoning", ...) -> reasoning
8. triples_client.query(s, TG_CONTAINS, ...) -> doc tracing (returns [])
9. prompt_client.prompt("kg-synthesis", ...) -> final answer
6. reranker_client.rerank(queries, documents, limit) -> scored edges
7. triples_client.query(s, TG_CONTAINS, ...) -> doc tracing (returns [])
8. prompt_client.prompt("kg-synthesis", ...) -> final answer
"""
prompt_client = AsyncMock()
embeddings_client = AsyncMock()
graph_embeddings_client = AsyncMock()
triples_client = AsyncMock()
reranker_client = AsyncMock()
# 1. Concept extraction
prompt_responses = {}
@ -116,7 +116,7 @@ def build_mock_clients():
EmbeddingMatch(entity=Term(type=IRI, iri=ENTITY_B)),
]
# 4. Triple queries (follow_edges_batch) - return our edges
# 4. Triple queries (hop_and_filter) - return our edges
kg_triples = [
make_schema_triple(*EDGE_1),
make_schema_triple(*EDGE_2),
@ -130,9 +130,18 @@ def build_mock_clients():
return [] # No labels found, will fall back to URI
triples_client.query.side_effect = mock_label_query
# 6+7. Edge scoring and reasoning: dynamically score/reason about
# whatever edges the query method sends us, since edge IDs are computed
# from str(Term) representations which include the full dataclass repr.
# 6. Reranker: select all documents with high scores
async def mock_rerank(queries, documents, limit):
results = []
for i, doc in enumerate(documents):
result = MagicMock()
result.document_id = doc["id"]
result.query_id = queries[0]["id"] if queries else "0"
result.score = 0.9 - (i * 0.1)
results.append(result)
return results[:limit]
reranker_client.rerank.side_effect = mock_rerank
synthesis_answer = "Quantum computing applies physics principles to computation."
async def mock_prompt(template_id, variables=None, **kwargs):
@ -141,26 +150,6 @@ def build_mock_clients():
response_type="text",
text=prompt_responses["extract-concepts"],
)
elif template_id == "kg-edge-scoring":
# Score all edges highly, using the IDs that GraphRag computed
edges = variables.get("knowledge", [])
return PromptResult(
response_type="jsonl",
objects=[
{"id": e["id"], "score": 10 - i}
for i, e in enumerate(edges)
],
)
elif template_id == "kg-edge-reasoning":
# Provide reasoning for each edge
edges = variables.get("knowledge", [])
return PromptResult(
response_type="jsonl",
objects=[
{"id": e["id"], "reasoning": f"Relevant edge {i}"}
for i, e in enumerate(edges)
],
)
elif template_id == "kg-synthesis":
return PromptResult(
response_type="text",
@ -170,7 +159,8 @@ def build_mock_clients():
prompt_client.prompt.side_effect = mock_prompt
return prompt_client, embeddings_client, graph_embeddings_client, triples_client
return (prompt_client, embeddings_client, graph_embeddings_client,
triples_client, reranker_client)
# ---------------------------------------------------------------------------
@ -197,7 +187,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0, # skip semantic pre-filter for simplicity
)
assert len(events) == 5, (
@ -222,7 +212,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
expected_types = [
@ -260,7 +250,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
uris = [e["explain_id"] for e in events]
@ -297,7 +287,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
q_uri = events[0]["explain_id"]
@ -320,7 +310,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
gnd_uri = events[1]["explain_id"]
@ -344,7 +334,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
exp_uri = events[2]["explain_id"]
@ -355,10 +345,10 @@ class TestGraphRagQueryProvenance:
assert int(t.o.value) > 0
@pytest.mark.asyncio
async def test_focus_has_selected_edges_with_reasoning(self):
async def test_focus_has_selected_edges_with_concept_and_score(self):
"""
The focus event should carry selected edges as quoted triples
with reasoning text.
with cross-encoder concept and score metadata.
"""
clients = build_mock_clients()
rag = GraphRag(*clients)
@ -371,7 +361,6 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
foc_uri = events[3]["explain_id"]
@ -387,11 +376,19 @@ class TestGraphRagQueryProvenance:
for t in edge_t:
assert t.o.triple is not None, "tg:edge object must be a quoted triple"
# Should have reasoning
reasoning = find_triples(foc_triples, TG_REASONING)
assert len(reasoning) > 0, "Focus should have reasoning for selected edges"
reasoning_texts = {t.o.value for t in reasoning}
assert any(r for r in reasoning_texts), "Reasoning should not be empty"
# Edge selections should be typed as EdgeSelection
edge_sel_uris = [t.o.iri for t in selected]
for uri in edge_sel_uris:
assert has_type(foc_triples, uri, TG_EDGE_SELECTION)
# Should have concept and score
concepts = find_triples(foc_triples, TG_CONCEPT)
assert len(concepts) > 0, "Focus should have tg:concept for selected edges"
scores = find_triples(foc_triples, TG_SCORE)
assert len(scores) > 0, "Focus should have tg:score for selected edges"
for t in scores:
float(t.o.value) # Should be parseable as float
@pytest.mark.asyncio
async def test_synthesis_is_answer_type(self):
@ -407,7 +404,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
syn_uri = events[4]["explain_id"]
@ -429,7 +426,7 @@ class TestGraphRagQueryProvenance:
result_text, usage = await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
assert result_text == "Quantum computing applies physics principles to computation."
@ -449,7 +446,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
parent_uri=parent,
)
@ -465,7 +462,7 @@ class TestGraphRagQueryProvenance:
result_text, usage = await rag.query(
query="What is quantum computing?",
edge_score_limit=0,
)
assert result_text == "Quantum computing applies physics principles to computation."
@ -484,7 +481,7 @@ class TestGraphRagQueryProvenance:
await rag.query(
query="What is quantum computing?",
explain_callback=explain_callback,
edge_score_limit=0,
)
for event in events:

View file

@ -0,0 +1,157 @@
"""
Unit tests for trustgraph.storage.kw_index.fts5.service the SQLite FTS5
keyword index. Covers the MATCH-expression sanitizer (raw user text is not
valid FTS5 syntax), exact-term retrieval for the motivating cases (dotted
clause numbers, error codes, hyphenated identifiers), chunk re-ingestion
replacing rather than duplicating, (workspace, collection) scoping, and
collection deletion.
"""
import tempfile
from pathlib import Path
import pytest
from unittest.mock import AsyncMock
from unittest import IsolatedAsyncioTestCase
from trustgraph.schema import Chunk, Metadata, KeywordIndexRequest
from trustgraph.storage.kw_index.fts5.service import (
Processor, to_match_query, _table,
)
class TestMatchQuerySanitizer:
def test_plain_words_are_quoted_and_or_joined(self):
assert to_match_query("return policy") == '"return" OR "policy"'
def test_dotted_and_hyphenated_terms_survive(self):
# Raw "7.3.2" is an FTS5 syntax error; "AURA-7" parses "-" as a
# column filter. Quoting neutralizes both.
assert to_match_query("clause 7.3.2 AURA-7") == (
'"clause" OR "7.3.2" OR "AURA-7"'
)
def test_embedded_quotes_are_escaped(self):
assert to_match_query('say "hello"') == '"say" OR """hello"""'
def test_empty_and_quote_only_queries_yield_none(self):
assert to_match_query("") is None
assert to_match_query(" ") is None
assert to_match_query('"') is None
def make_processor(index_path):
# A real file, not :memory: — the service holds separate write and read
# connections, which only share a database through the filesystem.
processor = Processor(
taskgroup=AsyncMock(),
id="test-kw-index",
index_path=index_path,
)
# Config-pushed collection state isn't wired in unit tests
processor.collection_exists = lambda workspace, collection: True
return processor
def chunk(chunk_id, text, collection="default"):
return Chunk(
metadata=Metadata(id="doc1", collection=collection),
chunk=text.encode("utf-8"),
document_id=chunk_id,
)
CHUNKS = [
("c1", "Clause 7.3.2 states that indemnification obligations survive."),
("c2", "Clause 7.3.1 covers limitation of liability."),
("c3", "Error E4032 occurs when the connection pool is exhausted."),
]
class TestFts5KeywordIndex(IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.processor = make_processor(str(Path(self._tmp.name) / "kw.db"))
for chunk_id, text in CHUNKS:
await self.processor.index_chunk("ws", chunk("ws-" + chunk_id, text))
async def asyncTearDown(self):
self.processor.db.close()
self.processor.read_db.close()
self._tmp.cleanup()
async def query(self, text, collection="default", limit=0):
return await self.processor.query_keyword_index(
"ws", KeywordIndexRequest(
query=text, limit=limit, collection=collection,
),
)
async def test_exact_dotted_term_matches_only_its_clause(self):
matches = await self.query("7.3.2")
assert [m.chunk_id for m in matches] == ["ws-c1"]
async def test_error_code_matches(self):
matches = await self.query("E4032")
assert [m.chunk_id for m in matches] == ["ws-c3"]
async def test_scores_are_higher_is_better(self):
matches = await self.query("clause indemnification")
assert matches[0].chunk_id == "ws-c1"
assert all(m.score > 0 for m in matches)
# c1 matches both terms so it must outrank c2
by_id = {m.chunk_id: m.score for m in matches}
assert by_id["ws-c1"] > by_id["ws-c2"]
async def test_reingesting_a_chunk_replaces_it(self):
await self.processor.index_chunk(
"ws", chunk("ws-c1", "Completely different content now.")
)
assert await self.query("indemnification 7.3.2") == []
matches = await self.query("completely different")
assert [m.chunk_id for m in matches] == ["ws-c1"]
async def test_collections_are_isolated(self):
await self.processor.index_chunk(
"ws", chunk("other-c1", "indemnification text", collection="other")
)
default_ids = [m.chunk_id for m in await self.query("indemnification")]
other_ids = [
m.chunk_id
for m in await self.query("indemnification", collection="other")
]
assert "other-c1" not in default_ids
assert other_ids == ["other-c1"]
async def test_workspaces_are_isolated(self):
matches = await self.processor.query_keyword_index(
"someone-else", KeywordIndexRequest(
query="indemnification", collection="default",
),
)
assert matches == []
async def test_unindexed_collection_returns_empty_not_error(self):
assert await self.query("anything", collection="never-written") == []
async def test_hostile_query_text_is_inert(self):
# FTS5 operators and SQL fragments arrive as quoted phrases
assert await self.query('body: DROP TABLE OR NOT NEAR(') == []
async def test_limit_is_applied(self):
matches = await self.query("clause", limit=1)
assert len(matches) == 1
async def test_delete_collection_drops_the_index(self):
await self.processor.delete_collection("ws", "default")
assert await self.query("clause") == []
async def test_dropped_message_when_collection_missing(self):
self.processor.collection_exists = lambda w, c: False
await self.processor.index_chunk(
"ws", chunk("ws-c9", "should be dropped")
)
self.processor.collection_exists = lambda w, c: True
assert await self.query("dropped") == []

View file

@ -527,7 +527,8 @@ class AsyncFlowInstance:
return result.get("response", "")
async def document_rag(self, query: str, collection: str,
doc_limit: int = 10, **kwargs: Any) -> str:
doc_limit: int = 10, fetch_limit: int = 0,
**kwargs: Any) -> str:
"""
Execute document-based RAG query (non-streaming).
@ -541,7 +542,9 @@ class AsyncFlowInstance:
Args:
query: User query text
collection: Collection identifier containing documents
doc_limit: Maximum number of document chunks to retrieve (default: 10)
doc_limit: Document chunks selected into the prompt (default: 10)
fetch_limit: Candidate chunks fetched from the vector store before
reranking (default: 0 = derive from doc_limit)
**kwargs: Additional service-specific parameters
Returns:
@ -564,6 +567,7 @@ class AsyncFlowInstance:
"query": query,
"collection": collection,
"doc-limit": doc_limit,
"fetch-limit": fetch_limit,
"streaming": False
}
request_data.update(kwargs)
@ -646,6 +650,16 @@ class AsyncFlowInstance:
return await self.request("embeddings", request_data)
async def rerank(self, queries: list, documents: list, limit: int = 10, **kwargs: Any):
request_data = {
"queries": queries,
"documents": documents,
"limit": limit,
}
request_data.update(kwargs)
return await self.request("reranker", request_data)
async def triples_query(self, s=None, p=None, o=None, collection=None, limit=100, **kwargs: Any):
"""
Query RDF triples using pattern matching.

View file

@ -94,7 +94,9 @@ class AsyncSocketClient:
if resp.get("type") == "auth-ok":
if not self._workspace_explicit:
self.workspace = resp.get("workspace", self.workspace)
self.workspace = resp.get(
"default_workspace", self.workspace,
)
elif resp.get("type") == "auth-failed":
await self._socket.close()
raise ProtocolException(
@ -377,12 +379,14 @@ class AsyncSocketFlowInstance:
yield chunk.content
async def document_rag(self, query: str, collection: str,
doc_limit: int = 10, streaming: bool = False, **kwargs):
doc_limit: int = 10, fetch_limit: int = 0,
streaming: bool = False, **kwargs):
"""Document RAG with optional streaming"""
request = {
"query": query,
"collection": collection,
"doc-limit": doc_limit,
"fetch-limit": fetch_limit,
"streaming": streaming
}
request.update(kwargs)
@ -441,6 +445,19 @@ class AsyncSocketFlowInstance:
return await self.client._send_request("embeddings", self.flow_id, request)
async def rerank(self, queries: list, documents: list, limit: int = 10,
**kwargs):
request = {
"queries": queries,
"documents": documents,
"limit": limit,
}
request.update(kwargs)
return await self.client._send_request(
"reranker", self.flow_id, request,
)
async def triples_query(self, s=None, p=None, o=None, collection=None, limit=100, **kwargs):
"""Triple pattern query"""
request = {"limit": limit}

View file

@ -18,6 +18,7 @@ TG_EDGE_COUNT = TG + "edgeCount"
TG_SELECTED_EDGE = TG + "selectedEdge"
TG_EDGE = TG + "edge"
TG_REASONING = TG + "reasoning"
TG_SCORE = TG + "score"
TG_DOCUMENT = TG + "document"
TG_CONCEPT = TG + "concept"
TG_ENTITY = TG + "entity"
@ -66,10 +67,12 @@ RDFS_LABEL = "http://www.w3.org/2000/01/rdf-schema#label"
@dataclass
class EdgeSelection:
"""A selected edge with reasoning from GraphRAG Focus step."""
"""A selected edge with cross-encoder metadata from GraphRAG Focus step."""
uri: str
edge: Optional[Dict[str, str]] = None # {"s": ..., "p": ..., "o": ...}
reasoning: str = ""
concept: str = ""
score: Optional[float] = None
@dataclass
@ -209,7 +212,7 @@ class Exploration(ExplainEntity):
@dataclass
class Focus(ExplainEntity):
"""Focus entity - selected edges with LLM reasoning (GraphRAG only)."""
"""Focus entity - selected edges with cross-encoder scoring (GraphRAG only)."""
selected_edge_uris: List[str] = field(default_factory=list)
edge_selections: List[EdgeSelection] = field(default_factory=list)
@ -418,14 +421,26 @@ def parse_edge_selection_triples(triples: List[Tuple[str, str, Any]]) -> EdgeSel
uri = triples[0][0] if triples else ""
edge = None
reasoning = ""
concept = ""
score = None
for s, p, o in triples:
if p == TG_EDGE and isinstance(o, dict):
edge = o
elif p == TG_REASONING:
reasoning = o
elif p == TG_CONCEPT:
concept = o
elif p == TG_SCORE:
try:
score = float(o)
except (ValueError, TypeError):
score = None
return EdgeSelection(uri=uri, edge=edge, reasoning=reasoning)
return EdgeSelection(
uri=uri, edge=edge, reasoning=reasoning,
concept=concept, score=score,
)
def extract_term_value(term: Dict[str, Any]) -> Any:

View file

@ -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(
@ -415,7 +418,7 @@ class FlowInstance:
def document_rag(
self, query,collection="default",
doc_limit=10,
doc_limit=10, fetch_limit=0,
):
"""
Execute document-based Retrieval-Augmented Generation (RAG) query.
@ -426,7 +429,9 @@ class FlowInstance:
Args:
query: Natural language query
collection: Collection identifier (default: "default")
doc_limit: Maximum document chunks to retrieve (default: 10)
doc_limit: Document chunks selected into the prompt (default: 10)
fetch_limit: Candidate chunks fetched from the vector store before
reranking (default: 0 = derive from doc_limit)
Returns:
str: Generated response incorporating document context
@ -447,6 +452,7 @@ class FlowInstance:
"query": query,
"collection": collection,
"doc-limit": doc_limit,
"fetch-limit": fetch_limit,
}
result = self.request(
@ -491,6 +497,19 @@ class FlowInstance:
input
)["vectors"]
def rerank(self, queries, documents, limit=10):
input = {
"queries": queries,
"documents": documents,
"limit": limit,
}
return self.request(
"service/reranker",
input
)
def graph_embeddings_query(self, text, collection, limit=10):
"""
Query knowledge graph entities using semantic similarity.

View file

@ -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", ""),

View file

@ -168,7 +168,9 @@ class SocketClient:
if resp.get("type") == "auth-ok":
if self.workspace == "default":
self.workspace = resp.get("workspace", self.workspace)
self.workspace = resp.get(
"default_workspace", self.workspace,
)
elif resp.get("type") == "auth-failed":
await self._socket.close()
raise ProtocolException(
@ -680,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]]:
@ -697,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)
@ -723,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."""
@ -735,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,
}
@ -750,6 +756,7 @@ class SocketFlowInstance:
query: str,
collection: str,
doc_limit: int = 10,
fetch_limit: int = 0,
streaming: bool = False,
**kwargs: Any
) -> Union[TextCompletionResult, Iterator[RAGChunk]]:
@ -762,6 +769,7 @@ class SocketFlowInstance:
"query": query,
"collection": collection,
"doc-limit": doc_limit,
"fetch-limit": fetch_limit,
"streaming": streaming
}
request.update(kwargs)
@ -783,6 +791,7 @@ class SocketFlowInstance:
query: str,
collection: str,
doc_limit: int = 10,
fetch_limit: int = 0,
**kwargs: Any
) -> Iterator[Union[RAGChunk, ProvenanceEvent]]:
"""Execute document-based RAG query with explainability support."""
@ -790,6 +799,7 @@ class SocketFlowInstance:
"query": query,
"collection": collection,
"doc-limit": doc_limit,
"fetch-limit": fetch_limit,
"streaming": True,
"explainable": True,
}
@ -883,6 +893,19 @@ class SocketFlowInstance:
return self.client._send_request_sync("embeddings", self.flow_id, request, False)
def rerank(self, queries: list, documents: list, limit: int = 10,
**kwargs: Any) -> Dict[str, Any]:
request = {
"queries": queries,
"documents": documents,
"limit": limit,
}
request.update(kwargs)
return self.client._send_request_sync(
"reranker", self.flow_id, request, False,
)
def triples_query(
self,
s: Optional[Union[str, Dict[str, Any]]] = None,

View file

@ -42,6 +42,11 @@ from . dynamic_tool_service import DynamicToolService
from . tool_service_client import ToolServiceClientSpec
from . agent_client import AgentClientSpec
from . structured_query_client import StructuredQueryClientSpec
from . reranker_client import RerankerClientSpec
from . reranker_service import RerankerService
from . keyword_index_service import KeywordIndexService
from . keyword_index_client import KeywordIndexClientSpec, KeywordIndexClient
from . row_embeddings_query_client import RowEmbeddingsQueryClientSpec
from . collection_config_handler import CollectionConfigHandler
from . audit_publisher import AuditPublisher

View file

@ -0,0 +1,42 @@
import json
import logging
from datetime import datetime, timezone
from uuid import uuid4
from . producer import Producer
from . metrics import ProducerMetrics
from trustgraph.schema import AuditEvent, audit_events_queue
logger = logging.getLogger(__name__)
class AuditPublisher:
def __init__(self, backend, component_name, processor_id=None):
self.component_name = component_name
self.producer = Producer(
backend=backend,
topic=audit_events_queue,
schema=AuditEvent,
metrics=ProducerMetrics(
processor=processor_id or component_name,
flow=None,
name="audit-events",
),
)
async def emit(self, event_type, payload):
event = AuditEvent(
schema_version=1,
event_id=str(uuid4()),
event_type=event_type,
timestamp=datetime.now(timezone.utc).isoformat(),
producer=self.component_name,
payload_json=json.dumps(payload),
)
try:
await self.producer.send(event)
except Exception as e:
logger.warning(f"Failed to emit audit event: {e}")

View file

@ -103,35 +103,19 @@ def resolve_cassandra_config(
host: Optional[str] = None,
username: Optional[str] = None,
password: Optional[str] = None,
default_keyspace: Optional[str] = None
default_keyspace: Optional[str] = None,
replication_factor: Optional[int] = None,
) -> Tuple[List[str], Optional[str], Optional[str], Optional[str], int]:
"""
Resolve Cassandra configuration from various sources.
Can accept either argparse args object or explicit parameters.
Converts host string to list format for Cassandra driver.
Args:
args: Optional argparse namespace with cassandra_host, cassandra_username, cassandra_password, cassandra_keyspace, cassandra_replication_factor
host: Optional explicit host parameter (overrides args)
username: Optional explicit username parameter (overrides args)
password: Optional explicit password parameter (overrides args)
default_keyspace: Optional default keyspace if not specified elsewhere
Returns:
tuple: (hosts_list, username, password, keyspace, replication_factor)
"""
# If args provided, extract values
keyspace = None
replication_factor = 1
if args is not None:
host = host or getattr(args, 'cassandra_host', None)
username = username or getattr(args, 'cassandra_username', None)
password = password or getattr(args, 'cassandra_password', None)
keyspace = getattr(args, 'cassandra_keyspace', None)
replication_factor = getattr(args, 'cassandra_replication_factor', 1)
replication_factor = replication_factor or getattr(
args, 'cassandra_replication_factor', None
)
# Apply defaults if still None
defaults = get_cassandra_defaults()
host = host or defaults['host']
username = username or defaults['username']

View file

@ -62,47 +62,48 @@ class IamClient(RequestResponse):
)
return resp.user
async def authenticate_anonymous(self, timeout=IAM_TIMEOUT):
async def authenticate_anonymous(self, timeout=IAM_TIMEOUT,
request_id="", client_ip=""):
"""Request anonymous access from the IAM regime.
Returns ``(user_id, workspace, roles)`` if the regime permits
anonymous access, or raises ``RuntimeError`` with error type
``auth-failed`` if it does not."""
Returns ``(user_id, default_workspace, roles)`` if the regime
permits anonymous access, or raises ``RuntimeError`` with
error type ``auth-failed`` if it does not."""
resp = await self._request(
operation="authenticate-anonymous",
request_id=request_id,
client_ip=client_ip,
timeout=timeout,
)
return (
resp.resolved_user_id,
resp.resolved_workspace,
resp.resolved_default_workspace,
list(resp.resolved_roles),
)
async def resolve_api_key(self, api_key, timeout=IAM_TIMEOUT):
async def resolve_api_key(self, api_key, timeout=IAM_TIMEOUT,
request_id="", client_ip=""):
"""Resolve a plaintext API key to its identity triple.
Returns ``(user_id, workspace, roles)`` or raises
Returns ``(user_id, default_workspace, roles)`` or raises
``RuntimeError`` with error type ``auth-failed`` if the key is
unknown / expired / revoked.
Note: the ``roles`` value is a regime-internal hint and is
not used by the gateway directly under the IAM contract;
all authorisation decisions go through ``authorise()``.
Returned here only for backward compatibility with callers
that haven't migrated."""
unknown / expired / revoked."""
resp = await self._request(
operation="resolve-api-key",
api_key=api_key,
request_id=request_id,
client_ip=client_ip,
timeout=timeout,
)
return (
resp.resolved_user_id,
resp.resolved_workspace,
resp.resolved_default_workspace,
list(resp.resolved_roles),
)
async def authorise(self, identity_handle, capability,
resource, parameters, timeout=IAM_TIMEOUT):
resource, parameters, timeout=IAM_TIMEOUT,
request_id="", client_ip=""):
"""Ask the IAM regime whether ``identity_handle`` may perform
``capability`` on ``resource`` given ``parameters``.
@ -117,6 +118,8 @@ class IamClient(RequestResponse):
capability=capability,
resource_json=json.dumps(resource or {}, sort_keys=True),
parameters_json=json.dumps(parameters or {}, sort_keys=True),
request_id=request_id,
client_ip=client_ip,
timeout=timeout,
)
return resp.decision_allow, resp.decision_ttl_seconds
@ -192,7 +195,7 @@ class IamClient(RequestResponse):
)
async def login(self, username, password, workspace="",
timeout=IAM_TIMEOUT):
timeout=IAM_TIMEOUT, request_id="", client_ip=""):
"""Validate credentials and return ``(jwt, expires_iso)``.
``workspace`` is optional; defaults at the server to the
OSS default workspace."""
@ -201,6 +204,8 @@ class IamClient(RequestResponse):
workspace=workspace,
username=username,
password=password,
request_id=request_id,
client_ip=client_ip,
timeout=timeout,
)
return resp.jwt, resp.jwt_expires

View file

@ -0,0 +1,44 @@
import logging
from . request_response_spec import RequestResponse, RequestResponseSpec
from .. schema import KeywordIndexRequest, KeywordIndexResponse
# Module logger
logger = logging.getLogger(__name__)
class KeywordIndexClient(RequestResponse):
async def query(self, query, limit=20, collection="default", timeout=30):
resp = await self.request(
KeywordIndexRequest(
query = query,
limit = limit,
collection = collection
),
timeout=timeout
)
logger.debug("Keyword index response: %s", resp)
if resp.error:
raise RuntimeError(resp.error.message)
# Return ChunkMatch objects with chunk_id and score
return resp.chunks
class KeywordIndexClientSpec(RequestResponseSpec):
def __init__(
self, request_name, response_name,
):
super(KeywordIndexClientSpec, self).__init__(
request_name = request_name,
request_schema = KeywordIndexRequest,
response_name = response_name,
response_schema = KeywordIndexResponse,
impl = KeywordIndexClient,
# Flow definitions predating the keyword index don't declare
# these topics; bind only where they exist so one stale
# definition can't wedge the processor.
optional = True,
)

View file

@ -0,0 +1,132 @@
"""
Keyword index service base class. A single service owns both sides of the
lexical index: it consumes Chunk messages off the ingestion stream (the last
message in the pipeline that still carries chunk text) and answers keyword
search requests over what it has indexed. Unlike the vector stores, ingest
and query are not split into two processors: the first backend (SQLite FTS5)
is a single-file index that cannot be shared between containers, so one
process must own it. Backends with a server (Elasticsearch/OpenSearch) can
still be split later behind the same schema.
"""
from __future__ import annotations
from argparse import ArgumentParser
import logging
from .. schema import Chunk
from .. schema import KeywordIndexRequest, KeywordIndexResponse
from .. schema import Error
from .. exceptions import TooManyRequests
from . flow_processor import FlowProcessor
from . consumer_spec import ConsumerSpec
from . producer_spec import ProducerSpec
# Module logger
logger = logging.getLogger(__name__)
default_ident = "kw-index"
default_concurrency = 10
class KeywordIndexService(FlowProcessor):
def __init__(self, **params):
id = params.get("id")
concurrency = params.get("concurrency", default_concurrency)
super(KeywordIndexService, self).__init__(
**params | { "id": id }
)
self.register_specification(
ConsumerSpec(
name = "input",
schema = Chunk,
handler = self.on_chunk,
)
)
self.register_specification(
ConsumerSpec(
name = "request",
schema = KeywordIndexRequest,
handler = self.on_request,
concurrency = concurrency,
)
)
self.register_specification(
ProducerSpec(
name = "response",
schema = KeywordIndexResponse,
)
)
async def on_chunk(self, msg, consumer, flow):
try:
request = msg.value()
# Workspace comes from the flow the message arrived on.
await self.index_chunk(flow.workspace, request)
except TooManyRequests as e:
raise e
except Exception as e:
logger.error(f"Exception in keyword index store: {e}", exc_info=True)
raise e
async def on_request(self, msg, consumer, flow):
try:
request = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
logger.debug(f"Handling keyword index query request {id}...")
chunks = await self.query_keyword_index(
flow.workspace, request,
)
logger.debug("Sending keyword index query response...")
r = KeywordIndexResponse(chunks=chunks, error=None)
await flow("response").send(r, properties={"id": id})
logger.debug("Keyword index query request completed")
except Exception as e:
logger.error(f"Exception in keyword index query service: {e}", exc_info=True)
logger.info("Sending error response...")
r = KeywordIndexResponse(
error=Error(
type = "keyword-index-query-error",
message = str(e),
),
chunks=[],
)
await flow("response").send(r, properties={"id": id})
@staticmethod
def add_args(parser: ArgumentParser) -> None:
FlowProcessor.add_args(parser)
parser.add_argument(
'-c', '--concurrency',
type=int,
default=default_concurrency,
help=f'Number of concurrent requests (default: {default_concurrency})'
)

View file

@ -157,21 +157,6 @@ class PromptClient(RequestResponse):
timeout = timeout,
)
async def kg_prompt(self, query, kg, timeout=600, streaming=False, chunk_callback=None):
return await self.prompt(
id = "kg-prompt",
variables = {
"query": query,
"knowledge": [
{ "s": v[0], "p": v[1], "o": v[2] }
for v in kg
]
},
timeout = timeout,
streaming = streaming,
chunk_callback = chunk_callback,
)
async def document_prompt(self, query, documents, timeout=600, streaming=False, chunk_callback=None):
return await self.prompt(
id = "document-prompt",

View file

@ -0,0 +1,87 @@
import os
import argparse
from typing import Optional, Any, Tuple
def get_qdrant_defaults() -> dict:
return {
'url': os.getenv('QDRANT_URL', 'http://localhost:6333'),
'api_key': os.getenv('QDRANT_API_KEY'),
'replication_factor': int(os.getenv('QDRANT_REPLICATION_FACTOR', '1')),
'shard_number': int(os.getenv('QDRANT_SHARD_NUMBER', '1')),
}
def add_qdrant_args(parser: argparse.ArgumentParser) -> None:
defaults = get_qdrant_defaults()
url_help = f"Qdrant URL (default: {defaults['url']})"
if 'QDRANT_URL' in os.environ:
url_help += " [from QDRANT_URL]"
api_key_help = "Qdrant API key"
if defaults['api_key']:
api_key_help += " (default: <set>)"
if 'QDRANT_API_KEY' in os.environ:
api_key_help += " [from QDRANT_API_KEY]"
replication_help = f"Qdrant collection replication factor (default: {defaults['replication_factor']})"
if 'QDRANT_REPLICATION_FACTOR' in os.environ:
replication_help += " [from QDRANT_REPLICATION_FACTOR]"
shard_help = f"Qdrant collection shard number (default: {defaults['shard_number']})"
if 'QDRANT_SHARD_NUMBER' in os.environ:
shard_help += " [from QDRANT_SHARD_NUMBER]"
parser.add_argument(
'--store-uri',
default=defaults['url'],
help=url_help,
)
parser.add_argument(
'--api-key',
default=defaults['api_key'],
help=api_key_help,
)
parser.add_argument(
'--qdrant-replication-factor',
type=int,
default=defaults['replication_factor'],
help=replication_help,
)
parser.add_argument(
'--qdrant-shard-number',
type=int,
default=defaults['shard_number'],
help=shard_help,
)
def resolve_qdrant_config(
args: Optional[Any] = None,
url: Optional[str] = None,
api_key: Optional[str] = None,
replication_factor: Optional[int] = None,
shard_number: Optional[int] = None,
) -> Tuple[str, Optional[str], int, int]:
if args is not None:
url = url or getattr(args, 'store_uri', None)
api_key = api_key or getattr(args, 'api_key', None)
replication_factor = replication_factor or getattr(
args, 'qdrant_replication_factor', None
)
shard_number = shard_number or getattr(
args, 'qdrant_shard_number', None
)
defaults = get_qdrant_defaults()
url = url or defaults['url']
api_key = api_key or defaults['api_key']
replication_factor = replication_factor or defaults['replication_factor']
shard_number = shard_number or defaults['shard_number']
return url, api_key, replication_factor, shard_number

View file

@ -109,16 +109,28 @@ class RequestResponse(Subscriber):
class RequestResponseSpec(Spec):
def __init__(
self, request_name, request_schema, response_name,
response_schema, impl=RequestResponse
response_schema, impl=RequestResponse, optional=False
):
self.request_name = request_name
self.request_schema = request_schema
self.response_name = response_name
self.response_schema = response_schema
self.impl = impl
self.optional = optional
def add(self, flow: Any, processor: Any, definition: dict[str, Any]) -> None:
# An optional client binds only when the flow definition declares
# its topics. Older definitions predating the topics would otherwise
# KeyError here during Flow construction, which wedges the whole
# processor in a start-flow retry loop; skipping instead leaves
# flow(name) returning None for the caller to handle per-request.
topics = definition.get("topics", {})
if self.optional and (
self.request_name not in topics
or self.response_name not in topics):
return
request_metrics = ProducerMetrics(
processor = flow.id, flow = flow.name, name = self.request_name
)

View file

@ -0,0 +1,43 @@
from . request_response_spec import RequestResponse, RequestResponseSpec
from .. schema import (
RerankerRequest, RerankerResponse,
RerankerQuery, RerankerDocument,
)
class RerankerClient(RequestResponse):
async def rerank(self, queries, documents, limit=10, timeout=300):
resp = await self.request(
RerankerRequest(
queries=[
RerankerQuery(query_id=q["id"], query_text=q["text"])
for q in queries
],
documents=[
RerankerDocument(
document_id=d["id"], document_text=d["text"]
)
for d in documents
],
limit=limit,
),
timeout=timeout
)
if resp.error:
raise RuntimeError(resp.error.message)
return resp.results
class RerankerClientSpec(RequestResponseSpec):
def __init__(
self, request_name, response_name,
):
super(RerankerClientSpec, self).__init__(
request_name = request_name,
request_schema = RerankerRequest,
response_name = response_name,
response_schema = RerankerResponse,
impl = RerankerClient,
)

View file

@ -0,0 +1,109 @@
from __future__ import annotations
from argparse import ArgumentParser
import logging
from .. schema import (
RerankerRequest, RerankerResponse, RerankerResult, Error,
)
from .. exceptions import TooManyRequests
from .. base import FlowProcessor, ConsumerSpec, ProducerSpec, ParameterSpec
logger = logging.getLogger(__name__)
default_ident = "reranker"
default_concurrency = 1
class RerankerService(FlowProcessor):
def __init__(self, **params):
id = params.get("id")
concurrency = params.get("concurrency", 1)
super(RerankerService, self).__init__(**params | {
"id": id,
"concurrency": concurrency,
})
self.register_specification(
ConsumerSpec(
name = "request",
schema = RerankerRequest,
handler = self.on_request,
concurrency = concurrency,
)
)
self.register_specification(
ProducerSpec(
name = "response",
schema = RerankerResponse
)
)
self.register_specification(
ParameterSpec(
name = "model",
)
)
async def on_request(self, msg, consumer, flow):
try:
request = msg.value()
id = msg.properties()["id"]
logger.debug(f"Handling reranker request {id}...")
model = flow("model")
results = await self.on_rerank(
request.queries, request.documents,
request.limit, model=model,
)
await flow("response").send(
RerankerResponse(
error = None,
results = results,
),
properties={"id": id}
)
logger.debug("Reranker request handled successfully")
except TooManyRequests as e:
raise e
except Exception as e:
logger.error(f"Exception in reranker service: {e}", exc_info=True)
logger.info("Sending error response...")
await flow.producer["response"].send(
RerankerResponse(
error=Error(
type = "reranker-error",
message = str(e),
),
results=[],
),
properties={"id": id}
)
@staticmethod
def add_args(parser: ArgumentParser) -> None:
parser.add_argument(
'-c', '--concurrency',
type=int,
default=default_concurrency,
help=f'Concurrent processing threads (default: {default_concurrency})'
)
FlowProcessor.add_args(parser)

View file

@ -140,20 +140,6 @@ class PromptClient(BaseClient):
timeout=timeout
)
def request_kg_prompt(self, query, kg, timeout=300):
return self.request(
id="kg-prompt",
variables={
"query": query,
"knowledge": [
{ "s": v[0], "p": v[1], "o": v[2] }
for v in kg
]
},
timeout=timeout
)
def request_document_prompt(self, query, documents, timeout=300):
return self.request(

View file

@ -27,6 +27,7 @@ from .translators.rows_query import RowsQueryRequestTranslator, RowsQueryRespons
from .translators.nlp_query import QuestionToStructuredQueryRequestTranslator, QuestionToStructuredQueryResponseTranslator
from .translators.structured_query import StructuredQueryRequestTranslator, StructuredQueryResponseTranslator
from .translators.diagnosis import StructuredDataDiagnosisRequestTranslator, StructuredDataDiagnosisResponseTranslator
from .translators.reranker import RerankerRequestTranslator, RerankerResponseTranslator
from .translators.collection import CollectionManagementRequestTranslator, CollectionManagementResponseTranslator
from .translators.sparql_query import SparqlQueryRequestTranslator, SparqlQueryResponseTranslator
@ -163,6 +164,12 @@ TranslatorRegistry.register_service(
SparqlQueryResponseTranslator()
)
TranslatorRegistry.register_service(
"reranker",
RerankerRequestTranslator(),
RerankerResponseTranslator()
)
# Register single-direction translators for document loading
TranslatorRegistry.register_request("document", DocumentTranslator())
TranslatorRegistry.register_request("text-document", TextDocumentTranslator())

View file

@ -20,3 +20,4 @@ from .embeddings_query import (
)
from .rows_query import RowsQueryRequestTranslator, RowsQueryResponseTranslator
from .diagnosis import StructuredDataDiagnosisRequestTranslator, StructuredDataDiagnosisResponseTranslator
from .reranker import RerankerRequestTranslator, RerankerResponseTranslator

View file

@ -5,6 +5,7 @@ from ...schema import (
UserInput, UserRecord,
WorkspaceInput, WorkspaceRecord,
ApiKeyInput, ApiKeyRecord,
GroupInput, GrantInput,
)
from .base import MessageTranslator
@ -43,12 +44,31 @@ def _api_key_input_from_dict(d):
)
def _group_input_from_dict(d):
if d is None:
return None
return GroupInput(
name=d.get("name", ""),
description=d.get("description", ""),
enabled=d.get("enabled", True),
)
def _grant_input_from_dict(d):
if d is None:
return None
return GrantInput(
capability=d.get("capability", ""),
workspace=d.get("workspace", ""),
)
def _user_record_to_dict(r):
if r is None:
return None
return {
"id": r.id,
"workspace": r.workspace,
"default_workspace": r.default_workspace,
"username": r.username,
"name": r.name,
"email": r.email,
@ -102,6 +122,15 @@ class IamRequestTranslator(MessageTranslator):
data.get("workspace_record")
),
key=_api_key_input_from_dict(data.get("key")),
group_id=data.get("group_id", ""),
member_type=data.get("member_type", ""),
member_id=data.get("member_id", ""),
group=_group_input_from_dict(data.get("group")),
grant=_grant_input_from_dict(data.get("grant")),
capability=data.get("capability", ""),
resource_json=data.get("resource_json", ""),
parameters_json=data.get("parameters_json", ""),
authorise_checks=data.get("authorise_checks", ""),
)
def encode(self, obj: IamRequest) -> Dict[str, Any]:
@ -109,6 +138,9 @@ class IamRequestTranslator(MessageTranslator):
for fname in (
"workspace", "actor", "user_id", "username", "key_id",
"api_key", "password", "new_password",
"group_id", "member_type", "member_id",
"capability", "resource_json", "parameters_json",
"authorise_checks",
):
v = getattr(obj, fname, "")
if v:
@ -135,6 +167,17 @@ class IamRequestTranslator(MessageTranslator):
"name": obj.key.name,
"expires": obj.key.expires,
}
if obj.group is not None:
result["group"] = {
"name": obj.group.name,
"description": obj.group.description,
"enabled": obj.group.enabled,
}
if obj.grant is not None:
result["grant"] = {
"capability": obj.grant.capability,
"workspace": obj.grant.workspace,
}
return result
@ -175,8 +218,8 @@ class IamResponseTranslator(MessageTranslator):
result["signing_key_public"] = obj.signing_key_public
if obj.resolved_user_id:
result["resolved_user_id"] = obj.resolved_user_id
if obj.resolved_workspace:
result["resolved_workspace"] = obj.resolved_workspace
if obj.resolved_default_workspace:
result["resolved_default_workspace"] = obj.resolved_default_workspace
if obj.resolved_roles:
result["resolved_roles"] = list(obj.resolved_roles)
if obj.temporary_password:
@ -190,6 +233,23 @@ class IamResponseTranslator(MessageTranslator):
# setup, so it can't be dropped by a truthy-only filter.
result["bootstrap_available"] = bool(obj.bootstrap_available)
# authorise / authorise-many outputs.
if obj.decision_allow:
result["decision_allow"] = obj.decision_allow
if obj.decision_ttl_seconds:
result["decision_ttl_seconds"] = obj.decision_ttl_seconds
if obj.decisions_json:
result["decisions_json"] = obj.decisions_json
# Enterprise IAM outputs.
for fname in (
"group_json", "groups_json", "members_json",
"grants_json", "effective_permissions_json",
):
v = getattr(obj, fname, "")
if v:
result[fname] = v
return result
def encode_with_completion(

View file

@ -0,0 +1,73 @@
from typing import Dict, Any, Tuple
from ...schema import (
RerankerRequest, RerankerResponse,
RerankerQuery, RerankerDocument, RerankerResult,
)
from .base import MessageTranslator
class RerankerRequestTranslator(MessageTranslator):
def decode(self, data: Dict[str, Any]) -> RerankerRequest:
return RerankerRequest(
queries=[
RerankerQuery(
query_id=q["query_id"],
query_text=q["query_text"],
)
for q in data.get("queries", [])
],
documents=[
RerankerDocument(
document_id=d["document_id"],
document_text=d["document_text"],
)
for d in data.get("documents", [])
],
limit=data.get("limit", 10),
)
def encode(self, obj: RerankerRequest) -> Dict[str, Any]:
return {
"queries": [
{"query_id": q.query_id, "query_text": q.query_text}
for q in obj.queries
],
"documents": [
{"document_id": d.document_id, "document_text": d.document_text}
for d in obj.documents
],
"limit": obj.limit,
}
class RerankerResponseTranslator(MessageTranslator):
def decode(self, data: Dict[str, Any]) -> RerankerResponse:
return RerankerResponse(
results=[
RerankerResult(
document_id=r["document_id"],
query_id=r["query_id"],
score=r["score"],
)
for r in data.get("results", [])
],
)
def encode(self, obj: RerankerResponse) -> Dict[str, Any]:
return {
"results": [
{
"document_id": r.document_id,
"query_id": r.query_id,
"score": r.score,
}
for r in obj.results
],
}
def encode_with_completion(
self, obj: RerankerResponse
) -> Tuple[Dict[str, Any], bool]:
return self.encode(obj), True

View file

@ -12,6 +12,7 @@ class DocumentRagRequestTranslator(MessageTranslator):
query=data["query"],
collection=data.get("collection", "default"),
doc_limit=int(data.get("doc-limit", 20)),
fetch_limit=int(data.get("fetch-limit", 0)),
streaming=data.get("streaming", False)
)
@ -20,6 +21,7 @@ class DocumentRagRequestTranslator(MessageTranslator):
"query": obj.query,
"collection": obj.collection,
"doc-limit": obj.doc_limit,
"fetch-limit": obj.fetch_limit,
"streaming": getattr(obj, "streaming", False)
}
@ -101,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)
)
@ -114,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)
}

View file

@ -64,6 +64,8 @@ from . uris import (
docrag_question_uri,
docrag_grounding_uri,
docrag_exploration_uri,
docrag_focus_uri,
chunk_selection_uri,
docrag_synthesis_uri,
)
@ -89,9 +91,13 @@ from . namespaces import (
TG_IMAGE_TYPE, TG_SUBGRAPH_TYPE,
# Query-time provenance predicates (GraphRAG)
TG_QUERY, TG_CONCEPT, TG_ENTITY,
TG_EDGE_COUNT, TG_SELECTED_EDGE, TG_REASONING,
TG_EDGE_COUNT, TG_SELECTED_EDGE, TG_REASONING, TG_SCORE,
# Edge selection entity type
TG_EDGE_SELECTION,
# Query-time provenance predicates (DocumentRAG)
TG_CHUNK_COUNT, TG_SELECTED_CHUNK,
# Chunk selection entity type
TG_CHUNK_SELECTION,
# Explainability entity types
TG_QUESTION, TG_GROUNDING, TG_EXPLORATION, TG_FOCUS, TG_SYNTHESIS,
TG_ANALYSIS, TG_CONCLUSION,
@ -130,6 +136,7 @@ from . triples import (
# Query-time provenance triple builders (DocumentRAG)
docrag_question_triples,
docrag_exploration_triples,
docrag_chunk_selection_triples,
docrag_synthesis_triples,
# Utility
set_graph,
@ -194,6 +201,8 @@ __all__ = [
"docrag_question_uri",
"docrag_grounding_uri",
"docrag_exploration_uri",
"docrag_focus_uri",
"chunk_selection_uri",
"docrag_synthesis_uri",
# Namespaces
"PROV", "PROV_ENTITY", "PROV_ACTIVITY", "PROV_AGENT",
@ -212,9 +221,13 @@ __all__ = [
"TG_CHUNK_TYPE", "TG_IMAGE_TYPE", "TG_SUBGRAPH_TYPE",
# Query-time provenance predicates (GraphRAG)
"TG_QUERY", "TG_CONCEPT", "TG_ENTITY",
"TG_EDGE_COUNT", "TG_SELECTED_EDGE", "TG_REASONING",
"TG_EDGE_COUNT", "TG_SELECTED_EDGE", "TG_REASONING", "TG_SCORE",
# Edge selection entity type
"TG_EDGE_SELECTION",
# Query-time provenance predicates (DocumentRAG)
"TG_CHUNK_COUNT", "TG_SELECTED_CHUNK",
# Chunk selection entity type
"TG_CHUNK_SELECTION",
# Explainability entity types
"TG_QUESTION", "TG_GROUNDING", "TG_EXPLORATION", "TG_FOCUS", "TG_SYNTHESIS",
"TG_ANALYSIS", "TG_CONCLUSION",
@ -250,6 +263,7 @@ __all__ = [
# Query-time provenance triple builders (DocumentRAG)
"docrag_question_triples",
"docrag_exploration_triples",
"docrag_chunk_selection_triples",
"docrag_synthesis_triples",
# Agent provenance triple builders
"agent_session_triples",

View file

@ -66,12 +66,19 @@ TG_EDGE_COUNT = TG + "edgeCount"
TG_SELECTED_EDGE = TG + "selectedEdge"
TG_EDGE = TG + "edge"
TG_REASONING = TG + "reasoning"
TG_SCORE = TG + "score"
TG_DOCUMENT = TG + "document" # Reference to document in librarian
# Edge selection entity type (cross-encoder scored edge in Focus)
TG_EDGE_SELECTION = TG + "EdgeSelection"
# Query-time provenance predicates (DocumentRAG)
TG_CHUNK_COUNT = TG + "chunkCount"
TG_SELECTED_CHUNK = TG + "selectedChunk"
# Chunk selection entity type (cross-encoder reranked chunk in Focus)
TG_CHUNK_SELECTION = TG + "ChunkSelection"
# Extraction provenance entity types
TG_DOCUMENT_TYPE = TG + "Document"
TG_PAGE_TYPE = TG + "Page"

View file

@ -24,10 +24,14 @@ from . namespaces import (
TG_ELEMENT_TYPES, TG_TABLE_COUNT, TG_IMAGE_COUNT,
# Query-time provenance predicates (GraphRAG)
TG_QUERY, TG_CONCEPT, TG_ENTITY,
TG_EDGE_COUNT, TG_SELECTED_EDGE, TG_EDGE, TG_REASONING,
TG_EDGE_COUNT, TG_SELECTED_EDGE, TG_EDGE, TG_REASONING, TG_SCORE,
TG_DOCUMENT,
# Edge selection entity type
TG_EDGE_SELECTION,
# Query-time provenance predicates (DocumentRAG)
TG_CHUNK_COUNT, TG_SELECTED_CHUNK,
# Chunk selection entity type
TG_CHUNK_SELECTION,
# Explainability entity types
TG_QUESTION, TG_GROUNDING, TG_EXPLORATION, TG_FOCUS, TG_SYNTHESIS,
# Unifying types
@ -38,7 +42,10 @@ from . namespaces import (
TG_IN_TOKEN, TG_OUT_TOKEN,
)
from . uris import activity_uri, agent_uri, subgraph_uri, edge_selection_uri
from . uris import (
activity_uri, agent_uri, subgraph_uri, edge_selection_uri,
chunk_selection_uri,
)
def set_graph(triples: List[Triple], graph: str) -> List[Triple]:
@ -536,10 +543,9 @@ def focus_triples(
_triple(focus_uri, PROV_WAS_DERIVED_FROM, _iri(exploration_uri)),
]
# Add each selected edge with its reasoning via intermediate entity
# Add each selected edge with metadata via intermediate entity
for idx, edge_info in enumerate(selected_edges_with_reasoning):
edge = edge_info.get("edge")
reasoning = edge_info.get("reasoning", "")
if edge:
s, p, o = edge
@ -552,13 +558,32 @@ def focus_triples(
_triple(focus_uri, TG_SELECTED_EDGE, _iri(edge_sel_uri))
)
# Type the edge selection entity
triples.append(
_triple(edge_sel_uri, RDF_TYPE, _iri(TG_EDGE_SELECTION))
)
# Attach quoted triple to edge selection entity
quoted = _quoted_triple(s, p, o)
triples.append(
Triple(s=_iri(edge_sel_uri), p=_iri(TG_EDGE), o=quoted)
)
# Attach reasoning to edge selection entity
# Structured cross-encoder metadata
concept = edge_info.get("concept")
if concept:
triples.append(
_triple(edge_sel_uri, TG_CONCEPT, _literal(concept))
)
score = edge_info.get("score")
if score is not None:
triples.append(
_triple(edge_sel_uri, TG_SCORE, _literal(str(score)))
)
# Legacy reasoning text (for non-cross-encoder callers)
reasoning = edge_info.get("reasoning", "")
if reasoning:
triples.append(
_triple(edge_sel_uri, TG_REASONING, _literal(reasoning))
@ -698,6 +723,75 @@ def docrag_exploration_triples(
return triples
def docrag_chunk_selection_triples(
focus_uri: str,
exploration_uri: str,
selected_chunks_with_scores: List[dict],
session_id: str,
) -> List[Triple]:
"""
Build triples for a document RAG focus entity (chunks selected by the
cross-encoder reranker).
Mirrors GraphRAG's focus_triples / tg:EdgeSelection pattern: a Focus entity
derived from exploration, with one ChunkSelection sub-entity per surviving
chunk carrying the chunk reference and the reranker score.
Structure:
<focus> a tg:Focus ; prov:wasDerivedFrom <exploration> .
<focus> tg:selectedChunk <chunk_sel_0> .
<chunk_sel_0> a tg:ChunkSelection .
<chunk_sel_0> tg:document <chunk_id> .
<chunk_sel_0> tg:score "0.97" .
Args:
focus_uri: URI of the focus entity (from docrag_focus_uri)
exploration_uri: URI of the parent exploration entity
selected_chunks_with_scores: List of dicts with 'chunk_id' and 'score'
session_id: Session UUID for generating chunk selection URIs
Returns:
List of Triple objects
"""
triples = [
_triple(focus_uri, RDF_TYPE, _iri(PROV_ENTITY)),
_triple(focus_uri, RDF_TYPE, _iri(TG_FOCUS)),
_triple(focus_uri, RDFS_LABEL, _literal("Chunk Selection")),
_triple(focus_uri, PROV_WAS_DERIVED_FROM, _iri(exploration_uri)),
]
for idx, chunk_info in enumerate(selected_chunks_with_scores):
chunk_id = chunk_info.get("chunk_id")
if not chunk_id:
continue
chunk_sel_uri = chunk_selection_uri(session_id, idx)
# Link focus to chunk selection entity
triples.append(
_triple(focus_uri, TG_SELECTED_CHUNK, _iri(chunk_sel_uri))
)
# Type the chunk selection entity
triples.append(
_triple(chunk_sel_uri, RDF_TYPE, _iri(TG_CHUNK_SELECTION))
)
# Reference the actual chunk (in librarian)
triples.append(
_triple(chunk_sel_uri, TG_DOCUMENT, _iri(chunk_id))
)
# Cross-encoder score
score = chunk_info.get("score")
if score is not None:
triples.append(
_triple(chunk_sel_uri, TG_SCORE, _literal(str(score)))
)
return triples
def docrag_synthesis_triples(
synthesis_uri: str,
exploration_uri: str,

View file

@ -309,6 +309,35 @@ def docrag_exploration_uri(session_id: str) -> str:
return f"urn:trustgraph:docrag:{session_id}/exploration"
def docrag_focus_uri(session_id: str) -> str:
"""
Generate URI for a document RAG focus entity (chunks selected by the
cross-encoder reranker).
Args:
session_id: The session UUID.
Returns:
URN in format: urn:trustgraph:docrag:{uuid}/focus
"""
return f"urn:trustgraph:docrag:{session_id}/focus"
def chunk_selection_uri(session_id: str, chunk_index: int) -> str:
"""
Generate URI for a chunk selection item (links a reranked chunk to its
score). Mirrors edge_selection_uri for GraphRAG.
Args:
session_id: The session UUID.
chunk_index: Index of this chunk in the selection (0-based).
Returns:
URN in format: urn:trustgraph:prov:chunk:{uuid}:{index}
"""
return f"urn:trustgraph:prov:chunk:{session_id}:{chunk_index}"
def docrag_synthesis_uri(session_id: str) -> str:
"""
Generate URI for a document RAG synthesis entity (final answer).

View file

@ -29,6 +29,8 @@ from . namespaces import (
TG_ANSWER_TYPE, TG_REFLECTION_TYPE, TG_THOUGHT_TYPE, TG_OBSERVATION_TYPE,
TG_DECOMPOSITION, TG_FINDING, TG_PLAN_TYPE, TG_STEP_RESULT,
TG_SUBAGENT_GOAL, TG_PLAN_STEP,
TG_EDGE_SELECTION, TG_SCORE,
TG_CHUNK_SELECTION,
)
@ -93,6 +95,8 @@ TG_CLASS_LABELS = [
_label_triple(TG_FINDING, "Finding"),
_label_triple(TG_PLAN_TYPE, "Plan"),
_label_triple(TG_STEP_RESULT, "Step Result"),
_label_triple(TG_EDGE_SELECTION, "Edge Selection"),
_label_triple(TG_CHUNK_SELECTION, "Chunk Selection"),
]
# TrustGraph predicate labels
@ -117,6 +121,7 @@ TG_PREDICATE_LABELS = [
_label_triple(TG_ENTITY, "entity"),
_label_triple(TG_SUBAGENT_GOAL, "subagent goal"),
_label_triple(TG_PLAN_STEP, "plan step"),
_label_triple(TG_SCORE, "score"),
]

View file

@ -16,3 +16,5 @@ from .collection import *
from .storage import *
from .tool_service import *
from .sparql_query import *
from .reranker import *
from .audit import *

View file

@ -0,0 +1,27 @@
from dataclasses import dataclass, field
from ..core.topic import queue
############################################################################
# Audit events — see docs/tech-specs/audit-events.md for the full spec.
#
# Transport: notify-class pub/sub (fire-and-forget, per-subscriber
# delivery). Producers are the API gateway and the IAM service.
# Consumers are optional enterprise components.
@dataclass
class AuditEvent:
schema_version: int = 1
event_id: str = ""
event_type: str = ""
timestamp: str = ""
producer: str = ""
payload_json: str = ""
audit_events_queue = queue('audit-events', cls='notify')
############################################################################

View file

@ -29,7 +29,7 @@ class UserInput:
@dataclass
class UserRecord:
id: str = ""
workspace: str = ""
default_workspace: str = ""
username: str = ""
name: str = ""
email: str = ""
@ -74,6 +74,21 @@ class ApiKeyRecord:
last_used: str = ""
# ---- Enterprise IAM types (additive) ----
@dataclass
class GroupInput:
name: str = ""
description: str = ""
enabled: bool = True
@dataclass
class GrantInput:
capability: str = ""
workspace: str = ""
@dataclass
class IamRequest:
operation: str = ""
@ -99,6 +114,17 @@ class IamRequest:
workspace_record: WorkspaceInput | None = None
key: ApiKeyInput | None = None
# ---- Enterprise IAM inputs (additive) ----
group_id: str = ""
member_type: str = ""
member_id: str = ""
group: GroupInput | None = None
grant: GrantInput | None = None
# ---- Audit context (informational, echoed into audit events) ----
request_id: str = ""
client_ip: str = ""
# ---- authorise / authorise-many inputs ----
# Capability string from the vocabulary in capabilities.md.
capability: str = ""
@ -138,7 +164,7 @@ class IamResponse:
# resolve-api-key
resolved_user_id: str = ""
resolved_workspace: str = ""
resolved_default_workspace: str = ""
resolved_roles: list[str] = field(default_factory=list)
# reset-password
@ -164,6 +190,14 @@ class IamResponse:
# authorise_checks.
decisions_json: str = ""
# ---- Enterprise IAM outputs (additive) ----
# JSON-serialised payloads for enterprise group/grant operations.
group_json: str = ""
groups_json: str = ""
members_json: str = ""
grants_json: str = ""
effective_permissions_json: str = ""
error: Error | None = None

View file

@ -6,17 +6,6 @@ from ..core.primitives import Error
# Prompt services, abstract the prompt generation
# extract-definitions:
# chunk -> definitions
# extract-relationships:
# chunk -> relationships
# kg-prompt:
# query, triples -> answer
# document-prompt:
# query, documents -> answer
# extract-rows
# schema, chunk -> rows
@dataclass
class PromptRequest:
id: str = ""

View file

@ -71,6 +71,27 @@ document_embeddings_response_queue = queue('document-embeddings', cls='response'
############################################################################
# Keyword index query - lexical (BM25) search over chunk text, the sparse
# counterpart to the doc embeddings query above. Matches share the ChunkMatch
# shape so both retrieval paths key on chunk_id; score is "higher is better"
# in both (BM25 rank scores are negated by the service to match).
@dataclass
class KeywordIndexRequest:
query: str = ""
limit: int = 0
collection: str = ""
@dataclass
class KeywordIndexResponse:
error: Error | None = None
chunks: list[ChunkMatch] = field(default_factory=list)
keyword_index_request_queue = queue('keyword-index', cls='request')
keyword_index_response_queue = queue('keyword-index', cls='response')
############################################################################
# Row embeddings query - for semantic/fuzzy matching on row index values
@dataclass

View file

@ -0,0 +1,35 @@
from dataclasses import dataclass, field
from ..core.primitives import Error
############################################################################
# Cross-encoder reranker
@dataclass
class RerankerQuery:
query_id: str = ""
query_text: str = ""
@dataclass
class RerankerDocument:
document_id: str = ""
document_text: str = ""
@dataclass
class RerankerRequest:
queries: list[RerankerQuery] = field(default_factory=list)
documents: list[RerankerDocument] = field(default_factory=list)
limit: int = 10
@dataclass
class RerankerResult:
document_id: str = ""
query_id: str = ""
score: float = 0.0
@dataclass
class RerankerResponse:
error: Error | None = None
results: list[RerankerResult] = field(default_factory=list)

View file

@ -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 = ""
@ -40,7 +41,10 @@ class GraphRagResponse:
class DocumentRagQuery:
query: str = ""
collection: str = ""
doc_limit: int = 0
doc_limit: int = 0 # docs selected into the synthesis prompt
fetch_limit: int = 0 # candidate pool fetched from the vector store
# before reranking (0 = derive from doc_limit;
# values below doc_limit are raised to it)
streaming: bool = False
@dataclass

View file

@ -10,7 +10,7 @@ description = "TrustGraph provides a means to run a pipeline of flexible AI proc
readme = "README.md"
requires-python = ">=3.8"
dependencies = [
"trustgraph-base>=2.5,<2.6",
"trustgraph-base>=2.7,<2.8",
"pulsar-client",
"prometheus-client",
"boto3",

View file

@ -10,7 +10,7 @@ description = "TrustGraph provides a means to run a pipeline of flexible AI proc
readme = "README.md"
requires-python = ">=3.8"
dependencies = [
"trustgraph-base>=2.5,<2.6",
"trustgraph-base>=2.7,<2.8",
"requests",
"pulsar-client",
"aiohttp",
@ -71,6 +71,7 @@ tg-invoke-rows-query = "trustgraph.cli.invoke_rows_query:main"
tg-invoke-sparql-query = "trustgraph.cli.invoke_sparql_query:main"
tg-invoke-row-embeddings = "trustgraph.cli.invoke_row_embeddings:main"
tg-invoke-prompt = "trustgraph.cli.invoke_prompt:main"
tg-invoke-reranker = "trustgraph.cli.invoke_reranker:main"
tg-invoke-structured-query = "trustgraph.cli.invoke_structured_query:main"
tg-load-doc-embeds = "trustgraph.cli.load_doc_embeds:main"
tg-load-kg-core = "trustgraph.cli.load_kg_core:main"
@ -115,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"

View file

@ -53,7 +53,7 @@ def main():
help="Auth token (default: $TRUSTGRAPH_TOKEN)",
)
parser.add_argument(
"--username", required=True, help="Username (unique in workspace)",
"--username", required=True, help="Username (globally unique)",
)
parser.add_argument(
"--password", default=None,
@ -75,10 +75,7 @@ def main():
)
parser.add_argument(
"-w", "--workspace", default=None,
help=(
"Target workspace (admin only; defaults to caller's "
"assigned workspace)"
),
help="Default workspace for the new user",
)
run_main(do_create_user, parser)

View file

@ -0,0 +1,301 @@
"""
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 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
import io
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")
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)
info.size = len(data)
info.mtime = int(time.time())
tar.addfile(info, io.BytesIO(data))
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/<c>/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"),
)
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():
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(
'-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:
export_workspace(
url=args.api_url,
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:
print("Exception:", e, flush=True)
sys.exit(1)
if __name__ == "__main__":
main()

View file

@ -0,0 +1,420 @@
"""
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.
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
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, Triple
from trustgraph.cli.nquads import parse_nquads
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, 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
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():
if not member.isfile():
continue
f = tar.extractfile(member)
if f is None:
continue
data = f.read()
name = member.name
if name == "manifest.json":
manifest = 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")
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"
)
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_config(api, entries, overwrite, dry_run):
"""Import config entries; returns (imported, skipped) counts."""
config = api.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(config.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)
elif values:
config.put(values)
return len(values), skipped
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():
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(
'-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 config 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(
'--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',
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,
flow_id=args.flow_id,
process=args.process,
process_collection=args.process_collection,
dry_run=args.dry_run,
token=args.token,
)
except Exception as e:
print("Exception:", e, flush=True)
sys.exit(1)
if __name__ == "__main__":
main()

View file

@ -21,10 +21,12 @@ default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
default_workspace = os.getenv("TRUSTGRAPH_WORKSPACE", "default")
default_collection = 'default'
default_doc_limit = 10
default_fetch_limit = 0
def question_explainable(
url, flow_id, question_text, collection, doc_limit, token=None, debug=False,
url, flow_id, question_text, collection, doc_limit, fetch_limit=0,
token=None, debug=False,
workspace="default",
):
"""Execute document RAG with explainability - shows provenance events inline."""
@ -39,6 +41,7 @@ def question_explainable(
query=question_text,
collection=collection,
doc_limit=doc_limit,
fetch_limit=fetch_limit,
):
if isinstance(item, RAGChunk):
# Print response content
@ -97,7 +100,7 @@ def question_explainable(
def question(
url, flow_id, question_text, collection, doc_limit,
url, flow_id, question_text, collection, doc_limit, fetch_limit=0,
streaming=True, token=None, explainable=False, debug=False,
show_usage=False, workspace="default",
):
@ -109,6 +112,7 @@ def question(
question_text=question_text,
collection=collection,
doc_limit=doc_limit,
fetch_limit=fetch_limit,
token=token,
debug=debug,
workspace=workspace,
@ -128,6 +132,7 @@ def question(
query=question_text,
collection=collection,
doc_limit=doc_limit,
fetch_limit=fetch_limit,
streaming=True
)
@ -155,6 +160,7 @@ def question(
query=question_text,
collection=collection,
doc_limit=doc_limit,
fetch_limit=fetch_limit,
)
print(result.text)
@ -214,7 +220,15 @@ def main():
'-d', '--doc-limit',
type=int,
default=default_doc_limit,
help=f'Document limit (default: {default_doc_limit})'
help=f'Documents selected into the prompt (default: {default_doc_limit})'
)
parser.add_argument(
'--fetch-limit',
type=int,
default=default_fetch_limit,
help='Candidate documents fetched from the vector store before '
'reranking (default: derive from doc-limit)'
)
parser.add_argument(
@ -251,6 +265,7 @@ def main():
question_text=args.question,
collection=args.collection,
doc_limit=args.doc_limit,
fetch_limit=args.fetch_limit,
streaming=not args.no_streaming,
token=args.token,
explainable=args.explainable,

View file

@ -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
@ -112,14 +115,13 @@ def _question_explainable_api(
if focus_full and focus_full.edge_selections:
for edge_sel in focus_full.edge_selections:
if edge_sel.edge:
# Resolve labels for edge components
s_label, p_label, o_label = explain_client.resolve_edge_labels(
edge_sel.edge, collection
)
print(f" Edge: ({s_label}, {p_label}, {o_label})", file=sys.stderr)
if edge_sel.reasoning:
r_short = edge_sel.reasoning[:100] + "..." if len(edge_sel.reasoning) > 100 else edge_sel.reasoning
print(f" Reason: {r_short}", file=sys.stderr)
if edge_sel.concept or edge_sel.score is not None:
score_str = f"{edge_sel.score:.4f}" if edge_sel.score is not None else "?"
print(f" Concept: {edge_sel.concept} Score: {score_str}", file=sys.stderr)
elif isinstance(entity, Synthesis):
print(f"\n [synthesis] {prov_id}", file=sys.stderr)
@ -139,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",
):
@ -157,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,
@ -181,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
)
@ -213,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)
@ -309,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',
@ -348,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,

View file

@ -0,0 +1,127 @@
"""
Invokes the reranker service to score and rank documents by relevance
to one or more queries.
"""
import argparse
import json
import os
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")
def query(url, flow_id, queries, documents, limit, token=None,
workspace="default"):
api = Api(url=url, token=token, workspace=workspace)
socket = api.socket()
flow = socket.flow(flow_id)
try:
query_objects = [
{"query_id": str(i), "query_text": q}
for i, q in enumerate(queries)
]
document_objects = [
{"document_id": str(i), "document_text": d}
for i, d in enumerate(documents)
]
result = flow.rerank(
queries=query_objects,
documents=document_objects,
limit=limit,
)
if "error" in result and result["error"]:
err = result["error"]
print(f"Error: [{err.get('type', '')}] {err.get('message', '')}")
return
for r in result.get("results", []):
doc_idx = int(r["document_id"])
query_idx = int(r["query_id"])
print(
f" {r['score']:.4f} | "
f"query: {queries[query_idx]} | "
f"doc: {documents[doc_idx]}"
)
finally:
socket.close()
def main():
parser = argparse.ArgumentParser(
prog='tg-invoke-reranker',
description=__doc__,
)
parser.add_argument(
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
parser.add_argument(
'-w', '--workspace',
default=default_workspace,
help=f'Workspace (default: {default_workspace})',
)
parser.add_argument(
'-f', '--flow-id',
default="default",
help=f'Flow ID (default: default)'
)
parser.add_argument(
'-l', '--limit',
type=int,
default=10,
help='Maximum number of results (default: 10)',
)
parser.add_argument(
'-q', '--query',
action='append',
required=True,
help='Query text (can be specified multiple times)',
)
parser.add_argument(
'documents',
nargs='+',
help='Documents to rerank',
)
args = parser.parse_args()
try:
query(
url=args.url,
flow_id=args.flow_id,
queries=args.query,
documents=args.documents,
limit=args.limit,
token=args.token,
workspace=args.workspace,
)
except Exception as e:
print("Exception:", e, flush=True)
if __name__ == "__main__":
main()

View file

@ -78,7 +78,7 @@ def load_structured_data(
logger.info("Step 1: Analyzing data to discover best matching schema...")
# Step 1: Auto-discover schema (reuse discover_schema logic)
discovered_schema = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, workspace=workspace)
discovered_schema = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, token=token, workspace=workspace)
if not discovered_schema:
logger.error("Failed to discover suitable schema automatically")
print("❌ Could not automatically determine the best schema for your data.")
@ -90,7 +90,7 @@ def load_structured_data(
# Step 2: Auto-generate descriptor
logger.info("Step 2: Generating descriptor configuration...")
auto_descriptor = _auto_generate_descriptor(api_url, input_file, discovered_schema, sample_chars, flow, logger, workspace=workspace)
auto_descriptor = _auto_generate_descriptor(api_url, input_file, discovered_schema, sample_chars, flow, logger, token=token, workspace=workspace)
if not auto_descriptor:
logger.error("Failed to generate descriptor automatically")
print("❌ Could not automatically generate descriptor configuration.")
@ -172,7 +172,7 @@ def load_structured_data(
logger.info(f"Sample chars: {sample_chars} characters")
# Use the helper function to discover schema (get raw response for display)
response = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, return_raw_response=True, workspace=workspace)
response = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, return_raw_response=True, token=token, workspace=workspace)
if response:
# Debug: print response type and content
@ -203,7 +203,7 @@ def load_structured_data(
# If no schema specified, discover it first
if not schema_name:
logger.info("No schema specified, auto-discovering...")
schema_name = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, workspace=workspace)
schema_name = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, token=token, workspace=workspace)
if not schema_name:
print("Error: Could not determine schema automatically.")
print("Please specify a schema using --schema-name or run --discover-schema first.")
@ -213,7 +213,7 @@ def load_structured_data(
logger.info(f"Target schema: {schema_name}")
# Generate descriptor using helper function
descriptor = _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, flow, logger, workspace=workspace)
descriptor = _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, flow, logger, token=token, workspace=workspace)
if descriptor:
# Output the generated descriptor
@ -603,7 +603,7 @@ def _send_to_trustgraph(rows, api_url, flow, batch_size=1000, token=None, worksp
# Helper functions for auto mode
def _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, return_raw_response=False, workspace="default"):
def _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, return_raw_response=False, token=None, workspace="default"):
"""Auto-discover the best matching schema for the input data
Args:
@ -626,7 +626,7 @@ def _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, retur
# Import API modules
from trustgraph.api import Api
from trustgraph.api.types import ConfigKey
api = Api(api_url, workspace=workspace)
api = Api(api_url, token=token, workspace=workspace)
config_api = api.config()
# Get available schemas
@ -707,7 +707,7 @@ def _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, retur
return None
def _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, flow, logger, workspace="default"):
def _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, flow, logger, token=None, workspace="default"):
"""Auto-generate descriptor configuration for the discovered schema"""
try:
# Read sample data
@ -717,7 +717,7 @@ def _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, fl
# Import API modules
from trustgraph.api import Api
from trustgraph.api.types import ConfigKey
api = Api(api_url, workspace=workspace)
api = Api(api_url, token=token, workspace=workspace)
config_api = api.config()
# Get schema definition

View file

@ -51,8 +51,8 @@ def main():
parser.add_argument(
"-w", "--workspace", default=None,
help=(
"Optional workspace to log in against. Defaults to "
"the user's assigned workspace."
"Override the default workspace for this session's JWT. "
"If omitted, uses the user's stored default workspace."
),
)
run_main(do_login, parser)

View file

@ -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):
"""<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 <graph-iri> 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))
]

View file

@ -203,9 +203,9 @@ def print_graphrag_text(trace, explain_client, flow, collection, api=None, show_
)
print(f" {i}. ({s_label}, {p_label}, {o_label})")
if edge_sel.reasoning:
r_short = edge_sel.reasoning[:100] + "..." if len(edge_sel.reasoning) > 100 else edge_sel.reasoning
print(f" Reasoning: {r_short}")
if edge_sel.concept or edge_sel.score is not None:
score_str = f"{edge_sel.score:.4f}" if edge_sel.score is not None else "?"
print(f" Concept: {edge_sel.concept} Score: {score_str}")
if show_provenance and edge_sel.edge:
provenance = trace_edge_provenance(
@ -519,7 +519,8 @@ def trace_to_dict(trace, trace_type):
"selected_edges": [
{
"edge": edge_sel.edge,
"reasoning": edge_sel.reasoning,
"concept": edge_sel.concept,
"score": edge_sel.score,
}
for edge_sel in focus.edge_selections
],

View file

@ -68,7 +68,7 @@ def do_update_user(args):
print(f"username : {rec.get('username', '')}")
print(f"name : {rec.get('name', '')}")
print(f"email : {rec.get('email', '')}")
print(f"workspace : {rec.get('workspace', '')}")
print(f"default_ws: {rec.get('default_workspace', '')}")
print(f"roles : {', '.join(rec.get('roles', []))}")
print(f"enabled : {'yes' if rec.get('enabled') else 'no'}")
print(
@ -114,7 +114,7 @@ def main():
"-w", "--workspace", default=None,
help=(
"Optional workspace integrity check — when supplied, "
"iam-svc verifies the target user's home workspace "
"iam-svc verifies the target user's default workspace "
"matches"
),
)

View file

@ -10,8 +10,8 @@ description = "HuggingFace embeddings support for TrustGraph."
readme = "README.md"
requires-python = ">=3.8"
dependencies = [
"trustgraph-base>=2.5,<2.6",
"trustgraph-flow>=2.5,<2.6",
"trustgraph-base>=2.7,<2.8",
"trustgraph-flow>=2.7,<2.8",
"torch",
"urllib3",
"transformers",

View file

@ -10,7 +10,7 @@ description = "TrustGraph provides a means to run a pipeline of flexible AI proc
readme = "README.md"
requires-python = ">=3.8"
dependencies = [
"trustgraph-base>=2.5,<2.6",
"trustgraph-base>=2.7,<2.8",
"aiohttp",
"anthropic",
"scylla-driver",
@ -19,6 +19,7 @@ dependencies = [
"faiss-cpu",
"falkordb",
"fastembed",
"flashrank",
"ibis",
"jsonschema",
"langchain",
@ -75,6 +76,7 @@ document-embeddings = "trustgraph.embeddings.document_embeddings:run"
document-rag = "trustgraph.retrieval.document_rag:run"
embeddings-fastembed = "trustgraph.embeddings.fastembed:run"
embeddings-ollama = "trustgraph.embeddings.ollama:run"
kw-index-fts5 = "trustgraph.storage.kw_index.fts5:run"
graph-embeddings-query-milvus = "trustgraph.query.graph_embeddings.milvus:run"
graph-embeddings-query-pinecone = "trustgraph.query.graph_embeddings.pinecone:run"
graph-embeddings-query-qdrant = "trustgraph.query.graph_embeddings.qdrant:run"
@ -83,6 +85,7 @@ graph-embeddings-write-pinecone = "trustgraph.storage.graph_embeddings.pinecone:
graph-embeddings-write-qdrant = "trustgraph.storage.graph_embeddings.qdrant:run"
graph-embeddings = "trustgraph.embeddings.graph_embeddings:run"
graph-rag = "trustgraph.retrieval.graph_rag:run"
reranker-flashrank = "trustgraph.reranker.flashrank:run"
kg-extract-agent = "trustgraph.extract.kg.agent:run"
kg-extract-definitions = "trustgraph.extract.kg.definitions:run"
kg-extract-rows = "trustgraph.extract.kg.rows:run"

View file

@ -18,6 +18,10 @@ description : str (default "Default")
Human-readable description passed to flow-svc.
parameters : dict (optional)
Optional parameter overrides passed to start-flow.
list_timeout : int (default 10)
Timeout in seconds for the list-flows request.
start_timeout : int (default 30)
Timeout in seconds for the start-flow request.
"""
from trustgraph.schema import FlowRequest
@ -34,6 +38,8 @@ class DefaultFlowStart(Initialiser):
blueprint=None,
description="Default",
parameters=None,
list_timeout=10,
start_timeout=30,
**kwargs,
):
super().__init__(**kwargs)
@ -46,6 +52,8 @@ class DefaultFlowStart(Initialiser):
self.blueprint = blueprint
self.description = description
self.parameters = dict(parameters) if parameters else {}
self.list_timeout = list_timeout
self.start_timeout = start_timeout
async def run(self, ctx, old_flag, new_flag):
@ -70,7 +78,7 @@ class DefaultFlowStart(Initialiser):
FlowRequest(
operation="list-flows",
),
timeout=10,
timeout=self.list_timeout,
)
if list_resp.error:
raise RuntimeError(
@ -99,7 +107,7 @@ class DefaultFlowStart(Initialiser):
description=self.description,
parameters=self.parameters,
),
timeout=30,
timeout=self.start_timeout,
)
if resp.error:
raise RuntimeError(

View file

@ -14,7 +14,9 @@ seed_file : str (required when source=="seed-file")
Path to a JSON seed file with the same shape TemplateSeed consumes.
overwrite : bool (default False)
On re-run (flag change), if True overwrite all keys; if False,
upsert-missing-only (preserves in-workspace customisations).
upsert-missing-only (preserves in-workspace customisations)
iam_timeout : int (default 10)
Timeout in seconds for the IAM create-workspace request.
Raises (in ``run``)
-------------------
@ -41,7 +43,9 @@ class WorkspaceInit(Initialiser):
source="template",
seed_file=None,
overwrite=False,
iam_timeout=10,
**kwargs,
):
super().__init__(**kwargs)
@ -59,6 +63,7 @@ class WorkspaceInit(Initialiser):
self.source = source
self.seed_file = seed_file
self.overwrite = overwrite
self.iam_timeout = iam_timeout
async def run(self, ctx, old_flag, new_flag):
await self._create_workspace(ctx)
@ -123,7 +128,7 @@ class WorkspaceInit(Initialiser):
enabled=True,
),
),
timeout=10,
timeout=self.iam_timeout,
)
if resp.error:
if resp.error.type == "duplicate":

View file

@ -83,7 +83,8 @@ class Processor(AsyncProcessor):
host=cassandra_host,
username=cassandra_username,
password=cassandra_password,
default_keyspace="config"
default_keyspace="config",
replication_factor=params.get("cassandra_replication_factor"),
)
# Store resolved configuration

View file

@ -61,7 +61,8 @@ class Processor(WorkspaceProcessor):
host=cassandra_host,
username=cassandra_username,
password=cassandra_password,
default_keyspace="knowledge"
default_keyspace="knowledge",
replication_factor=params.get("cassandra_replication_factor"),
)
self.cassandra_host = hosts

Some files were not shown because too many files have changed in this diff Show more