Update gen.py

gen.py

@@ -789,6 +789,10 @@ async def generate_text(
         
                     yield chunk
 
+        holder = request.scope.get("_stream_holder")
+        if holder is not None:
+            holder["generator"] = event_generator
+
         return StreamingResponse(
             event_generator(),
             media_type="text/event-stream",
@@ -945,15 +949,7 @@ def return_models_openai():
 
 import time
 import uuid
-
 def _normalize_responses_input(input_field) -> list[dict]:
-    """
-    Coerce the Responses API `input` field into a standard messages[] list.
-
-    Accepted shapes:
-      • str                          → [{"role":"user","content":"..."}]
-      • list of message-like dicts   → pass through, normalising content parts
-    """
     if isinstance(input_field, str):
         return [{"role": "user", "content": input_field}]
 
@@ -982,9 +978,6 @@ def _normalize_responses_input(input_field) -> list[dict]:
 
 
 def _wrap_responses_output(chat_payload: dict, model_name: str) -> dict:
-    """
-    Wrap a standard chat-completions JSON response into the Responses API shape.
-    """
     choices = chat_payload.get("choices", [])
     output = []
 
@@ -1033,10 +1026,6 @@ def _wrap_responses_output(chat_payload: dict, model_name: str) -> dict:
 def _wrap_responses_stream_chunk(
     line: str, response_id: str, model_name: str, sent_created: bool
 ) -> tuple[str, bool, bool]:
-    """
-    Returns (translated_sse_string, sent_created, is_done).
-    is_done=True means the upstream [DONE] was processed and response.completed was emitted.
-    """
     if not line.startswith("data:"):
         return line + "\n", sent_created, False
 
@@ -1141,22 +1130,6 @@ async def create_response(
     authorization: Optional[str] = Header(None),
     x_client_id: Optional[str] = Header(None),
 ):
-    """
-    OpenAI Responses API-compatible endpoint.
-
-    Accepts the Responses API request shape, normalises it into the chat
-    completions format, routes it through the existing generate_text logic,
-    and wraps the result back into the Responses API shape.
-
-    Supported fields:
-      • input       (str | list)  — required
-      • model       (str)         — accepted but ignored (router decides)
-      • stream      (bool)        — optional, default False
-      • tools       (list)        — optional, forwarded as-is
-      • tool_choice (str|dict)    — optional, forwarded as-is
-      • temperature (float)       — optional, forwarded
-      • max_output_tokens (int)   — mapped to max_tokens
-    """
     body = await request.json()
 
     input_field = body.get("input")
@@ -1178,8 +1151,13 @@ async def create_response(
 
     raw_body = json.dumps(chat_body).encode()
 
+    # shared dict passed by reference through scope — generate_text writes
+    # the generator factory here, create_response reads it back after the await
+    stream_holder = {}
+
     scope = dict(request.scope)
     scope["path"] = "/gen/chat/completions"
+    scope["_stream_holder"] = stream_holder
     scope["headers"] = [
         (k, v) for k, v in request.scope["headers"]
         if k.lower() not in (b"content-length",)
@@ -1192,37 +1170,40 @@ async def create_response(
 
     response = await generate_text(sub_request, authorization, x_client_id)
 
+    # --- streaming path ---
     if chat_body.get("stream"):
         response_id = f"resp_{uuid.uuid4().hex[:24]}"
         model_label = MODEL_MAP.get("lightning", "lightning")
 
+        raw_generator = stream_holder.get("generator")
+        if raw_generator is None:
+            raise HTTPException(500, "Stream generator not captured")
+
         async def responses_stream():
             sent_created = False
             completed = False
             buffer = ""
-        
-            async for chunk in response.body_iterator:
+
+            async for chunk in raw_generator():
                 if isinstance(chunk, bytes):
                     chunk = chunk.decode("utf-8", errors="replace")
                 buffer += chunk
-        
+
                 while "\n" in buffer:
                     line, buffer = buffer.split("\n", 1)
                     line = line.rstrip("\r")
                     if not line:
                         continue
-        
+
                     translated, new_sent, is_done = _wrap_responses_stream_chunk(
                         line, response_id, model_label, sent_created
                     )
                     sent_created = new_sent
-        
                     if is_done:
                         completed = True
-        
                     if translated:
                         yield translated
-        
+
             if buffer.strip():
                 translated, new_sent, is_done = _wrap_responses_stream_chunk(
                     buffer.strip(), response_id, model_label, sent_created
@@ -1232,11 +1213,8 @@ async def create_response(
                     completed = True
                 if translated:
                     yield translated
-        
-            print(f"[RESPONSES STREAM] generator exhausted. sent_created={sent_created} completed={completed}")
-        
+
             if sent_created and not completed:
-                print("[RESPONSES STREAM] guard firing — upstream closed without [DONE]")
                 done_event = json.dumps({
                     "type": "response.completed",
                     "response": {
@@ -1260,6 +1238,7 @@ async def create_response(
             },
         )
 
+    # --- non-streaming path ---
     if hasattr(response, "body"):
         raw = response.body
     else: