Python API implements streaming interfaces (#577)

* Tech spec

* Python CLI utilities updated to use the API including streaming features

* Added type safety to Python API

* Completed missing auth token support in CLI

tests/unit/test_python_api_client.py

@@ -0,0 +1,446 @@
+"""
+Unit tests for TrustGraph Python API client library
+
+These tests use mocks and do not require a running server.
+"""
+
+import pytest
+from unittest.mock import Mock, patch, MagicMock, call
+import json
+
+from trustgraph.api import (
+    Api,
+    Triple,
+    AgentThought,
+    AgentObservation,
+    AgentAnswer,
+    RAGChunk,
+)
+
+
+class TestApiInstantiation:
+    """Test Api class instantiation and configuration"""
+
+    def test_api_instantiation_defaults(self):
+        """Test Api with default parameters"""
+        api = Api()
+        assert api.url == "http://localhost:8088/api/v1/"
+        assert api.timeout == 60
+        assert api.token is None
+
+    def test_api_instantiation_with_url(self):
+        """Test Api with custom URL"""
+        api = Api(url="http://test-server:9000/")
+        assert api.url == "http://test-server:9000/api/v1/"
+
+    def test_api_instantiation_with_url_trailing_slash(self):
+        """Test Api adds trailing slash if missing"""
+        api = Api(url="http://test-server:9000")
+        assert api.url == "http://test-server:9000/api/v1/"
+
+    def test_api_instantiation_with_token(self):
+        """Test Api with authentication token"""
+        api = Api(token="test-token-123")
+        assert api.token == "test-token-123"
+
+    def test_api_instantiation_with_timeout(self):
+        """Test Api with custom timeout"""
+        api = Api(timeout=120)
+        assert api.timeout == 120
+
+
+class TestApiLazyInitialization:
+    """Test lazy initialization of client components"""
+
+    def test_socket_client_lazy_init(self):
+        """Test socket client is created on first access"""
+        api = Api(url="http://test/", token="token")
+
+        assert api._socket_client is None
+        socket = api.socket()
+        assert api._socket_client is not None
+        assert socket is api._socket_client
+
+        # Second access returns same instance
+        socket2 = api.socket()
+        assert socket2 is socket
+
+    def test_bulk_client_lazy_init(self):
+        """Test bulk client is created on first access"""
+        api = Api(url="http://test/")
+
+        assert api._bulk_client is None
+        bulk = api.bulk()
+        assert api._bulk_client is not None
+
+    def test_async_flow_lazy_init(self):
+        """Test async flow is created on first access"""
+        api = Api(url="http://test/")
+
+        assert api._async_flow is None
+        async_flow = api.async_flow()
+        assert api._async_flow is not None
+
+    def test_metrics_lazy_init(self):
+        """Test metrics client is created on first access"""
+        api = Api(url="http://test/")
+
+        assert api._metrics is None
+        metrics = api.metrics()
+        assert api._metrics is not None
+
+
+class TestApiContextManager:
+    """Test context manager functionality"""
+
+    def test_sync_context_manager(self):
+        """Test synchronous context manager"""
+        with Api(url="http://test/") as api:
+            assert api is not None
+            assert isinstance(api, Api)
+        # Should exit cleanly
+
+    @pytest.mark.asyncio
+    async def test_async_context_manager(self):
+        """Test asynchronous context manager"""
+        async with Api(url="http://test/") as api:
+            assert api is not None
+            assert isinstance(api, Api)
+        # Should exit cleanly
+
+
+class TestFlowClient:
+    """Test Flow client functionality"""
+
+    @patch('requests.post')
+    def test_flow_list(self, mock_post):
+        """Test listing flows"""
+        mock_post.return_value.status_code = 200
+        mock_post.return_value.json.return_value = {"flow-ids": ["flow1", "flow2"]}
+
+        api = Api(url="http://test/")
+        flows = api.flow().list()
+
+        assert flows == ["flow1", "flow2"]
+        assert mock_post.called
+
+    @patch('requests.post')
+    def test_flow_list_with_token(self, mock_post):
+        """Test flow listing includes auth token"""
+        mock_post.return_value.status_code = 200
+        mock_post.return_value.json.return_value = {"flow-ids": []}
+
+        api = Api(url="http://test/", token="my-token")
+        api.flow().list()
+
+        # Verify Authorization header was set
+        call_args = mock_post.call_args
+        headers = call_args[1]['headers'] if 'headers' in call_args[1] else {}
+        assert 'Authorization' in headers
+        assert headers['Authorization'] == 'Bearer my-token'
+
+    @patch('requests.post')
+    def test_flow_get(self, mock_post):
+        """Test getting flow definition"""
+        flow_def = {"name": "test-flow", "description": "Test"}
+        mock_post.return_value.status_code = 200
+        mock_post.return_value.json.return_value = {"flow": json.dumps(flow_def)}
+
+        api = Api(url="http://test/")
+        result = api.flow().get("test-flow")
+
+        assert result == flow_def
+
+    def test_flow_instance_creation(self):
+        """Test creating flow instance"""
+        api = Api(url="http://test/")
+        flow_instance = api.flow().id("my-flow")
+
+        assert flow_instance is not None
+        assert flow_instance.id == "my-flow"
+
+    def test_flow_instance_has_methods(self):
+        """Test flow instance has expected methods"""
+        api = Api(url="http://test/")
+        flow_instance = api.flow().id("my-flow")
+
+        expected_methods = [
+            'text_completion', 'agent', 'graph_rag', 'document_rag',
+            'graph_embeddings_query', 'embeddings', 'prompt',
+            'triples_query', 'objects_query'
+        ]
+
+        for method in expected_methods:
+            assert hasattr(flow_instance, method), f"Missing method: {method}"
+
+
+class TestSocketClient:
+    """Test WebSocket client functionality"""
+
+    def test_socket_client_url_conversion_http(self):
+        """Test HTTP URL converted to WebSocket"""
+        api = Api(url="http://test-server:8088/")
+        socket = api.socket()
+
+        assert socket.url.startswith("ws://")
+        assert "test-server" in socket.url
+
+    def test_socket_client_url_conversion_https(self):
+        """Test HTTPS URL converted to secure WebSocket"""
+        api = Api(url="https://test-server:8088/")
+        socket = api.socket()
+
+        assert socket.url.startswith("wss://")
+
+    def test_socket_client_token_passed(self):
+        """Test token is passed to socket client"""
+        api = Api(url="http://test/", token="socket-token")
+        socket = api.socket()
+
+        assert socket.token == "socket-token"
+
+    def test_socket_flow_instance(self):
+        """Test creating socket flow instance"""
+        api = Api(url="http://test/")
+        socket = api.socket()
+        flow_instance = socket.flow("test-flow")
+
+        assert flow_instance is not None
+        assert flow_instance.flow_id == "test-flow"
+
+    def test_socket_flow_has_methods(self):
+        """Test socket flow instance has expected methods"""
+        api = Api(url="http://test/")
+        flow_instance = api.socket().flow("test-flow")
+
+        expected_methods = [
+            'agent', 'text_completion', 'graph_rag', 'document_rag',
+            'prompt', 'graph_embeddings_query', 'embeddings',
+            'triples_query', 'objects_query', 'mcp_tool'
+        ]
+
+        for method in expected_methods:
+            assert hasattr(flow_instance, method), f"Missing method: {method}"
+
+
+class TestBulkClient:
+    """Test bulk operations client"""
+
+    def test_bulk_client_url_conversion(self):
+        """Test bulk client uses WebSocket URL"""
+        api = Api(url="http://test/")
+        bulk = api.bulk()
+
+        assert bulk.url.startswith("ws://")
+
+    def test_bulk_client_has_import_methods(self):
+        """Test bulk client has import methods"""
+        api = Api(url="http://test/")
+        bulk = api.bulk()
+
+        import_methods = [
+            'import_triples',
+            'import_graph_embeddings',
+            'import_document_embeddings',
+            'import_entity_contexts',
+            'import_objects'
+        ]
+
+        for method in import_methods:
+            assert hasattr(bulk, method), f"Missing method: {method}"
+
+    def test_bulk_client_has_export_methods(self):
+        """Test bulk client has export methods"""
+        api = Api(url="http://test/")
+        bulk = api.bulk()
+
+        export_methods = [
+            'export_triples',
+            'export_graph_embeddings',
+            'export_document_embeddings',
+            'export_entity_contexts'
+        ]
+
+        for method in export_methods:
+            assert hasattr(bulk, method), f"Missing method: {method}"
+
+
+class TestMetricsClient:
+    """Test metrics client"""
+
+    @patch('requests.get')
+    def test_metrics_get(self, mock_get):
+        """Test getting metrics"""
+        mock_get.return_value.status_code = 200
+        mock_get.return_value.text = "# HELP metric_name\nmetric_name 42"
+
+        api = Api(url="http://test/")
+        metrics_text = api.metrics().get()
+
+        assert "metric_name" in metrics_text
+        assert mock_get.called
+
+    @patch('requests.get')
+    def test_metrics_with_token(self, mock_get):
+        """Test metrics request includes token"""
+        mock_get.return_value.status_code = 200
+        mock_get.return_value.text = "metrics"
+
+        api = Api(url="http://test/", token="metrics-token")
+        api.metrics().get()
+
+        # Verify token in headers
+        call_args = mock_get.call_args
+        headers = call_args[1].get('headers', {})
+        assert 'Authorization' in headers
+
+
+class TestStreamingTypes:
+    """Test streaming chunk types"""
+
+    def test_agent_thought_creation(self):
+        """Test creating AgentThought chunk"""
+        chunk = AgentThought(content="thinking...", end_of_message=False)
+
+        assert chunk.content == "thinking..."
+        assert chunk.end_of_message is False
+        assert chunk.chunk_type == "thought"
+
+    def test_agent_observation_creation(self):
+        """Test creating AgentObservation chunk"""
+        chunk = AgentObservation(content="observing...", end_of_message=False)
+
+        assert chunk.content == "observing..."
+        assert chunk.chunk_type == "observation"
+
+    def test_agent_answer_creation(self):
+        """Test creating AgentAnswer chunk"""
+        chunk = AgentAnswer(
+            content="answer",
+            end_of_message=True,
+            end_of_dialog=True
+        )
+
+        assert chunk.content == "answer"
+        assert chunk.end_of_message is True
+        assert chunk.end_of_dialog is True
+        assert chunk.chunk_type == "final-answer"
+
+    def test_rag_chunk_creation(self):
+        """Test creating RAGChunk"""
+        chunk = RAGChunk(
+            content="response chunk",
+            end_of_stream=False,
+            error=None
+        )
+
+        assert chunk.content == "response chunk"
+        assert chunk.end_of_stream is False
+        assert chunk.error is None
+
+    def test_rag_chunk_with_error(self):
+        """Test RAGChunk with error"""
+        error_dict = {"type": "error", "message": "failed"}
+        chunk = RAGChunk(
+            content="",
+            end_of_stream=True,
+            error=error_dict
+        )
+
+        assert chunk.error == error_dict
+
+
+class TestTripleType:
+    """Test Triple data structure"""
+
+    def test_triple_creation(self):
+        """Test creating Triple"""
+        triple = Triple(s="subject", p="predicate", o="object")
+
+        assert triple.s == "subject"
+        assert triple.p == "predicate"
+        assert triple.o == "object"
+
+    def test_triple_with_uris(self):
+        """Test Triple with URI values"""
+        triple = Triple(
+            s="http://example.org/entity1",
+            p="http://example.org/relation",
+            o="http://example.org/entity2"
+        )
+
+        assert triple.s.startswith("http://")
+        assert triple.p.startswith("http://")
+        assert triple.o.startswith("http://")
+
+
+class TestAsyncClients:
+    """Test async client availability"""
+
+    def test_async_flow_creation(self):
+        """Test creating async flow client"""
+        api = Api(url="http://test/")
+        async_flow = api.async_flow()
+
+        assert async_flow is not None
+
+    def test_async_socket_creation(self):
+        """Test creating async socket client"""
+        api = Api(url="http://test/")
+        async_socket = api.async_socket()
+
+        assert async_socket is not None
+        assert async_socket.url.startswith("ws://")
+
+    def test_async_bulk_creation(self):
+        """Test creating async bulk client"""
+        api = Api(url="http://test/")
+        async_bulk = api.async_bulk()
+
+        assert async_bulk is not None
+
+    def test_async_metrics_creation(self):
+        """Test creating async metrics client"""
+        api = Api(url="http://test/")
+        async_metrics = api.async_metrics()
+
+        assert async_metrics is not None
+
+
+class TestErrorHandling:
+    """Test error handling"""
+
+    @patch('requests.post')
+    def test_protocol_exception_on_non_200(self, mock_post):
+        """Test ProtocolException raised on non-200 status"""
+        from trustgraph.api.exceptions import ProtocolException
+
+        mock_post.return_value.status_code = 500
+
+        api = Api(url="http://test/")
+
+        with pytest.raises(ProtocolException):
+            api.flow().list()
+
+    @patch('requests.post')
+    def test_application_exception_on_error_response(self, mock_post):
+        """Test ApplicationException on error in response"""
+        from trustgraph.api.exceptions import ApplicationException
+
+        mock_post.return_value.status_code = 200
+        mock_post.return_value.json.return_value = {
+            "error": {
+                "type": "ValidationError",
+                "message": "Invalid input"
+            }
+        }
+
+        api = Api(url="http://test/")
+
+        with pytest.raises(ApplicationException):
+            api.flow().list()
+
+
+# Run tests with: pytest tests/unit/test_python_api_client.py -v
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

trustgraph-base/trustgraph/api/__init__.py

@@ -1,3 +1,80 @@
 
-from . api import *
+# Core API
+from .api import Api
+
+# Flow clients
+from .flow import Flow, FlowInstance
+from .async_flow import AsyncFlow, AsyncFlowInstance
+
+# WebSocket clients
+from .socket_client import SocketClient, SocketFlowInstance
+from .async_socket_client import AsyncSocketClient, AsyncSocketFlowInstance
+
+# Bulk operation clients
+from .bulk_client import BulkClient
+from .async_bulk_client import AsyncBulkClient
+
+# Metrics clients
+from .metrics import Metrics
+from .async_metrics import AsyncMetrics
+
+# Types
+from .types import (
+    Triple,
+    ConfigKey,
+    ConfigValue,
+    DocumentMetadata,
+    ProcessingMetadata,
+    CollectionMetadata,
+    StreamingChunk,
+    AgentThought,
+    AgentObservation,
+    AgentAnswer,
+    RAGChunk,
+)
+
+# Exceptions
+from .exceptions import ProtocolException, ApplicationException
+
+__all__ = [
+    # Core API
+    "Api",
+
+    # Flow clients
+    "Flow",
+    "FlowInstance",
+    "AsyncFlow",
+    "AsyncFlowInstance",
+
+    # WebSocket clients
+    "SocketClient",
+    "SocketFlowInstance",
+    "AsyncSocketClient",
+    "AsyncSocketFlowInstance",
+
+    # Bulk operation clients
+    "BulkClient",
+    "AsyncBulkClient",
+
+    # Metrics clients
+    "Metrics",
+    "AsyncMetrics",
+
+    # Types
+    "Triple",
+    "ConfigKey",
+    "ConfigValue",
+    "DocumentMetadata",
+    "ProcessingMetadata",
+    "CollectionMetadata",
+    "StreamingChunk",
+    "AgentThought",
+    "AgentObservation",
+    "AgentAnswer",
+    "RAGChunk",
+
+    # Exceptions
+    "ProtocolException",
+    "ApplicationException",
+]
 

trustgraph-base/trustgraph/api/api.py

@@ -3,6 +3,7 @@ import requests
 import json
 import base64
 import time
+from typing import Optional
 
 from . library import Library
 from . flow import Flow
@@ -26,7 +27,7 @@ def check_error(response):
 
 class Api:
 
-    def __init__(self, url="http://localhost:8088/", timeout=60):
+    def __init__(self, url="http://localhost:8088/", timeout=60, token: Optional[str] = None):
 
         self.url = url
 
@@ -36,6 +37,16 @@ class Api:
         self.url += "api/v1/"
 
         self.timeout = timeout
+        self.token = token
+
+        # Lazy initialization for new clients
+        self._socket_client = None
+        self._bulk_client = None
+        self._async_flow = None
+        self._async_socket_client = None
+        self._async_bulk_client = None
+        self._metrics = None
+        self._async_metrics = None
 
     def flow(self):
         return Flow(api=self)
@@ -50,8 +61,12 @@ class Api:
 
         url = f"{self.url}{path}"
 
+        headers = {}
+        if self.token:
+            headers["Authorization"] = f"Bearer {self.token}"
+
         # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=request, timeout=self.timeout)
+        resp = requests.post(url, json=request, timeout=self.timeout, headers=headers)
 
         # Should be a 200 status code
         if resp.status_code != 200:
@@ -72,3 +87,96 @@ class Api:
 
     def collection(self):
         return Collection(self)
+
+    # New synchronous methods
+    def socket(self):
+        """Synchronous WebSocket-based interface for streaming operations"""
+        if self._socket_client is None:
+            from . socket_client import SocketClient
+            # Extract base URL (remove api/v1/ suffix)
+            base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/")
+            self._socket_client = SocketClient(base_url, self.timeout, self.token)
+        return self._socket_client
+
+    def bulk(self):
+        """Synchronous bulk operations interface for import/export"""
+        if self._bulk_client is None:
+            from . bulk_client import BulkClient
+            # Extract base URL (remove api/v1/ suffix)
+            base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/")
+            self._bulk_client = BulkClient(base_url, self.timeout, self.token)
+        return self._bulk_client
+
+    def metrics(self):
+        """Synchronous metrics interface"""
+        if self._metrics is None:
+            from . metrics import Metrics
+            # Extract base URL (remove api/v1/ suffix)
+            base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/")
+            self._metrics = Metrics(base_url, self.timeout, self.token)
+        return self._metrics
+
+    # New asynchronous methods
+    def async_flow(self):
+        """Asynchronous REST-based flow interface"""
+        if self._async_flow is None:
+            from . async_flow import AsyncFlow
+            self._async_flow = AsyncFlow(self.url, self.timeout, self.token)
+        return self._async_flow
+
+    def async_socket(self):
+        """Asynchronous WebSocket-based interface for streaming operations"""
+        if self._async_socket_client is None:
+            from . async_socket_client import AsyncSocketClient
+            # Extract base URL (remove api/v1/ suffix)
+            base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/")
+            self._async_socket_client = AsyncSocketClient(base_url, self.timeout, self.token)
+        return self._async_socket_client
+
+    def async_bulk(self):
+        """Asynchronous bulk operations interface for import/export"""
+        if self._async_bulk_client is None:
+            from . async_bulk_client import AsyncBulkClient
+            # Extract base URL (remove api/v1/ suffix)
+            base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/")
+            self._async_bulk_client = AsyncBulkClient(base_url, self.timeout, self.token)
+        return self._async_bulk_client
+
+    def async_metrics(self):
+        """Asynchronous metrics interface"""
+        if self._async_metrics is None:
+            from . async_metrics import AsyncMetrics
+            # Extract base URL (remove api/v1/ suffix)
+            base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/")
+            self._async_metrics = AsyncMetrics(base_url, self.timeout, self.token)
+        return self._async_metrics
+
+    # Resource management
+    def close(self):
+        """Close all synchronous connections"""
+        if self._socket_client:
+            self._socket_client.close()
+        if self._bulk_client:
+            self._bulk_client.close()
+
+    async def aclose(self):
+        """Close all asynchronous connections"""
+        if self._async_socket_client:
+            await self._async_socket_client.aclose()
+        if self._async_bulk_client:
+            await self._async_bulk_client.aclose()
+        if self._async_flow:
+            await self._async_flow.aclose()
+
+    # Context manager support
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        await self.aclose()

trustgraph-base/trustgraph/api/async_bulk_client.py

@@ -0,0 +1,131 @@
+
+import json
+import websockets
+from typing import Optional, AsyncIterator, Dict, Any, Iterator
+
+from . types import Triple
+
+
+class AsyncBulkClient:
+    """Asynchronous bulk operations client"""
+
+    def __init__(self, url: str, timeout: int, token: Optional[str]) -> None:
+        self.url: str = self._convert_to_ws_url(url)
+        self.timeout: int = timeout
+        self.token: Optional[str] = token
+
+    def _convert_to_ws_url(self, url: str) -> str:
+        """Convert HTTP URL to WebSocket URL"""
+        if url.startswith("http://"):
+            return url.replace("http://", "ws://", 1)
+        elif url.startswith("https://"):
+            return url.replace("https://", "wss://", 1)
+        elif url.startswith("ws://") or url.startswith("wss://"):
+            return url
+        else:
+            return f"ws://{url}"
+
+    async def import_triples(self, flow: str, triples: AsyncIterator[Triple], **kwargs: Any) -> None:
+        """Bulk import triples via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/triples"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for triple in triples:
+                message = {
+                    "s": triple.s,
+                    "p": triple.p,
+                    "o": triple.o
+                }
+                await websocket.send(json.dumps(message))
+
+    async def export_triples(self, flow: str, **kwargs: Any) -> AsyncIterator[Triple]:
+        """Bulk export triples via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/export/triples"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for raw_message in websocket:
+                data = json.loads(raw_message)
+                yield Triple(
+                    s=data.get("s", ""),
+                    p=data.get("p", ""),
+                    o=data.get("o", "")
+                )
+
+    async def import_graph_embeddings(self, flow: str, embeddings: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
+        """Bulk import graph embeddings via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/graph-embeddings"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for embedding in embeddings:
+                await websocket.send(json.dumps(embedding))
+
+    async def export_graph_embeddings(self, flow: str, **kwargs: Any) -> AsyncIterator[Dict[str, Any]]:
+        """Bulk export graph embeddings via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/export/graph-embeddings"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for raw_message in websocket:
+                yield json.loads(raw_message)
+
+    async def import_document_embeddings(self, flow: str, embeddings: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
+        """Bulk import document embeddings via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/document-embeddings"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for embedding in embeddings:
+                await websocket.send(json.dumps(embedding))
+
+    async def export_document_embeddings(self, flow: str, **kwargs: Any) -> AsyncIterator[Dict[str, Any]]:
+        """Bulk export document embeddings via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/export/document-embeddings"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for raw_message in websocket:
+                yield json.loads(raw_message)
+
+    async def import_entity_contexts(self, flow: str, contexts: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
+        """Bulk import entity contexts via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/entity-contexts"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for context in contexts:
+                await websocket.send(json.dumps(context))
+
+    async def export_entity_contexts(self, flow: str, **kwargs: Any) -> AsyncIterator[Dict[str, Any]]:
+        """Bulk export entity contexts via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/export/entity-contexts"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for raw_message in websocket:
+                yield json.loads(raw_message)
+
+    async def import_objects(self, flow: str, objects: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
+        """Bulk import objects via WebSocket"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/objects"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for obj in objects:
+                await websocket.send(json.dumps(obj))
+
+    async def aclose(self) -> None:
+        """Close connections"""
+        # Cleanup handled by context managers
+        pass

trustgraph-base/trustgraph/api/async_flow.py

@@ -0,0 +1,245 @@
+
+import aiohttp
+import json
+from typing import Optional, Dict, Any, List
+
+from . exceptions import ProtocolException, ApplicationException
+
+
+def check_error(response):
+    if "error" in response:
+        try:
+            msg = response["error"]["message"]
+            tp = response["error"]["type"]
+        except:
+            raise ApplicationException(response["error"])
+
+        raise ApplicationException(f"{tp}: {msg}")
+
+
+class AsyncFlow:
+    """Asynchronous REST-based flow interface"""
+
+    def __init__(self, url: str, timeout: int, token: Optional[str]) -> None:
+        self.url: str = url
+        self.timeout: int = timeout
+        self.token: Optional[str] = token
+
+    async def request(self, path: str, request_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Make async HTTP request to Gateway API"""
+        url = f"{self.url}{path}"
+
+        headers = {"Content-Type": "application/json"}
+        if self.token:
+            headers["Authorization"] = f"Bearer {self.token}"
+
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.post(url, json=request_data, headers=headers) as resp:
+                if resp.status != 200:
+                    raise ProtocolException(f"Status code {resp.status}")
+
+                try:
+                    obj = await resp.json()
+                except:
+                    raise ProtocolException(f"Expected JSON response")
+
+                check_error(obj)
+                return obj
+
+    async def list(self) -> List[str]:
+        """List all flows"""
+        result = await self.request("flow", {"operation": "list-flows"})
+        return result.get("flow-ids", [])
+
+    async def get(self, id: str) -> Dict[str, Any]:
+        """Get flow definition"""
+        result = await self.request("flow", {
+            "operation": "get-flow",
+            "flow-id": id
+        })
+        return json.loads(result.get("flow", "{}"))
+
+    async def start(self, class_name: str, id: str, description: str, parameters: Optional[Dict] = None):
+        """Start a flow"""
+        request_data = {
+            "operation": "start-flow",
+            "flow-id": id,
+            "class-name": class_name,
+            "description": description
+        }
+        if parameters:
+            request_data["parameters"] = json.dumps(parameters)
+
+        await self.request("flow", request_data)
+
+    async def stop(self, id: str):
+        """Stop a flow"""
+        await self.request("flow", {
+            "operation": "stop-flow",
+            "flow-id": id
+        })
+
+    async def list_classes(self) -> List[str]:
+        """List flow classes"""
+        result = await self.request("flow", {"operation": "list-classes"})
+        return result.get("class-names", [])
+
+    async def get_class(self, class_name: str) -> Dict[str, Any]:
+        """Get flow class definition"""
+        result = await self.request("flow", {
+            "operation": "get-class",
+            "class-name": class_name
+        })
+        return json.loads(result.get("class-definition", "{}"))
+
+    async def put_class(self, class_name: str, definition: Dict[str, Any]):
+        """Create/update flow class"""
+        await self.request("flow", {
+            "operation": "put-class",
+            "class-name": class_name,
+            "class-definition": json.dumps(definition)
+        })
+
+    async def delete_class(self, class_name: str):
+        """Delete flow class"""
+        await self.request("flow", {
+            "operation": "delete-class",
+            "class-name": class_name
+        })
+
+    def id(self, flow_id: str):
+        """Get async flow instance"""
+        return AsyncFlowInstance(self, flow_id)
+
+    async def aclose(self) -> None:
+        """Close connection (cleanup handled by aiohttp session)"""
+        pass
+
+
+class AsyncFlowInstance:
+    """Asynchronous REST flow instance"""
+
+    def __init__(self, flow: AsyncFlow, flow_id: str):
+        self.flow = flow
+        self.flow_id = flow_id
+
+    async def request(self, service: str, request_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Make request to flow-scoped service"""
+        return await self.flow.request(f"flow/{self.flow_id}/service/{service}", request_data)
+
+    async def agent(self, question: str, user: str, state: Optional[Dict] = None,
+                    group: Optional[str] = None, history: Optional[List] = None, **kwargs: Any) -> Dict[str, Any]:
+        """Execute agent (non-streaming, use async_socket for streaming)"""
+        request_data = {
+            "question": question,
+            "user": user,
+            "streaming": False  # REST doesn't support streaming
+        }
+        if state is not None:
+            request_data["state"] = state
+        if group is not None:
+            request_data["group"] = group
+        if history is not None:
+            request_data["history"] = history
+        request_data.update(kwargs)
+
+        return await self.request("agent", request_data)
+
+    async def text_completion(self, system: str, prompt: str, **kwargs: Any) -> str:
+        """Text completion (non-streaming, use async_socket for streaming)"""
+        request_data = {
+            "system": system,
+            "prompt": prompt,
+            "streaming": False
+        }
+        request_data.update(kwargs)
+
+        result = await self.request("text-completion", request_data)
+        return result.get("response", "")
+
+    async def graph_rag(self, question: str, user: str, collection: str,
+                        max_subgraph_size: int = 1000, max_subgraph_count: int = 5,
+                        max_entity_distance: int = 3, **kwargs: Any) -> str:
+        """Graph RAG (non-streaming, use async_socket for streaming)"""
+        request_data = {
+            "question": question,
+            "user": user,
+            "collection": collection,
+            "max-subgraph-size": max_subgraph_size,
+            "max-subgraph-count": max_subgraph_count,
+            "max-entity-distance": max_entity_distance,
+            "streaming": False
+        }
+        request_data.update(kwargs)
+
+        result = await self.request("graph-rag", request_data)
+        return result.get("response", "")
+
+    async def document_rag(self, question: str, user: str, collection: str,
+                           doc_limit: int = 10, **kwargs: Any) -> str:
+        """Document RAG (non-streaming, use async_socket for streaming)"""
+        request_data = {
+            "question": question,
+            "user": user,
+            "collection": collection,
+            "doc-limit": doc_limit,
+            "streaming": False
+        }
+        request_data.update(kwargs)
+
+        result = await self.request("document-rag", request_data)
+        return result.get("response", "")
+
+    async def graph_embeddings_query(self, text: str, user: str, collection: str, limit: int = 10, **kwargs: Any):
+        """Query graph embeddings for semantic search"""
+        request_data = {
+            "text": text,
+            "user": user,
+            "collection": collection,
+            "limit": limit
+        }
+        request_data.update(kwargs)
+
+        return await self.request("graph-embeddings", request_data)
+
+    async def embeddings(self, text: str, **kwargs: Any):
+        """Generate text embeddings"""
+        request_data = {"text": text}
+        request_data.update(kwargs)
+
+        return await self.request("embeddings", request_data)
+
+    async def triples_query(self, s=None, p=None, o=None, user=None, collection=None, limit=100, **kwargs: Any):
+        """Triple pattern query"""
+        request_data = {"limit": limit}
+        if s is not None:
+            request_data["s"] = str(s)
+        if p is not None:
+            request_data["p"] = str(p)
+        if o is not None:
+            request_data["o"] = str(o)
+        if user is not None:
+            request_data["user"] = user
+        if collection is not None:
+            request_data["collection"] = collection
+        request_data.update(kwargs)
+
+        return await self.request("triples", request_data)
+
+    async def objects_query(self, query: str, user: str, collection: str, variables: Optional[Dict] = None,
+                            operation_name: Optional[str] = None, **kwargs: Any):
+        """GraphQL query"""
+        request_data = {
+            "query": query,
+            "user": user,
+            "collection": collection
+        }
+        if variables:
+            request_data["variables"] = variables
+        if operation_name:
+            request_data["operationName"] = operation_name
+        request_data.update(kwargs)
+
+        return await self.request("objects", request_data)

trustgraph-base/trustgraph/api/async_metrics.py

@@ -0,0 +1,33 @@
+
+import aiohttp
+from typing import Optional, Dict
+
+
+class AsyncMetrics:
+    """Asynchronous metrics client"""
+
+    def __init__(self, url: str, timeout: int, token: Optional[str]) -> None:
+        self.url: str = url
+        self.timeout: int = timeout
+        self.token: Optional[str] = token
+
+    async def get(self) -> str:
+        """Get Prometheus metrics as text"""
+        url: str = f"{self.url}/api/metrics"
+
+        headers: Dict[str, str] = {}
+        if self.token:
+            headers["Authorization"] = f"Bearer {self.token}"
+
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.get(url, headers=headers) as resp:
+                if resp.status != 200:
+                    raise Exception(f"Status code {resp.status}")
+
+                return await resp.text()
+
+    async def aclose(self) -> None:
+        """Close connections"""
+        pass

trustgraph-base/trustgraph/api/async_socket_client.py

@@ -0,0 +1,335 @@
+
+import json
+import websockets
+from typing import Optional, Dict, Any, AsyncIterator, Union
+
+from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk
+from . exceptions import ProtocolException, ApplicationException
+
+
+class AsyncSocketClient:
+    """Asynchronous WebSocket client"""
+
+    def __init__(self, url: str, timeout: int, token: Optional[str]):
+        self.url = self._convert_to_ws_url(url)
+        self.timeout = timeout
+        self.token = token
+        self._request_counter = 0
+
+    def _convert_to_ws_url(self, url: str) -> str:
+        """Convert HTTP URL to WebSocket URL"""
+        if url.startswith("http://"):
+            return url.replace("http://", "ws://", 1)
+        elif url.startswith("https://"):
+            return url.replace("https://", "wss://", 1)
+        elif url.startswith("ws://") or url.startswith("wss://"):
+            return url
+        else:
+            # Assume ws://
+            return f"ws://{url}"
+
+    def flow(self, flow_id: str):
+        """Get async flow instance for WebSocket operations"""
+        return AsyncSocketFlowInstance(self, flow_id)
+
+    async def _send_request(self, service: str, flow: Optional[str], request: Dict[str, Any]):
+        """Async WebSocket request implementation (non-streaming)"""
+        # Generate unique request ID
+        self._request_counter += 1
+        request_id = f"req-{self._request_counter}"
+
+        # Build WebSocket URL with optional token
+        ws_url = f"{self.url}/api/v1/socket"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        # Build request message
+        message = {
+            "id": request_id,
+            "service": service,
+            "request": request
+        }
+        if flow:
+            message["flow"] = flow
+
+        # Connect and send request
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            await websocket.send(json.dumps(message))
+
+            # Wait for single response
+            raw_message = await websocket.recv()
+            response = json.loads(raw_message)
+
+            if response.get("id") != request_id:
+                raise ProtocolException(f"Response ID mismatch")
+
+            if "error" in response:
+                raise ApplicationException(response["error"])
+
+            if "response" not in response:
+                raise ProtocolException(f"Missing response in message")
+
+            return response["response"]
+
+    async def _send_request_streaming(self, service: str, flow: Optional[str], request: Dict[str, Any]):
+        """Async WebSocket request implementation (streaming)"""
+        # Generate unique request ID
+        self._request_counter += 1
+        request_id = f"req-{self._request_counter}"
+
+        # Build WebSocket URL with optional token
+        ws_url = f"{self.url}/api/v1/socket"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        # Build request message
+        message = {
+            "id": request_id,
+            "service": service,
+            "request": request
+        }
+        if flow:
+            message["flow"] = flow
+
+        # Connect and send request
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            await websocket.send(json.dumps(message))
+
+            # Yield chunks as they arrive
+            async for raw_message in websocket:
+                response = json.loads(raw_message)
+
+                if response.get("id") != request_id:
+                    continue  # Ignore messages for other requests
+
+                if "error" in response:
+                    raise ApplicationException(response["error"])
+
+                if "response" in response:
+                    resp = response["response"]
+
+                    # Parse different chunk types
+                    chunk = self._parse_chunk(resp)
+                    yield chunk
+
+                    # Check if this is the final chunk
+                    if resp.get("end_of_stream") or resp.get("end_of_dialog") or response.get("complete"):
+                        break
+
+    def _parse_chunk(self, resp: Dict[str, Any]):
+        """Parse response chunk into appropriate type"""
+        chunk_type = resp.get("chunk_type")
+
+        if chunk_type == "thought":
+            return AgentThought(
+                content=resp.get("content", ""),
+                end_of_message=resp.get("end_of_message", False)
+            )
+        elif chunk_type == "observation":
+            return AgentObservation(
+                content=resp.get("content", ""),
+                end_of_message=resp.get("end_of_message", False)
+            )
+        elif chunk_type == "final-answer":
+            return AgentAnswer(
+                content=resp.get("content", ""),
+                end_of_message=resp.get("end_of_message", False),
+                end_of_dialog=resp.get("end_of_dialog", False)
+            )
+        else:
+            # RAG-style chunk (or generic chunk)
+            return RAGChunk(
+                content=resp.get("chunk", ""),
+                end_of_stream=resp.get("end_of_stream", False),
+                error=resp.get("error")
+            )
+
+    async def aclose(self):
+        """Close WebSocket connection"""
+        # Cleanup handled by context manager
+        pass
+
+
+class AsyncSocketFlowInstance:
+    """Asynchronous WebSocket flow instance"""
+
+    def __init__(self, client: AsyncSocketClient, flow_id: str):
+        self.client = client
+        self.flow_id = flow_id
+
+    async def agent(self, question: str, user: str, state: Optional[Dict[str, Any]] = None,
+                    group: Optional[str] = None, history: Optional[list] = None,
+                    streaming: bool = False, **kwargs) -> Union[Dict[str, Any], AsyncIterator]:
+        """Agent with optional streaming"""
+        request = {
+            "question": question,
+            "user": user,
+            "streaming": streaming
+        }
+        if state is not None:
+            request["state"] = state
+        if group is not None:
+            request["group"] = group
+        if history is not None:
+            request["history"] = history
+        request.update(kwargs)
+
+        if streaming:
+            return self.client._send_request_streaming("agent", self.flow_id, request)
+        else:
+            return await self.client._send_request("agent", self.flow_id, request)
+
+    async def text_completion(self, system: str, prompt: str, streaming: bool = False, **kwargs):
+        """Text completion with optional streaming"""
+        request = {
+            "system": system,
+            "prompt": prompt,
+            "streaming": streaming
+        }
+        request.update(kwargs)
+
+        if streaming:
+            return self._text_completion_streaming(request)
+        else:
+            result = await self.client._send_request("text-completion", self.flow_id, request)
+            return result.get("response", "")
+
+    async def _text_completion_streaming(self, request):
+        """Helper for streaming text completion"""
+        async for chunk in self.client._send_request_streaming("text-completion", self.flow_id, request):
+            if hasattr(chunk, 'content'):
+                yield chunk.content
+
+    async def graph_rag(self, question: str, user: str, collection: str,
+                        max_subgraph_size: int = 1000, max_subgraph_count: int = 5,
+                        max_entity_distance: int = 3, streaming: bool = False, **kwargs):
+        """Graph RAG with optional streaming"""
+        request = {
+            "question": question,
+            "user": user,
+            "collection": collection,
+            "max-subgraph-size": max_subgraph_size,
+            "max-subgraph-count": max_subgraph_count,
+            "max-entity-distance": max_entity_distance,
+            "streaming": streaming
+        }
+        request.update(kwargs)
+
+        if streaming:
+            return self._graph_rag_streaming(request)
+        else:
+            result = await self.client._send_request("graph-rag", self.flow_id, request)
+            return result.get("response", "")
+
+    async def _graph_rag_streaming(self, request):
+        """Helper for streaming graph RAG"""
+        async for chunk in self.client._send_request_streaming("graph-rag", self.flow_id, request):
+            if hasattr(chunk, 'content'):
+                yield chunk.content
+
+    async def document_rag(self, question: str, user: str, collection: str,
+                           doc_limit: int = 10, streaming: bool = False, **kwargs):
+        """Document RAG with optional streaming"""
+        request = {
+            "question": question,
+            "user": user,
+            "collection": collection,
+            "doc-limit": doc_limit,
+            "streaming": streaming
+        }
+        request.update(kwargs)
+
+        if streaming:
+            return self._document_rag_streaming(request)
+        else:
+            result = await self.client._send_request("document-rag", self.flow_id, request)
+            return result.get("response", "")
+
+    async def _document_rag_streaming(self, request):
+        """Helper for streaming document RAG"""
+        async for chunk in self.client._send_request_streaming("document-rag", self.flow_id, request):
+            if hasattr(chunk, 'content'):
+                yield chunk.content
+
+    async def prompt(self, id: str, variables: Dict[str, str], streaming: bool = False, **kwargs):
+        """Execute prompt with optional streaming"""
+        request = {
+            "id": id,
+            "variables": variables,
+            "streaming": streaming
+        }
+        request.update(kwargs)
+
+        if streaming:
+            return self._prompt_streaming(request)
+        else:
+            result = await self.client._send_request("prompt", self.flow_id, request)
+            return result.get("response", "")
+
+    async def _prompt_streaming(self, request):
+        """Helper for streaming prompt"""
+        async for chunk in self.client._send_request_streaming("prompt", self.flow_id, request):
+            if hasattr(chunk, 'content'):
+                yield chunk.content
+
+    async def graph_embeddings_query(self, text: str, user: str, collection: str, limit: int = 10, **kwargs):
+        """Query graph embeddings for semantic search"""
+        request = {
+            "text": text,
+            "user": user,
+            "collection": collection,
+            "limit": limit
+        }
+        request.update(kwargs)
+
+        return await self.client._send_request("graph-embeddings", self.flow_id, request)
+
+    async def embeddings(self, text: str, **kwargs):
+        """Generate text embeddings"""
+        request = {"text": text}
+        request.update(kwargs)
+
+        return await self.client._send_request("embeddings", self.flow_id, request)
+
+    async def triples_query(self, s=None, p=None, o=None, user=None, collection=None, limit=100, **kwargs):
+        """Triple pattern query"""
+        request = {"limit": limit}
+        if s is not None:
+            request["s"] = str(s)
+        if p is not None:
+            request["p"] = str(p)
+        if o is not None:
+            request["o"] = str(o)
+        if user is not None:
+            request["user"] = user
+        if collection is not None:
+            request["collection"] = collection
+        request.update(kwargs)
+
+        return await self.client._send_request("triples", self.flow_id, request)
+
+    async def objects_query(self, query: str, user: str, collection: str, variables: Optional[Dict] = None,
+                            operation_name: Optional[str] = None, **kwargs):
+        """GraphQL query"""
+        request = {
+            "query": query,
+            "user": user,
+            "collection": collection
+        }
+        if variables:
+            request["variables"] = variables
+        if operation_name:
+            request["operationName"] = operation_name
+        request.update(kwargs)
+
+        return await self.client._send_request("objects", self.flow_id, request)
+
+    async def mcp_tool(self, name: str, parameters: Dict[str, Any], **kwargs):
+        """Execute MCP tool"""
+        request = {
+            "name": name,
+            "parameters": parameters
+        }
+        request.update(kwargs)
+
+        return await self.client._send_request("mcp-tool", self.flow_id, request)

trustgraph-base/trustgraph/api/bulk_client.py

@@ -0,0 +1,270 @@
+
+import json
+import asyncio
+import websockets
+from typing import Optional, Iterator, Dict, Any, Coroutine
+
+from . types import Triple
+from . exceptions import ProtocolException
+
+
+class BulkClient:
+    """Synchronous bulk operations client"""
+
+    def __init__(self, url: str, timeout: int, token: Optional[str]) -> None:
+        self.url: str = self._convert_to_ws_url(url)
+        self.timeout: int = timeout
+        self.token: Optional[str] = token
+
+    def _convert_to_ws_url(self, url: str) -> str:
+        """Convert HTTP URL to WebSocket URL"""
+        if url.startswith("http://"):
+            return url.replace("http://", "ws://", 1)
+        elif url.startswith("https://"):
+            return url.replace("https://", "wss://", 1)
+        elif url.startswith("ws://") or url.startswith("wss://"):
+            return url
+        else:
+            return f"ws://{url}"
+
+    def _run_async(self, coro: Coroutine[Any, Any, Any]) -> Any:
+        """Run async coroutine synchronously"""
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        return loop.run_until_complete(coro)
+
+    def import_triples(self, flow: str, triples: Iterator[Triple], **kwargs: Any) -> None:
+        """Bulk import triples via WebSocket"""
+        self._run_async(self._import_triples_async(flow, triples))
+
+    async def _import_triples_async(self, flow: str, triples: Iterator[Triple]) -> None:
+        """Async implementation of triple import"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/triples"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            for triple in triples:
+                message = {
+                    "s": triple.s,
+                    "p": triple.p,
+                    "o": triple.o
+                }
+                await websocket.send(json.dumps(message))
+
+    def export_triples(self, flow: str, **kwargs: Any) -> Iterator[Triple]:
+        """Bulk export triples via WebSocket"""
+        async_gen = self._export_triples_async(flow)
+
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        try:
+            while True:
+                try:
+                    triple = loop.run_until_complete(async_gen.__anext__())
+                    yield triple
+                except StopAsyncIteration:
+                    break
+        finally:
+            try:
+                loop.run_until_complete(async_gen.aclose())
+            except:
+                pass
+
+    async def _export_triples_async(self, flow: str) -> Iterator[Triple]:
+        """Async implementation of triple export"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/export/triples"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for raw_message in websocket:
+                data = json.loads(raw_message)
+                yield Triple(
+                    s=data.get("s", ""),
+                    p=data.get("p", ""),
+                    o=data.get("o", "")
+                )
+
+    def import_graph_embeddings(self, flow: str, embeddings: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
+        """Bulk import graph embeddings via WebSocket"""
+        self._run_async(self._import_graph_embeddings_async(flow, embeddings))
+
+    async def _import_graph_embeddings_async(self, flow: str, embeddings: Iterator[Dict[str, Any]]) -> None:
+        """Async implementation of graph embeddings import"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/graph-embeddings"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            for embedding in embeddings:
+                await websocket.send(json.dumps(embedding))
+
+    def export_graph_embeddings(self, flow: str, **kwargs: Any) -> Iterator[Dict[str, Any]]:
+        """Bulk export graph embeddings via WebSocket"""
+        async_gen = self._export_graph_embeddings_async(flow)
+
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        try:
+            while True:
+                try:
+                    embedding = loop.run_until_complete(async_gen.__anext__())
+                    yield embedding
+                except StopAsyncIteration:
+                    break
+        finally:
+            try:
+                loop.run_until_complete(async_gen.aclose())
+            except:
+                pass
+
+    async def _export_graph_embeddings_async(self, flow: str) -> Iterator[Dict[str, Any]]:
+        """Async implementation of graph embeddings export"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/export/graph-embeddings"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for raw_message in websocket:
+                yield json.loads(raw_message)
+
+    def import_document_embeddings(self, flow: str, embeddings: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
+        """Bulk import document embeddings via WebSocket"""
+        self._run_async(self._import_document_embeddings_async(flow, embeddings))
+
+    async def _import_document_embeddings_async(self, flow: str, embeddings: Iterator[Dict[str, Any]]) -> None:
+        """Async implementation of document embeddings import"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/document-embeddings"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            for embedding in embeddings:
+                await websocket.send(json.dumps(embedding))
+
+    def export_document_embeddings(self, flow: str, **kwargs: Any) -> Iterator[Dict[str, Any]]:
+        """Bulk export document embeddings via WebSocket"""
+        async_gen = self._export_document_embeddings_async(flow)
+
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        try:
+            while True:
+                try:
+                    embedding = loop.run_until_complete(async_gen.__anext__())
+                    yield embedding
+                except StopAsyncIteration:
+                    break
+        finally:
+            try:
+                loop.run_until_complete(async_gen.aclose())
+            except:
+                pass
+
+    async def _export_document_embeddings_async(self, flow: str) -> Iterator[Dict[str, Any]]:
+        """Async implementation of document embeddings export"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/export/document-embeddings"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for raw_message in websocket:
+                yield json.loads(raw_message)
+
+    def import_entity_contexts(self, flow: str, contexts: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
+        """Bulk import entity contexts via WebSocket"""
+        self._run_async(self._import_entity_contexts_async(flow, contexts))
+
+    async def _import_entity_contexts_async(self, flow: str, contexts: Iterator[Dict[str, Any]]) -> None:
+        """Async implementation of entity contexts import"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/entity-contexts"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            for context in contexts:
+                await websocket.send(json.dumps(context))
+
+    def export_entity_contexts(self, flow: str, **kwargs: Any) -> Iterator[Dict[str, Any]]:
+        """Bulk export entity contexts via WebSocket"""
+        async_gen = self._export_entity_contexts_async(flow)
+
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        try:
+            while True:
+                try:
+                    context = loop.run_until_complete(async_gen.__anext__())
+                    yield context
+                except StopAsyncIteration:
+                    break
+        finally:
+            try:
+                loop.run_until_complete(async_gen.aclose())
+            except:
+                pass
+
+    async def _export_entity_contexts_async(self, flow: str) -> Iterator[Dict[str, Any]]:
+        """Async implementation of entity contexts export"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/export/entity-contexts"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            async for raw_message in websocket:
+                yield json.loads(raw_message)
+
+    def import_objects(self, flow: str, objects: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
+        """Bulk import objects via WebSocket"""
+        self._run_async(self._import_objects_async(flow, objects))
+
+    async def _import_objects_async(self, flow: str, objects: Iterator[Dict[str, Any]]) -> None:
+        """Async implementation of objects import"""
+        ws_url = f"{self.url}/api/v1/flow/{flow}/import/objects"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            for obj in objects:
+                await websocket.send(json.dumps(obj))
+
+    def close(self) -> None:
+        """Close connections"""
+        # Cleanup handled by context managers
+        pass

trustgraph-base/trustgraph/api/flow.py

@@ -211,6 +211,21 @@ class FlowInstance:
             input
         )["vectors"]
 
+    def graph_embeddings_query(self, text, user, collection, limit=10):
+
+        # Query graph embeddings for semantic search
+        input = {
+            "text": text,
+            "user": user,
+            "collection": collection,
+            "limit": limit
+        }
+
+        return self.request(
+            "service/graph-embeddings",
+            input
+        )
+
     def prompt(self, id, variables):
 
         input = {

trustgraph-base/trustgraph/api/metrics.py

@@ -0,0 +1,27 @@
+
+import requests
+from typing import Optional, Dict
+
+
+class Metrics:
+    """Synchronous metrics client"""
+
+    def __init__(self, url: str, timeout: int, token: Optional[str]) -> None:
+        self.url: str = url
+        self.timeout: int = timeout
+        self.token: Optional[str] = token
+
+    def get(self) -> str:
+        """Get Prometheus metrics as text"""
+        url: str = f"{self.url}/api/metrics"
+
+        headers: Dict[str, str] = {}
+        if self.token:
+            headers["Authorization"] = f"Bearer {self.token}"
+
+        resp = requests.get(url, timeout=self.timeout, headers=headers)
+
+        if resp.status_code != 200:
+            raise Exception(f"Status code {resp.status_code}")
+
+        return resp.text

trustgraph-base/trustgraph/api/socket_client.py

@@ -0,0 +1,445 @@
+
+import json
+import asyncio
+import websockets
+from typing import Optional, Dict, Any, Iterator, Union, List
+from threading import Lock
+
+from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, StreamingChunk
+from . exceptions import ProtocolException, ApplicationException
+
+
+class SocketClient:
+    """Synchronous WebSocket client (wraps async websockets library)"""
+
+    def __init__(self, url: str, timeout: int, token: Optional[str]) -> None:
+        self.url: str = self._convert_to_ws_url(url)
+        self.timeout: int = timeout
+        self.token: Optional[str] = token
+        self._connection: Optional[Any] = None
+        self._request_counter: int = 0
+        self._lock: Lock = Lock()
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+
+    def _convert_to_ws_url(self, url: str) -> str:
+        """Convert HTTP URL to WebSocket URL"""
+        if url.startswith("http://"):
+            return url.replace("http://", "ws://", 1)
+        elif url.startswith("https://"):
+            return url.replace("https://", "wss://", 1)
+        elif url.startswith("ws://") or url.startswith("wss://"):
+            return url
+        else:
+            # Assume ws://
+            return f"ws://{url}"
+
+    def flow(self, flow_id: str) -> "SocketFlowInstance":
+        """Get flow instance for WebSocket operations"""
+        return SocketFlowInstance(self, flow_id)
+
+    def _send_request_sync(
+        self,
+        service: str,
+        flow: Optional[str],
+        request: Dict[str, Any],
+        streaming: bool = False
+    ) -> Union[Dict[str, Any], Iterator[StreamingChunk]]:
+        """Synchronous wrapper around async WebSocket communication"""
+        # Create event loop if needed
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # If loop is running (e.g., in Jupyter), create new loop
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        if streaming:
+            # For streaming, we need to return an iterator
+            # Create a generator that runs async code
+            return self._streaming_generator(service, flow, request, loop)
+        else:
+            # For non-streaming, just run the async code and return result
+            return loop.run_until_complete(self._send_request_async(service, flow, request))
+
+    def _streaming_generator(
+        self,
+        service: str,
+        flow: Optional[str],
+        request: Dict[str, Any],
+        loop: asyncio.AbstractEventLoop
+    ) -> Iterator[StreamingChunk]:
+        """Generator that yields streaming chunks"""
+        async_gen = self._send_request_async_streaming(service, flow, request)
+
+        try:
+            while True:
+                try:
+                    chunk = loop.run_until_complete(async_gen.__anext__())
+                    yield chunk
+                except StopAsyncIteration:
+                    break
+        finally:
+            # Clean up async generator
+            try:
+                loop.run_until_complete(async_gen.aclose())
+            except:
+                pass
+
+    async def _send_request_async(
+        self,
+        service: str,
+        flow: Optional[str],
+        request: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Async implementation of WebSocket request (non-streaming)"""
+        # Generate unique request ID
+        with self._lock:
+            self._request_counter += 1
+            request_id = f"req-{self._request_counter}"
+
+        # Build WebSocket URL with optional token
+        ws_url = f"{self.url}/api/v1/socket"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        # Build request message
+        message = {
+            "id": request_id,
+            "service": service,
+            "request": request
+        }
+        if flow:
+            message["flow"] = flow
+
+        # Connect and send request
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            await websocket.send(json.dumps(message))
+
+            # Wait for single response
+            raw_message = await websocket.recv()
+            response = json.loads(raw_message)
+
+            if response.get("id") != request_id:
+                raise ProtocolException(f"Response ID mismatch")
+
+            if "error" in response:
+                raise ApplicationException(response["error"])
+
+            if "response" not in response:
+                raise ProtocolException(f"Missing response in message")
+
+            return response["response"]
+
+    async def _send_request_async_streaming(
+        self,
+        service: str,
+        flow: Optional[str],
+        request: Dict[str, Any]
+    ) -> Iterator[StreamingChunk]:
+        """Async implementation of WebSocket request (streaming)"""
+        # Generate unique request ID
+        with self._lock:
+            self._request_counter += 1
+            request_id = f"req-{self._request_counter}"
+
+        # Build WebSocket URL with optional token
+        ws_url = f"{self.url}/api/v1/socket"
+        if self.token:
+            ws_url = f"{ws_url}?token={self.token}"
+
+        # Build request message
+        message = {
+            "id": request_id,
+            "service": service,
+            "request": request
+        }
+        if flow:
+            message["flow"] = flow
+
+        # Connect and send request
+        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
+            await websocket.send(json.dumps(message))
+
+            # Yield chunks as they arrive
+            async for raw_message in websocket:
+                response = json.loads(raw_message)
+
+                if response.get("id") != request_id:
+                    continue  # Ignore messages for other requests
+
+                if "error" in response:
+                    raise ApplicationException(response["error"])
+
+                if "response" in response:
+                    resp = response["response"]
+
+                    # Parse different chunk types
+                    chunk = self._parse_chunk(resp)
+                    yield chunk
+
+                    # Check if this is the final chunk
+                    if resp.get("end_of_stream") or resp.get("end_of_dialog") or response.get("complete"):
+                        break
+
+    def _parse_chunk(self, resp: Dict[str, Any]) -> StreamingChunk:
+        """Parse response chunk into appropriate type"""
+        chunk_type = resp.get("chunk_type")
+
+        if chunk_type == "thought":
+            return AgentThought(
+                content=resp.get("content", ""),
+                end_of_message=resp.get("end_of_message", False)
+            )
+        elif chunk_type == "observation":
+            return AgentObservation(
+                content=resp.get("content", ""),
+                end_of_message=resp.get("end_of_message", False)
+            )
+        elif chunk_type == "final-answer":
+            return AgentAnswer(
+                content=resp.get("content", ""),
+                end_of_message=resp.get("end_of_message", False),
+                end_of_dialog=resp.get("end_of_dialog", False)
+            )
+        else:
+            # RAG-style chunk (or generic chunk)
+            return RAGChunk(
+                content=resp.get("chunk", ""),
+                end_of_stream=resp.get("end_of_stream", False),
+                error=resp.get("error")
+            )
+
+    def close(self) -> None:
+        """Close WebSocket connection"""
+        # Cleanup handled by context manager in async code
+        pass
+
+
+class SocketFlowInstance:
+    """Synchronous WebSocket flow instance with same interface as REST FlowInstance"""
+
+    def __init__(self, client: SocketClient, flow_id: str) -> None:
+        self.client: SocketClient = client
+        self.flow_id: str = flow_id
+
+    def agent(
+        self,
+        question: str,
+        user: str,
+        state: Optional[Dict[str, Any]] = None,
+        group: Optional[str] = None,
+        history: Optional[List[Dict[str, Any]]] = None,
+        streaming: bool = False,
+        **kwargs: Any
+    ) -> Union[Dict[str, Any], Iterator[StreamingChunk]]:
+        """Agent with optional streaming"""
+        request = {
+            "question": question,
+            "user": user,
+            "streaming": streaming
+        }
+        if state is not None:
+            request["state"] = state
+        if group is not None:
+            request["group"] = group
+        if history is not None:
+            request["history"] = history
+        request.update(kwargs)
+
+        return self.client._send_request_sync("agent", self.flow_id, request, streaming)
+
+    def text_completion(self, system: str, prompt: str, streaming: bool = False, **kwargs) -> Union[str, Iterator[str]]:
+        """Text completion with optional streaming"""
+        request = {
+            "system": system,
+            "prompt": prompt,
+            "streaming": streaming
+        }
+        request.update(kwargs)
+
+        result = self.client._send_request_sync("text-completion", self.flow_id, request, streaming)
+
+        if streaming:
+            # For text completion, yield just the content
+            for chunk in result:
+                if hasattr(chunk, 'content'):
+                    yield chunk.content
+        else:
+            return result.get("response", "")
+
+    def graph_rag(
+        self,
+        question: str,
+        user: str,
+        collection: str,
+        max_subgraph_size: int = 1000,
+        max_subgraph_count: int = 5,
+        max_entity_distance: int = 3,
+        streaming: bool = False,
+        **kwargs: Any
+    ) -> Union[str, Iterator[str]]:
+        """Graph RAG with optional streaming"""
+        request = {
+            "question": question,
+            "user": user,
+            "collection": collection,
+            "max-subgraph-size": max_subgraph_size,
+            "max-subgraph-count": max_subgraph_count,
+            "max-entity-distance": max_entity_distance,
+            "streaming": streaming
+        }
+        request.update(kwargs)
+
+        result = self.client._send_request_sync("graph-rag", self.flow_id, request, streaming)
+
+        if streaming:
+            for chunk in result:
+                if hasattr(chunk, 'content'):
+                    yield chunk.content
+        else:
+            return result.get("response", "")
+
+    def document_rag(
+        self,
+        question: str,
+        user: str,
+        collection: str,
+        doc_limit: int = 10,
+        streaming: bool = False,
+        **kwargs: Any
+    ) -> Union[str, Iterator[str]]:
+        """Document RAG with optional streaming"""
+        request = {
+            "question": question,
+            "user": user,
+            "collection": collection,
+            "doc-limit": doc_limit,
+            "streaming": streaming
+        }
+        request.update(kwargs)
+
+        result = self.client._send_request_sync("document-rag", self.flow_id, request, streaming)
+
+        if streaming:
+            for chunk in result:
+                if hasattr(chunk, 'content'):
+                    yield chunk.content
+        else:
+            return result.get("response", "")
+
+    def prompt(
+        self,
+        id: str,
+        variables: Dict[str, str],
+        streaming: bool = False,
+        **kwargs: Any
+    ) -> Union[str, Iterator[str]]:
+        """Execute prompt with optional streaming"""
+        request = {
+            "id": id,
+            "variables": variables,
+            "streaming": streaming
+        }
+        request.update(kwargs)
+
+        result = self.client._send_request_sync("prompt", self.flow_id, request, streaming)
+
+        if streaming:
+            for chunk in result:
+                if hasattr(chunk, 'content'):
+                    yield chunk.content
+        else:
+            return result.get("response", "")
+
+    def graph_embeddings_query(
+        self,
+        text: str,
+        user: str,
+        collection: str,
+        limit: int = 10,
+        **kwargs: Any
+    ) -> Dict[str, Any]:
+        """Query graph embeddings for semantic search"""
+        request = {
+            "text": text,
+            "user": user,
+            "collection": collection,
+            "limit": limit
+        }
+        request.update(kwargs)
+
+        return self.client._send_request_sync("graph-embeddings", self.flow_id, request, False)
+
+    def embeddings(self, text: str, **kwargs: Any) -> Dict[str, Any]:
+        """Generate text embeddings"""
+        request = {"text": text}
+        request.update(kwargs)
+
+        return self.client._send_request_sync("embeddings", self.flow_id, request, False)
+
+    def triples_query(
+        self,
+        s: Optional[str] = None,
+        p: Optional[str] = None,
+        o: Optional[str] = None,
+        user: Optional[str] = None,
+        collection: Optional[str] = None,
+        limit: int = 100,
+        **kwargs: Any
+    ) -> Dict[str, Any]:
+        """Triple pattern query"""
+        request = {"limit": limit}
+        if s is not None:
+            request["s"] = str(s)
+        if p is not None:
+            request["p"] = str(p)
+        if o is not None:
+            request["o"] = str(o)
+        if user is not None:
+            request["user"] = user
+        if collection is not None:
+            request["collection"] = collection
+        request.update(kwargs)
+
+        return self.client._send_request_sync("triples", self.flow_id, request, False)
+
+    def objects_query(
+        self,
+        query: str,
+        user: str,
+        collection: str,
+        variables: Optional[Dict[str, Any]] = None,
+        operation_name: Optional[str] = None,
+        **kwargs: Any
+    ) -> Dict[str, Any]:
+        """GraphQL query"""
+        request = {
+            "query": query,
+            "user": user,
+            "collection": collection
+        }
+        if variables:
+            request["variables"] = variables
+        if operation_name:
+            request["operationName"] = operation_name
+        request.update(kwargs)
+
+        return self.client._send_request_sync("objects", self.flow_id, request, False)
+
+    def mcp_tool(
+        self,
+        name: str,
+        parameters: Dict[str, Any],
+        **kwargs: Any
+    ) -> Dict[str, Any]:
+        """Execute MCP tool"""
+        request = {
+            "name": name,
+            "parameters": parameters
+        }
+        request.update(kwargs)
+
+        return self.client._send_request_sync("mcp-tool", self.flow_id, request, False)

trustgraph-base/trustgraph/api/types.py

@@ -1,7 +1,7 @@
 
 import dataclasses
 import datetime
-from typing import List
+from typing import List, Optional, Dict, Any
 from .. knowledge import hash, Uri, Literal
 
 @dataclasses.dataclass
@@ -51,3 +51,33 @@ class CollectionMetadata:
     tags : List[str]
     created_at : str
     updated_at : str
+
+# Streaming chunk types
+
+@dataclasses.dataclass
+class StreamingChunk:
+    """Base class for streaming chunks"""
+    content: str
+    end_of_message: bool = False
+
+@dataclasses.dataclass
+class AgentThought(StreamingChunk):
+    """Agent reasoning chunk"""
+    chunk_type: str = "thought"
+
+@dataclasses.dataclass
+class AgentObservation(StreamingChunk):
+    """Agent tool observation chunk"""
+    chunk_type: str = "observation"
+
+@dataclasses.dataclass
+class AgentAnswer(StreamingChunk):
+    """Agent final answer chunk"""
+    chunk_type: str = "final-answer"
+    end_of_dialog: bool = False
+
+@dataclasses.dataclass
+class RAGChunk(StreamingChunk):
+    """RAG streaming chunk"""
+    end_of_stream: bool = False
+    error: Optional[Dict[str, str]] = None

trustgraph-cli/trustgraph/cli/delete_config_item.py

@@ -8,10 +8,11 @@ from trustgraph.api import Api
 from trustgraph.api.types import ConfigKey
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def delete_config_item(url, config_type, key):
+def delete_config_item(url, config_type, key, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     config_key = ConfigKey(type=config_type, key=key)
     api.delete([config_key])
@@ -43,6 +44,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -51,6 +58,7 @@ def main():
             url=args.api_url,
             config_type=args.type,
             key=args.key,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/get_config_item.py

@@ -9,10 +9,11 @@ from trustgraph.api import Api
 from trustgraph.api.types import ConfigKey
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def get_config_item(url, config_type, key, format_type):
+def get_config_item(url, config_type, key, format_type, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     config_key = ConfigKey(type=config_type, key=key)
     values = api.get([config_key])
@@ -59,6 +60,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -68,6 +75,7 @@ def main():
             config_type=args.type,
             key=args.key,
             format_type=args.format,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/get_kg_core.py

@@ -14,6 +14,7 @@ import msgpack
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
 default_user = 'trustgraph'
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 def write_triple(f, data):
     msg = (
@@ -51,13 +52,16 @@ def write_ge(f, data):
     )
     f.write(msgpack.packb(msg, use_bin_type=True))
 
-async def fetch(url, user, id, output):
+async def fetch(url, user, id, output, token=None):
 
     if not url.endswith("/"):
         url += "/"
 
     url = url + "api/v1/socket"
 
+    if token:
+        url = f"{url}?token={token}"
+
     mid = str(uuid.uuid4())
 
     async with connect(url) as ws:
@@ -138,6 +142,12 @@ def main():
         help=f'Output file'
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -148,6 +158,7 @@ def main():
                 user = args.user,
                 id = args.id,
                 output = args.output,
+                token = args.token,
             )
         )
 

trustgraph-cli/trustgraph/cli/invoke_agent.py

@@ -5,12 +5,10 @@ Uses the agent service to answer a question
 import argparse
 import os
 import textwrap
-import uuid
-import asyncio
-import json
-from websockets.asyncio.client import connect
+from trustgraph.api import Api
 
-default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 
@@ -99,79 +97,47 @@ def output(text, prefix="> ", width=78):
     )
     print(out)
 
-async def question(
+def question(
         url, question, flow_id, user, collection,
-        plan=None, state=None, group=None, verbose=False, streaming=True
+        plan=None, state=None, group=None, verbose=False, streaming=True,
+        token=None
 ):
 
-    if not url.endswith("/"):
-        url += "/"
-
-    url = url + "api/v1/socket"
-
     if verbose:
         output(wrap(question), "\U00002753 ")
         print()
 
-    # Track last chunk type and current outputter for streaming
-    last_chunk_type = None
-    current_outputter = None
+    # Create API client
+    api = Api(url=url, token=token)
+    socket = api.socket()
+    flow = socket.flow(flow_id)
 
-    def think(x):
-        if verbose:
-            output(wrap(x), "\U0001f914 ")
-            print()
+    # Prepare request parameters
+    request_params = {
+        "question": question,
+        "user": user,
+        "streaming": streaming,
+    }
 
-    def observe(x):
-        if verbose:
-            output(wrap(x), "\U0001f4a1 ")
-            print()
+    # Only add optional fields if they have values
+    if state is not None:
+        request_params["state"] = state
+    if group is not None:
+        request_params["group"] = group
 
-    mid = str(uuid.uuid4())
+    try:
+        # Call agent
+        response = flow.agent(**request_params)
 
-    async with connect(url) as ws:
+        # Handle streaming response
+        if streaming:
+            # Track last chunk type and current outputter for streaming
+            last_chunk_type = None
+            current_outputter = None
 
-        req = {
-            "id": mid,
-            "service": "agent",
-            "flow": flow_id,
-            "request": {
-                "question": question,
-                "user": user,
-                "history": [],
-                "streaming": streaming
-            }
-        }
-
-        # Only add optional fields if they have values
-        if state is not None:
-            req["request"]["state"] = state
-        if group is not None:
-            req["request"]["group"] = group
-
-        req = json.dumps(req)
-
-        await ws.send(req)
-
-        while True:
-
-            msg = await ws.recv()
-
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                print("Ignore message")
-                continue
-
-            response = obj["response"]
-
-            # Handle streaming format (new format with chunk_type)
-            if "chunk_type" in response:
-                chunk_type = response["chunk_type"]
-                content = response.get("content", "")
+            for chunk in response:
+                chunk_type = chunk.chunk_type
+                content = chunk.content
 
                 # Check if we're switching to a new message type
                 if last_chunk_type != chunk_type:
@@ -195,33 +161,27 @@ async def question(
                 # Output the chunk
                 if current_outputter:
                     current_outputter.output(content)
-                elif chunk_type == "answer":
+                elif chunk_type == "final-answer":
                     print(content, end="", flush=True)
-            else:
-                # Handle legacy format (backward compatibility)
-                if "thought" in response:
-                    think(response["thought"])
 
-                if "observation" in response:
-                    observe(response["observation"])
+            # Close any remaining outputter
+            if current_outputter:
+                current_outputter.__exit__(None, None, None)
+                current_outputter = None
+            # Add final newline if we were outputting answer
+            elif last_chunk_type == "final-answer":
+                print()
 
-                if "answer" in response:
-                    print(response["answer"])
+        else:
+            # Non-streaming response
+            if "answer" in response:
+                print(response["answer"])
+            if "error" in response:
+                raise RuntimeError(response["error"])
 
-                if "error" in response:
-                    raise RuntimeError(response["error"])
-
-            if obj["complete"]:
-                # Close any remaining outputter
-                if current_outputter:
-                    current_outputter.__exit__(None, None, None)
-                    current_outputter = None
-                # Add final newline if we were outputting answer
-                elif last_chunk_type == "answer":
-                    print()
-                break
-
-        await ws.close()
+    finally:
+        # Clean up socket connection
+        socket.close()
 
 def main():
 
@@ -236,6 +196,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-f', '--flow-id',
         default="default",
@@ -292,19 +258,18 @@ def main():
 
     try:
 
-        asyncio.run(
-            question(
-                url = args.url,
-                flow_id = args.flow_id,
-                question = args.question,
-                user = args.user,
-                collection = args.collection,
-                plan = args.plan,
-                state = args.state,
-                group = args.group,
-                verbose = args.verbose,
-                streaming = not args.no_streaming,
-            )
+        question(
+            url = args.url,
+            flow_id = args.flow_id,
+            question = args.question,
+            user = args.user,
+            collection = args.collection,
+            plan = args.plan,
+            state = args.state,
+            group = args.group,
+            verbose = args.verbose,
+            streaming = not args.no_streaming,
+            token = args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/invoke_document_rag.py

@@ -4,89 +4,50 @@ Uses the DocumentRAG service to answer a question
 
 import argparse
 import os
-import asyncio
-import json
-import uuid
-from websockets.asyncio.client import connect
 from trustgraph.api import Api
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 default_doc_limit = 10
 
-async def question_streaming(url, flow_id, question, user, collection, doc_limit):
-    """Streaming version using websockets"""
+def question(url, flow_id, question, user, collection, doc_limit, streaming=True, token=None):
 
-    # Convert http:// to ws://
-    if url.startswith('http://'):
-        url = 'ws://' + url[7:]
-    elif url.startswith('https://'):
-        url = 'wss://' + url[8:]
+    # Create API client
+    api = Api(url=url, token=token)
 
-    if not url.endswith("/"):
-        url += "/"
+    if streaming:
+        # Use socket client for streaming
+        socket = api.socket()
+        flow = socket.flow(flow_id)
 
-    url = url + "api/v1/socket"
+        try:
+            response = flow.document_rag(
+                question=question,
+                user=user,
+                collection=collection,
+                doc_limit=doc_limit,
+                streaming=True
+            )
 
-    mid = str(uuid.uuid4())
+            # Stream output
+            for chunk in response:
+                print(chunk.content, end="", flush=True)
+            print()  # Final newline
 
-    async with connect(url) as ws:
-        req = {
-            "id": mid,
-            "service": "document-rag",
-            "flow": flow_id,
-            "request": {
-                "query": question,
-                "user": user,
-                "collection": collection,
-                "doc-limit": doc_limit,
-                "streaming": True
-            }
-        }
-
-        req = json.dumps(req)
-        await ws.send(req)
-
-        while True:
-            msg = await ws.recv()
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                print("Ignore message")
-                continue
-
-            response = obj["response"]
-
-            # Handle streaming format (chunk)
-            if "chunk" in response:
-                chunk = response["chunk"]
-                print(chunk, end="", flush=True)
-            elif "response" in response:
-                # Final response with complete text
-                # Already printed via chunks, just add newline
-                pass
-
-            if obj["complete"]:
-                print()  # Final newline
-                break
-
-        await ws.close()
-
-def question_non_streaming(url, flow_id, question, user, collection, doc_limit):
-    """Non-streaming version using HTTP API"""
-
-    api = Api(url).flow().id(flow_id)
-
-    resp = api.document_rag(
-        question=question, user=user, collection=collection,
-        doc_limit=doc_limit,
-    )
-
-    print(resp)
+        finally:
+            socket.close()
+    else:
+        # Use REST API for non-streaming
+        flow = api.flow().id(flow_id)
+        resp = flow.document_rag(
+            question=question,
+            user=user,
+            collection=collection,
+            doc_limit=doc_limit,
+        )
+        print(resp)
 
 def main():
 
@@ -101,6 +62,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-f', '--flow-id',
         default="default",
@@ -127,6 +94,7 @@ def main():
 
     parser.add_argument(
         '-d', '--doc-limit',
+        type=int,
         default=default_doc_limit,
         help=f'Document limit (default: {default_doc_limit})'
     )
@@ -141,30 +109,20 @@ def main():
 
     try:
 
-        if not args.no_streaming:
-            asyncio.run(
-                question_streaming(
-                    url=args.url,
-                    flow_id=args.flow_id,
-                    question=args.question,
-                    user=args.user,
-                    collection=args.collection,
-                    doc_limit=args.doc_limit,
-                )
-            )
-        else:
-            question_non_streaming(
-                url=args.url,
-                flow_id=args.flow_id,
-                question=args.question,
-                user=args.user,
-                collection=args.collection,
-                doc_limit=args.doc_limit,
-            )
+        question(
+            url=args.url,
+            flow_id=args.flow_id,
+            question=args.question,
+            user=args.user,
+            collection=args.collection,
+            doc_limit=args.doc_limit,
+            streaming=not args.no_streaming,
+            token=args.token,
+        )
 
     except Exception as e:
 
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/invoke_graph_rag.py

@@ -4,13 +4,10 @@ Uses the GraphRAG service to answer a question
 
 import argparse
 import os
-import asyncio
-import json
-import uuid
-from websockets.asyncio.client import connect
 from trustgraph.api import Api
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 default_entity_limit = 50
@@ -18,89 +15,51 @@ default_triple_limit = 30
 default_max_subgraph_size = 150
 default_max_path_length = 2
 
-async def question_streaming(
+def question(
         url, flow_id, question, user, collection, entity_limit, triple_limit,
-        max_subgraph_size, max_path_length
+        max_subgraph_size, max_path_length, streaming=True, token=None
 ):
-    """Streaming version using websockets"""
 
-    # Convert http:// to ws://
-    if url.startswith('http://'):
-        url = 'ws://' + url[7:]
-    elif url.startswith('https://'):
-        url = 'wss://' + url[8:]
+    # Create API client
+    api = Api(url=url, token=token)
 
-    if not url.endswith("/"):
-        url += "/"
+    if streaming:
+        # Use socket client for streaming
+        socket = api.socket()
+        flow = socket.flow(flow_id)
 
-    url = url + "api/v1/socket"
+        try:
+            response = flow.graph_rag(
+                question=question,
+                user=user,
+                collection=collection,
+                entity_limit=entity_limit,
+                triple_limit=triple_limit,
+                max_subgraph_size=max_subgraph_size,
+                max_path_length=max_path_length,
+                streaming=True
+            )
 
-    mid = str(uuid.uuid4())
+            # Stream output
+            for chunk in response:
+                print(chunk.content, end="", flush=True)
+            print()  # Final newline
 
-    async with connect(url) as ws:
-        req = {
-            "id": mid,
-            "service": "graph-rag",
-            "flow": flow_id,
-            "request": {
-                "query": question,
-                "user": user,
-                "collection": collection,
-                "entity-limit": entity_limit,
-                "triple-limit": triple_limit,
-                "max-subgraph-size": max_subgraph_size,
-                "max-path-length": max_path_length,
-                "streaming": True
-            }
-        }
-
-        req = json.dumps(req)
-        await ws.send(req)
-
-        while True:
-            msg = await ws.recv()
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                print("Ignore message")
-                continue
-
-            response = obj["response"]
-
-            # Handle streaming format (chunk)
-            if "chunk" in response:
-                chunk = response["chunk"]
-                print(chunk, end="", flush=True)
-            elif "response" in response:
-                # Final response with complete text
-                # Already printed via chunks, just add newline
-                pass
-
-            if obj["complete"]:
-                print()  # Final newline
-                break
-
-        await ws.close()
-
-def question_non_streaming(
-        url, flow_id, question, user, collection, entity_limit, triple_limit,
-        max_subgraph_size, max_path_length
-):
-    """Non-streaming version using HTTP API"""
-
-    api = Api(url).flow().id(flow_id)
-
-    resp = api.graph_rag(
-        question=question, user=user, collection=collection,
-        entity_limit=entity_limit, triple_limit=triple_limit,
-        max_subgraph_size=max_subgraph_size,
-        max_path_length=max_path_length
-    )
-
-    print(resp)
+        finally:
+            socket.close()
+    else:
+        # Use REST API for non-streaming
+        flow = api.flow().id(flow_id)
+        resp = flow.graph_rag(
+            question=question,
+            user=user,
+            collection=collection,
+            entity_limit=entity_limit,
+            triple_limit=triple_limit,
+            max_subgraph_size=max_subgraph_size,
+            max_path_length=max_path_length
+        )
+        print(resp)
 
 def main():
 
@@ -115,6 +74,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-f', '--flow-id',
         default="default",
@@ -141,24 +106,28 @@ def main():
 
     parser.add_argument(
         '-e', '--entity-limit',
+        type=int,
         default=default_entity_limit,
         help=f'Entity limit (default: {default_entity_limit})'
     )
 
     parser.add_argument(
-        '-t', '--triple-limit',
+        '--triple-limit',
+        type=int,
         default=default_triple_limit,
         help=f'Triple limit (default: {default_triple_limit})'
     )
 
     parser.add_argument(
         '-s', '--max-subgraph-size',
+        type=int,
         default=default_max_subgraph_size,
         help=f'Max subgraph size (default: {default_max_subgraph_size})'
     )
 
     parser.add_argument(
         '-p', '--max-path-length',
+        type=int,
         default=default_max_path_length,
         help=f'Max path length (default: {default_max_path_length})'
     )
@@ -173,36 +142,23 @@ def main():
 
     try:
 
-        if not args.no_streaming:
-            asyncio.run(
-                question_streaming(
-                    url=args.url,
-                    flow_id=args.flow_id,
-                    question=args.question,
-                    user=args.user,
-                    collection=args.collection,
-                    entity_limit=args.entity_limit,
-                    triple_limit=args.triple_limit,
-                    max_subgraph_size=args.max_subgraph_size,
-                    max_path_length=args.max_path_length,
-                )
-            )
-        else:
-            question_non_streaming(
-                url=args.url,
-                flow_id=args.flow_id,
-                question=args.question,
-                user=args.user,
-                collection=args.collection,
-                entity_limit=args.entity_limit,
-                triple_limit=args.triple_limit,
-                max_subgraph_size=args.max_subgraph_size,
-                max_path_length=args.max_path_length,
-            )
+        question(
+            url=args.url,
+            flow_id=args.flow_id,
+            question=args.question,
+            user=args.user,
+            collection=args.collection,
+            entity_limit=args.entity_limit,
+            triple_limit=args.triple_limit,
+            max_subgraph_size=args.max_subgraph_size,
+            max_path_length=args.max_path_length,
+            streaming=not args.no_streaming,
+            token=args.token,
+        )
 
     except Exception as e:
 
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/invoke_llm.py

@@ -5,64 +5,39 @@ and user prompt.  Both arguments are required.
 
 import argparse
 import os
-import json
-import uuid
-import asyncio
-from websockets.asyncio.client import connect
+from trustgraph.api import Api
 
-default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-async def query(url, flow_id, system, prompt, streaming=True):
+def query(url, flow_id, system, prompt, streaming=True, token=None):
 
-    if not url.endswith("/"):
-        url += "/"
+    # Create API client
+    api = Api(url=url, token=token)
+    socket = api.socket()
+    flow = socket.flow(flow_id)
 
-    url = url + "api/v1/socket"
+    try:
+        # Call text completion
+        response = flow.text_completion(
+            system=system,
+            prompt=prompt,
+            streaming=streaming
+        )
 
-    mid = str(uuid.uuid4())
+        if streaming:
+            # Stream output to stdout without newline
+            for chunk in response:
+                print(chunk.content, end="", flush=True)
+            # Add final newline after streaming
+            print()
+        else:
+            # Non-streaming: print complete response
+            print(response)
 
-    async with connect(url) as ws:
-
-        req = {
-            "id": mid,
-            "service": "text-completion",
-            "flow": flow_id,
-            "request": {
-                "system": system,
-                "prompt": prompt,
-                "streaming": streaming
-            }
-        }
-
-        await ws.send(json.dumps(req))
-
-        while True:
-
-            msg = await ws.recv()
-
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                continue
-
-            if "response" in obj["response"]:
-                if streaming:
-                    # Stream output to stdout without newline
-                    print(obj["response"]["response"], end="", flush=True)
-                else:
-                    # Non-streaming: print complete response
-                    print(obj["response"]["response"])
-
-            if obj["complete"]:
-                if streaming:
-                    # Add final newline after streaming
-                    print()
-                break
-
-        await ws.close()
+    finally:
+        # Clean up socket connection
+        socket.close()
 
 def main():
 
@@ -77,6 +52,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         'system',
         nargs=1,
@@ -105,17 +86,18 @@ def main():
 
     try:
 
-        asyncio.run(query(
+        query(
             url=args.url,
             flow_id=args.flow_id,
             system=args.system[0],
             prompt=args.prompt[0],
-            streaming=not args.no_streaming
-        ))
+            streaming=not args.no_streaming,
+            token=args.token,
+        )
 
     except Exception as e:
 
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/invoke_prompt.py

@@ -10,76 +10,61 @@ using key=value arguments on the command line, and these replace
 import argparse
 import os
 import json
-import uuid
-import asyncio
-from websockets.asyncio.client import connect
+from trustgraph.api import Api
 
-default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-async def query(url, flow_id, template_id, variables, streaming=True):
+def query(url, flow_id, template_id, variables, streaming=True, token=None):
 
-    if not url.endswith("/"):
-        url += "/"
+    # Create API client
+    api = Api(url=url, token=token)
+    socket = api.socket()
+    flow = socket.flow(flow_id)
 
-    url = url + "api/v1/socket"
+    try:
+        # Call prompt
+        response = flow.prompt(
+            id=template_id,
+            variables=variables,
+            streaming=streaming
+        )
 
-    mid = str(uuid.uuid4())
+        if streaming:
+            full_response = {"text": "", "object": ""}
 
-    async with connect(url) as ws:
+            # Stream output
+            for chunk in response:
+                content = chunk.content
+                if content:
+                    print(content, end="", flush=True)
+                    full_response["text"] += content
 
-        req = {
-            "id": mid,
-            "service": "prompt",
-            "flow": flow_id,
-            "request": {
-                "id": template_id,
-                "variables": variables,
-                "streaming": streaming
-            }
-        }
+                # Check if this is an object response (JSON)
+                if hasattr(chunk, 'object') and chunk.object:
+                    full_response["object"] = chunk.object
 
-        await ws.send(json.dumps(req))
+            # Handle final output
+            if full_response["text"]:
+                # Add final newline after streaming text
+                print()
+            elif full_response["object"]:
+                # Print JSON object (pretty-printed)
+                print(json.dumps(json.loads(full_response["object"]), indent=4))
 
-        full_response = {"text": "", "object": ""}
-
-        while True:
-
-            msg = await ws.recv()
-
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                continue
-
-            response = obj["response"]
-
-            # Handle text responses (streaming)
-            if "text" in response and response["text"]:
-                if streaming:
-                    # Stream output to stdout without newline
-                    print(response["text"], end="", flush=True)
-                    full_response["text"] += response["text"]
-                else:
-                    # Non-streaming: print complete response
+        else:
+            # Non-streaming: handle response
+            if isinstance(response, str):
+                print(response)
+            elif isinstance(response, dict):
+                if "text" in response:
                     print(response["text"])
+                elif "object" in response:
+                    print(json.dumps(json.loads(response["object"]), indent=4))
 
-            # Handle object responses (JSON, never streamed)
-            if "object" in response and response["object"]:
-                full_response["object"] = response["object"]
-
-            if obj["complete"]:
-                if streaming and full_response["text"]:
-                    # Add final newline after streaming text
-                    print()
-                elif full_response["object"]:
-                    # Print JSON object (pretty-printed)
-                    print(json.dumps(json.loads(full_response["object"]), indent=4))
-                break
-
-        await ws.close()
+    finally:
+        # Clean up socket connection
+        socket.close()
 
 def main():
 
@@ -94,6 +79,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-f', '--flow-id',
         default="default",
@@ -135,17 +126,18 @@ specified multiple times''',
 
     try:
 
-        asyncio.run(query(
+        query(
             url=args.url,
             flow_id=args.flow_id,
             template_id=args.id[0],
             variables=variables,
-            streaming=not args.no_streaming
-        ))
+            streaming=not args.no_streaming,
+            token=args.token,
+        )
 
     except Exception as e:
 
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/list_config_items.py

@@ -8,10 +8,11 @@ import json
 from trustgraph.api import Api
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def list_config_items(url, config_type, format_type):
+def list_config_items(url, config_type, format_type, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     keys = api.list(config_type)
 
@@ -47,6 +48,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -55,6 +62,7 @@ def main():
             url=args.api_url,
             config_type=args.type,
             format_type=args.format,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/load_knowledge.py

@@ -2,18 +2,17 @@
 Loads triples and entity contexts into the knowledge graph.
 """
 
-import asyncio
 import argparse
 import os
 import time
 import rdflib
-import json
-from websockets.asyncio.client import connect
-from typing import List, Dict, Any
+from typing import Iterator, Tuple
 
+from trustgraph.api import Api, Triple
 from trustgraph.log_level import LogLevel
 
-default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 
@@ -26,108 +25,114 @@ class KnowledgeLoader:
             user,
             collection,
             document_id,
-            url = default_url,
+            url=default_url,
+            token=None,
     ):
-
-        if not url.endswith("/"):
-            url += "/"
-
-        self.triples_url = url + f"api/v1/flow/{flow}/import/triples"
-        self.entity_contexts_url = url + f"api/v1/flow/{flow}/import/entity-contexts"
-
         self.files = files
+        self.flow = flow
         self.user = user
         self.collection = collection
         self.document_id = document_id
+        self.url = url
+        self.token = token
 
-    async def run(self):
-
-        try:
-            # Load triples first
-            async with connect(self.triples_url) as ws:
-                for file in self.files:
-                    await self.load_triples(file, ws)
-
-            # Then load entity contexts
-            async with connect(self.entity_contexts_url) as ws:
-                for file in self.files:
-                    await self.load_entity_contexts(file, ws)
-
-        except Exception as e:
-            print(e, flush=True)
-
-    async def load_triples(self, file, ws):
+    def load_triples_from_file(self, file) -> Iterator[Triple]:
+        """Generator that yields Triple objects from a Turtle file"""
 
         g = rdflib.Graph()
         g.parse(file, format="turtle")
 
-        def Value(value, is_uri):
-            return { "v": value, "e": is_uri }
-
         for e in g:
-            s = Value(value=str(e[0]), is_uri=True)
-            p = Value(value=str(e[1]), is_uri=True)
-            if type(e[2]) == rdflib.term.URIRef:
-                o = Value(value=str(e[2]), is_uri=True)
+            # Extract subject, predicate, object
+            s_value = str(e[0])
+            p_value = str(e[1])
+
+            # Check if object is a URI or literal
+            if isinstance(e[2], rdflib.term.URIRef):
+                o_value = str(e[2])
+                o_is_uri = True
             else:
-                o = Value(value=str(e[2]), is_uri=False)
+                o_value = str(e[2])
+                o_is_uri = False
 
-            req = {
-                "metadata": {
-                    "id": self.document_id,
-                    "metadata": [],
-                    "user": self.user,
-                    "collection": self.collection
-                },
-                "triples": [
-                    {
-                        "s": s,
-                        "p": p,
-                        "o": o,
-                    }
-                ]
-            }
+            # Create Triple object
+            # Note: The Triple dataclass has 's', 'p', 'o' fields as strings
+            # The API will handle the metadata wrapping
+            yield Triple(s=s_value, p=p_value, o=o_value)
 
-            await ws.send(json.dumps(req))
-
-    async def load_entity_contexts(self, file, ws):
-        """
-        Load entity contexts by extracting entities from the RDF graph
-        and generating contextual descriptions based on their relationships.
-        """
+    def load_entity_contexts_from_file(self, file) -> Iterator[Tuple[str, str]]:
+        """Generator that yields (entity, context) tuples from a Turtle file"""
 
         g = rdflib.Graph()
         g.parse(file, format="turtle")
 
         for s, p, o in g:
-            # If object is a URI, do nothing
+            # If object is a URI, skip (we only want literal contexts)
             if isinstance(o, rdflib.term.URIRef):
                 continue
-                
-            # If object is a literal, create entity context for subject with literal as context
+
+            # If object is a literal, create entity context for subject
             s_str = str(s)
             o_str = str(o)
-            
-            req = {
-                "metadata": {
-                    "id": self.document_id,
-                    "metadata": [],
-                    "user": self.user,
-                    "collection": self.collection
-                },
-                "entities": [
-                    {
-                        "entity": {
-                            "v": s_str,
-                            "e": True
-                        },
-                        "context": o_str
+
+            yield (s_str, o_str)
+
+    def run(self):
+        """Load triples and entity contexts using Python API"""
+
+        try:
+            # Create API client
+            api = Api(url=self.url, token=self.token)
+            bulk = api.bulk()
+
+            # Load triples from all files
+            print("Loading triples...")
+            for file in self.files:
+                print(f"  Processing {file}...")
+                triples = self.load_triples_from_file(file)
+
+                bulk.import_triples(
+                    flow=self.flow,
+                    triples=triples,
+                    metadata={
+                        "id": self.document_id,
+                        "metadata": [],
+                        "user": self.user,
+                        "collection": self.collection
                     }
-                ]
-            }
+                )
 
-            await ws.send(json.dumps(req))
+            print("Triples loaded.")
 
+            # Load entity contexts from all files
+            print("Loading entity contexts...")
+            for file in self.files:
+                print(f"  Processing {file}...")
+
+                # Convert tuples to the format expected by import_entity_contexts
+                def entity_context_generator():
+                    for entity, context in self.load_entity_contexts_from_file(file):
+                        yield {
+                            "entity": {"v": entity, "e": True},
+                            "context": context
+                        }
+
+                bulk.import_entity_contexts(
+                    flow=self.flow,
+                    entities=entity_context_generator(),
+                    metadata={
+                        "id": self.document_id,
+                        "metadata": [],
+                        "user": self.user,
+                        "collection": self.collection
+                    }
+                )
+
+            print("Entity contexts loaded.")
+
+        except Exception as e:
+            print(f"Error: {e}", flush=True)
+            raise
 
 def main():
 
@@ -142,6 +147,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-i', '--document-id',
         required=True,
@@ -166,7 +177,6 @@ def main():
         help=f'Collection ID (default: {default_collection})'
     )
 
-
     parser.add_argument(
         'files', nargs='+',
         help=f'Turtle files to load'
@@ -178,15 +188,16 @@ def main():
 
         try:
             loader = KnowledgeLoader(
-                document_id = args.document_id,
-                url = args.api_url,
-                flow = args.flow_id,
-                files = args.files,
-                user = args.user,
-                collection = args.collection,
+                document_id=args.document_id,
+                url=args.api_url,
+                token=args.token,
+                flow=args.flow_id,
+                files=args.files,
+                user=args.user,
+                collection=args.collection,
             )
 
-            asyncio.run(loader.run())
+            loader.run()
 
             print("Triples and entity contexts loaded.")
             break
@@ -199,4 +210,4 @@ def main():
         time.sleep(10)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/load_sample_documents.py

@@ -13,6 +13,7 @@ from trustgraph.api.types import hash, Uri, Literal, Triple
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
 default_user = 'trustgraph'
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 
 from requests.adapters import HTTPAdapter
@@ -655,10 +656,10 @@ documents = [
 class Loader:
 
     def __init__(
-            self, url, user
+            self, url, user, token=None
     ):
 
-        self.api = Api(url).library()
+        self.api = Api(url, token=token).library()
         self.user = user
 
     def load(self, documents):
@@ -719,6 +720,12 @@ def main():
         help=f'User ID (default: {default_user})'
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -726,6 +733,7 @@ def main():
         p = Loader(
             url=args.url,
             user=args.user,
+            token=args.token,
         )
 
         p.load(documents)

trustgraph-cli/trustgraph/cli/load_structured_data.py

@@ -22,6 +22,7 @@ import logging
 logger = logging.getLogger(__name__)
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 
 def load_structured_data(
@@ -41,7 +42,8 @@ def load_structured_data(
     user: str = 'trustgraph',
     collection: str = 'default',
     dry_run: bool = False,
-    verbose: bool = False
+    verbose: bool = False,
+    token: str = None
 ):
     """
     Load structured data using a descriptor configuration.
@@ -133,9 +135,9 @@ def load_structured_data(
                 
                 # Get batch size from descriptor
                 batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000)
-                
+
                 # Send to TrustGraph using shared function
-                imported_count = _send_to_trustgraph(output_objects, api_url, flow, batch_size)
+                imported_count = _send_to_trustgraph(output_objects, api_url, flow, batch_size, token=token)
                 
                 # Summary
                 format_info = descriptor.get('format', {})
@@ -288,10 +290,10 @@ def load_structured_data(
         
         # Get batch size from descriptor or use default
         batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000)
-        
+
         # Send to TrustGraph
         print(f"🚀 Importing {len(output_records)} records to TrustGraph...")
-        imported_count = _send_to_trustgraph(output_records, api_url, flow, batch_size)
+        imported_count = _send_to_trustgraph(output_records, api_url, flow, batch_size, token=token)
         
         # Get summary info from descriptor
         format_info = descriptor.get('format', {})
@@ -571,66 +573,30 @@ def _process_data_pipeline(input_file, descriptor_file, user, collection, sample
     return output_records, descriptor
 
 
-def _send_to_trustgraph(objects, api_url, flow, batch_size=1000):
-    """Send ExtractedObject records to TrustGraph using WebSocket"""
-    import json
-    import asyncio
-    from websockets.asyncio.client import connect
-    
+def _send_to_trustgraph(objects, api_url, flow, batch_size=1000, token=None):
+    """Send ExtractedObject records to TrustGraph using Python API"""
+    from trustgraph.api import Api
+
     try:
-        # Construct objects import URL similar to load_knowledge pattern
-        if not api_url.endswith("/"):
-            api_url += "/"
-        
-        # Convert HTTP URL to WebSocket URL if needed
-        ws_url = api_url.replace("http://", "ws://").replace("https://", "wss://")
-        objects_url = ws_url + f"api/v1/flow/{flow}/import/objects"
-        
-        logger.info(f"Connecting to objects import endpoint: {objects_url}")
-        
-        async def import_objects():
-            async with connect(objects_url) as ws:
-                imported_count = 0
-                
-                for record in objects:
-                    try:
-                        # Send individual ExtractedObject records
-                        await ws.send(json.dumps(record))
-                        imported_count += 1
-                        
-                        if imported_count % 100 == 0:
-                            logger.debug(f"Imported {imported_count}/{len(objects)} records...")
-                            
-                    except Exception as e:
-                        logger.error(f"Failed to send record {imported_count + 1}: {e}")
-                        print(f"❌ Failed to send record {imported_count + 1}: {e}")
-                
-                logger.info(f"Successfully imported {imported_count} records to TrustGraph")
-                return imported_count
-        
-        # Run the async import
-        imported_count = asyncio.run(import_objects())
-        
-        # Summary
         total_records = len(objects)
-        failed_count = total_records - imported_count
-        
+        logger.info(f"Importing {total_records} records to TrustGraph...")
+
+        # Use Python API bulk import
+        api = Api(api_url, token=token)
+        bulk = api.bulk()
+
+        bulk.import_objects(flow=flow, objects=iter(objects))
+
+        logger.info(f"Successfully imported {total_records} records to TrustGraph")
+
+        # Summary
         print(f"\n📊 Import Summary:")
         print(f"- Total records: {total_records}")
-        print(f"- Successfully imported: {imported_count}")
-        print(f"- Failed: {failed_count}")
-        
-        if failed_count > 0:
-            print(f"⚠️  {failed_count} records failed to import. Check logs for details.")
-        else:
-            print("✅ All records imported successfully!")
-            
-        return imported_count
-        
-    except ImportError as e:
-        logger.error(f"Failed to import required modules: {e}")
-        print(f"Error: Required modules not available - {e}")
-        raise
+        print(f"- Successfully imported: {total_records}")
+        print("✅ All records imported successfully!")
+
+        return total_records
+
     except Exception as e:
         logger.error(f"Failed to import data to TrustGraph: {e}")
         print(f"Import failed: {e}")
@@ -1024,7 +990,13 @@ For more information on the descriptor format, see:
         '--error-file',
         help='Path to write error records (optional)'
     )
-    
+
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
     
     # Input validation
@@ -1077,7 +1049,8 @@ For more information on the descriptor format, see:
             user=args.user,
             collection=args.collection,
             dry_run=args.dry_run,
-            verbose=args.verbose
+            verbose=args.verbose,
+            token=args.token
         )
     except FileNotFoundError as e:
         print(f"Error: File not found - {e}", file=sys.stderr)

trustgraph-cli/trustgraph/cli/load_turtle.py

@@ -1,18 +1,18 @@
 """
-Loads triples into the knowledge graph.
+Loads triples into the knowledge graph from Turtle files.
 """
 
-import asyncio
 import argparse
 import os
 import time
 import rdflib
-import json
-from websockets.asyncio.client import connect
+from typing import Iterator
 
+from trustgraph.api import Api, Triple
 from trustgraph.log_level import LogLevel
 
-default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 
@@ -25,67 +25,67 @@ class Loader:
             user,
             collection,
             document_id,
-            url = default_url,
+            url=default_url,
+            token=None,
     ):
-
-        if not url.endswith("/"):
-            url += "/"
-
-        url = url + f"api/v1/flow/{flow}/import/triples"
-
-        self.url = url
-
         self.files = files
+        self.flow = flow
         self.user = user
         self.collection = collection
         self.document_id = document_id
+        self.url = url
+        self.token = token
 
-    async def run(self):
-
-        try:
-
-            async with connect(self.url) as ws:
-                for file in self.files:
-                    await self.load_file(file, ws)
-
-        except Exception as e:
-            print(e, flush=True)
-
-    async def load_file(self, file, ws):
+    def load_triples_from_file(self, file) -> Iterator[Triple]:
+        """Generator that yields Triple objects from a Turtle file"""
 
         g = rdflib.Graph()
         g.parse(file, format="turtle")
 
-        def Value(value, is_uri):
-            return { "v": value, "e": is_uri }
-
-        triples = []
-
         for e in g:
-            s = Value(value=str(e[0]), is_uri=True)
-            p = Value(value=str(e[1]), is_uri=True)
-            if type(e[2]) == rdflib.term.URIRef:
-                o = Value(value=str(e[2]), is_uri=True)
+            # Extract subject, predicate, object
+            s_value = str(e[0])
+            p_value = str(e[1])
+
+            # Check if object is a URI or literal
+            if isinstance(e[2], rdflib.term.URIRef):
+                o_value = str(e[2])
             else:
-                o = Value(value=str(e[2]), is_uri=False)
+                o_value = str(e[2])
 
-            req = {
-                "metadata": {
-                    "id": self.document_id,
-                    "metadata": [],
-                    "user": self.user,
-                    "collection": self.collection
-                },
-                "triples": [
-                    {
-                        "s": s,
-                        "p": p,
-                        "o": o,
+            # Create Triple object
+            yield Triple(s=s_value, p=p_value, o=o_value)
+
+    def run(self):
+        """Load triples using Python API"""
+
+        try:
+            # Create API client
+            api = Api(url=self.url, token=self.token)
+            bulk = api.bulk()
+
+            # Load triples from all files
+            print("Loading triples...")
+            for file in self.files:
+                print(f"  Processing {file}...")
+                triples = self.load_triples_from_file(file)
+
+                bulk.import_triples(
+                    flow=self.flow,
+                    triples=triples,
+                    metadata={
+                        "id": self.document_id,
+                        "metadata": [],
+                        "user": self.user,
+                        "collection": self.collection
                     }
-                ]
-            }
+                )
 
-            await ws.send(json.dumps(req))
+            print("Triples loaded.")
+
+        except Exception as e:
+            print(f"Error: {e}", flush=True)
+            raise
 
 def main():
 
@@ -100,6 +100,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-i', '--document-id',
         required=True,
@@ -134,16 +140,17 @@ def main():
     while True:
 
         try:
-            p = Loader(
-                document_id = args.document_id,
-                url = args.api_url,
-                flow = args.flow_id,
-                files = args.files,
-                user = args.user,
-                collection = args.collection,
+            loader = Loader(
+                document_id=args.document_id,
+                url=args.api_url,
+                token=args.token,
+                flow=args.flow_id,
+                files=args.files,
+                user=args.user,
+                collection=args.collection,
             )
 
-            asyncio.run(p.run())
+            loader.run()
 
             print("File loaded.")
             break
@@ -156,4 +163,4 @@ def main():
         time.sleep(10)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/put_config_item.py

@@ -9,10 +9,11 @@ from trustgraph.api import Api
 from trustgraph.api.types import ConfigValue
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def put_config_item(url, config_type, key, value):
+def put_config_item(url, config_type, key, value, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     config_value = ConfigValue(type=config_type, key=key, value=value)
     api.put([config_value])
@@ -56,6 +57,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -70,6 +77,7 @@ def main():
             config_type=args.type,
             key=args.key,
             value=value,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/put_flow_class.py

@@ -9,10 +9,11 @@ from trustgraph.api import Api
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def put_flow_class(url, class_name, config):
+def put_flow_class(url, class_name, config, token=None):
 
-    api = Api(url)
+    api = Api(url, token=token)
 
     class_names = api.flow().put_class(class_name, config)
 
@@ -29,6 +30,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-n', '--class-name',
         help=f'Flow class name',
@@ -47,6 +54,7 @@ def main():
             url=args.api_url,
             class_name=args.class_name,
             config=json.loads(args.config),
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/put_kg_core.py

@@ -13,6 +13,7 @@ import msgpack
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
 default_user = 'trustgraph'
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 def read_message(unpacked, id, user):
     
@@ -47,13 +48,16 @@ def read_message(unpacked, id, user):
     else:
         raise RuntimeError("Unpacked unexpected messsage type", unpacked[0])
 
-async def put(url, user, id, input):
+async def put(url, user, id, input, token=None):
 
     if not url.endswith("/"):
         url += "/"
 
     url = url + "api/v1/socket"
 
+    if token:
+        url = f"{url}?token={token}"
+
     async with connect(url) as ws:
 
 
@@ -160,6 +164,12 @@ def main():
         help=f'Input file'
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -170,6 +180,7 @@ def main():
                 user = args.user,
                 id = args.id,
                 input = args.input,
+                token = args.token,
             )
         )
 

trustgraph-cli/trustgraph/cli/remove_library_document.py

@@ -10,11 +10,12 @@ from trustgraph.api import Api
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
 default_user = 'trustgraph'
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 
-def remove_doc(url, user, id):
+def remove_doc(url, user, id, token=None):
 
-    api = Api(url).library()
+    api = Api(url, token=token).library()
 
     api.remove_document(user=user, id=id)
 
@@ -43,11 +44,17 @@ def main():
         help=f'Document ID'
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
-        remove_doc(args.url, args.user, args.identifier)
+        remove_doc(args.url, args.user, args.identifier, token=args.token)
 
     except Exception as e:
 

trustgraph-cli/trustgraph/cli/set_collection.py

@@ -9,10 +9,11 @@ from trustgraph.api import Api
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
 default_user = "trustgraph"
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def set_collection(url, user, collection, name, description, tags):
+def set_collection(url, user, collection, name, description, tags, token=None):
 
-    api = Api(url).collection()
+    api = Api(url, token=token).collection()
 
     result = api.update_collection(
         user=user,
@@ -82,6 +83,12 @@ def main():
         help='Collection tags (can be specified multiple times)'
     )
 
+    parser.add_argument(
+        '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -92,7 +99,8 @@ def main():
             collection = args.collection,
             name = args.name,
             description = args.description,
-            tags = args.tags
+            tags = args.tags,
+            token = args.token
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/set_mcp_tool.py

@@ -20,6 +20,7 @@ import textwrap
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 def set_mcp_tool(
         url : str,
@@ -27,9 +28,10 @@ def set_mcp_tool(
         remote_name : str,
         tool_url : str,
         auth_token : str = None,
+        token : str = None,
 ):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     # Build the MCP tool configuration
     config = {
@@ -72,6 +74,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-i', '--id',
         required=True,
@@ -116,7 +124,8 @@ def main():
             id=args.id,
             remote_name=remote_name,
             tool_url=args.tool_url,
-            auth_token=args.auth_token
+            auth_token=args.auth_token,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/set_prompt.py

@@ -10,10 +10,11 @@ import tabulate
 import textwrap
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def set_system(url, system):
+def set_system(url, system, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     api.put([
         ConfigValue(type="prompt", key="system", value=json.dumps(system))
@@ -21,9 +22,9 @@ def set_system(url, system):
 
     print("System prompt set.")
 
-def set_prompt(url, id, prompt, response, schema):
+def set_prompt(url, id, prompt, response, schema, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     values = api.get([
         ConfigKey(type="prompt", key="template-index")
@@ -71,6 +72,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '--id',
         help=f'Prompt ID',
@@ -103,9 +110,9 @@ def main():
         if args.system:
             if args.id or args.prompt or args.schema or args.response:
                 raise RuntimeError("Can't use --system with other args")
-                
+
             set_system(
-                url=args.api_url, system=args.system
+                url=args.api_url, system=args.system, token=args.token
             )
 
         else:
@@ -130,7 +137,7 @@ def main():
 
             set_prompt(
                 url=args.api_url, id=args.id, prompt=args.prompt,
-                response=args.response, schema=schobj
+                response=args.response, schema=schobj, token=args.token
             )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/set_token_costs.py

@@ -10,10 +10,11 @@ import tabulate
 import textwrap
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def set_costs(api_url, model, input_costs, output_costs):
+def set_costs(api_url, model, input_costs, output_costs, token=None):
 
-    api = Api(api_url).config()
+    api = Api(api_url, token=token).config()
 
     api.put([
         ConfigValue(
@@ -95,6 +96,12 @@ def main():
         help=f'Input costs in $ per 1M tokens',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:

trustgraph-cli/trustgraph/cli/set_tool.py

@@ -26,6 +26,7 @@ import textwrap
 import dataclasses
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 @dataclasses.dataclass
 class Argument:
@@ -67,9 +68,10 @@ def set_tool(
         group : List[str],
         state : str,
         applicable_states : List[str],
+        token : str = None,
 ):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     values = api.get([
         ConfigKey(type="agent", key="tool-index")
@@ -156,6 +158,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '--id',
         help=f'Unique tool identifier',
@@ -257,6 +265,7 @@ def main():
             group=args.group,
             state=args.state,
             applicable_states=args.applicable_states,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_config.py

@@ -8,10 +8,11 @@ from trustgraph.api import Api
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def show_config(url):
+def show_config(url, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     config, version = api.all()
 
@@ -31,12 +32,19 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
         show_config(
             url=args.api_url,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_flow_classes.py

@@ -9,6 +9,7 @@ from trustgraph.api import Api, ConfigKey
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 def format_parameters(params_metadata, config_api):
     """
@@ -57,9 +58,9 @@ def format_parameters(params_metadata, config_api):
 
     return "\n".join(param_list)
 
-def show_flow_classes(url):
+def show_flow_classes(url, token=None):
 
-    api = Api(url)
+    api = Api(url, token=token)
     flow_api = api.flow()
     config_api = api.config()
 
@@ -106,12 +107,19 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
         show_flow_classes(
             url=args.api_url,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_flow_state.py

@@ -9,10 +9,11 @@ import os
 
 default_metrics_url = "http://localhost:8088/api/metrics"
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def dump_status(metrics_url, api_url, flow_id):
+def dump_status(metrics_url, api_url, flow_id, token=None):
 
-    api = Api(api_url).flow()
+    api = Api(api_url, token=token).flow()
 
     flow = api.get(flow_id)
     class_name = flow["class-name"]
@@ -77,11 +78,17 @@ def main():
         help=f'Metrics URL (default: {default_metrics_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
-        dump_status(args.metrics_url, args.api_url, args.flow_id)
+        dump_status(args.metrics_url, args.api_url, args.flow_id, token=args.token)
 
     except Exception as e:
 

trustgraph-cli/trustgraph/cli/show_flows.py

@@ -9,6 +9,7 @@ from trustgraph.api import Api, ConfigKey
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 def get_interface(config_api, i):
 
@@ -128,9 +129,9 @@ def format_parameters(flow_params, class_params_metadata, config_api):
 
     return "\n".join(param_list) if param_list else "None"
 
-def show_flows(url):
+def show_flows(url, token=None):
 
-    api = Api(url)
+    api = Api(url, token=token)
     config_api = api.config()
     flow_api = api.flow()
 
@@ -199,12 +200,19 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
         show_flows(
             url=args.api_url,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_graph.py

@@ -9,10 +9,11 @@ from trustgraph.api import Api
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
 default_user = 'trustgraph'
 default_collection = 'default'
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def show_graph(url, flow_id, user, collection):
+def show_graph(url, flow_id, user, collection, token=None):
 
-    api = Api(url).flow().id(flow_id)
+    api = Api(url, token=token).flow().id(flow_id)
 
     rows = api.triples_query(
         user=user, collection=collection,
@@ -53,6 +54,12 @@ def main():
         help=f'Collection ID (default: {default_collection})'
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
@@ -62,6 +69,7 @@ def main():
             flow_id = args.flow_id,
             user = args.user,
             collection = args.collection,
+            token = args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_kg_cores.py

@@ -9,10 +9,11 @@ from trustgraph.api import Api, ConfigKey
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def show_cores(url, user):
+def show_cores(url, user, token=None):
 
-    api = Api(url).knowledge()
+    api = Api(url, token=token).knowledge()
 
     ids = api.list_kg_cores()
 
@@ -35,6 +36,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-U', '--user',
         default="trustgraph",
@@ -46,7 +53,9 @@ def main():
     try:
 
         show_cores(
-            url=args.api_url, user=args.user
+            url=args.api_url,
+            user=args.user,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_library_documents.py

@@ -9,11 +9,12 @@ from trustgraph.api import Api, ConfigKey
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = "trustgraph"
 
-def show_docs(url, user):
+def show_docs(url, user, token=None):
 
-    api = Api(url).library()
+    api = Api(url, token=token).library()
 
     docs = api.get_documents(user=user)
 
@@ -52,6 +53,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-U', '--user',
         default=default_user,
@@ -63,7 +70,9 @@ def main():
     try:
 
         show_docs(
-            url = args.api_url, user = args.user
+            url = args.api_url,
+            user = args.user,
+            token = args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_library_processing.py

@@ -9,10 +9,11 @@ import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
 default_user = "trustgraph"
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def show_procs(url, user):
+def show_procs(url, user, token=None):
 
-    api = Api(url).library()
+    api = Api(url, token=token).library()
 
     procs = api.get_processings(user = user)
 
@@ -57,12 +58,18 @@ def main():
         help=f'User ID (default: {default_user})'
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
         show_procs(
-            url = args.api_url, user = args.user
+            url = args.api_url, user = args.user, token = args.token
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_mcp_tools.py

@@ -10,10 +10,11 @@ import tabulate
 import textwrap
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def show_config(url):
+def show_config(url, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     values = api.get_values(type="mcp")
 
@@ -57,12 +58,19 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
         show_config(
             url=args.api_url,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_parameter_types.py

@@ -13,6 +13,7 @@ from trustgraph.api import Api, ConfigKey
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
 def format_enum_values(enum_list):
     """
@@ -75,11 +76,11 @@ def format_constraints(param_type_def):
 
     return ", ".join(constraints) if constraints else "None"
 
-def show_parameter_types(url):
+def show_parameter_types(url, token=None):
     """
     Show all parameter type definitions
     """
-    api = Api(url)
+    api = Api(url, token=token)
     config_api = api.config()
 
     # Get list of all parameter types
@@ -145,6 +146,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-t', '--type',
         help='Show only the specified parameter type',
@@ -155,19 +162,19 @@ def main():
     try:
         if args.type:
             # Show specific parameter type
-            show_specific_parameter_type(args.api_url, args.type)
+            show_specific_parameter_type(args.api_url, args.type, args.token)
         else:
             # Show all parameter types
-            show_parameter_types(args.api_url)
+            show_parameter_types(args.api_url, args.token)
 
     except Exception as e:
         print("Exception:", e, flush=True)
 
-def show_specific_parameter_type(url, param_type_name):
+def show_specific_parameter_type(url, param_type_name, token=None):
     """
     Show a specific parameter type definition
     """
-    api = Api(url)
+    api = Api(url, token=token)
     config_api = api.config()
 
     try:

trustgraph-cli/trustgraph/cli/show_prompts.py

@@ -10,10 +10,11 @@ import tabulate
 import textwrap
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def show_config(url):
+def show_config(url, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     values = api.get([
         ConfigKey(type="prompt", key="system"),
@@ -78,12 +79,19 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
         show_config(
             url=args.api_url,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_token_costs.py

@@ -12,10 +12,11 @@ import textwrap
 tabulate.PRESERVE_WHITESPACE = True
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def show_config(url):
+def show_config(url, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     models = api.list("token-costs")
 
@@ -61,12 +62,19 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
         show_config(
             url=args.api_url,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/show_tools.py

@@ -17,10 +17,11 @@ import tabulate
 import textwrap
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def show_config(url):
+def show_config(url, token=None):
 
-    api = Api(url).config()
+    api = Api(url, token=token).config()
 
     values = api.get_values(type="tool")
 
@@ -100,12 +101,19 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     args = parser.parse_args()
 
     try:
 
         show_config(
             url=args.api_url,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/start_flow.py

@@ -17,10 +17,11 @@ from trustgraph.api import Api
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def start_flow(url, class_name, flow_id, description, parameters=None):
+def start_flow(url, class_name, flow_id, description, parameters=None, token=None):
 
-    api = Api(url).flow()
+    api = Api(url, token=token).flow()
 
     api.start(
         class_name = class_name,
@@ -42,6 +43,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-n', '--class-name',
         required=True,
@@ -112,6 +119,7 @@ def main():
             flow_id = args.flow_id,
             description = args.description,
             parameters = parameters,
+            token = args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/start_library_processing.py

@@ -9,13 +9,14 @@ from trustgraph.api import Api, ConfigKey
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = "trustgraph"
 
 def start_processing(
-        url, user, document_id, id, flow, collection, tags
+        url, user, document_id, id, flow, collection, tags, token=None
 ):
 
-    api = Api(url).library()
+    api = Api(url, token=token).library()
 
     if tags:
         tags = tags.split(",")
@@ -44,6 +45,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-U', '--user',
         default=default_user,
@@ -90,7 +97,8 @@ def main():
             id = args.id,
             flow = args.flow_id,
             collection = args.collection,
-            tags = args.tags
+            tags = args.tags,
+            token = args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/stop_flow.py

@@ -9,10 +9,11 @@ from trustgraph.api import Api
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def stop_flow(url, flow_id):
+def stop_flow(url, flow_id, token=None):
 
-    api = Api(url).flow()
+    api = Api(url, token=token).flow()
 
     api.stop(id = flow_id)
 
@@ -29,6 +30,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-i', '--flow-id',
         required=True,
@@ -42,6 +49,7 @@ def main():
         stop_flow(
             url=args.api_url,
             flow_id=args.flow_id,
+            token=args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/stop_library_processing.py

@@ -10,13 +10,14 @@ from trustgraph.api import Api, ConfigKey
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = "trustgraph"
 
 def stop_processing(
-        url, user, id
+        url, user, id, token=None
 ):
 
-    api = Api(url).library()
+    api = Api(url, token=token).library()
 
     api.stop_processing(user = user, id = id)
 
@@ -33,6 +34,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-U', '--user',
         default=default_user,
@@ -53,6 +60,7 @@ def main():
             url = args.api_url,
             user = args.user,
             id = args.id,
+            token = args.token,
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/unload_kg_core.py

@@ -11,12 +11,13 @@ from trustgraph.api import Api
 import json
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_flow = "default"
 default_collection = "default"
 
-def unload_kg_core(url, user, id, flow):
+def unload_kg_core(url, user, id, flow, token=None):
 
-    api = Api(url).knowledge()
+    api = Api(url, token=token).knowledge()
 
     class_names = api.unload_kg_core(user = user, id = id, flow=flow)
 
@@ -33,6 +34,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-U', '--user',
         default="trustgraph",
@@ -60,6 +67,7 @@ def main():
             user=args.user,
             id=args.id,
             flow=args.flow_id,
+            token=args.token,
         )
 
     except Exception as e: