mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-08 21:02:12 +02:00
Compare commits
63 commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0374098ee9 | ||
|
|
f76f2abe08 | ||
|
|
2bdc930b2a | ||
|
|
e5206bddd0 | ||
|
|
5a1b42ab33 | ||
|
|
a84cffa485 | ||
|
|
0b6bae02f5 | ||
|
|
d5f3b6d9f6 | ||
|
|
08f69007c9 | ||
|
|
cfbd5b9079 | ||
|
|
12ea6d46bc | ||
|
|
4ac6bc0986 | ||
|
|
8d630ef204 | ||
|
|
d2a4abb7c5 | ||
|
|
68e816e65c | ||
|
|
76c4763b9b | ||
|
|
508d0bb5c1 | ||
|
|
c05296376e | ||
|
|
f04ae5331d | ||
|
|
db7fdbc652 | ||
|
|
4aaa1ce915 | ||
|
|
9cf7dcb578 | ||
|
|
6c9a545a06 | ||
|
|
f18d48dc39 | ||
|
|
6887076ce0 | ||
|
|
55e2a2a3ce | ||
|
|
11ca7c89c4 | ||
|
|
656ca430b9 | ||
|
|
f20b50cfb2 | ||
|
|
04c5921687 | ||
|
|
01cc8dbc64 | ||
|
|
1aa9549912 | ||
|
|
5cb4f83afa | ||
|
|
0a828379be | ||
|
|
16f8cfd972 | ||
|
|
a3df4f62bb | ||
|
|
09b8a1d347 | ||
|
|
fa264ded46 | ||
|
|
cae931409a | ||
|
|
6b0475e315 | ||
|
|
cb0ad1a450 | ||
|
|
fc0ecc770a | ||
|
|
345da375b1 | ||
|
|
0ba1eeeda0 | ||
|
|
eb1e38d7d0 | ||
|
|
b8770a6005 | ||
|
|
28802a644a | ||
|
|
8797d9d9ff | ||
|
|
627c669097 | ||
|
|
8b0619e5d8 | ||
|
|
e3f9f8c357 | ||
|
|
81d57826c8 | ||
|
|
79d7ef6a90 | ||
|
|
28a51c244f | ||
|
|
fa5ebe2393 | ||
|
|
e1c9351454 | ||
|
|
dbc21c0bb9 | ||
|
|
08bfec1539 | ||
|
|
4913f8c2eb | ||
|
|
97453d9b83 | ||
|
|
6dfa47aac8 | ||
|
|
dcee842455 | ||
|
|
36eadbda3a |
153 changed files with 12718 additions and 2507 deletions
2
.github/workflows/pull-request.yaml
vendored
2
.github/workflows/pull-request.yaml
vendored
|
|
@ -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
218
README.dev-install.md
Normal 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
293
README.md
|
|
@ -3,52 +3,104 @@
|
|||
|
||||
<img src="TG-fullname-logo.svg" width=100% />
|
||||
|
||||
[](https://pypi.org/project/trustgraph/) [](LICENSE) 
|
||||
[](https://pypi.org/project/trustgraph/)  
|
||||
[](https://discord.gg/sQMwkRz5GX) [](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?
|
||||
|
||||
[](https://www.youtube.com/watch?v=gZjlt5WcWB4)
|
||||
|
||||
## Watch Context Graphs in Action
|
||||
## Watch Building Real Agents from a Context Graph
|
||||
|
||||
[](https://www.youtube.com/watch?v=sWc7mkhITIo)
|
||||
[](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.
|
||||
|
||||
### What’s 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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
402
docs/tech-specs/audit-events.md
Normal file
402
docs/tech-specs/audit-events.md
Normal 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).
|
||||
|
|
@ -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` |
|
||||
|
||||
|
|
|
|||
564
docs/tech-specs/graph-rag-semantic-filter.md
Normal file
564
docs/tech-specs/graph-rag-semantic-filter.md
Normal 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 1–10 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 1–10 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
|
||||
query–document 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 query–document
|
||||
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 1–10 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.0–1.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
2603
install_trustgraph.sh
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 gateway→IAM audit
|
||||
chain."""
|
||||
names = _field_names(IamRequest)
|
||||
assert "request_id" in names
|
||||
assert "client_ip" in names
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
90
tests/unit/test_base/test_audit_publisher.py
Normal file
90
tests/unit/test_base/test_audit_publisher.py
Normal 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
|
||||
|
|
@ -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
|
||||
81
tests/unit/test_base/test_optional_request_response_spec.py
Normal file
81
tests/unit/test_base/test_optional_request_response_spec.py
Normal 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": {}},
|
||||
)
|
||||
|
|
@ -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"""
|
||||
|
|
|
|||
136
tests/unit/test_base/test_qdrant_config.py
Normal file
136
tests/unit/test_base/test_qdrant_config.py
Normal 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
|
||||
54
tests/unit/test_bootstrap/test_default_flow_start.py
Normal file
54
tests/unit/test_bootstrap/test_default_flow_start.py
Normal 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
|
||||
13
tests/unit/test_bootstrap/test_workspace_init.py
Normal file
13
tests/unit/test_bootstrap/test_workspace_init.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
97
tests/unit/test_cli/test_nquads.py
Normal file
97
tests/unit/test_cli/test_nquads.py
Normal 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
|
||||
449
tests/unit/test_cli/test_workspace_bundle_commands.py
Normal file
449
tests/unit/test_cli/test_workspace_bundle_commands.py
Normal 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()
|
||||
|
|
@ -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:
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
198
tests/unit/test_gateway/test_audit_middleware.py
Normal file
198
tests/unit/test_gateway/test_audit_middleware.py
Normal 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
|
||||
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
|
|
|||
345
tests/unit/test_iam/test_iam_audit_events.py
Normal file
345
tests/unit/test_iam/test_iam_audit_events.py
Normal 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()
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
211
tests/unit/test_retrieval/test_document_rag_hybrid.py
Normal file
211
tests/unit/test_retrieval/test_document_rag_hybrid.py
Normal 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,
|
||||
)
|
||||
550
tests/unit/test_retrieval/test_document_rag_rerank.py
Normal file
550
tests/unit/test_retrieval/test_document_rag_rerank.py
Normal 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]
|
||||
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
395
tests/unit/test_retrieval/test_graph_rag_direction_aware_text.py
Normal file
395
tests/unit/test_retrieval/test_graph_rag_direction_aware_text.py
Normal 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]
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
157
tests/unit/test_storage/test_kw_index_fts5_storage.py
Normal file
157
tests/unit/test_storage/test_kw_index_fts5_storage.py
Normal 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") == []
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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", ""),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
42
trustgraph-base/trustgraph/base/audit_publisher.py
Normal file
42
trustgraph-base/trustgraph/base/audit_publisher.py
Normal 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}")
|
||||
|
|
@ -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']
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
44
trustgraph-base/trustgraph/base/keyword_index_client.py
Normal file
44
trustgraph-base/trustgraph/base/keyword_index_client.py
Normal 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,
|
||||
)
|
||||
132
trustgraph-base/trustgraph/base/keyword_index_service.py
Normal file
132
trustgraph-base/trustgraph/base/keyword_index_service.py
Normal 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})'
|
||||
)
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
87
trustgraph-base/trustgraph/base/qdrant_config.py
Normal file
87
trustgraph-base/trustgraph/base/qdrant_config.py
Normal 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
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
43
trustgraph-base/trustgraph/base/reranker_client.py
Normal file
43
trustgraph-base/trustgraph/base/reranker_client.py
Normal 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,
|
||||
)
|
||||
109
trustgraph-base/trustgraph/base/reranker_service.py
Normal file
109
trustgraph-base/trustgraph/base/reranker_service.py
Normal 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)
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -20,3 +20,4 @@ from .embeddings_query import (
|
|||
)
|
||||
from .rows_query import RowsQueryRequestTranslator, RowsQueryResponseTranslator
|
||||
from .diagnosis import StructuredDataDiagnosisRequestTranslator, StructuredDataDiagnosisResponseTranslator
|
||||
from .reranker import RerankerRequestTranslator, RerankerResponseTranslator
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
73
trustgraph-base/trustgraph/messaging/translators/reranker.py
Normal file
73
trustgraph-base/trustgraph/messaging/translators/reranker.py
Normal 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
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -16,3 +16,5 @@ from .collection import *
|
|||
from .storage import *
|
||||
from .tool_service import *
|
||||
from .sparql_query import *
|
||||
from .reranker import *
|
||||
from .audit import *
|
||||
27
trustgraph-base/trustgraph/schema/services/audit.py
Normal file
27
trustgraph-base/trustgraph/schema/services/audit.py
Normal 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')
|
||||
|
||||
############################################################################
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = ""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
35
trustgraph-base/trustgraph/schema/services/reranker.py
Normal file
35
trustgraph-base/trustgraph/schema/services/reranker.py
Normal 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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
301
trustgraph-cli/trustgraph/cli/export_workspace.py
Normal file
301
trustgraph-cli/trustgraph/cli/export_workspace.py
Normal 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()
|
||||
420
trustgraph-cli/trustgraph/cli/import_workspace.py
Normal file
420
trustgraph-cli/trustgraph/cli/import_workspace.py
Normal 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()
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
127
trustgraph-cli/trustgraph/cli/invoke_reranker.py
Normal file
127
trustgraph-cli/trustgraph/cli/invoke_reranker.py
Normal 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()
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
137
trustgraph-cli/trustgraph/cli/nquads.py
Normal file
137
trustgraph-cli/trustgraph/cli/nquads.py
Normal 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))
|
||||
]
|
||||
|
|
@ -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
|
||||
],
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Add a link
Reference in a new issue