diff --git a/trustgraph-base/trustgraph/api/__init__.py b/trustgraph-base/trustgraph/api/__init__.py index daa6a964..7a3fc86d 100644 --- a/trustgraph-base/trustgraph/api/__init__.py +++ b/trustgraph-base/trustgraph/api/__init__.py @@ -1,3 +1,80 @@ -from . api import * +# Core API +from .api import Api + +# Flow clients +from .flow import Flow, FlowInstance +from .async_flow import AsyncFlow, AsyncFlowInstance + +# WebSocket clients +from .socket_client import SocketClient, SocketFlowInstance +from .async_socket_client import AsyncSocketClient, AsyncSocketFlowInstance + +# Bulk operation clients +from .bulk_client import BulkClient +from .async_bulk_client import AsyncBulkClient + +# Metrics clients +from .metrics import Metrics +from .async_metrics import AsyncMetrics + +# Types +from .types import ( + Triple, + ConfigKey, + ConfigValue, + DocumentMetadata, + ProcessingMetadata, + CollectionMetadata, + StreamingChunk, + AgentThought, + AgentObservation, + AgentAnswer, + RAGChunk, +) + +# Exceptions +from .exceptions import ProtocolException, ApplicationException + +__all__ = [ + # Core API + "Api", + + # Flow clients + "Flow", + "FlowInstance", + "AsyncFlow", + "AsyncFlowInstance", + + # WebSocket clients + "SocketClient", + "SocketFlowInstance", + "AsyncSocketClient", + "AsyncSocketFlowInstance", + + # Bulk operation clients + "BulkClient", + "AsyncBulkClient", + + # Metrics clients + "Metrics", + "AsyncMetrics", + + # Types + "Triple", + "ConfigKey", + "ConfigValue", + "DocumentMetadata", + "ProcessingMetadata", + "CollectionMetadata", + "StreamingChunk", + "AgentThought", + "AgentObservation", + "AgentAnswer", + "RAGChunk", + + # Exceptions + "ProtocolException", + "ApplicationException", +] diff --git a/trustgraph-base/trustgraph/api/api.py b/trustgraph-base/trustgraph/api/api.py index b0bae8ce..d1f07513 100644 --- a/trustgraph-base/trustgraph/api/api.py +++ b/trustgraph-base/trustgraph/api/api.py @@ -3,6 +3,7 @@ import requests import json import base64 import time +from typing import Optional from . library import Library from . flow import Flow @@ -26,7 +27,7 @@ def check_error(response): class Api: - def __init__(self, url="http://localhost:8088/", timeout=60): + def __init__(self, url="http://localhost:8088/", timeout=60, token: Optional[str] = None): self.url = url @@ -36,6 +37,16 @@ class Api: self.url += "api/v1/" self.timeout = timeout + self.token = token + + # Lazy initialization for new clients + self._socket_client = None + self._bulk_client = None + self._async_flow = None + self._async_socket_client = None + self._async_bulk_client = None + self._metrics = None + self._async_metrics = None def flow(self): return Flow(api=self) @@ -50,8 +61,12 @@ class Api: url = f"{self.url}{path}" + headers = {} + if self.token: + headers["Authorization"] = f"Bearer {self.token}" + # Invoke the API, input is passed as JSON - resp = requests.post(url, json=request, timeout=self.timeout) + resp = requests.post(url, json=request, timeout=self.timeout, headers=headers) # Should be a 200 status code if resp.status_code != 200: @@ -72,3 +87,96 @@ class Api: def collection(self): return Collection(self) + + # New synchronous methods + def socket(self): + """Synchronous WebSocket-based interface for streaming operations""" + if self._socket_client is None: + from . socket_client import SocketClient + # Extract base URL (remove api/v1/ suffix) + base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/") + self._socket_client = SocketClient(base_url, self.timeout, self.token) + return self._socket_client + + def bulk(self): + """Synchronous bulk operations interface for import/export""" + if self._bulk_client is None: + from . bulk_client import BulkClient + # Extract base URL (remove api/v1/ suffix) + base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/") + self._bulk_client = BulkClient(base_url, self.timeout, self.token) + return self._bulk_client + + def metrics(self): + """Synchronous metrics interface""" + if self._metrics is None: + from . metrics import Metrics + # Extract base URL (remove api/v1/ suffix) + base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/") + self._metrics = Metrics(base_url, self.timeout, self.token) + return self._metrics + + # New asynchronous methods + def async_flow(self): + """Asynchronous REST-based flow interface""" + if self._async_flow is None: + from . async_flow import AsyncFlow + self._async_flow = AsyncFlow(self.url, self.timeout, self.token) + return self._async_flow + + def async_socket(self): + """Asynchronous WebSocket-based interface for streaming operations""" + if self._async_socket_client is None: + from . async_socket_client import AsyncSocketClient + # Extract base URL (remove api/v1/ suffix) + base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/") + self._async_socket_client = AsyncSocketClient(base_url, self.timeout, self.token) + return self._async_socket_client + + def async_bulk(self): + """Asynchronous bulk operations interface for import/export""" + if self._async_bulk_client is None: + from . async_bulk_client import AsyncBulkClient + # Extract base URL (remove api/v1/ suffix) + base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/") + self._async_bulk_client = AsyncBulkClient(base_url, self.timeout, self.token) + return self._async_bulk_client + + def async_metrics(self): + """Asynchronous metrics interface""" + if self._async_metrics is None: + from . async_metrics import AsyncMetrics + # Extract base URL (remove api/v1/ suffix) + base_url = self.url.rsplit("api/v1/", 1)[0].rstrip("/") + self._async_metrics = AsyncMetrics(base_url, self.timeout, self.token) + return self._async_metrics + + # Resource management + def close(self): + """Close all synchronous connections""" + if self._socket_client: + self._socket_client.close() + if self._bulk_client: + self._bulk_client.close() + + async def aclose(self): + """Close all asynchronous connections""" + if self._async_socket_client: + await self._async_socket_client.aclose() + if self._async_bulk_client: + await self._async_bulk_client.aclose() + if self._async_flow: + await self._async_flow.aclose() + + # Context manager support + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + await self.aclose() diff --git a/trustgraph-base/trustgraph/api/async_bulk_client.py b/trustgraph-base/trustgraph/api/async_bulk_client.py new file mode 100644 index 00000000..cf18cb3a --- /dev/null +++ b/trustgraph-base/trustgraph/api/async_bulk_client.py @@ -0,0 +1,131 @@ + +import json +import websockets +from typing import Optional, AsyncIterator, Dict, Any + +from . types import Triple + + +class AsyncBulkClient: + """Asynchronous bulk operations 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 + + 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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): + """Close connections""" + # Cleanup handled by context managers + pass diff --git a/trustgraph-base/trustgraph/api/async_flow.py b/trustgraph-base/trustgraph/api/async_flow.py new file mode 100644 index 00000000..65e39476 --- /dev/null +++ b/trustgraph-base/trustgraph/api/async_flow.py @@ -0,0 +1,245 @@ + +import aiohttp +import json +from typing import Optional, Dict, Any, List + +from . exceptions import ProtocolException, ApplicationException + + +def check_error(response): + if "error" in response: + try: + msg = response["error"]["message"] + tp = response["error"]["type"] + except: + raise ApplicationException(response["error"]) + + raise ApplicationException(f"{tp}: {msg}") + + +class AsyncFlow: + """Asynchronous REST-based flow interface""" + + def __init__(self, url: str, timeout: int, token: Optional[str]): + self.url = url + self.timeout = timeout + self.token = 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): + """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) -> 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) -> 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) -> 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) -> 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): + """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): + """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): + """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): + """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) diff --git a/trustgraph-base/trustgraph/api/async_metrics.py b/trustgraph-base/trustgraph/api/async_metrics.py new file mode 100644 index 00000000..3a00926f --- /dev/null +++ b/trustgraph-base/trustgraph/api/async_metrics.py @@ -0,0 +1,29 @@ + +import aiohttp +from typing import Optional + + +class AsyncMetrics: + """Asynchronous metrics client""" + + def __init__(self, url: str, timeout: int, token: Optional[str]): + self.url = url + self.timeout = timeout + self.token = token + + async def get(self) -> str: + """Get Prometheus metrics as text""" + url = f"{self.url}/api/metrics" + + headers = {} + 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() diff --git a/trustgraph-base/trustgraph/api/async_socket_client.py b/trustgraph-base/trustgraph/api/async_socket_client.py new file mode 100644 index 00000000..4491d025 --- /dev/null +++ b/trustgraph-base/trustgraph/api/async_socket_client.py @@ -0,0 +1,296 @@ + +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], streaming: bool = False): + """Async WebSocket request implementation""" + # 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)) + + if streaming: + # 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 + else: + # 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"] + + 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("agent", self.flow_id, request, True) + else: + return await self.client._send_request("agent", self.flow_id, request, False) + + 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: + async for chunk in self.client._send_request("text-completion", self.flow_id, request, True): + if hasattr(chunk, 'content'): + yield chunk.content + else: + result = await self.client._send_request("text-completion", self.flow_id, request, False) + 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, 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: + async for chunk in self.client._send_request("graph-rag", self.flow_id, request, True): + if hasattr(chunk, 'content'): + yield chunk.content + else: + result = await self.client._send_request("graph-rag", self.flow_id, request, False) + return result.get("response", "") + + 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: + async for chunk in self.client._send_request("document-rag", self.flow_id, request, True): + if hasattr(chunk, 'content'): + yield chunk.content + else: + result = await self.client._send_request("document-rag", self.flow_id, request, False) + return result.get("response", "") + + 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: + async for chunk in self.client._send_request("prompt", self.flow_id, request, True): + if hasattr(chunk, 'content'): + yield chunk.content + else: + result = await self.client._send_request("prompt", self.flow_id, request, False) + return result.get("response", "") + + 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, False) + + 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, False) + + 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, False) + + 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, False) + + 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, False) diff --git a/trustgraph-base/trustgraph/api/bulk_client.py b/trustgraph-base/trustgraph/api/bulk_client.py new file mode 100644 index 00000000..aa5648da --- /dev/null +++ b/trustgraph-base/trustgraph/api/bulk_client.py @@ -0,0 +1,272 @@ + +import json +import asyncio +import websockets +from typing import Optional, Iterator, Dict, Any + +from . types import Triple +from . exceptions import ProtocolException + + +class BulkClient: + """Synchronous bulk operations 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 + + 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): + """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) -> 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]): + """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) -> 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 Runtime + +Error: + 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): + """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) -> 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]]): + """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) -> 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): + """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) -> 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]]): + """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) -> 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): + """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) -> 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]]): + """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) -> 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): + """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) -> 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]]): + """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): + """Close connections""" + # Cleanup handled by context managers + pass diff --git a/trustgraph-base/trustgraph/api/flow.py b/trustgraph-base/trustgraph/api/flow.py index 0214a4bd..f7cda5c2 100644 --- a/trustgraph-base/trustgraph/api/flow.py +++ b/trustgraph-base/trustgraph/api/flow.py @@ -211,6 +211,21 @@ class FlowInstance: input )["vectors"] + def graph_embeddings_query(self, text, user, collection, limit=10): + + # Query graph embeddings for semantic search + input = { + "text": text, + "user": user, + "collection": collection, + "limit": limit + } + + return self.request( + "service/graph-embeddings", + input + ) + def prompt(self, id, variables): input = { diff --git a/trustgraph-base/trustgraph/api/metrics.py b/trustgraph-base/trustgraph/api/metrics.py new file mode 100644 index 00000000..4dcd93e0 --- /dev/null +++ b/trustgraph-base/trustgraph/api/metrics.py @@ -0,0 +1,27 @@ + +import requests +from typing import Optional + + +class Metrics: + """Synchronous metrics client""" + + def __init__(self, url: str, timeout: int, token: Optional[str]): + self.url = url + self.timeout = timeout + self.token = token + + def get(self) -> str: + """Get Prometheus metrics as text""" + url = f"{self.url}/api/metrics" + + headers = {} + 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 diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py new file mode 100644 index 00000000..eb571096 --- /dev/null +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -0,0 +1,343 @@ + +import json +import asyncio +import websockets +from typing import Optional, Dict, Any, Iterator, Union +from threading import Lock + +from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk +from . exceptions import ProtocolException, ApplicationException + + +class SocketClient: + """Synchronous WebSocket client (wraps async websockets library)""" + + 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._connection = None + self._request_counter = 0 + self._lock = Lock() + self._loop = 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): + """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): + """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, False)) + + def _streaming_generator(self, service: str, flow: Optional[str], request: Dict[str, Any], loop): + """Generator that yields streaming chunks""" + async_gen = self._send_request_async(service, flow, request, True) + + 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], streaming: bool): + """Async implementation of WebSocket request""" + # 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)) + + if streaming: + # 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 + else: + # 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"] + + 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") + ) + + def close(self): + """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): + self.client = client + self.flow_id = flow_id + + 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], Iterator]: + """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) -> 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) -> 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) -> 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): + """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): + """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=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 self.client._send_request_sync("triples", self.flow_id, request, False) + + 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 self.client._send_request_sync("objects", self.flow_id, request, False) + + def mcp_tool(self, name: str, parameters: Dict[str, Any], **kwargs): + """Execute MCP tool""" + request = { + "name": name, + "parameters": parameters + } + request.update(kwargs) + + return self.client._send_request_sync("mcp-tool", self.flow_id, request, False) diff --git a/trustgraph-base/trustgraph/api/types.py b/trustgraph-base/trustgraph/api/types.py index 71b438f6..016f2c7a 100644 --- a/trustgraph-base/trustgraph/api/types.py +++ b/trustgraph-base/trustgraph/api/types.py @@ -1,7 +1,7 @@ import dataclasses import datetime -from typing import List +from typing import List, Optional, Dict, Any from .. knowledge import hash, Uri, Literal @dataclasses.dataclass @@ -51,3 +51,33 @@ class CollectionMetadata: tags : List[str] created_at : str updated_at : str + +# Streaming chunk types + +@dataclasses.dataclass +class StreamingChunk: + """Base class for streaming chunks""" + content: str + end_of_message: bool = False + +@dataclasses.dataclass +class AgentThought(StreamingChunk): + """Agent reasoning chunk""" + chunk_type: str = "thought" + +@dataclasses.dataclass +class AgentObservation(StreamingChunk): + """Agent tool observation chunk""" + chunk_type: str = "observation" + +@dataclasses.dataclass +class AgentAnswer(StreamingChunk): + """Agent final answer chunk""" + chunk_type: str = "final-answer" + end_of_dialog: bool = False + +@dataclasses.dataclass +class RAGChunk(StreamingChunk): + """RAG streaming chunk""" + end_of_stream: bool = False + error: Optional[Dict[str, str]] = None