diff --git a/README.md b/README.md index b66edc70..c366a3d9 100644 --- a/README.md +++ b/README.md @@ -11,11 +11,11 @@ trustgraph-ai%2Ftrustgraph | Trendshift -# The semantic deployment platform +# The agent runtime platform -TrustGraph is a comprehensive semantic infrastructure for agents 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 deterministic agent workloads. +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 @@ -99,21 +99,23 @@ For a browser based configuration, try the [Configuration Terminal](https://conf - [**Developer APIs and CLI**](https://docs.trustgraph.ai/reference) - [**Deployment Guides**](https://docs.trustgraph.ai/deployment) -## Context Graph UI +## Workbench -Image +The **Workbench** provides tools for all major features of TrustGraph. The **Workbench** is on port `8888` by default. -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 -- **Flow Management** — Flow creation and detail views with configurable parameters, temperature controls, and grouped storage layout -- **Workspace UX** — Workspace selection and management surfaced directly in the interface -- **Prompt Editor** — A dedicated prompt editing workflow +- **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 ## TypeScript Library for UIs diff --git a/docs/tech-specs/knowledge-core-completeness.md b/docs/tech-specs/knowledge-core-completeness.md deleted file mode 100644 index 3ccb41f0..00000000 --- a/docs/tech-specs/knowledge-core-completeness.md +++ /dev/null @@ -1,535 +0,0 @@ ---- -layout: default -title: "Knowledge Core Completeness" -parent: "Tech Specs" ---- - -# Knowledge Core Completeness - -## Overview - -Knowledge cores are portable snapshots of extracted knowledge: triples, graph -embeddings, and document embeddings stored in Cassandra's `knowledge` keyspace. -They can be downloaded as files, transferred between TrustGraph instances, and -loaded back into vector and graph stores. - -Recent additions to TrustGraph — explainability/provenance and named graphs — -were not carried through to the knowledge core system. This means that -exporting and re-importing a core loses provenance links, graph assignments, -and source material, breaking the explainability chain. - -This specification addresses three gaps: - -1. **Named graphs not stored** — The `g` (graph name) field on triples is - silently dropped when writing to the core store and comes back as `None` - on read. -2. **Provenance triples not captured** — Provenance triples (PROV-O) are - generated during extraction and flow to graph stores, but never enter - the knowledge core store. It is unclear whether they arrive at the store - in the correct form. -3. **Source material not included** — Documents, text pages, and chunks in - the librarian's bucket store are not part of the core. After loading a - core on a different instance, provenance links to source material point - at nothing. - -## Goals - -- **Self-contained cores**: A downloaded knowledge core file contains - everything needed to reconstruct the full knowledge graph including - provenance and source attribution on a fresh instance. -- **Named graph preservation**: Round-tripping a core preserves graph - assignments on all triples. -- **Backward compatibility**: Existing core files (without graph names or - source material) can still be uploaded and loaded. New fields are optional - on import. -- **No change to core identity**: A core is still identified by its document - ID. The additional data is associated with the same core ID. -- **Minimal file format changes**: Extend the existing msgpack record format - with new record types rather than restructuring existing ones. - -## Background - -### Current Lifecycle - -``` -Extraction pipeline - │ - ├─ triples ──────────────────► knowledge core store (Cassandra) - ├─ graph embeddings ─────────► knowledge core store (Cassandra) - ├─ document embeddings ──────► knowledge core store (Cassandra) - ├─ provenance triples ───────► graph store (only) - └─ source documents ─────────► librarian bucket store (only) - -Download: Cassandra ──► knowledge manager ──► API gateway ──► client file -Upload: client file ──► API gateway ──► knowledge manager ──► Cassandra -Load: Cassandra ──► knowledge manager ──► Pulsar topics ──► graph/vector stores -``` - -### Current Core File Format (msgpack) - -A core file is a sequence of concatenated msgpack records. Each record is a -2-element tuple: `(type_tag, payload)`. - -| Type tag | Payload | Description | -|----------|---------|-------------| -| `"t"` | `{"m": {id, root, collection}, "t": [triple_dicts]}` | Triple batch | -| `"ge"` | `{"m": {id, root, collection}, "e": [{entity, vector}]}` | Graph embedding batch | - -### What's Missing - -#### Named Graphs - -The `Triple` dataclass has a `g: str | None` field (graph name IRI), used to -separate provenance graphs (`urn:graph:source`, `urn:graph:retrieval`) from -the default graph. However: - -- **Cassandra schema** (`knowledge.triples` table): stores a 6-tuple per - triple `(s_val, s_is_uri, p_val, p_is_uri, o_val, o_is_uri)` — no graph - field. -- **`add_triples()`** (`tables/knowledge.py:231`): destructures only `s`, - `p`, `o` — `g` is discarded. -- **`get_triples()`** (`tables/knowledge.py:396`): reconstructs `Triple` - with `g` defaulting to `None`. -- **Core file format**: triple dicts do not include a graph field. - -#### Provenance Triples - -Provenance triples are generated in the extraction pipeline -(`trustgraph-base/trustgraph/provenance/triples.py`) and published to graph -store topics. They use named graphs (`urn:graph:source`, -`urn:graph:retrieval`) and PROV-O vocabulary. - -The knowledge core store processor (`storage/knowledge/store.py`) listens on -`triples-input` and `graph-embeddings-input`. Whether provenance triples -arrive on the same `triples-input` topic or a separate one needs -verification. Even if they do arrive, the graph name would be lost (per -above). - -#### Source Material - -The librarian stores the full document hierarchy in a separate system: - -- **Blob store** (S3/MinIO): original documents, text pages, chunks — - keyed by object UUID under `doc/{object_id}`. -- **Cassandra `library` keyspace**: document metadata including `id`, - `kind` (MIME type), `title`, `parent_id`, `document_type` - (`source`/`extracted`), `object_id` (blob reference). - -Provenance triples link extracted facts back to chunk/page/document IDs. -Those IDs resolve through the librarian. When a core is loaded on a -different instance, the librarian has no matching documents, so the entire -provenance chain is broken. - -### Key Source Files - -| Component | File | Purpose | -|-----------|------|---------| -| Core Cassandra schema | `trustgraph-flow/trustgraph/tables/knowledge.py` | Table definitions, read/write | -| Core manager | `trustgraph-flow/trustgraph/cores/knowledge.py` | API operations, load-to-store | -| Core store processor | `trustgraph-flow/trustgraph/storage/knowledge/store.py` | Extraction → Cassandra | -| CLI download | `trustgraph-cli/trustgraph/cli/get_kg_core.py` | Core → msgpack file | -| CLI upload | `trustgraph-cli/trustgraph/cli/put_kg_core.py` | Msgpack file → core | -| CLI load | `trustgraph-cli/trustgraph/cli/load_kg_core.py` | Core → graph/vector stores | -| API client | `trustgraph-base/trustgraph/api/knowledge.py` | Client-side knowledge API | -| Triple schema | `trustgraph-base/trustgraph/schema/core/primitives.py` | Triple dataclass with `g` field | -| Provenance generation | `trustgraph-base/trustgraph/provenance/triples.py` | PROV-O triple creation | -| Librarian | `trustgraph-flow/trustgraph/librarian/librarian.py` | Document storage service | -| Library tables | `trustgraph-flow/trustgraph/tables/library.py` | Document metadata in Cassandra | -| Blob store | `trustgraph-flow/trustgraph/librarian/blob_store.py` | S3/MinIO object storage | - -## Technical Design - -### Change 1: Named Graph Field in Core Storage - -#### Cassandra Schema - -Extend the `triples` tuple from 6 to 7 elements, adding the graph name: - -``` -triples list> -``` - -**Migration**: The schema change uses `ALTER TABLE` or is handled by -creating a new table version. Existing rows with 6-element tuples must be -handled gracefully on read — if the tuple has 6 elements, treat graph as -default. - -#### Write Path (`add_triples`) - -Change `tables/knowledge.py:add_triples()` to include `triple.g`: - -```python -triples = [ - ( - *term_to_tuple(v.s), *term_to_tuple(v.p), *term_to_tuple(v.o), - v.g or "" - ) - for v in m.triples -] -``` - -#### Read Path (`get_triples`) - -Change `tables/knowledge.py:get_triples()` to restore the graph name: - -```python -Triple( - s = tuple_to_term(elt[0], elt[1]), - p = tuple_to_term(elt[2], elt[3]), - o = tuple_to_term(elt[4], elt[5]), - g = elt[6] if len(elt) > 6 and elt[6] else None, -) -``` - -The `len(elt) > 6` guard provides backward compatibility with existing -6-element rows. - -#### Core File Format - -Extend triple dicts in the `"t"` record to include the graph name: - -```python -# In get_kg_core.py write_triple — each triple dict gains "g" key -{"s": ..., "p": ..., "o": ..., "g": "urn:graph:source"} -``` - -On read (`put_kg_core.py`), treat missing `"g"` key as default graph for -backward compatibility with old core files. - -### Change 2: Provenance Triples in Cores - -#### Investigation Required - -Before implementation, verify: - -1. Whether provenance triples arrive on the `triples-input` topic that the - knowledge core store processor already listens on. -2. If not, which topic they use, and whether the store processor should - subscribe to it. - -#### If provenance triples already arrive at the store - -The only change needed is Change 1 (named graphs) — the provenance triples -are already being stored, just without their graph name. Once graph names -are preserved, provenance triples will round-trip correctly. - -#### If provenance triples do NOT arrive at the store - -Two options: - -**Option A — Route provenance to the existing store topic**: Configure the -flow so provenance triples are published to the same `triples-input` topic. -This is the simpler approach and keeps the store processor unchanged. - -**Option B — Add a subscription**: Add a new `ConsumerSpec` in the store -processor for the provenance topic. This keeps provenance routing -independent but adds complexity. - -Recommendation: Option A, unless there is a reason provenance triples are -intentionally kept off the core store topic. - -### Change 3: Source Material in Cores - -This is the largest change. The goal is that when a core is loaded on a -fresh instance, provenance links to source material resolve. - -#### Architecture - -Source material is **not stored in the knowledge core tables**. It lives in -the librarian (Cassandra `library` keyspace + S3/MinIO blob store) and is -fetched on demand via the librarian's existing service API. - -The knowledge manager acts as a **client of the librarian service** — it -calls the librarian's request/response API over pub/sub to retrieve document -metadata and content. It does not access the library's Cassandra tables or -blob store directly. - -#### Transport - -The librarian's pub/sub API already handles chunking of large documents. -This chunking is designed to be websocket-friendly, so library content -flowing through the API gateway to external clients does not require -re-chunking. The API gateway remains a transport layer. - -``` -Download: - Knowledge manager ──pub/sub──► Librarian (fetch metadata + content) - Knowledge manager ──pub/sub──► API gateway ──websocket──► Client - -Upload: - Client ──websocket──► API gateway ──pub/sub──► Knowledge manager - Knowledge manager ──pub/sub──► Librarian (store metadata + content) -``` - -#### What to Include - -The provenance chain links facts → chunks → pages → documents. For the -chain to resolve, the core must include: - -1. **Document metadata** — the library record for each document in the - hierarchy (id, kind, title, parent_id, document_type, etc.) -2. **Document content** — the blob data for each document (original file, - extracted text pages, text chunks) - -Including the full hierarchy is necessary because: -- A user viewing provenance needs to traverse fact → chunk → page → document -- The chunk text is needed to show what text a fact was extracted from -- The page text provides broader context -- The original document is needed for full source attribution - -#### Size Implications - -Source material will significantly increase core file sizes. A rough model: - -| Component | Typical size per document | -|-----------|-------------------------| -| Triples + embeddings (current) | 1-10 MB | -| Chunk text (all chunks) | ~same as original document | -| Page text (all pages) | ~same as original document | -| Original document (PDF, etc.) | Varies widely (KB to hundreds of MB) | - -For a 10 MB PDF, the core could grow from ~5 MB to ~25 MB (original + -derived text + existing data). For large document sets, cores could become -very large. - -**Decision needed**: Whether to include original documents or just derived -text (pages + chunks). Including only derived text still allows provenance -display but loses the ability to serve the original file. - -#### New Core File Record Types - -Add new msgpack record types for library content: - -| Type tag | Payload | Description | -|----------|---------|-------------| -| `"lm"` | `{"id", "kind", "title", "parent_id", "document_type", "comments", "tags", "metadata"}` | Library document metadata | -| `"lb"` | `{"id", "data"}` | Library document blob content (chunked by pub/sub layer) | - -These are emitted after the existing `"t"` and `"ge"` records during -download and processed during upload. - -#### Download Path - -Extend `KnowledgeManager.get_kg_core()` to: - -1. Stream triples and graph embeddings from the core store (existing - behavior). -2. Use the librarian service API to retrieve documents associated with - this core ID: - a. Fetch the root document metadata and content. - b. Use `list-children` to discover child documents (pages, chunks). - c. Recursively fetch metadata and content for each child. -3. Stream each document as `"lm"` (metadata) and `"lb"` (content) records. - -The knowledge manager gains the librarian service as a pub/sub dependency. -Large document content is chunked by the librarian's existing pub/sub -transport — the knowledge manager receives and forwards these chunks without -buffering the full blob in memory. - -#### Upload Path - -Extend `KnowledgeManager.put_kg_core()` to handle the new record types: - -1. For `"lm"` records: call the librarian service API to create/update - the document metadata. -2. For `"lb"` records: call the librarian service API to store the - document content. - -Parent-child relationships are preserved because `parent_id` is stored in -the metadata. Documents should be processed in hierarchy order (parent -before child) to satisfy any ordering constraints. - -#### Load Path - -The load path (`_load_kg_core`) publishes triples and embeddings to Pulsar -topics for ingestion into graph/vector stores. Source material does not need -to flow through the load path — it is already in the librarian after the -upload step and can be accessed directly by services that need it. - -No changes to the load path for source material. - -#### CLI Changes - -**`tg-get-kg-core`**: Add handling for `"lm"` and `"lb"` record types in -the file writer. - -**`tg-put-kg-core`**: Add handling for `"lm"` and `"lb"` record types in -the file reader. Send library records to the knowledge manager alongside -triple/embedding records. - -#### Associating Documents with Cores - -The core ID is `metadata.root`, which is the root document ID from the -librarian. This provides a natural join: the core's root document and all -its children (pages, chunks) are the source material for that core. - -The librarian's `list-children` API provides the child documents. A -recursive traversal from the root document collects the full hierarchy. - -### API Changes - -#### KnowledgeResponse Schema - -Add optional fields to `KnowledgeResponse` for library data: - -```python -@dataclass -class KnowledgeResponse: - error: Error | None = None - ids: list | None = None - eos: bool = False - triples: Triples | None = None - graph_embeddings: GraphEmbeddings | None = None - document_embeddings: DocumentEmbeddings | None = None - library_metadata: LibraryMetadata | None = None # new - library_blob: LibraryBlob | None = None # new -``` - -#### New Schema Types - -```python -@dataclass -class LibraryMetadata: - id: str - kind: str | None = None - title: str | None = None - parent_id: str | None = None - document_type: str | None = None - comments: str | None = None - tags: list[str] | None = None - metadata: list[Triple] | None = None - -@dataclass -class LibraryBlob: - id: str - data: bytes -``` - -#### Socket API - -The existing streaming protocol for `get-kg-core` / `put-kg-core` carries -these new fields naturally — responses already stream multiple record types. - -### Dependencies Between Changes - -``` -Change 1 (named graphs) ◄── Change 2 depends on this - │ - └── Change 2 (provenance triples) - │ - └── Change 3 (source material) is independent -``` - -Change 1 is a prerequisite for Change 2 (provenance triples use named -graphs). Change 3 is independent and can be implemented in parallel. - -## Security Considerations - -- **Workspace isolation**: Core download/upload must respect workspace - boundaries. Source material from the librarian must only be included if - it belongs to the same workspace as the core. This is already enforced - by the existing workspace-scoped queries. -- **Large blob transfer**: Streaming large documents through the API - is handled by the librarian's existing pub/sub chunking, which is - designed to be websocket-friendly. No additional chunking layer is - needed. -- **Cross-instance trust**: When uploading a core from an external source, - the library content should be treated as untrusted input. Document - metadata and blob content should be validated before insertion. - -## Performance Considerations - -- **Core file size**: Including source material will significantly increase - core file sizes. Consider adding a flag to download/upload commands to - optionally exclude source material for use cases where only the knowledge - graph is needed. -- **Streaming**: All paths already use streaming (paged Cassandra queries, - msgpack record-at-a-time). Library content should follow the same pattern. -- **Cassandra schema migration**: Changing the tuple width in the `triples` - table requires careful handling. Cassandra frozen tuples cannot be altered - in place — a migration strategy is needed (see Migration Plan). - -## Testing Strategy - -- **Unit tests**: Triple round-trip with graph name (write → read → - verify `g` field preserved). Backward compatibility with 6-element tuples. -- **Integration tests**: Full lifecycle — extract with provenance → download - core → upload to fresh instance → load → verify provenance chain resolves. -- **File format tests**: Read old-format core files (no graph name, no - library records) and verify they load without error. -- **Library inclusion tests**: Download core with source material → upload → - verify documents accessible through librarian. - -## Migration Plan - -### Cassandra Schema - -The `triples` table stores tuples in a `list>` column. Cassandra -does not support altering the type of an existing column. Options: - -**Option A — New table**: Create a `triples_v2` table with the 7-element -tuple. Migrate data from `triples` to `triples_v2`. The read path checks -both tables during a transition period, then the old table is dropped. - -**Option B — Dual read**: Keep the existing table. The read path handles -both 6-element and 7-element tuples by checking length. New writes use -7-element tuples. This works if Cassandra accepts variable-length tuples in -a list — **needs verification**. - -**Option C — Separate graph column**: Instead of extending the tuple, add a -parallel `graphs list` column where `graphs[i]` corresponds to -`triples[i]`. This avoids tuple migration entirely but requires keeping the -two lists in sync. - -Recommendation: Verify Option B first (simplest). Fall back to Option A if -Cassandra rejects mixed tuple lengths. - -### Core File Format - -Backward compatible by design: -- Old files lack `"g"` in triple dicts and have no `"lm"`/`"lb"` records → - handled by defaults. -- New files read by old code → old code ignores unknown record types (the - existing `read_message` raises on unknown types, so this needs a small - fix to skip unknown types gracefully). - -## Open Questions - -1. **Provenance topic routing**: Do provenance triples currently arrive at - the `triples-input` topic consumed by the knowledge core store? If not, - what topic are they on? - -2. **Include original documents?**: Should cores include the original - uploaded document (e.g. PDF), or only derived text (pages + chunks)? - Including originals makes cores fully self-contained but potentially - very large. Excluding them preserves provenance text display but loses - the ability to serve the original file. - -3. **Optional source material**: Should there be a flag on download/upload - to include or exclude source material? This would let users choose - between compact cores (knowledge only) and complete cores (knowledge + - sources). - -4. **Cassandra tuple migration**: Can Cassandra handle mixed-length tuples - in a `list>` column, or is a table migration required? - -5. **Document embedding cores**: DE cores are managed alongside KG cores. - Do they need the same treatment (source material inclusion)? The - document embeddings reference chunk IDs — the same provenance chain - applies. - -6. **Core versioning**: Should the core file include a version marker so - readers can distinguish old-format from new-format files without - trial-and-error parsing? - -## References - -- Extraction-time provenance: `docs/tech-specs/extraction-time-provenance.md` -- Query-time explainability: `docs/tech-specs/query-time-explainability.md` -- Agent explainability: `docs/tech-specs/agent-explainability.md` -- Data ownership model: `docs/tech-specs/data-ownership-model.md` diff --git a/tests/unit/test_base/test_cassandra_config.py b/tests/unit/test_base/test_cassandra_config.py index fe8a8379..a291434d 100644 --- a/tests/unit/test_base/test_cassandra_config.py +++ b/tests/unit/test_base/test_cassandra_config.py @@ -409,57 +409,4 @@ class TestEdgeCases: assert hosts == ['mixed-host'] assert username is None # Stays None - assert password == 'mixed-pass' - - -class TestReplicationFactorParamPath: - - def test_explicit_kwarg(self): - with patch.dict(os.environ, {}, clear=True): - _, _, _, _, rf = resolve_cassandra_config( - replication_factor=3, - ) - assert rf == 3 - - def test_kwarg_overrides_env(self): - with patch.dict(os.environ, {'CASSANDRA_REPLICATION_FACTOR': '5'}, clear=True): - _, _, _, _, rf = resolve_cassandra_config( - replication_factor=3, - ) - assert rf == 3 - - def test_env_fallback_when_kwarg_none(self): - with patch.dict(os.environ, {'CASSANDRA_REPLICATION_FACTOR': '5'}, clear=True): - _, _, _, _, rf = resolve_cassandra_config( - replication_factor=None, - ) - assert rf == 5 - - def test_default_when_no_kwarg_no_env(self): - with patch.dict(os.environ, {}, clear=True): - _, _, _, _, rf = resolve_cassandra_config() - assert rf == 1 - - def test_params_dict_path(self): - with patch.dict(os.environ, {}, clear=True): - params = {'cassandra_replication_factor': 3} - _, _, _, _, rf = resolve_cassandra_config( - replication_factor=params.get('cassandra_replication_factor'), - ) - assert rf == 3 - - def test_params_dict_overrides_env(self): - with patch.dict(os.environ, {'CASSANDRA_REPLICATION_FACTOR': '5'}, clear=True): - params = {'cassandra_replication_factor': 3} - _, _, _, _, rf = resolve_cassandra_config( - replication_factor=params.get('cassandra_replication_factor'), - ) - assert rf == 3 - - def test_params_dict_missing_falls_to_env(self): - with patch.dict(os.environ, {'CASSANDRA_REPLICATION_FACTOR': '5'}, clear=True): - params = {} - _, _, _, _, rf = resolve_cassandra_config( - replication_factor=params.get('cassandra_replication_factor'), - ) - assert rf == 5 \ No newline at end of file + assert password == 'mixed-pass' \ No newline at end of file diff --git a/tests/unit/test_base/test_qdrant_config.py b/tests/unit/test_base/test_qdrant_config.py deleted file mode 100644 index dbbe4214..00000000 --- a/tests/unit/test_base/test_qdrant_config.py +++ /dev/null @@ -1,136 +0,0 @@ - -import os -import pytest -from unittest.mock import patch - -from trustgraph.base.qdrant_config import ( - get_qdrant_defaults, - resolve_qdrant_config, -) - - -class TestGetQdrantDefaults: - - def test_defaults_with_no_env_vars(self): - with patch.dict(os.environ, {}, clear=True): - defaults = get_qdrant_defaults() - assert defaults['url'] == 'http://localhost:6333' - assert defaults['api_key'] is None - assert defaults['replication_factor'] == 1 - assert defaults['shard_number'] == 1 - - def test_defaults_from_env(self): - env = { - 'QDRANT_URL': 'http://qdrant:6333', - 'QDRANT_API_KEY': 'secret', - 'QDRANT_REPLICATION_FACTOR': '3', - 'QDRANT_SHARD_NUMBER': '5', - } - with patch.dict(os.environ, env, clear=True): - defaults = get_qdrant_defaults() - assert defaults['url'] == 'http://qdrant:6333' - assert defaults['api_key'] == 'secret' - assert defaults['replication_factor'] == 3 - assert defaults['shard_number'] == 5 - - -class TestResolveQdrantConfig: - - def test_defaults(self): - with patch.dict(os.environ, {}, clear=True): - url, api_key, rf, sn = resolve_qdrant_config() - assert url == 'http://localhost:6333' - assert api_key is None - assert rf == 1 - assert sn == 1 - - def test_explicit_kwargs(self): - with patch.dict(os.environ, {}, clear=True): - url, api_key, rf, sn = resolve_qdrant_config( - url='http://custom:6333', - api_key='key', - replication_factor=3, - shard_number=5, - ) - assert url == 'http://custom:6333' - assert api_key == 'key' - assert rf == 3 - assert sn == 5 - - def test_kwargs_override_env(self): - env = { - 'QDRANT_URL': 'http://env:6333', - 'QDRANT_REPLICATION_FACTOR': '10', - 'QDRANT_SHARD_NUMBER': '10', - } - with patch.dict(os.environ, env, clear=True): - url, _, rf, sn = resolve_qdrant_config( - url='http://explicit:6333', - replication_factor=3, - shard_number=5, - ) - assert url == 'http://explicit:6333' - assert rf == 3 - assert sn == 5 - - def test_env_fallback_when_kwargs_none(self): - env = { - 'QDRANT_URL': 'http://env:6333', - 'QDRANT_REPLICATION_FACTOR': '3', - 'QDRANT_SHARD_NUMBER': '5', - } - with patch.dict(os.environ, env, clear=True): - url, _, rf, sn = resolve_qdrant_config() - assert url == 'http://env:6333' - assert rf == 3 - assert sn == 5 - - def test_params_dict_path(self): - with patch.dict(os.environ, {}, clear=True): - params = { - 'store_uri': 'http://params:6333', - 'api_key': 'pkey', - 'qdrant_replication_factor': 3, - 'qdrant_shard_number': 5, - } - url, api_key, rf, sn = resolve_qdrant_config( - url=params.get('store_uri'), - api_key=params.get('api_key'), - replication_factor=params.get('qdrant_replication_factor'), - shard_number=params.get('qdrant_shard_number'), - ) - assert url == 'http://params:6333' - assert api_key == 'pkey' - assert rf == 3 - assert sn == 5 - - def test_params_dict_overrides_env(self): - env = { - 'QDRANT_REPLICATION_FACTOR': '10', - 'QDRANT_SHARD_NUMBER': '10', - } - with patch.dict(os.environ, env, clear=True): - params = { - 'qdrant_replication_factor': 3, - 'qdrant_shard_number': 5, - } - _, _, rf, sn = resolve_qdrant_config( - replication_factor=params.get('qdrant_replication_factor'), - shard_number=params.get('qdrant_shard_number'), - ) - assert rf == 3 - assert sn == 5 - - def test_params_dict_missing_falls_to_env(self): - env = { - 'QDRANT_REPLICATION_FACTOR': '3', - 'QDRANT_SHARD_NUMBER': '5', - } - with patch.dict(os.environ, env, clear=True): - params = {} - _, _, rf, sn = resolve_qdrant_config( - replication_factor=params.get('qdrant_replication_factor'), - shard_number=params.get('qdrant_shard_number'), - ) - assert rf == 3 - assert sn == 5 diff --git a/tests/unit/test_cores/test_knowledge_manager.py b/tests/unit/test_cores/test_knowledge_manager.py index 7797c9be..8f73dcc6 100644 --- a/tests/unit/test_cores/test_knowledge_manager.py +++ b/tests/unit/test_cores/test_knowledge_manager.py @@ -11,12 +11,7 @@ from unittest.mock import AsyncMock, Mock, patch, MagicMock from unittest.mock import call from trustgraph.cores.knowledge import KnowledgeManager -from trustgraph.schema import ( - KnowledgeResponse, Triples, GraphEmbeddings, Metadata, Triple, Term, - EntityEmbeddings, IRI, LITERAL, - LibraryMetadata, LibraryBlob, - LibrarianResponse, DocumentMetadata, -) +from trustgraph.schema import KnowledgeResponse, Triples, GraphEmbeddings, Metadata, Triple, Term, EntityEmbeddings, IRI, LITERAL @pytest.fixture @@ -378,252 +373,11 @@ class TestKnowledgeManagerOtherMethods: mock_respond = AsyncMock() await knowledge_manager.delete_kg_core(mock_request, mock_respond, "test-user") - + # Verify table store was called correctly knowledge_manager.table_store.delete_kg_core.assert_called_once_with("test-user", "test-doc-id") - + # Verify response mock_respond.assert_called_once() response = mock_respond.call_args[0][0] - assert response.error is None - - -class TestKnowledgeManagerLibraryDownload: - """Test get_kg_core streaming of library documents.""" - - @pytest.fixture - def manager_with_librarian(self, mock_flow_config): - with patch('trustgraph.cores.knowledge.KnowledgeTableStore'): - mock_librarian = AsyncMock() - manager = KnowledgeManager( - cassandra_host=["localhost"], - cassandra_username="test_user", - cassandra_password="test_pass", - keyspace="test_keyspace", - flow_config=mock_flow_config, - librarian=mock_librarian, - ) - manager.table_store = AsyncMock() - return manager - - @pytest.mark.asyncio - async def test_get_kg_core_streams_library_docs(self, manager_with_librarian): - mock_request = Mock() - mock_request.id = "root-doc" - mock_respond = AsyncMock() - - manager_with_librarian.table_store.get_triples = AsyncMock() - manager_with_librarian.table_store.get_graph_embeddings = AsyncMock() - - root_meta = DocumentMetadata( - id="root-doc", kind="application/pdf", title="Test PDF", - document_type="source", - ) - child_meta = DocumentMetadata( - id="chunk-1", kind="text/plain", title="Chunk 1", - parent_id="root-doc", document_type="chunk", - ) - - manager_with_librarian.librarian.fetch_document_metadata.return_value = root_meta - manager_with_librarian.librarian.request.return_value = LibrarianResponse( - document_metadatas=[child_meta], - ) - manager_with_librarian.librarian.fetch_document_content.side_effect = [ - b"cm9vdCBjb250ZW50", - b"Y2h1bmsgY29udGVudA==", - ] - - await manager_with_librarian.get_kg_core( - mock_request, mock_respond, "test-user" - ) - - responses = [c[0][0] for c in mock_respond.call_args_list] - - lm_responses = [r for r in responses if r.library_metadata is not None] - lb_responses = [r for r in responses if r.library_blob is not None] - eos_responses = [r for r in responses if r.eos is True] - - assert len(lm_responses) == 2 - assert lm_responses[0].library_metadata.id == "root-doc" - assert lm_responses[0].library_metadata.document_type == "source" - assert lm_responses[1].library_metadata.id == "chunk-1" - assert lm_responses[1].library_metadata.parent_id == "root-doc" - - assert len(lb_responses) == 2 - assert lb_responses[0].library_blob.id == "root-doc" - assert lb_responses[0].library_blob.data == b"cm9vdCBjb250ZW50" - assert lb_responses[1].library_blob.id == "chunk-1" - - assert len(eos_responses) == 1 - - @pytest.mark.asyncio - async def test_get_kg_core_no_librarian_skips_library(self, mock_flow_config): - with patch('trustgraph.cores.knowledge.KnowledgeTableStore'): - manager = KnowledgeManager( - cassandra_host=["localhost"], - cassandra_username="u", cassandra_password="p", - keyspace="ks", flow_config=mock_flow_config, - ) - manager.table_store = AsyncMock() - manager.table_store.get_triples = AsyncMock() - manager.table_store.get_graph_embeddings = AsyncMock() - - mock_request = Mock() - mock_request.id = "doc-1" - mock_respond = AsyncMock() - - await manager.get_kg_core(mock_request, mock_respond, "w") - - responses = [c[0][0] for c in mock_respond.call_args_list] - assert all(r.library_metadata is None for r in responses) - assert all(r.library_blob is None for r in responses) - - @pytest.mark.asyncio - async def test_get_kg_core_librarian_metadata_failure_is_graceful( - self, manager_with_librarian, - ): - mock_request = Mock() - mock_request.id = "missing-doc" - mock_respond = AsyncMock() - - manager_with_librarian.table_store.get_triples = AsyncMock() - manager_with_librarian.table_store.get_graph_embeddings = AsyncMock() - manager_with_librarian.librarian.fetch_document_metadata.side_effect = ( - RuntimeError("not found") - ) - - await manager_with_librarian.get_kg_core( - mock_request, mock_respond, "test-user" - ) - - responses = [c[0][0] for c in mock_respond.call_args_list] - assert all(r.library_metadata is None for r in responses) - assert any(r.eos for r in responses) - - -class TestKnowledgeManagerLibraryUpload: - """Test put_kg_core handling of library metadata and blob records.""" - - @pytest.fixture - def manager_with_librarian(self, mock_flow_config): - with patch('trustgraph.cores.knowledge.KnowledgeTableStore'): - mock_librarian = AsyncMock() - manager = KnowledgeManager( - cassandra_host=["localhost"], - cassandra_username="u", cassandra_password="p", - keyspace="ks", flow_config=mock_flow_config, - librarian=mock_librarian, - ) - manager.table_store = AsyncMock() - return manager - - @pytest.mark.asyncio - async def test_put_metadata_then_blob_calls_librarian( - self, manager_with_librarian, - ): - mock_respond = AsyncMock() - manager_with_librarian.librarian.request.return_value = LibrarianResponse() - - # First call: metadata - req_meta = Mock() - req_meta.triples = None - req_meta.graph_embeddings = None - req_meta.library_metadata = LibraryMetadata( - id="doc-1", kind="application/pdf", title="Test", - document_type="source", - ) - req_meta.library_blob = None - await manager_with_librarian.put_kg_core(req_meta, mock_respond, "ws") - - # Metadata is buffered, librarian not called yet - manager_with_librarian.librarian.request.assert_not_called() - - # Second call: blob - req_blob = Mock() - req_blob.triples = None - req_blob.graph_embeddings = None - req_blob.library_metadata = None - req_blob.library_blob = LibraryBlob( - id="doc-1", data=b"dGVzdA==", - ) - await manager_with_librarian.put_kg_core(req_blob, mock_respond, "ws") - - # Now librarian should have been called with add-document - manager_with_librarian.librarian.request.assert_called_once() - call_args = manager_with_librarian.librarian.request.call_args[0][0] - assert call_args.operation == "add-document" - assert call_args.document_metadata.id == "doc-1" - assert call_args.document_metadata.kind == "application/pdf" - assert call_args.content == b"dGVzdA==" - - @pytest.mark.asyncio - async def test_put_child_document_uses_add_child_operation( - self, manager_with_librarian, - ): - mock_respond = AsyncMock() - manager_with_librarian.librarian.request.return_value = LibrarianResponse() - - req_meta = Mock() - req_meta.triples = None - req_meta.graph_embeddings = None - req_meta.library_metadata = LibraryMetadata( - id="chunk-1", kind="text/plain", title="Chunk", - parent_id="doc-1", document_type="chunk", - ) - req_meta.library_blob = None - await manager_with_librarian.put_kg_core(req_meta, mock_respond, "ws") - - req_blob = Mock() - req_blob.triples = None - req_blob.graph_embeddings = None - req_blob.library_metadata = None - req_blob.library_blob = LibraryBlob(id="chunk-1", data=b"Y2h1bms=") - await manager_with_librarian.put_kg_core(req_blob, mock_respond, "ws") - - call_args = manager_with_librarian.librarian.request.call_args[0][0] - assert call_args.operation == "add-child-document" - assert call_args.document_metadata.parent_id == "doc-1" - - @pytest.mark.asyncio - async def test_put_blob_without_metadata_logs_warning( - self, manager_with_librarian, - ): - mock_respond = AsyncMock() - - req_blob = Mock() - req_blob.triples = None - req_blob.graph_embeddings = None - req_blob.library_metadata = None - req_blob.library_blob = LibraryBlob(id="orphan", data=b"data") - await manager_with_librarian.put_kg_core(req_blob, mock_respond, "ws") - - # Librarian should not be called for orphan blob - manager_with_librarian.librarian.request.assert_not_called() - - @pytest.mark.asyncio - async def test_put_existing_document_is_graceful( - self, manager_with_librarian, - ): - mock_respond = AsyncMock() - manager_with_librarian.librarian.request.side_effect = RuntimeError( - "Document already exists" - ) - - req_meta = Mock() - req_meta.triples = None - req_meta.graph_embeddings = None - req_meta.library_metadata = LibraryMetadata( - id="doc-1", kind="application/pdf", title="Test", - document_type="source", - ) - req_meta.library_blob = None - await manager_with_librarian.put_kg_core(req_meta, mock_respond, "ws") - - req_blob = Mock() - req_blob.triples = None - req_blob.graph_embeddings = None - req_blob.library_metadata = None - req_blob.library_blob = LibraryBlob(id="doc-1", data=b"data") - await manager_with_librarian.put_kg_core(req_blob, mock_respond, "ws") - - # Should not raise — "already exists" is handled gracefully \ No newline at end of file + assert response.error is None \ No newline at end of file diff --git a/tests/unit/test_decoding/test_pdf_decoder.py b/tests/unit/test_decoding/test_pdf_decoder.py index 641a9d78..04807b20 100644 --- a/tests/unit/test_decoding/test_pdf_decoder.py +++ b/tests/unit/test_decoding/test_pdf_decoder.py @@ -49,7 +49,7 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase): async def test_on_message_success(self, mock_pdf_loader_class, mock_producer, mock_consumer): """Test successful PDF processing""" # Mock PDF content - pdf_content = b"%PDF-1.7\nfake pdf content" + pdf_content = b"fake pdf content" pdf_base64 = base64.b64encode(pdf_content).decode('utf-8') # Mock PyPDFLoader @@ -88,55 +88,13 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase): # Verify triples were sent for each page (provenance) assert mock_triples_flow.send.call_count == 2 - @patch('trustgraph.base.librarian_client.Consumer') - @patch('trustgraph.base.librarian_client.Producer') - @patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader') - @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) - async def test_on_message_rejects_librarian_content_that_is_not_pdf(self, mock_pdf_loader_class, mock_producer, mock_consumer): - """Test rejecting non-PDF content before invoking the PDF loader""" - html_content = b"Not found" - html_base64 = base64.b64encode(html_content) - - mock_metadata = Metadata(id="test-doc") - mock_document = Document(metadata=mock_metadata, document_id="doc-123") - mock_msg = MagicMock() - mock_msg.value.return_value = mock_document - - mock_output_flow = AsyncMock() - mock_triples_flow = AsyncMock() - mock_flow = MagicMock(side_effect=lambda name: { - "output": mock_output_flow, - "triples": mock_triples_flow, - }.get(name)) - mock_flow.librarian.fetch_document_metadata = AsyncMock( - return_value=MagicMock(kind="application/pdf") - ) - mock_flow.librarian.fetch_document_content = AsyncMock( - return_value=html_base64 - ) - mock_flow.librarian.save_child_document = AsyncMock() - - config = { - 'id': 'test-pdf-decoder', - 'taskgroup': AsyncMock() - } - - processor = Processor(**config) - - await processor.on_message(mock_msg, None, mock_flow) - - mock_pdf_loader_class.assert_not_called() - mock_output_flow.send.assert_not_called() - mock_triples_flow.send.assert_not_called() - mock_flow.librarian.save_child_document.assert_not_called() - @patch('trustgraph.base.librarian_client.Consumer') @patch('trustgraph.base.librarian_client.Producer') @patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader') @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) async def test_on_message_empty_pdf(self, mock_pdf_loader_class, mock_producer, mock_consumer): """Test handling of empty PDF""" - pdf_content = b"%PDF-1.7\nfake pdf content" + pdf_content = b"fake pdf content" pdf_base64 = base64.b64encode(pdf_content).decode('utf-8') mock_loader = MagicMock() @@ -168,7 +126,7 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase): @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) async def test_on_message_unicode_content(self, mock_pdf_loader_class, mock_producer, mock_consumer): """Test handling of unicode content in PDF""" - pdf_content = b"%PDF-1.7\nfake pdf content" + pdf_content = b"fake pdf content" pdf_base64 = base64.b64encode(pdf_content).decode('utf-8') mock_loader = MagicMock() diff --git a/tests/unit/test_query/test_rows_cassandra_query.py b/tests/unit/test_query/test_rows_cassandra_query.py index fb385f43..b61500a4 100644 --- a/tests/unit/test_query/test_rows_cassandra_query.py +++ b/tests/unit/test_query/test_rows_cassandra_query.py @@ -333,8 +333,8 @@ class TestUnifiedTableQueries: """Test queries against the unified rows table""" @pytest.mark.asyncio - @patch('trustgraph.query.rows.cassandra.service.async_execute_paged', new_callable=AsyncMock) - async def test_query_with_index_match(self, mock_async_execute_paged): + @patch('trustgraph.query.rows.cassandra.service.async_execute', new_callable=AsyncMock) + async def test_query_with_index_match(self, mock_async_execute): """Test query execution with matching index""" processor = MagicMock() processor.session = MagicMock() @@ -344,10 +344,10 @@ class TestUnifiedTableQueries: processor.find_matching_index = Processor.find_matching_index.__get__(processor, Processor) processor.query_cassandra = Processor.query_cassandra.__get__(processor, Processor) - # Mock async_execute_paged to return test data (list of pages) + # Mock async_execute to return test data mock_row = MagicMock() mock_row.data = {"id": "123", "name": "Test Product", "category": "electronics"} - mock_async_execute_paged.return_value = [[mock_row]] + mock_async_execute.return_value = [mock_row] schema = RowSchema( name="products", @@ -370,10 +370,10 @@ class TestUnifiedTableQueries: # Verify Cassandra was connected and queried processor.connect_cassandra.assert_called_once() - mock_async_execute_paged.assert_called_once() + mock_async_execute.assert_called_once() # Verify query structure - should query unified rows table - call_args = mock_async_execute_paged.call_args + call_args = mock_async_execute.call_args query = call_args[0][1] params = call_args[0][2] @@ -394,8 +394,8 @@ class TestUnifiedTableQueries: assert results[0]["category"] == "electronics" @pytest.mark.asyncio - @patch('trustgraph.query.rows.cassandra.service.async_scan', new_callable=AsyncMock) - async def test_query_without_index_match(self, mock_async_scan): + @patch('trustgraph.query.rows.cassandra.service.async_execute', new_callable=AsyncMock) + async def test_query_without_index_match(self, mock_async_execute): """Test query execution without matching index (scan mode)""" processor = MagicMock() processor.session = MagicMock() @@ -406,10 +406,12 @@ class TestUnifiedTableQueries: processor._matches_filters = Processor._matches_filters.__get__(processor, Processor) processor.query_cassandra = Processor.query_cassandra.__get__(processor, Processor) - # Mock async_scan to return filtered test data + # Mock async_execute to return test data mock_row1 = MagicMock() mock_row1.data = {"id": "1", "name": "Product A", "price": "100"} - mock_async_scan.return_value = [mock_row1] + mock_row2 = MagicMock() + mock_row2.data = {"id": "2", "name": "Product B", "price": "200"} + mock_async_execute.return_value = [mock_row1, mock_row2] schema = RowSchema( name="products", @@ -430,16 +432,13 @@ class TestUnifiedTableQueries: limit=10 ) - # Verify async_scan was called - mock_async_scan.assert_called_once() - - # Verify query structure - call_args = mock_async_scan.call_args + # Query should use ALLOW FILTERING for scan + call_args = mock_async_execute.call_args query = call_args[0][1] assert "ALLOW FILTERING" in query - # Should return filtered results + # Should post-filter results assert len(results) == 1 assert results[0]["name"] == "Product A" diff --git a/tests/unit/test_reliability/test_null_embedding_protection.py b/tests/unit/test_reliability/test_null_embedding_protection.py index 41d0f88b..dbe06b40 100644 --- a/tests/unit/test_reliability/test_null_embedding_protection.py +++ b/tests/unit/test_reliability/test_null_embedding_protection.py @@ -259,8 +259,6 @@ class TestGraphEmbeddingsNullProtection: proc.collection_exists = MagicMock(return_value=True) proc._cache_lock = asyncio.Lock() proc._known_collections = set() - proc.replication_factor = 1 - proc.shard_number = 1 msg = MagicMock() msg.metadata.collection = "graphs" diff --git a/tests/unit/test_tables/test_knowledge_table_store.py b/tests/unit/test_tables/test_knowledge_table_store.py index 2d058733..9a0b55c4 100644 --- a/tests/unit/test_tables/test_knowledge_table_store.py +++ b/tests/unit/test_tables/test_knowledge_table_store.py @@ -155,7 +155,7 @@ class TestGetTriples: @pytest.mark.asyncio @patch('trustgraph.tables.knowledge.async_execute_paged', new_callable=AsyncMock) async def test_row_converts_to_triples(self, mock_async_execute_paged): - # row[3] is a list of (s_val, s_uri, p_val, p_uri, o_val, o_uri, graph) + # row[3] is a list of (s_val, s_uri, p_val, p_uri, o_val, o_uri) fake_row = ( None, None, None, [ @@ -163,7 +163,6 @@ class TestGetTriples: "http://example.org/alice", True, "http://example.org/knows", True, "http://example.org/bob", True, - "urn:graph:source", ), ], ) @@ -192,33 +191,3 @@ class TestGetTriples: assert t.s.iri == "http://example.org/alice" assert t.p.iri == "http://example.org/knows" assert t.o.iri == "http://example.org/bob" - assert t.g == "urn:graph:source" - - @pytest.mark.asyncio - @patch('trustgraph.tables.knowledge.async_execute_paged', new_callable=AsyncMock) - async def test_empty_graph_name_becomes_none(self, mock_async_execute_paged): - fake_row = ( - None, None, None, - [ - ( - "http://example.org/alice", True, - "http://example.org/knows", True, - "http://example.org/bob", True, - "", - ), - ], - ) - - store = _make_store() - store.cassandra = Mock() - store.get_triples_stmt = Mock() - mock_async_execute_paged.return_value = [[fake_row]] - - received = [] - - async def receiver(msg): - received.append(msg) - - await store.get_triples("w", "d", receiver) - - assert received[0].triples[0].g is None diff --git a/tests/unit/test_translators/test_knowledge_translator_roundtrip.py b/tests/unit/test_translators/test_knowledge_translator_roundtrip.py index af128f23..437b83c8 100644 --- a/tests/unit/test_translators/test_knowledge_translator_roundtrip.py +++ b/tests/unit/test_translators/test_knowledge_translator_roundtrip.py @@ -1,6 +1,5 @@ """ -Round-trip unit tests for KnowledgeRequestTranslator and -KnowledgeResponseTranslator. +Round-trip unit tests for KnowledgeRequestTranslator. Regression coverage: a previous version of the decode side constructed EntityEmbeddings(vectors=...) — the schema field is `vector` (singular), @@ -16,13 +15,9 @@ Triples breaks the test. import pytest -from trustgraph.messaging.translators.knowledge import ( - KnowledgeRequestTranslator, - KnowledgeResponseTranslator, -) +from trustgraph.messaging.translators.knowledge import KnowledgeRequestTranslator from trustgraph.schema import ( KnowledgeRequest, - KnowledgeResponse, GraphEmbeddings, EntityEmbeddings, Triples, @@ -30,8 +25,6 @@ from trustgraph.schema import ( Metadata, Term, IRI, - LibraryMetadata, - LibraryBlob, ) @@ -152,161 +145,3 @@ class TestKnowledgeRequestTranslatorTriples: assert t.s.iri == "http://example.org/alice" assert t.p.iri == "http://example.org/knows" assert t.o.iri == "http://example.org/bob" - - -class TestKnowledgeRequestTranslatorLibrary: - - def test_roundtrip_preserves_library_metadata(self, translator): - request = KnowledgeRequest( - operation="put-kg-core", - id="doc-1", - library_metadata=LibraryMetadata( - id="doc-1", - kind="application/pdf", - title="Test Document", - parent_id="", - document_type="source", - comments="test comments", - tags=["tag1", "tag2"], - ), - ) - - encoded = translator.encode(request) - assert "library-metadata" in encoded - lm = encoded["library-metadata"] - assert lm["id"] == "doc-1" - assert lm["kind"] == "application/pdf" - assert lm["title"] == "Test Document" - assert lm["parent-id"] == "" - assert lm["document-type"] == "source" - assert lm["comments"] == "test comments" - assert lm["tags"] == ["tag1", "tag2"] - - decoded = translator.decode(encoded) - assert decoded.library_metadata is not None - assert decoded.library_metadata.id == "doc-1" - assert decoded.library_metadata.kind == "application/pdf" - assert decoded.library_metadata.title == "Test Document" - assert decoded.library_metadata.parent_id == "" - assert decoded.library_metadata.document_type == "source" - assert decoded.library_metadata.comments == "test comments" - assert decoded.library_metadata.tags == ["tag1", "tag2"] - - def test_roundtrip_preserves_child_document_metadata(self, translator): - request = KnowledgeRequest( - operation="put-kg-core", - id="doc-1", - library_metadata=LibraryMetadata( - id="chunk-1", - kind="text/plain", - title="Chunk 1", - parent_id="doc-1", - document_type="chunk", - ), - ) - - encoded = translator.encode(request) - decoded = translator.decode(encoded) - - assert decoded.library_metadata.parent_id == "doc-1" - assert decoded.library_metadata.document_type == "chunk" - - def test_roundtrip_preserves_library_blob(self, translator): - request = KnowledgeRequest( - operation="put-kg-core", - id="doc-1", - library_blob=LibraryBlob( - id="doc-1", - data=b"SGVsbG8gV29ybGQ=", - ), - ) - - encoded = translator.encode(request) - assert "library-blob" in encoded - assert encoded["library-blob"]["id"] == "doc-1" - assert encoded["library-blob"]["data"] == "SGVsbG8gV29ybGQ=" - - decoded = translator.decode(encoded) - assert decoded.library_blob is not None - assert decoded.library_blob.id == "doc-1" - assert decoded.library_blob.data == "SGVsbG8gV29ybGQ=" - - def test_absent_library_fields_decode_as_none(self, translator): - decoded = translator.decode({ - "operation": "get-kg-core", - "id": "doc-1", - }) - assert decoded.library_metadata is None - assert decoded.library_blob is None - - -class TestKnowledgeResponseTranslatorLibrary: - - @pytest.fixture - def response_translator(self): - return KnowledgeResponseTranslator() - - def test_encode_library_metadata(self, response_translator): - response = KnowledgeResponse( - ids=None, - library_metadata=LibraryMetadata( - id="doc-1", - kind="application/pdf", - title="Test", - parent_id="", - document_type="source", - comments="", - tags=[], - ), - ) - encoded = response_translator.encode(response) - assert "library-metadata" in encoded - assert encoded["library-metadata"]["id"] == "doc-1" - assert encoded["library-metadata"]["kind"] == "application/pdf" - assert encoded["library-metadata"]["document-type"] == "source" - - def test_encode_library_blob_bytes_to_string(self, response_translator): - response = KnowledgeResponse( - ids=None, - library_blob=LibraryBlob( - id="doc-1", - data=b"dGVzdCBkYXRh", - ), - ) - encoded = response_translator.encode(response) - assert "library-blob" in encoded - assert encoded["library-blob"]["id"] == "doc-1" - assert encoded["library-blob"]["data"] == "dGVzdCBkYXRh" - assert isinstance(encoded["library-blob"]["data"], str) - - def test_encode_library_blob_string_passthrough(self, response_translator): - response = KnowledgeResponse( - ids=None, - library_blob=LibraryBlob( - id="doc-1", - data="already-a-string", - ), - ) - encoded = response_translator.encode(response) - assert encoded["library-blob"]["data"] == "already-a-string" - - def test_library_metadata_is_not_final(self, response_translator): - response = KnowledgeResponse( - ids=None, - library_metadata=LibraryMetadata(id="doc-1"), - ) - _, is_final = response_translator.encode_with_completion(response) - assert is_final is False - - def test_library_blob_is_not_final(self, response_translator): - response = KnowledgeResponse( - ids=None, - library_blob=LibraryBlob(id="doc-1", data=b"data"), - ) - _, is_final = response_translator.encode_with_completion(response) - assert is_final is False - - def test_eos_is_final(self, response_translator): - response = KnowledgeResponse(eos=True) - _, is_final = response_translator.encode_with_completion(response) - assert is_final is True diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index 91bc67a1..b88d0c78 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -502,7 +502,6 @@ class SocketClient: def put_kg_core( self, id: str, triples=None, graph_embeddings=None, - library_metadata=None, library_blob=None, ) -> Dict[str, Any]: request = { "operation": "put-kg-core", @@ -513,10 +512,6 @@ class SocketClient: request["triples"] = triples if graph_embeddings is not None: request["graph-embeddings"] = graph_embeddings - if library_metadata is not None: - request["library-metadata"] = library_metadata - if library_blob is not None: - request["library-blob"] = library_blob return self._send_request_sync("knowledge", None, request) def get_de_core(self, id: str) -> Iterator[Dict[str, Any]]: diff --git a/trustgraph-base/trustgraph/base/cassandra_config.py b/trustgraph-base/trustgraph/base/cassandra_config.py index b2e36fbd..78505c68 100644 --- a/trustgraph-base/trustgraph/base/cassandra_config.py +++ b/trustgraph-base/trustgraph/base/cassandra_config.py @@ -103,19 +103,35 @@ def resolve_cassandra_config( host: Optional[str] = None, username: Optional[str] = None, password: Optional[str] = None, - default_keyspace: Optional[str] = None, - replication_factor: Optional[int] = None, + default_keyspace: Optional[str] = None ) -> Tuple[List[str], Optional[str], Optional[str], Optional[str], int]: + """ + Resolve Cassandra configuration from various sources. + + Can accept either argparse args object or explicit parameters. + Converts host string to list format for Cassandra driver. + + Args: + args: Optional argparse namespace with cassandra_host, cassandra_username, cassandra_password, cassandra_keyspace, cassandra_replication_factor + host: Optional explicit host parameter (overrides args) + username: Optional explicit username parameter (overrides args) + password: Optional explicit password parameter (overrides args) + default_keyspace: Optional default keyspace if not specified elsewhere + + Returns: + tuple: (hosts_list, username, password, keyspace, replication_factor) + """ + # If args provided, extract values keyspace = None + replication_factor = 1 if args is not None: host = host or getattr(args, 'cassandra_host', None) username = username or getattr(args, 'cassandra_username', None) password = password or getattr(args, 'cassandra_password', None) keyspace = getattr(args, 'cassandra_keyspace', None) - replication_factor = replication_factor or getattr( - args, 'cassandra_replication_factor', None - ) + replication_factor = getattr(args, 'cassandra_replication_factor', 1) + # Apply defaults if still None defaults = get_cassandra_defaults() host = host or defaults['host'] username = username or defaults['username'] diff --git a/trustgraph-base/trustgraph/base/qdrant_config.py b/trustgraph-base/trustgraph/base/qdrant_config.py deleted file mode 100644 index f3e015ca..00000000 --- a/trustgraph-base/trustgraph/base/qdrant_config.py +++ /dev/null @@ -1,87 +0,0 @@ - -import os -import argparse -from typing import Optional, Any, Tuple - - -def get_qdrant_defaults() -> dict: - return { - 'url': os.getenv('QDRANT_URL', 'http://localhost:6333'), - 'api_key': os.getenv('QDRANT_API_KEY'), - 'replication_factor': int(os.getenv('QDRANT_REPLICATION_FACTOR', '1')), - 'shard_number': int(os.getenv('QDRANT_SHARD_NUMBER', '1')), - } - - -def add_qdrant_args(parser: argparse.ArgumentParser) -> None: - defaults = get_qdrant_defaults() - - url_help = f"Qdrant URL (default: {defaults['url']})" - if 'QDRANT_URL' in os.environ: - url_help += " [from QDRANT_URL]" - - api_key_help = "Qdrant API key" - if defaults['api_key']: - api_key_help += " (default: )" - if 'QDRANT_API_KEY' in os.environ: - api_key_help += " [from QDRANT_API_KEY]" - - replication_help = f"Qdrant collection replication factor (default: {defaults['replication_factor']})" - if 'QDRANT_REPLICATION_FACTOR' in os.environ: - replication_help += " [from QDRANT_REPLICATION_FACTOR]" - - shard_help = f"Qdrant collection shard number (default: {defaults['shard_number']})" - if 'QDRANT_SHARD_NUMBER' in os.environ: - shard_help += " [from QDRANT_SHARD_NUMBER]" - - parser.add_argument( - '--store-uri', - default=defaults['url'], - help=url_help, - ) - - parser.add_argument( - '--api-key', - default=defaults['api_key'], - help=api_key_help, - ) - - parser.add_argument( - '--qdrant-replication-factor', - type=int, - default=defaults['replication_factor'], - help=replication_help, - ) - - parser.add_argument( - '--qdrant-shard-number', - type=int, - default=defaults['shard_number'], - help=shard_help, - ) - - -def resolve_qdrant_config( - args: Optional[Any] = None, - url: Optional[str] = None, - api_key: Optional[str] = None, - replication_factor: Optional[int] = None, - shard_number: Optional[int] = None, -) -> Tuple[str, Optional[str], int, int]: - if args is not None: - url = url or getattr(args, 'store_uri', None) - api_key = api_key or getattr(args, 'api_key', None) - replication_factor = replication_factor or getattr( - args, 'qdrant_replication_factor', None - ) - shard_number = shard_number or getattr( - args, 'qdrant_shard_number', None - ) - - defaults = get_qdrant_defaults() - url = url or defaults['url'] - api_key = api_key or defaults['api_key'] - replication_factor = replication_factor or defaults['replication_factor'] - shard_number = shard_number or defaults['shard_number'] - - return url, api_key, replication_factor, shard_number diff --git a/trustgraph-base/trustgraph/messaging/translators/knowledge.py b/trustgraph-base/trustgraph/messaging/translators/knowledge.py index 3f09b41b..3830bf59 100644 --- a/trustgraph-base/trustgraph/messaging/translators/knowledge.py +++ b/trustgraph-base/trustgraph/messaging/translators/knowledge.py @@ -2,8 +2,7 @@ from typing import Dict, Any, Tuple, Optional from ...schema import ( KnowledgeRequest, KnowledgeResponse, Triples, GraphEmbeddings, DocumentEmbeddings, ChunkEmbeddings, - Metadata, EntityEmbeddings, - LibraryMetadata, LibraryBlob, + Metadata, EntityEmbeddings ) from .base import MessageTranslator from .primitives import ValueTranslator, SubgraphTranslator @@ -62,27 +61,6 @@ class KnowledgeRequestTranslator(MessageTranslator): ] ) - library_metadata = None - if "library-metadata" in data: - lm = data["library-metadata"] - library_metadata = LibraryMetadata( - id=lm.get("id", ""), - kind=lm.get("kind", ""), - title=lm.get("title", ""), - parent_id=lm.get("parent-id", ""), - document_type=lm.get("document-type", ""), - comments=lm.get("comments", ""), - tags=lm.get("tags", []), - ) - - library_blob = None - if "library-blob" in data: - lb = data["library-blob"] - library_blob = LibraryBlob( - id=lb.get("id", ""), - data=lb.get("data", b""), - ) - return KnowledgeRequest( operation=data.get("operation"), id=data.get("id"), @@ -91,8 +69,6 @@ class KnowledgeRequestTranslator(MessageTranslator): triples=triples, graph_embeddings=graph_embeddings, document_embeddings=document_embeddings, - library_metadata=library_metadata, - library_blob=library_blob, ) def encode(self, obj: KnowledgeRequest) -> Dict[str, Any]: @@ -149,26 +125,6 @@ class KnowledgeRequestTranslator(MessageTranslator): ], } - if obj.library_metadata: - result["library-metadata"] = { - "id": obj.library_metadata.id, - "kind": obj.library_metadata.kind, - "title": obj.library_metadata.title, - "parent-id": obj.library_metadata.parent_id, - "document-type": obj.library_metadata.document_type, - "comments": obj.library_metadata.comments, - "tags": obj.library_metadata.tags, - } - - if obj.library_blob: - data = obj.library_blob.data - if isinstance(data, bytes): - data = data.decode("utf-8") - result["library-blob"] = { - "id": obj.library_blob.id, - "data": data, - } - return result @@ -238,32 +194,6 @@ class KnowledgeResponseTranslator(MessageTranslator): } } - # Streaming library metadata response - if obj.library_metadata: - return { - "library-metadata": { - "id": obj.library_metadata.id, - "kind": obj.library_metadata.kind, - "title": obj.library_metadata.title, - "parent-id": obj.library_metadata.parent_id, - "document-type": obj.library_metadata.document_type, - "comments": obj.library_metadata.comments, - "tags": obj.library_metadata.tags, - } - } - - # Streaming library blob response - if obj.library_blob: - data = obj.library_blob.data - if isinstance(data, bytes): - data = data.decode("utf-8") - return { - "library-blob": { - "id": obj.library_blob.id, - "data": data, - } - } - # End of stream marker if obj.eos is True: return {"eos": True} @@ -279,9 +209,7 @@ class KnowledgeResponseTranslator(MessageTranslator): is_final = ( obj.ids is not None or # List response obj.eos is True or # End of stream - (not obj.triples and not obj.graph_embeddings - and not obj.document_embeddings - and not obj.library_metadata and not obj.library_blob) # Empty response + (not obj.triples and not obj.graph_embeddings and not obj.document_embeddings) # Empty response ) return response, is_final \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/knowledge/knowledge.py b/trustgraph-base/trustgraph/schema/knowledge/knowledge.py index 4353065b..a3879103 100644 --- a/trustgraph-base/trustgraph/schema/knowledge/knowledge.py +++ b/trustgraph-base/trustgraph/schema/knowledge/knowledge.py @@ -21,21 +21,6 @@ from .embeddings import GraphEmbeddings, DocumentEmbeddings # <- () # <- (error) -@dataclass -class LibraryMetadata: - id: str = "" - kind: str = "" - title: str = "" - parent_id: str = "" - document_type: str = "" - comments: str = "" - tags: list[str] = field(default_factory=list) - -@dataclass -class LibraryBlob: - id: str = "" - data: bytes = b"" - @dataclass class KnowledgeRequest: # get-kg-core, delete-kg-core, list-kg-cores, put-kg-core @@ -59,10 +44,6 @@ class KnowledgeRequest: # put-de-core document_embeddings: DocumentEmbeddings | None = None - # put-kg-core (source material) - library_metadata: LibraryMetadata | None = None - library_blob: LibraryBlob | None = None - @dataclass class KnowledgeResponse: error: Error | None = None @@ -71,8 +52,6 @@ class KnowledgeResponse: triples: Triples | None = None graph_embeddings: GraphEmbeddings | None = None document_embeddings: DocumentEmbeddings | None = None - library_metadata: LibraryMetadata | None = None - library_blob: LibraryBlob | None = None knowledge_request_queue = queue('knowledge', cls='request') knowledge_response_queue = queue('knowledge', cls='response') diff --git a/trustgraph-cli/trustgraph/cli/get_kg_core.py b/trustgraph-cli/trustgraph/cli/get_kg_core.py index 2ff1a3cc..b4f37b81 100644 --- a/trustgraph-cli/trustgraph/cli/get_kg_core.py +++ b/trustgraph-cli/trustgraph/cli/get_kg_core.py @@ -47,31 +47,6 @@ def write_ge(f, data): ) f.write(msgpack.packb(msg, use_bin_type=True)) -def write_library_metadata(f, data): - msg = ( - "lm", - { - "i": data["id"], - "k": data.get("kind", ""), - "t": data.get("title", ""), - "p": data.get("parent-id", ""), - "d": data.get("document-type", ""), - "c": data.get("comments", ""), - "g": data.get("tags", []), - } - ) - f.write(msgpack.packb(msg, use_bin_type=True)) - -def write_library_blob(f, data): - msg = ( - "lb", - { - "i": data["id"], - "d": data.get("data", b""), - } - ) - f.write(msgpack.packb(msg, use_bin_type=True)) - def fetch(url, workspace, id, output, token=None): api = Api(url=url, token=token, workspace=workspace) @@ -80,8 +55,6 @@ def fetch(url, workspace, id, output, token=None): try: ge = 0 t = 0 - lm = 0 - lb = 0 with open(output, "wb") as f: @@ -95,15 +68,7 @@ def fetch(url, workspace, id, output, token=None): ge += 1 write_ge(f, response["graph-embeddings"]) - if "library-metadata" in response: - lm += 1 - write_library_metadata(f, response["library-metadata"]) - - if "library-blob" in response: - lb += 1 - write_library_blob(f, response["library-blob"]) - - print(f"Got: {t} triple, {ge} GE, {lm} library metadata, {lb} library blob messages.") + print(f"Got: {t} triple, {ge} GE messages.") finally: socket.close() diff --git a/trustgraph-cli/trustgraph/cli/load_structured_data.py b/trustgraph-cli/trustgraph/cli/load_structured_data.py index 5649a5ae..dccf548e 100644 --- a/trustgraph-cli/trustgraph/cli/load_structured_data.py +++ b/trustgraph-cli/trustgraph/cli/load_structured_data.py @@ -78,7 +78,7 @@ def load_structured_data( logger.info("Step 1: Analyzing data to discover best matching schema...") # Step 1: Auto-discover schema (reuse discover_schema logic) - discovered_schema = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, token=token, workspace=workspace) + discovered_schema = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, workspace=workspace) if not discovered_schema: logger.error("Failed to discover suitable schema automatically") print("❌ Could not automatically determine the best schema for your data.") @@ -90,7 +90,7 @@ def load_structured_data( # Step 2: Auto-generate descriptor logger.info("Step 2: Generating descriptor configuration...") - auto_descriptor = _auto_generate_descriptor(api_url, input_file, discovered_schema, sample_chars, flow, logger, token=token, workspace=workspace) + auto_descriptor = _auto_generate_descriptor(api_url, input_file, discovered_schema, sample_chars, flow, logger, workspace=workspace) if not auto_descriptor: logger.error("Failed to generate descriptor automatically") print("❌ Could not automatically generate descriptor configuration.") @@ -172,7 +172,7 @@ def load_structured_data( logger.info(f"Sample chars: {sample_chars} characters") # Use the helper function to discover schema (get raw response for display) - response = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, return_raw_response=True, token=token, workspace=workspace) + response = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, return_raw_response=True, workspace=workspace) if response: # Debug: print response type and content @@ -203,7 +203,7 @@ def load_structured_data( # If no schema specified, discover it first if not schema_name: logger.info("No schema specified, auto-discovering...") - schema_name = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, token=token, workspace=workspace) + schema_name = _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, workspace=workspace) if not schema_name: print("Error: Could not determine schema automatically.") print("Please specify a schema using --schema-name or run --discover-schema first.") @@ -213,7 +213,7 @@ def load_structured_data( logger.info(f"Target schema: {schema_name}") # Generate descriptor using helper function - descriptor = _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, flow, logger, token=token, workspace=workspace) + descriptor = _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, flow, logger, workspace=workspace) if descriptor: # Output the generated descriptor @@ -603,7 +603,7 @@ def _send_to_trustgraph(rows, api_url, flow, batch_size=1000, token=None, worksp # Helper functions for auto mode -def _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, return_raw_response=False, token=None, workspace="default"): +def _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, return_raw_response=False, workspace="default"): """Auto-discover the best matching schema for the input data Args: @@ -626,7 +626,7 @@ def _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, retur # Import API modules from trustgraph.api import Api from trustgraph.api.types import ConfigKey - api = Api(api_url, token=token, workspace=workspace) + api = Api(api_url, workspace=workspace) config_api = api.config() # Get available schemas @@ -707,7 +707,7 @@ def _auto_discover_schema(api_url, input_file, sample_chars, flow, logger, retur return None -def _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, flow, logger, token=None, workspace="default"): +def _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, flow, logger, workspace="default"): """Auto-generate descriptor configuration for the discovered schema""" try: # Read sample data @@ -717,7 +717,7 @@ def _auto_generate_descriptor(api_url, input_file, schema_name, sample_chars, fl # Import API modules from trustgraph.api import Api from trustgraph.api.types import ConfigKey - api = Api(api_url, token=token, workspace=workspace) + api = Api(api_url, workspace=workspace) config_api = api.config() # Get schema definition diff --git a/trustgraph-cli/trustgraph/cli/put_kg_core.py b/trustgraph-cli/trustgraph/cli/put_kg_core.py index f4e0b3dd..fe0981a5 100644 --- a/trustgraph-cli/trustgraph/cli/put_kg_core.py +++ b/trustgraph-cli/trustgraph/cli/put_kg_core.py @@ -40,23 +40,6 @@ def read_message(unpacked, id): }, "triples": msg["t"], } - elif unpacked[0] == "lm": - msg = unpacked[1] - return "lm", { - "id": msg["i"], - "kind": msg.get("k", ""), - "title": msg.get("t", ""), - "parent-id": msg.get("p", ""), - "document-type": msg.get("d", ""), - "comments": msg.get("c", ""), - "tags": msg.get("g", []), - } - elif unpacked[0] == "lb": - msg = unpacked[1] - return "lb", { - "id": msg["i"], - "data": msg.get("d", b""), - } else: raise RuntimeError("Unpacked unexpected messsage type", unpacked[0]) @@ -68,8 +51,6 @@ def put(url, workspace, id, input, token=None): try: ge = 0 t = 0 - lm = 0 - lb = 0 with open(input, "rb") as f: @@ -92,18 +73,10 @@ def put(url, workspace, id, input, token=None): t += 1 socket.put_kg_core(id, triples=msg) - elif kind == "lm": - lm += 1 - socket.put_kg_core(id, library_metadata=msg) - - elif kind == "lb": - lb += 1 - socket.put_kg_core(id, library_blob=msg) - else: raise RuntimeError("Unexpected message kind", kind) - print(f"Put: {t} triple, {ge} GE, {lm} library metadata, {lb} library blob messages.") + print(f"Put: {t} triple, {ge} GE messages.") finally: socket.close() diff --git a/trustgraph-flow/trustgraph/config/service/service.py b/trustgraph-flow/trustgraph/config/service/service.py index 725f1106..c5fac198 100644 --- a/trustgraph-flow/trustgraph/config/service/service.py +++ b/trustgraph-flow/trustgraph/config/service/service.py @@ -83,8 +83,7 @@ class Processor(AsyncProcessor): host=cassandra_host, username=cassandra_username, password=cassandra_password, - default_keyspace="config", - replication_factor=params.get("cassandra_replication_factor"), + default_keyspace="config" ) # Store resolved configuration diff --git a/trustgraph-flow/trustgraph/cores/knowledge.py b/trustgraph-flow/trustgraph/cores/knowledge.py index 6f017c43..f1fa53f5 100644 --- a/trustgraph-flow/trustgraph/cores/knowledge.py +++ b/trustgraph-flow/trustgraph/cores/knowledge.py @@ -1,7 +1,6 @@ from .. schema import KnowledgeResponse, Error, Triples, GraphEmbeddings -from .. schema import DocumentEmbeddings, LibraryMetadata, LibraryBlob -from .. schema import LibrarianRequest, DocumentMetadata +from .. schema import DocumentEmbeddings from .. knowledge import hash from .. exceptions import RequestError from .. tables.knowledge import KnowledgeTableStore @@ -19,7 +18,7 @@ class KnowledgeManager: def __init__( self, cassandra_host, cassandra_username, cassandra_password, - keyspace, flow_config, librarian=None, replication_factor=1, + keyspace, flow_config, replication_factor=1, ): self.table_store = KnowledgeTableStore( @@ -27,9 +26,6 @@ class KnowledgeManager: replication_factor ) - self.librarian = librarian - self._pending_library_metadata = {} - self.loader_queue = asyncio.Queue(maxsize=20) self.background_task = None self.flow_config = flow_config @@ -90,9 +86,6 @@ class KnowledgeManager: publish_ge, ) - if self.librarian: - await self._stream_library_docs(request.id, respond) - logger.debug("Knowledge core retrieval complete") await respond( @@ -129,12 +122,6 @@ class KnowledgeManager: workspace, request.graph_embeddings ) - if request.library_metadata and self.librarian: - await self._put_library_metadata(request.library_metadata, workspace) - - if request.library_blob and self.librarian: - await self._put_library_blob(request.library_blob, workspace) - await respond( KnowledgeResponse( error = None, @@ -263,112 +250,6 @@ class KnowledgeManager: await self.loader_queue.put((request, respond, workspace)) - async def _stream_library_docs(self, document_id, respond): - - try: - root_meta = await self.librarian.fetch_document_metadata( - document_id - ) - except Exception as e: - logger.warning(f"Could not fetch library metadata for {document_id}: {e}") - return - - if root_meta is None: - return - - await self._stream_one_doc(root_meta, respond) - - try: - resp = await self.librarian.request( - LibrarianRequest( - operation="list-children", - document_id=document_id, - ) - ) - except Exception as e: - logger.warning(f"Could not list children for {document_id}: {e}") - return - - for child_meta in resp.document_metadatas: - await self._stream_one_doc(child_meta, respond) - - async def _stream_one_doc(self, doc_meta, respond): - - lm = LibraryMetadata( - id=doc_meta.id, - kind=doc_meta.kind, - title=doc_meta.title, - parent_id=doc_meta.parent_id, - document_type=doc_meta.document_type, - comments=doc_meta.comments, - tags=doc_meta.tags or [], - ) - - await respond( - KnowledgeResponse(library_metadata=lm) - ) - - try: - content = await self.librarian.fetch_document_content( - doc_meta.id - ) - except Exception as e: - logger.warning(f"Could not fetch content for {doc_meta.id}: {e}") - return - - await respond( - KnowledgeResponse( - library_blob=LibraryBlob( - id=doc_meta.id, - data=content, - ) - ) - ) - - async def _put_library_metadata(self, lm, workspace): - self._pending_library_metadata[lm.id] = lm - - async def _put_library_blob(self, lb, workspace): - - lm = self._pending_library_metadata.pop(lb.id, None) - if lm is None: - logger.warning( - f"Received library blob for {lb.id} with no preceding metadata" - ) - return - - doc_meta = DocumentMetadata( - id=lm.id, - kind=lm.kind, - title=lm.title, - parent_id=lm.parent_id, - document_type=lm.document_type, - comments=lm.comments, - tags=lm.tags or [], - ) - - if lm.parent_id: - operation = "add-child-document" - else: - operation = "add-document" - - try: - await self.librarian.request( - LibrarianRequest( - operation=operation, - document_id=lm.id, - document_metadata=doc_meta, - content=lb.data, - ) - ) - except RuntimeError as e: - if "already exists" in str(e): - logger.debug(f"Library document {lm.id} already exists, skipping") - else: - logger.warning(f"Could not save library document {lm.id}: {e}") - except Exception as e: - logger.warning(f"Could not save library document {lm.id}: {e}") - async def core_loader(self): logger.info("Knowledge background processor running...") diff --git a/trustgraph-flow/trustgraph/cores/service.py b/trustgraph-flow/trustgraph/cores/service.py index 5c50c207..a04e42ca 100755 --- a/trustgraph-flow/trustgraph/cores/service.py +++ b/trustgraph-flow/trustgraph/cores/service.py @@ -12,7 +12,6 @@ import logging from .. base import WorkspaceProcessor, Consumer, Producer, Publisher, Subscriber from .. base import ConsumerMetrics, ProducerMetrics from .. base.cassandra_config import add_cassandra_args, resolve_cassandra_config -from .. base import LibrarianClient from .. schema import KnowledgeRequest, KnowledgeResponse, Error from .. schema import knowledge_request_queue, knowledge_response_queue @@ -61,8 +60,7 @@ class Processor(WorkspaceProcessor): host=cassandra_host, username=cassandra_username, password=cassandra_password, - default_keyspace="knowledge", - replication_factor=params.get("cassandra_replication_factor"), + default_keyspace="knowledge" ) self.cassandra_host = hosts @@ -79,17 +77,12 @@ class Processor(WorkspaceProcessor): } ) - self.librarian_client = LibrarianClient( - id=id, backend=self.pubsub, taskgroup=self.taskgroup, - ) - self.knowledge = KnowledgeManager( cassandra_host = self.cassandra_host, cassandra_username = self.cassandra_username, cassandra_password = self.cassandra_password, keyspace = keyspace, flow_config = self, - librarian = self.librarian_client, replication_factor = replication_factor, ) @@ -163,7 +156,6 @@ class Processor(WorkspaceProcessor): async def start(self): await super(Processor, self).start() - await self.librarian_client.start() async def on_knowledge_config(self, workspace, config, version): diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index ae393028..ca242265 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -32,10 +32,6 @@ logger = logging.getLogger(__name__) default_ident = "document-decoder" -def _looks_like_pdf(content): - return content.lstrip().startswith(b"%PDF-") - - class Processor(FlowProcessor): def __init__(self, **params): @@ -98,37 +94,33 @@ class Processor(FlowProcessor): ) return - # Check if we should fetch from librarian or use inline data - if v.document_id: - # Fetch from librarian via Pulsar - logger.info(f"Fetching document {v.document_id} from librarian...") - - content = await flow.librarian.fetch_document_content( - document_id=v.document_id, - - ) - - # Content is base64 encoded - if isinstance(content, str): - content = content.encode('utf-8') - decoded_content = base64.b64decode(content) - - logger.info(f"Fetched {len(decoded_content)} bytes from librarian") - else: - # Use inline data (backward compatibility) - decoded_content = base64.b64decode(v.data) - - if not _looks_like_pdf(decoded_content): - logger.error( - f"Document {v.metadata.id} is not valid PDF content. " - f"Ignoring document." - ) - return - - with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as fp: + with tempfile.NamedTemporaryFile(delete_on_close=False, suffix='.pdf') as fp: temp_path = fp.name - fp.write(decoded_content) - fp.close() + + # Check if we should fetch from librarian or use inline data + if v.document_id: + # Fetch from librarian via Pulsar + logger.info(f"Fetching document {v.document_id} from librarian...") + fp.close() + + content = await flow.librarian.fetch_document_content( + document_id=v.document_id, + + ) + + # Content is base64 encoded + if isinstance(content, str): + content = content.encode('utf-8') + decoded_content = base64.b64decode(content) + + with open(temp_path, 'wb') as f: + f.write(decoded_content) + + logger.info(f"Fetched {len(decoded_content)} bytes from librarian") + else: + # Use inline data (backward compatibility) + fp.write(base64.b64decode(v.data)) + fp.close() global PyPDFLoader if PyPDFLoader is None: diff --git a/trustgraph-flow/trustgraph/direct/cassandra_kg.py b/trustgraph-flow/trustgraph/direct/cassandra_kg.py index f1e4a577..d7abd1a9 100644 --- a/trustgraph-flow/trustgraph/direct/cassandra_kg.py +++ b/trustgraph-flow/trustgraph/direct/cassandra_kg.py @@ -6,7 +6,7 @@ import logging from cassandra.cluster import Cluster from cassandra.auth import PlainTextAuthProvider from cassandra.query import BatchStatement, SimpleStatement -import ssl +from ssl import SSLContext, PROTOCOL_TLSv1_2 from ..tables.cassandra_async import async_execute @@ -41,15 +41,13 @@ class KnowledgeGraph: def __init__( self, hosts=None, - keyspace="trustgraph", username=None, password=None, - replication_factor=1, + keyspace="trustgraph", username=None, password=None ): if hosts is None: hosts = ["localhost"] self.keyspace = keyspace - self.replication_factor = replication_factor self.username = username # 7-table schema for quads with full query pattern support @@ -70,7 +68,7 @@ class KnowledgeGraph: self.collection_metadata_table = "collection_metadata" if username and password: - ssl_context = ssl.create_default_context() + ssl_context = SSLContext(PROTOCOL_TLSv1_2) auth_provider = PlainTextAuthProvider(username=username, password=password) self.cluster = Cluster(hosts, auth_provider=auth_provider, ssl_context=ssl_context) else: @@ -94,7 +92,7 @@ class KnowledgeGraph: create keyspace if not exists {self.keyspace} with replication = {{ 'class' : 'SimpleStrategy', - 'replication_factor' : {self.replication_factor} + 'replication_factor' : 1 }}; """) @@ -541,15 +539,13 @@ class EntityCentricKnowledgeGraph: def __init__( self, hosts=None, - keyspace="trustgraph", username=None, password=None, - replication_factor=1, + keyspace="trustgraph", username=None, password=None ): if hosts is None: hosts = ["localhost"] self.keyspace = keyspace - self.replication_factor = replication_factor self.username = username # 2-table entity-centric schema @@ -560,7 +556,7 @@ class EntityCentricKnowledgeGraph: self.collection_metadata_table = "collection_metadata" if username and password: - ssl_context = ssl.create_default_context() + ssl_context = SSLContext(PROTOCOL_TLSv1_2) auth_provider = PlainTextAuthProvider(username=username, password=password) self.cluster = Cluster(hosts, auth_provider=auth_provider, ssl_context=ssl_context) else: @@ -584,7 +580,7 @@ class EntityCentricKnowledgeGraph: create keyspace if not exists {self.keyspace} with replication = {{ 'class' : 'SimpleStrategy', - 'replication_factor' : {self.replication_factor} + 'replication_factor' : 1 }}; """) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/core_export.py b/trustgraph-flow/trustgraph/gateway/dispatch/core_export.py index 90080cc4..6696afbe 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/core_export.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/core_export.py @@ -73,39 +73,6 @@ class CoreExport: enc = msgpack.packb(msg) await response.write(enc) - if "library-metadata" in resp: - - data = resp["library-metadata"] - msg = ( - "lm", - { - "i": data["id"], - "k": data.get("kind", ""), - "t": data.get("title", ""), - "p": data.get("parent-id", ""), - "d": data.get("document-type", ""), - "c": data.get("comments", ""), - "g": data.get("tags", []), - } - ) - - enc = msgpack.packb(msg) - await response.write(enc) - - if "library-blob" in resp: - - data = resp["library-blob"] - msg = ( - "lb", - { - "i": data["id"], - "d": data.get("data", b""), - } - ) - - enc = msgpack.packb(msg, use_bin_type=True) - await response.write(enc) - await kr.process( { "operation": "get-kg-core", diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/core_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/core_import.py index bf660def..d03d4efd 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/core_import.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/core_import.py @@ -79,39 +79,6 @@ class CoreImport: await kr.process(msg) - elif unpacked[0] == "lm": - msg = unpacked[1] - msg = { - "operation": "put-kg-core", - "workspace": workspace, - "id": id, - "library-metadata": { - "id": msg["i"], - "kind": msg.get("k", ""), - "title": msg.get("t", ""), - "parent-id": msg.get("p", ""), - "document-type": msg.get("d", ""), - "comments": msg.get("c", ""), - "tags": msg.get("g", []), - } - } - - await kr.process(msg) - - elif unpacked[0] == "lb": - msg = unpacked[1] - msg = { - "operation": "put-kg-core", - "workspace": workspace, - "id": id, - "library-blob": { - "id": msg["i"], - "data": msg.get("d", b""), - } - } - - await kr.process(msg) - except Exception as e: logger.error(f"Core import exception: {e}", exc_info=True) await error(str(e)) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/mux.py b/trustgraph-flow/trustgraph/gateway/dispatch/mux.py index 9b119f8e..73bbb1f3 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/mux.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/mux.py @@ -4,8 +4,6 @@ import queue import uuid import logging -from ..capabilities import PUBLIC, AUTHENTICATED - # Module logger logger = logging.getLogger(__name__) @@ -158,41 +156,37 @@ class Mux: }) return - # Resolve workspace (default-fill from the caller's - # bound workspace). Workspace resolution applies to all - # operations regardless of capability level. + # Resolve workspace first (default-fill from the caller's + # bound workspace), then ask the regime to authorise the + # service-level capability against the matched + # operation's resource shape. try: await enforce_workspace(data, self.identity, self.auth) if isinstance(inner, dict): await enforce_workspace(inner, self.identity, self.auth) - # Authorisation: capability sentinels short-circuit - # the regime call; capability strings go through - # authorise(). - if op.capability not in (PUBLIC, AUTHENTICATED): - if data.get("flow"): - resource = { - "workspace": data.get("workspace", ""), - "flow": data.get("flow", ""), - } - parameters = {} - else: - # Build a minimal RequestContext so the matched - # operation's own extractors decide resource - # and parameters — same path the HTTP - # endpoints take. - from ..registry import RequestContext - ctx = RequestContext( - body=inner if isinstance(inner, dict) else {}, - match_info={}, - identity=self.identity, - ) - resource = op.extract_resource(ctx) - parameters = op.extract_parameters(ctx) - - await self.auth.authorise( - self.identity, op.capability, resource, parameters, + if data.get("flow"): + resource = { + "workspace": data.get("workspace", ""), + "flow": data.get("flow", ""), + } + parameters = {} + else: + # Build a minimal RequestContext so the matched + # operation's own extractors decide resource and + # parameters — same path the HTTP endpoints take. + from ..registry import RequestContext + ctx = RequestContext( + body=inner if isinstance(inner, dict) else {}, + match_info={}, + identity=self.identity, ) + resource = op.extract_resource(ctx) + parameters = op.extract_parameters(ctx) + + await self.auth.authorise( + self.identity, op.capability, resource, parameters, + ) except _web.HTTPNotFound: await self.ws.send_json({ "id": request_id, diff --git a/trustgraph-flow/trustgraph/iam/service/service.py b/trustgraph-flow/trustgraph/iam/service/service.py index b2f3976d..8ce22757 100644 --- a/trustgraph-flow/trustgraph/iam/service/service.py +++ b/trustgraph-flow/trustgraph/iam/service/service.py @@ -101,7 +101,6 @@ class Processor(AsyncProcessor): username=cassandra_username, password=cassandra_password, default_keyspace="iam", - replication_factor=params.get("cassandra_replication_factor"), ) self.cassandra_host = hosts diff --git a/trustgraph-flow/trustgraph/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py index 4d3efbfb..cc5efdae 100755 --- a/trustgraph-flow/trustgraph/librarian/service.py +++ b/trustgraph-flow/trustgraph/librarian/service.py @@ -8,7 +8,6 @@ import asyncio import base64 import json import logging -import os from datetime import datetime from .. base import WorkspaceProcessor, Consumer, Producer, Publisher, Subscriber @@ -55,16 +54,6 @@ default_object_store_access_key = "object-user" default_object_store_secret_key = "object-password" default_object_store_use_ssl = False default_object_store_region = None - -# Environment variables consulted as a fallback when the -# corresponding params field is not set in the processor-group YAML -# or via CLI. Intended for K8s Secret / env-var injection so -# credentials never have to live in the YAML (and thus in git). -ENV_OBJECT_STORE_ENDPOINT = "OBJECT_STORE_ENDPOINT" -ENV_OBJECT_STORE_ACCESS_KEY = "OBJECT_STORE_ACCESS_KEY" -ENV_OBJECT_STORE_SECRET_KEY = "OBJECT_STORE_SECRET_KEY" -ENV_OBJECT_STORE_USE_SSL = "OBJECT_STORE_USE_SSL" -ENV_OBJECT_STORE_REGION = "OBJECT_STORE_REGION" default_cassandra_host = "cassandra" default_min_chunk_size = 1 # No minimum by default (for Garage) @@ -100,36 +89,22 @@ class Processor(WorkspaceProcessor): "config_response_queue", default_config_response_queue ) - # Resolve object-store config. Precedence: explicit params - # (CLI / processor-group YAML) → environment variable → - # hardcoded default. The env-var path lets K8s Secrets feed - # credentials without them appearing in the YAML. - object_store_endpoint = ( - params.get("object_store_endpoint") - or os.environ.get(ENV_OBJECT_STORE_ENDPOINT) - or default_object_store_endpoint + object_store_endpoint = params.get("object_store_endpoint", default_object_store_endpoint) + object_store_access_key = params.get( + "object_store_access_key", + default_object_store_access_key ) - object_store_access_key = ( - params.get("object_store_access_key") - or os.environ.get(ENV_OBJECT_STORE_ACCESS_KEY) - or default_object_store_access_key + object_store_secret_key = params.get( + "object_store_secret_key", + default_object_store_secret_key ) - object_store_secret_key = ( - params.get("object_store_secret_key") - or os.environ.get(ENV_OBJECT_STORE_SECRET_KEY) - or default_object_store_secret_key + object_store_use_ssl = params.get( + "object_store_use_ssl", + default_object_store_use_ssl ) - object_store_use_ssl = params.get("object_store_use_ssl") - if object_store_use_ssl is None: - env_ssl = os.environ.get(ENV_OBJECT_STORE_USE_SSL) - if env_ssl is not None: - object_store_use_ssl = env_ssl.lower() in ("true", "1", "yes") - else: - object_store_use_ssl = default_object_store_use_ssl - object_store_region = ( - params.get("object_store_region") - or os.environ.get(ENV_OBJECT_STORE_REGION) - or default_object_store_region + object_store_region = params.get( + "object_store_region", + default_object_store_region ) min_chunk_size = params.get( @@ -146,8 +121,7 @@ class Processor(WorkspaceProcessor): host=cassandra_host, username=cassandra_username, password=cassandra_password, - default_keyspace="librarian", - replication_factor=params.get("cassandra_replication_factor"), + default_keyspace="librarian" ) # Store resolved configuration diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py index de25a139..f6770744 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py @@ -12,33 +12,31 @@ from qdrant_client import QdrantClient from .... schema import DocumentEmbeddingsResponse, ChunkMatch from .... schema import Error from .... base import DocumentEmbeddingsQueryService -from .... base.qdrant_config import add_qdrant_args, resolve_qdrant_config # Module logger logger = logging.getLogger(__name__) default_ident = "doc-embeddings-query" +default_store_uri = 'http://localhost:6333' + class Processor(DocumentEmbeddingsQueryService): def __init__(self, **params): - store_uri = params.get("store_uri") - api_key = params.get("api_key") + store_uri = params.get("store_uri", default_store_uri) - url, api_key, _, _ = resolve_qdrant_config( - url=store_uri, - api_key=api_key, - ) + #optional api key + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { - "store_uri": url, + "store_uri": store_uri, "api_key": api_key, } ) - self.qdrant = QdrantClient(url=url, api_key=api_key) + self.qdrant = QdrantClient(url=store_uri, api_key=api_key) async def query_document_embeddings(self, workspace, msg): @@ -87,7 +85,18 @@ class Processor(DocumentEmbeddingsQueryService): def add_args(parser): DocumentEmbeddingsQueryService.add_args(parser) - add_qdrant_args(parser) + + parser.add_argument( + '-t', '--store-uri', + default=default_store_uri, + help=f'Qdrant store URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'API key for qdrant (default: None)' + ) def run(): diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index aa93925d..167130c9 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py @@ -12,32 +12,31 @@ from qdrant_client import QdrantClient from .... schema import GraphEmbeddingsResponse, EntityMatch from .... schema import Error, Term, IRI, LITERAL from .... base import GraphEmbeddingsQueryService -from .... base.qdrant_config import add_qdrant_args, resolve_qdrant_config # Module logger logger = logging.getLogger(__name__) default_ident = "graph-embeddings-query" +default_store_uri = 'http://localhost:6333' + class Processor(GraphEmbeddingsQueryService): def __init__(self, **params): - store_uri = params.get("store_uri") - api_key = params.get("api_key") + store_uri = params.get("store_uri", default_store_uri) - url, api_key, _, _ = resolve_qdrant_config( - url=store_uri, api_key=api_key, - ) + #optional api key + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { - "store_uri": url, + "store_uri": store_uri, "api_key": api_key, } ) - self.qdrant = QdrantClient(url=url, api_key=api_key) + self.qdrant = QdrantClient(url=store_uri, api_key=api_key) def create_value(self, ent): if ent.startswith("http://") or ent.startswith("https://"): @@ -105,7 +104,18 @@ class Processor(GraphEmbeddingsQueryService): def add_args(parser): GraphEmbeddingsQueryService.add_args(parser) - add_qdrant_args(parser) + + parser.add_argument( + '-t', '--store-uri', + default=default_store_uri, + help=f'Qdrant store URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'API key for qdrant (default: None)' + ) def run(): diff --git a/trustgraph-flow/trustgraph/query/ontology/sparql_cassandra.py b/trustgraph-flow/trustgraph/query/ontology/sparql_cassandra.py index a9005ee4..b7f0f423 100644 --- a/trustgraph-flow/trustgraph/query/ontology/sparql_cassandra.py +++ b/trustgraph-flow/trustgraph/query/ontology/sparql_cassandra.py @@ -116,7 +116,7 @@ class CassandraTripleStore(Store if RDFLIB_AVAILABLE else object): # Create keyspace self.session.execute(f""" CREATE KEYSPACE IF NOT EXISTS {self.keyspace} - WITH replication = {{'class': 'SimpleStrategy', 'replication_factor': {self.cassandra_config.get('replication_factor', 1)}}} + WITH replication = {{'class': 'SimpleStrategy', 'replication_factor': 1}} """) # Create triples table optimized for SPARQL queries diff --git a/trustgraph-flow/trustgraph/query/row_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/row_embeddings/qdrant/service.py index 7e1a5851..1534c044 100644 --- a/trustgraph-flow/trustgraph/query/row_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/row_embeddings/qdrant/service.py @@ -19,12 +19,12 @@ from .... schema import ( RowIndexMatch, Error ) from .... base import FlowProcessor, ConsumerSpec, ProducerSpec -from .... base.qdrant_config import add_qdrant_args, resolve_qdrant_config # Module logger logger = logging.getLogger(__name__) default_ident = "row-embeddings-query" +default_store_uri = 'http://localhost:6333' default_concurrency = 10 @@ -35,17 +35,13 @@ class Processor(FlowProcessor): id = params.get("id", default_ident) concurrency = params.get("concurrency", default_concurrency) - store_uri = params.get("store_uri") - api_key = params.get("api_key") - - url, api_key, _, _ = resolve_qdrant_config( - url=store_uri, api_key=api_key, - ) + store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { "id": id, - "store_uri": url, + "store_uri": store_uri, "api_key": api_key, } ) @@ -66,7 +62,7 @@ class Processor(FlowProcessor): ) ) - self.qdrant = QdrantClient(url=url, api_key=api_key) + self.qdrant = QdrantClient(url=store_uri, api_key=api_key) def sanitize_name(self, name: str) -> str: """Sanitize names for Qdrant collection naming""" @@ -196,9 +192,21 @@ class Processor(FlowProcessor): @staticmethod def add_args(parser): + """Add command-line arguments""" FlowProcessor.add_args(parser) - add_qdrant_args(parser) + + parser.add_argument( + '-t', '--store-uri', + default=default_store_uri, + help=f'Qdrant store URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help='API key for Qdrant (default: None)' + ) parser.add_argument( '-c', '--concurrency', diff --git a/trustgraph-flow/trustgraph/query/rows/cassandra/service.py b/trustgraph-flow/trustgraph/query/rows/cassandra/service.py index f9868d67..7157daae 100644 --- a/trustgraph-flow/trustgraph/query/rows/cassandra/service.py +++ b/trustgraph-flow/trustgraph/query/rows/cassandra/service.py @@ -24,7 +24,7 @@ from .... schema import RowsQueryRequest, RowsQueryResponse, GraphQLError from .... schema import Error, RowSchema, Field as SchemaField from .... base import FlowProcessor, ConsumerSpec, ProducerSpec from .... base.cassandra_config import add_cassandra_args, resolve_cassandra_config -from .... tables.cassandra_async import async_execute, async_execute_paged, async_scan +from .... tables.cassandra_async import async_execute from ... graphql import GraphQLSchemaBuilder, SortDirection @@ -180,7 +180,7 @@ class Processor(FlowProcessor): description=field_def.get("description", ""), required=field_def.get("required", False), enum_values=field_def.get("enum", []), - indexed=field_def.get("indexed", False), + indexed=field_def.get("indexed", False) ) fields.append(field) @@ -232,8 +232,6 @@ class Processor(FlowProcessor): for index_name in index_names: if index_name in filters: value = filters[index_name] - if value == "" or value is None: - continue # Single field index -> single element list index_value = [str(value)] return (index_name, index_value) @@ -284,13 +282,11 @@ class Processor(FlowProcessor): query += f" LIMIT {limit}" try: - pages = await async_execute_paged( - self.session, query, params - ) - for page in pages: - for row in page: - row_dict = dict(row.data) if row.data else {} - results.append(row_dict) + rows = await async_execute(self.session, query, params) + for row in rows: + # Convert data map to dict with proper field names + row_dict = dict(row.data) if row.data else {} + results.append(row_dict) except Exception as e: logger.error(f"Failed to query rows: {e}", exc_info=True) raise @@ -312,6 +308,8 @@ class Processor(FlowProcessor): # Query using the first index (arbitrary choice for scan) primary_index = index_names[0] + # We need to scan all values for this index + # This requires ALLOW FILTERING or a different approach query = f""" SELECT data, source FROM {safe_keyspace}.rows WHERE collection = %s @@ -322,18 +320,17 @@ class Processor(FlowProcessor): params = [collection, schema_name, primary_index] try: - def row_filter(row): - row_dict = dict(row.data) if row.data else {} - return self._matches_filters(row_dict, filters, row_schema) + rows = await async_execute(self.session, query, params) - matched_rows = await async_scan( - self.session, query, params, - row_filter=row_filter, - limit=limit, - ) - for row in matched_rows: + for row in rows: row_dict = dict(row.data) if row.data else {} - results.append(row_dict) + + # Apply post-filters + if self._matches_filters(row_dict, filters, row_schema): + results.append(row_dict) + + if limit and len(results) >= limit: + break except Exception as e: logger.error(f"Failed to scan rows: {e}", exc_info=True) @@ -366,7 +363,7 @@ class Processor(FlowProcessor): # Parse filter key for operator if '_' in filter_key: parts = filter_key.rsplit('_', 1) - if parts[1] in ['gt', 'gte', 'lt', 'lte', 'contains', 'in', 'not', 'startsWith', 'endsWith', 'not_in']: + if parts[1] in ['gt', 'gte', 'lt', 'lte', 'contains', 'in']: field_name = parts[0] operator = parts[1] else: @@ -403,18 +400,6 @@ class Processor(FlowProcessor): elif operator == 'in': if str(row_value) not in [str(v) for v in filter_value]: return False - elif operator == 'not': - if str(row_value) == str(filter_value): - return False - elif operator == 'startsWith': - if not str(row_value).startswith(str(filter_value)): - return False - elif operator == 'endsWith': - if not str(row_value).endswith(str(filter_value)): - return False - elif operator == 'not_in': - if str(row_value) in [str(v) for v in filter_value]: - return False except (ValueError, TypeError): return False diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index 08d88849..2bfef99c 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -14,36 +14,29 @@ from qdrant_client.models import Distance, VectorParams from .... base import DocumentEmbeddingsStoreService, CollectionConfigHandler from .... base import AsyncProcessor, Consumer, Producer from .... base import ConsumerMetrics, ProducerMetrics -from .... base.qdrant_config import add_qdrant_args, resolve_qdrant_config # Module logger logger = logging.getLogger(__name__) default_ident = "doc-embeddings-write" +default_store_uri = 'http://localhost:6333' + class Processor(CollectionConfigHandler, DocumentEmbeddingsStoreService): def __init__(self, **params): - store_uri = params.get("store_uri") - api_key = params.get("api_key") - - url, api_key, replication_factor, shard_number = resolve_qdrant_config( - url=store_uri, api_key=api_key, - replication_factor=params.get("qdrant_replication_factor"), - shard_number=params.get("qdrant_shard_number"), - ) + store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { - "store_uri": url, + "store_uri": store_uri, "api_key": api_key, } ) - self.qdrant = QdrantClient(url=url, api_key=api_key) - self.replication_factor = replication_factor - self.shard_number = shard_number + self.qdrant = QdrantClient(url=store_uri, api_key=api_key) self._cache_lock = asyncio.Lock() self._known_collections: set[str] = set() @@ -68,8 +61,6 @@ class Processor(CollectionConfigHandler, DocumentEmbeddingsStoreService): vectors_config=VectorParams( size=dim, distance=Distance.COSINE ), - replication_factor=self.replication_factor, - shard_number=self.shard_number, ) self._known_collections.add(collection_name) @@ -118,7 +109,18 @@ class Processor(CollectionConfigHandler, DocumentEmbeddingsStoreService): def add_args(parser): DocumentEmbeddingsStoreService.add_args(parser) - add_qdrant_args(parser) + + parser.add_argument( + '-t', '--store-uri', + default=default_store_uri, + help=f'Qdrant URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'Qdrant API key (default: None)' + ) async def create_collection(self, workspace: str, collection: str, metadata: dict): """ diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index b6072bdc..13dcdba8 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -14,7 +14,6 @@ from qdrant_client.models import Distance, VectorParams from .... base import GraphEmbeddingsStoreService, CollectionConfigHandler from .... base import AsyncProcessor, Consumer, Producer from .... base import ConsumerMetrics, ProducerMetrics -from .... base.qdrant_config import add_qdrant_args, resolve_qdrant_config from .... schema import IRI, LITERAL # Module logger @@ -30,34 +29,29 @@ def get_term_value(term): elif term.type == LITERAL: return term.value else: + # For blank nodes or other types, use id or value return term.id or term.value default_ident = "graph-embeddings-write" +default_store_uri = 'http://localhost:6333' + class Processor(CollectionConfigHandler, GraphEmbeddingsStoreService): def __init__(self, **params): - store_uri = params.get("store_uri") - api_key = params.get("api_key") - - url, api_key, replication_factor, shard_number = resolve_qdrant_config( - url=store_uri, api_key=api_key, - replication_factor=params.get("qdrant_replication_factor"), - shard_number=params.get("qdrant_shard_number"), - ) + store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { - "store_uri": url, + "store_uri": store_uri, "api_key": api_key, } ) - self.qdrant = QdrantClient(url=url, api_key=api_key) - self.replication_factor = replication_factor - self.shard_number = shard_number + self.qdrant = QdrantClient(url=store_uri, api_key=api_key) self._cache_lock = asyncio.Lock() self._known_collections: set[str] = set() @@ -82,8 +76,6 @@ class Processor(CollectionConfigHandler, GraphEmbeddingsStoreService): vectors_config=VectorParams( size=dim, distance=Distance.COSINE ), - replication_factor=self.replication_factor, - shard_number=self.shard_number, ) self._known_collections.add(collection_name) @@ -136,7 +128,18 @@ class Processor(CollectionConfigHandler, GraphEmbeddingsStoreService): def add_args(parser): GraphEmbeddingsStoreService.add_args(parser) - add_qdrant_args(parser) + + parser.add_argument( + '-t', '--store-uri', + default=default_store_uri, + help=f'Qdrant store URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'Qdrant API key' + ) async def create_collection(self, workspace: str, collection: str, metadata: dict): """ diff --git a/trustgraph-flow/trustgraph/storage/knowledge/store.py b/trustgraph-flow/trustgraph/storage/knowledge/store.py index f6e12a85..162a4057 100644 --- a/trustgraph-flow/trustgraph/storage/knowledge/store.py +++ b/trustgraph-flow/trustgraph/storage/knowledge/store.py @@ -27,8 +27,7 @@ class Processor(FlowProcessor): host=params.get("cassandra_host"), username=params.get("cassandra_username"), password=params.get("cassandra_password"), - default_keyspace='knowledge', - replication_factor=params.get("cassandra_replication_factor"), + default_keyspace='knowledge' ) super(Processor, self).__init__( diff --git a/trustgraph-flow/trustgraph/storage/row_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/row_embeddings/qdrant/write.py index 4c65edb1..a01629c5 100644 --- a/trustgraph-flow/trustgraph/storage/row_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/row_embeddings/qdrant/write.py @@ -27,12 +27,12 @@ from qdrant_client.models import PointStruct, Distance, VectorParams from .... schema import RowEmbeddings from .... base import FlowProcessor, ConsumerSpec from .... base import CollectionConfigHandler -from .... base.qdrant_config import add_qdrant_args, resolve_qdrant_config # Module logger logger = logging.getLogger(__name__) default_ident = "row-embeddings-write" +default_store_uri = 'http://localhost:6333' class Processor(CollectionConfigHandler, FlowProcessor): @@ -41,19 +41,13 @@ class Processor(CollectionConfigHandler, FlowProcessor): id = params.get("id", default_ident) - store_uri = params.get("store_uri") - api_key = params.get("api_key") - - url, api_key, replication_factor, shard_number = resolve_qdrant_config( - url=store_uri, api_key=api_key, - replication_factor=params.get("qdrant_replication_factor"), - shard_number=params.get("qdrant_shard_number"), - ) + store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { "id": id, - "store_uri": url, + "store_uri": store_uri, "api_key": api_key, } ) @@ -69,9 +63,7 @@ class Processor(CollectionConfigHandler, FlowProcessor): # Register config handler for collection management self.register_config_handler(self.on_collection_config, types=["collection"]) - self.qdrant = QdrantClient(url=url, api_key=api_key) - self.replication_factor = replication_factor - self.shard_number = shard_number + self.qdrant = QdrantClient(url=store_uri, api_key=api_key) self._cache_lock = asyncio.Lock() self._known_collections: set[str] = set() @@ -111,8 +103,6 @@ class Processor(CollectionConfigHandler, FlowProcessor): size=dimension, distance=Distance.COSINE ), - replication_factor=self.replication_factor, - shard_number=self.shard_number, ) self._known_collections.add(collection_name) @@ -259,9 +249,21 @@ class Processor(CollectionConfigHandler, FlowProcessor): @staticmethod def add_args(parser): + """Add command-line arguments""" FlowProcessor.add_args(parser) - add_qdrant_args(parser) + + parser.add_argument( + '-t', '--store-uri', + default=default_store_uri, + help=f'Qdrant URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help='Qdrant API key (default: None)' + ) def run(): diff --git a/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py b/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py index 31fc41a7..65eeee06 100755 --- a/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py @@ -47,18 +47,16 @@ class Processor(CollectionConfigHandler, FlowProcessor): cassandra_password = params.get("cassandra_password") # Resolve configuration with environment variable fallback - hosts, username, password, keyspace, replication_factor = resolve_cassandra_config( + hosts, username, password, keyspace, _ = resolve_cassandra_config( host=cassandra_host, username=cassandra_username, - password=cassandra_password, - replication_factor=params.get("cassandra_replication_factor"), + password=cassandra_password ) # Store resolved configuration with proper names self.cassandra_host = hosts # Store as list self.cassandra_username = username self.cassandra_password = password - self.replication_factor = replication_factor # Config key for schemas self.config_key = params.get("config_type", "schema") @@ -172,7 +170,7 @@ class Processor(CollectionConfigHandler, FlowProcessor): description=field_def.get("description", ""), required=field_def.get("required", False), enum_values=field_def.get("enum", []), - indexed=field_def.get("indexed", False), + indexed=field_def.get("indexed", False) ) fields.append(field) @@ -234,7 +232,7 @@ class Processor(CollectionConfigHandler, FlowProcessor): CREATE KEYSPACE IF NOT EXISTS {safe_keyspace} WITH REPLICATION = {{ 'class': 'SimpleStrategy', - 'replication_factor': {self.replication_factor} + 'replication_factor': 1 }} """ diff --git a/trustgraph-flow/trustgraph/tables/cassandra_async.py b/trustgraph-flow/trustgraph/tables/cassandra_async.py index fe410a26..205ed6b9 100644 --- a/trustgraph-flow/trustgraph/tables/cassandra_async.py +++ b/trustgraph-flow/trustgraph/tables/cassandra_async.py @@ -80,14 +80,14 @@ def _set_exception_if_pending(fut, exc): fut.set_exception(exc) -async def async_execute_paged(session, query, parameters=None, fetch_size=5000): +async def async_execute_paged(session, query, parameters=None, fetch_size=100): """Execute a CQL query with page-by-page iteration. Uses synchronous session.execute() inside run_in_executor so that the driver's ResultSet paging works correctly without materialising the entire result set in memory. - Returns all pages as a list of lists. + Yields one page of rows at a time (as a list). """ loop = asyncio.get_running_loop() @@ -111,50 +111,3 @@ async def async_execute_paged(session, query, parameters=None, fetch_size=5000): return await loop.run_in_executor( None, _fetch_all_pages ) - - -async def async_scan( - session, query, parameters=None, row_filter=None, - limit=None, fetch_size=5000, -): - """Scan a CQL query page-by-page, applying a filter and limit. - - Only matching rows accumulate in memory. Each page is discarded - after processing, so peak memory is bounded by fetch_size plus - the number of matching rows (capped by limit). - - Args: - session: cassandra.cluster.Session - query: CQL statement string - parameters: bind params - row_filter: callable(row) -> bool, or None to accept all - limit: max results to return, or None for unlimited - fetch_size: rows per Cassandra page fetch - - Returns: - List of matching rows. - """ - loop = asyncio.get_running_loop() - - if isinstance(query, str): - stmt = SimpleStatement(query, fetch_size=fetch_size) - else: - stmt = query - stmt.fetch_size = fetch_size - - def _scan(): - results = [] - result_set = session.execute(stmt, parameters) - while True: - for row in result_set.current_rows: - if row_filter is None or row_filter(row): - results.append(row) - if limit and len(results) >= limit: - return results - if result_set.has_more_pages: - result_set.fetch_next_page() - else: - break - return results - - return await loop.run_in_executor(None, _scan) diff --git a/trustgraph-flow/trustgraph/tables/config.py b/trustgraph-flow/trustgraph/tables/config.py index c87cb3b5..74ceb6f4 100644 --- a/trustgraph-flow/trustgraph/tables/config.py +++ b/trustgraph-flow/trustgraph/tables/config.py @@ -4,7 +4,7 @@ from .. schema import Metadata, GraphEmbeddings from cassandra.cluster import Cluster from cassandra.auth import PlainTextAuthProvider -import ssl +from ssl import SSLContext, PROTOCOL_TLSv1_2 import uuid import time @@ -33,7 +33,7 @@ class ConfigTableStore: cassandra_host = [h.strip() for h in cassandra_host.split(',')] if cassandra_username and cassandra_password: - ssl_context = ssl.create_default_context() + ssl_context = SSLContext(PROTOCOL_TLSv1_2) auth_provider = PlainTextAuthProvider( username=cassandra_username, password=cassandra_password ) diff --git a/trustgraph-flow/trustgraph/tables/iam.py b/trustgraph-flow/trustgraph/tables/iam.py index b60e9cff..d7bf5e3d 100644 --- a/trustgraph-flow/trustgraph/tables/iam.py +++ b/trustgraph-flow/trustgraph/tables/iam.py @@ -15,7 +15,7 @@ import logging from cassandra.cluster import Cluster from cassandra.auth import PlainTextAuthProvider -import ssl +from ssl import SSLContext, PROTOCOL_TLSv1_2 from . cassandra_async import async_execute @@ -39,7 +39,7 @@ class IamTableStore: cassandra_host = [h.strip() for h in cassandra_host.split(",")] if cassandra_username and cassandra_password: - ssl_context = ssl.create_default_context() + ssl_context = SSLContext(PROTOCOL_TLSv1_2) auth_provider = PlainTextAuthProvider( username=cassandra_username, password=cassandra_password, ) diff --git a/trustgraph-flow/trustgraph/tables/knowledge.py b/trustgraph-flow/trustgraph/tables/knowledge.py index 53a12b35..6a23731b 100644 --- a/trustgraph-flow/trustgraph/tables/knowledge.py +++ b/trustgraph-flow/trustgraph/tables/knowledge.py @@ -23,7 +23,7 @@ def tuple_to_term(value, is_uri): else: return Term(type=LITERAL, value=value) from cassandra.auth import PlainTextAuthProvider -import ssl +from ssl import SSLContext, PROTOCOL_TLSv1_2 import uuid import time @@ -50,7 +50,7 @@ class KnowledgeTableStore: cassandra_host = [h.strip() for h in cassandra_host.split(',')] if cassandra_username and cassandra_password: - ssl_context = ssl.create_default_context() + ssl_context = SSLContext(PROTOCOL_TLSv1_2) auth_provider = PlainTextAuthProvider( username=cassandra_username, password=cassandra_password ) @@ -98,8 +98,7 @@ class KnowledgeTableStore: text, boolean, text, boolean, text, boolean >>, triples list>, PRIMARY KEY ((workspace, document_id), id) ); @@ -235,8 +234,7 @@ class KnowledgeTableStore: triples = [ ( - *term_to_tuple(v.s), *term_to_tuple(v.p), *term_to_tuple(v.o), - v.g or "" + *term_to_tuple(v.s), *term_to_tuple(v.p), *term_to_tuple(v.o) ) for v in m.triples ] @@ -418,7 +416,6 @@ class KnowledgeTableStore: s = tuple_to_term(elt[0], elt[1]), p = tuple_to_term(elt[2], elt[3]), o = tuple_to_term(elt[4], elt[5]), - g = elt[6] if elt[6] else None, ) for elt in row[3] ] diff --git a/trustgraph-flow/trustgraph/tables/library.py b/trustgraph-flow/trustgraph/tables/library.py index 5094e103..58486f0e 100644 --- a/trustgraph-flow/trustgraph/tables/library.py +++ b/trustgraph-flow/trustgraph/tables/library.py @@ -24,7 +24,7 @@ from .. exceptions import RequestError from cassandra.cluster import Cluster from cassandra.auth import PlainTextAuthProvider from cassandra.query import BatchStatement -import ssl +from ssl import SSLContext, PROTOCOL_TLSv1_2 import uuid import time @@ -53,7 +53,7 @@ class LibraryTableStore: cassandra_host = [h.strip() for h in cassandra_host.split(',')] if cassandra_username and cassandra_password: - ssl_context = ssl.create_default_context() + ssl_context = SSLContext(PROTOCOL_TLSv1_2) auth_provider = PlainTextAuthProvider( username=cassandra_username, password=cassandra_password ) diff --git a/trustgraph-mcp/trustgraph/mcp_server/mcp.py b/trustgraph-mcp/trustgraph/mcp_server/mcp.py index 11b975b2..7378db64 100755 --- a/trustgraph-mcp/trustgraph/mcp_server/mcp.py +++ b/trustgraph-mcp/trustgraph/mcp_server/mcp.py @@ -8,180 +8,71 @@ import logging import json import uuid import argparse -from dataclasses import dataclass, field +from dataclasses import dataclass from collections.abc import AsyncIterator from functools import partial from mcp.server.fastmcp import FastMCP, Context -from mcp.server.auth.provider import AccessToken, TokenVerifier -from mcp.server.auth.middleware.auth_context import get_access_token +from mcp.types import TextContent +from websockets.asyncio.client import connect from trustgraph.base.logging import add_logging_args, setup_logging -from . tg_socket import WebSocketManager, _token_key - -logger = logging.getLogger(__name__) - - -# Wire-format Term type codes (match TermTranslator compact keys) -_TERM_TYPES = { - "iri": "i", - "literal": "l", - "blank": "b", -} - - -def _make_term(value: str, term_type: str) -> dict: - """Build a compact-key Term dict for the gateway wire format. - - Args: - value: The term value (IRI string, literal text, or blank node id). - term_type: One of "iri", "literal", "blank". - """ - t = _TERM_TYPES.get(term_type) - if t is None: - raise ValueError( - f"Unknown term type '{term_type}' — " - f"expected one of: {', '.join(_TERM_TYPES)}" - ) - - if t == "i": - return {"t": t, "i": value} - elif t == "l": - return {"t": t, "v": value} - elif t == "b": - return {"t": t, "d": value} - return {"t": t} - -# ── Security boundary: MCP client → MCP server ── -# The MCP client authenticates to this server via a Bearer token in the -# HTTP Authorization header. The SDK's auth middleware extracts and -# verifies the token before any tool handler runs. -# -# We implement a pass-through TokenVerifier: the gateway is the real -# authority, so we accept any non-empty Bearer token here and forward -# it to the gateway for validation. The gateway's in-band auth -# protocol and IAM regime decide whether the token is valid. -# -# This means an invalid token will connect to the MCP server but will -# fail when the first WebSocket auth frame is sent to the gateway. -# That is intentional — the gateway is the single source of truth. - - -class PassthroughTokenVerifier(TokenVerifier): - """Accept any non-empty Bearer token and forward it downstream. - - The TrustGraph gateway is the authority for token validation, not - this MCP server. We store the raw token in the AccessToken so that - tool handlers can retrieve it via ``get_access_token().token`` and - forward it to the gateway. - """ - - async def verify_token(self, token: str) -> AccessToken | None: - if not token: - return None - return AccessToken( - token=token, - client_id="mcp-caller", - scopes=[], - ) - +from . tg_socket import WebSocketManager @dataclass class AppContext: - sockets: dict[str, WebSocketManager] = field(default_factory=dict) - websocket_url: str = "" - + sockets: dict[str, WebSocketManager] + websocket_url: str + gateway_token: str @asynccontextmanager -async def app_lifespan( - server: FastMCP, - websocket_url: str = "ws://api-gateway:8088/api/v1/socket", -) -> AsyncIterator[AppContext]: - """Manage per-server state: the pool of per-caller WebSocket - connections to the gateway.""" +async def app_lifespan(server: FastMCP, websocket_url: str = "ws://api-gateway:8088/api/v1/socket", gateway_token: str = "") -> AsyncIterator[AppContext]: - sockets: dict[str, WebSocketManager] = {} + """ + Manage application lifecycle with type-safe context + """ + + # Initialize on startup + sockets = {} try: - yield AppContext(sockets=sockets, websocket_url=websocket_url) + yield AppContext(sockets=sockets, websocket_url=websocket_url, gateway_token=gateway_token) finally: - logger.info("Shutting down — closing %d WebSocket(s)", len(sockets)) + # Cleanup on shutdown + logging.info("Shutting down context") - for key, manager in sockets.items(): - try: - await manager.stop() - except Exception as e: - logger.warning("Error closing socket %s: %s", key, e) + for k, manager in sockets.items(): + logging.info(f"Closing socket for {k}") + await manager.stop() - logger.info("Shutdown complete") + logging.info("Shutdown complete") - -def _require_token() -> str: - """Extract the caller's Bearer token from the MCP auth context. - - Raises RuntimeError if no token is present (the caller did not - authenticate). - """ - # ── Security boundary: token extraction ── - # get_access_token() reads the contextvar set by the SDK's - # AuthContextMiddleware. The token was placed there by - # PassthroughTokenVerifier.verify_token() and is the raw Bearer - # value from the MCP client's Authorization header. - access = get_access_token() - if access is None or not access.token: - raise RuntimeError( - "Authentication required — send a Bearer token in the " - "Authorization header" - ) - return access.token - - -async def get_socket_manager(ctx, token): - """Return (or create) an authenticated WebSocket for this token. - - Each unique token gets its own WebSocket connection so that - gateway-side identity, workspace binding, and capability scoping - are preserved per caller. - """ +async def get_socket_manager(ctx): lifespan_context = ctx.request_context.lifespan_context sockets = lifespan_context.sockets websocket_url = lifespan_context.websocket_url + gateway_token = lifespan_context.gateway_token - key = _token_key(token) + if "default" in sockets: + logging.info("Return existing socket manager") + return sockets["default"] - if key in sockets: - manager = sockets[key] - if manager.socket is not None: - return manager - # Socket was closed (e.g. server-side timeout) — reconnect. - del sockets[key] + logging.info(f"Opening socket to {websocket_url}...") - logger.info("Opening authenticated WebSocket to %s …", websocket_url) + # Create manager with empty pending requests + manager = WebSocketManager(websocket_url, token=gateway_token) - manager = WebSocketManager(websocket_url, token=token) + # Start reader task with the proper manager await manager.start() - # Verify the token is valid by calling whoami. This confirms the - # gateway accepted the token and gives us the caller's identity. - try: - identity = await manager.whoami() - logger.info( - "WebSocket ready — caller: %s", - identity.get("handle", "unknown"), - ) - except Exception as e: - await manager.stop() - raise RuntimeError( - f"Token rejected by gateway (whoami failed): {e}" - ) from e + sockets["default"] = manager - sockets[key] = manager + logging.info("Return new socket manager") return manager - @dataclass class EmbeddingsResponse: vectors: List[List[float]] @@ -291,23 +182,10 @@ class PutConfigResponse: class DeleteConfigResponse: pass -@dataclass -class SparqlQueryResponse: - query_type: str - variables: List[str] - bindings: List[Dict[str, Any]] - ask_result: bool - triples: List[Dict[str, Any]] - -@dataclass -class GraphQLQueryResponse: - data: Any - errors: List[Dict[str, Any]] - @dataclass class GetPromptsResponse: prompts: List[str] - + @dataclass class GetPromptResponse: prompt: Dict[str, Any] @@ -316,61 +194,31 @@ class GetPromptResponse: class GetSystemPromptResponse: prompt: str - class McpServer: - def __init__( - self, - host: str = "0.0.0.0", - port: int = 8000, - websocket_url: str = "ws://api-gateway:8088/api/v1/socket", - auth_issuer: str = "", - auth_resource_url: str = "", - ): + def __init__(self, host: str = "0.0.0.0", port: int = 8000, websocket_url: str = "ws://api-gateway:8088/api/v1/socket", gateway_token: str = ""): self.host = host self.port = port self.websocket_url = websocket_url + self.gateway_token = gateway_token - lifespan_with_url = partial( - app_lifespan, websocket_url=websocket_url, - ) - - # ── Security: MCP-level auth configuration ── - # The SDK requires AuthSettings whenever a token_verifier is - # present. The issuer_url tells MCP clients where to obtain - # tokens; resource_server_url identifies this server in OAuth - # protected-resource metadata. - # - # The PassthroughTokenVerifier accepts any non-empty Bearer - # token — real validation happens at the gateway. This is - # intentional: the gateway is the single source of truth for - # identity and capability checks. - from mcp.server.auth.settings import AuthSettings - - auth_settings = AuthSettings( - issuer_url=auth_issuer or f"http://{host}:{port}", - resource_server_url=auth_resource_url or f"http://{host}:{port}", - ) - + # Create a partial function to pass websocket_url to app_lifespan + lifespan_with_url = partial(app_lifespan, websocket_url=websocket_url, gateway_token=gateway_token) + self.mcp = FastMCP( - "TrustGraph", - dependencies=["trustgraph-base"], - host=self.host, - port=self.port, + "TrustGraph", dependencies=["trustgraph-base"], + host=self.host, port=self.port, lifespan=lifespan_with_url, - token_verifier=PassthroughTokenVerifier(), - auth=auth_settings, ) self._register_tools() - + def _register_tools(self): """Register all MCP tools""" + # Register all the tools that were previously registered globally self.mcp.tool()(self.embeddings) self.mcp.tool()(self.text_completion) self.mcp.tool()(self.graph_rag) self.mcp.tool()(self.agent) self.mcp.tool()(self.triples_query) - self.mcp.tool()(self.sparql_query) - self.mcp.tool()(self.graphql_query) self.mcp.tool()(self.graph_embeddings_query) self.mcp.tool()(self.get_config_all) self.mcp.tool()(self.get_config) @@ -395,69 +243,67 @@ class McpServer: self.mcp.tool()(self.load_document) self.mcp.tool()(self.remove_document) self.mcp.tool()(self.add_processing) - + def run(self): """Run the MCP server""" self.mcp.run(transport="streamable-http") - async def _get_manager(self, ctx): - """Get an authenticated WebSocket manager for the current caller. - - Extracts the Bearer token from the MCP auth context and returns - a per-token WebSocket connection to the gateway. - """ - token = _require_token() - return await get_socket_manager(ctx, token) - async def embeddings( self, - texts: List[str], + text: str, flow_id: str | None = None, - workspace: str | None = None, ctx: Context = None, ) -> EmbeddingsResponse: """ - Generate vector embeddings for the given texts using TrustGraph's embedding models. - + Generate vector embeddings for the given text using TrustGraph's embedding models. + This tool converts text into high-dimensional vectors that capture semantic meaning, enabling similarity searches, clustering, and other vector-based operations. - + Args: - texts: List of input texts to convert into embeddings. Each text can be a - sentence, paragraph, or document. + text: The input text to convert into embeddings. Can be a sentence, paragraph, + or document. The text will be processed by the configured embedding model. flow_id: Optional flow identifier to use for processing (default: "default"). Different flows may use different embedding models or configurations. - workspace: Optional workspace to query. If omitted, uses the caller's - default workspace. - + Returns: - EmbeddingsResponse containing a list of vectors, one per input text. + EmbeddingsResponse containing a list of vectors. Each vector is a list of floats + representing the text's semantic embedding in the model's vector space. + + Example usage: + - Convert a query into embeddings for similarity search + - Generate embeddings for documents before storing them + - Create embeddings for comparison with existing knowledge """ - logger.info("Embeddings request") + logging.info("Embeddings request made") if flow_id is None: flow_id = "default" - manager = await self._get_manager(ctx) + manager = await get_socket_manager(ctx, "trustgraph") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Computing embeddings via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + if ctx is None: + raise RuntimeError("No context provided") - request_data = {"texts": texts} - - gen = manager.request( - "embeddings", request_data, flow_id, workspace=workspace, + await ctx.session.send_log_message( + level="info", + data=f"Computing embeddings via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, ) + # Send websocket request + request_data = {"text": text} + logging.info("making request") + + gen = manager.request("embeddings", request_data, flow_id) + async for response in gen: + + # Extract vectors from response vectors = response.get("vectors", [[]]) break - + return EmbeddingsResponse(vectors=vectors) async def text_completion( @@ -465,47 +311,62 @@ class McpServer: prompt: str, system: str | None = None, flow_id: str | None = None, - workspace: str | None = None, ctx: Context = None, ) -> TextCompletionResponse: """ Generate text completions using TrustGraph's language models. - + + This tool sends prompts to configured language models and returns generated text. + It supports both user prompts and system instructions for controlling generation. + Args: prompt: The main prompt or question to send to the language model. + This is the primary input that guides the model's response. system: Optional system prompt that sets the context, role, or behavior - for the AI assistant. - flow_id: Optional flow identifier (default: "default"). - workspace: Optional workspace to query. If omitted, uses the caller's - default workspace. - + for the AI assistant (e.g., "You are a helpful coding assistant"). + System prompts influence how the model interprets and responds. + flow_id: Optional flow identifier (default: "default"). Different flows + may use different models, parameters, or processing pipelines. + Returns: TextCompletionResponse containing the generated text response from the model. + + Example usage: + - Ask questions and get AI-generated answers + - Generate code, documentation, or creative content + - Perform text analysis, summarization, or transformation tasks + - Use system prompts to control tone, style, or domain expertise """ if system is None: system = "" if flow_id is None: flow_id = "default" - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Generating text completion via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + # Use websocket if context is available + logging.info("Text completion request made via websocket") - request_data = {"system": system, "prompt": prompt} + manager = await get_socket_manager(ctx, "trustgraph") - gen = manager.request( - "text-completion", request_data, flow_id, workspace=workspace, + await ctx.session.send_log_message( + level="info", + data=f"Generating text completion via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, ) + # Send websocket request + request_data = {"system": system, "prompt": prompt} + + gen = manager.request("text-completion", request_data, flow_id) + async for response in gen: + + # Extract vectors from response text = response.get("response", "") break - + return TextCompletionResponse(response=text) async def graph_rag( @@ -517,43 +378,58 @@ class McpServer: max_subgraph_size: int | None = None, max_path_length: int | None = None, flow_id: str | None = None, - workspace: str | None = None, ctx: Context = None, ) -> GraphRagResponse: """ Perform Graph-based Retrieval Augmented Generation (GraphRAG) queries. - + GraphRAG combines knowledge graph traversal with language model generation to provide - contextually rich answers. - + contextually rich answers. It explores relationships between entities to build relevant + context before generating responses. + Args: question: The question or query to answer using the knowledge graph. + The system will find relevant entities and relationships to inform the response. collection: Knowledge collection to query (default: "default"). + Different collections may contain domain-specific knowledge. entity_limit: Maximum number of entities to retrieve during graph traversal. + Higher limits provide more context but increase processing time. triple_limit: Maximum number of relationship triples to consider. + Controls the depth of relationship exploration. max_subgraph_size: Maximum size of the subgraph to extract for context. + Larger subgraphs provide richer context but use more resources. max_path_length: Maximum path length to traverse in the knowledge graph. + Longer paths can discover distant but relevant relationships. flow_id: Processing flow to use (default: "default"). - workspace: Optional workspace to query. If omitted, uses the caller's - default workspace. - + Returns: GraphRagResponse containing the generated answer informed by knowledge graph context. + + Example usage: + - Answer complex questions requiring multi-hop reasoning + - Explore relationships between entities in your knowledge base + - Generate responses grounded in structured knowledge + - Perform research queries across connected information """ if collection is None: collection = "default" if flow_id is None: flow_id = "default" - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Processing GraphRAG query via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("GraphRAG request made via websocket") + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Processing GraphRAG query via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) + + # Build request data with all parameters request_data = { "query": question } @@ -564,19 +440,20 @@ class McpServer: if max_subgraph_size: request_data["max_subgraph_size"] = max_subgraph_size if max_path_length: request_data["max_path_length"] = max_path_length - gen = manager.request( - "graph-rag", request_data, flow_id, workspace=workspace, - ) + gen = manager.request("graph-rag", request_data, flow_id) text_chunks = [] async for response in gen: + # Handle new message format with message_type message_type = response.get("message_type", "chunk") + # Only collect text from chunk messages if message_type == "chunk": chunk_text = response.get("response", "") if chunk_text: text_chunks.append(chunk_text) + # Check if session is complete if response.get("end_of_session"): break @@ -587,447 +464,404 @@ class McpServer: question: str, collection: str | None = None, flow_id: str | None = None, - workspace: str | None = None, ctx: Context = None, ) -> AgentResponse: """ Execute intelligent agent queries with reasoning and tool usage capabilities. - + + The agent can perform complex multi-step reasoning, use tools, and provide + detailed thought processes. It's designed for tasks requiring planning, + analysis, and iterative problem-solving. + Args: - question: The question or task for the agent to solve. + question: The question or task for the agent to solve. Can be complex + queries requiring multiple steps, analysis, or tool usage. collection: Knowledge collection the agent can access (default: "default"). - flow_id: Agent workflow to use (default: "default"). - workspace: Optional workspace to query. If omitted, uses the caller's - default workspace. - + Determines what information and tools are available. + flow_id: Agent workflow to use (default: "default"). Different flows + may have different capabilities, tools, or reasoning strategies. + Returns: AgentResponse containing the final answer after the agent's reasoning process. + During execution, you'll see intermediate thoughts and observations. + + Example usage: + - Solve complex analytical problems requiring multiple steps + - Perform research tasks across multiple information sources + - Handle queries that need tool usage and decision-making + - Get detailed explanations of reasoning processes + + Note: This tool provides real-time updates on the agent's thinking process + through log messages, so you can follow its reasoning steps. """ if collection is None: collection = "default" if flow_id is None: flow_id = "default" - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Processing agent query via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Agent request made via websocket") + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Processing agent query via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) + + # Build request data with all parameters request_data = { "question": question } if collection: request_data["collection"] = collection - gen = manager.request( - "agent", request_data, flow_id, workspace=workspace, - ) + gen = manager.request("agent", request_data, flow_id) async for response in gen: - logger.debug("Agent response: %s", response) + logging.debug(f"Agent response: {response}") - if ctx: - if "thought" in response: - await ctx.session.send_log_message( - level="info", - data=f"Thinking: {response['thought']}", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + if "thought" in response: + await ctx.session.send_log_message( + level="info", + data=f"Thinking: {response['thought']}", + logger="notification_stream", + related_request_id=ctx.request_id, + ) - if "observation" in response: - await ctx.session.send_log_message( - level="info", - data=f"Observation: {response['observation']}", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + if "observation" in response: + await ctx.session.send_log_message( + level="info", + data=f"Observation: {response['observation']}", + logger="notification_stream", + related_request_id=ctx.request_id, + ) + # Extract vectors from response if "answer" in response: answer = response.get("answer", "") return AgentResponse(answer=answer) async def triples_query( self, - s: str | None = None, - s_type: str | None = None, - p: str | None = None, - p_type: str | None = None, - o: str | None = None, - o_type: str | None = None, - collection: str | None = None, - graph: str | None = None, + s_v: str | None = None, + s_e: bool | None = None, + p_v: str | None = None, + p_e: bool | None = None, + o_v: str | None = None, + o_e: bool | None = None, limit: int | None = None, flow_id: str | None = None, - workspace: str | None = None, ctx: Context = None, ) -> TriplesQueryResponse: """ Query knowledge graph triples using subject-predicate-object patterns. - - Each of s, p, o is an RDF term value. Use the corresponding _type - parameter to specify the term kind: - - "iri" (default for s and p): an IRI / entity reference - - "literal" (default for o): a plain literal value - - "blank": a blank node identifier - + + Knowledge graphs store information as triples (subject, predicate, object). + This tool allows flexible querying by specifying any combination of these + components, with wildcards for unspecified parts. + Args: - s: Subject value to match. Leave None for wildcard. - s_type: Subject term type: "iri" (default), "literal", or "blank". - p: Predicate value to match. Leave None for wildcard. - p_type: Predicate term type: "iri" (default), "literal", or "blank". - o: Object value to match. Leave None for wildcard. - o_type: Object term type: "iri", "literal" (default), or "blank". - collection: Knowledge collection to query (default: "default"). - graph: Named graph IRI to restrict the query. None = default graph, - "*" = all graphs. + s_v: Subject value to match (e.g., "John", "Apple Inc."). Leave None for wildcard. + s_e: Whether subject should be treated as an entity (True) or literal (False). + p_v: Predicate/relationship value (e.g., "works_for", "type_of"). Leave None for wildcard. + p_e: Whether predicate should be treated as an entity (True) or literal (False). + o_v: Object value to match (e.g., "Engineer", "Company"). Leave None for wildcard. + o_e: Whether object should be treated as an entity (True) or literal (False). limit: Maximum number of triples to return (default: 20). flow_id: Processing flow identifier (default: "default"). - workspace: Optional workspace to query. If omitted, uses the caller's - default workspace. - + Returns: TriplesQueryResponse containing matching triples from the knowledge graph. + + Example queries: + - Find all relationships for an entity: s_v="John", others None + - Find all instances of a relationship: p_v="works_for", others None + - Find specific facts: s_v="John", p_v="works_for", o_v=None + - Explore entity types: p_v="type_of", others None + + Use this for: + - Exploring knowledge graph structure + - Finding specific facts or relationships + - Discovering connections between entities + - Validating or debugging knowledge content """ if flow_id is None: flow_id = "default" if limit is None: limit = 20 - if collection is None: collection = "default" - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Processing triples query via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Triples query request made via websocket") - request_data = { - "limit": limit, - "collection": collection, - } + manager = await get_socket_manager(ctx, "trustgraph") - if s is not None: - request_data["s"] = _make_term(s, s_type or "iri") - - if p is not None: - request_data["p"] = _make_term(p, p_type or "iri") - - if o is not None: - request_data["o"] = _make_term(o, o_type or "literal") - - if graph is not None: - request_data["g"] = graph - - gen = manager.request( - "triples", request_data, flow_id, workspace=workspace, + await ctx.session.send_log_message( + level="info", + data=f"Processing triples query via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, ) + # Build request data with Value objects + request_data = { + "limit": limit + } + + # Add subject if provided + if s_v is not None: + request_data["s"] = {"v": s_v, "e": s_e } + + # Add predicate if provided + if p_v is not None: + request_data["p"] = {"v": p_v, "e": p_e } + + # Add object if provided + if o_v is not None: + request_data["o"] = {"v": o_v, "e": o_e } + + gen = manager.request("triples", request_data, flow_id) + async for response in gen: + # Extract response data triples = response.get("response", []) break - + return TriplesQueryResponse(triples=triples) - async def sparql_query( - self, - query: str, - collection: str | None = None, - limit: int | None = None, - flow_id: str | None = None, - workspace: str | None = None, - ctx: Context = None, - ) -> SparqlQueryResponse: - """ - Execute a SPARQL query against the knowledge graph. - - Supports SELECT, ASK, CONSTRUCT, and DESCRIBE query forms. - - Args: - query: SPARQL query string (e.g. "SELECT ?s ?p ?o WHERE { ?s ?p ?o } LIMIT 10"). - collection: Knowledge collection to query (default: "default"). - limit: Safety limit on number of results (default: 10000). - flow_id: Processing flow identifier (default: "default"). - workspace: Optional workspace to query. If omitted, uses the caller's - default workspace. - - Returns: - SparqlQueryResponse containing the query results. The structure depends - on query type: - - SELECT: variables (column names) and bindings (rows of Term values) - - ASK: ask_result (boolean) - - CONSTRUCT/DESCRIBE: triples - """ - - if collection is None: collection = "default" - if flow_id is None: flow_id = "default" - if limit is None: limit = 10000 - - manager = await self._get_manager(ctx) - - if ctx: - await ctx.session.send_log_message( - level="info", - data="Processing SPARQL query via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) - - request_data = { - "query": query, - "collection": collection, - "limit": limit, - } - - gen = manager.request( - "sparql", request_data, flow_id, workspace=workspace, - ) - - async for response in gen: - query_type = response.get("query-type", "") - return SparqlQueryResponse( - query_type=query_type, - variables=response.get("variables", []), - bindings=response.get("bindings", []), - ask_result=response.get("ask-result", False), - triples=response.get("triples", []), - ) - - async def graphql_query( - self, - query: str, - collection: str | None = None, - variables: Dict[str, Any] | None = None, - operation_name: str | None = None, - flow_id: str | None = None, - workspace: str | None = None, - ctx: Context = None, - ) -> GraphQLQueryResponse: - """ - Execute a GraphQL query against structured data (rows). - - Queries structured data schemas that have been loaded into TrustGraph. - The available types and fields depend on the schemas configured in the - target workspace. - - Args: - query: GraphQL query string (e.g. '{ customers(where: {status: {eq: "active"}}) { id name } }'). - collection: Data collection to query (default: "default"). - variables: Optional GraphQL variables as a dict. - operation_name: Optional operation name for multi-operation documents. - flow_id: Processing flow identifier (default: "default"). - workspace: Optional workspace to query. If omitted, uses the caller's - default workspace. - - Returns: - GraphQLQueryResponse containing data (the query result) and errors - (any GraphQL field-level errors). - """ - - if collection is None: collection = "default" - if flow_id is None: flow_id = "default" - - manager = await self._get_manager(ctx) - - if ctx: - await ctx.session.send_log_message( - level="info", - data="Processing GraphQL query via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) - - request_data = { - "query": query, - "collection": collection, - "variables": variables or {}, - } - - if operation_name is not None: - request_data["operation_name"] = operation_name - - gen = manager.request( - "rows", request_data, flow_id, workspace=workspace, - ) - - async for response in gen: - return GraphQLQueryResponse( - data=response.get("data"), - errors=response.get("errors", []), - ) - async def graph_embeddings_query( self, vectors: List[List[float]], limit: int | None = None, flow_id: str | None = None, - workspace: str | None = None, ctx: Context = None, ) -> GraphEmbeddingsQueryResponse: """ Find entities in the knowledge graph using vector similarity search. - + + This tool performs semantic search by comparing embedding vectors to find + the most similar entities in the knowledge graph. It's useful for finding + conceptually related information even when exact text matches don't exist. + Args: - vectors: List of embedding vectors to search with. + vectors: List of embedding vectors to search with. Each vector should be + a list of floats representing semantic embeddings (typically from + the embeddings tool). Multiple vectors can be provided for batch queries. limit: Maximum number of similar entities to return (default: 20). + Higher limits provide more results but may include less relevant matches. flow_id: Processing flow identifier (default: "default"). - workspace: Optional workspace to query. If omitted, uses the caller's - default workspace. - + Returns: - GraphEmbeddingsQueryResponse containing entities ranked by similarity. + GraphEmbeddingsQueryResponse containing entities ranked by similarity to the + input vectors, along with similarity scores and entity metadata. + + Example workflow: + 1. Use the 'embeddings' tool to convert text to vectors + 2. Use this tool to find similar entities in the knowledge graph + 3. Explore the returned entities for relevant information + + Use this for: + - Semantic search across knowledge entities + - Finding conceptually similar content + - Discovering related entities without exact keyword matches + - Building recommendation systems based on entity similarity """ if flow_id is None: flow_id = "default" if limit is None: limit = 20 - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Processing graph embeddings query via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Graph embeddings query request made via websocket") + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Processing graph embeddings query via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) + + # Build request data request_data = { "vectors": vectors, "limit": limit } - gen = manager.request( - "graph-embeddings", request_data, flow_id, workspace=workspace, - ) + gen = manager.request("graph-embeddings", request_data, flow_id) async for response in gen: + # Extract entities from response entities = response.get("entities", []) break - + return GraphEmbeddingsQueryResponse(entities=entities) async def get_config_all( self, - workspace: str | None = None, ctx: Context = None, ) -> ConfigResponse: """ Retrieve the complete TrustGraph system configuration. - - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + + This tool returns all configuration settings for the TrustGraph system, + including model configurations, API keys, flow definitions, and system parameters. + Returns: - ConfigResponse containing the full configuration as a nested dictionary. + ConfigResponse containing the full configuration as a nested dictionary + with all system settings, organized by category (e.g., models, flows, storage). + + Use this for: + - Inspecting current system configuration + - Debugging configuration issues + - Understanding available models and settings + - Auditing system setup and parameters """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving all configuration via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get config all request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving all configuration via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "config" } - gen = manager.request("config", request_data, None, workspace=workspace) + gen = manager.request("config", request_data, None) async for response in gen: config = response.get("config", {}) break - + return ConfigResponse(config=config) async def get_config( self, keys: List[Dict[str, str]], - workspace: str | None = None, ctx: Context = None, ) -> ConfigGetResponse: """ Retrieve specific configuration values by key. - + + This tool allows you to fetch specific configuration settings without + retrieving the entire configuration. Useful for checking particular + settings or API keys. + Args: - keys: List of configuration keys to retrieve. Each key should be a dict with - 'type' and 'key' fields. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + keys: List of configuration keys to retrieve. Each key should be a dict with: + - 'type': Configuration category (e.g., 'llm', 'embeddings', 'storage') + - 'key': Specific setting name within that category + Returns: ConfigGetResponse containing the requested configuration values. + + Example keys: + - {'type': 'llm', 'key': 'openai.model'} + - {'type': 'embeddings', 'key': 'default.model'} + - {'type': 'storage', 'key': 'database.url'} + + Use this for: + - Checking specific model configurations + - Validating API key settings + - Inspecting individual system parameters """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving specific configuration via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get config request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving specific configuration via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "get", "keys": keys } - gen = manager.request("config", request_data, None, workspace=workspace) + gen = manager.request("config", request_data, None) async for response in gen: values = response.get("values", []) break - + return ConfigGetResponse(values=values) async def put_config( self, values: List[Dict[str, str]], - workspace: str | None = None, ctx: Context = None, ) -> PutConfigResponse: """ Update system configuration values. - + + This tool allows you to modify TrustGraph system settings, such as + model parameters, API keys, and system behavior configurations. + Args: - values: List of configuration updates. Each should be a dict with - 'type', 'key', and 'value' fields. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + values: List of configuration updates. Each update should be a dict with: + - 'type': Configuration category (e.g., 'llm', 'embeddings') + - 'key': Specific setting name to update + - 'value': New value for the setting + Returns: PutConfigResponse confirming the configuration update. + + Example updates: + - {'type': 'llm', 'key': 'openai.model', 'value': 'gpt-4'} + - {'type': 'embeddings', 'key': 'batch_size', 'value': '100'} + + Use this for: + - Switching between different models + - Updating API credentials + - Modifying system behavior parameters + - Configuring processing settings + + Note: Configuration changes may require system restart to take effect. """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Updating configuration via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Put config request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Updating configuration via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "put", "values": values } - gen = manager.request("config", request_data, None, workspace=workspace) + gen = manager.request("config", request_data, None) async for response in gen: return PutConfigResponse() @@ -1035,73 +869,97 @@ class McpServer: async def delete_config( self, keys: List[Dict[str, str]], - workspace: str | None = None, ctx: Context = None, ) -> DeleteConfigResponse: """ Delete specific configuration entries from the system. - + + This tool removes configuration settings, reverting them to system defaults + or disabling specific features. + Args: - keys: List of configuration keys to delete. Each should be a dict with - 'type' and 'key' fields. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + keys: List of configuration keys to delete. Each key should be a dict with: + - 'type': Configuration category (e.g., 'llm', 'embeddings') + - 'key': Specific setting name to remove + Returns: DeleteConfigResponse confirming the deletion. + + Use this for: + - Removing custom model configurations + - Clearing API credentials + - Resetting settings to defaults + - Cleaning up obsolete configurations + + Warning: Deleting essential configuration may cause system functionality + to be disabled until properly reconfigured. """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Deleting configuration via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Delete config request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Deleting configuration via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "delete", "keys": keys } - gen = manager.request("config", request_data, None, workspace=workspace) + gen = manager.request("config", request_data, None) async for response in gen: return DeleteConfigResponse() async def get_prompts( self, - workspace: str | None = None, ctx: Context = None, ) -> GetPromptsResponse: """ List all available prompt templates in the system. - - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + + Prompt templates are reusable prompts that can be used with language models + for consistent behavior across different queries and use cases. + Returns: GetPromptsResponse containing a list of available prompt template IDs. + Each ID can be used with get_prompt to retrieve the full template. + + Use this for: + - Discovering available prompt templates + - Exploring pre-configured prompts for different tasks + - Finding templates for specific use cases + - Understanding what prompt options are available """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving prompt templates via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get prompts request made via websocket") + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving prompt templates via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) + + # First get all config request_data = { "operation": "config" } - gen = manager.request("config", request_data, None, workspace=workspace) + gen = manager.request("config", request_data, None) async for response in gen: config = response.get("config", {}) @@ -1113,36 +971,49 @@ class McpServer: async def get_prompt( self, prompt_id: str, - workspace: str | None = None, ctx: Context = None, ) -> GetPromptResponse: """ Retrieve a specific prompt template by ID. - + + Prompt templates contain structured prompts with placeholders, instructions, + and metadata for specific tasks or domains. + Args: prompt_id: The unique identifier of the prompt template to retrieve. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + Use get_prompts to see available template IDs. + Returns: - GetPromptResponse containing the complete prompt template. + GetPromptResponse containing the complete prompt template with its + structure, placeholders, and usage instructions. + + Use this for: + - Examining prompt template structure + - Understanding how to use specific templates + - Copying or modifying existing prompts + - Learning prompt engineering patterns """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Retrieving prompt template '{prompt_id}' via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get prompt request made via websocket") + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving prompt template '{prompt_id}' via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) + + # First get all config request_data = { "operation": "config" } - gen = manager.request("config", request_data, None, workspace=workspace) + gen = manager.request("config", request_data, None) async for response in gen: config = response.get("config", {}) @@ -1154,35 +1025,44 @@ class McpServer: async def get_system_prompt( self, - workspace: str | None = None, ctx: Context = None, ) -> GetSystemPromptResponse: """ Retrieve the current system prompt configuration. - - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + + The system prompt defines the default behavior, personality, and instructions + for language models across the TrustGraph system. + Returns: - GetSystemPromptResponse containing the system prompt text. + GetSystemPromptResponse containing the system prompt text and configuration. + + Use this for: + - Understanding default AI behavior settings + - Checking current system-wide prompt configuration + - Auditing AI personality and instruction settings + - Debugging unexpected AI responses """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving system prompt via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get system prompt request made via websocket") + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving system prompt via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) + + # First get all config request_data = { "operation": "config" } - gen = manager.request("config", request_data, None, workspace=workspace) + gen = manager.request("config", request_data, None) async for response in gen: config = response.get("config", {}) @@ -1193,39 +1073,51 @@ class McpServer: async def get_token_costs( self, - workspace: str | None = None, ctx: Context = None, ) -> ConfigTokenCostsResponse: """ Retrieve token pricing information for all configured AI models. - - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + + This tool provides cost information for input and output tokens across + different language models, helping with budget planning and cost optimization. + Returns: - ConfigTokenCostsResponse containing pricing data for each model. + ConfigTokenCostsResponse containing pricing data for each model including: + - Model name/identifier + - Input token cost (per token) + - Output token cost (per token) + + Use this for: + - Estimating costs for different models + - Choosing cost-effective models for tasks + - Budget planning and cost analysis + - Monitoring and optimizing AI spending """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving token costs via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get token costs request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving token costs via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "getvalues", "type": "token-costs" } - gen = manager.request("config", request_data, None, workspace=workspace) + gen = manager.request("config", request_data, None) async for response in gen: values = response.get("values", []) + # Transform to match TypeScript API format costs = [] for item in values: try: @@ -1238,89 +1130,106 @@ class McpServer: except (json.JSONDecodeError, AttributeError): continue break - + return ConfigTokenCostsResponse(costs=costs) async def get_knowledge_cores( self, - workspace: str | None = None, ctx: Context = None, ) -> KnowledgeCoresResponse: """ List all available knowledge graph cores in the current workspace. - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. + Knowledge cores are packaged collections of structured knowledge that can + be loaded into the system for querying and reasoning. They contain entities, + relationships, and facts organized as knowledge graphs. Returns: KnowledgeCoresResponse containing a list of available knowledge core IDs. + + Use this for: + - Discovering available knowledge collections + - Understanding what knowledge domains are accessible + - Planning which cores to load for specific tasks + - Managing knowledge resources """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving knowledge graph cores via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get knowledge cores request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving knowledge graph cores via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "list-kg-cores", } - gen = manager.request( - "knowledge", request_data, None, workspace=workspace, - ) + gen = manager.request("knowledge", request_data, None) async for response in gen: ids = response.get("ids", []) break - + return KnowledgeCoresResponse(ids=ids) async def delete_kg_core( self, core_id: str, - workspace: str | None = None, ctx: Context = None, ) -> DeleteKgCoreResponse: """ Permanently delete a knowledge graph core. + This operation removes a knowledge core from storage. Use with caution + as this action cannot be undone. + Args: core_id: Unique identifier of the knowledge core to delete. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. Returns: DeleteKgCoreResponse confirming the deletion. + + Use this for: + - Cleaning up obsolete knowledge cores + - Removing test or experimental data + - Managing storage space + - Maintaining organized knowledge collections + + Warning: This permanently deletes the knowledge core and all its data. """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Deleting knowledge graph core '{core_id}' via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Delete KG core request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Deleting knowledge graph core '{core_id}' via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "delete-kg-core", "id": core_id, } - gen = manager.request( - "knowledge", request_data, None, workspace=workspace, - ) + gen = manager.request("knowledge", request_data, None) async for response in gen: break - + return DeleteKgCoreResponse() async def load_kg_core( @@ -1328,34 +1237,46 @@ class McpServer: core_id: str, flow: str, collection: str | None = None, - workspace: str | None = None, ctx: Context = None, ) -> LoadKgCoreResponse: """ Load a knowledge graph core into the active system for querying. + This operation makes a knowledge core available for GraphRAG queries, + triple searches, and other knowledge-based operations. + Args: core_id: Unique identifier of the knowledge core to load. - flow: Processing flow to use for loading the core. - collection: Target collection name (default: "default"). - workspace: Optional workspace. If omitted, uses the caller's - default workspace. + flow: Processing flow to use for loading the core. Different flows + may apply different processing, indexing, or optimization steps. + collection: Target collection name (default: "default"). The loaded + knowledge will be available under this collection name. Returns: LoadKgCoreResponse confirming the core has been loaded. + + Use this for: + - Making knowledge cores available for queries + - Switching between different knowledge domains + - Loading domain-specific knowledge for tasks + - Preparing knowledge for GraphRAG operations """ if collection is None: collection = "default" - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Loading knowledge graph core '{core_id}' via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Load KG core request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Loading knowledge graph core '{core_id}' via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "load-kg-core", @@ -1364,241 +1285,292 @@ class McpServer: "collection": collection } - gen = manager.request( - "knowledge", request_data, None, workspace=workspace, - ) + gen = manager.request("knowledge", request_data, None) async for response in gen: break - + return LoadKgCoreResponse() async def get_kg_core( self, core_id: str, - workspace: str | None = None, ctx: Context = None, ) -> GetKgCoreResponse: """ Download and retrieve the complete content of a knowledge graph core. + This tool streams the entire content of a knowledge core, returning all + entities, relationships, and metadata. Due to potentially large data sizes, + the content is streamed in chunks. + Args: core_id: Unique identifier of the knowledge core to retrieve. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. Returns: GetKgCoreResponse containing all chunks of the knowledge core data. + Each chunk contains part of the knowledge graph structure. + + Use this for: + - Examining knowledge core content and structure + - Debugging knowledge graph data + - Exporting knowledge for backup or analysis + - Understanding the scope and quality of knowledge + + Note: Large knowledge cores may take significant time to download. + Progress updates are provided through log messages during streaming. """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Retrieving knowledge graph core '{core_id}' via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get KG core request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving knowledge graph core '{core_id}' via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "get-kg-core", "id": core_id, } + # Collect all streaming responses chunks = [] - gen = manager.request( - "knowledge", request_data, None, workspace=workspace, - ) + gen = manager.request("knowledge", request_data, None) async for response in gen: + # Check for end of stream if response.get("eos", False): - if ctx: - await ctx.session.send_log_message( - level="info", - data="Completed streaming KG core data", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + await ctx.session.send_log_message( + level="info", + data=f"Completed streaming KG core data", + logger="notification_stream", + related_request_id=ctx.request_id, + ) break else: chunks.append(response) - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Received KG core chunk ({len(chunks)} chunks so far)", - logger="notification_stream", - related_request_id=ctx.request_id, - ) - + await ctx.session.send_log_message( + level="info", + data=f"Received KG core chunk ({len(chunks)} chunks so far)", + logger="notification_stream", + related_request_id=ctx.request_id, + ) + return GetKgCoreResponse(chunks=chunks) async def get_flows( self, - workspace: str | None = None, ctx: Context = None, ) -> FlowsResponse: """ List all available processing flows in the system. - - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + + Flows define processing pipelines for different types of operations + (e.g., document processing, knowledge extraction, query handling). + Each flow encapsulates a specific workflow with configured steps. + Returns: FlowsResponse containing a list of available flow identifiers. + + Use this for: + - Discovering available processing workflows + - Understanding what processing options are available + - Choosing appropriate flows for specific tasks + - Planning workflow-based operations """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving available flows via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get flows request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving available flows via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "list-flows" } - gen = manager.request( - "flow", request_data, None, workspace=workspace, - ) + gen = manager.request("flow", request_data, None) async for response in gen: flow_ids = response.get("flow-ids", []) break - + return FlowsResponse(flow_ids=flow_ids) async def get_flow( self, flow_id: str, - workspace: str | None = None, ctx: Context = None, ) -> FlowResponse: """ Retrieve the complete definition of a specific processing flow. - + + This tool returns the detailed configuration, steps, and parameters + of a processing flow, showing how it processes data and what operations it performs. + Args: flow_id: Unique identifier of the flow to retrieve. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + Returns: - FlowResponse containing the complete flow definition. + FlowResponse containing the complete flow definition including: + - Flow configuration and parameters + - Processing steps and their order + - Input/output specifications + - Dependencies and requirements + + Use this for: + - Understanding how specific flows work + - Debugging flow processing issues + - Learning flow configuration patterns + - Customizing or duplicating flows """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Retrieving flow definition for '{flow_id}' via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get flow request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving flow definition for '{flow_id}' via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "get-flow", "flow-id": flow_id, } - gen = manager.request( - "flow", request_data, None, workspace=workspace, - ) + gen = manager.request("flow", request_data, None) async for response in gen: flow_data = response.get("flow", "{}") + # Parse JSON flow definition as done in TypeScript flow = json.loads(flow_data) if isinstance(flow_data, str) else flow_data break - + return FlowResponse(flow=flow) async def get_flow_classes( self, - workspace: str | None = None, ctx: Context = None, ) -> FlowClassesResponse: """ List all available flow class templates. - - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + + Flow classes are templates that define types of processing workflows. + They serve as blueprints for creating specific flow instances with + customized parameters. + Returns: FlowClassesResponse containing a list of available flow class names. + + Use this for: + - Discovering available flow templates + - Understanding what types of processing are supported + - Planning new flow creation + - Exploring system capabilities """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving flow classes via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get flow classes request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving flow classes via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "list-classes" } - gen = manager.request( - "flow", request_data, None, workspace=workspace, - ) + gen = manager.request("flow", request_data, None) async for response in gen: class_names = response.get("class-names", []) break - + return FlowClassesResponse(class_names=class_names) async def get_flow_class( self, class_name: str, - workspace: str | None = None, ctx: Context = None, ) -> FlowClassResponse: """ Retrieve the definition of a specific flow class template. - + + Flow classes define the structure, parameters, and capabilities of + flow types. This tool returns the class specification including + configurable parameters and processing logic. + Args: class_name: Name of the flow class to retrieve. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + Returns: - FlowClassResponse containing the flow class definition. + FlowClassResponse containing the flow class definition with: + - Class parameters and configuration options + - Processing capabilities and requirements + - Usage instructions and examples + + Use this for: + - Understanding flow class capabilities + - Learning how to configure new flows + - Troubleshooting flow creation issues + - Exploring advanced flow features """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Retrieving flow class definition for '{class_name}' via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get flow class request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving flow class definition for '{class_name}' via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "get-class", "class-name": class_name } - gen = manager.request( - "flow", request_data, None, workspace=workspace, - ) + gen = manager.request("flow", request_data, None) async for response in gen: class_def_data = response.get("class-definition", "{}") + # Parse JSON class definition as done in TypeScript class_definition = json.loads(class_def_data) if isinstance(class_def_data, str) else class_def_data break - + return FlowClassResponse(class_definition=class_definition) async def start_flow( @@ -1606,32 +1578,43 @@ class McpServer: flow_id: str, class_name: str, description: str, - workspace: str | None = None, ctx: Context = None, ) -> StartFlowResponse: """ Create and start a new processing flow instance. - + + This tool creates a new flow based on a flow class template and starts + it running. The flow will begin processing according to its configuration. + Args: flow_id: Unique identifier for the new flow instance. class_name: Flow class template to use for creating the flow. + Use get_flow_classes to see available classes. description: Human-readable description of the flow's purpose. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + Returns: StartFlowResponse confirming the flow has been started. + + Use this for: + - Creating new processing workflows + - Starting automated processing tasks + - Launching background operations + - Initiating data processing pipelines """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Starting flow '{flow_id}' with class '{class_name}' via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Start flow request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Starting flow '{flow_id}' with class '{class_name}' via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "start-flow", @@ -1640,135 +1623,162 @@ class McpServer: "description": description } - gen = manager.request( - "flow", request_data, None, workspace=workspace, - ) + gen = manager.request("flow", request_data, None) async for response in gen: break - + return StartFlowResponse() async def stop_flow( self, flow_id: str, - workspace: str | None = None, ctx: Context = None, ) -> StopFlowResponse: """ Stop a running flow instance. - + + This tool gracefully stops a running flow, allowing it to complete + current operations before shutting down. + Args: flow_id: Unique identifier of the flow instance to stop. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. - + Returns: StopFlowResponse confirming the flow has been stopped. + + Use this for: + - Stopping unwanted or completed flows + - Managing system resources + - Interrupting long-running processes + - Maintaining flow lifecycle """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Stopping flow '{flow_id}' via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Stop flow request made via websocket") + + manager = await get_socket_manager(ctx, "trustgraph") + + await ctx.session.send_log_message( + level="info", + data=f"Stopping flow '{flow_id}' via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "stop-flow", "flow-id": flow_id } - gen = manager.request( - "flow", request_data, None, workspace=workspace, - ) + gen = manager.request("flow", request_data, None) async for response in gen: break - + return StopFlowResponse() async def get_documents( self, - workspace: str | None = None, ctx: Context = None, ) -> DocumentsResponse: """ List all documents stored in the TrustGraph document library. - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. + This tool returns metadata for all documents that have been uploaded + to the system, including their processing status and properties. Returns: - DocumentsResponse containing metadata for each document. + DocumentsResponse containing metadata for each document including: + - Document ID and title + - Upload timestamp + - MIME type and size information + - Tags and custom metadata + - Processing status + + Use this for: + - Browsing available documents + - Managing document collections + - Finding documents by metadata + - Auditing document storage """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving documents list via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get documents request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving documents list via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "list-documents", } - gen = manager.request( - "librarian", request_data, None, workspace=workspace, - ) + gen = manager.request("librarian", request_data, None) async for response in gen: document_metadatas = response.get("document-metadatas", []) break - + return DocumentsResponse(document_metadatas=document_metadatas) async def get_processing( self, - workspace: str | None = None, ctx: Context = None, ) -> ProcessingResponse: """ List all documents currently in the processing queue. - Args: - workspace: Optional workspace. If omitted, uses the caller's - default workspace. + This tool shows documents that are being processed or waiting to be + processed, along with their processing status and configuration. Returns: - ProcessingResponse containing processing metadata. + ProcessingResponse containing processing metadata including: + - Processing job ID and document ID + - Processing flow and status + - Target collection + - Timestamp and progress information + + Use this for: + - Monitoring document processing progress + - Debugging processing issues + - Managing processing queues + - Understanding system workload """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Retrieving processing list via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Get processing request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Retrieving processing list via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "list-processing", } - gen = manager.request( - "librarian", request_data, None, workspace=workspace, - ) + gen = manager.request("librarian", request_data, None) async for response in gen: processing_metadatas = response.get("processing-metadatas", []) break - + return ProcessingResponse(processing_metadatas=processing_metadatas) async def load_document( @@ -1780,39 +1790,50 @@ class McpServer: title: str = "", comments: str = "", tags: List[str] | None = None, - workspace: str | None = None, ctx: Context = None, ) -> LoadDocumentResponse: """ Upload a document to the TrustGraph document library. + This tool stores documents with rich metadata for later processing, + search, and knowledge extraction. Documents can be text files, PDFs, + or other supported formats. + Args: document: The document content as a string. For binary files, this should be base64-encoded content. document_id: Optional unique identifier. If not provided, one will be generated. metadata: Optional list of custom metadata key-value pairs. - mime_type: MIME type of the document. + mime_type: MIME type of the document (e.g., 'text/plain', 'application/pdf'). title: Human-readable title for the document. comments: Optional description or notes about the document. - tags: List of tags for categorizing the document. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. + tags: List of tags for categorizing and finding the document. Returns: LoadDocumentResponse confirming the document has been stored. + + Use this for: + - Adding new documents to the knowledge base + - Storing reference materials and data sources + - Building document collections for processing + - Importing external content for analysis """ if tags is None: tags = [] - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data="Loading document to library via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Load document request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Loading document to library via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) import time timestamp = int(time.time()) @@ -1831,55 +1852,63 @@ class McpServer: "content": document } - gen = manager.request( - "librarian", request_data, None, workspace=workspace, - ) + gen = manager.request("librarian", request_data, None) async for response in gen: break - + return LoadDocumentResponse() async def remove_document( self, document_id: str, - workspace: str | None = None, ctx: Context = None, ) -> RemoveDocumentResponse: """ Permanently remove a document from the library. + This operation deletes a document and all its associated metadata. + Use with caution as this action cannot be undone. + Args: document_id: Unique identifier of the document to remove. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. Returns: RemoveDocumentResponse confirming the document has been deleted. + + Use this for: + - Cleaning up obsolete or incorrect documents + - Managing storage space + - Removing sensitive or inappropriate content + - Maintaining organized document collections + + Warning: This permanently deletes the document and all its metadata. """ - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Removing document '{document_id}' from library via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Remove document request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Removing document '{document_id}' from library via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) request_data = { "operation": "remove-document", "document-id": document_id, } - gen = manager.request( - "librarian", request_data, None, workspace=workspace, - ) + gen = manager.request("librarian", request_data, None) async for response in gen: break - + return RemoveDocumentResponse() async def add_processing( @@ -1889,37 +1918,53 @@ class McpServer: flow: str, collection: str | None = None, tags: List[str] | None = None, - workspace: str | None = None, ctx: Context = None, ) -> AddProcessingResponse: """ Queue a document for processing through a specific workflow. + This tool adds a document to the processing queue where it will be + processed by the specified flow to extract knowledge, create embeddings, + or perform other analysis operations. + Args: processing_id: Unique identifier for this processing job. document_id: ID of the document to process (must exist in library). - flow: Processing flow to use. + flow: Processing flow to use. Different flows perform different + types of analysis (e.g., knowledge extraction, summarization). collection: Target collection for processed knowledge (default: "default"). + Results will be stored under this collection name. tags: Optional tags for categorizing this processing job. - workspace: Optional workspace. If omitted, uses the caller's - default workspace. Returns: AddProcessingResponse confirming the document has been queued. + + Use this for: + - Processing uploaded documents into knowledge + - Extracting entities and relationships from text + - Creating searchable embeddings + - Converting documents into structured knowledge + + Note: Processing may take time depending on document size and flow complexity. + Use get_processing to monitor progress. """ if collection is None: collection = "default" if tags is None: tags = [] - manager = await self._get_manager(ctx) + if ctx is None: + raise RuntimeError("No context provided") - if ctx: - await ctx.session.send_log_message( - level="info", - data=f"Adding document '{document_id}' to processing queue via websocket...", - logger="notification_stream", - related_request_id=ctx.request_id, - ) + logging.info("Add processing request made via websocket") + + manager = await get_socket_manager(ctx) + + await ctx.session.send_log_message( + level="info", + data=f"Adding document '{document_id}' to processing queue via websocket...", + logger="notification_stream", + related_request_id=ctx.request_id, + ) import time timestamp = int(time.time()) @@ -1936,61 +1981,38 @@ class McpServer: } } - gen = manager.request( - "librarian", request_data, None, workspace=workspace, - ) + gen = manager.request("librarian", request_data, None) async for response in gen: break - + return AddProcessingResponse() - def main(): parser = argparse.ArgumentParser(description='TrustGraph MCP Server') - parser.add_argument( - '--host', default='0.0.0.0', - help='Host to bind to (default: 0.0.0.0)', - ) - parser.add_argument( - '--port', type=int, default=8000, - help='Port to bind to (default: 8000)', - ) - parser.add_argument( - '--websocket-url', - default='ws://api-gateway:8088/api/v1/socket', - help='WebSocket URL for the TrustGraph gateway', - ) - parser.add_argument( - '--auth-issuer', - default=os.environ.get("AUTH_ISSUER", ""), - help='OAuth issuer URL for MCP auth metadata discovery', - ) - parser.add_argument( - '--auth-resource-url', - default=os.environ.get("AUTH_RESOURCE_URL", ""), - help='Resource server URL for OAuth protected resource metadata', - ) + parser.add_argument('--host', default='0.0.0.0', help='Host to bind to (default: 0.0.0.0)') + parser.add_argument('--port', type=int, default=8000, help='Port to bind to (default: 8000)') + parser.add_argument('--websocket-url', default='ws://api-gateway:8088/api/v1/socket', help='WebSocket URL to connect to (default: ws://api-gateway:8088/api/v1/socket)') + # Add logging arguments add_logging_args(parser) args = parser.parse_args() + # Setup logging before creating server setup_logging(vars(args)) - server = McpServer( - host=args.host, - port=args.port, - websocket_url=args.websocket_url, - auth_issuer=args.auth_issuer, - auth_resource_url=args.auth_resource_url, - ) + # Read gateway auth token from environment + gateway_token = os.environ.get("GATEWAY_SECRET", "") + + # Create and run the MCP server + server = McpServer(host=args.host, port=args.port, websocket_url=args.websocket_url, gateway_token=gateway_token) server.run() - def run(): + """Legacy function for backward compatibility""" main() - if __name__ == "__main__": main() + diff --git a/trustgraph-mcp/trustgraph/mcp_server/tg_socket.py b/trustgraph-mcp/trustgraph/mcp_server/tg_socket.py index 9fbf7459..bff8ae75 100644 --- a/trustgraph-mcp/trustgraph/mcp_server/tg_socket.py +++ b/trustgraph-mcp/trustgraph/mcp_server/tg_socket.py @@ -1,110 +1,49 @@ +from dataclasses import dataclass from websockets.asyncio.client import connect +from urllib.parse import urlencode, urlparse, urlunparse, parse_qs import asyncio import logging import json import uuid -import hashlib - -logger = logging.getLogger(__name__) - - -def _token_key(token): - """Derive a dict key from a token without storing the raw secret.""" - return hashlib.sha256(token.encode()).hexdigest()[:16] - +import time class WebSocketManager: - """Manages an authenticated WebSocket connection to the TrustGraph - gateway on behalf of a single caller. - Each caller token gets its own WebSocketManager so that gateway-side - identity, workspace, and capability scoping are preserved end-to-end. - """ - - def __init__(self, url, token): + def __init__(self, url, token=None): self.url = url - # ── Security boundary: token storage ── - # This is the MCP caller's Bearer token, forwarded verbatim to - # the gateway. It MUST NOT be logged, persisted, or shared - # across callers. It is held only for the lifetime of this - # connection so that re-auth (e.g. after a reconnect) is - # possible. self.token = token self.socket = None - self.identity = None - self.last_used = None + + # FIXME: authentication is broken. The /api/v1/socket endpoint uses + # in-band auth (first-frame protocol via the Mux dispatcher), not + # query-parameter tokens. This query-string token is silently ignored. + # Fix: after connect(), send an auth frame with the bearer token as + # the first message, matching the gateway's in-band auth protocol. + def _build_url(self): + if not self.token: + return self.url + parsed = urlparse(self.url) + params = parse_qs(parsed.query) + params["token"] = [self.token] + new_query = urlencode(params, doseq=True) + return urlunparse(parsed._replace(query=new_query)) async def start(self): - """Connect and authenticate via the gateway's in-band auth - protocol. Raises on auth failure.""" - - # ── Security boundary: MCP server → gateway ── - # The WebSocket connects to the gateway and authenticates using - # the caller's Bearer token via the in-band first-frame auth - # protocol. The token belongs to the MCP client — we forward - # it as-is and never interpret its contents. - self.socket = await connect(self.url) + self.socket = await connect(self._build_url()) self.pending_requests = {} self.running = True - - await self._authenticate() - self.reader_task = asyncio.create_task(self.reader()) - async def _authenticate(self): - """Send in-band auth frame and wait for auth-ok / auth-failed. - - The gateway expects ``{"type": "auth", "token": "..."}`` as the - first frame on a new WebSocket. Any service frame sent before - auth-ok is rejected. - """ - await self.socket.send(json.dumps({ - "type": "auth", - "token": self.token, - })) - - response_text = await asyncio.wait_for(self.socket.recv(), 10) - response = json.loads(response_text) - - if response.get("type") == "auth-ok": - logger.info( - "WebSocket authenticated, default workspace: %s", - response.get("workspace"), - ) - return - - # Auth failed — close immediately, do not leave an - # unauthenticated socket open. - await self.socket.close() - self.socket = None - - if response.get("type") == "auth-failed": - raise RuntimeError( - "Gateway rejected the authentication token" - ) - - raise RuntimeError( - f"Unexpected auth response type: {response.get('type')}" - ) - - async def whoami(self): - """Verify the token by calling the gateway's whoami endpoint. - Returns the identity dict and caches it on ``self.identity``. - """ - gen = self.request("iam", {"operation": "whoami"}, flow_id=None) - async for response in gen: - self.identity = response - return response - async def stop(self): self.running = False - if hasattr(self, "reader_task"): - await self.reader_task + await self.reader_task async def reader(self): - """Background task: read WebSocket frames and route them to the - correct pending-request queue by ``id``.""" + """ + Background task to read websocket responses and route to correct + request + """ while self.running: try: @@ -120,21 +59,23 @@ class WebSocketManager: request_id = response.get("id") if request_id and request_id in self.pending_requests: + # Put the response in the queue queue = self.pending_requests[request_id] await queue.put(response) else: - logger.warning( - "Response for unknown request ID: %s", request_id + logging.warning( + f"Response for unknown request ID: {request_id}" ) except Exception as e: - logger.error("Error in websocket reader: %s", e) + logging.error(f"Error in websocket reader: {e}") + # Put error in all pending queues for queue in self.pending_requests.values(): try: await queue.put({"error": str(e)}) - except Exception: + except: pass self.pending_requests.clear() @@ -145,29 +86,25 @@ class WebSocketManager: async def request( self, service, request_data, flow_id="default", - workspace=None, ): - """Send a request via WebSocket and yield responses. - - Args: - service: Gateway service name (e.g. "graph-rag", "config"). - request_data: Inner request payload. - flow_id: Optional flow identifier. ``None`` omits the field - (workspace-level services don't use flows). - workspace: Optional workspace override. When ``None`` the - gateway uses the caller's default workspace. + """ + Send a request via websocket and handle single or streaming responses """ - import time - self.last_used = time.monotonic() - + # Generate unique request ID request_id = f"{uuid.uuid4()}" + # Determine if this service streams responses + streaming_services = {"agent"} + is_streaming = service in streaming_services + + # Create a queue for all responses (streaming and single) response_queue = asyncio.Queue() self.pending_requests[request_id] = response_queue try: + # Build request message message = { "id": request_id, "service": service, @@ -177,16 +114,7 @@ class WebSocketManager: if flow_id is not None: message["flow"] = flow_id - # ── Security boundary: workspace scoping ── - # When the caller supplies a workspace, we set it on the - # message envelope. The gateway's enforce_workspace() - # validates that the authenticated identity is permitted - # to access the target workspace — we MUST NOT skip or - # override that check. When workspace is None, the - # gateway default-fills from the identity's bound workspace. - if workspace is not None: - message["workspace"] = workspace - + # Send request await self.socket.send(json.dumps(message)) while self.running: @@ -199,17 +127,19 @@ class WebSocketManager: continue if "error" in response: - if isinstance(response["error"], dict): - raise RuntimeError( - response["error"].get("message", str(response["error"])) - ) + if "message" in response["error"]: + raise RuntimeError(response["error"]["text"]) else: raise RuntimeError(str(response["error"])) yield response["response"] - if response.get("complete"): - break + if "complete" in response: + if response["complete"]: + break - finally: + except Exception as e: + # Clean up on error self.pending_requests.pop(request_id, None) + raise e +