pygpt-net 2.6.28__py3-none-any.whl → 2.6.30__py3-none-any.whl

pygpt_net/controller/chat/stream.py

@@ -6,12 +6,14 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.08.19 07:00:00                  #
+# Updated Date: 2025.08.28 20:00:00                  #
 # ================================================== #
 
 import base64
 import io
-from typing import Optional, Literal
+import json
+from dataclasses import dataclass, field
+from typing import Optional, Literal, Any
 
 from PySide6.QtCore import QObject, Signal, Slot, QRunnable
 
@@ -44,9 +46,11 @@ ChunkType = Literal[
     "api_completion",
     "langchain_chat",
     "llama_chat",
+    "google",
     "raw",
 ]
 
+
 class WorkerSignals(QObject):
     """
     Defines the signals available from a running worker thread.
@@ -59,6 +63,31 @@ class WorkerSignals(QObject):
     eventReady = Signal(object)
 
 
+@dataclass
+class WorkerState:
+    """Holds mutable state for the streaming loop."""
+    output_parts: list[str] = field(default_factory=list)
+    output_tokens: int = 0
+    begin: bool = True
+    error: Optional[Exception] = None
+    fn_args_buffers: dict[str, io.StringIO] = field(default_factory=dict)
+    citations: Optional[list] = field(default_factory=list)
+    image_paths: list[str] = field(default_factory=list)
+    files: list[dict] = field(default_factory=list)
+    img_path: Optional[str] = None
+    is_image: bool = False
+    has_google_inline_image: bool = False
+    is_code: bool = False
+    force_func_call: bool = False
+    stopped: bool = False
+    chunk_type: ChunkType = "raw"
+    generator: Any = None
+    usage_vendor: Optional[str] = None
+    usage_payload: dict = field(default_factory=dict)
+    google_stream_ref: Any = None
+    tool_calls: list[dict] = field(default_factory=list)
+
+
 class StreamWorker(QRunnable):
     def __init__(self, ctx: CtxItem, window, parent=None):
         super().__init__()
@@ -78,336 +107,974 @@ class StreamWorker(QRunnable):
         emit_error = self.signals.errorOccurred.emit
         emit_end = self.signals.end.emit
 
-        output_parts = []
-        output_tokens = 0
-        begin = True
-        error = None
-        fn_args_buffers: dict[str, io.StringIO] = {}
-        citations: Optional[list] = []
-        files = []
-        img_path = core.image.gen_unique_path(ctx)
-        is_image = False
-        is_code = False
-        force_func_call = False
-        stopped = False
-        chunk_type: ChunkType = "raw"
-        generator = self.stream
-
-        base_data = {
-            "meta": ctx.meta,
-            "ctx": ctx,
-        }
+        state = WorkerState()
+        state.generator = self.stream
+        state.img_path = core.image.gen_unique_path(ctx)
+
+        base_data = {"meta": ctx.meta, "ctx": ctx}
         emit_event(RenderEvent(RenderEvent.STREAM_BEGIN, base_data))
 
-        tool_calls = []
         try:
-            if generator is not None:
-                for chunk in generator:
-                    if ctrl.kernel.stopped():
-                        if hasattr(generator, 'close'):
-                            try:
-                                generator.close()
-                            except Exception:
-                                pass
-                        elif hasattr(generator, 'cancel'):
-                            try:
-                                generator.cancel()
-                            except Exception:
-                                pass
-                        elif hasattr(generator, 'stop'):
-                            try:
-                                generator.stop()
-                            except Exception:
-                                pass
-                        ctx.msg_id = None
-                        stopped = True
+            if state.generator is not None:
+                for chunk in state.generator:
+                    # cooperative stop
+                    if self._should_stop(ctrl, state, ctx):
                         break
 
-                    if error is not None:
+                    # if error flagged, stop early
+                    if state.error is not None:
                         ctx.msg_id = None
-                        stopped = True
+                        state.stopped = True
                         break
 
                     etype: Optional[EventType] = None
-                    response = None
 
+                    # detect chunk type
                     if ctx.use_responses_api:
                         if hasattr(chunk, 'type'):
                             etype = chunk.type  # type: ignore[assignment]
-                            chunk_type = "api_chat_responses"
+                            state.chunk_type = "api_chat_responses"
                         else:
                             continue
                     else:
-                        if (hasattr(chunk, 'choices')
-                                and chunk.choices
-                                and hasattr(chunk.choices[0], 'delta')
-                                and chunk.choices[0].delta is not None):
-                            chunk_type = "api_chat"
-                        elif (hasattr(chunk, 'choices')
-                              and chunk.choices
-                              and hasattr(chunk.choices[0], 'text')
-                              and chunk.choices[0].text is not None):
-                            chunk_type = "api_completion"
-                        elif hasattr(chunk, 'content') and chunk.content is not None:
-                            chunk_type = "langchain_chat"
-                        elif hasattr(chunk, 'delta') and chunk.delta is not None:
-                            chunk_type = "llama_chat"
-                        else:
-                            chunk_type = "raw"
-
-                    if chunk_type == "api_chat":
-                        citations = None
-                        delta = chunk.choices[0].delta
-                        if delta and delta.content is not None:
-                            if citations is None and hasattr(chunk, 'citations') and chunk.citations is not None:
-                                citations = chunk.citations
-                                ctx.urls = citations
-                            response = delta.content
-
-                        if delta and delta.tool_calls:
-                            for tool_chunk in delta.tool_calls:
-                                if tool_chunk.index is None:
-                                    tool_chunk.index = 0
-                                if len(tool_calls) <= tool_chunk.index:
-                                    tool_calls.append(
-                                        {
-                                            "id": "",
-                                            "type": "function",
-                                            "function": {"name": "", "arguments": ""}
-                                        }
-                                    )
-                                tool_call = tool_calls[tool_chunk.index]
-                                if getattr(tool_chunk, "id", None):
-                                    tool_call["id"] += tool_chunk.id
-                                if getattr(tool_chunk.function, "name", None):
-                                    tool_call["function"]["name"] += tool_chunk.function.name
-                                if getattr(tool_chunk.function, "arguments", None):
-                                    tool_call["function"]["arguments"] += tool_chunk.function.arguments
-
-                    elif chunk_type == "api_chat_responses":
-                        if etype == "response.completed":
-                            for item in chunk.response.output:
-                                if item.type == "mcp_list_tools":
-                                    core.gpt.responses.mcp_tools = item.tools
-                                elif item.type == "mcp_call":
-                                    call = {
-                                        "id": item.id,
-                                        "type": "mcp_call",
-                                        "approval_request_id": item.approval_request_id,
-                                        "arguments": item.arguments,
-                                        "error": item.error,
-                                        "name": item.name,
-                                        "output": item.output,
-                                        "server_label": item.server_label,
-                                    }
-                                    tool_calls.append({
-                                        "id": item.id,
-                                        "call_id": "",
-                                        "type": "function",
-                                        "function": {"name": item.name, "arguments": item.arguments}
-                                    })
-                                    ctx.extra["mcp_call"] = call
-                                    core.ctx.update_item(ctx)
-                                elif item.type == "mcp_approval_request":
-                                    call = {
-                                        "id": item.id,
-                                        "type": "mcp_call",
-                                        "arguments": item.arguments,
-                                        "name": item.name,
-                                        "server_label": item.server_label,
-                                    }
-                                    ctx.extra["mcp_approval_request"] = call
-                                    core.ctx.update_item(ctx)
-
-                        elif etype == "response.output_text.delta":
-                            response = chunk.delta
-
-                        elif etype == "response.output_item.added" and chunk.item.type == "function_call":
-                            tool_calls.append({
-                                "id": chunk.item.id,
-                                "call_id": chunk.item.call_id,
-                                "type": "function",
-                                "function": {"name": chunk.item.name, "arguments": ""}
-                            })
-                            fn_args_buffers[chunk.item.id] = io.StringIO()
-                        elif etype == "response.function_call_arguments.delta":
-                            buf = fn_args_buffers.get(chunk.item_id)
-                            if buf is not None:
-                                buf.write(chunk.delta)
-                        elif etype == "response.function_call_arguments.done":
-                            buf = fn_args_buffers.pop(chunk.item_id, None)
-                            if buf is not None:
-                                try:
-                                    args_val = buf.getvalue()
-                                finally:
-                                    buf.close()
-                                for tc in tool_calls:
-                                    if tc["id"] == chunk.item_id:
-                                        tc["function"]["arguments"] = args_val
-                                        break
-
-                        elif etype == "response.output_text.annotation.added":
-                            ann = chunk.annotation
-                            if ann['type'] == "url_citation":
-                                if citations is None:
-                                    citations = []
-                                url_citation = ann['url']
-                                citations.append(url_citation)
-                                ctx.urls = citations
-                            elif ann['type'] == "container_file_citation":
-                                files.append({
-                                    "container_id": ann['container_id'],
-                                    "file_id": ann['file_id'],
-                                })
-
-                        elif etype == "response.reasoning_summary_text.delta":
-                            response = chunk.delta
-
-                        elif etype == "response.output_item.done":
-                            tool_calls, has_calls = core.gpt.computer.handle_stream_chunk(ctx, chunk, tool_calls)
-                            if has_calls:
-                                force_func_call = True
-
-                        elif etype == "response.code_interpreter_call_code.delta":
-                            if not is_code:
-                                response = "\n\n**Code interpreter**\n```python\n" + chunk.delta
-                                is_code = True
-                            else:
-                                response = chunk.delta
-                        elif etype == "response.code_interpreter_call_code.done":
-                            response = "\n\n```\n-----------\n"
-
-                        elif etype == "response.image_generation_call.partial_image":
-                            image_base64 = chunk.partial_image_b64
-                            image_bytes = base64.b64decode(image_base64)
-                            with open(img_path, "wb") as f:
-                                f.write(image_bytes)
-                            del image_bytes
-                            is_image = True
-
-                        elif etype == "response.created":
-                            ctx.msg_id = str(chunk.response.id)
-                            core.ctx.update_item(ctx)
-
-                        elif etype in {"response.done", "response.failed", "error"}:
-                            pass
-
-                    elif chunk_type == "api_completion":
-                        choice0 = chunk.choices[0]
-                        if choice0.text is not None:
-                            response = choice0.text
-
-                    elif chunk_type == "langchain_chat":
-                        if chunk.content is not None:
-                            response = str(chunk.content)
-
-                    elif chunk_type == "llama_chat":
-                        if chunk.delta is not None:
-                            response = str(chunk.delta)
-                        tool_chunks = getattr(chunk.message, "additional_kwargs", {}).get("tool_calls", [])
-                        if tool_chunks:
-                            for tool_chunk in tool_chunks:
-                                id_val = getattr(tool_chunk, "call_id", None) or getattr(tool_chunk, "id", None)
-                                name = getattr(tool_chunk, "name", None) or getattr(getattr(tool_chunk, "function", None), "name", None)
-                                args = getattr(tool_chunk, "arguments", None)
-                                if args is None:
-                                    f = getattr(tool_chunk, "function", None)
-                                    args = getattr(f, "arguments", None) if f else None
-                                if id_val:
-                                    if not args:
-                                        args = "{}"
-                                    tool_call = {
-                                        "id": id_val,
-                                        "type": "function",
-                                        "function": {"name": name, "arguments": args}
-                                    }
-                                    tool_calls.clear()
-                                    tool_calls.append(tool_call)
+                        state.chunk_type = self._detect_chunk_type(chunk)
 
-                    else:
-                        if chunk is not None:
-                            response = chunk if isinstance(chunk, str) else str(chunk)
+                    # process chunk according to type
+                    response = self._process_chunk(ctx, core, state, chunk, etype)
 
-                    if response is not None and response != "" and not stopped:
-                        if begin and response == "":
-                            continue
-                        output_parts.append(response)
-                        output_tokens += 1
-                        emit_event(
-                            RenderEvent(
-                                RenderEvent.STREAM_APPEND,
-                                {
-                                    "meta": ctx.meta,
-                                    "ctx": ctx,
-                                    "chunk": response,
-                                    "begin": begin,
-                                },
-                            )
-                        )
-                        begin = False
+                    # emit response delta if present
+                    if response is not None and response != "" and not state.stopped:
+                        self._append_response(ctx, state, response, emit_event)
 
+                    # free per-iteration ref
                     chunk = None
 
-                if tool_calls:
-                    ctx.force_call = force_func_call
-                    core.debug.info("[chat] Tool calls found, unpacking...")
-                    core.command.unpack_tool_calls_chunks(ctx, tool_calls)
-
-                if is_image:
-                    core.debug.info("[chat] Image generation call found")
-                    ctx.images = [img_path]
+                # after loop: handle tool-calls and images assembly
+                self._handle_after_loop(ctx, core, state)
 
         except Exception as e:
-            error = e
+            state.error = e
 
         finally:
-            output = "".join(output_parts)
-            output_parts.clear()
-            del output_parts
+            self._finalize(ctx, core, state, emit_end, emit_error)
 
-            if has_unclosed_code_tag(output):
-                output += "\n```"
+    # ------------ Orchestration helpers ------------
 
-            if generator and hasattr(generator, 'close'):
-                try:
-                    generator.close()
-                except Exception:
-                    pass
+    def _should_stop(self, ctrl, state: WorkerState, ctx: CtxItem) -> bool:
+        """
+        Checks external stop signal and attempts to stop the generator gracefully.
 
-            del generator
-            self.stream = None
+        :param ctrl: Controller with stop signal
+        :param state: WorkerState
+        :param ctx: CtxItem
+        :return: True if stopped, False otherwise
+        """
+        if not ctrl.kernel.stopped():
+            return False
+
+        gen = state.generator
+        if gen is not None:
+            # Try common stop methods without raising
+            for meth in ("close", "cancel", "stop"):
+                if hasattr(gen, meth):
+                    try:
+                        getattr(gen, meth)()
+                    except Exception:
+                        pass
+
+        ctx.msg_id = None
+        state.stopped = True
+        return True
+
+    def _detect_chunk_type(self, chunk) -> ChunkType:
+        """
+        Detects chunk type for various providers/SDKs.
 
-            ctx.output = output
-            ctx.set_tokens(ctx.input_tokens, output_tokens)
-            core.ctx.update_item(ctx)
+        :param chunk: The chunk object from the stream
+        :return: Detected ChunkType
+        """
+        if (hasattr(chunk, 'choices')
+                and chunk.choices
+                and hasattr(chunk.choices[0], 'delta')
+                and chunk.choices[0].delta is not None):
+            return "api_chat"
+        if (hasattr(chunk, 'choices')
+                and chunk.choices
+                and hasattr(chunk.choices[0], 'text')
+                and chunk.choices[0].text is not None):
+            return "api_completion"
+        if hasattr(chunk, 'content') and chunk.content is not None:
+            return "langchain_chat"
+        if hasattr(chunk, 'delta') and chunk.delta is not None:
+            return "llama_chat"
+        if hasattr(chunk, "candidates"):  # Google python-genai chunk
+            return "google"
+        return "raw"
+
+    def _append_response(
+            self,
+            ctx: CtxItem,
+            state: WorkerState,
+            response: str,
+            emit_event
+    ):
+        """
+        Appends response delta and emits STREAM_APPEND event.
+
+        Skips empty initial chunks if state.begin is True.
+
+        :param ctx: CtxItem
+        :param state: WorkerState
+        :param response: Response delta string
+        :param emit_event: Function to emit RenderEvent
+        """
+        if state.begin and response == "":
+            return
+        state.output_parts.append(response)
+        state.output_tokens += 1
+        emit_event(
+            RenderEvent(
+                RenderEvent.STREAM_APPEND,
+                {
+                    "meta": ctx.meta,
+                    "ctx": ctx,
+                    "chunk": response,
+                    "begin": state.begin,
+                },
+            )
+        )
+        state.begin = False
+
+    def _handle_after_loop(self, ctx: CtxItem, core, state: WorkerState):
+        """
+        Post-loop handling for tool calls and images assembly.
+
+        :param ctx: CtxItem
+        :param core: Core instance
+        :param state: WorkerState
+        """
+        if state.tool_calls:
+            ctx.force_call = state.force_func_call
+            core.debug.info("[chat] Tool calls found, unpacking...")
+            # Ensure function.arguments is JSON string
+            for tc in state.tool_calls:
+                fn = tc.get("function") or {}
+                if isinstance(fn.get("arguments"), dict):
+                    fn["arguments"] = json.dumps(fn["arguments"], ensure_ascii=False)
+            core.command.unpack_tool_calls_chunks(ctx, state.tool_calls)
+
+        # OpenAI partial image assembly
+        if state.is_image and state.img_path:
+            core.debug.info("[chat] OpenAI partial image assembled")
+            ctx.images = [state.img_path]
+
+        # Google inline images
+        if state.image_paths:
+            core.debug.info("[chat] Google inline images found")
+            if not isinstance(ctx.images, list) or not ctx.images:
+                ctx.images = list(state.image_paths)
+            else:
+                seen = set(ctx.images)
+                for p in state.image_paths:
+                    if p not in seen:
+                        ctx.images.append(p)
+                        seen.add(p)
+
+    def _finalize(self, ctx: CtxItem, core, state: WorkerState, emit_end, emit_error):
+        """
+        Finalize stream: build output, usage, tokens, files, errors, cleanup.
+
+        :param ctx: CtxItem
+        :param core: Core instance
+        :param state: WorkerState
+        :param emit_end: Function to emit end signal
+        """
+        # Build final output
+        output = "".join(state.output_parts)
+        state.output_parts.clear()
+
+        if has_unclosed_code_tag(output):
+            output += "\n```"
+
+        # Attempt to resolve Google usage from the stream object if missing
+        if (state.usage_vendor is None or state.usage_vendor == "google") and not state.usage_payload and state.generator is not None:
+            try:
+                if hasattr(state.generator, "resolve"):
+                    state.generator.resolve()
+                    um = getattr(state.generator, "usage_metadata", None)
+                    if um:
+                        self._capture_google_usage(state, um)
+            except Exception:
+                pass
+
+        # Close generator if possible
+        gen = state.generator
+        if gen and hasattr(gen, 'close'):
+            try:
+                gen.close()
+            except Exception:
+                pass
+
+        self.stream = None
+        ctx.output = output
+
+        # Tokens usage
+        if state.usage_payload:
+            in_tok_final = state.usage_payload.get("in")
+            out_tok_final = state.usage_payload.get("out")
+
+            if in_tok_final is None:
+                in_tok_final = ctx.input_tokens if ctx.input_tokens is not None else 0
+            if out_tok_final is None:
+                out_tok_final = state.output_tokens
+
+            ctx.set_tokens(in_tok_final, out_tok_final)
+
+            # Attach usage details in ctx.extra for debugging
+            try:
+                if not isinstance(ctx.extra, dict):
+                    ctx.extra = {}
+                ctx.extra["usage"] = {
+                    "vendor": state.usage_vendor,
+                    "input_tokens": in_tok_final,
+                    "output_tokens": out_tok_final,
+                    "reasoning_tokens": state.usage_payload.get("reasoning", 0),
+                    "total_reported": state.usage_payload.get("total"),
+                }
+            except Exception:
+                pass
+        else:
+            # Fallback when usage is not available
+            ctx.set_tokens(ctx.input_tokens if ctx.input_tokens is not None else 0, state.output_tokens)
+
+        core.ctx.update_item(ctx)
+
+        # OpenAI only: download container files if present
+        if state.files and not state.stopped:
+            core.debug.info("[chat] Container files found, downloading...")
+            try:
+                core.api.openai.container.download_files(ctx, state.files)
+            except Exception as e:
+                core.debug.error(f"[chat] Error downloading container files: {e}")
+
+        # Emit error and end
+        if state.error:
+            emit_error(state.error)
+        emit_end(ctx)
+
+        # Cleanup local buffers
+        for _buf in state.fn_args_buffers.values():
+            try:
+                _buf.close()
+            except Exception:
+                pass
+        state.fn_args_buffers.clear()
+        state.files.clear()
+        state.tool_calls.clear()
+        if state.citations is not None and state.citations is not ctx.urls:
+            state.citations.clear()
+        state.citations = None
+
+        # Worker cleanup (signals etc.)
+        self.cleanup()
+
+    # ------------ Chunk processors ------------
+
+    def _process_chunk(
+            self,
+            ctx: CtxItem,
+            core,
+            state: WorkerState,
+            chunk,
+            etype: Optional[EventType]
+    ) -> Optional[str]:
+        """
+        Dispatches processing to concrete provider-specific processing.
+
+        :param ctx: CtxItem
+        :param core: Core instance
+        :param state: WorkerState
+        :param chunk: The chunk object from the stream
+        :param etype: Optional event type for Responses API
+        :return: Response delta string or None
+        """
+        t = state.chunk_type
+        if t == "api_chat":
+            return self._process_api_chat(ctx, state, chunk)
+        if t == "api_chat_responses":
+            return self._process_api_chat_responses(ctx, core, state, chunk, etype)
+        if t == "api_completion":
+            return self._process_api_completion(chunk)
+        if t == "langchain_chat":
+            return self._process_langchain_chat(chunk)
+        if t == "llama_chat":
+            return self._process_llama_chat(state, chunk)
+        if t == "google":
+            return self._process_google_chunk(ctx, core, state, chunk)
+        # raw fallback
+        return self._process_raw(chunk)
+
+    def _process_api_chat(
+            self,
+            ctx: CtxItem,
+            state: WorkerState,
+            chunk
+    ) -> Optional[str]:
+        """
+        OpenAI Chat Completions stream delta.
+
+        Handles text deltas, citations, and streamed tool_calls.
+
+        :param ctx: CtxItem
+        :param state: WorkerState
+        :param chunk: The chunk object from the stream
+        :return: Response delta string or None
+        """
+        response = None
+        state.citations = None  # as in original, reset to None for this type
+
+        delta = chunk.choices[0].delta if getattr(chunk, "choices", None) else None
+        if delta and getattr(delta, "content", None) is not None:
+            if state.citations is None and hasattr(chunk, 'citations') and chunk.citations is not None:
+                state.citations = chunk.citations
+                ctx.urls = state.citations
+            response = delta.content
+
+        # Accumulate streamed tool_calls
+        if delta and getattr(delta, "tool_calls", None):
+            for tool_chunk in delta.tool_calls:
+                if tool_chunk.index is None:
+                    tool_chunk.index = 0
+                if len(state.tool_calls) <= tool_chunk.index:
+                    state.tool_calls.append(
+                        {
+                            "id": "",
+                            "type": "function",
+                            "function": {"name": "", "arguments": ""}
+                        }
+                    )
+                tool_call = state.tool_calls[tool_chunk.index]
+                if getattr(tool_chunk, "id", None):
+                    tool_call["id"] += tool_chunk.id
+                if getattr(getattr(tool_chunk, "function", None), "name", None):
+                    tool_call["function"]["name"] += tool_chunk.function.name
+                if getattr(getattr(tool_chunk, "function", None), "arguments", None):
+                    tool_call["function"]["arguments"] += tool_chunk.function.arguments
+
+        # Capture usage (if available on final chunk with include_usage=True)
+        try:
+            u = getattr(chunk, "usage", None)
+            if u:
+                self._capture_openai_usage(state, u)
+        except Exception:
+            pass
+
+        return response
+
+    def _process_api_chat_responses(
+            self,
+            ctx: CtxItem,
+            core,
+            state: WorkerState,
+            chunk,
+            etype: Optional[EventType]
+    ) -> Optional[str]:
+        """
+        OpenAI Responses API stream events
+
+        Handles various event types including text deltas, tool calls, citations, images, and usage.
 
-            output = None
+        :param ctx: CtxItem
+        :param core: Core instance
+        :param state: WorkerState
+        :param chunk: The chunk object from the stream
+        :param etype: EventType string
+        :return: Response delta string or None
+        """
+        response = None
 
-            if files and not stopped:
-                core.debug.info("[chat] Container files found, downloading...")
+        if etype == "response.completed":
+            # usage on final response
+            try:
+                u = getattr(chunk.response, "usage", None)
+                if u:
+                    self._capture_openai_usage(state, u)
+            except Exception:
+                pass
+
+            for item in chunk.response.output:
+                if item.type == "mcp_list_tools":
+                    core.api.openai.responses.mcp_tools = item.tools
+                elif item.type == "mcp_call":
+                    call = {
+                        "id": item.id,
+                        "type": "mcp_call",
+                        "approval_request_id": item.approval_request_id,
+                        "arguments": item.arguments,
+                        "error": item.error,
+                        "name": item.name,
+                        "output": item.output,
+                        "server_label": item.server_label,
+                    }
+                    state.tool_calls.append({
+                        "id": item.id,
+                        "call_id": "",
+                        "type": "function",
+                        "function": {"name": item.name, "arguments": item.arguments}
+                    })
+                    ctx.extra["mcp_call"] = call
+                    core.ctx.update_item(ctx)
+                elif item.type == "mcp_approval_request":
+                    call = {
+                        "id": item.id,
+                        "type": "mcp_call",
+                        "arguments": item.arguments,
+                        "name": item.name,
+                        "server_label": item.server_label,
+                    }
+                    ctx.extra["mcp_approval_request"] = call
+                    core.ctx.update_item(ctx)
+
+        elif etype == "response.output_text.delta":
+            response = chunk.delta
+
+        elif etype == "response.output_item.added" and chunk.item.type == "function_call":
+            state.tool_calls.append({
+                "id": chunk.item.id,
+                "call_id": chunk.item.call_id,
+                "type": "function",
+                "function": {"name": chunk.item.name, "arguments": ""}
+            })
+            state.fn_args_buffers[chunk.item.id] = io.StringIO()
+
+        elif etype == "response.function_call_arguments.delta":
+            buf = state.fn_args_buffers.get(chunk.item_id)
+            if buf is not None:
+                buf.write(chunk.delta)
+
+        elif etype == "response.function_call_arguments.done":
+            buf = state.fn_args_buffers.pop(chunk.item_id, None)
+            if buf is not None:
                 try:
-                    core.gpt.container.download_files(ctx, files)
-                except Exception as e:
-                    core.debug.error(f"[chat] Error downloading container files: {e}")
+                    args_val = buf.getvalue()
+                finally:
+                    buf.close()
+                for tc in state.tool_calls:
+                    if tc["id"] == chunk.item_id:
+                        tc["function"]["arguments"] = args_val
+                        break
+
+        elif etype == "response.output_text.annotation.added":
+            ann = chunk.annotation
+            if ann['type'] == "url_citation":
+                if state.citations is None:
+                    state.citations = []
+                url_citation = ann['url']
+                state.citations.append(url_citation)
+                ctx.urls = state.citations
+            elif ann['type'] == "container_file_citation":
+                state.files.append({
+                    "container_id": ann['container_id'],
+                    "file_id": ann['file_id'],
+                })
+
+        elif etype == "response.reasoning_summary_text.delta":
+            response = chunk.delta
+
+        elif etype == "response.output_item.done":
+            # Delegate to computer handler which may add tool calls
+            tool_calls, has_calls = core.api.openai.computer.handle_stream_chunk(ctx, chunk, state.tool_calls)
+            state.tool_calls = tool_calls
+            if has_calls:
+                state.force_func_call = True
+
+        elif etype == "response.code_interpreter_call_code.delta":
+            if not state.is_code:
+                response = "\n\n**Code interpreter**\n```python\n" + chunk.delta
+                state.is_code = True
+            else:
+                response = chunk.delta
+
+        elif etype == "response.code_interpreter_call_code.done":
+            response = "\n\n```\n-----------\n"
+
+        elif etype == "response.image_generation_call.partial_image":
+            image_base64 = chunk.partial_image_b64
+            image_bytes = base64.b64decode(image_base64)
+            if state.img_path:
+                with open(state.img_path, "wb") as f:
+                    f.write(image_bytes)
+            del image_bytes
+            state.is_image = True
+
+        elif etype == "response.created":
+            ctx.msg_id = str(chunk.response.id)
+            core.ctx.update_item(ctx)
+
+        elif etype in {"response.done", "response.failed", "error"}:
+            pass
+
+        return response
+
+    def _process_api_completion(self, chunk) -> Optional[str]:
+        """
+        OpenAI Completions stream delta.
+
+        :param chunk: The chunk object from the stream
+        :return: Response delta string or None
+        """
+        if getattr(chunk, "choices", None):
+            choice0 = chunk.choices[0]
+            if getattr(choice0, "text", None) is not None:
+                return choice0.text
+        return None
+
+    def _process_langchain_chat(self, chunk) -> Optional[str]:
+        """
+        LangChain chat streaming delta.
+
+        :param chunk: The chunk object from the stream
+        :return: Response delta string or None
+        """
+        if getattr(chunk, "content", None) is not None:
+            return str(chunk.content)
+        return None
+
+    def _process_llama_chat(self, state: WorkerState, chunk) -> Optional[str]:
+        """
+        Llama chat streaming delta with optional tool call extraction.
+
+        :param state: WorkerState
+        :param chunk: The chunk object from the stream
+        :return: Response delta string or None
+        """
+        response = None
+        if getattr(chunk, "delta", None) is not None:
+            response = str(chunk.delta)
+
+        tool_chunks = getattr(getattr(chunk, "message", None), "additional_kwargs", {}).get("tool_calls", [])
+        if tool_chunks:
+            for tool_chunk in tool_chunks:
+                id_val = getattr(tool_chunk, "call_id", None) or getattr(tool_chunk, "id", None)
+                name = getattr(tool_chunk, "name", None) or getattr(getattr(tool_chunk, "function", None), "name", None)
+                args = getattr(tool_chunk, "arguments", None)
+                if args is None:
+                    f = getattr(tool_chunk, "function", None)
+                    args = getattr(f, "arguments", None) if f else None
+                if id_val:
+                    if not args:
+                        args = "{}"
+                    tool_call = {
+                        "id": id_val,
+                        "type": "function",
+                        "function": {"name": name, "arguments": args}
+                    }
+                    state.tool_calls.clear()
+                    state.tool_calls.append(tool_call)
+
+        return response
+
+    def _process_google_chunk(self, ctx: CtxItem, core, state: WorkerState, chunk) -> Optional[str]:
+        """
+        Google python-genai streaming chunk.
+
+        Handles text, tool calls, inline images, code execution parts, citations, and usage.
 
-            if error:
-                emit_error(error)
+        :param ctx: CtxItem
+        :param core: Core instance
+        :param state: WorkerState
+        :param chunk: The chunk object from the stream
+        :return: Response delta string or None
+        """
+        response_parts: list[str] = []
+
+        # Keep a reference to stream object for resolve() later if needed
+        if state.google_stream_ref is None:
+            state.google_stream_ref = state.generator
+
+        # Try to capture usage from this chunk (usage_metadata)
+        try:
+            um = getattr(chunk, "usage_metadata", None)
+            if um:
+                self._capture_google_usage(state, um)
+        except Exception:
+            pass
+
+        # 1) Plain text delta (if present)
+        t = None
+        try:
+            t = getattr(chunk, "text", None)
+            if t:
+                response_parts.append(t)
+        except Exception:
+            pass
+
+        # 2) Tool calls (function_calls property preferred)
+        fc_list = []
+        try:
+            fc_list = getattr(chunk, "function_calls", None) or []
+        except Exception:
+            fc_list = []
+
+        new_calls = []
+
+        def _to_plain_dict(obj):
+            """
+            Best-effort conversion of SDK objects to plain dict/list.
+            """
+            try:
+                if hasattr(obj, "to_json_dict"):
+                    return obj.to_json_dict()
+                if hasattr(obj, "model_dump"):
+                    return obj.model_dump()
+                if hasattr(obj, "to_dict"):
+                    return obj.to_dict()
+            except Exception:
+                pass
+            if isinstance(obj, dict):
+                return {k: _to_plain_dict(v) for k, v in obj.items()}
+            if isinstance(obj, (list, tuple)):
+                return [_to_plain_dict(x) for x in obj]
+            return obj
+
+        if fc_list:
+            for fc in fc_list:
+                name = getattr(fc, "name", "") or ""
+                args_obj = getattr(fc, "args", {}) or {}
+                args_dict = _to_plain_dict(args_obj) or {}
+                new_calls.append({
+                    "id": getattr(fc, "id", "") or "",
+                    "type": "function",
+                    "function": {
+                        "name": name,
+                        "arguments": json.dumps(args_dict, ensure_ascii=False),
+                    }
+                })
+        else:
+            # Fallback: read from candidates -> parts[].function_call
+            try:
+                cands = getattr(chunk, "candidates", None) or []
+                for cand in cands:
+                    content = getattr(cand, "content", None)
+                    parts = getattr(content, "parts", None) or []
+                    for p in parts:
+                        fn = getattr(p, "function_call", None)
+                        if not fn:
+                            continue
+                        name = getattr(fn, "name", "") or ""
+                        args_obj = getattr(fn, "args", {}) or {}
+                        args_dict = _to_plain_dict(args_obj) or {}
+                        new_calls.append({
+                            "id": getattr(fn, "id", "") or "",
+                            "type": "function",
+                            "function": {
+                                "name": name,
+                                "arguments": json.dumps(args_dict, ensure_ascii=False),
+                            }
+                        })
+            except Exception:
+                pass
+
+        # De-duplicate tool calls and mark force flag if any found
+        if new_calls:
+            seen = {(tc["function"]["name"], tc["function"]["arguments"]) for tc in state.tool_calls}
+            for tc in new_calls:
+                key = (tc["function"]["name"], tc["function"]["arguments"])
+                if key not in seen:
+                    state.tool_calls.append(tc)
+                    seen.add(key)
+            state.force_func_call = True
+
+        # 3) Inspect candidates for code execution parts, inline images, and citations
+        try:
+            cands = getattr(chunk, "candidates", None) or []
+            for cand in cands:
+                content = getattr(cand, "content", None)
+                parts = getattr(content, "parts", None) or []
+
+                for p in parts:
+                    # Code execution: executable code part -> open or append within fenced block
+                    ex = getattr(p, "executable_code", None)
+                    if ex:
+                        lang = (getattr(ex, "language", None) or "python").strip() or "python"
+                        code_txt = (
+                            getattr(ex, "code", None) or
+                            getattr(ex, "program", None) or
+                            getattr(ex, "source", None) or
+                            ""
+                        )
+                        if code_txt is None:
+                            code_txt = ""
+                        if not state.is_code:
+                            response_parts.append(f"\n\n**Code interpreter**\n```{lang.lower()}\n{code_txt}")
+                            state.is_code = True
+                        else:
+                            response_parts.append(str(code_txt))
+
+                    # Code execution result -> close fenced block (output will be streamed as normal text if provided)
+                    cer = getattr(p, "code_execution_result", None)
+                    if cer:
+                        if state.is_code:
+                            response_parts.append("\n\n```\n-----------\n")
+                            state.is_code = False
+                        # Note: We do not append execution outputs here to avoid duplicating chunk.text.
+
+                    # Inline image blobs
+                    blob = getattr(p, "inline_data", None)
+                    if blob:
+                        mime = (getattr(blob, "mime_type", "") or "").lower()
+                        if mime.startswith("image/"):
+                            data = getattr(blob, "data", None)
+                            if data:
+                                # inline_data.data may be bytes or base64-encoded string
+                                if isinstance(data, (bytes, bytearray)):
+                                    img_bytes = bytes(data)
+                                else:
+                                    img_bytes = base64.b64decode(data)
+                                save_path = core.image.gen_unique_path(ctx)
+                                with open(save_path, "wb") as f:
+                                    f.write(img_bytes)
+                                if not isinstance(ctx.images, list):
+                                    ctx.images = []
+                                ctx.images.append(save_path)
+                                state.image_paths.append(save_path)
+                                state.has_google_inline_image = True
+
+                    # File data that points to externally hosted image (http/https)
+                    fdata = getattr(p, "file_data", None)
+                    if fdata:
+                        uri = getattr(fdata, "file_uri", None) or getattr(fdata, "uri", None)
+                        mime = (getattr(fdata, "mime_type", "") or "").lower()
+                        if uri and mime.startswith("image/") and (uri.startswith("http://") or uri.startswith("https://")):
+                            if ctx.urls is None:
+                                ctx.urls = []
+                            ctx.urls.append(uri)
+
+            # Collect citations (web search URLs) if present in candidates metadata
+            self._collect_google_citations(ctx, state, chunk)
+
+        except Exception:
+            # Never break stream on extraction failures
+            pass
+
+        # Combine all response parts
+        response = "".join(response_parts) if response_parts else None
+        return response
+
+    def _process_raw(self, chunk) -> Optional[str]:
+        """
+        Raw chunk fallback.
 
-            emit_end(ctx)
+        :param chunk: The chunk object from the stream
+        :return: String representation of chunk or None
+        """
+        if chunk is not None:
+            return chunk if isinstance(chunk, str) else str(chunk)
+        return None
 
-            for _buf in fn_args_buffers.values():
+    # ------------ Usage helpers ------------
+
+    def _safe_get(self, obj, path: str):
+        """
+        Dot-path getter for dicts and objects.
+
+        :param obj: dict or object
+        :param path: Dot-separated path string
+        """
+        cur = obj
+        for seg in path.split("."):
+            if cur is None:
+                return None
+            if isinstance(cur, dict):
+                cur = cur.get(seg)
+            else:
+                # Support numeric indices for lists like candidates.0...
+                if seg.isdigit() and isinstance(cur, (list, tuple)):
+                    idx = int(seg)
+                    if 0 <= idx < len(cur):
+                        cur = cur[idx]
+                    else:
+                        return None
+                else:
+                    cur = getattr(cur, seg, None)
+        return cur
+
+    def _as_int(self, val):
+        """
+        Coerce to int if possible, else None.
+
+        :param val: Any value
+        :return: int or None
+        """
+        if val is None:
+            return None
+        try:
+            return int(val)
+        except Exception:
+            try:
+                return int(float(val))
+            except Exception:
+                return None
+
+    def _capture_openai_usage(self, state: WorkerState, u_obj):
+        """
+        Extract usage for OpenAI; include reasoning tokens in output if available.
+
+        :param state: WorkerState
+        :param u_obj: Usage object from OpenAI response
+        """
+        if not u_obj:
+            return
+        state.usage_vendor = "openai"
+        in_tok = self._as_int(self._safe_get(u_obj, "input_tokens")) or self._as_int(self._safe_get(u_obj, "prompt_tokens"))
+        out_tok = self._as_int(self._safe_get(u_obj, "output_tokens")) or self._as_int(self._safe_get(u_obj, "completion_tokens"))
+        total = self._as_int(self._safe_get(u_obj, "total_tokens"))
+        reasoning = (
+            self._as_int(self._safe_get(u_obj, "output_tokens_details.reasoning_tokens")) or
+            self._as_int(self._safe_get(u_obj, "completion_tokens_details.reasoning_tokens")) or
+            self._as_int(self._safe_get(u_obj, "reasoning_tokens")) or
+            0
+        )
+        out_with_reason = (out_tok or 0) + (reasoning or 0)
+        state.usage_payload = {"in": in_tok, "out": out_with_reason, "reasoning": reasoning or 0, "total": total}
+
+    def _capture_google_usage(self, state: WorkerState, um_obj):
+        """
+        Extract usage for Google python-genai; prefer total - prompt to include reasoning.
+
+        :param state: WorkerState
+        :param um_obj: Usage metadata object from Google chunk
+        """
+        if not um_obj:
+            return
+        state.usage_vendor = "google"
+        prompt = (
+            self._as_int(self._safe_get(um_obj, "prompt_token_count")) or
+            self._as_int(self._safe_get(um_obj, "prompt_tokens")) or
+            self._as_int(self._safe_get(um_obj, "input_tokens"))
+        )
+        total = (
+            self._as_int(self._safe_get(um_obj, "total_token_count")) or
+            self._as_int(self._safe_get(um_obj, "total_tokens"))
+        )
+        candidates = (
+            self._as_int(self._safe_get(um_obj, "candidates_token_count")) or
+            self._as_int(self._safe_get(um_obj, "output_tokens"))
+        )
+        reasoning = (
+            self._as_int(self._safe_get(um_obj, "candidates_reasoning_token_count")) or
+            self._as_int(self._safe_get(um_obj, "reasoning_tokens")) or 0
+        )
+        if total is not None and prompt is not None:
+            out_total = max(0, total - prompt)
+        else:
+            out_total = candidates
+        state.usage_payload = {"in": prompt, "out": out_total, "reasoning": reasoning or 0, "total": total}
+
+    def _collect_google_citations(self, ctx: CtxItem, state: WorkerState, chunk: Any):
+        """
+        Collect web citations (URLs) from Google GenAI stream.
+
+        Tries multiple known locations (grounding metadata and citation metadata)
+        in a defensive manner to remain compatible with SDK changes.
+        """
+        try:
+            cands = getattr(chunk, "candidates", None) or []
+        except Exception:
+            cands = []
+
+        if not isinstance(state.citations, list):
+            state.citations = []
+
+        # Helper to add URLs with de-duplication
+        def _add_url(url: Optional[str]):
+            if not url or not isinstance(url, str):
+                return
+            url = url.strip()
+            if not (url.startswith("http://") or url.startswith("https://")):
+                return
+            # Initialize ctx.urls if needed
+            if ctx.urls is None:
+                ctx.urls = []
+            if url not in state.citations:
+                state.citations.append(url)
+            if url not in ctx.urls:
+                ctx.urls.append(url)
+
+        # Candidate-level metadata extraction
+        for cand in cands:
+            # Grounding metadata (web search attributions)
+            gm = self._safe_get(cand, "grounding_metadata") or self._safe_get(cand, "groundingMetadata")
+            if gm:
+                atts = self._safe_get(gm, "grounding_attributions") or self._safe_get(gm, "groundingAttributions") or []
                 try:
-                    _buf.close()
+                    for att in atts or []:
+                        # Try several common paths for URI
+                        for path in (
+                            "web.uri",
+                            "web.url",
+                            "source.web.uri",
+                            "source.web.url",
+                            "source.uri",
+                            "source.url",
+                            "uri",
+                            "url",
+                        ):
+                            _add_url(self._safe_get(att, path))
                 except Exception:
                     pass
-            fn_args_buffers.clear()
-            files.clear()
-            tool_calls.clear()
-            if citations is not None and citations is not ctx.urls:
-                citations.clear()
-            citations = None
-
-            self.cleanup()
+                # Also check search entry point
+                for path in (
+                    "search_entry_point.uri",
+                    "search_entry_point.url",
+                    "searchEntryPoint.uri",
+                    "searchEntryPoint.url",
+                    "search_entry_point.rendered_content_uri",
+                    "searchEntryPoint.rendered_content_uri",
+                ):
+                    _add_url(self._safe_get(gm, path))
+
+            # Citation metadata (legacy and alt paths)
+            cm = self._safe_get(cand, "citation_metadata") or self._safe_get(cand, "citationMetadata")
+            if cm:
+                cit_arrays = (
+                    self._safe_get(cm, "citation_sources") or
+                    self._safe_get(cm, "citationSources") or
+                    self._safe_get(cm, "citations") or []
+                )
+                try:
+                    for cit in cit_arrays or []:
+                        for path in ("uri", "url", "source.uri", "source.url", "web.uri", "web.url"):
+                            _add_url(self._safe_get(cit, path))
+                except Exception:
+                    pass
+
+            # Part-level citation metadata
+            try:
+                parts = self._safe_get(cand, "content.parts") or []
+                for p in parts:
+                    # Per-part citation metadata
+                    pcm = self._safe_get(p, "citation_metadata") or self._safe_get(p, "citationMetadata")
+                    if pcm:
+                        arr = (
+                            self._safe_get(pcm, "citation_sources") or
+                            self._safe_get(pcm, "citationSources") or
+                            self._safe_get(pcm, "citations") or []
+                        )
+                        for cit in arr or []:
+                            for path in ("uri", "url", "source.uri", "source.url", "web.uri", "web.url"):
+                                _add_url(self._safe_get(cit, path))
+                    # Per-part grounding attributions (rare)
+                    gpa = self._safe_get(p, "grounding_attributions") or self._safe_get(p, "groundingAttributions") or []
+                    for att in gpa or []:
+                        for path in ("web.uri", "web.url", "source.web.uri", "source.web.url", "uri", "url"):
+                            _add_url(self._safe_get(att, path))
+            except Exception:
+                pass
+
+        # Bind to ctx on first discovery for compatibility with other parts of the app
+        if state.citations and (ctx.urls is None or not ctx.urls):
+            ctx.urls = list(state.citations)
 
     def cleanup(self):
         """Cleanup resources after worker execution."""