mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-05-03 20:32:38 +02:00
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:
parent
b957004db9
commit
01aeede78b
53 changed files with 4489 additions and 715 deletions
1508
docs/tech-specs/python-api-refactor.md
Normal file
1508
docs/tech-specs/python-api-refactor.md
Normal file
File diff suppressed because it is too large
Load diff
446
tests/unit/test_python_api_client.py
Normal file
446
tests/unit/test_python_api_client.py
Normal 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"])
|
||||||
|
|
@ -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",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
131
trustgraph-base/trustgraph/api/async_bulk_client.py
Normal file
131
trustgraph-base/trustgraph/api/async_bulk_client.py
Normal 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
|
||||||
245
trustgraph-base/trustgraph/api/async_flow.py
Normal file
245
trustgraph-base/trustgraph/api/async_flow.py
Normal 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)
|
||||||
33
trustgraph-base/trustgraph/api/async_metrics.py
Normal file
33
trustgraph-base/trustgraph/api/async_metrics.py
Normal 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
|
||||||
335
trustgraph-base/trustgraph/api/async_socket_client.py
Normal file
335
trustgraph-base/trustgraph/api/async_socket_client.py
Normal 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)
|
||||||
270
trustgraph-base/trustgraph/api/bulk_client.py
Normal file
270
trustgraph-base/trustgraph/api/bulk_client.py
Normal 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
|
||||||
|
|
@ -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 = {
|
||||||
|
|
|
||||||
27
trustgraph-base/trustgraph/api/metrics.py
Normal file
27
trustgraph-base/trustgraph/api/metrics.py
Normal 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
|
||||||
445
trustgraph-base/trustgraph/api/socket_client.py
Normal file
445
trustgraph-base/trustgraph/api/socket_client.py
Normal 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)
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue