mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-14 07:42:11 +02:00
Compare commits
63 commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b1d1833234 | ||
|
|
40f01c123b | ||
|
|
9136526863 | ||
|
|
f106ae2103 | ||
|
|
e9c6a850ad | ||
|
|
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 | ||
|
|
28a51c244f | ||
|
|
fa5ebe2393 | ||
|
|
97453d9b83 | ||
|
|
6dfa47aac8 | ||
|
|
dcee842455 | ||
|
|
36eadbda3a |
192 changed files with 15731 additions and 2410 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
|
||||
|
|
|
|||
3
Makefile
3
Makefile
|
|
@ -57,7 +57,8 @@ container-bedrock container-vertexai \
|
|||
container-hf container-ocr \
|
||||
container-unstructured container-mcp
|
||||
|
||||
some-containers: container-base container-flow container-unstructured
|
||||
some-containers: container-base container-flow
|
||||
# container-unstructured
|
||||
|
||||
push:
|
||||
${DOCKER} push ${CONTAINER_BASE}/trustgraph-base:${VERSION}
|
||||
|
|
|
|||
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.
|
||||
|
|
|
|||
|
|
@ -29,11 +29,14 @@ def main():
|
|||
("trustgraph", "Base trustgraph package"),
|
||||
("trustgraph.base", "Base classes"),
|
||||
("trustgraph.base.llm_service", "LLM service base class"),
|
||||
("trustgraph.base.image_to_text_service", "Image-to-text service base class"),
|
||||
("trustgraph.schema", "Schema definitions"),
|
||||
("trustgraph.exceptions", "Exception classes"),
|
||||
("trustgraph.model", "Model package"),
|
||||
("trustgraph.model.text_completion", "Text completion package"),
|
||||
("trustgraph.model.text_completion.vertexai", "VertexAI package"),
|
||||
("trustgraph.model.image_to_text", "Image-to-text package"),
|
||||
("trustgraph.model.image_to_text.openai", "Image-to-text OpenAI package"),
|
||||
]
|
||||
|
||||
success_count = 0
|
||||
|
|
|
|||
|
|
@ -111,7 +111,7 @@ Processors that talk to external LLMs or APIs read their credentials
|
|||
from env vars, same as in the per-container deployment:
|
||||
|
||||
- `OPENAI_TOKEN`, `OPENAI_BASE_URL` — for `text-completion` /
|
||||
`text-completion-rag`
|
||||
`text-completion-rag` / `image-to-text`
|
||||
|
||||
Export whatever your particular `group.yaml` needs before running.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
16
dev-tools/proc-group/groups/image-to-text.yaml
Normal file
16
dev-tools/proc-group/groups/image-to-text.yaml
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
# Image-to-text. Outbound vision-model calls. Isolated for the same
|
||||
# reason as the LLM group: the upstream API is the most likely thing to
|
||||
# need restart (provider changes, model changes, API flakiness).
|
||||
|
||||
_defaults: &defaults
|
||||
pubsub_backend: rabbitmq
|
||||
rabbitmq_host: localhost
|
||||
log_level: INFO
|
||||
|
||||
processors:
|
||||
|
||||
- class: trustgraph.model.image_to_text.openai.Processor
|
||||
params:
|
||||
<<: *defaults
|
||||
id: image-to-text
|
||||
max_output: 4096
|
||||
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` |
|
||||
|
||||
|
|
|
|||
|
|
@ -79,13 +79,17 @@ Interfaces can take two forms:
|
|||
"embeddings": {
|
||||
"request": "non-persistent://tg/request/{workspace}:embeddings:{class}",
|
||||
"response": "non-persistent://tg/response/{workspace}:embeddings:{class}"
|
||||
},
|
||||
"image-to-text": {
|
||||
"request": "non-persistent://tg/request/{workspace}:image-to-text:{class}",
|
||||
"response": "non-persistent://tg/response/{workspace}:image-to-text:{class}"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Types of Interfaces:**
|
||||
- **Entry Points**: Where external systems inject data (`document-load`, `agent`)
|
||||
- **Service Interfaces**: Request/response patterns for services (`embeddings`, `text-completion`)
|
||||
- **Service Interfaces**: Request/response patterns for services (`embeddings`, `text-completion`, `image-to-text`)
|
||||
- **Data Interfaces**: Fire-and-forget data flow connection points (`triples-store`, `entity-contexts-load`)
|
||||
|
||||
### 4. Parameters Section
|
||||
|
|
|
|||
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.
|
||||
|
|
@ -167,6 +167,11 @@ Values are absent (not zero) when token counts are unavailable.
|
|||
## GraphRagResponse Schema
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class Source:
|
||||
uri: str = "" # Source document URI
|
||||
title: str = "" # Document title (empty when the document has none)
|
||||
|
||||
@dataclass
|
||||
class GraphRagResponse:
|
||||
error: Error | None = None
|
||||
|
|
@ -177,6 +182,10 @@ class GraphRagResponse:
|
|||
explain_triples: list[Triple] = field(default_factory=list)
|
||||
message_type: str = "" # "chunk" or "explain"
|
||||
end_of_session: bool = False
|
||||
in_token: int | None = None
|
||||
out_token: int | None = None
|
||||
model: str | None = None
|
||||
sources: list[Source] = field(default_factory=list)
|
||||
```
|
||||
|
||||
### Message Types
|
||||
|
|
@ -227,6 +236,15 @@ Selected edges can be traced back to source documents:
|
|||
2. Follow `prov:wasDerivedFrom` chain to root document
|
||||
3. Each step in chain: chunk → page → document
|
||||
|
||||
### Source References in the Response
|
||||
|
||||
GraphRAG performs this walk on every query to enrich the synthesis
|
||||
prompt with document metadata. The same walk also produces structured
|
||||
`sources` entries (`uri` plus `title` from `dc:title`/`rdfs:label`),
|
||||
deduplicated and sorted by URI, attached to the final response message
|
||||
(`end_of_session=True`) at no additional query cost. Clients can display
|
||||
citations without re-running the traversal against the knowledge graph.
|
||||
|
||||
### Cassandra Quoted Triple Support
|
||||
|
||||
The Cassandra query service supports matching quoted triples:
|
||||
|
|
|
|||
263
docs/tech-specs/structured-output.md
Normal file
263
docs/tech-specs/structured-output.md
Normal file
|
|
@ -0,0 +1,263 @@
|
|||
---
|
||||
layout: default
|
||||
title: "Structured Output: LLM-Native JSON Schema Enforcement"
|
||||
parent: "Tech Specs"
|
||||
---
|
||||
|
||||
# Structured Output: LLM-Native JSON Schema Enforcement
|
||||
|
||||
## Problem / Opportunity
|
||||
|
||||
TrustGraph's knowledge-graph pipeline relies on LLMs to produce
|
||||
structured JSON output — entity extractions, relationship triples,
|
||||
topic classifications, and other schema-governed artefacts. Today,
|
||||
the correct structure is requested via natural-language instructions
|
||||
embedded in the prompt template: the prompt describes the expected
|
||||
JSON shape, and the system parses the LLM's free-text response,
|
||||
hoping it conforms.
|
||||
|
||||
This approach has several weaknesses:
|
||||
|
||||
1. **Fragile parsing.** LLM responses may include markdown fencing,
|
||||
preamble text, trailing commentary, or minor schema violations
|
||||
(missing fields, wrong types, extra keys). Every consumer must
|
||||
tolerate or work around these deviations, adding defensive code
|
||||
and retry logic.
|
||||
|
||||
2. **Wasted tokens and latency.** A significant portion of each
|
||||
prompt is spent describing the output format in prose. When the
|
||||
model deviates, retries consume additional tokens and add
|
||||
end-to-end latency.
|
||||
|
||||
3. **Silent data-quality issues.** Malformed responses that pass
|
||||
lenient parsing can inject bad data into the knowledge graph —
|
||||
wrong types, truncated lists, invented field names — without
|
||||
raising errors.
|
||||
|
||||
4. **Untapped LLM capability.** Most modern LLMs (OpenAI, Google
|
||||
Gemini, Anthropic Claude, Ollama-hosted models via llama.cpp)
|
||||
support *structured output* or *guided decoding*: the caller
|
||||
supplies a JSON schema and the model constrains token selection at
|
||||
the logit level to guarantee schema-valid output. TrustGraph
|
||||
already defines the required JSON schemas inside its prompt
|
||||
definitions but does not pass them through to the LLM backend.
|
||||
|
||||
### Opportunity
|
||||
|
||||
By threading the existing JSON schemas from prompt definitions
|
||||
through the text-completion service to each LLM backend's native
|
||||
structured-output API, TrustGraph can:
|
||||
|
||||
- **Guarantee valid output** on every call — no parsing heuristics,
|
||||
no retries for format errors.
|
||||
- **Reduce prompt size** by removing prose format instructions that
|
||||
the schema makes redundant.
|
||||
- **Improve data quality** in the knowledge graph by eliminating an
|
||||
entire class of silent ingestion errors.
|
||||
- **Simplify service code** by removing defensive JSON extraction and
|
||||
validation logic from every consumer.
|
||||
|
||||
## Scope
|
||||
|
||||
Prompt definitions declare a `response-type` of `"text"`, `"json"`,
|
||||
or `"jsonl"`. Structured output applies only to prompts that produce
|
||||
machine-readable output (`"json"` and `"jsonl"`).
|
||||
|
||||
JSONL presents a compatibility challenge: LLM structured-output APIs
|
||||
enforce a single top-level JSON schema, but JSONL prompts ask the
|
||||
model to emit one JSON object per line — a format that is not itself
|
||||
valid JSON. Converting JSONL prompts to request a JSON array would
|
||||
conflict with the prompt text and sacrifice truncation resilience
|
||||
(partial JSONL is recoverable line-by-line; a truncated array is
|
||||
broken JSON).
|
||||
|
||||
This spec takes a three-phase approach:
|
||||
|
||||
- **Phase 1** — plumb schemas through to LLM backends with automatic
|
||||
compatibility detection; non-compliant schemas fall back to the
|
||||
current free-text path.
|
||||
- **Phase 2** — fix up non-compliant schemas so more prompts benefit.
|
||||
- **Phase 3** — address JSONL prompts.
|
||||
|
||||
---
|
||||
|
||||
## Phase 1 — Structured Output with Automatic Fallback
|
||||
|
||||
### Design
|
||||
|
||||
Phase 1 threads the JSON schema from the prompt definition through
|
||||
the text-completion service to the LLM backend's native
|
||||
structured-output API. Only prompts with `response-type: "json"` are
|
||||
candidates.
|
||||
|
||||
Not all existing schemas are compatible with LLM structured-output
|
||||
APIs. Rather than require schema changes up front, Phase 1 includes
|
||||
a **runtime compatibility check**: if a schema passes, structured
|
||||
output is used; if not, the prompt falls back to the current
|
||||
free-text path with post-hoc validation. This makes the feature
|
||||
safe to deploy immediately.
|
||||
|
||||
### Strict-Mode Schema Requirements
|
||||
|
||||
LLM providers impose constraints beyond standard JSON Schema
|
||||
validation. A schema is considered strict-mode compatible when:
|
||||
|
||||
- Every `object` has `additionalProperties: false`.
|
||||
- Every property defined in `properties` appears in `required`.
|
||||
Optional fields use a nullable type (e.g. `"type": ["string", "null"]`)
|
||||
instead of omitting the key from `required`.
|
||||
- No `minimum`, `maximum`, `minLength`, `maxLength`, or `pattern`
|
||||
constraints (unsupported by most providers' constrained decoders).
|
||||
- No open-ended objects (`{"type": "object"}` without `properties`).
|
||||
- A schema is present and non-null.
|
||||
|
||||
### Runtime Compatibility Check
|
||||
|
||||
`PromptManager` (or a shared utility) inspects each schema at load
|
||||
time against the strict-mode rules above. The result is a boolean
|
||||
flag per prompt: `structured_output_eligible`.
|
||||
|
||||
- **Eligible** — `response_format` and `schema` are set on the
|
||||
`TextCompletionRequest`; the LLM enforces the schema at generation
|
||||
time.
|
||||
- **Not eligible** — request is sent without schema fields; the
|
||||
current free-text parsing and `jsonschema.validate()` path is used.
|
||||
|
||||
This is a per-prompt decision, not a global switch.
|
||||
|
||||
### Text-Completion Request Changes
|
||||
|
||||
`TextCompletionRequest` gains two optional fields:
|
||||
|
||||
```
|
||||
TextCompletionRequest:
|
||||
system: str
|
||||
prompt: str
|
||||
streaming: bool
|
||||
response_format: str | None # "json" or None (default)
|
||||
schema: dict | None # JSON Schema object or None
|
||||
```
|
||||
|
||||
When `response_format` is `"json"` and `schema` is provided, the LLM
|
||||
backend MUST use its native structured-output mechanism. When either
|
||||
field is absent or null, behaviour is unchanged.
|
||||
|
||||
### LLM Backend Mapping
|
||||
|
||||
Each backend maps `response_format` + `schema` to its provider's
|
||||
native API:
|
||||
|
||||
| Backend | API mechanism |
|
||||
|------------|-------------------------------------------------------|
|
||||
| OpenAI | `response_format={"type": "json_schema", "json_schema": {"name": "...", "schema": ...}}` |
|
||||
| Claude | `tool_use` with a single tool whose `input_schema` is the target schema, or the `response_format` parameter when available |
|
||||
| Gemini | `response_mime_type="application/json"` + `response_schema=...` |
|
||||
| Ollama | `format="json"` + schema in the `format` field (llama.cpp guided decoding) |
|
||||
| Llamafile | `response_format={"type": "json_object"}` + schema |
|
||||
|
||||
Backends that do not support schema-level enforcement (e.g. older
|
||||
Ollama versions) fall back to `response_format=json` without a schema
|
||||
and rely on post-hoc validation.
|
||||
|
||||
### Current Prompt Compatibility
|
||||
|
||||
Of the nine `response-type: "json"` prompts, two are strict-mode
|
||||
compatible today:
|
||||
|
||||
| Prompt | Status | Issue |
|
||||
|--------------------------|-----------|------------------------------------|
|
||||
| `schema-selection` | Ready | — |
|
||||
| `supervisor-decompose` | Ready | — |
|
||||
| `plan-create` | Fixable | Optional fields not in `required` |
|
||||
| `graphql-generation` | Blocked | Open-ended `variables` object; `minimum`/`maximum` on `confidence` |
|
||||
| `plan-step-execute` | Blocked | Open-ended `arguments` object |
|
||||
| `diagnose-structured-data` | No schema | — |
|
||||
| `diagnose-xml` | No schema | — |
|
||||
| `diagnose-json` | No schema | — |
|
||||
| `diagnose-csv` | No schema | — |
|
||||
|
||||
### What Does Not Change
|
||||
|
||||
- Prompt templates and their text content.
|
||||
- The `"text"` and `"jsonl"` response-type paths.
|
||||
- The `TextCompletionResponse` schema.
|
||||
- Post-hoc validation (retained as a defence-in-depth measure).
|
||||
|
||||
---
|
||||
|
||||
## Phase 2 — Schema Remediation
|
||||
|
||||
Phase 2 expands structured-output coverage by fixing schemas that
|
||||
failed the Phase 1 compatibility check.
|
||||
|
||||
### Fixable Schemas
|
||||
|
||||
**`plan-create`** — `tool_hint` and `depends_on` are optional
|
||||
(present in `properties` but absent from `required`). Fix: add both
|
||||
to `required` and change their types to nullable:
|
||||
|
||||
```json
|
||||
"tool_hint": {"type": ["string", "null"]},
|
||||
"depends_on": {
|
||||
"type": ["array", "null"],
|
||||
"items": {"type": "integer"}
|
||||
}
|
||||
```
|
||||
|
||||
### Schemas Requiring Design Decisions
|
||||
|
||||
**`graphql-generation`** — Two issues:
|
||||
|
||||
- `variables` is an open-ended object (`"additionalProperties": true`)
|
||||
with no fixed properties. Constrained decoding cannot handle
|
||||
arbitrary keys. Options: remove `variables` from the schema and
|
||||
accept it as free-form text within a wrapper, or restructure as a
|
||||
JSON-encoded string field.
|
||||
- `confidence` uses `"minimum": 0.0, "maximum": 1.0`. Fix: remove
|
||||
the numeric bounds; accept any number and clamp in application code
|
||||
if needed.
|
||||
|
||||
**`plan-step-execute`** — `arguments` is an open-ended object with no
|
||||
fixed properties. Same constraint as `graphql-generation.variables`.
|
||||
|
||||
### Missing Schemas
|
||||
|
||||
The four `diagnose-*` prompts have `response-type: "json"` but no
|
||||
schema. Adding schemas for these prompts would bring them into
|
||||
structured-output scope. This requires defining the expected
|
||||
response shape for each diagnostic prompt.
|
||||
|
||||
---
|
||||
|
||||
## Phase 3 (Future) — Structured Output for JSONL Prompts
|
||||
|
||||
JSONL prompts ask the LLM to emit multiple JSON objects, one per
|
||||
line. Each object is validated individually against the prompt's
|
||||
schema. The current approach is tolerant of truncation and
|
||||
malformed lines — useful properties for large extractions.
|
||||
|
||||
### Options
|
||||
|
||||
**Option A — Array wrapper.** Change the prompt text to request a
|
||||
JSON array instead of JSONL. Wrap the schema as
|
||||
`{"type": "array", "items": <existing-schema>}` and use structured
|
||||
output. Trade-off: loses line-by-line truncation resilience; requires
|
||||
updating every JSONL prompt template.
|
||||
|
||||
**Option B — Structured output per chunk.** Split the input so each
|
||||
text-completion call produces a single JSON object, then aggregate.
|
||||
Trade-off: more LLM calls; higher latency and cost; may not suit
|
||||
prompts that extract variable-length lists from a single chunk.
|
||||
|
||||
**Option C — Hybrid.** Use structured output with the array-wrapped
|
||||
schema but retain the post-hoc JSONL parser as a fallback for
|
||||
backends that do not support structured output or when the response
|
||||
is truncated. Trade-off: two code paths to maintain.
|
||||
|
||||
**Option D — Status quo.** Leave JSONL prompts on the free-text path
|
||||
with post-hoc validation. Structured output for `"json"` prompts
|
||||
already covers the most schema-sensitive cases; JSONL extraction is
|
||||
inherently more tolerant of partial results.
|
||||
|
||||
Phase 3 design will be selected after earlier phases are deployed and
|
||||
real-world structured-output behaviour is observed across backends.
|
||||
2603
install_trustgraph.sh
Normal file
2603
install_trustgraph.sh
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -0,0 +1,27 @@
|
|||
type: object
|
||||
description: |
|
||||
Image-to-text request - describe an image using a vision model.
|
||||
|
||||
The image payload is base64-encoded; raw binary is not accepted.
|
||||
required:
|
||||
- image
|
||||
- mime_type
|
||||
properties:
|
||||
image:
|
||||
type: string
|
||||
description: Base64-encoded image payload
|
||||
example: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==
|
||||
mime_type:
|
||||
type: string
|
||||
description: MIME type of the image
|
||||
example: image/png
|
||||
prompt:
|
||||
type: string
|
||||
description: |
|
||||
Instruction for the vision model. Optional; the service applies a
|
||||
default prompt ("Describe this image") when omitted or empty.
|
||||
example: List the objects visible in this image.
|
||||
system:
|
||||
type: string
|
||||
description: Optional system prompt that sets behavior and context for the vision model
|
||||
example: You are an expert at describing photographs for visually impaired users.
|
||||
|
|
@ -0,0 +1,21 @@
|
|||
type: object
|
||||
description: Image-to-text response
|
||||
required:
|
||||
- description
|
||||
properties:
|
||||
description:
|
||||
type: string
|
||||
description: Generated description of the image
|
||||
example: A red square on a white background.
|
||||
in_token:
|
||||
type: integer
|
||||
description: Number of input tokens consumed
|
||||
example: 245
|
||||
out_token:
|
||||
type: integer
|
||||
description: Number of output tokens generated
|
||||
example: 32
|
||||
model:
|
||||
type: string
|
||||
description: Model used to describe the image
|
||||
example: gpt-5-mini
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -23,6 +23,13 @@ properties:
|
|||
description: Provenance triples for this explain event (inline, no follow-up query needed)
|
||||
items:
|
||||
$ref: '../common/Triple.yaml'
|
||||
sources:
|
||||
type: array
|
||||
description: |
|
||||
Source documents the answer was derived from, deduplicated and sorted
|
||||
by URI. Present on the final message (end_of_session true).
|
||||
items:
|
||||
$ref: './Source.yaml'
|
||||
end_of_stream:
|
||||
type: boolean
|
||||
description: Indicates LLM response stream is complete
|
||||
|
|
|
|||
15
specs/api/components/schemas/rag/Source.yaml
Normal file
15
specs/api/components/schemas/rag/Source.yaml
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
type: object
|
||||
description: |
|
||||
Source document reference. Produced by tracing the graph edges used for
|
||||
retrieval back to the documents they were extracted from.
|
||||
properties:
|
||||
uri:
|
||||
type: string
|
||||
description: Source document URI
|
||||
example: urn:document:5a90a175-9906-4dcb-b482-a8c1b6cbf9e0
|
||||
title:
|
||||
type: string
|
||||
description: Document title (empty when the document has none)
|
||||
example: Quantum Mechanics Primer
|
||||
required:
|
||||
- uri
|
||||
|
|
@ -52,7 +52,7 @@ info:
|
|||
Workspace context comes from the authenticated token.
|
||||
|
||||
Accessed via `/api/v1/flow/{flow}/service/{kind}`:
|
||||
- AI services: agent, text-completion, prompt, RAG (document/graph)
|
||||
- AI services: agent, text-completion, image-to-text, prompt, RAG (document/graph)
|
||||
- Embeddings: embeddings, graph-embeddings, document-embeddings
|
||||
- Query: triples, rows, nlp-query, structured-query, sparql-query, row-embeddings
|
||||
- Data loading: text-load, document-load
|
||||
|
|
@ -136,6 +136,8 @@ paths:
|
|||
$ref: './paths/flow/graph-rag.yaml'
|
||||
/api/v1/flow/{flow}/service/text-completion:
|
||||
$ref: './paths/flow/text-completion.yaml'
|
||||
/api/v1/flow/{flow}/service/image-to-text:
|
||||
$ref: './paths/flow/image-to-text.yaml'
|
||||
/api/v1/flow/{flow}/service/prompt:
|
||||
$ref: './paths/flow/prompt.yaml'
|
||||
/api/v1/flow/{flow}/service/embeddings:
|
||||
|
|
|
|||
|
|
@ -118,6 +118,9 @@ post:
|
|||
- Quantum information theory
|
||||
- Computational complexity theory
|
||||
end-of-stream: false
|
||||
sources:
|
||||
- uri: urn:document:5a90a175-9906-4dcb-b482-a8c1b6cbf9e0
|
||||
title: Quantum Mechanics Primer
|
||||
streamingChunk:
|
||||
summary: Streaming response chunk
|
||||
value:
|
||||
|
|
|
|||
87
specs/api/paths/flow/image-to-text.yaml
Normal file
87
specs/api/paths/flow/image-to-text.yaml
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
post:
|
||||
tags:
|
||||
- Flow Services
|
||||
summary: Image to text - describe images with a vision model
|
||||
description: |
|
||||
Describe an image using a vision-capable LLM.
|
||||
|
||||
This is a **flow-scoped** service. It requires a flow instance
|
||||
and operates within the workspace associated with the
|
||||
authenticated bearer token.
|
||||
|
||||
## Image-to-Text Overview
|
||||
|
||||
Converts image content into a text description:
|
||||
- General image description ("what is in this picture?")
|
||||
- Targeted extraction via a custom prompt (objects, text, layout)
|
||||
- Behavior shaping via an optional system prompt
|
||||
|
||||
## Request Fields
|
||||
|
||||
- **image**: Base64-encoded image payload (raw binary is not accepted)
|
||||
- **mime_type**: Image MIME type, e.g. `image/png`, `image/jpeg`
|
||||
- **prompt**: Optional instruction; defaults to "Describe this image"
|
||||
- **system**: Optional system prompt for behavior and constraints
|
||||
|
||||
## Token Counting
|
||||
|
||||
Response includes token usage:
|
||||
- `in_token`: Input tokens (prompt + image)
|
||||
- `out_token`: Generated tokens
|
||||
- Useful for cost tracking and optimization
|
||||
|
||||
## Availability
|
||||
|
||||
Image-to-text is an optional service: it is only available in flows
|
||||
whose blueprint defines the `image-to-text` interface. Invoking it
|
||||
on a flow without the interface returns an error.
|
||||
|
||||
operationId: imageToTextService
|
||||
security:
|
||||
- bearerAuth: []
|
||||
parameters:
|
||||
- name: flow
|
||||
in: path
|
||||
required: true
|
||||
schema:
|
||||
type: string
|
||||
description: Flow instance ID
|
||||
example: my-flow
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '../../components/schemas/image-to-text/ImageToTextRequest.yaml'
|
||||
examples:
|
||||
basicDescription:
|
||||
summary: Basic image description with default prompt
|
||||
value:
|
||||
image: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==
|
||||
mime_type: image/png
|
||||
targetedExtraction:
|
||||
summary: Targeted extraction with custom prompts
|
||||
value:
|
||||
image: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==
|
||||
mime_type: image/png
|
||||
prompt: List all text visible in this image.
|
||||
system: You are an OCR assistant. Reply with the extracted text only.
|
||||
responses:
|
||||
'200':
|
||||
description: Successful response
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '../../components/schemas/image-to-text/ImageToTextResponse.yaml'
|
||||
examples:
|
||||
imageDescription:
|
||||
summary: Image description with token usage
|
||||
value:
|
||||
description: A single blue pixel on a transparent background.
|
||||
in_token: 245
|
||||
out_token: 32
|
||||
model: gpt-5-mini
|
||||
'401':
|
||||
$ref: '../../components/responses/Unauthorized.yaml'
|
||||
'500':
|
||||
$ref: '../../components/responses/Error.yaml'
|
||||
|
|
@ -43,7 +43,7 @@ info:
|
|||
- config, flow, librarian, knowledge, collection-management
|
||||
|
||||
**Flow-Scoped Services** (require `flow` parameter, workspace from token):
|
||||
- agent, text-completion, prompt, document-rag, graph-rag
|
||||
- agent, text-completion, image-to-text, prompt, document-rag, graph-rag
|
||||
- embeddings, graph-embeddings, document-embeddings
|
||||
- triples, rows, nlp-query, structured-query, sparql-query, structured-diag, row-embeddings
|
||||
- text-load, document-load, mcp-tool
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ payload:
|
|||
- $ref: './requests/DocumentRagRequest.yaml'
|
||||
- $ref: './requests/GraphRagRequest.yaml'
|
||||
- $ref: './requests/TextCompletionRequest.yaml'
|
||||
- $ref: './requests/ImageToTextRequest.yaml'
|
||||
- $ref: './requests/PromptRequest.yaml'
|
||||
- $ref: './requests/EmbeddingsRequest.yaml'
|
||||
- $ref: './requests/McpToolRequest.yaml'
|
||||
|
|
|
|||
|
|
@ -0,0 +1,28 @@
|
|||
type: object
|
||||
description: WebSocket request for image-to-text service (flow-scoped service)
|
||||
required:
|
||||
- id
|
||||
- service
|
||||
- flow
|
||||
- request
|
||||
properties:
|
||||
id:
|
||||
type: string
|
||||
description: Unique request identifier
|
||||
service:
|
||||
type: string
|
||||
const: image-to-text
|
||||
description: Service identifier for image-to-text service
|
||||
flow:
|
||||
type: string
|
||||
description: Flow ID
|
||||
request:
|
||||
$ref: '../../../../api/components/schemas/image-to-text/ImageToTextRequest.yaml'
|
||||
examples:
|
||||
- id: req-1
|
||||
service: image-to-text
|
||||
flow: my-flow
|
||||
request:
|
||||
image: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==
|
||||
mime_type: image/png
|
||||
prompt: Describe this image
|
||||
|
|
@ -25,9 +25,10 @@ properties:
|
|||
|
||||
Global services: iam
|
||||
Workspace-scoped services: config, flow, librarian, knowledge, collection-management
|
||||
Flow-scoped services: agent, text-completion, prompt, document-rag, graph-rag,
|
||||
embeddings, graph-embeddings, document-embeddings, triples, objects,
|
||||
nlp-query, structured-query, structured-diag, text-load, document-load, mcp-tool
|
||||
Flow-scoped services: agent, text-completion, image-to-text, prompt,
|
||||
document-rag, graph-rag, embeddings, graph-embeddings, document-embeddings,
|
||||
triples, objects, nlp-query, structured-query, structured-diag, text-load,
|
||||
document-load, mcp-tool
|
||||
examples:
|
||||
- config
|
||||
- agent
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from unittest.mock import MagicMock
|
|||
|
||||
from trustgraph.schema import (
|
||||
TextCompletionRequest, TextCompletionResponse,
|
||||
ImageToTextRequest, ImageToTextResponse,
|
||||
DocumentRagQuery, DocumentRagResponse,
|
||||
AgentRequest, AgentResponse, AgentStep,
|
||||
Chunk, Triple, Triples, Term, Error,
|
||||
|
|
@ -30,6 +31,10 @@ def schema_registry():
|
|||
"TextCompletionRequest": TextCompletionRequest,
|
||||
"TextCompletionResponse": TextCompletionResponse,
|
||||
|
||||
# Image to Text
|
||||
"ImageToTextRequest": ImageToTextRequest,
|
||||
"ImageToTextResponse": ImageToTextResponse,
|
||||
|
||||
# Document RAG
|
||||
"DocumentRagQuery": DocumentRagQuery,
|
||||
"DocumentRagResponse": DocumentRagResponse,
|
||||
|
|
@ -70,6 +75,20 @@ def sample_message_data():
|
|||
"out_token": 100,
|
||||
"model": "gpt-3.5-turbo"
|
||||
},
|
||||
"ImageToTextRequest": {
|
||||
# The image field carries base64 ASCII text end-to-end
|
||||
"image": "aW1hZ2UtYnl0ZXM=",
|
||||
"mime_type": "image/png",
|
||||
"prompt": "Describe this image",
|
||||
"system": "You are a helpful assistant."
|
||||
},
|
||||
"ImageToTextResponse": {
|
||||
"error": None,
|
||||
"description": "A single blue pixel on a white background.",
|
||||
"in_token": 245,
|
||||
"out_token": 32,
|
||||
"model": "gpt-5-mini"
|
||||
},
|
||||
"DocumentRagQuery": {
|
||||
"query": "What is artificial intelligence?",
|
||||
"collection": "test_collection",
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from pulsar.schema import Record
|
|||
|
||||
from trustgraph.schema import (
|
||||
TextCompletionRequest, TextCompletionResponse,
|
||||
ImageToTextRequest, ImageToTextResponse,
|
||||
DocumentRagQuery, DocumentRagResponse,
|
||||
AgentRequest, AgentResponse, AgentStep,
|
||||
Chunk, Triple, Triples, Term, Error,
|
||||
|
|
@ -117,6 +118,111 @@ class TestTextCompletionMessageContracts:
|
|||
assert error_response.response is None
|
||||
|
||||
|
||||
@pytest.mark.contract
|
||||
class TestImageToTextMessageContracts:
|
||||
"""Contract tests for Image to Text message schemas"""
|
||||
|
||||
def test_image_to_text_request_schema_contract(self, sample_message_data):
|
||||
"""Test ImageToTextRequest schema contract"""
|
||||
# Arrange
|
||||
request_data = sample_message_data["ImageToTextRequest"]
|
||||
|
||||
# Act & Assert
|
||||
assert validate_schema_contract(ImageToTextRequest, request_data)
|
||||
|
||||
# Test required fields
|
||||
request = ImageToTextRequest(**request_data)
|
||||
assert hasattr(request, 'image')
|
||||
assert hasattr(request, 'mime_type')
|
||||
assert hasattr(request, 'prompt')
|
||||
assert hasattr(request, 'system')
|
||||
# The image field carries base64 ASCII text end-to-end
|
||||
assert isinstance(request.image, str)
|
||||
assert isinstance(request.mime_type, str)
|
||||
|
||||
def test_image_to_text_response_schema_contract(self, sample_message_data):
|
||||
"""Test ImageToTextResponse schema contract"""
|
||||
# Arrange
|
||||
response_data = sample_message_data["ImageToTextResponse"]
|
||||
|
||||
# Act & Assert
|
||||
assert validate_schema_contract(ImageToTextResponse, response_data)
|
||||
|
||||
# Test required fields
|
||||
response = ImageToTextResponse(**response_data)
|
||||
assert hasattr(response, 'error')
|
||||
assert hasattr(response, 'description')
|
||||
assert hasattr(response, 'in_token')
|
||||
assert hasattr(response, 'out_token')
|
||||
assert hasattr(response, 'model')
|
||||
|
||||
def test_image_to_text_request_serialization_contract(self, sample_message_data):
|
||||
"""Test ImageToTextRequest serialization/deserialization contract"""
|
||||
# Arrange
|
||||
request_data = sample_message_data["ImageToTextRequest"]
|
||||
|
||||
# Act & Assert
|
||||
assert serialize_deserialize_test(ImageToTextRequest, request_data)
|
||||
|
||||
def test_image_to_text_response_serialization_contract(self, sample_message_data):
|
||||
"""Test ImageToTextResponse serialization/deserialization contract"""
|
||||
# Arrange
|
||||
response_data = sample_message_data["ImageToTextResponse"]
|
||||
|
||||
# Act & Assert
|
||||
assert serialize_deserialize_test(ImageToTextResponse, response_data)
|
||||
|
||||
def test_image_to_text_request_field_constraints(self):
|
||||
"""Test ImageToTextRequest field type constraints"""
|
||||
# Test valid data
|
||||
valid_request = ImageToTextRequest(
|
||||
image="aW1hZ2UtYnl0ZXM=",
|
||||
mime_type="image/jpeg",
|
||||
prompt="What is in this picture?",
|
||||
system="You are helpful."
|
||||
)
|
||||
assert valid_request.image == "aW1hZ2UtYnl0ZXM="
|
||||
assert valid_request.mime_type == "image/jpeg"
|
||||
assert valid_request.prompt == "What is in this picture?"
|
||||
assert valid_request.system == "You are helpful."
|
||||
|
||||
# Prompt and system are optional
|
||||
minimal_request = ImageToTextRequest(
|
||||
image="aW1hZ2UtYnl0ZXM=",
|
||||
mime_type="image/png"
|
||||
)
|
||||
assert minimal_request.prompt == ""
|
||||
assert minimal_request.system == ""
|
||||
|
||||
def test_image_to_text_response_field_constraints(self):
|
||||
"""Test ImageToTextResponse field type constraints"""
|
||||
# Test valid response with no error
|
||||
valid_response = ImageToTextResponse(
|
||||
error=None,
|
||||
description="A red square.",
|
||||
in_token=245,
|
||||
out_token=32,
|
||||
model="gpt-5-mini"
|
||||
)
|
||||
assert valid_response.error is None
|
||||
assert valid_response.description == "A red square."
|
||||
assert valid_response.in_token == 245
|
||||
assert valid_response.out_token == 32
|
||||
assert valid_response.model == "gpt-5-mini"
|
||||
|
||||
# Test response with error
|
||||
error_response = ImageToTextResponse(
|
||||
error=Error(type="image-to-text-error", message="Vision backend failed"),
|
||||
description=None,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=None
|
||||
)
|
||||
assert error_response.error is not None
|
||||
assert error_response.error.type == "image-to-text-error"
|
||||
assert error_response.description is None
|
||||
|
||||
|
||||
@pytest.mark.contract
|
||||
class TestDocumentRagMessageContracts:
|
||||
"""Contract tests for Document RAG message schemas"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ based on message fields like end_of_stream and end_of_dialog.
|
|||
import pytest
|
||||
|
||||
from trustgraph.schema import (
|
||||
GraphRagResponse, DocumentRagResponse, AgentResponse, Error
|
||||
GraphRagResponse, DocumentRagResponse, AgentResponse, Error, Source
|
||||
)
|
||||
from trustgraph.messaging import TranslatorRegistry
|
||||
|
||||
|
|
@ -110,6 +110,57 @@ class TestRAGTranslatorCompletionFlags:
|
|||
assert response_dict["end_of_stream"] is True
|
||||
assert response_dict["end_of_session"] is False
|
||||
|
||||
def test_graph_rag_translator_encodes_sources_on_final_message(self):
|
||||
"""
|
||||
Test that GraphRagResponseTranslator encodes source references
|
||||
as uri/title dicts on the final message.
|
||||
"""
|
||||
# Arrange
|
||||
translator = TranslatorRegistry.get_response_translator("graph-rag")
|
||||
response = GraphRagResponse(
|
||||
response="A small domesticated mammal.",
|
||||
message_type="chunk",
|
||||
end_of_stream=True,
|
||||
end_of_session=True,
|
||||
error=None,
|
||||
sources=[
|
||||
Source(uri="urn:document:alpha",
|
||||
title="Quantum Mechanics Primer"),
|
||||
Source(uri="urn:document:beta", title=""),
|
||||
]
|
||||
)
|
||||
|
||||
# Act
|
||||
response_dict, is_final = translator.encode_with_completion(response)
|
||||
|
||||
# Assert
|
||||
assert is_final is True
|
||||
assert response_dict["sources"] == [
|
||||
{"uri": "urn:document:alpha",
|
||||
"title": "Quantum Mechanics Primer"},
|
||||
{"uri": "urn:document:beta", "title": ""},
|
||||
]
|
||||
|
||||
def test_graph_rag_translator_omits_empty_sources(self):
|
||||
"""
|
||||
Test that the sources key is omitted when there are no sources.
|
||||
"""
|
||||
# Arrange
|
||||
translator = TranslatorRegistry.get_response_translator("graph-rag")
|
||||
response = GraphRagResponse(
|
||||
response="Chunk 1",
|
||||
message_type="chunk",
|
||||
end_of_stream=False,
|
||||
end_of_session=False,
|
||||
error=None
|
||||
)
|
||||
|
||||
# Act
|
||||
response_dict, is_final = translator.encode_with_completion(response)
|
||||
|
||||
# Assert
|
||||
assert "sources" not in response_dict
|
||||
|
||||
def test_document_rag_translator_is_final_with_end_of_session_true(self):
|
||||
"""
|
||||
Test that DocumentRagResponseTranslator returns is_final=True
|
||||
|
|
|
|||
|
|
@ -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,11 +171,11 @@ 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
|
||||
response, usage, sources = response
|
||||
assert response is not None
|
||||
assert isinstance(response, str)
|
||||
assert "machine learning" in response.lower()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -123,7 +127,7 @@ class TestGraphRagStreaming:
|
|||
)
|
||||
|
||||
# Assert
|
||||
response, usage = response
|
||||
response, usage, sources = response
|
||||
assert_streaming_chunks_valid(collector.chunks, min_chunks=1)
|
||||
assert_callback_invoked(AsyncMock(call_count=len(collector.chunks)), min_calls=1)
|
||||
|
||||
|
|
@ -171,8 +175,8 @@ class TestGraphRagStreaming:
|
|||
)
|
||||
|
||||
# Assert - Results should be equivalent
|
||||
non_streaming_text, _ = non_streaming_response
|
||||
streaming_text, _ = streaming_response
|
||||
non_streaming_text, _, _ = non_streaming_response
|
||||
streaming_text, _, _ = streaming_response
|
||||
assert streaming_text == non_streaming_text
|
||||
assert len(streaming_chunks) > 0
|
||||
assert "".join(streaming_chunks) == streaming_text
|
||||
|
|
@ -212,7 +216,7 @@ class TestGraphRagStreaming:
|
|||
|
||||
# Assert - Should complete without error
|
||||
assert response is not None
|
||||
response_text, usage = response
|
||||
response_text, usage, sources = response
|
||||
assert isinstance(response_text, str)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ class TestPromptStreaming:
|
|||
" of", " artificial", " intelligence", "."
|
||||
]
|
||||
|
||||
async def streaming_text_completion_stream(system, prompt, handler, timeout=600):
|
||||
async def streaming_text_completion_stream(system, prompt, handler, timeout=600, response_format=None, schema=None):
|
||||
"""Simulate streaming text completion via text_completion_stream"""
|
||||
for i, chunk_text in enumerate(chunks):
|
||||
response = TextCompletionResponse(
|
||||
|
|
@ -58,7 +58,7 @@ class TestPromptStreaming:
|
|||
model="test-model",
|
||||
)
|
||||
|
||||
async def non_streaming_text_completion(system, prompt, timeout=600):
|
||||
async def non_streaming_text_completion(system, prompt, timeout=600, response_format=None, schema=None):
|
||||
"""Simulate non-streaming text completion"""
|
||||
full_text = "Machine learning is a field of artificial intelligence."
|
||||
return TextCompletionResult(
|
||||
|
|
@ -230,7 +230,7 @@ class TestPromptStreaming:
|
|||
# Mock text completion client that raises an error
|
||||
text_completion_client = AsyncMock()
|
||||
|
||||
async def failing_stream(system, prompt, handler, timeout=600):
|
||||
async def failing_stream(system, prompt, handler, timeout=600, response_format=None, schema=None):
|
||||
raise RuntimeError("Text completion error")
|
||||
|
||||
text_completion_client.text_completion_stream = AsyncMock(
|
||||
|
|
@ -316,7 +316,7 @@ class TestPromptStreaming:
|
|||
# Mock text completion that sends empty chunks
|
||||
text_completion_client = AsyncMock()
|
||||
|
||||
async def empty_streaming(system, prompt, handler, timeout=600):
|
||||
async def empty_streaming(system, prompt, handler, timeout=600, response_format=None, schema=None):
|
||||
# Send empty chunk followed by final marker
|
||||
await handler(TextCompletionResponse(
|
||||
response="",
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
100
tests/unit/test_api/test_rag_chunk_sources.py
Normal file
100
tests/unit/test_api/test_rag_chunk_sources.py
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
"""
|
||||
Tests that the socket clients propagate the sources field from the
|
||||
wire format to RAGChunk, and that graph_rag results carry it.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from trustgraph.api.socket_client import SocketClient
|
||||
from trustgraph.api.async_socket_client import AsyncSocketClient
|
||||
from trustgraph.api.types import RAGChunk
|
||||
|
||||
WIRE_SOURCES = [
|
||||
{"uri": "urn:document:alpha", "title": "Quantum Mechanics Primer"},
|
||||
{"uri": "urn:document:beta", "title": ""},
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
# We only need _parse_chunk — don't connect
|
||||
c = object.__new__(SocketClient)
|
||||
return c
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_client():
|
||||
c = object.__new__(AsyncSocketClient)
|
||||
return c
|
||||
|
||||
|
||||
class TestParseChunkSources:
|
||||
|
||||
def test_final_chunk_carries_sources(self, client):
|
||||
resp = {
|
||||
"message_type": "chunk",
|
||||
"response": "",
|
||||
"end_of_stream": False,
|
||||
"end_of_session": True,
|
||||
"sources": WIRE_SOURCES,
|
||||
}
|
||||
chunk = client._parse_chunk(resp)
|
||||
assert isinstance(chunk, RAGChunk)
|
||||
assert chunk.sources == WIRE_SOURCES
|
||||
|
||||
def test_intermediate_chunk_has_empty_sources(self, client):
|
||||
resp = {
|
||||
"message_type": "chunk",
|
||||
"response": "partial text",
|
||||
"end_of_stream": False,
|
||||
}
|
||||
chunk = client._parse_chunk(resp)
|
||||
assert isinstance(chunk, RAGChunk)
|
||||
assert chunk.sources == []
|
||||
|
||||
def test_async_final_chunk_carries_sources(self, async_client):
|
||||
resp = {
|
||||
"message_type": "chunk",
|
||||
"response": "",
|
||||
"end_of_session": True,
|
||||
"sources": WIRE_SOURCES,
|
||||
}
|
||||
chunk = async_client._parse_chunk(resp)
|
||||
assert isinstance(chunk, RAGChunk)
|
||||
assert chunk.sources == WIRE_SOURCES
|
||||
|
||||
def test_async_intermediate_chunk_has_empty_sources(self, async_client):
|
||||
resp = {
|
||||
"message_type": "chunk",
|
||||
"response": "partial text",
|
||||
}
|
||||
chunk = async_client._parse_chunk(resp)
|
||||
assert isinstance(chunk, RAGChunk)
|
||||
assert chunk.sources == []
|
||||
|
||||
|
||||
class TestRestGraphRagSources:
|
||||
|
||||
def test_graph_rag_result_carries_sources(self):
|
||||
from trustgraph.api.flow import FlowInstance
|
||||
|
||||
instance = object.__new__(FlowInstance)
|
||||
instance.request = lambda path, request: {
|
||||
"response": "The answer.",
|
||||
"sources": WIRE_SOURCES,
|
||||
}
|
||||
|
||||
result = instance.graph_rag(query="What is quantum computing?")
|
||||
assert result.text == "The answer."
|
||||
assert result.sources == WIRE_SOURCES
|
||||
|
||||
def test_graph_rag_result_defaults_to_empty_sources(self):
|
||||
from trustgraph.api.flow import FlowInstance
|
||||
|
||||
instance = object.__new__(FlowInstance)
|
||||
instance.request = lambda path, request: {
|
||||
"response": "The answer.",
|
||||
}
|
||||
|
||||
result = instance.graph_rag(query="What is quantum computing?")
|
||||
assert result.sources == []
|
||||
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
|
||||
224
tests/unit/test_base/test_image_to_text_service.py
Normal file
224
tests/unit/test_base/test_image_to_text_service.py
Normal file
|
|
@ -0,0 +1,224 @@
|
|||
"""
|
||||
Unit tests for the ImageToTextService base class
|
||||
Following the same pattern as the LLM service parameter tests
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest import IsolatedAsyncioTestCase
|
||||
|
||||
from trustgraph.base.image_to_text_service import (
|
||||
ImageToTextService, ImageDescriptionResult,
|
||||
)
|
||||
from trustgraph.base import ParameterSpec, ConsumerSpec, ProducerSpec
|
||||
from trustgraph.schema import ImageToTextRequest, ImageToTextResponse
|
||||
from trustgraph.exceptions import TooManyRequests
|
||||
|
||||
|
||||
class MockAsyncProcessor:
|
||||
def __init__(self, **params):
|
||||
self.config_handlers = []
|
||||
self.id = params.get('id', 'test-service')
|
||||
self.specifications = []
|
||||
|
||||
|
||||
class TestImageToTextService(IsolatedAsyncioTestCase):
|
||||
"""Test image-to-text service base class functionality"""
|
||||
|
||||
def make_service(self):
|
||||
config = {
|
||||
'id': 'test-image-to-text-service',
|
||||
'concurrency': 1,
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
return ImageToTextService(**config)
|
||||
|
||||
def make_message(self, image="aW1hZ2U=", mime_type="image/png",
|
||||
prompt="Describe this image", system="Be concise"):
|
||||
mock_message = MagicMock()
|
||||
mock_message.value.return_value = MagicMock()
|
||||
mock_message.value.return_value.image = image
|
||||
mock_message.value.return_value.mime_type = mime_type
|
||||
mock_message.value.return_value.prompt = prompt
|
||||
mock_message.value.return_value.system = system
|
||||
mock_message.properties.return_value = {"id": "test-id"}
|
||||
return mock_message
|
||||
|
||||
def make_flow(self, model="vision-model"):
|
||||
mock_response_producer = AsyncMock()
|
||||
|
||||
mock_flow = MagicMock()
|
||||
mock_flow.name = "test-flow"
|
||||
mock_flow.side_effect = lambda param: {
|
||||
"model": model,
|
||||
"response": mock_response_producer,
|
||||
}.get(param)
|
||||
|
||||
mock_error_producer = AsyncMock()
|
||||
mock_flow.producer = {"response": mock_error_producer}
|
||||
|
||||
return mock_flow, mock_response_producer, mock_error_producer
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
def test_specification_registration(self):
|
||||
"""Test that the service registers request/response/model specs"""
|
||||
# Act
|
||||
service = self.make_service()
|
||||
|
||||
# Assert
|
||||
consumer_specs = {spec.name: spec for spec in service.specifications
|
||||
if isinstance(spec, ConsumerSpec)}
|
||||
producer_specs = {spec.name: spec for spec in service.specifications
|
||||
if isinstance(spec, ProducerSpec)}
|
||||
param_specs = {spec.name: spec for spec in service.specifications
|
||||
if isinstance(spec, ParameterSpec)}
|
||||
|
||||
assert "request" in consumer_specs
|
||||
assert consumer_specs["request"].schema == ImageToTextRequest
|
||||
assert "response" in producer_specs
|
||||
assert producer_specs["response"].schema == ImageToTextResponse
|
||||
assert "model" in param_specs
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_dispatches_to_describe_image(self):
|
||||
"""Test that on_request dispatches request fields to describe_image"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(return_value=ImageDescriptionResult(
|
||||
text="A cat on a mat",
|
||||
in_token=10,
|
||||
out_token=5,
|
||||
model="vision-model"
|
||||
))
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, mock_producer, _ = self.make_flow()
|
||||
|
||||
# Act
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# Assert
|
||||
service.describe_image.assert_called_once()
|
||||
call_args = service.describe_image.call_args
|
||||
|
||||
assert call_args[0][0] == "aW1hZ2U=" # image
|
||||
assert call_args[0][1] == "image/png" # mime_type
|
||||
assert call_args[0][2] == "Describe this image" # prompt
|
||||
assert call_args[0][3] == "Be concise" # system
|
||||
assert call_args[0][4] == "vision-model" # model
|
||||
|
||||
# Verify flow was queried for the model parameter
|
||||
mock_flow.assert_any_call("model")
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_formats_response(self):
|
||||
"""Test that on_request propagates description/tokens/model"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(return_value=ImageDescriptionResult(
|
||||
text="A cat on a mat",
|
||||
in_token=10,
|
||||
out_token=5,
|
||||
model="vision-model"
|
||||
))
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, mock_producer, _ = self.make_flow()
|
||||
|
||||
# Act
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# Assert
|
||||
mock_producer.send.assert_called_once()
|
||||
response = mock_producer.send.call_args[0][0]
|
||||
properties = mock_producer.send.call_args[1]["properties"]
|
||||
|
||||
assert response.error is None
|
||||
assert response.description == "A cat on a mat"
|
||||
assert response.in_token == 10
|
||||
assert response.out_token == 5
|
||||
assert response.model == "vision-model"
|
||||
assert properties == {"id": "test-id"}
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_handles_missing_model_parameter(self):
|
||||
"""Test that on_request passes None model when flow has none"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(return_value=ImageDescriptionResult(
|
||||
text="A cat on a mat",
|
||||
in_token=10,
|
||||
out_token=5,
|
||||
model="default-model"
|
||||
))
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, mock_producer, _ = self.make_flow(model=None)
|
||||
|
||||
# Act
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# Assert
|
||||
service.describe_image.assert_called_once()
|
||||
call_args = service.describe_image.call_args
|
||||
|
||||
assert call_args[0][4] is None # model (will use processor default)
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_error_produces_structured_error(self):
|
||||
"""Test that a backend exception produces a structured error response"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(side_effect=Exception("Test error"))
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, _, mock_error_producer = self.make_flow()
|
||||
|
||||
# Act
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# Assert
|
||||
mock_error_producer.send.assert_called_once()
|
||||
error_response = mock_error_producer.send.call_args[0][0]
|
||||
|
||||
assert error_response.error is not None
|
||||
assert error_response.error.type == "image-to-text-error"
|
||||
assert "Test error" in error_response.error.message
|
||||
assert error_response.description is None
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_reraises_too_many_requests(self):
|
||||
"""Test that TooManyRequests is re-raised for the retry machinery"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(side_effect=TooManyRequests())
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, mock_producer, mock_error_producer = self.make_flow()
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(TooManyRequests):
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# No response of any kind should have been sent
|
||||
mock_producer.send.assert_not_called()
|
||||
mock_error_producer.send.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main([__file__])
|
||||
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"""
|
||||
|
|
|
|||
268
tests/unit/test_base/test_schema_compatibility.py
Normal file
268
tests/unit/test_base/test_schema_compatibility.py
Normal file
|
|
@ -0,0 +1,268 @@
|
|||
"""
|
||||
Unit tests for schema_compatibility.is_strict_mode_compatible
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from trustgraph.base.schema_compatibility import is_strict_mode_compatible
|
||||
|
||||
|
||||
class TestIsStrictModeCompatible:
|
||||
|
||||
def test_none_schema(self):
|
||||
assert is_strict_mode_compatible(None) is False
|
||||
|
||||
def test_empty_dict(self):
|
||||
assert is_strict_mode_compatible({}) is True
|
||||
|
||||
def test_simple_string(self):
|
||||
assert is_strict_mode_compatible({"type": "string"}) is True
|
||||
|
||||
def test_compliant_object(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"age": {"type": "integer"},
|
||||
},
|
||||
"required": ["name", "age"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_missing_additional_properties(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {"name": {"type": "string"}},
|
||||
"required": ["name"],
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_additional_properties_true(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {"name": {"type": "string"}},
|
||||
"required": ["name"],
|
||||
"additionalProperties": True,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_property_not_in_required(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"nickname": {"type": "string"},
|
||||
},
|
||||
"required": ["name"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_open_ended_object_no_properties(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_implicit_object_with_properties_key(self):
|
||||
schema = {
|
||||
"properties": {
|
||||
"x": {"type": "number"},
|
||||
},
|
||||
"required": ["x"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_nested_object_compliant(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"address": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"street": {"type": "string"},
|
||||
},
|
||||
"required": ["street"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"required": ["address"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_nested_object_non_compliant(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"metadata": {
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
"required": ["metadata"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_array_with_compliant_items(self):
|
||||
schema = {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {"type": "integer"},
|
||||
},
|
||||
"required": ["id"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_array_with_non_compliant_items(self):
|
||||
schema = {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {"id": {"type": "integer"}},
|
||||
},
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_array_with_simple_items(self):
|
||||
schema = {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_oneof_all_compliant(self):
|
||||
schema = {
|
||||
"oneOf": [
|
||||
{"type": "string"},
|
||||
{"type": "integer"},
|
||||
]
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_oneof_with_non_compliant_branch(self):
|
||||
schema = {
|
||||
"oneOf": [
|
||||
{"type": "string"},
|
||||
{"type": "object"},
|
||||
]
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_anyof(self):
|
||||
schema = {
|
||||
"anyOf": [
|
||||
{"type": "string"},
|
||||
{"type": "number"},
|
||||
]
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_allof(self):
|
||||
schema = {
|
||||
"allOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"a": {"type": "string"}},
|
||||
"required": ["a"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
]
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_unsupported_minimum(self):
|
||||
schema = {"type": "integer", "minimum": 0}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_maximum(self):
|
||||
schema = {"type": "integer", "maximum": 100}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_pattern(self):
|
||||
schema = {"type": "string", "pattern": "^[a-z]+$"}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_min_length(self):
|
||||
schema = {"type": "string", "minLength": 1}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_max_length(self):
|
||||
schema = {"type": "string", "maxLength": 255}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_min_items(self):
|
||||
schema = {"type": "array", "items": {"type": "string"}, "minItems": 1}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_max_items(self):
|
||||
schema = {"type": "array", "items": {"type": "string"}, "maxItems": 10}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_exclusive_minimum(self):
|
||||
schema = {"type": "number", "exclusiveMinimum": 0}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_min_max_properties(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {"a": {"type": "string"}},
|
||||
"required": ["a"],
|
||||
"additionalProperties": False,
|
||||
"minProperties": 1,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_unsupported_constraint_inside_nested_property(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"score": {"type": "integer", "minimum": 0, "maximum": 100},
|
||||
},
|
||||
"required": ["score"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
|
||||
def test_nullable_property(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": ["string", "null"]},
|
||||
},
|
||||
"required": ["name"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_realistic_compliant_schema(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"action": {"type": "string"},
|
||||
"services": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"required": ["action", "services"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is True
|
||||
|
||||
def test_realistic_non_compliant_optional_field(self):
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"action": {"type": "string"},
|
||||
"reason": {"type": "string"},
|
||||
},
|
||||
"required": ["action"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
assert is_strict_mode_compatible(schema) is False
|
||||
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:
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,79 @@
|
|||
"""
|
||||
Unit tests for entity URI normalization.
|
||||
|
||||
Covers ASCII behaviour and, critically, non-ASCII (e.g. Chinese) entity
|
||||
names and types, which must be preserved rather than stripped so that
|
||||
distinct entities do not collapse onto a single URI.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from trustgraph.extract.kg.ontology.entity_normalizer import (
|
||||
normalize_entity_name,
|
||||
normalize_type_identifier,
|
||||
build_entity_uri,
|
||||
EntityRegistry,
|
||||
)
|
||||
|
||||
|
||||
class TestNormalizeEntityNameAscii:
|
||||
"""ASCII normalization must keep working as before."""
|
||||
|
||||
def test_lowercases_and_hyphenates(self):
|
||||
assert normalize_entity_name("Cornish pasty") == "cornish-pasty"
|
||||
|
||||
def test_strips_punctuation(self):
|
||||
assert normalize_entity_name("beef!!!") == "beef"
|
||||
|
||||
def test_collapses_and_trims_hyphens(self):
|
||||
assert normalize_entity_name(" beef stew ") == "beef-stew"
|
||||
|
||||
|
||||
class TestNormalizeEntityNameUnicode:
|
||||
"""Non-ASCII names must be preserved, not stripped to empty."""
|
||||
|
||||
def test_chinese_name_preserved(self):
|
||||
# Previously the ASCII-only filter deleted every CJK character,
|
||||
# collapsing the name to "".
|
||||
assert normalize_entity_name("苹果") == "苹果"
|
||||
|
||||
def test_chinese_name_with_space(self):
|
||||
assert normalize_entity_name("有机 苹果") == "有机-苹果"
|
||||
|
||||
def test_distinct_chinese_names_stay_distinct(self):
|
||||
assert normalize_entity_name("苹果") != normalize_entity_name("香蕉")
|
||||
|
||||
def test_mixed_ascii_and_chinese(self):
|
||||
assert normalize_entity_name("iPhone 手机") == "iphone-手机"
|
||||
|
||||
|
||||
class TestNormalizeTypeIdentifierUnicode:
|
||||
def test_ascii_prefixed_type(self):
|
||||
assert normalize_type_identifier("fo/Recipe") == "fo-recipe"
|
||||
|
||||
def test_chinese_type_preserved(self):
|
||||
assert normalize_type_identifier("食物") == "食物"
|
||||
|
||||
|
||||
class TestBuildEntityUriUnicode:
|
||||
def test_distinct_chinese_entities_get_distinct_uris(self):
|
||||
# This is the core regression: with ASCII-only normalization both
|
||||
# names collapsed to the same URI, merging unrelated entities.
|
||||
uri_apple = build_entity_uri("苹果", "水果", "food")
|
||||
uri_banana = build_entity_uri("香蕉", "水果", "food")
|
||||
assert uri_apple != uri_banana
|
||||
assert uri_apple.endswith("苹果")
|
||||
assert uri_banana.endswith("香蕉")
|
||||
|
||||
def test_uri_is_not_bare_prefix(self):
|
||||
uri = build_entity_uri("苹果", "水果", "food")
|
||||
# Must not collapse to ".../food/水果-" or ".../food/-"
|
||||
assert not uri.endswith("-")
|
||||
|
||||
|
||||
class TestEntityRegistryUnicode:
|
||||
def test_distinct_chinese_entities_not_deduplicated(self):
|
||||
registry = EntityRegistry("food")
|
||||
uri_apple = registry.get_or_create_uri("苹果", "水果")
|
||||
uri_banana = registry.get_or_create_uri("香蕉", "水果")
|
||||
assert uri_apple != uri_banana
|
||||
assert registry.size() == 2
|
||||
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
|
||||
|
|
|
|||
|
|
@ -588,6 +588,13 @@ class TestDispatcherManager:
|
|||
with pytest.raises(RuntimeError, match="This kind not supported by flow"):
|
||||
await manager.invoke_flow_service("data", "responder", "default", "test_flow", "agent")
|
||||
|
||||
def test_request_response_dispatchers_include_image_to_text(self):
|
||||
"""image-to-text must be registered as a request/response service"""
|
||||
from trustgraph.gateway.dispatch.manager import request_response_dispatchers
|
||||
from trustgraph.gateway.dispatch.image_to_text import ImageToTextRequestor
|
||||
|
||||
assert request_response_dispatchers["image-to-text"] is ImageToTextRequestor
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoke_flow_service_invalid_kind(self):
|
||||
"""Test invoke_flow_service with invalid kind"""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
3
tests/unit/test_image_to_text/__init__.py
Normal file
3
tests/unit/test_image_to_text/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
"""
|
||||
Unit tests for image-to-text services
|
||||
"""
|
||||
372
tests/unit/test_image_to_text/test_openai_processor.py
Normal file
372
tests/unit/test_image_to_text/test_openai_processor.py
Normal file
|
|
@ -0,0 +1,372 @@
|
|||
"""
|
||||
Unit tests for trustgraph.model.image_to_text.openai
|
||||
Following the same successful pattern as the text completion OpenAI tests
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest import IsolatedAsyncioTestCase
|
||||
|
||||
# Import the service under test
|
||||
from trustgraph.model.image_to_text.openai.service import Processor
|
||||
from trustgraph.base import ImageDescriptionResult
|
||||
from trustgraph.exceptions import TooManyRequests, LlmError
|
||||
|
||||
SAMPLE_IMAGE = base64.b64encode(b"image-data").decode("utf-8")
|
||||
|
||||
|
||||
class TestOpenAIImageToTextProcessor(IsolatedAsyncioTestCase):
|
||||
"""Test OpenAI image-to-text processor functionality"""
|
||||
|
||||
def make_config(self, **overrides):
|
||||
config = {
|
||||
'model': 'test-vision-model',
|
||||
'api_key': 'test-api-key',
|
||||
'url': 'https://api.openai.com/v1',
|
||||
'max_output': 4096,
|
||||
'concurrency': 1,
|
||||
'taskgroup': AsyncMock(),
|
||||
'id': 'test-processor'
|
||||
}
|
||||
config.update(overrides)
|
||||
return config
|
||||
|
||||
def make_response(self, content="An image description",
|
||||
prompt_tokens=20, completion_tokens=12):
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = content
|
||||
mock_response.usage.prompt_tokens = prompt_tokens
|
||||
mock_response.usage.completion_tokens = completion_tokens
|
||||
return mock_response
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_processor_initialization_basic(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test basic processor initialization"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
# Act
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Assert
|
||||
assert processor.default_model == 'test-vision-model'
|
||||
assert processor.max_output == 4096
|
||||
assert hasattr(processor, 'openai')
|
||||
mock_openai_class.assert_called_once_with(base_url='https://api.openai.com/v1', api_key='test-api-key')
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_processor_initialization_with_defaults(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test processor initialization with default values"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
# Only provide required fields, should use defaults
|
||||
config = {
|
||||
'api_key': 'test-api-key',
|
||||
'concurrency': 1,
|
||||
'taskgroup': AsyncMock(),
|
||||
'id': 'test-processor'
|
||||
}
|
||||
|
||||
# Act
|
||||
processor = Processor(**config)
|
||||
|
||||
# Assert
|
||||
assert processor.default_model == 'gpt-5-mini' # default_model
|
||||
assert processor.max_output == 4096 # default_max_output
|
||||
mock_openai_class.assert_called_once_with(base_url='https://api.openai.com/v1', api_key='test-api-key')
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_processor_initialization_without_api_key(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test processor initialization without API key uses placeholder"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
# Act
|
||||
processor = Processor(**self.make_config(api_key=None))
|
||||
|
||||
# Assert
|
||||
mock_openai_class.assert_called_once_with(
|
||||
base_url='https://api.openai.com/v1', api_key='not-set'
|
||||
)
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_openai_client_initialization_without_base_url(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test OpenAI client initialization without base_url"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
# Act
|
||||
processor = Processor(**self.make_config(url=None))
|
||||
|
||||
# Assert - should be called without base_url when it's None
|
||||
mock_openai_class.assert_called_once_with(api_key='test-api-key')
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_success(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test successful image description with data-URI message shape"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act
|
||||
result = await processor.describe_image(
|
||||
SAMPLE_IMAGE, "image/png", "What is in this image?", "",
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, ImageDescriptionResult)
|
||||
assert result.text == "An image description"
|
||||
assert result.in_token == 20
|
||||
assert result.out_token == 12
|
||||
assert result.model == 'test-vision-model'
|
||||
|
||||
# Verify the OpenAI API call
|
||||
mock_openai_client.chat.completions.create.assert_called_once_with(
|
||||
model='test-vision-model',
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What is in this image?"
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{SAMPLE_IMAGE}"
|
||||
}
|
||||
}
|
||||
]
|
||||
}],
|
||||
max_completion_tokens=4096
|
||||
)
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_default_prompt(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that an empty prompt falls back to the default prompt"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "", "")
|
||||
|
||||
# Assert
|
||||
call_args = mock_openai_client.chat.completions.create.call_args
|
||||
text_block = call_args[1]['messages'][0]['content'][0]
|
||||
|
||||
assert text_block['text'] == "Describe this image"
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_system_prompt(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that a system prompt is prepended to the user prompt"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act
|
||||
await processor.describe_image(
|
||||
SAMPLE_IMAGE, "image/png", "What is this?", "You are terse",
|
||||
)
|
||||
|
||||
# Assert
|
||||
call_args = mock_openai_client.chat.completions.create.call_args
|
||||
text_block = call_args[1]['messages'][0]['content'][0]
|
||||
|
||||
assert text_block['text'] == "You are terse\n\nWhat is this?"
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_model_override(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test model parameter override functionality"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act - Override model at runtime
|
||||
await processor.describe_image(
|
||||
SAMPLE_IMAGE, "image/png", "Describe", "",
|
||||
model="other-vision-model",
|
||||
)
|
||||
|
||||
# Assert
|
||||
call_kwargs = mock_openai_client.chat.completions.create.call_args.kwargs
|
||||
assert call_kwargs['model'] == 'other-vision-model'
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_no_override_uses_default(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that no model override uses the processor default"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act
|
||||
await processor.describe_image(
|
||||
SAMPLE_IMAGE, "image/png", "Describe", "", model=None,
|
||||
)
|
||||
|
||||
# Assert
|
||||
call_kwargs = mock_openai_client.chat.completions.create.call_args.kwargs
|
||||
assert call_kwargs['model'] == 'test-vision-model'
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_rate_limit_error(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test rate limit error handling"""
|
||||
# Arrange
|
||||
from openai import RateLimitError
|
||||
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.side_effect = RateLimitError("Rate limit exceeded", response=MagicMock(), body=None)
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(TooManyRequests):
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "Describe", "")
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_rate_limit_unrecoverable(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that unrecoverable rate limit codes raise RuntimeError"""
|
||||
# Arrange
|
||||
from openai import RateLimitError
|
||||
|
||||
body = {
|
||||
"error": {
|
||||
"code": "insufficient_quota",
|
||||
"message": "You exceeded your current quota",
|
||||
}
|
||||
}
|
||||
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.side_effect = RateLimitError("Rate limit exceeded", response=MagicMock(), body=body)
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(RuntimeError, match="insufficient_quota"):
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "Describe", "")
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_internal_server_error(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that InternalServerError is mapped to retryable LlmError"""
|
||||
# Arrange
|
||||
from openai import InternalServerError
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 503
|
||||
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.side_effect = InternalServerError("Service unavailable", response=mock_response, body=None)
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(LlmError):
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "Describe", "")
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_generic_exception(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test handling of generic exceptions"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.side_effect = Exception("API connection error")
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(Exception, match="API connection error"):
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "Describe", "")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main([__file__])
|
||||
|
|
@ -107,6 +107,7 @@ class TestGraphRagDagStructure:
|
|||
embeddings_client = AsyncMock()
|
||||
graph_embeddings_client = AsyncMock()
|
||||
triples_client = AsyncMock()
|
||||
reranker_client = AsyncMock()
|
||||
|
||||
embeddings_client.embed.return_value = [[0.1, 0.2]]
|
||||
graph_embeddings_client.query.return_value = [
|
||||
|
|
@ -121,27 +122,22 @@ class TestGraphRagDagStructure:
|
|||
]
|
||||
triples_client.query.return_value = []
|
||||
|
||||
result = MagicMock()
|
||||
result.document_id = "0"
|
||||
result.query_id = "0"
|
||||
result.score = 0.95
|
||||
reranker_client.rerank.return_value = [result]
|
||||
|
||||
async def mock_prompt(template_id, variables=None, **kwargs):
|
||||
if template_id == "extract-concepts":
|
||||
return PromptResult(response_type="text", text="concept")
|
||||
elif template_id == "kg-edge-scoring":
|
||||
edges = variables.get("knowledge", [])
|
||||
return PromptResult(
|
||||
response_type="jsonl",
|
||||
objects=[{"id": e["id"], "score": 10} for e in edges],
|
||||
)
|
||||
elif template_id == "kg-edge-reasoning":
|
||||
edges = variables.get("knowledge", [])
|
||||
return PromptResult(
|
||||
response_type="jsonl",
|
||||
objects=[{"id": e["id"], "reasoning": "relevant"} for e in edges],
|
||||
)
|
||||
elif template_id == "kg-synthesis":
|
||||
return PromptResult(response_type="text", text="Answer.")
|
||||
return PromptResult(response_type="text", text="")
|
||||
|
||||
prompt_client.prompt.side_effect = mock_prompt
|
||||
return prompt_client, embeddings_client, graph_embeddings_client, triples_client
|
||||
return (prompt_client, embeddings_client, graph_embeddings_client,
|
||||
triples_client, reranker_client)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dag_chain(self, mock_clients):
|
||||
|
|
@ -152,7 +148,7 @@ class TestGraphRagDagStructure:
|
|||
events.append({"explain_id": explain_id, "triples": triples})
|
||||
|
||||
await rag.query(
|
||||
query="test", explain_callback=explain_cb, edge_score_limit=0,
|
||||
query="test", explain_callback=explain_cb,
|
||||
)
|
||||
|
||||
dag = _collect_events(events)
|
||||
|
|
|
|||
|
|
@ -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,10 +552,10 @@ class TestQuery:
|
|||
collection="test_collection",
|
||||
entity_limit=25,
|
||||
triple_limit=15,
|
||||
explain_callback=collect_provenance
|
||||
explain_callback=collect_provenance,
|
||||
)
|
||||
|
||||
response_text, usage = response
|
||||
response_text, usage, sources = response
|
||||
assert response_text == expected_response
|
||||
|
||||
# 5 events: question, grounding, exploration, focus, synthesis
|
||||
|
|
@ -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]
|
||||
|
|
@ -14,13 +14,16 @@ from dataclasses import dataclass
|
|||
from trustgraph.retrieval.graph_rag.graph_rag import GraphRag, edge_id
|
||||
from trustgraph.schema import Triple as SchemaTriple, Term, IRI, LITERAL
|
||||
from trustgraph.base import PromptResult
|
||||
from trustgraph.base.triples_client import Triple as ClientTriple
|
||||
from trustgraph.knowledge import Uri, Literal
|
||||
|
||||
from trustgraph.provenance.namespaces import (
|
||||
RDF_TYPE, PROV_ENTITY, PROV_WAS_DERIVED_FROM,
|
||||
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,
|
||||
TG_CONTAINS, DC_TITLE, RDFS_LABEL,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -91,17 +94,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 +119,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 +133,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 +153,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 +162,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 +190,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 +215,7 @@ class TestGraphRagQueryProvenance:
|
|||
await rag.query(
|
||||
query="What is quantum computing?",
|
||||
explain_callback=explain_callback,
|
||||
edge_score_limit=0,
|
||||
|
||||
)
|
||||
|
||||
expected_types = [
|
||||
|
|
@ -260,7 +253,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 +290,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 +313,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 +337,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 +348,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 +364,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 +379,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 +407,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"]
|
||||
|
|
@ -426,10 +426,10 @@ class TestGraphRagQueryProvenance:
|
|||
async def explain_callback(triples, explain_id):
|
||||
events.append({"triples": triples, "explain_id": explain_id})
|
||||
|
||||
result_text, usage = await rag.query(
|
||||
result_text, usage, sources = 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 +449,7 @@ class TestGraphRagQueryProvenance:
|
|||
await rag.query(
|
||||
query="What is quantum computing?",
|
||||
explain_callback=explain_callback,
|
||||
edge_score_limit=0,
|
||||
|
||||
parent_uri=parent,
|
||||
)
|
||||
|
||||
|
|
@ -463,9 +463,9 @@ class TestGraphRagQueryProvenance:
|
|||
clients = build_mock_clients()
|
||||
rag = GraphRag(*clients)
|
||||
|
||||
result_text, usage = await rag.query(
|
||||
result_text, usage, sources = await rag.query(
|
||||
query="What is quantum computing?",
|
||||
edge_score_limit=0,
|
||||
|
||||
)
|
||||
|
||||
assert result_text == "Quantum computing applies physics principles to computation."
|
||||
|
|
@ -484,7 +484,7 @@ class TestGraphRagQueryProvenance:
|
|||
await rag.query(
|
||||
query="What is quantum computing?",
|
||||
explain_callback=explain_callback,
|
||||
edge_score_limit=0,
|
||||
|
||||
)
|
||||
|
||||
for event in events:
|
||||
|
|
@ -493,3 +493,165 @@ class TestGraphRagQueryProvenance:
|
|||
f"Triple {t.s.iri} {t.p.iri} should be in "
|
||||
f"urn:graph:retrieval, got {t.g}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Source document tracing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Provenance chains served by the mock triples client:
|
||||
# EDGE_1, EDGE_2 -> SUBGRAPH_A -> chunk/a -> page/a -> DOC_ALPHA
|
||||
# EDGE_3 -> SUBGRAPH_B -> chunk/b -> page/b -> DOC_BETA + DOC_GAMMA
|
||||
SUBGRAPH_A = "http://trustgraph.ai/sg/aaa"
|
||||
SUBGRAPH_B = "http://trustgraph.ai/sg/bbb"
|
||||
DOC_ALPHA = "urn:document:alpha"
|
||||
DOC_BETA = "urn:document:beta"
|
||||
DOC_GAMMA = "urn:document:gamma"
|
||||
TG_MIME_TYPE = "http://trustgraph.ai/ns/provenance/mimeType"
|
||||
|
||||
DERIVATIONS = {
|
||||
SUBGRAPH_A: ["http://trustgraph.ai/chunk/a"],
|
||||
"http://trustgraph.ai/chunk/a": ["http://trustgraph.ai/page/a"],
|
||||
"http://trustgraph.ai/page/a": [DOC_ALPHA],
|
||||
SUBGRAPH_B: ["http://trustgraph.ai/chunk/b"],
|
||||
"http://trustgraph.ai/chunk/b": ["http://trustgraph.ai/page/b"],
|
||||
"http://trustgraph.ai/page/b": [DOC_BETA, DOC_GAMMA],
|
||||
}
|
||||
|
||||
# alpha has both dc:title and rdfs:label (dc:title preferred), beta has
|
||||
# only rdfs:label (fallback), gamma has no title at all (empty string)
|
||||
DOC_METADATA = {
|
||||
DOC_ALPHA: [
|
||||
ClientTriple(Uri(DOC_ALPHA), Uri(RDFS_LABEL),
|
||||
Literal("alpha label")),
|
||||
ClientTriple(Uri(DOC_ALPHA), Uri(DC_TITLE),
|
||||
Literal("Quantum Mechanics Primer")),
|
||||
ClientTriple(Uri(DOC_ALPHA), Uri(TG_MIME_TYPE),
|
||||
Literal("application/pdf")),
|
||||
],
|
||||
DOC_BETA: [
|
||||
ClientTriple(Uri(DOC_BETA), Uri(RDFS_LABEL),
|
||||
Literal("Physics Notes")),
|
||||
],
|
||||
DOC_GAMMA: [
|
||||
ClientTriple(Uri(DOC_GAMMA), Uri(TG_MIME_TYPE),
|
||||
Literal("text/plain")),
|
||||
],
|
||||
}
|
||||
|
||||
EXPECTED_SOURCES = [
|
||||
{"uri": DOC_ALPHA, "title": "Quantum Mechanics Primer"},
|
||||
{"uri": DOC_BETA, "title": "Physics Notes"},
|
||||
{"uri": DOC_GAMMA, "title": ""},
|
||||
]
|
||||
|
||||
# Total triples_client.query calls query() makes against the graph above:
|
||||
# 6 label lookups + 3 tg:contains + 9 wasDerivedFrom + 3 doc metadata.
|
||||
# Sources are built from the same fetches, so this total must not grow.
|
||||
EXPECTED_TRIPLES_QUERY_CALLS = 21
|
||||
|
||||
|
||||
def build_source_tracing_clients(fail_tracing=False):
|
||||
"""Like build_mock_clients, but the triples client also serves the
|
||||
tg:contains + prov:wasDerivedFrom chains and document metadata."""
|
||||
(prompt_client, embeddings_client, graph_embeddings_client,
|
||||
triples_client, reranker_client) = build_mock_clients()
|
||||
|
||||
def subgraph_for(quoted):
|
||||
t = quoted.triple
|
||||
if t.p.iri == "http://schema.org/relatedTo":
|
||||
return SUBGRAPH_A
|
||||
return SUBGRAPH_A if t.s.iri == ENTITY_A else SUBGRAPH_B
|
||||
|
||||
async def mock_query(s=None, p=None, o=None, limit=1,
|
||||
user=None, collection=None, g=None):
|
||||
if p == TG_CONTAINS and o is not None:
|
||||
if fail_tracing:
|
||||
raise RuntimeError("triple store unavailable")
|
||||
sg = subgraph_for(o)
|
||||
return [ClientTriple(Uri(sg), Uri(TG_CONTAINS), o)]
|
||||
if p == PROV_WAS_DERIVED_FROM:
|
||||
return [
|
||||
ClientTriple(Uri(str(s)), Uri(PROV_WAS_DERIVED_FROM),
|
||||
Uri(target))
|
||||
for target in DERIVATIONS.get(str(s), [])
|
||||
]
|
||||
if p is None and str(s) in DOC_METADATA:
|
||||
return DOC_METADATA[str(s)]
|
||||
return [] # Label lookups: fall back to URI
|
||||
|
||||
triples_client.query.side_effect = mock_query
|
||||
|
||||
return (prompt_client, embeddings_client, graph_embeddings_client,
|
||||
triples_client, reranker_client)
|
||||
|
||||
|
||||
class TestGraphRagSourceTracing:
|
||||
"""query() should return structured source references built from the
|
||||
provenance walk it already performs."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_returns_sources(self):
|
||||
"""Sources are deduplicated, uri-sorted, titled where possible."""
|
||||
clients = build_source_tracing_clients()
|
||||
rag = GraphRag(*clients)
|
||||
|
||||
resp, usage, sources = await rag.query(
|
||||
query="What is quantum computing?",
|
||||
)
|
||||
|
||||
assert resp == (
|
||||
"Quantum computing applies physics principles to computation."
|
||||
)
|
||||
assert sources == EXPECTED_SOURCES
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sources_add_zero_triple_queries(self):
|
||||
"""Building sources must not add any triple-store queries."""
|
||||
clients = build_source_tracing_clients()
|
||||
triples_client = clients[3]
|
||||
rag = GraphRag(*clients)
|
||||
|
||||
resp, usage, sources = await rag.query(
|
||||
query="What is quantum computing?",
|
||||
)
|
||||
|
||||
assert sources == EXPECTED_SOURCES
|
||||
assert triples_client.query.call_count == (
|
||||
EXPECTED_TRIPLES_QUERY_CALLS
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_doc_metadata_still_reaches_synthesis_prompt(self):
|
||||
"""The kg-synthesis prompt context keeps the document edges."""
|
||||
clients = build_source_tracing_clients()
|
||||
prompt_client = clients[0]
|
||||
rag = GraphRag(*clients)
|
||||
|
||||
await rag.query(query="What is quantum computing?")
|
||||
|
||||
synthesis_calls = [
|
||||
c for c in prompt_client.prompt.call_args_list
|
||||
if c.args[0] == "kg-synthesis"
|
||||
]
|
||||
assert len(synthesis_calls) == 1
|
||||
knowledge = synthesis_calls[0].kwargs["variables"]["knowledge"]
|
||||
assert {
|
||||
"s": DOC_ALPHA, "p": DC_TITLE,
|
||||
"o": "Quantum Mechanics Primer",
|
||||
} in knowledge
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tracing_failure_degrades_to_empty_sources(self):
|
||||
"""A failing walk yields empty sources, answer unaffected."""
|
||||
clients = build_source_tracing_clients(fail_tracing=True)
|
||||
rag = GraphRag(*clients)
|
||||
|
||||
resp, usage, sources = await rag.query(
|
||||
query="What is quantum computing?",
|
||||
)
|
||||
|
||||
assert resp == (
|
||||
"Quantum computing applies physics principles to computation."
|
||||
)
|
||||
assert sources == []
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import pytest
|
|||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
|
||||
from trustgraph.retrieval.graph_rag.rag import Processor
|
||||
from trustgraph.schema import GraphRagQuery, GraphRagResponse
|
||||
from trustgraph.schema import GraphRagQuery, GraphRagResponse, Source
|
||||
|
||||
|
||||
class TestGraphRagService:
|
||||
|
|
@ -44,7 +44,7 @@ class TestGraphRagService:
|
|||
await explain_callback([], "urn:trustgraph:prov:retrieval:test")
|
||||
await explain_callback([], "urn:trustgraph:prov:selection:test")
|
||||
await explain_callback([], "urn:trustgraph:prov:answer:test")
|
||||
return "A small domesticated mammal.", {"in_token": None, "out_token": None, "model": None}
|
||||
return "A small domesticated mammal.", {"in_token": None, "out_token": None, "model": None}, []
|
||||
|
||||
mock_rag_instance.query.side_effect = mock_query
|
||||
|
||||
|
|
@ -93,6 +93,7 @@ class TestGraphRagService:
|
|||
assert chunk_msg.response == "A small domesticated mammal."
|
||||
assert chunk_msg.end_of_stream is True
|
||||
assert chunk_msg.end_of_session is True
|
||||
assert chunk_msg.sources == []
|
||||
|
||||
# Verify provenance triples were sent to provenance queue
|
||||
assert mock_provenance_producer.send.call_count == 4
|
||||
|
|
@ -180,7 +181,7 @@ class TestGraphRagService:
|
|||
|
||||
async def mock_query(**kwargs):
|
||||
# Don't call explain_callback
|
||||
return "Response text", {"in_token": None, "out_token": None, "model": None}
|
||||
return "Response text", {"in_token": None, "out_token": None, "model": None}, []
|
||||
|
||||
mock_rag_instance.query.side_effect = mock_query
|
||||
|
||||
|
|
@ -219,3 +220,112 @@ class TestGraphRagService:
|
|||
assert chunk_msg.response == "Response text"
|
||||
assert chunk_msg.end_of_stream is True
|
||||
assert chunk_msg.end_of_session is True
|
||||
|
||||
@patch('trustgraph.retrieval.graph_rag.rag.GraphRag')
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_final_message_carries_sources(
|
||||
self, mock_graph_rag_class):
|
||||
"""
|
||||
Test that the non-streaming response carries the source references
|
||||
returned by the query.
|
||||
"""
|
||||
processor = Processor(
|
||||
taskgroup=MagicMock(),
|
||||
id="test-processor",
|
||||
)
|
||||
|
||||
mock_rag_instance = AsyncMock()
|
||||
mock_graph_rag_class.return_value = mock_rag_instance
|
||||
|
||||
async def mock_query(**kwargs):
|
||||
return "Answer.", \
|
||||
{"in_token": None, "out_token": None, "model": None}, \
|
||||
[
|
||||
{"uri": "urn:document:alpha",
|
||||
"title": "Quantum Mechanics Primer"},
|
||||
{"uri": "urn:document:beta", "title": ""},
|
||||
]
|
||||
|
||||
mock_rag_instance.query.side_effect = mock_query
|
||||
|
||||
msg = MagicMock()
|
||||
msg.value.return_value = GraphRagQuery(
|
||||
query="Test query",
|
||||
collection="default",
|
||||
streaming=False
|
||||
)
|
||||
msg.properties.return_value = {"id": "test-id"}
|
||||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
|
||||
mock_response_producer = AsyncMock()
|
||||
flow.side_effect = lambda service_name: mock_response_producer
|
||||
|
||||
# Execute
|
||||
await processor.on_request(msg, consumer, flow)
|
||||
|
||||
# Final (only) message carries the sources
|
||||
chunk_msg = mock_response_producer.send.call_args_list[0][0][0]
|
||||
assert chunk_msg.end_of_session is True
|
||||
assert chunk_msg.sources == [
|
||||
Source(uri="urn:document:alpha",
|
||||
title="Quantum Mechanics Primer"),
|
||||
Source(uri="urn:document:beta", title=""),
|
||||
]
|
||||
|
||||
@patch('trustgraph.retrieval.graph_rag.rag.GraphRag')
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_final_message_carries_sources(
|
||||
self, mock_graph_rag_class):
|
||||
"""
|
||||
Test that in streaming mode only the final end_of_session message
|
||||
carries the source references.
|
||||
"""
|
||||
processor = Processor(
|
||||
taskgroup=MagicMock(),
|
||||
id="test-processor",
|
||||
)
|
||||
|
||||
mock_rag_instance = AsyncMock()
|
||||
mock_graph_rag_class.return_value = mock_rag_instance
|
||||
|
||||
async def mock_query(**kwargs):
|
||||
chunk_callback = kwargs.get('chunk_callback')
|
||||
await chunk_callback("Streamed answer.", True)
|
||||
return "Streamed answer.", \
|
||||
{"in_token": None, "out_token": None, "model": None}, \
|
||||
[{"uri": "urn:document:alpha", "title": "Primer"}]
|
||||
|
||||
mock_rag_instance.query.side_effect = mock_query
|
||||
|
||||
msg = MagicMock()
|
||||
msg.value.return_value = GraphRagQuery(
|
||||
query="Test query",
|
||||
collection="default",
|
||||
streaming=True
|
||||
)
|
||||
msg.properties.return_value = {"id": "test-id"}
|
||||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
|
||||
mock_response_producer = AsyncMock()
|
||||
flow.side_effect = lambda service_name: mock_response_producer
|
||||
|
||||
# Execute
|
||||
await processor.on_request(msg, consumer, flow)
|
||||
|
||||
# 2 messages: streamed chunk, then end_of_session close
|
||||
assert mock_response_producer.send.call_count == 2
|
||||
|
||||
chunk_msg = mock_response_producer.send.call_args_list[0][0][0]
|
||||
assert chunk_msg.end_of_session is False
|
||||
assert chunk_msg.sources == []
|
||||
|
||||
final_msg = mock_response_producer.send.call_args_list[1][0][0]
|
||||
assert final_msg.end_of_session is True
|
||||
assert final_msg.sources == [
|
||||
Source(uri="urn:document:alpha", title="Primer"),
|
||||
]
|
||||
|
||||
|
|
|
|||
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") == []
|
||||
|
|
@ -0,0 +1,130 @@
|
|||
"""
|
||||
Round-trip unit tests for ImageToTextRequestTranslator and
|
||||
ImageToTextResponseTranslator.
|
||||
|
||||
The image field carries base64 text end-to-end (raw binary can't ride
|
||||
the JSON wire format), so the request decode is THE validation point
|
||||
for image payloads entering the system: invalid base64 must be
|
||||
rejected at the gateway, before anything is queued.
|
||||
|
||||
Image-to-text is non-streaming, so encode_with_completion must always
|
||||
report the response as final.
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
|
||||
from trustgraph.messaging.translators.image_to_text import (
|
||||
ImageToTextRequestTranslator,
|
||||
ImageToTextResponseTranslator,
|
||||
)
|
||||
from trustgraph.schema import (
|
||||
ImageToTextRequest,
|
||||
ImageToTextResponse,
|
||||
)
|
||||
|
||||
|
||||
IMAGE_BYTES = b"\x89PNG\r\n\x1a\nfake-image-payload"
|
||||
IMAGE_B64 = base64.b64encode(IMAGE_BYTES).decode("utf-8")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def request_translator():
|
||||
return ImageToTextRequestTranslator()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def response_translator():
|
||||
return ImageToTextResponseTranslator()
|
||||
|
||||
|
||||
class TestImageToTextRequestTranslator:
|
||||
|
||||
def test_decode_full_request(self, request_translator):
|
||||
decoded = request_translator.decode({
|
||||
"image": IMAGE_B64,
|
||||
"mime_type": "image/png",
|
||||
"prompt": "What is shown here?",
|
||||
"system": "You are an art critic",
|
||||
})
|
||||
|
||||
assert isinstance(decoded, ImageToTextRequest)
|
||||
assert decoded.image == IMAGE_B64
|
||||
assert decoded.mime_type == "image/png"
|
||||
assert decoded.prompt == "What is shown here?"
|
||||
assert decoded.system == "You are an art critic"
|
||||
|
||||
def test_decode_defaults_optional_fields(self, request_translator):
|
||||
"""prompt/system are optional; the backend supplies the default prompt."""
|
||||
decoded = request_translator.decode({
|
||||
"image": IMAGE_B64,
|
||||
"mime_type": "image/jpeg",
|
||||
})
|
||||
|
||||
assert decoded.prompt == ""
|
||||
assert decoded.system == ""
|
||||
|
||||
def test_decode_rejects_invalid_base64(self, request_translator):
|
||||
with pytest.raises(ValueError):
|
||||
request_translator.decode({
|
||||
"image": "this is !!! not *** base64",
|
||||
"mime_type": "image/png",
|
||||
})
|
||||
|
||||
def test_roundtrip_is_lossless(self, request_translator):
|
||||
request = ImageToTextRequest(
|
||||
image=IMAGE_B64,
|
||||
mime_type="image/png",
|
||||
prompt="Describe this image",
|
||||
system="Be terse",
|
||||
)
|
||||
|
||||
encoded = request_translator.encode(request)
|
||||
decoded = request_translator.decode(encoded)
|
||||
|
||||
assert decoded.image == IMAGE_B64
|
||||
assert decoded.mime_type == "image/png"
|
||||
assert decoded.prompt == "Describe this image"
|
||||
assert decoded.system == "Be terse"
|
||||
|
||||
|
||||
class TestImageToTextResponseTranslator:
|
||||
|
||||
def test_encode_full_response(self, response_translator):
|
||||
response = ImageToTextResponse(
|
||||
error=None,
|
||||
description="A cat sitting on a mat",
|
||||
in_token=100,
|
||||
out_token=20,
|
||||
model="test-model",
|
||||
)
|
||||
|
||||
encoded = response_translator.encode(response)
|
||||
|
||||
assert encoded == {
|
||||
"description": "A cat sitting on a mat",
|
||||
"in_token": 100,
|
||||
"out_token": 20,
|
||||
"model": "test-model",
|
||||
}
|
||||
|
||||
def test_encode_omits_absent_token_fields(self, response_translator):
|
||||
response = ImageToTextResponse(description="A dog")
|
||||
|
||||
encoded = response_translator.encode(response)
|
||||
|
||||
assert encoded == {"description": "A dog"}
|
||||
|
||||
def test_encode_with_completion_always_final(self, response_translator):
|
||||
"""Image-to-text is non-streaming: every response is final."""
|
||||
response = ImageToTextResponse(description="A dog")
|
||||
|
||||
result, is_final = response_translator.encode_with_completion(response)
|
||||
|
||||
assert result == {"description": "A dog"}
|
||||
assert is_final is True
|
||||
|
||||
def test_decode_not_implemented(self, response_translator):
|
||||
with pytest.raises(NotImplementedError):
|
||||
response_translator.decode({"description": "A dog"})
|
||||
|
|
@ -107,6 +107,7 @@ from .types import (
|
|||
AgentAnswer,
|
||||
RAGChunk,
|
||||
TextCompletionResult,
|
||||
ImageToTextResult,
|
||||
ProvenanceEvent,
|
||||
)
|
||||
|
||||
|
|
@ -186,6 +187,7 @@ __all__ = [
|
|||
"AgentAnswer",
|
||||
"RAGChunk",
|
||||
"TextCompletionResult",
|
||||
"ImageToTextResult",
|
||||
"ProvenanceEvent",
|
||||
|
||||
# Exceptions
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ def check_error(response):
|
|||
try:
|
||||
msg = response["error"]["message"]
|
||||
tp = response["error"]["type"]
|
||||
except:
|
||||
except KeyError:
|
||||
raise ApplicationException(response["error"])
|
||||
|
||||
raise ApplicationException(f"{tp}: {msg}")
|
||||
|
|
@ -208,7 +208,7 @@ class Api:
|
|||
try:
|
||||
# Parse the response as JSON
|
||||
object = resp.json()
|
||||
except:
|
||||
except ValueError:
|
||||
raise ProtocolException(f"Expected JSON response")
|
||||
|
||||
check_error(object)
|
||||
|
|
|
|||
|
|
@ -12,9 +12,10 @@ AsyncSocketClient instead.
|
|||
|
||||
import aiohttp
|
||||
import json
|
||||
import base64
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
from . types import TextCompletionResult
|
||||
from . types import TextCompletionResult, ImageToTextResult
|
||||
|
||||
from . exceptions import ProtocolException, ApplicationException
|
||||
|
||||
|
|
@ -24,7 +25,7 @@ def check_error(response):
|
|||
try:
|
||||
msg = response["error"]["message"]
|
||||
tp = response["error"]["type"]
|
||||
except:
|
||||
except KeyError:
|
||||
raise ApplicationException(response["error"])
|
||||
|
||||
raise ApplicationException(f"{tp}: {msg}")
|
||||
|
|
@ -87,7 +88,7 @@ class AsyncFlow:
|
|||
|
||||
try:
|
||||
obj = await resp.json()
|
||||
except:
|
||||
except (ValueError, aiohttp.ContentTypeError):
|
||||
raise ProtocolException(f"Expected JSON response")
|
||||
|
||||
check_error(obj)
|
||||
|
|
@ -476,6 +477,56 @@ class AsyncFlowInstance:
|
|||
model=result.get("model"),
|
||||
)
|
||||
|
||||
async def image_to_text(self, image: bytes, mime_type: str,
|
||||
prompt: Optional[str] = None,
|
||||
system: Optional[str] = None,
|
||||
**kwargs: Any) -> ImageToTextResult:
|
||||
"""
|
||||
Describe an image using the image-to-text service (non-streaming).
|
||||
|
||||
Args:
|
||||
image: Image content as bytes
|
||||
mime_type: Image MIME type (e.g. "image/jpeg")
|
||||
prompt: Optional user prompt (backend default used if None)
|
||||
system: Optional system prompt
|
||||
**kwargs: Additional service-specific parameters
|
||||
|
||||
Returns:
|
||||
ImageToTextResult: Result with text, in_token, out_token, model
|
||||
|
||||
Example:
|
||||
```python
|
||||
async_flow = await api.async_flow()
|
||||
flow = async_flow.id("default")
|
||||
|
||||
with open("photo.jpg", "rb") as f:
|
||||
result = await flow.image_to_text(
|
||||
image=f.read(),
|
||||
mime_type="image/jpeg",
|
||||
)
|
||||
print(result.text)
|
||||
print(f"Tokens: {result.in_token} in, {result.out_token} out")
|
||||
```
|
||||
"""
|
||||
# The image rides the JSON wire format as base64 text
|
||||
request_data = {
|
||||
"image": base64.b64encode(image).decode("utf-8"),
|
||||
"mime_type": mime_type,
|
||||
}
|
||||
if prompt is not None:
|
||||
request_data["prompt"] = prompt
|
||||
if system is not None:
|
||||
request_data["system"] = system
|
||||
request_data.update(kwargs)
|
||||
|
||||
result = await self.request("image-to-text", request_data)
|
||||
return ImageToTextResult(
|
||||
text=result.get("description", ""),
|
||||
in_token=result.get("in_token"),
|
||||
out_token=result.get("out_token"),
|
||||
model=result.get("model"),
|
||||
)
|
||||
|
||||
async def graph_rag(self, query: str, collection: str,
|
||||
max_subgraph_size: int = 1000, max_subgraph_count: int = 5,
|
||||
max_entity_distance: int = 3, **kwargs: Any) -> str:
|
||||
|
|
@ -527,7 +578,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 +593,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 +618,7 @@ class AsyncFlowInstance:
|
|||
"query": query,
|
||||
"collection": collection,
|
||||
"doc-limit": doc_limit,
|
||||
"fetch-limit": fetch_limit,
|
||||
"streaming": False
|
||||
}
|
||||
request_data.update(kwargs)
|
||||
|
|
@ -646,6 +701,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.
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
|
||||
import json
|
||||
import base64
|
||||
import asyncio
|
||||
import websockets
|
||||
from typing import Optional, Dict, Any, AsyncIterator, Union
|
||||
|
||||
from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, TextCompletionResult
|
||||
from . exceptions import ProtocolException, ApplicationException
|
||||
from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, TextCompletionResult, ImageToTextResult
|
||||
from . exceptions import ProtocolException, ApplicationException, raise_from_error_dict
|
||||
|
||||
|
||||
class AsyncSocketClient:
|
||||
|
|
@ -94,7 +95,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(
|
||||
|
|
@ -139,7 +142,7 @@ class AsyncSocketClient:
|
|||
for queue in self._pending.values():
|
||||
try:
|
||||
await queue.put({"error": str(e)})
|
||||
except:
|
||||
except asyncio.CancelledError :
|
||||
pass
|
||||
finally:
|
||||
self._connected = False
|
||||
|
|
@ -265,6 +268,7 @@ class AsyncSocketClient:
|
|||
in_token=resp.get("in_token"),
|
||||
out_token=resp.get("out_token"),
|
||||
model=resp.get("model"),
|
||||
sources=resp.get("sources", []),
|
||||
)
|
||||
|
||||
async def aclose(self):
|
||||
|
|
@ -350,6 +354,38 @@ class AsyncSocketFlowInstance:
|
|||
if isinstance(chunk, RAGChunk):
|
||||
yield chunk
|
||||
|
||||
async def image_to_text(self, image: bytes, mime_type: str,
|
||||
prompt: Optional[str] = None,
|
||||
system: Optional[str] = None,
|
||||
**kwargs) -> ImageToTextResult:
|
||||
"""Describe an image using the image-to-text service (non-streaming).
|
||||
|
||||
Returns an ImageToTextResult with the description text and token counts.
|
||||
"""
|
||||
# The image rides the JSON wire format as base64 text
|
||||
request = {
|
||||
"image": base64.b64encode(image).decode("utf-8"),
|
||||
"mime_type": mime_type,
|
||||
}
|
||||
if prompt is not None:
|
||||
request["prompt"] = prompt
|
||||
if system is not None:
|
||||
request["system"] = system
|
||||
request.update(kwargs)
|
||||
|
||||
result = await self.client._send_request("image-to-text", self.flow_id, request)
|
||||
|
||||
# Service errors arrive inside the response body
|
||||
if isinstance(result, dict) and result.get("error"):
|
||||
raise_from_error_dict(result["error"])
|
||||
|
||||
return ImageToTextResult(
|
||||
text=result.get("description", ""),
|
||||
in_token=result.get("in_token"),
|
||||
out_token=result.get("out_token"),
|
||||
model=result.get("model"),
|
||||
)
|
||||
|
||||
async def graph_rag(self, query: str, collection: str,
|
||||
max_subgraph_size: int = 1000, max_subgraph_count: int = 5,
|
||||
max_entity_distance: int = 3, streaming: bool = False, **kwargs):
|
||||
|
|
@ -377,12 +413,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 +479,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}
|
||||
|
|
|
|||
|
|
@ -201,7 +201,7 @@ class BulkClient:
|
|||
finally:
|
||||
try:
|
||||
loop.run_until_complete(async_gen.aclose())
|
||||
except:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _export_triples_async(self, flow: str) -> Iterator[Triple]:
|
||||
|
|
@ -299,7 +299,7 @@ class BulkClient:
|
|||
finally:
|
||||
try:
|
||||
loop.run_until_complete(async_gen.aclose())
|
||||
except:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _export_graph_embeddings_async(self, flow: str) -> Iterator[Dict[str, Any]]:
|
||||
|
|
@ -393,7 +393,7 @@ class BulkClient:
|
|||
finally:
|
||||
try:
|
||||
loop.run_until_complete(async_gen.aclose())
|
||||
except:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _export_document_embeddings_async(self, flow: str) -> Iterator[Dict[str, Any]]:
|
||||
|
|
@ -517,7 +517,7 @@ class BulkClient:
|
|||
finally:
|
||||
try:
|
||||
loop.run_until_complete(async_gen.aclose())
|
||||
except:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _export_entity_contexts_async(self, flow: str) -> Iterator[Dict[str, Any]]:
|
||||
|
|
|
|||
|
|
@ -254,7 +254,7 @@ class Config:
|
|||
)
|
||||
for v in object["values"]
|
||||
]
|
||||
except:
|
||||
except (KeyError, TypeError):
|
||||
raise ProtocolException(f"Response not formatted correctly")
|
||||
|
||||
def get_values_all_workspaces(self, type):
|
||||
|
|
@ -330,6 +330,6 @@ class Config:
|
|||
|
||||
try:
|
||||
return object["config"], object["version"]
|
||||
except:
|
||||
except KeyError:
|
||||
raise ProtocolException(f"Response not formatted correctly")
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import base64
|
|||
|
||||
from .. knowledge import hash, Uri, Literal, QuotedTriple
|
||||
from .. schema import IRI, LITERAL, TRIPLE
|
||||
from . types import Triple, TextCompletionResult
|
||||
from . types import Triple, TextCompletionResult, ImageToTextResult
|
||||
from . exceptions import ProtocolException
|
||||
|
||||
|
||||
|
|
@ -296,6 +296,54 @@ class FlowInstance:
|
|||
model=result.get("model"),
|
||||
)
|
||||
|
||||
def image_to_text(self, image, mime_type, prompt=None, system=None):
|
||||
"""
|
||||
Describe an image using the flow's image-to-text service.
|
||||
|
||||
Args:
|
||||
image: Image content as bytes
|
||||
mime_type: Image MIME type (e.g. "image/jpeg")
|
||||
prompt: Optional user prompt (backend default used if None)
|
||||
system: Optional system prompt
|
||||
|
||||
Returns:
|
||||
ImageToTextResult: Result with text, in_token, out_token, model
|
||||
|
||||
Example:
|
||||
```python
|
||||
flow = api.flow().id("default")
|
||||
with open("photo.jpg", "rb") as f:
|
||||
result = flow.image_to_text(
|
||||
image=f.read(),
|
||||
mime_type="image/jpeg",
|
||||
)
|
||||
print(result.text)
|
||||
print(f"Tokens: {result.in_token} in, {result.out_token} out")
|
||||
```
|
||||
"""
|
||||
|
||||
# The image rides the JSON wire format as base64 text
|
||||
input = {
|
||||
"image": base64.b64encode(image).decode("utf-8"),
|
||||
"mime_type": mime_type,
|
||||
}
|
||||
if prompt is not None:
|
||||
input["prompt"] = prompt
|
||||
if system is not None:
|
||||
input["system"] = system
|
||||
|
||||
result = self.request(
|
||||
"service/image-to-text",
|
||||
input
|
||||
)
|
||||
|
||||
return ImageToTextResult(
|
||||
text=result.get("description", ""),
|
||||
in_token=result.get("in_token"),
|
||||
out_token=result.get("out_token"),
|
||||
model=result.get("model"),
|
||||
)
|
||||
|
||||
def agent(self, question,state=None, group=None, history=None):
|
||||
"""
|
||||
Execute an agent operation with reasoning and tool use capabilities.
|
||||
|
|
@ -357,6 +405,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 +422,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 +449,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(
|
||||
|
|
@ -411,11 +462,12 @@ class FlowInstance:
|
|||
in_token=result.get("in_token"),
|
||||
out_token=result.get("out_token"),
|
||||
model=result.get("model"),
|
||||
sources=result.get("sources", []),
|
||||
)
|
||||
|
||||
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 +478,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 +501,7 @@ class FlowInstance:
|
|||
"query": query,
|
||||
"collection": collection,
|
||||
"doc-limit": doc_limit,
|
||||
"fetch-limit": fetch_limit,
|
||||
}
|
||||
|
||||
result = self.request(
|
||||
|
|
@ -491,6 +546,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", ""),
|
||||
|
|
|
|||
|
|
@ -9,13 +9,14 @@ multiplexes requests by ID.
|
|||
"""
|
||||
|
||||
import json
|
||||
import base64
|
||||
import asyncio
|
||||
import websockets
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
from typing import Optional, Dict, Any, Iterator, Union, List
|
||||
from threading import Lock
|
||||
|
||||
from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, StreamingChunk, ProvenanceEvent, TextCompletionResult
|
||||
from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, StreamingChunk, ProvenanceEvent, TextCompletionResult, ImageToTextResult
|
||||
from . exceptions import ProtocolException, raise_from_error_dict
|
||||
|
||||
|
||||
|
|
@ -168,7 +169,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(
|
||||
|
|
@ -449,6 +452,7 @@ class SocketClient:
|
|||
in_token=resp.get("in_token"),
|
||||
out_token=resp.get("out_token"),
|
||||
model=resp.get("model"),
|
||||
sources=resp.get("sources", []),
|
||||
)
|
||||
|
||||
def _build_provenance_event(self, resp: Dict[str, Any]) -> ProvenanceEvent:
|
||||
|
|
@ -670,6 +674,38 @@ class SocketFlowInstance:
|
|||
if isinstance(chunk, RAGChunk):
|
||||
yield chunk
|
||||
|
||||
def image_to_text(self, image: bytes, mime_type: str,
|
||||
prompt: Optional[str] = None,
|
||||
system: Optional[str] = None,
|
||||
**kwargs: Any) -> ImageToTextResult:
|
||||
"""Describe an image using the image-to-text service (non-streaming).
|
||||
|
||||
Returns an ImageToTextResult with the description text and token counts.
|
||||
"""
|
||||
# The image rides the JSON wire format as base64 text
|
||||
request = {
|
||||
"image": base64.b64encode(image).decode("utf-8"),
|
||||
"mime_type": mime_type,
|
||||
}
|
||||
if prompt is not None:
|
||||
request["prompt"] = prompt
|
||||
if system is not None:
|
||||
request["system"] = system
|
||||
request.update(kwargs)
|
||||
|
||||
result = self.client._send_request_sync("image-to-text", self.flow_id, request, False)
|
||||
|
||||
# Service errors arrive inside the response body
|
||||
if isinstance(result, dict) and result.get("error"):
|
||||
raise_from_error_dict(result["error"])
|
||||
|
||||
return ImageToTextResult(
|
||||
text=result.get("description", ""),
|
||||
in_token=result.get("in_token"),
|
||||
out_token=result.get("out_token"),
|
||||
model=result.get("model"),
|
||||
)
|
||||
|
||||
def graph_rag(
|
||||
self,
|
||||
query: str,
|
||||
|
|
@ -680,6 +716,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 +734,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)
|
||||
|
|
@ -711,6 +749,7 @@ class SocketFlowInstance:
|
|||
in_token=result.get("in_token"),
|
||||
out_token=result.get("out_token"),
|
||||
model=result.get("model"),
|
||||
sources=result.get("sources", []),
|
||||
)
|
||||
|
||||
def graph_rag_explain(
|
||||
|
|
@ -723,6 +762,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 +775,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 +791,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 +804,7 @@ class SocketFlowInstance:
|
|||
"query": query,
|
||||
"collection": collection,
|
||||
"doc-limit": doc_limit,
|
||||
"fetch-limit": fetch_limit,
|
||||
"streaming": streaming
|
||||
}
|
||||
request.update(kwargs)
|
||||
|
|
@ -783,6 +826,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 +834,7 @@ class SocketFlowInstance:
|
|||
"query": query,
|
||||
"collection": collection,
|
||||
"doc-limit": doc_limit,
|
||||
"fetch-limit": fetch_limit,
|
||||
"streaming": True,
|
||||
"explainable": True,
|
||||
}
|
||||
|
|
@ -883,6 +928,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,
|
||||
|
|
|
|||
|
|
@ -205,6 +205,8 @@ class RAGChunk(StreamingChunk):
|
|||
in_token: Input token count (populated on the final chunk, 0 otherwise)
|
||||
out_token: Output token count (populated on the final chunk, 0 otherwise)
|
||||
model: Model identifier (populated on the final chunk, empty otherwise)
|
||||
sources: Source document references as uri/title dicts (populated
|
||||
on the final chunk, empty otherwise)
|
||||
message_type: Always "rag"
|
||||
"""
|
||||
message_type: str = "rag"
|
||||
|
|
@ -213,6 +215,7 @@ class RAGChunk(StreamingChunk):
|
|||
in_token: Optional[int] = None
|
||||
out_token: Optional[int] = None
|
||||
model: Optional[str] = None
|
||||
sources: List[Dict[str, str]] = dataclasses.field(default_factory=list)
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TextCompletionResult:
|
||||
|
|
@ -228,6 +231,27 @@ class TextCompletionResult:
|
|||
in_token: Input token count (None if not available)
|
||||
out_token: Output token count (None if not available)
|
||||
model: Model identifier (None if not available)
|
||||
sources: Source document references as uri/title dicts (graph RAG
|
||||
only, empty otherwise)
|
||||
"""
|
||||
text: Optional[str]
|
||||
in_token: Optional[int] = None
|
||||
out_token: Optional[int] = None
|
||||
model: Optional[str] = None
|
||||
sources: List[Dict[str, str]] = dataclasses.field(default_factory=list)
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ImageToTextResult:
|
||||
"""
|
||||
Result from an image-to-text request.
|
||||
|
||||
Returned by image_to_text(). Non-streaming only.
|
||||
|
||||
Attributes:
|
||||
text: Text description of the image
|
||||
in_token: Input token count (None if not available)
|
||||
out_token: Output token count (None if not available)
|
||||
model: Model identifier (None if not available)
|
||||
"""
|
||||
text: Optional[str]
|
||||
in_token: Optional[int] = None
|
||||
|
|
|
|||
|
|
@ -42,6 +42,13 @@ 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 . image_to_text_service import ImageToTextService, ImageDescriptionResult
|
||||
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
|
||||
from . schema_compatibility import is_strict_mode_compatible
|
||||
|
||||
|
|
|
|||
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}")
|
||||
|
|
@ -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
|
||||
|
|
|
|||
170
trustgraph-base/trustgraph/base/image_to_text_service.py
Normal file
170
trustgraph-base/trustgraph/base/image_to_text_service.py
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
"""
|
||||
Image-to-text description base class
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from argparse import ArgumentParser
|
||||
|
||||
import logging
|
||||
from prometheus_client import Histogram, Info
|
||||
|
||||
from .. schema import ImageToTextRequest, ImageToTextResponse, Error
|
||||
from .. exceptions import TooManyRequests
|
||||
from .. base import FlowProcessor, ConsumerSpec, ProducerSpec, ParameterSpec
|
||||
|
||||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
default_ident = "image-to-text"
|
||||
default_concurrency = 1
|
||||
|
||||
class ImageDescriptionResult:
|
||||
def __init__(
|
||||
self, text = None, in_token = None, out_token = None,
|
||||
model = None,
|
||||
):
|
||||
self.text = text
|
||||
self.in_token = in_token
|
||||
self.out_token = out_token
|
||||
self.model = model
|
||||
__slots__ = ["text", "in_token", "out_token", "model"]
|
||||
|
||||
class ImageToTextService(FlowProcessor):
|
||||
"""
|
||||
Extensible service processing image description requests.
|
||||
|
||||
This class handles the core logic of dispatching image-to-text
|
||||
requests to integrated underlying vision model providers
|
||||
(e.g. OpenAI).
|
||||
"""
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
id = params.get("id", default_ident)
|
||||
concurrency = params.get("concurrency", 1)
|
||||
|
||||
super(ImageToTextService, self).__init__(**params | {
|
||||
"id": id,
|
||||
"concurrency": concurrency,
|
||||
})
|
||||
|
||||
self.register_specification(
|
||||
ConsumerSpec(
|
||||
name = "request",
|
||||
schema = ImageToTextRequest,
|
||||
handler = self.on_request,
|
||||
concurrency = concurrency,
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ProducerSpec(
|
||||
name = "response",
|
||||
schema = ImageToTextResponse
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ParameterSpec(
|
||||
name = "model",
|
||||
)
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "image_to_text_metric"):
|
||||
__class__.image_to_text_metric = Histogram(
|
||||
'image_to_text_duration',
|
||||
'Image-to-text duration (seconds)',
|
||||
["id", "flow"],
|
||||
buckets=[
|
||||
0.25, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
|
||||
8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
|
||||
17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0, 25.0,
|
||||
30.0, 35.0, 40.0, 45.0, 50.0, 60.0, 80.0, 100.0,
|
||||
120.0
|
||||
]
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "image_to_text_model_metric"):
|
||||
__class__.image_to_text_model_metric = Info(
|
||||
'image_to_text_model',
|
||||
'Image-to-text model',
|
||||
["processor", "flow"]
|
||||
)
|
||||
|
||||
async def on_request(self, msg, consumer, flow):
|
||||
|
||||
try:
|
||||
|
||||
request = msg.value()
|
||||
|
||||
# Sender-produced ID
|
||||
|
||||
id = msg.properties()["id"]
|
||||
|
||||
model = flow("model")
|
||||
|
||||
with __class__.image_to_text_metric.labels(
|
||||
id=self.id,
|
||||
flow=f"{flow.name}-{consumer.name}",
|
||||
).time():
|
||||
|
||||
response = await self.describe_image(
|
||||
request.image, request.mime_type,
|
||||
request.prompt, request.system, model,
|
||||
)
|
||||
|
||||
await flow("response").send(
|
||||
ImageToTextResponse(
|
||||
error=None,
|
||||
description=response.text,
|
||||
in_token=response.in_token,
|
||||
out_token=response.out_token,
|
||||
model=response.model,
|
||||
),
|
||||
properties={"id": id}
|
||||
)
|
||||
|
||||
__class__.image_to_text_model_metric.labels(
|
||||
processor = self.id,
|
||||
flow = flow.name
|
||||
).info({
|
||||
"model": str(model) if model is not None else "",
|
||||
})
|
||||
|
||||
except TooManyRequests as e:
|
||||
raise e
|
||||
|
||||
except Exception as e:
|
||||
|
||||
# Apart from rate limits, treat all exceptions as unrecoverable
|
||||
|
||||
logger.error(f"Image-to-text service exception: {e}", exc_info=True)
|
||||
|
||||
logger.debug("Sending error response...")
|
||||
|
||||
await flow.producer["response"].send(
|
||||
ImageToTextResponse(
|
||||
error=Error(
|
||||
type = "image-to-text-error",
|
||||
message = str(e),
|
||||
),
|
||||
description=None,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=None,
|
||||
),
|
||||
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)
|
||||
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})'
|
||||
)
|
||||
|
|
@ -126,6 +126,8 @@ class LlmService(FlowProcessor):
|
|||
|
||||
# Check if streaming is requested and supported
|
||||
streaming = getattr(request, 'streaming', False)
|
||||
response_format = getattr(request, 'response_format', None)
|
||||
schema = getattr(request, 'schema', None)
|
||||
|
||||
if streaming and self.supports_streaming():
|
||||
|
||||
|
|
@ -136,7 +138,8 @@ class LlmService(FlowProcessor):
|
|||
).time():
|
||||
|
||||
async for chunk in self.generate_content_stream(
|
||||
request.system, request.prompt, model, temperature
|
||||
request.system, request.prompt, model, temperature,
|
||||
response_format=response_format, schema=schema,
|
||||
):
|
||||
await flow("response").send(
|
||||
TextCompletionResponse(
|
||||
|
|
@ -159,7 +162,8 @@ class LlmService(FlowProcessor):
|
|||
).time():
|
||||
|
||||
response = await self.generate_content(
|
||||
request.system, request.prompt, model, temperature
|
||||
request.system, request.prompt, model, temperature,
|
||||
response_format=response_format, schema=schema,
|
||||
)
|
||||
|
||||
await flow("response").send(
|
||||
|
|
@ -215,7 +219,10 @@ class LlmService(FlowProcessor):
|
|||
"""
|
||||
return False
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
async def generate_content_stream(
|
||||
self, system, prompt, model=None, temperature=None,
|
||||
response_format=None, schema=None,
|
||||
):
|
||||
"""
|
||||
Override in subclass to implement streaming.
|
||||
Should yield LlmChunk objects.
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
90
trustgraph-base/trustgraph/base/schema_compatibility.py
Normal file
90
trustgraph-base/trustgraph/base/schema_compatibility.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def is_strict_mode_compatible(schema):
|
||||
"""
|
||||
Check whether a JSON schema is compatible with LLM structured-output
|
||||
strict mode. Returns True if the schema can be passed directly to
|
||||
providers like OpenAI, vLLM, etc.
|
||||
"""
|
||||
|
||||
if schema is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
_check_node(schema)
|
||||
return True
|
||||
except _IncompatibleSchema as e:
|
||||
logger.debug("Schema not strict-mode compatible: %s", e)
|
||||
return False
|
||||
|
||||
|
||||
class _IncompatibleSchema(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _check_node(node):
|
||||
|
||||
if not isinstance(node, dict):
|
||||
return
|
||||
|
||||
node_type = node.get("type")
|
||||
|
||||
if node_type == "object" or (
|
||||
node_type is None and "properties" in node
|
||||
):
|
||||
_check_object(node)
|
||||
|
||||
if node_type == "array":
|
||||
items = node.get("items")
|
||||
if items:
|
||||
_check_node(items)
|
||||
|
||||
for keyword in ("oneOf", "anyOf", "allOf"):
|
||||
for child in node.get(keyword, []):
|
||||
_check_node(child)
|
||||
|
||||
_check_unsupported_constraints(node)
|
||||
|
||||
|
||||
def _check_object(node):
|
||||
|
||||
props = node.get("properties")
|
||||
if props is None:
|
||||
raise _IncompatibleSchema(
|
||||
"object without properties (open-ended)"
|
||||
)
|
||||
|
||||
if node.get("additionalProperties") is not False:
|
||||
raise _IncompatibleSchema(
|
||||
"object missing additionalProperties: false"
|
||||
)
|
||||
|
||||
required = set(node.get("required", []))
|
||||
for key in props:
|
||||
if key not in required:
|
||||
raise _IncompatibleSchema(
|
||||
f"property '{key}' not in required"
|
||||
)
|
||||
|
||||
for value in props.values():
|
||||
_check_node(value)
|
||||
|
||||
|
||||
UNSUPPORTED_KEYWORDS = {
|
||||
"minimum", "maximum", "exclusiveMinimum", "exclusiveMaximum",
|
||||
"minLength", "maxLength", "pattern",
|
||||
"minItems", "maxItems",
|
||||
"minProperties", "maxProperties",
|
||||
}
|
||||
|
||||
|
||||
def _check_unsupported_constraints(node):
|
||||
found = UNSUPPORTED_KEYWORDS & node.keys()
|
||||
if found:
|
||||
raise _IncompatibleSchema(
|
||||
f"unsupported constraints: {', '.join(sorted(found))}"
|
||||
)
|
||||
|
|
@ -14,11 +14,15 @@ class TextCompletionResult:
|
|||
|
||||
class TextCompletionClient(RequestResponse):
|
||||
|
||||
async def text_completion(self, system, prompt, timeout=600):
|
||||
async def text_completion(
|
||||
self, system, prompt, timeout=600,
|
||||
response_format=None, schema=None,
|
||||
):
|
||||
|
||||
resp = await self.request(
|
||||
TextCompletionRequest(
|
||||
system = system, prompt = prompt, streaming = False
|
||||
system=system, prompt=prompt, streaming=False,
|
||||
response_format=response_format, schema=schema,
|
||||
),
|
||||
timeout=timeout
|
||||
)
|
||||
|
|
@ -35,6 +39,7 @@ class TextCompletionClient(RequestResponse):
|
|||
|
||||
async def text_completion_stream(
|
||||
self, system, prompt, handler, timeout=600,
|
||||
response_format=None, schema=None,
|
||||
):
|
||||
"""
|
||||
Streaming text completion. `handler` is an async callable invoked
|
||||
|
|
@ -54,7 +59,8 @@ class TextCompletionClient(RequestResponse):
|
|||
|
||||
final = await self.request(
|
||||
TextCompletionRequest(
|
||||
system = system, prompt = prompt, streaming = True
|
||||
system=system, prompt=prompt, streaming=True,
|
||||
response_format=response_format, schema=schema,
|
||||
),
|
||||
recipient=on_chunk,
|
||||
timeout=timeout,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from .translators import *
|
|||
from .translators.agent import AgentRequestTranslator, AgentResponseTranslator
|
||||
from .translators.embeddings import EmbeddingsRequestTranslator, EmbeddingsResponseTranslator
|
||||
from .translators.text_completion import TextCompletionRequestTranslator, TextCompletionResponseTranslator
|
||||
from .translators.image_to_text import ImageToTextRequestTranslator, ImageToTextResponseTranslator
|
||||
from .translators.retrieval import (
|
||||
DocumentRagRequestTranslator, DocumentRagResponseTranslator,
|
||||
GraphRagRequestTranslator, GraphRagResponseTranslator
|
||||
|
|
@ -27,6 +28,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
|
||||
|
||||
|
|
@ -49,6 +51,12 @@ TranslatorRegistry.register_service(
|
|||
TextCompletionResponseTranslator()
|
||||
)
|
||||
|
||||
TranslatorRegistry.register_service(
|
||||
"image-to-text",
|
||||
ImageToTextRequestTranslator(),
|
||||
ImageToTextResponseTranslator()
|
||||
)
|
||||
|
||||
TranslatorRegistry.register_service(
|
||||
"document-rag",
|
||||
DocumentRagRequestTranslator(),
|
||||
|
|
@ -163,6 +171,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(
|
||||
|
|
|
|||
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