ws -> sse for stability

Author: binary-husky

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_client.py

@@ -10,7 +10,10 @@
 from openai.types.chat.chat_completion import ChatCompletion
 
 import pickle
-import websockets
+import httpx
+import logging
+logging.getLogger("httpx").setLevel(logging.WARNING)
+
 import base64
 import json
 
@@ -96,7 +99,7 @@ def should_terminate(self) -> bool:
 
     def begin_service(self):
         """
-        Starts the websocket service loop.
+        Starts the SSE service loop.
         """
         t = threading.Thread(target=lambda: asyncio.run(self._ensure_service_loop()), daemon=True)
         t.start()
@@ -112,7 +115,7 @@ async def _ensure_service_loop(self):
     async def _service_loop(self):
         """
         In fact this is not a service,
-        it is a client that pretends to be a service, by interacting with a local websocket interchange server.
+        it is a client that pretends to be a service, by interacting with a local interchange server via SSE.
 
         This design is for efficiency
         """
@@ -121,43 +124,66 @@ async def _service_loop(self):
 
         port = os.getenv("AJET_DAT_INTERCHANGE_PORT")
         assert port is not None, "AJET_DAT_INTERCHANGE_PORT env var must be set"
-        uri = f"ws://127.0.0.1:{port}/hook/context_tracker_client_listen"
 
-        async with websockets.connect(uri, ping_timeout=3600, open_timeout=16) as websocket:
+        base_url = f"http://127.0.0.1:{port}"
+        listen_url = f"{base_url}/hook/context_tracker_client_listen"
+        response_url = f"{base_url}/hook/context_tracker_client_response"
+        key = f"episode_uuid:{self.episode_uuid}"
+
+        async with httpx.AsyncClient(timeout=None) as client:
+            try:
+                async with client.stream("GET", listen_url, params={"episode_uuid": self.episode_uuid}, timeout=None) as response:
+                    async for line in response.aiter_lines():
+                        if self.should_terminate:
+                            break
+
+                        if not line.strip():
+                            continue
+
+                        if line.startswith(":"): # keepalive
+                            continue
+
+                        if line.startswith("data: "):
+                            data = line[6:].strip()
+                            if not data:
+                                continue
+
+                            try:
+                                parsed_msg = InterchangeCompletionRequest.model_validate_json(data)
+
+                                result = await self.llm_infer(
+                                    req=parsed_msg.completion_request,
+                                    timeline_uuid=parsed_msg.timeline_uuid,
+                                    agent_name=parsed_msg.agent_name,
+                                    target_tag=parsed_msg.target_tag,
+                                    episode_uuid=parsed_msg.episode_uuid,
+                                )
+
+                                # Send response back
+                                await client.post(
+                                    response_url,
+                                    params={"key": key},
+                                    content=pickle.dumps(result),
+                                    headers={"Content-Type": "application/octet-stream"}
+                                )
+
+                            except Exception as e:
+                                logger.error(f"Error processing SSE event: {e}")
+                                continue
+
+            except httpx.RequestError as e:
+                 logger.warning(f"SSE connection error: {e}")
+                 raise # Let ensure_service_loop handle restart
+
+            # Send terminate signal if we are exiting cleanly
             try:
-                # Send initialization parameters
-                # Sending as a list [agent_name, target_tag, episode_uuid] to match "input (a,b,c)" structure
-                await websocket.send(pickle.dumps(f"episode_uuid:{self.episode_uuid}"))
-
-                while not self.should_terminate:
-
-                    try:
-                        # wait message from ajet/tuner_lib/weight_tuner/experimental/as_oai_model_server.py
-                        parsed_msg_str: str = pickle.loads(
-                            await asyncio.wait_for(websocket.recv(decode=False), timeout=0.25)
-                        )
-                        parsed_msg:InterchangeCompletionRequest = InterchangeCompletionRequest(**json.loads(parsed_msg_str))
-
-                        response = await self.llm_infer(
-                            req=parsed_msg.completion_request,
-                            timeline_uuid=parsed_msg.timeline_uuid,
-                            agent_name=parsed_msg.agent_name,
-                            target_tag=parsed_msg.target_tag,
-                            episode_uuid=parsed_msg.episode_uuid,
-                        )
-                        await websocket.send(pickle.dumps(response))
-
-                    except asyncio.TimeoutError:
-                        # 0.25s timeout, loop back to check should_terminate
-                        continue
-                    except websockets.exceptions.ConnectionClosed:
-                        logger.warning("Websocket connection closed by server")
-                        return # Exit inner loop to reconnect or finish
-
-                await websocket.send(pickle.dumps("terminate"))
-
-            except (OSError, IOError) as e:
-                logger.warning(f"Websocket connection error: {e}")
+                await client.post(
+                    response_url,
+                    params={"key": key},
+                    content=pickle.dumps("terminate"),
+                    headers={"Content-Type": "application/octet-stream"}
+                )
+            except:
                 pass