From 66170834f89b8783d765c02c1bcfd9ea407b8e4b Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Tue, 25 Nov 2025 23:51:52 +0000
Subject: [PATCH] Fixing some streamingg

---
 .../trustgraph/base/prompt_client.py          | 67 ++++++++++++++-----
 1 file changed, 52 insertions(+), 15 deletions(-)

diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py
index 568a25a3..ccdf8e9c 100644
--- a/trustgraph-base/trustgraph/base/prompt_client.py
+++ b/trustgraph-base/trustgraph/base/prompt_client.py
@@ -8,24 +8,61 @@ class PromptClient(RequestResponse):
 
     async def prompt(self, id, variables, timeout=600, streaming=False):
 
-        resp = await self.request(
-            PromptRequest(
-                id = id,
-                terms = {
-                    k: json.dumps(v)
-                    for k, v in variables.items()
-                },
-                streaming = streaming
-            ),
-            timeout=timeout
-        )
+        if not streaming:
+            # Non-streaming path
+            resp = await self.request(
+                PromptRequest(
+                    id = id,
+                    terms = {
+                        k: json.dumps(v)
+                        for k, v in variables.items()
+                    },
+                    streaming = False
+                ),
+                timeout=timeout
+            )
 
-        if resp.error:
-            raise RuntimeError(resp.error.message)
+            if resp.error:
+                raise RuntimeError(resp.error.message)
 
-        if resp.text: return resp.text
+            if resp.text: return resp.text
 
-        return json.loads(resp.object)
+            return json.loads(resp.object)
+
+        else:
+            # Streaming path - collect all chunks
+            full_text = ""
+            full_object = None
+
+            async def collect_chunks(resp):
+                nonlocal full_text, full_object
+
+                if resp.error:
+                    raise RuntimeError(resp.error.message)
+
+                if resp.text:
+                    full_text += resp.text
+                elif resp.object:
+                    full_object = resp.object
+
+                return getattr(resp, 'end_of_stream', False)
+
+            await self.request(
+                PromptRequest(
+                    id = id,
+                    terms = {
+                        k: json.dumps(v)
+                        for k, v in variables.items()
+                    },
+                    streaming = True
+                ),
+                recipient=collect_chunks,
+                timeout=timeout
+            )
+
+            if full_text: return full_text
+
+            return json.loads(full_object)
 
     async def extract_definitions(self, text, timeout=600):
         return await self.prompt(