Fix Python streaming SDK issues (#580)

* Fix verify CLI issues

* Fixing content mechanisms in API

* Fixing error handling

* Fixing invoke_prompt, invoke_llm, invoke_agent

trustgraph-base/trustgraph/api/__init__.py

@@ -34,7 +34,26 @@ from .types import (
 )
 
 # Exceptions
-from .exceptions import ProtocolException, ApplicationException
+from .exceptions import (
+    ProtocolException,
+    TrustGraphException,
+    AgentError,
+    ConfigError,
+    DocumentRagError,
+    FlowError,
+    GatewayError,
+    GraphRagError,
+    LLMError,
+    LoadError,
+    LookupError,
+    NLPQueryError,
+    ObjectsQueryError,
+    RequestError,
+    StructuredQueryError,
+    UnexpectedError,
+    # Legacy alias
+    ApplicationException,
+)
 
 __all__ = [
     # Core API
@@ -75,6 +94,21 @@ __all__ = [
 
     # Exceptions
     "ProtocolException",
-    "ApplicationException",
+    "TrustGraphException",
+    "AgentError",
+    "ConfigError",
+    "DocumentRagError",
+    "FlowError",
+    "GatewayError",
+    "GraphRagError",
+    "LLMError",
+    "LoadError",
+    "LookupError",
+    "NLPQueryError",
+    "ObjectsQueryError",
+    "RequestError",
+    "StructuredQueryError",
+    "UnexpectedError",
+    "ApplicationException",  # Legacy alias
 ]
 

trustgraph-base/trustgraph/api/async_socket_client.py

@@ -130,18 +130,26 @@ class AsyncSocketClient:
                 content=resp.get("content", ""),
                 end_of_message=resp.get("end_of_message", False)
             )
-        elif chunk_type == "final-answer":
+        elif chunk_type == "answer" or 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)
             )
+        elif chunk_type == "action":
+            # Agent action chunks - treat as thoughts for display purposes
+            return AgentThought(
+                content=resp.get("content", ""),
+                end_of_message=resp.get("end_of_message", False)
+            )
         else:
             # RAG-style chunk (or generic chunk)
+            # Text-completion uses "response" field, RAG uses "chunk" field, Prompt uses "text" field
+            content = resp.get("response", resp.get("chunk", resp.get("text", "")))
             return RAGChunk(
-                content=resp.get("chunk", ""),
+                content=content,
                 end_of_stream=resp.get("end_of_stream", False),
-                error=resp.get("error")
+                error=None  # Errors are always thrown, never stored
             )
 
     async def aclose(self):

trustgraph-base/trustgraph/api/exceptions.py

@@ -1,6 +1,134 @@
+"""
+TrustGraph API Exceptions
 
+Exception hierarchy for errors returned by TrustGraph services.
+Each service error type maps to a specific exception class.
+"""
+
+# Protocol-level exceptions (communication errors)
 class ProtocolException(Exception):
+    """Raised when WebSocket protocol errors occur"""
     pass
 
-class ApplicationException(Exception):
+
+# Base class for all TrustGraph application errors
+class TrustGraphException(Exception):
+    """Base class for all TrustGraph service errors"""
+    def __init__(self, message: str, error_type: str = None):
+        super().__init__(message)
+        self.message = message
+        self.error_type = error_type
+
+
+# Service-specific exceptions
+class AgentError(TrustGraphException):
+    """Agent service error"""
     pass
+
+
+class ConfigError(TrustGraphException):
+    """Configuration service error"""
+    pass
+
+
+class DocumentRagError(TrustGraphException):
+    """Document RAG retrieval error"""
+    pass
+
+
+class FlowError(TrustGraphException):
+    """Flow management error"""
+    pass
+
+
+class GatewayError(TrustGraphException):
+    """API Gateway error"""
+    pass
+
+
+class GraphRagError(TrustGraphException):
+    """Graph RAG retrieval error"""
+    pass
+
+
+class LLMError(TrustGraphException):
+    """LLM service error"""
+    pass
+
+
+class LoadError(TrustGraphException):
+    """Data loading error"""
+    pass
+
+
+class LookupError(TrustGraphException):
+    """Lookup/search error"""
+    pass
+
+
+class NLPQueryError(TrustGraphException):
+    """NLP query service error"""
+    pass
+
+
+class ObjectsQueryError(TrustGraphException):
+    """Objects query service error"""
+    pass
+
+
+class RequestError(TrustGraphException):
+    """Request processing error"""
+    pass
+
+
+class StructuredQueryError(TrustGraphException):
+    """Structured query service error"""
+    pass
+
+
+class UnexpectedError(TrustGraphException):
+    """Unexpected/unknown error"""
+    pass
+
+
+# Mapping from error type string to exception class
+ERROR_TYPE_MAPPING = {
+    "agent-error": AgentError,
+    "config-error": ConfigError,
+    "document-rag-error": DocumentRagError,
+    "flow-error": FlowError,
+    "gateway-error": GatewayError,
+    "graph-rag-error": GraphRagError,
+    "llm-error": LLMError,
+    "load-error": LoadError,
+    "lookup-error": LookupError,
+    "nlp-query-error": NLPQueryError,
+    "objects-query-error": ObjectsQueryError,
+    "request-error": RequestError,
+    "structured-query-error": StructuredQueryError,
+    "unexpected-error": UnexpectedError,
+}
+
+
+def raise_from_error_dict(error_dict: dict) -> None:
+    """
+    Raise appropriate exception from TrustGraph error dictionary.
+
+    Args:
+        error_dict: Dictionary with 'type' and 'message' keys
+
+    Raises:
+        Appropriate TrustGraphException subclass based on error type
+    """
+    error_type = error_dict.get("type", "unexpected-error")
+    message = error_dict.get("message", "Unknown error")
+
+    # Look up exception class, default to UnexpectedError
+    exception_class = ERROR_TYPE_MAPPING.get(error_type, UnexpectedError)
+
+    # Raise the appropriate exception
+    raise exception_class(message, error_type)
+
+
+# Legacy exception for backwards compatibility
+ApplicationException = TrustGraphException

trustgraph-base/trustgraph/api/socket_client.py

@@ -6,7 +6,7 @@ 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
+from . exceptions import ProtocolException, raise_from_error_dict
 
 
 class SocketClient:
@@ -126,7 +126,7 @@ class SocketClient:
                 raise ProtocolException(f"Response ID mismatch")
 
             if "error" in response:
-                raise ApplicationException(response["error"])
+                raise_from_error_dict(response["error"])
 
             if "response" not in response:
                 raise ProtocolException(f"Missing response in message")
@@ -171,11 +171,15 @@ class SocketClient:
                     continue  # Ignore messages for other requests
 
                 if "error" in response:
-                    raise ApplicationException(response["error"])
+                    raise_from_error_dict(response["error"])
 
                 if "response" in response:
                     resp = response["response"]
 
+                    # Check for errors in response chunks
+                    if "error" in resp:
+                        raise_from_error_dict(resp["error"])
+
                     # Parse different chunk types
                     chunk = self._parse_chunk(resp)
                     yield chunk
@@ -198,18 +202,26 @@ class SocketClient:
                 content=resp.get("content", ""),
                 end_of_message=resp.get("end_of_message", False)
             )
-        elif chunk_type == "final-answer":
+        elif chunk_type == "answer" or 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)
             )
+        elif chunk_type == "action":
+            # Agent action chunks - treat as thoughts for display purposes
+            return AgentThought(
+                content=resp.get("content", ""),
+                end_of_message=resp.get("end_of_message", False)
+            )
         else:
             # RAG-style chunk (or generic chunk)
+            # Text-completion uses "response" field, RAG uses "chunk" field, Prompt uses "text" field
+            content = resp.get("response", resp.get("chunk", resp.get("text", "")))
             return RAGChunk(
-                content=resp.get("chunk", ""),
+                content=content,
                 end_of_stream=resp.get("end_of_stream", False),
-                error=resp.get("error")
+                error=None  # Errors are always thrown, never stored
             )
 
     def close(self) -> None:

trustgraph-base/trustgraph/api/types.py

@@ -79,5 +79,6 @@ class AgentAnswer(StreamingChunk):
 @dataclasses.dataclass
 class RAGChunk(StreamingChunk):
     """RAG streaming chunk"""
+    chunk_type: str = "rag"
     end_of_stream: bool = False
     error: Optional[Dict[str, str]] = None