codebind 0.6.2__py3-none-any.whl → 0.6.3__py3-none-any.whl

codebind/conversation.py

@@ -9,9 +9,16 @@ from dataclasses import dataclass, field
 from typing import Any, Protocol
 from uuid import uuid4
 
-from langchain_core.messages import BaseMessage, HumanMessage, message_to_dict, messages_from_dict
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    ToolMessage,
+    message_to_dict,
+    messages_from_dict,
+)
 
-_VERSION = 2
+from .images import PreparedImage, prepare_mime_image
 
 
 class ConversationStore(Protocol):
@@ -35,28 +42,39 @@ class Conversation:
 
     @classmethod
     def from_dict(cls, value: dict[str, Any]) -> Conversation:
-        if value.get("version") != _VERSION:
-            raise ValueError("Unsupported Codebind conversation version")
+        if not isinstance(value, dict):
+            raise ValueError("Invalid Codebind conversation")
         identifier = value.get("id")
         events = value.get("events")
         revision = value.get("revision")
         cwd = value.get("cwd")
-        if not isinstance(identifier, str) or not isinstance(events, list):
+        if not isinstance(identifier, str) or not identifier or not isinstance(events, list):
             raise ValueError("Invalid Codebind conversation")
-        if not isinstance(revision, int) or revision != len(events):
+        if type(revision) is not int or revision != len(events):
             raise ValueError("Invalid Codebind conversation revision")
         if any(
-            not isinstance(event, dict) or event.get("sequence") != sequence
+            not isinstance(event, dict)
+            or type(event.get("sequence")) is not int
+            or event["sequence"] != sequence
+            or event.get("type") not in {"message", "notebook", "turn_started", "turn_finished"}
+            or not isinstance(event.get("turn_id"), str)
+            or (
+                event.get("type") == "turn_finished"
+                and event.get("status") not in {"completed", "cancelled", "failed", "interrupted"}
+            )
             for sequence, event in enumerate(events, start=1)
         ):
-            raise ValueError("Invalid Codebind conversation sequence")
+            raise ValueError("Invalid Codebind conversation event")
         if not isinstance(cwd, str):
-            cwd = ""
-        return cls(identifier, revision, deepcopy(events), cwd)
+            raise ValueError("Invalid Codebind conversation directory")
+        conversation = cls(identifier, revision, deepcopy(events), cwd)
+        conversation._validate_tool_history()
+        conversation.notebook()
+        conversation.messages()
+        return conversation
 
     def as_dict(self) -> dict[str, Any]:
         return {
-            "version": _VERSION,
             "id": self.id,
             "revision": self.revision,
             "cwd": self.cwd,
@@ -64,21 +82,83 @@ class Conversation:
         }
 
     def append(self, event_type: str, **data: Any) -> dict[str, Any]:
+        if event_type == "turn_finished" and self.pending_tool_calls({data.get("turn_id")}):
+            raise ValueError("Cannot finish a turn with unanswered tool calls")
         event = {"sequence": self.revision + 1, "type": event_type, **data}
         self.events.append(event)
         self.revision += 1
         return event
 
     def append_message(self, message: BaseMessage, turn_id: str) -> dict[str, Any]:
+        if isinstance(message, ToolMessage):
+            if message.tool_call_id not in dict(self.pending_tool_calls({turn_id})):
+                raise ValueError("Tool result has no matching unanswered call")
+        elif isinstance(message, AIMessage):
+            identifiers = [call.get("id") for call in message.tool_calls]
+            if any(not isinstance(value, str) or not value for value in identifiers):
+                raise ValueError("Assistant tool calls require IDs")
+            if len(identifiers) != len(set(identifiers)):
+                raise ValueError("Assistant tool call IDs must be unique")
         return self.append("message", turn_id=turn_id, message=message_to_dict(message))
 
+    def _validate_tool_history(self) -> None:
+        pending: dict[str, str] = {}
+        seen: set[str] = set()
+        for event in self.events:
+            turn_id = event.get("turn_id")
+            if event.get("type") == "turn_finished":
+                if turn_id in pending.values():
+                    raise ValueError("Finished turn contains unanswered tool calls")
+                continue
+            if event.get("type") != "message":
+                continue
+            message = event.get("message")
+            if not isinstance(message, dict) or not isinstance(turn_id, str):
+                raise ValueError("Invalid Codebind message event")
+            data = message.get("data")
+            if not isinstance(data, dict):
+                raise ValueError("Invalid Codebind message data")
+            if message.get("type") == "ai":
+                calls = data.get("tool_calls") or []
+                if not isinstance(calls, list):
+                    raise ValueError("Invalid assistant tool calls")
+                additional = data.get("additional_kwargs") or {}
+                if not isinstance(additional, dict):
+                    raise ValueError("Invalid assistant response metadata")
+                response_items = additional.get("response_items")
+                if isinstance(response_items, list):
+                    raw_calls = [
+                        item.get("call_id")
+                        for item in response_items
+                        if isinstance(item, dict) and item.get("type") == "function_call"
+                    ]
+                    parsed_calls = [
+                        call.get("id") if isinstance(call, dict) else None for call in calls
+                    ]
+                    if raw_calls != parsed_calls:
+                        raise ValueError(
+                            "Assistant tool calls disagree with provider response items"
+                        )
+                for call in calls:
+                    identifier = call.get("id") if isinstance(call, dict) else None
+                    if not isinstance(identifier, str) or not identifier or identifier in seen:
+                        raise ValueError("Invalid or repeated assistant tool call ID")
+                    seen.add(identifier)
+                    pending[identifier] = turn_id
+            elif message.get("type") == "tool":
+                identifier = data.get("tool_call_id")
+                if not isinstance(identifier, str) or pending.get(identifier) != turn_id:
+                    raise ValueError("Tool result has no matching unanswered call")
+                del pending[identifier]
+
     def append_notebook(
         self,
         cells: list[dict[str, Any]],
         turn_id: str,
     ) -> dict[str, Any] | None:
         current = self.notebook()
-        incoming = _normalize_notebook(cells)
+        images: list[tuple[str, PreparedImage]] = []
+        incoming = _normalize_notebook(cells, images)
         has_notebook = any(event.get("type") == "notebook" for event in self.events)
 
         if not has_notebook:
@@ -109,11 +189,19 @@ class Conversation:
             )
             + "</notebook-context>"
         )
+        previous_images = _image_digests(current)
+        blocks: list[dict[str, str]] = [{"type": "text", "text": content}]
+        for label, image in images:
+            if previous_images.get(label) == image.source_sha256:
+                continue
+            blocks.append({"type": "text", "text": f"Displayed image from {label}:"})
+            blocks.append(image.content_block())
+        message = HumanMessage(content_blocks=blocks) if len(blocks) > 1 else HumanMessage(content)
         return self.append(
             "notebook",
             turn_id=turn_id,
             **payload,
-            message=message_to_dict(HumanMessage(content)),
+            message=message_to_dict(message),
         )
 
     def rollback(self, event: dict[str, Any]) -> None:
@@ -226,7 +314,43 @@ class NotebookConversationStore:
         await self.bridge.save_conversation(value)
 
 
-def _normalize_notebook(value: Any) -> list[dict[str, Any]]:
+def _canonical_mime_bundle(
+    value: dict[str, Any],
+    label: str,
+    images: list[tuple[str, PreparedImage]] | None,
+) -> dict[str, Any]:
+    data = {key: deepcopy(item) for key, item in value.items() if not key.startswith("image/")}
+    try:
+        image = prepare_mime_image(value)
+    except ValueError as error:
+        data["image_omitted"] = str(error)
+    else:
+        if image is not None:
+            data["image"] = image.descriptor()
+            if images is not None:
+                images.append((label, image))
+    return data
+
+
+def _image_digests(cells: list[dict[str, Any]]) -> dict[str, str]:
+    digests: dict[str, str] = {}
+    for cell in cells:
+        identifier = cell["id"]
+        for index, output in enumerate(cell.get("outputs", [])):
+            descriptor = (output.get("data") or {}).get("image")
+            if isinstance(descriptor, dict) and isinstance(descriptor.get("source_sha256"), str):
+                digests[f"cell {identifier} output {index + 1}"] = descriptor["source_sha256"]
+        for name, bundle in cell.get("attachments", {}).items():
+            descriptor = bundle.get("image")
+            if isinstance(descriptor, dict) and isinstance(descriptor.get("source_sha256"), str):
+                digests[f"cell {identifier} attachment {name}"] = descriptor["source_sha256"]
+    return digests
+
+
+def _normalize_notebook(
+    value: Any,
+    images: list[tuple[str, PreparedImage]] | None = None,
+) -> list[dict[str, Any]]:
     if not isinstance(value, list):
         raise ValueError("Invalid Codebind notebook snapshot")
     cells: list[dict[str, Any]] = []
@@ -255,7 +379,34 @@ def _normalize_notebook(value: Any) -> list[dict[str, Any]]:
             ):
                 raise ValueError("Invalid Codebind code cell")
             cell["execution_count"] = execution_count
-            cell["outputs"] = deepcopy(outputs)
+            normalized_outputs: list[dict[str, Any]] = []
+            for index, output in enumerate(outputs):
+                if not isinstance(output, dict):
+                    raise ValueError("Invalid Codebind cell output")
+                normalized = deepcopy(output)
+                data = output.get("data")
+                if isinstance(data, dict):
+                    normalized["data"] = _canonical_mime_bundle(
+                        data,
+                        f"cell {identifier} output {index + 1}",
+                        images,
+                    )
+                normalized_outputs.append(normalized)
+            cell["outputs"] = normalized_outputs
+        attachments = raw_cell.get("attachments")
+        if attachments is not None:
+            if not isinstance(attachments, dict):
+                raise ValueError("Invalid Codebind cell attachments")
+            normalized_attachments: dict[str, Any] = {}
+            for name, bundle in attachments.items():
+                if not isinstance(name, str) or not isinstance(bundle, dict):
+                    raise ValueError("Invalid Codebind cell attachment")
+                normalized_attachments[name] = _canonical_mime_bundle(
+                    bundle,
+                    f"cell {identifier} attachment {name}",
+                    images,
+                )
+            cell["attachments"] = normalized_attachments
         cells.append(cell)
     return cells
 

codebind/execution.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import sys
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass
 from typing import Any, cast
 
 from IPython.core.displaypub import DisplayPublisher
@@ -13,9 +13,13 @@ from IPython.core.interactiveshell import InteractiveShell
 from IPython.utils.capture import capture_output
 
 from .display import display_cell
+from .images import PreparedImage, prepare_mime_image
 from .jupyter import JupyterLabBridge
 
 
+_MAXIMUM_IMAGES_PER_REPORT = 32
+
+
 @dataclass(frozen=True, slots=True)
 class ExecutionReport:
     """The model-facing result of one IPython cell."""
@@ -26,10 +30,21 @@ class ExecutionReport:
     result: str | None
     displays: tuple[str, ...]
     error: dict[str, str] | None
+    images: tuple[PreparedImage, ...] = ()
+    image_omissions: tuple[str, ...] = ()
 
     def as_dict(self) -> dict[str, Any]:
-        """Return a JSON-serializable representation."""
-        return asdict(self)
+        """Return the text report and image descriptors without repeating image bytes."""
+        return {
+            "ok": self.ok,
+            "stdout": self.stdout,
+            "stderr": self.stderr,
+            "result": self.result,
+            "displays": self.displays,
+            "error": self.error,
+            "images": [image.descriptor() for image in self.images],
+            "image_omissions": self.image_omissions,
+        }
 
 
 class _ExecutionDisplayHook:
@@ -156,6 +171,7 @@ def _captured_execution(
             return
 
         previous_displayhook = sys.displayhook
+        previous_trap_hook = shell.display_trap.hook
         previous_showtraceback = shell.showtraceback
         previous_showsyntaxerror = shell.showsyntaxerror
         if on_output is not None:
@@ -169,12 +185,15 @@ def _captured_execution(
                     on_clear or (lambda wait: None),
                 ),
             )
-        sys.displayhook = _ExecutionDisplayHook(shell, expression_outputs, on_output)
+        execution_displayhook = _ExecutionDisplayHook(shell, expression_outputs, on_output)
+        shell.display_trap.hook = execution_displayhook
+        sys.displayhook = execution_displayhook
         shell.showtraceback = lambda *args, **kwargs: None
         shell.showsyntaxerror = lambda *args, **kwargs: None
         try:
             yield captured
         finally:
+            shell.display_trap.hook = previous_trap_hook
             sys.displayhook = previous_displayhook
             shell.showtraceback = previous_showtraceback
             shell.showsyntaxerror = previous_showsyntaxerror
@@ -221,7 +240,7 @@ class IPythonExecutor:
         else:
             captured.show()
 
-        return self._report(result, captured)
+        return self._report(result, captured, expression_outputs)
 
     async def aexecute(self, cell: str) -> ExecutionReport:
         """Execute an async-capable cell and replay its native rich output."""
@@ -258,7 +277,7 @@ class IPythonExecutor:
         else:
             captured.show()
 
-        return self._report(result, captured)
+        return self._report(result, captured, expression_outputs)
 
     def _stream_callbacks(
         self,
@@ -279,15 +298,39 @@ class IPythonExecutor:
 
         return on_output, on_clear
 
-    @staticmethod
-    def _report(result: Any, captured: Any) -> ExecutionReport:
-        """Build the model-facing text projection of an IPython execution."""
+    def _report(
+        self,
+        result: Any,
+        captured: Any,
+        expression_outputs: list[dict[str, Any]],
+    ) -> ExecutionReport:
+        """Build the model-facing text and image projection of an IPython execution."""
 
         displays: list[str] = []
-        for output in captured.outputs:
-            data = getattr(output, "data", None)
+        images: list[PreparedImage] = []
+        omissions: list[str] = []
+        bundles = [getattr(output, "data", None) for output in captured.outputs]
+        bundles.extend(output.get("data") for output in expression_outputs)
+        if not expression_outputs and result.result is not None:
+            data, _ = self.shell.display_formatter.format(result.result)
+            if any(key.startswith("image/") for key in data):
+                bundles.append(data)
+        for index, data in enumerate(bundles):
             if isinstance(data, dict) and "text/plain" in data:
                 displays.append(str(data["text/plain"]))
+            if not isinstance(data, dict):
+                continue
+            if len(images) >= _MAXIMUM_IMAGES_PER_REPORT:
+                if any(key.startswith("image/") for key in data):
+                    omissions.append(f"Display {index + 1}: image count limit reached")
+                continue
+            try:
+                image = prepare_mime_image(data)
+            except ValueError as error:
+                omissions.append(f"Display {index + 1}: {error}")
+            else:
+                if image is not None:
+                    images.append(image)
 
         exception = result.error_before_exec or result.error_in_exec
         error = (
@@ -302,6 +345,8 @@ class IPythonExecutor:
             result=repr(result.result) if result.result is not None else None,
             displays=tuple(displays),
             error=error,
+            images=tuple(images),
+            image_omissions=tuple(omissions),
         )
 
     def _notebook_outputs(

codebind/extension.py

@@ -10,23 +10,37 @@ from langchain_core.language_models import BaseChatModel
 
 from .configuration import load_configuration, load_models
 from .conversation import MemoryConversationStore, NotebookConversationStore
-from .jupyter import JupyterLabBridge
+from .jupyter import JupyterLabBridge, TARGET_NAME
 from .session import Session
 
 
 _STATE_ATTRIBUTE = "_codebind_extension_state"
 
 
+def _close_state(state: dict[str, Any]) -> None:
+    manager = state.get("comm_manager")
+    target = state.get("comm_target")
+    if manager is not None and getattr(manager, "targets", {}).get(TARGET_NAME) is target:
+        manager.unregister_target(TARGET_NAME, target)
+    load_task = state.get("load_task")
+    if isinstance(load_task, asyncio.Task):
+        load_task.cancel()
+    session = state.get("session")
+    if isinstance(session, Session) and session.bridge is not None:
+        session.bridge.close()
+    discard_model = state.get("discard_model")
+    if callable(discard_model):
+        try:
+            asyncio.get_running_loop().create_task(discard_model())
+        except RuntimeError:
+            asyncio.run(discard_model())
+
+
 def load_ipython_extension(ipython: InteractiveShell) -> None:
     """Load Codebind into the active IPython session."""
     previous = getattr(ipython, _STATE_ATTRIBUTE, None)
     if isinstance(previous, dict):
-        previous_load = previous.get("load_task")
-        if isinstance(previous_load, asyncio.Task):
-            previous_load.cancel()
-        previous_session = previous.get("session")
-        if isinstance(previous_session, Session) and previous_session.bridge is not None:
-            previous_session.bridge.close()
+        _close_state(previous)
 
     configuration = load_configuration()
     models = load_models()
@@ -38,72 +52,103 @@ def load_ipython_extension(ipython: InteractiveShell) -> None:
             model = models.chat(configuration.model, **configuration.parameters)
         return model
 
-    def discard_model(selected: BaseChatModel) -> None:
+    async def discard_model() -> None:
         nonlocal model
-        if model is selected:
-            model = None
-
-    bridge = JupyterLabBridge.connect(ipython)
-    store = NotebookConversationStore(bridge) if bridge is not None else MemoryConversationStore()
-    session = Session(shell=ipython, bridge=bridge, store=store)
-
-    async def answer_question(question: str, notebook: list[dict[str, Any]]) -> None:
-        selected = get_model()
-        try:
-            await session.asend(question, selected, notebook=notebook)
-        except BaseException:
-            discard_model(selected)
-            raise
+        selected = model
+        model = None
+        close = getattr(selected, "aclose", None)
+        if callable(close):
+            try:
+                await close()
+            except Exception:
+                pass
+
+    manager = getattr(getattr(ipython, "kernel", None), "comm_manager", None)
+    state: dict[str, Any] = {
+        "session": None,
+        "bridge": None,
+        "load_task": None,
+        "get_model": get_model,
+        "discard_model": discard_model,
+        "comm_manager": manager,
+        "comm_target": None,
+    }
+
+    if manager is None:
+        state["session"] = Session(shell=ipython, store=MemoryConversationStore())
+    else:
+
+        def accept_comm(comm: Any, message: dict[str, Any]) -> None:
+            data = message.get("content", {}).get("data", {})
+            if not isinstance(data, dict):
+                raise ValueError("Invalid Codebind notebook connection")
+            conversation = data.get("conversation")
+            if conversation is not None and not isinstance(conversation, dict):
+                raise ValueError("Invalid Codebind conversation")
+
+            previous_task = state.get("load_task")
+            if isinstance(previous_task, asyncio.Task):
+                previous_task.cancel()
+            previous_bridge = state.get("bridge")
+            if isinstance(previous_bridge, JupyterLabBridge):
+                previous_bridge.close()
+
+            bridge = JupyterLabBridge(comm, conversation)
+            session = Session(
+                shell=ipython,
+                bridge=bridge,
+                store=NotebookConversationStore(bridge),
+            )
+
+            async def answer_question(question: str, notebook: list[dict[str, Any]]) -> None:
+                selected = get_model()
+                try:
+                    await session.asend(question, selected, notebook=notebook)
+                except BaseException:
+                    await discard_model()
+                    raise
+
+            async def prepare_session() -> None:
+                try:
+                    await session.aload()
+                except asyncio.CancelledError:
+                    raise
+                except Exception as error:
+                    if bridge.ready:
+                        bridge.report_session_ready(error)
+                else:
+                    if bridge.ready:
+                        bridge.report_session_ready()
+
+            bridge.handle_questions(answer_question)
+            state["bridge"] = bridge
+            state["session"] = session
+            state["load_task"] = asyncio.get_running_loop().create_task(prepare_session())
+
+        manager.register_target(TARGET_NAME, accept_comm)
+        state["comm_target"] = accept_comm
 
     def question_magic(line: str, cell: str | None = None) -> None:
+        session = state.get("session")
+        if not isinstance(session, Session):
+            session = Session(shell=ipython, store=MemoryConversationStore())
+            state["session"] = session
         selected = get_model()
         try:
             session.send(cell if cell is not None else line, selected)
         except BaseException:
-            discard_model(selected)
+            asyncio.run(discard_model())
             raise
 
-    load_task: asyncio.Task[None] | None = None
-    if bridge is not None:
-        bridge.handle_questions(answer_question)
-
-        async def prepare_session() -> None:
-            try:
-                await session.aload()
-            except asyncio.CancelledError:
-                raise
-            except Exception as error:
-                bridge.report_session_ready(error)
-            else:
-                bridge.report_session_ready()
-
-        try:
-            load_task = asyncio.get_running_loop().create_task(prepare_session())
-        except RuntimeError:
-            pass
     ipython.register_magic_function(question_magic, "line_cell", "question")
-    setattr(
-        ipython,
-        _STATE_ATTRIBUTE,
-        {
-            "session": session,
-            "bridge": bridge,
-            "load_task": load_task,
-            "get_model": get_model,
-        },
-    )
+    setattr(ipython, _STATE_ATTRIBUTE, state)
 
 
 def unload_ipython_extension(ipython: InteractiveShell) -> None:
     """Unload Codebind without touching the user namespace."""
     state: Any = getattr(ipython, _STATE_ATTRIBUTE, None)
     if isinstance(state, dict):
-        load_task = state.get("load_task")
-        if isinstance(load_task, asyncio.Task):
-            load_task.cancel()
-        session = state.get("session")
-        if isinstance(session, Session) and session.bridge is not None:
-            session.bridge.close()
+        _close_state(state)
         delattr(ipython, _STATE_ATTRIBUTE)
     ipython.magics_manager.magics["line"].pop("question", None)
     ipython.magics_manager.magics["cell"].pop("question", None)