pythinker-code 2.0.0__py3-none-any.whl

pythinker_code/soul/context.py

@@ -0,0 +1,339 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any, cast
+
+import aiofiles
+import aiofiles.os
+from pydantic import ValidationError
+from pythinker_core.message import Message
+
+from pythinker_code.soul.compaction import estimate_text_tokens
+from pythinker_code.soul.message import system
+from pythinker_code.utils.logging import logger
+from pythinker_code.utils.path import next_available_rotation
+
+
+class Context:
+    def __init__(self, file_backend: Path):
+        self._file_backend = file_backend
+        self._history: list[Message] = []
+        self._token_count: int = 0
+        self._pending_token_estimate: int = 0
+        self._next_checkpoint_id: int = 0
+        """The ID of the next checkpoint, starting from 0, incremented after each checkpoint."""
+        self._system_prompt: str | None = None
+
+    async def restore(self) -> bool:
+        logger.debug("Restoring context from file: {file_backend}", file_backend=self._file_backend)
+        if self._history:
+            logger.error("The context storage is already modified")
+            raise RuntimeError("The context storage is already modified")
+        if not self._file_backend.exists():
+            logger.debug("No context file found, skipping restoration")
+            return False
+        if self._file_backend.stat().st_size == 0:
+            logger.debug("Empty context file, skipping restoration")
+            return False
+
+        messages_after_last_usage: list[Message] = []
+        async with aiofiles.open(self._file_backend, encoding="utf-8", errors="replace") as f:
+            line_no = 0
+            async for line in f:
+                line_no += 1
+                if not line.strip():
+                    continue
+                line_json = self._parse_context_line(
+                    line,
+                    file_backend=self._file_backend,
+                    line_no=line_no,
+                )
+                if line_json is None:
+                    continue
+                self._apply_context_record(
+                    line_json,
+                    history=self._history,
+                    messages_after_last_usage=messages_after_last_usage,
+                    file_backend=self._file_backend,
+                    line_no=line_no,
+                )
+
+        self._pending_token_estimate = estimate_text_tokens(messages_after_last_usage)
+        return True
+
+    @property
+    def history(self) -> Sequence[Message]:
+        return self._history
+
+    @property
+    def token_count(self) -> int:
+        return self._token_count
+
+    @property
+    def token_count_with_pending(self) -> int:
+        return self._token_count + self._pending_token_estimate
+
+    @property
+    def n_checkpoints(self) -> int:
+        return self._next_checkpoint_id
+
+    @property
+    def system_prompt(self) -> str | None:
+        return self._system_prompt
+
+    @property
+    def file_backend(self) -> Path:
+        return self._file_backend
+
+    async def write_system_prompt(self, prompt: str) -> None:
+        """Write the system prompt as the first record of the context file.
+
+        If the file is empty, writes it directly. If the file already has content
+        (e.g. a legacy session without system prompt), prepends it atomically via a
+        temporary file to avoid corruption on crash and avoid loading the entire file
+        into memory.
+        """
+        prompt_line = json.dumps({"role": "_system_prompt", "content": prompt}) + "\n"
+
+        def _write_system_prompt_sync() -> None:
+            if not self._file_backend.exists() or self._file_backend.stat().st_size == 0:
+                self._file_backend.write_text(prompt_line, encoding="utf-8")
+                return
+
+            tmp_path = self._file_backend.with_suffix(".tmp")
+            with (
+                tmp_path.open("w", encoding="utf-8") as tmp_f,
+                self._file_backend.open(encoding="utf-8") as src_f,
+            ):
+                tmp_f.write(prompt_line)
+                while True:
+                    chunk = src_f.read(64 * 1024)
+                    if not chunk:
+                        break
+                    tmp_f.write(chunk)
+            tmp_path.replace(self._file_backend)
+
+        await asyncio.to_thread(_write_system_prompt_sync)
+
+        self._system_prompt = prompt
+
+    async def checkpoint(self, add_user_message: bool):
+        checkpoint_id = self._next_checkpoint_id
+        self._next_checkpoint_id += 1
+        logger.debug("Checkpointing, ID: {id}", id=checkpoint_id)
+
+        async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
+            await f.write(json.dumps({"role": "_checkpoint", "id": checkpoint_id}) + "\n")
+        if add_user_message:
+            await self.append_message(
+                Message(role="user", content=[system(f"CHECKPOINT {checkpoint_id}")])
+            )
+
+    async def revert_to(self, checkpoint_id: int):
+        """
+        Revert the context to the specified checkpoint.
+        After this, the specified checkpoint and all subsequent content will be
+        removed from the context. File backend will be rotated.
+
+        Args:
+            checkpoint_id (int): The ID of the checkpoint to revert to. 0 is the first checkpoint.
+
+        Raises:
+            ValueError: When the checkpoint does not exist.
+            RuntimeError: When no available rotation path is found.
+        """
+
+        logger.debug("Reverting checkpoint, ID: {id}", id=checkpoint_id)
+        if checkpoint_id >= self._next_checkpoint_id:
+            logger.error("Checkpoint {checkpoint_id} does not exist", checkpoint_id=checkpoint_id)
+            raise ValueError(f"Checkpoint {checkpoint_id} does not exist")
+
+        # rotate the context file
+        rotated_file_path = await next_available_rotation(self._file_backend)
+        if rotated_file_path is None:
+            logger.error("No available rotation path found")
+            raise RuntimeError("No available rotation path found")
+        await aiofiles.os.replace(self._file_backend, rotated_file_path)
+        logger.debug(
+            "Rotated context file: {rotated_file_path}", rotated_file_path=rotated_file_path
+        )
+
+        # restore the context until the specified checkpoint
+        self._history.clear()
+        self._token_count = 0
+        self._next_checkpoint_id = 0
+        self._system_prompt = None
+        messages_after_last_usage: list[Message] = []
+        async with (
+            aiofiles.open(rotated_file_path, encoding="utf-8", errors="replace") as old_file,
+            aiofiles.open(self._file_backend, "w", encoding="utf-8") as new_file,
+        ):
+            line_no = 0
+            async for line in old_file:
+                line_no += 1
+                if not line.strip():
+                    continue
+
+                line_json = self._parse_context_line(
+                    line,
+                    file_backend=rotated_file_path,
+                    line_no=line_no,
+                )
+                if line_json is None:
+                    continue
+                if line_json.get("role") == "_checkpoint" and line_json.get("id") == checkpoint_id:
+                    break
+
+                keep_line = self._apply_context_record(
+                    line_json,
+                    history=self._history,
+                    messages_after_last_usage=messages_after_last_usage,
+                    file_backend=rotated_file_path,
+                    line_no=line_no,
+                )
+                if keep_line:
+                    await new_file.write(line)
+
+        self._pending_token_estimate = estimate_text_tokens(messages_after_last_usage)
+
+    async def clear(self):
+        """
+        Clear the context history.
+        This is almost equivalent to revert_to(0), but without relying on the assumption
+        that the first checkpoint exists.
+        File backend will be rotated.
+
+        Raises:
+            RuntimeError: When no available rotation path is found.
+        """
+
+        logger.debug("Clearing context")
+
+        # rotate the context file
+        rotated_file_path = await next_available_rotation(self._file_backend)
+        if rotated_file_path is None:
+            logger.error("No available rotation path found")
+            raise RuntimeError("No available rotation path found")
+        await aiofiles.os.replace(self._file_backend, rotated_file_path)
+        self._file_backend.touch()
+        logger.debug(
+            "Rotated context file: {rotated_file_path}", rotated_file_path=rotated_file_path
+        )
+
+        self._history.clear()
+        self._token_count = 0
+        self._pending_token_estimate = 0
+        self._next_checkpoint_id = 0
+        self._system_prompt = None
+
+    async def append_message(self, message: Message | Sequence[Message]):
+        logger.debug("Appending message(s) to context: {message}", message=message)
+        messages = [message] if isinstance(message, Message) else message
+        self._history.extend(messages)
+        self._pending_token_estimate += estimate_text_tokens(messages)
+
+        async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
+            for message in messages:
+                await f.write(message.model_dump_json(exclude_none=True) + "\n")
+
+    async def update_token_count(self, token_count: int):
+        logger.debug("Updating token count in context: {token_count}", token_count=token_count)
+        self._token_count = token_count
+        self._pending_token_estimate = 0
+
+        async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
+            await f.write(json.dumps({"role": "_usage", "token_count": token_count}) + "\n")
+
+    def _parse_context_line(
+        self,
+        line: str,
+        *,
+        file_backend: Path,
+        line_no: int,
+    ) -> dict[str, Any] | None:
+        try:
+            line_json = json.loads(line, strict=False)
+        except json.JSONDecodeError as exc:
+            logger.warning(
+                "Skipping malformed context line {line_no} in {file}: {error}",
+                line_no=line_no,
+                file=file_backend,
+                error=exc,
+            )
+            return None
+        if not isinstance(line_json, dict):
+            logger.warning(
+                "Skipping non-object context line {line_no} in {file}",
+                line_no=line_no,
+                file=file_backend,
+            )
+            return None
+        return cast(dict[str, Any], line_json)
+
+    def _apply_context_record(
+        self,
+        line_json: dict[str, Any],
+        *,
+        history: list[Message],
+        messages_after_last_usage: list[Message],
+        file_backend: Path,
+        line_no: int,
+    ) -> bool:
+        role = line_json.get("role")
+        if not isinstance(role, str):
+            logger.warning(
+                "Skipping context line {line_no} in {file}: missing or invalid role",
+                line_no=line_no,
+                file=file_backend,
+            )
+            return False
+        if role == "_system_prompt":
+            content = line_json.get("content")
+            if not isinstance(content, str):
+                logger.warning(
+                    "Skipping invalid system prompt line {line_no} in {file}",
+                    line_no=line_no,
+                    file=file_backend,
+                )
+                return False
+            self._system_prompt = content
+            return True
+        if role == "_usage":
+            token_count = line_json.get("token_count")
+            if not isinstance(token_count, int):
+                logger.warning(
+                    "Skipping invalid usage line {line_no} in {file}",
+                    line_no=line_no,
+                    file=file_backend,
+                )
+                return False
+            self._token_count = token_count
+            messages_after_last_usage.clear()
+            return True
+        if role == "_checkpoint":
+            checkpoint_id = line_json.get("id")
+            if not isinstance(checkpoint_id, int):
+                logger.warning(
+                    "Skipping invalid checkpoint line {line_no} in {file}",
+                    line_no=line_no,
+                    file=file_backend,
+                )
+                return False
+            self._next_checkpoint_id = checkpoint_id + 1
+            return True
+        try:
+            message = Message.model_validate(line_json)
+        except ValidationError as exc:
+            logger.warning(
+                "Skipping invalid context message line {line_no} in {file}: {error}",
+                line_no=line_no,
+                file=file_backend,
+                error=exc,
+            )
+            return False
+        history.append(message)
+        messages_after_last_usage.append(message)
+        return True

pythinker_code/soul/denwarenji.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class DMail(BaseModel):
+    message: str = Field(description="The message to send.")
+    checkpoint_id: int = Field(description="The checkpoint to send the message back to.", ge=0)
+    # TODO: allow restoring filesystem state to the checkpoint
+
+
+class DenwaRenjiError(Exception):
+    pass
+
+
+class DenwaRenji:
+    def __init__(self):
+        self._pending_dmail: DMail | None = None
+        self._n_checkpoints: int = 0
+
+    def send_dmail(self, dmail: DMail):
+        """Send a D-Mail. Intended to be called by the SendDMail tool."""
+        if self._pending_dmail is not None:
+            raise DenwaRenjiError("Only one D-Mail can be sent at a time")
+        if dmail.checkpoint_id < 0:
+            raise DenwaRenjiError("The checkpoint ID can not be negative")
+        if dmail.checkpoint_id >= self._n_checkpoints:
+            raise DenwaRenjiError("There is no checkpoint with the given ID")
+        self._pending_dmail = dmail
+
+    def set_n_checkpoints(self, n_checkpoints: int):
+        """Set the number of checkpoints. Intended to be called by the soul."""
+        self._n_checkpoints = n_checkpoints
+
+    def fetch_pending_dmail(self) -> DMail | None:
+        """Fetch a pending D-Mail. Intended to be called by the soul."""
+        pending_dmail = self._pending_dmail
+        self._pending_dmail = None
+        return pending_dmail

pythinker_code/soul/dynamic_injection.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pythinker_core.message import Message
+
+from pythinker_code.notifications import is_notification_message
+
+if TYPE_CHECKING:
+    from pythinker_code.soul.pythinkersoul import PythinkerSoul
+
+
+@dataclass(frozen=True, slots=True)
+class DynamicInjection:
+    """A dynamic prompt content to be injected before an LLM step."""
+
+    type: str  # identifier, e.g. "plan_mode"
+    content: str  # text content (will be wrapped in <system-reminder> tags)
+
+
+class DynamicInjectionProvider(ABC):
+    """Base class for dynamic injection providers.
+
+    Called before each LLM step. Implementations handle their own throttling.
+    Providers can access all runtime state via the ``soul`` parameter
+    (context_usage, runtime, config, etc.).
+    """
+
+    @abstractmethod
+    async def get_injections(
+        self,
+        history: Sequence[Message],
+        soul: PythinkerSoul,
+    ) -> list[DynamicInjection]: ...
+
+    async def on_context_compacted(self) -> None:
+        """Called after the context is compacted (history is rebuilt).
+
+        Override to reset internal throttling state when prior injections
+        may have been collapsed into the compaction summary and are no
+        longer literally present in history. Default is a no-op.
+        """
+        return None
+
+    async def on_auto_changed(self, enabled: bool) -> None:
+        """Called when auto mode is toggled at runtime.
+
+        Override to reset internal throttling state when a mode-specific
+        reminder should be eligible to fire again after a user toggle.
+        """
+        _ = enabled
+        return None
+
+
+def normalize_history(history: Sequence[Message]) -> list[Message]:
+    """Merge adjacent user messages to produce a clean API input sequence.
+
+    Dynamic injections are stored as standalone user messages in history;
+    normalization merges them into the adjacent user message.
+
+    Only ``user`` role messages are merged. Assistant and tool messages
+    are never merged because their ``tool_calls`` / ``tool_call_id``
+    fields form linked pairs that must stay intact.
+    """
+    if not history:
+        return []
+
+    result: list[Message] = []
+    for msg in history:
+        if (
+            result
+            and result[-1].role == msg.role
+            and msg.role == "user"
+            and not is_notification_message(result[-1])
+            and not is_notification_message(msg)
+        ):
+            merged_content = list(result[-1].content) + list(msg.content)
+            result[-1] = Message(role="user", content=merged_content)
+        else:
+            result.append(msg)
+    return result

pythinker_code/soul/dynamic_injections/auto_mode.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from pythinker_core.message import Message
+
+from pythinker_code.soul.dynamic_injection import DynamicInjection, DynamicInjectionProvider
+
+if TYPE_CHECKING:
+    from pythinker_code.soul.pythinkersoul import PythinkerSoul
+
+_AUTO_INJECTION_TYPE = "auto_mode"
+
+_AUTO_PROMPT = (
+    "You are running in auto mode. No user is present to answer "
+    "questions or approve actions. All tool calls are auto-approved by "
+    "the harness.\n"
+    "- Do NOT call AskUserQuestion — it will be auto-dismissed with no "
+    "answer, wasting a turn. Make your best judgment and proceed.\n"
+    "- You CAN use EnterPlanMode / ExitPlanMode normally. They will be "
+    "auto-approved. Planning still helps you think before acting; use "
+    "it for non-trivial tasks, then exit and execute.\n"
+    "- Finish the user's request end-to-end in this run. Do not defer "
+    "decisions to a human."
+)
+
+AUTO_DISABLED_REMINDER = (
+    "Auto mode is now disabled. The user is back at the terminal and CAN answer "
+    "AskUserQuestion.\n"
+    "- Ignore any earlier auto mode reminders that said no user is present or "
+    "that you must not call AskUserQuestion.\n"
+    "- AskUserQuestion is available again when a decision genuinely changes "
+    "your next action. Do not ask routine confirmations or progress check-ins.\n"
+    "- Tool calls are no longer auto-approved by auto mode. They may still be "
+    "auto-approved if yolo mode remains active."
+)
+
+
+class AutoModeInjectionProvider(DynamicInjectionProvider):
+    """Injects auto-mode (no user present) guidance for the model."""
+
+    def __init__(self) -> None:
+        self._injected: bool = False
+
+    async def get_injections(
+        self,
+        history: Sequence[Message],
+        soul: PythinkerSoul,
+    ) -> list[DynamicInjection]:
+        _ = history
+        if not soul.is_auto:
+            return []
+        if not soul.is_auto_flag:
+            return []
+
+        if self._injected:
+            return []
+        self._injected = True
+        return [DynamicInjection(type=_AUTO_INJECTION_TYPE, content=_AUTO_PROMPT)]
+
+    async def on_context_compacted(self) -> None:
+        # Compaction rewrites history; the prior auto-mode reminder may have
+        # been summarized away, so let the next auto step restate the
+        # constraint.
+        self._injected = False
+
+    async def on_auto_changed(self, enabled: bool) -> None:
+        # A runtime toggle changes the latest truth about user presence.
+        # Re-arm so the next LLM step can inject the current auto guidance.
+        _ = enabled
+        self._injected = False