diff --git a/statek/__init__.py b/statek/__init__.py index e824b47..e0f04bf 100644 --- a/statek/__init__.py +++ b/statek/__init__.py @@ -52,7 +52,7 @@ from .executors.post_processor import FinalCheck, PostProcessor from .system import (tool, subtask, docstr, get_any, get_all, error_handler, is_valid_error_handler, docs_style, find_sub_task_handler, - find_tools) + get_pending_notifications, find_tools) from .shared_context import ( ContextCategory, ContextCategoryDict, @@ -70,6 +70,7 @@ complete_sub_task, create_new_job, create_sub_task, + send_notification_to, ) from .runner import start_statek, start_statek_async from . import task @@ -187,6 +188,8 @@ def init(settings: Optional[StatekSettings] = None, restricted: bool = True) -> "create_new_job", "create_sub_task", "find_sub_task_handler", + "get_pending_notifications", + "send_notification_to", "ContextCategory", "ContextCategoryDict", "ContextVar", diff --git a/statek/executors/chat_log_item.py b/statek/executors/chat_log_item.py index 7f4fc36..b313b16 100644 --- a/statek/executors/chat_log_item.py +++ b/statek/executors/chat_log_item.py @@ -124,6 +124,18 @@ class SubTaskLogItem(ChatLogItem): handler: "SubTaskHandler" +@db0.memo(no_default_tags=True) +@dataclass(kw_only=True) +class NotificationLogItem(ChatLogItem): + """An arbitrary plain-text message delivered between jobs.""" + + message: str + + def __llm_repr__(self) -> str: + """Show the message without recursively formatting its saved tool result.""" + return f"NotificationLogItem(message={self.message!r})" + + @db0.memo(no_default_tags=True) @dataclass class UserLogItem: diff --git a/statek/executors/job.py b/statek/executors/job.py index 2d4b2c4..fc1da78 100644 --- a/statek/executors/job.py +++ b/statek/executors/job.py @@ -42,6 +42,7 @@ from statek.executors.chat_log_item import ( ChatLogItem, LLM_LogItem, + NotificationLogItem, PostProcessedItem, ReminderLogItem, SubTaskLogItem, @@ -64,6 +65,7 @@ build_warmup_code, extract_dialog, get_current_job, + format_default_llm_repr, parse_tool_log, parse_warmup_block, perm_ctx_get, @@ -693,8 +695,8 @@ def add_locals(self, **kwargs) -> None: """Inject additional values into this job's Python local state.""" self.py_env.update_locals(**kwargs) - def _pending_notifications(self) -> List[SubTaskLogItem]: - """Return pending subtask notifications, initializing old jobs lazily.""" + def _pending_notifications(self) -> List[Union[SubTaskLogItem, NotificationLogItem]]: + """Return pending chat-log notifications, initializing old jobs lazily.""" if not hasattr(self, "_Job__pending_chat_log") or self.__pending_chat_log is None: self.__pending_chat_log = [] return self.__pending_chat_log @@ -702,7 +704,7 @@ def _pending_notifications(self) -> List[SubTaskLogItem]: def find_sub_task_handler(self, id: Optional[Any] = None) -> Optional["SubTaskHandler"]: # pylint: disable=redefined-builtin """Find a pending or logged subtask handler by id, or the most recent one.""" for item in reversed(self._pending_notifications()): - if id is None or item.handler.id == id: + if isinstance(item, SubTaskLogItem) and (id is None or item.handler.id == id): return item.handler for item in reversed(self.chat_log): if isinstance(item, SubTaskLogItem) and (id is None or item.handler.id == id): @@ -730,10 +732,15 @@ def _chat_log_finalized_for_notifications(self) -> bool: return True def _process_pending_notifications(self) -> bool: - """Move pending subtask notifications to chat_log at safe boundaries.""" + """Move finalized pending notifications into the ordered chat history.""" pending = self._pending_notifications() if not pending or not self._chat_log_finalized_for_notifications(): return False + console_pos = len(self.py_env.console) if self.py_env.console else 0 + for item in pending: + item.console_pos = console_pos + if isinstance(item, NotificationLogItem) and item.tool_log is None: + item.tool_log = format_default_llm_repr([item]) self.chat_log.extend(pending) self.__pending_chat_log = [] return True @@ -757,6 +764,24 @@ def push_notification(self, handler: "SubTaskHandler") -> None: return self._pending_notifications().append(item) + def send_notification(self, message: str) -> None: + """Record a plain-text notification for this job's next safe boundary. + + A completed job resumes with the notification in its history. Other + jobs queue it without disturbing their pending execution or future. + """ + item = NotificationLogItem( + console_pos=len(self.py_env.console) if self.py_env.console else 0, + message=message, + ) + item.tool_log = format_default_llm_repr([item]) + if self.status == JobStatus.DONE: # pylint: disable=no-member + self.chat_log.append(item) + self.set_status(JobStatus.STARTED) # pylint: disable=no-member + self.py_env.exit_status = None + return + self._pending_notifications().append(item) + def _current_model_pricing(self): """Return pricing for the concrete model used by the current LLM request.""" metadata = self.job_def.metadata or {} @@ -1062,6 +1087,16 @@ def _subtask_tool_call(idx: int, handler: "SubTaskHandler") -> CallSpecWrapper: kwargs={"code": Job._subtask_lookup_code(handler)}, ) + @staticmethod + def _notification_tool_call(idx: int) -> CallSpecWrapper: + """Return the synthetic python_cli call for one delivered text message.""" + return CallSpecWrapper( + id=f"STATEK-NOTIFICATION-{idx:03d}", + func_name="python_cli", + args=None, + kwargs={"code": "print(get_pending_notifications())"}, + ) + def get_next_prompt(self) -> str: """ Generate the next prompt to be included in the LLM chat. @@ -1237,6 +1272,7 @@ def _build_chat_history( PostProcessedItem, ReminderLogItem, SubTaskLogItem, + NotificationLogItem, )): next_pos = items[j].console_pos break @@ -1319,6 +1355,25 @@ def _yield_pushes(lo: int, hi: int) -> Iterable[ChatHistoryItem]: ) continue + if isinstance(item, NotificationLogItem): + call = self._notification_tool_call(idx) + yield ChatHistoryItem( + role=ChatRole.ASSISTANT, + content=None, + content_src=ContentSource.SYSTEM, + tool_calls=[call], + ) + result = item.tool_log if item.tool_log is not None else format_default_llm_repr([item]) + if isinstance(result, ToolError): + result = result.err_message + yield ChatHistoryItem( + role=ChatRole.TOOL, + content=result, + content_src=ContentSource.CONSOLE, + tool_calls=[call], + ) + continue + # Determine the source content of the assistant turn. if isinstance(item, WarmupLogItem): block_num = item.warmup_block_num diff --git a/statek/executors/utils.py b/statek/executors/utils.py index 865ed7f..309f53f 100644 --- a/statek/executors/utils.py +++ b/statek/executors/utils.py @@ -938,7 +938,7 @@ def _log_pending_console(job: Job): def _process_pending_notification_follow_up(job: Job, harness) -> bool: - """Process pending subtask notifications before a job is finalized.""" + """Process pending notifications before a job is finalized.""" if not job._process_pending_notifications(): # pylint: disable=protected-access return False job.py_env.exit_status = None @@ -1211,6 +1211,7 @@ def _cli_error_append(cli_idx, text): # Step 12: Get next request parameters — log pending console batch first _log_pending_console(job) + job._process_pending_notifications() # pylint: disable=protected-access request_difficulty = job.get_current_difficulty() request = job.get_next_request(difficulty=request_difficulty) request["metadata"] = dict(request["metadata"] or {}) diff --git a/statek/system.py b/statek/system.py index c777a80..cccf8ba 100644 --- a/statek/system.py +++ b/statek/system.py @@ -26,6 +26,7 @@ from .dbzero_restricted import as_unrestricted from .future import get_any_future, get_all_future, FutureResult from .docstring import parse_tool_docstring, format_docstring +from .executors.chat_log_item import ChatLogItem from .utils import find_locals, get_current_agent_name, get_current_job, _statek_ctx_scope @@ -451,6 +452,22 @@ def wrapper(*args, **kwargs): return _decorate(f) +@tool(system=True) +def get_pending_notifications(**kwargs) -> Iterable[ChatLogItem]: # pylint: disable=unused-argument + """Return an ordered snapshot of the current job's unprocessed notifications. + + The snapshot contains pending chat log items, including subtask results, + without moving them into the job's history or consuming the queue. + + Raises: + RuntimeError: When called without a current job. + """ + job = get_current_job() + if job is None: + raise RuntimeError("get_pending_notifications requires a current job") + return list(job._pending_notifications()) # pylint: disable=protected-access + + @tool(system=True) def find_sub_task_handler(id: Optional[Any] = None, **kwargs) -> Any: # pylint: disable=redefined-builtin,unused-argument """Return a subtask handler from the current job notification history. diff --git a/statek/task.py b/statek/task.py index ec2c184..3590f15 100644 --- a/statek/task.py +++ b/statek/task.py @@ -410,6 +410,30 @@ def complete(self, result: Optional[Any] = None, error: Optional[str] = None) -> parent_job.push_notification(self) +@tool(system=True) +def send_notification_to(target: SubTaskHandler | Job, message: str, **kwargs) -> None: # pylint: disable=unused-argument + """Send plain text to an existing job or subtask handler's job. + + The message is recorded now but delivered at the recipient's next safe + execution boundary, not by interrupting its current code or awaited future. + A completed job is restarted to handle the new notification. + + Args: + target: An existing Job or SubTaskHandler for the recipient job. + message: Plain-text message to deliver without executing it as code. + + Raises: + TypeError: If target is not a Job or SubTaskHandler. + """ + if isinstance(target, SubTaskHandler): + job = target.job + elif isinstance(target, Job): + job = target + else: + raise TypeError("target must be a Job or SubTaskHandler") + job.send_notification(message) + + def complete_sub_task(result: Optional[Any] = None, error: Optional[str] = None) -> None: """Complete the subtask handler registered in the current child job locals.""" job = get_current_job() diff --git a/statek/utils.py b/statek/utils.py index 177a7e6..61ee003 100644 --- a/statek/utils.py +++ b/statek/utils.py @@ -24,6 +24,7 @@ from dataclasses import dataclass, is_dataclass, fields as dataclass_fields from datetime import datetime from decimal import Decimal +from types import UnionType from typing import ( Any, Callable, @@ -639,8 +640,8 @@ def _type_to_json_schema(type_hint) -> str: origin = get_origin(type_hint) args = get_args(type_hint) - # Optional[X] = Union[X, None] and other Union types - if origin is Union: + # Optional[X] = Union[X, None]; support both typing.Union and X | Y. + if origin in (Union, UnionType): non_none = [a for a in args if a is not _NONE_TYPE] return _type_to_json_schema(non_none[0]) if non_none else "null" diff --git a/tests/test_job.py b/tests/test_job.py index 5ab46dc..59a4a3e 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -9,7 +9,7 @@ from unittest.mock import patch, MagicMock import dbzero as db0 import pytest -from tests.conftest import create_chat_log_item, set_warmup_positions +from tests.conftest import DB0_DIR, create_chat_log_item, set_warmup_positions from statek.executors.job import ( DialogItem, Job, @@ -33,6 +33,7 @@ from statek.agents.dialog_agent import RecurringReminder from statek.executors.chat_log_item import ( LLM_LogItem, + NotificationLogItem, PostProcessedItem, ReminderLogItem, SubTaskLogItem, @@ -41,10 +42,11 @@ ) from statek.executors.post_processor import PostProcessor from statek.executors.utils import _log_pending_console +from statek.future import FutureResult from statek.settings import ChatStyle, LLM_API_Settings from statek.locale import StatekLocale, StatekLangCode, StatekCountryCode from statek.prompt_config import make_system_prompt, parse_system_prompt -from statek.utils import CodeBlock, CallSpec, _statek_ctx_scope +from statek.utils import CodeBlock, CallSpec, _statek_ctx_scope, format_default_llm_repr from statek.system import find_sub_task_handler from statek.task import SubTaskHandler, TaskError @@ -3056,6 +3058,216 @@ def test_push_notification_active_buffers_until_last_item_finalized(self, job_fa assert job._process_pending_notifications() is False # pylint: disable=protected-access +# pylint: disable=protected-access +class TestTextNotificationHistory: + """A plain-text notification is a durable, synthetic tool-result event.""" + + def test_notification_is_persisted_and_formatted_without_recursive_tool_log(self, job_factory): + """The message is required and the saved result does not expand itself.""" + job = job_factory() + item = NotificationLogItem(console_pos=0, message="Żółw\n