release/v2.2 -> master (#733)

trustgraph-flow/trustgraph/gateway/dispatch/mux.py

@@ -33,9 +33,12 @@ class Mux:
 
     async def receive(self, msg):
 
+        request_id = None
+
         try:
 
             data = msg.json()
+            request_id = data.get("id")
 
             if "request" not in data:
                 raise RuntimeError("Bad message")
@@ -51,7 +54,13 @@ class Mux:
 
         except Exception as e:
             logger.error(f"Receive exception: {str(e)}", exc_info=True)
-            await self.ws.send_json({"error": str(e)})
+            error_resp = {
+                "error": {"message": str(e), "type": "error"},
+                "complete": True,
+            }
+            if request_id:
+                error_resp["id"] = request_id
+            await self.ws.send_json(error_resp)
 
     async def maybe_tidy_workers(self, workers):
 
@@ -97,12 +106,12 @@ class Mux:
             })
 
         worker = asyncio.create_task(
-            self.request_task(request, responder, flow, svc)
+            self.request_task(id, request, responder, flow, svc)
         )
 
         workers.append(worker)
 
-    async def request_task(self, request, responder, flow, svc):
+    async def request_task(self, id, request, responder, flow, svc):
 
         try:
 
@@ -119,7 +128,11 @@ class Mux:
                 )
 
         except Exception as e:
-            await self.ws.send_json({"error": str(e)})
+            await self.ws.send_json({
+                "id": id,
+                "error": {"message": str(e), "type": "error"},
+                "complete": True,
+            })
 
     async def run(self):
 
@@ -143,7 +156,11 @@ class Mux:
             except Exception as e:
                 # This is an internal working error, may not be recoverable
                 logger.error(f"Run prepare exception: {e}", exc_info=True)
-                await self.ws.send_json({"id": id, "error": str(e)})
+                await self.ws.send_json({
+                    "id": id,
+                    "error": {"message": str(e), "type": "error"},
+                    "complete": True,
+                })
                 self.running.stop()
 
                 if self.ws:
@@ -160,7 +177,11 @@ class Mux:
 
             except Exception as e:
                 logger.error(f"Exception in mux: {e}", exc_info=True)
-                await self.ws.send_json({"error": str(e)})
+                await self.ws.send_json({
+                    "id": id,
+                    "error": {"message": str(e), "type": "error"},
+                    "complete": True,
+                })
 
         self.running.stop()
 

trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py

@@ -20,7 +20,7 @@ default_ident = "text-completion"
 
 default_temperature = 0.0
 default_max_output = 4192
-default_api = "2024-12-01-preview"
+default_api = os.getenv("AZURE_API_VERSION", "2024-12-01-preview")
 default_endpoint = os.getenv("AZURE_ENDPOINT", None)
 default_token = os.getenv("AZURE_TOKEN", None)
 default_model = os.getenv("AZURE_MODEL", None)
@@ -90,7 +90,7 @@ class Processor(LlmService):
                     }
                 ],
                 temperature=effective_temperature,
-                max_tokens=self.max_output,
+                max_completion_tokens=self.max_output,
                 top_p=1,
             )
 
@@ -159,7 +159,7 @@ class Processor(LlmService):
                     }
                 ],
                 temperature=effective_temperature,
-                max_tokens=self.max_output,
+                max_completion_tokens=self.max_output,
                 top_p=1,
                 stream=True,
                 stream_options={"include_usage": True}

trustgraph-flow/trustgraph/model/text_completion/openai/llm.py

@@ -86,7 +86,7 @@ class Processor(LlmService):
                     }
                 ],
                 temperature=effective_temperature,
-                max_tokens=self.max_output,
+                max_completion_tokens=self.max_output,
             )
             
             inputtokens = resp.usage.prompt_tokens
@@ -152,7 +152,7 @@ class Processor(LlmService):
                     }
                 ],
                 temperature=effective_temperature,
-                max_tokens=self.max_output,
+                max_completion_tokens=self.max_output,
                 stream=True,
                 stream_options={"include_usage": True}
             )