Schema changes for prompt service

trustgraph-base/trustgraph/messaging/translators/prompt.py

@@ -53,6 +53,13 @@ class PromptResponseTranslator(MessageTranslator):
         # Always include end_of_stream flag for streaming support
         result["end_of_stream"] = getattr(obj, "end_of_stream", False)
 
+        if obj.in_token:
+            result["in_token"] = obj.in_token
+        if obj.out_token:
+            result["out_token"] = obj.out_token
+        if obj.model:
+            result["model"] = obj.model
+
         return result
     
     def encode_with_completion(self, obj: PromptResponse) -> Tuple[Dict[str, Any], bool]:

trustgraph-base/trustgraph/schema/services/prompt.py

@@ -41,4 +41,9 @@ class PromptResponse:
     # Indicates final message in stream
     end_of_stream: bool = False
 
+    # Token usage from the underlying text completion
+    in_token: int = 0
+    out_token: int = 0
+    model: str = ""
+
 ############################################################################