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
This commit is contained in:
cybermaggedon 2025-12-04 17:38:57 +00:00 committed by GitHub
parent b957004db9
commit 01aeede78b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
53 changed files with 4489 additions and 715 deletions

File diff suppressed because it is too large Load diff

View file

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

View file

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

View file

@ -3,6 +3,7 @@ import requests
import json import json
import base64 import base64
import time import time
from typing import Optional
from . library import Library from . library import Library
from . flow import Flow from . flow import Flow
@ -26,7 +27,7 @@ def check_error(response):
class Api: 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 self.url = url
@ -36,6 +37,16 @@ class Api:
self.url += "api/v1/" self.url += "api/v1/"
self.timeout = timeout 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): def flow(self):
return Flow(api=self) return Flow(api=self)
@ -50,8 +61,12 @@ class Api:
url = f"{self.url}{path}" url = f"{self.url}{path}"
headers = {}
if self.token:
headers["Authorization"] = f"Bearer {self.token}"
# Invoke the API, input is passed as JSON # 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 # Should be a 200 status code
if resp.status_code != 200: if resp.status_code != 200:
@ -72,3 +87,96 @@ class Api:
def collection(self): def collection(self):
return 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()

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -211,6 +211,21 @@ class FlowInstance:
input input
)["vectors"] )["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): def prompt(self, id, variables):
input = { input = {

View file

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

View file

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

View file

@ -1,7 +1,7 @@
import dataclasses import dataclasses
import datetime import datetime
from typing import List from typing import List, Optional, Dict, Any
from .. knowledge import hash, Uri, Literal from .. knowledge import hash, Uri, Literal
@dataclasses.dataclass @dataclasses.dataclass
@ -51,3 +51,33 @@ class CollectionMetadata:
tags : List[str] tags : List[str]
created_at : str created_at : str
updated_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

View file

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

View file

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

View file

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

View file

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

View file

@ -4,89 +4,50 @@ Uses the DocumentRAG service to answer a question
import argparse import argparse
import os import os
import asyncio
import json
import uuid
from websockets.asyncio.client import connect
from trustgraph.api import Api from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
default_user = 'trustgraph' default_user = 'trustgraph'
default_collection = 'default' default_collection = 'default'
default_doc_limit = 10 default_doc_limit = 10
async def question_streaming(url, flow_id, question, user, collection, doc_limit): def question(url, flow_id, question, user, collection, doc_limit, streaming=True, token=None):
"""Streaming version using websockets"""
# Convert http:// to ws:// # Create API client
if url.startswith('http://'): api = Api(url=url, token=token)
url = 'ws://' + url[7:]
elif url.startswith('https://'):
url = 'wss://' + url[8:]
if not url.endswith("/"): if streaming:
url += "/" # 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: finally:
req = { socket.close()
"id": mid, else:
"service": "document-rag", # Use REST API for non-streaming
"flow": flow_id, flow = api.flow().id(flow_id)
"request": { resp = flow.document_rag(
"query": question, question=question,
"user": user, user=user,
"collection": collection, collection=collection,
"doc-limit": doc_limit, doc_limit=doc_limit,
"streaming": True )
} print(resp)
}
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)
def main(): def main():
@ -101,6 +62,12 @@ def main():
help=f'API URL (default: {default_url})', help=f'API URL (default: {default_url})',
) )
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
parser.add_argument( parser.add_argument(
'-f', '--flow-id', '-f', '--flow-id',
default="default", default="default",
@ -127,6 +94,7 @@ def main():
parser.add_argument( parser.add_argument(
'-d', '--doc-limit', '-d', '--doc-limit',
type=int,
default=default_doc_limit, default=default_doc_limit,
help=f'Document limit (default: {default_doc_limit})' help=f'Document limit (default: {default_doc_limit})'
) )
@ -141,30 +109,20 @@ def main():
try: try:
if not args.no_streaming: question(
asyncio.run( url=args.url,
question_streaming( flow_id=args.flow_id,
url=args.url, question=args.question,
flow_id=args.flow_id, user=args.user,
question=args.question, collection=args.collection,
user=args.user, doc_limit=args.doc_limit,
collection=args.collection, streaming=not args.no_streaming,
doc_limit=args.doc_limit, token=args.token,
) )
)
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,
)
except Exception as e: except Exception as e:
print("Exception:", e, flush=True) print("Exception:", e, flush=True)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

View file

@ -4,13 +4,10 @@ Uses the GraphRAG service to answer a question
import argparse import argparse
import os import os
import asyncio
import json
import uuid
from websockets.asyncio.client import connect
from trustgraph.api import Api from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
default_user = 'trustgraph' default_user = 'trustgraph'
default_collection = 'default' default_collection = 'default'
default_entity_limit = 50 default_entity_limit = 50
@ -18,89 +15,51 @@ default_triple_limit = 30
default_max_subgraph_size = 150 default_max_subgraph_size = 150
default_max_path_length = 2 default_max_path_length = 2
async def question_streaming( def question(
url, flow_id, question, user, collection, entity_limit, triple_limit, 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:// # Create API client
if url.startswith('http://'): api = Api(url=url, token=token)
url = 'ws://' + url[7:]
elif url.startswith('https://'):
url = 'wss://' + url[8:]
if not url.endswith("/"): if streaming:
url += "/" # 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: finally:
req = { socket.close()
"id": mid, else:
"service": "graph-rag", # Use REST API for non-streaming
"flow": flow_id, flow = api.flow().id(flow_id)
"request": { resp = flow.graph_rag(
"query": question, question=question,
"user": user, user=user,
"collection": collection, collection=collection,
"entity-limit": entity_limit, entity_limit=entity_limit,
"triple-limit": triple_limit, triple_limit=triple_limit,
"max-subgraph-size": max_subgraph_size, max_subgraph_size=max_subgraph_size,
"max-path-length": max_path_length, max_path_length=max_path_length
"streaming": True )
} print(resp)
}
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)
def main(): def main():
@ -115,6 +74,12 @@ def main():
help=f'API URL (default: {default_url})', help=f'API URL (default: {default_url})',
) )
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
parser.add_argument( parser.add_argument(
'-f', '--flow-id', '-f', '--flow-id',
default="default", default="default",
@ -141,24 +106,28 @@ def main():
parser.add_argument( parser.add_argument(
'-e', '--entity-limit', '-e', '--entity-limit',
type=int,
default=default_entity_limit, default=default_entity_limit,
help=f'Entity limit (default: {default_entity_limit})' help=f'Entity limit (default: {default_entity_limit})'
) )
parser.add_argument( parser.add_argument(
'-t', '--triple-limit', '--triple-limit',
type=int,
default=default_triple_limit, default=default_triple_limit,
help=f'Triple limit (default: {default_triple_limit})' help=f'Triple limit (default: {default_triple_limit})'
) )
parser.add_argument( parser.add_argument(
'-s', '--max-subgraph-size', '-s', '--max-subgraph-size',
type=int,
default=default_max_subgraph_size, default=default_max_subgraph_size,
help=f'Max subgraph size (default: {default_max_subgraph_size})' help=f'Max subgraph size (default: {default_max_subgraph_size})'
) )
parser.add_argument( parser.add_argument(
'-p', '--max-path-length', '-p', '--max-path-length',
type=int,
default=default_max_path_length, default=default_max_path_length,
help=f'Max path length (default: {default_max_path_length})' help=f'Max path length (default: {default_max_path_length})'
) )
@ -173,36 +142,23 @@ def main():
try: try:
if not args.no_streaming: question(
asyncio.run( url=args.url,
question_streaming( flow_id=args.flow_id,
url=args.url, question=args.question,
flow_id=args.flow_id, user=args.user,
question=args.question, collection=args.collection,
user=args.user, entity_limit=args.entity_limit,
collection=args.collection, triple_limit=args.triple_limit,
entity_limit=args.entity_limit, max_subgraph_size=args.max_subgraph_size,
triple_limit=args.triple_limit, max_path_length=args.max_path_length,
max_subgraph_size=args.max_subgraph_size, streaming=not args.no_streaming,
max_path_length=args.max_path_length, token=args.token,
) )
)
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,
)
except Exception as e: except Exception as e:
print("Exception:", e, flush=True) print("Exception:", e, flush=True)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

View file

@ -5,64 +5,39 @@ and user prompt. Both arguments are required.
import argparse import argparse
import os import os
import json from trustgraph.api import Api
import uuid
import asyncio
from websockets.asyncio.client import connect
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("/"): # Create API client
url += "/" 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: finally:
# Clean up socket connection
req = { socket.close()
"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()
def main(): def main():
@ -77,6 +52,12 @@ def main():
help=f'API URL (default: {default_url})', help=f'API URL (default: {default_url})',
) )
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
parser.add_argument( parser.add_argument(
'system', 'system',
nargs=1, nargs=1,
@ -105,17 +86,18 @@ def main():
try: try:
asyncio.run(query( query(
url=args.url, url=args.url,
flow_id=args.flow_id, flow_id=args.flow_id,
system=args.system[0], system=args.system[0],
prompt=args.prompt[0], prompt=args.prompt[0],
streaming=not args.no_streaming streaming=not args.no_streaming,
)) token=args.token,
)
except Exception as e: except Exception as e:
print("Exception:", e, flush=True) print("Exception:", e, flush=True)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

View file

@ -10,76 +10,61 @@ using key=value arguments on the command line, and these replace
import argparse import argparse
import os import os
import json import json
import uuid from trustgraph.api import Api
import asyncio
from websockets.asyncio.client import connect
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("/"): # Create API client
url += "/" 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 = { # Check if this is an object response (JSON)
"id": mid, if hasattr(chunk, 'object') and chunk.object:
"service": "prompt", full_response["object"] = chunk.object
"flow": flow_id,
"request": {
"id": template_id,
"variables": variables,
"streaming": streaming
}
}
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": ""} else:
# Non-streaming: handle response
while True: if isinstance(response, str):
print(response)
msg = await ws.recv() elif isinstance(response, dict):
if "text" in response:
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
print(response["text"]) print(response["text"])
elif "object" in response:
print(json.dumps(json.loads(response["object"]), indent=4))
# Handle object responses (JSON, never streamed) finally:
if "object" in response and response["object"]: # Clean up socket connection
full_response["object"] = response["object"] socket.close()
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()
def main(): def main():
@ -94,6 +79,12 @@ def main():
help=f'API URL (default: {default_url})', help=f'API URL (default: {default_url})',
) )
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
parser.add_argument( parser.add_argument(
'-f', '--flow-id', '-f', '--flow-id',
default="default", default="default",
@ -135,17 +126,18 @@ specified multiple times''',
try: try:
asyncio.run(query( query(
url=args.url, url=args.url,
flow_id=args.flow_id, flow_id=args.flow_id,
template_id=args.id[0], template_id=args.id[0],
variables=variables, variables=variables,
streaming=not args.no_streaming streaming=not args.no_streaming,
)) token=args.token,
)
except Exception as e: except Exception as e:
print("Exception:", e, flush=True) print("Exception:", e, flush=True)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

View file

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

View file

@ -2,18 +2,17 @@
Loads triples and entity contexts into the knowledge graph. Loads triples and entity contexts into the knowledge graph.
""" """
import asyncio
import argparse import argparse
import os import os
import time import time
import rdflib import rdflib
import json from typing import Iterator, Tuple
from websockets.asyncio.client import connect
from typing import List, Dict, Any
from trustgraph.api import Api, Triple
from trustgraph.log_level import LogLevel 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_user = 'trustgraph'
default_collection = 'default' default_collection = 'default'
@ -26,108 +25,114 @@ class KnowledgeLoader:
user, user,
collection, collection,
document_id, 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.files = files
self.flow = flow
self.user = user self.user = user
self.collection = collection self.collection = collection
self.document_id = document_id self.document_id = document_id
self.url = url
self.token = token
async def run(self): def load_triples_from_file(self, file) -> Iterator[Triple]:
"""Generator that yields Triple objects from a Turtle file"""
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):
g = rdflib.Graph() g = rdflib.Graph()
g.parse(file, format="turtle") g.parse(file, format="turtle")
def Value(value, is_uri):
return { "v": value, "e": is_uri }
for e in g: for e in g:
s = Value(value=str(e[0]), is_uri=True) # Extract subject, predicate, object
p = Value(value=str(e[1]), is_uri=True) s_value = str(e[0])
if type(e[2]) == rdflib.term.URIRef: p_value = str(e[1])
o = Value(value=str(e[2]), is_uri=True)
# 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: else:
o = Value(value=str(e[2]), is_uri=False) o_value = str(e[2])
o_is_uri = False
req = { # Create Triple object
"metadata": { # Note: The Triple dataclass has 's', 'p', 'o' fields as strings
"id": self.document_id, # The API will handle the metadata wrapping
"metadata": [], yield Triple(s=s_value, p=p_value, o=o_value)
"user": self.user,
"collection": self.collection
},
"triples": [
{
"s": s,
"p": p,
"o": o,
}
]
}
await ws.send(json.dumps(req)) def load_entity_contexts_from_file(self, file) -> Iterator[Tuple[str, str]]:
"""Generator that yields (entity, context) tuples from a Turtle file"""
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.
"""
g = rdflib.Graph() g = rdflib.Graph()
g.parse(file, format="turtle") g.parse(file, format="turtle")
for s, p, o in g: 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): if isinstance(o, rdflib.term.URIRef):
continue 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) s_str = str(s)
o_str = str(o) o_str = str(o)
req = { yield (s_str, o_str)
"metadata": {
"id": self.document_id, def run(self):
"metadata": [], """Load triples and entity contexts using Python API"""
"user": self.user,
"collection": self.collection try:
}, # Create API client
"entities": [ api = Api(url=self.url, token=self.token)
{ bulk = api.bulk()
"entity": {
"v": s_str, # Load triples from all files
"e": True print("Loading triples...")
}, for file in self.files:
"context": o_str 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(): def main():
@ -142,6 +147,12 @@ def main():
help=f'API URL (default: {default_url})', help=f'API URL (default: {default_url})',
) )
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
parser.add_argument( parser.add_argument(
'-i', '--document-id', '-i', '--document-id',
required=True, required=True,
@ -166,7 +177,6 @@ def main():
help=f'Collection ID (default: {default_collection})' help=f'Collection ID (default: {default_collection})'
) )
parser.add_argument( parser.add_argument(
'files', nargs='+', 'files', nargs='+',
help=f'Turtle files to load' help=f'Turtle files to load'
@ -178,15 +188,16 @@ def main():
try: try:
loader = KnowledgeLoader( loader = KnowledgeLoader(
document_id = args.document_id, document_id=args.document_id,
url = args.api_url, url=args.api_url,
flow = args.flow_id, token=args.token,
files = args.files, flow=args.flow_id,
user = args.user, files=args.files,
collection = args.collection, user=args.user,
collection=args.collection,
) )
asyncio.run(loader.run()) loader.run()
print("Triples and entity contexts loaded.") print("Triples and entity contexts loaded.")
break break
@ -199,4 +210,4 @@ def main():
time.sleep(10) time.sleep(10)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

View file

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

View file

@ -22,6 +22,7 @@ import logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
def load_structured_data( def load_structured_data(
@ -41,7 +42,8 @@ def load_structured_data(
user: str = 'trustgraph', user: str = 'trustgraph',
collection: str = 'default', collection: str = 'default',
dry_run: bool = False, dry_run: bool = False,
verbose: bool = False verbose: bool = False,
token: str = None
): ):
""" """
Load structured data using a descriptor configuration. Load structured data using a descriptor configuration.
@ -133,9 +135,9 @@ def load_structured_data(
# Get batch size from descriptor # Get batch size from descriptor
batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000) batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000)
# Send to TrustGraph using shared function # 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 # Summary
format_info = descriptor.get('format', {}) format_info = descriptor.get('format', {})
@ -288,10 +290,10 @@ def load_structured_data(
# Get batch size from descriptor or use default # Get batch size from descriptor or use default
batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000) batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000)
# Send to TrustGraph # Send to TrustGraph
print(f"🚀 Importing {len(output_records)} records 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 # Get summary info from descriptor
format_info = descriptor.get('format', {}) format_info = descriptor.get('format', {})
@ -571,66 +573,30 @@ def _process_data_pipeline(input_file, descriptor_file, user, collection, sample
return output_records, descriptor return output_records, descriptor
def _send_to_trustgraph(objects, api_url, flow, batch_size=1000): def _send_to_trustgraph(objects, api_url, flow, batch_size=1000, token=None):
"""Send ExtractedObject records to TrustGraph using WebSocket""" """Send ExtractedObject records to TrustGraph using Python API"""
import json from trustgraph.api import Api
import asyncio
from websockets.asyncio.client import connect
try: 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) 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"\n📊 Import Summary:")
print(f"- Total records: {total_records}") print(f"- Total records: {total_records}")
print(f"- Successfully imported: {imported_count}") print(f"- Successfully imported: {total_records}")
print(f"- Failed: {failed_count}") print("✅ All records imported successfully!")
if failed_count > 0: return total_records
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
except Exception as e: except Exception as e:
logger.error(f"Failed to import data to TrustGraph: {e}") logger.error(f"Failed to import data to TrustGraph: {e}")
print(f"Import failed: {e}") print(f"Import failed: {e}")
@ -1024,7 +990,13 @@ For more information on the descriptor format, see:
'--error-file', '--error-file',
help='Path to write error records (optional)' 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() args = parser.parse_args()
# Input validation # Input validation
@ -1077,7 +1049,8 @@ For more information on the descriptor format, see:
user=args.user, user=args.user,
collection=args.collection, collection=args.collection,
dry_run=args.dry_run, dry_run=args.dry_run,
verbose=args.verbose verbose=args.verbose,
token=args.token
) )
except FileNotFoundError as e: except FileNotFoundError as e:
print(f"Error: File not found - {e}", file=sys.stderr) print(f"Error: File not found - {e}", file=sys.stderr)

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -10,10 +10,11 @@ import tabulate
import textwrap import textwrap
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') 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([ api.put([
ConfigValue( ConfigValue(
@ -95,6 +96,12 @@ def main():
help=f'Input costs in $ per 1M tokens', 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() args = parser.parse_args()
try: try:

View file

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

View file

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

View file

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

View file

@ -9,10 +9,11 @@ import os
default_metrics_url = "http://localhost:8088/api/metrics" default_metrics_url = "http://localhost:8088/api/metrics"
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') 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) flow = api.get(flow_id)
class_name = flow["class-name"] class_name = flow["class-name"]
@ -77,11 +78,17 @@ def main():
help=f'Metrics URL (default: {default_metrics_url})', 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() args = parser.parse_args()
try: 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: except Exception as e:

View file

@ -9,6 +9,7 @@ from trustgraph.api import Api, ConfigKey
import json import json
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
def get_interface(config_api, i): 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" 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() config_api = api.config()
flow_api = api.flow() flow_api = api.flow()
@ -199,12 +200,19 @@ def main():
help=f'API URL (default: {default_url})', help=f'API URL (default: {default_url})',
) )
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
args = parser.parse_args() args = parser.parse_args()
try: try:
show_flows( show_flows(
url=args.api_url, url=args.api_url,
token=args.token,
) )
except Exception as e: except Exception as e:

View file

@ -9,10 +9,11 @@ from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = 'trustgraph' default_user = 'trustgraph'
default_collection = 'default' 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( rows = api.triples_query(
user=user, collection=collection, user=user, collection=collection,
@ -53,6 +54,12 @@ def main():
help=f'Collection ID (default: {default_collection})' 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() args = parser.parse_args()
try: try:
@ -62,6 +69,7 @@ def main():
flow_id = args.flow_id, flow_id = args.flow_id,
user = args.user, user = args.user,
collection = args.collection, collection = args.collection,
token = args.token,
) )
except Exception as e: except Exception as e:

View file

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

View file

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

View file

@ -9,10 +9,11 @@ import json
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = "trustgraph" 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) procs = api.get_processings(user = user)
@ -57,12 +58,18 @@ def main():
help=f'User ID (default: {default_user})' 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() args = parser.parse_args()
try: try:
show_procs( show_procs(
url = args.api_url, user = args.user url = args.api_url, user = args.user, token = args.token
) )
except Exception as e: except Exception as e:

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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