pythinker-code 2.0.0__py3-none-any.whl

pythinker_code/wire/file.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import json
+import time
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import aiofiles
+from pydantic import BaseModel, ConfigDict, ValidationError
+
+from pythinker_code.utils.logging import logger
+from pythinker_code.wire.protocol import WIRE_PROTOCOL_LEGACY_VERSION, WIRE_PROTOCOL_VERSION
+from pythinker_code.wire.types import WireMessage, WireMessageEnvelope
+
+
+class WireFileMetadata(BaseModel):
+    """Metadata header stored as the first line in wire.jsonl."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    type: Literal["metadata"] = "metadata"
+    protocol_version: str
+
+
+class WireMessageRecord(BaseModel):
+    """The persisted record of a `WireMessage`."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    timestamp: float
+    message: WireMessageEnvelope
+
+    @classmethod
+    def from_wire_message(cls, msg: WireMessage, *, timestamp: float) -> WireMessageRecord:
+        return cls(timestamp=timestamp, message=WireMessageEnvelope.from_wire_message(msg))
+
+    def to_wire_message(self) -> WireMessage:
+        return self.message.to_wire_message()
+
+
+def parse_wire_file_metadata(line: str) -> WireFileMetadata | None:
+    """Parse a wire file metadata line; return None if the line is not metadata."""
+    try:
+        return WireFileMetadata.model_validate_json(line)
+    except (ValidationError, ValueError):
+        return None
+
+
+def parse_wire_file_line(line: str) -> WireFileMetadata | WireMessageRecord:
+    """Parse a wire file line into metadata or a message record."""
+    metadata = parse_wire_file_metadata(line)
+    if metadata is not None:
+        return metadata
+    return WireMessageRecord.model_validate_json(line)
+
+
+@dataclass(slots=True)
+class WireFile:
+    path: Path
+    protocol_version: str = WIRE_PROTOCOL_VERSION
+
+    def __post_init__(self) -> None:
+        if self.path.exists():
+            version = _load_protocol_version(self.path)
+            self.protocol_version = version if version is not None else WIRE_PROTOCOL_LEGACY_VERSION
+        else:
+            self.protocol_version = WIRE_PROTOCOL_VERSION
+
+    def __str__(self) -> str:
+        return str(self.path)
+
+    @property
+    def version(self) -> str:
+        return self.protocol_version
+
+    def is_empty(self) -> bool:
+        if not self.path.exists():
+            return True
+        try:
+            with self.path.open(encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    if parse_wire_file_metadata(line) is not None:
+                        continue
+                    return False
+        except OSError:
+            logger.exception("Failed to read wire file {file}:", file=self.path)
+            return False
+        return True
+
+    async def iter_records(self) -> AsyncIterator[WireMessageRecord]:
+        if not self.path.exists():
+            return
+        try:
+            async with aiofiles.open(self.path, encoding="utf-8") as f:
+                async for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        parsed = parse_wire_file_line(line)
+                    except Exception:
+                        logger.exception(
+                            "Failed to parse line in wire file {file}:", file=self.path
+                        )
+                        continue
+                    if isinstance(parsed, WireFileMetadata):
+                        continue
+                    yield parsed
+        except Exception:
+            logger.exception("Failed to read wire file {file}:", file=self.path)
+
+    async def append_message(self, msg: WireMessage, *, timestamp: float | None = None) -> None:
+        record = WireMessageRecord.from_wire_message(
+            msg,
+            timestamp=time.time() if timestamp is None else timestamp,
+        )
+        await self.append_record(record)
+
+    async def append_record(self, record: WireMessageRecord) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        needs_header = not self.path.exists() or self.path.stat().st_size == 0
+        async with aiofiles.open(self.path, mode="a", encoding="utf-8") as f:
+            if needs_header:
+                metadata = WireFileMetadata(protocol_version=self.protocol_version)
+                await f.write(_dump_line(metadata))
+            await f.write(_dump_line(record))
+
+
+def _dump_line(model: BaseModel) -> str:
+    return json.dumps(model.model_dump(mode="json"), ensure_ascii=False) + "\n"
+
+
+def _load_protocol_version(path: Path) -> str | None:
+    try:
+        with path.open(encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                metadata = parse_wire_file_metadata(line)
+                if metadata is None:
+                    return None
+                return metadata.protocol_version
+    except OSError:
+        logger.exception("Failed to read wire file {file}:", file=path)
+    return None

pythinker_code/wire/jsonrpc.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    TypeAdapter,
+    field_serializer,
+    field_validator,
+    model_serializer,
+)
+from pythinker_core.utils.typing import JsonType
+
+from pythinker_code.wire.serde import serialize_wire_message
+from pythinker_code.wire.types import (
+    ContentPart,
+    Event,
+    Request,
+    is_event,
+    is_request,
+)
+
+
+class _MessageBase(BaseModel):
+    jsonrpc: Literal["2.0"] = "2.0"
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class JSONRPCErrorObject(BaseModel):
+    code: int
+    message: str
+    data: JsonType | None = None
+
+
+class JSONRPCMessage(_MessageBase):
+    """The generic JSON-RPC message format used for validation."""
+
+    method: str | None = None
+    id: str | None = None
+    params: JsonType | None = None
+    result: JsonType | None = None
+    error: JSONRPCErrorObject | None = None
+
+    def method_is_inbound(self) -> bool:
+        return self.method in JSONRPC_IN_METHODS
+
+    def is_request(self) -> bool:
+        return self.method is not None and self.id is not None
+
+    def is_notification(self) -> bool:
+        return self.method is not None and self.id is None
+
+    def is_response(self) -> bool:
+        return self.method is None and self.id is not None
+
+
+class JSONRPCSuccessResponse(_MessageBase):
+    id: str
+    result: JsonType
+
+
+class JSONRPCErrorResponse(_MessageBase):
+    id: str
+    error: JSONRPCErrorObject
+
+
+class JSONRPCErrorResponseNullableID(_MessageBase):
+    id: str | None
+    error: JSONRPCErrorObject
+
+
+class ClientInfo(BaseModel):
+    name: str
+    version: str | None = None
+
+
+class ExternalTool(BaseModel):
+    name: str
+    description: str
+    parameters: dict[str, JsonType]
+
+
+class ClientCapabilities(BaseModel):
+    """Capabilities declared by the Wire client during initialization."""
+
+    supports_question: bool = False
+    """Whether the client can handle QuestionRequest messages."""
+    supports_plan_mode: bool = False
+    """Whether the client supports plan mode (EnterPlanMode / ExitPlanMode)."""
+
+
+class WireHookSubscription(BaseModel):
+    """Hook event subscription from the wire client."""
+
+    id: str
+    """Unique subscription ID — referenced in HookRequest."""
+    event: str
+    """Which event to subscribe to."""
+    matcher: str = ""
+    """Regex filter. Empty matches everything."""
+    timeout: int = 30
+    """Seconds to wait for client response."""
+
+
+class JSONRPCInitializeMessage(_MessageBase):
+    class Params(BaseModel):
+        protocol_version: str
+        client: ClientInfo | None = None
+        external_tools: list[ExternalTool] | None = None
+        hooks: list[WireHookSubscription] | None = None
+        capabilities: ClientCapabilities | None = None
+
+    method: Literal["initialize"] = "initialize"
+    id: str
+    params: Params
+
+
+class JSONRPCPromptMessage(_MessageBase):
+    class Params(BaseModel):
+        user_input: str | list[ContentPart]
+
+    method: Literal["prompt"] = "prompt"
+    id: str
+    params: Params
+
+    @model_serializer()
+    def _serialize(self) -> dict[str, Any]:
+        raise NotImplementedError("Prompt message serialization is not implemented.")
+
+
+class JSONRPCReplayMessage(_MessageBase):
+    method: Literal["replay"] = "replay"
+    id: str
+    params: JsonType | None = None
+
+
+class JSONRPCSteerMessage(_MessageBase):
+    class Params(BaseModel):
+        user_input: str | list[ContentPart]
+
+    method: Literal["steer"] = "steer"
+    id: str
+    params: Params
+
+    @model_serializer()
+    def _serialize(self) -> dict[str, Any]:
+        raise NotImplementedError("Steer message serialization is not implemented.")
+
+
+class _SetPlanModeParams(BaseModel):
+    enabled: bool
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class JSONRPCSetPlanModeMessage(_MessageBase):
+    method: Literal["set_plan_mode"] = "set_plan_mode"
+    id: str
+    params: _SetPlanModeParams
+
+
+class JSONRPCCancelMessage(_MessageBase):
+    method: Literal["cancel"] = "cancel"
+    id: str
+    params: JsonType | None = None
+
+    @model_serializer()
+    def _serialize(self) -> dict[str, Any]:
+        raise NotImplementedError("Cancel message serialization is not implemented.")
+
+
+class JSONRPCEventMessage(_MessageBase):
+    method: Literal["event"] = "event"
+    params: Event
+
+    @field_serializer("params")
+    def _serialize_params(self, params: Event) -> dict[str, JsonType]:
+        return serialize_wire_message(params)
+
+    @field_validator("params", mode="before")
+    @classmethod
+    def _validate_params(cls, value: Any) -> Event:
+        if is_event(value):
+            return value
+        raise NotImplementedError("Event message deserialization is not implemented.")
+
+
+class JSONRPCRequestMessage(_MessageBase):
+    method: Literal["request"] = "request"
+    id: str
+    params: Request
+
+    @field_serializer("params")
+    def _serialize_params(self, params: Request) -> dict[str, JsonType]:
+        return serialize_wire_message(params)
+
+    @field_validator("params", mode="before")
+    @classmethod
+    def _validate_params(cls, value: Any) -> Request:
+        if is_request(value):
+            return value
+        raise NotImplementedError("Request message deserialization is not implemented.")
+
+
+type JSONRPCInMessage = (
+    JSONRPCSuccessResponse
+    | JSONRPCErrorResponse
+    | JSONRPCInitializeMessage
+    | JSONRPCPromptMessage
+    | JSONRPCSteerMessage
+    | JSONRPCReplayMessage
+    | JSONRPCSetPlanModeMessage
+    | JSONRPCCancelMessage
+)
+JSONRPCInMessageAdapter = TypeAdapter[JSONRPCInMessage](JSONRPCInMessage)
+JSONRPC_IN_METHODS = {"initialize", "prompt", "steer", "replay", "set_plan_mode", "cancel"}
+
+type JSONRPCOutMessage = (
+    JSONRPCSuccessResponse
+    | JSONRPCErrorResponse
+    | JSONRPCErrorResponseNullableID
+    | JSONRPCEventMessage
+    | JSONRPCRequestMessage
+)
+JSONRPC_OUT_METHODS = {"event", "request"}
+
+
+class ErrorCodes:
+    # Predefined JSON-RPC 2.0 error codes
+    PARSE_ERROR = -32700
+    """Invalid JSON was received by the server."""
+    INVALID_REQUEST = -32600
+    """The JSON sent is not a valid Request object."""
+    METHOD_NOT_FOUND = -32601
+    """The method does not exist / is not available."""
+    INVALID_PARAMS = -32602
+    """Invalid method parameter(s)."""
+    INTERNAL_ERROR = -32603
+    """Internal JSON-RPC error."""
+
+    INVALID_STATE = -32000
+    """The server is in an invalid state to process the request."""
+    LLM_NOT_SET = -32001
+    """The LLM is not set."""
+    LLM_NOT_SUPPORTED = -32002
+    """The specified LLM is not supported."""
+    CHAT_PROVIDER_ERROR = -32003
+    """There was an error from the chat provider."""
+    AUTH_EXPIRED = -32004
+    """Authentication has expired; user should re-login."""
+
+
+class Statuses:
+    FINISHED = "finished"
+    """The agent run has finished successfully."""
+    CANCELLED = "cancelled"
+    """The agent run was cancelled by the user."""
+    MAX_STEPS_REACHED = "max_steps_reached"
+    """The agent run reached the maximum number of steps."""
+    STEERED = "steered"
+    """A steer message was queued for injection into the active turn."""

pythinker_code/wire/protocol.py

@@ -0,0 +1,2 @@
+WIRE_PROTOCOL_VERSION: str = "1.9"
+WIRE_PROTOCOL_LEGACY_VERSION: str = "1.1"

pythinker_code/wire/root_hub.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from pythinker_code.utils.aioqueue import Queue
+from pythinker_code.utils.broadcast import BroadcastQueue
+from pythinker_code.wire.types import WireMessage
+
+
+class RootWireHub:
+    """Session-level broadcast hub for out-of-turn wire messages."""
+
+    def __init__(self) -> None:
+        self._queue = BroadcastQueue[WireMessage]()
+
+    def subscribe(self) -> Queue[WireMessage]:
+        return self._queue.subscribe()
+
+    def unsubscribe(self, queue: Queue[WireMessage]) -> None:
+        self._queue.unsubscribe(queue)
+
+    async def publish(self, msg: WireMessage) -> None:
+        await self._queue.publish(msg)
+
+    def publish_nowait(self, msg: WireMessage) -> None:
+        self._queue.publish_nowait(msg)
+
+    def shutdown(self) -> None:
+        self._queue.shutdown()

pythinker_code/wire/serde.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pythinker_core.utils.typing import JsonType
+
+from pythinker_code.wire.types import WireMessage, WireMessageEnvelope
+
+
+def serialize_wire_message(msg: WireMessage) -> dict[str, JsonType]:
+    """
+    Convert a `WireMessage` into a jsonifiable dict.
+    """
+    envelope = WireMessageEnvelope.from_wire_message(msg)
+    return envelope.model_dump(mode="json")
+
+
+def deserialize_wire_message(data: dict[str, JsonType] | Any) -> WireMessage:
+    """
+    Convert a jsonifiable dict into a `WireMessage`.
+
+    Raises:
+        ValueError: If the message type is unknown or the payload is invalid.
+    """
+    envelope = WireMessageEnvelope.model_validate(data)
+    return envelope.to_wire_message()