python-corekit 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

corekit/llm/streaming.py

@@ -0,0 +1,149 @@
+"""
+Think-tag streaming helpers for llama.cpp-style ``<think>`` content.
+"""
+
+import re
+from collections.abc import Generator, Iterator
+from enum import Enum
+from typing import NamedTuple
+
+from corekit.llm.events import StreamEvent, TextEvent, ThinkingEvent
+
+__all__ = [
+    "THINK_PATTERN",
+    "ModelStreamingState",
+    "ThinkingMarker",
+    "ThinkingMarkerPair",
+    "strip_thinking",
+]
+
+THINK_PATTERN = re.compile(r"<think>.*?</think>\s*", re.DOTALL)
+
+
+def strip_thinking(text: str) -> str:
+    """
+    Remove complete ``<think>…</think>`` blocks from a finished string.
+    """
+    return THINK_PATTERN.sub("", text).strip()
+
+
+class ThinkingMarkerPair(NamedTuple):
+    start: str
+    end: str
+
+
+class ThinkingMarker(Enum):
+    """
+    Delimiter pairs for thinking regions in model content.
+    """
+
+    XML_THINK = ThinkingMarkerPair(start="<think>", end="</think>")
+
+    @property
+    def start(self) -> str:
+        return self.value.start
+
+    @property
+    def end(self) -> str:
+        return self.value.end
+
+
+class ModelStreamingState:
+    """
+    Tracks ``<think>…</think>`` (or other marker) parse state across chunks.
+    """
+
+    def __init__(self, marker: ThinkingMarker = ThinkingMarker.XML_THINK) -> None:
+        self.marker = marker
+        self.inside = False
+        self.text_buffer = ""
+        self.think_buffer = ""
+        self._complete_buffer = ""
+        self._pending = ""  # held-back prefix of an open/close marker
+
+    def format(self) -> str:
+        return (
+            f"ModelStreamingState(text={self.text_buffer!r}, "
+            f"think={self.think_buffer!r}, complete={self._complete_buffer!r})"
+        )
+
+    def feed(self, delta: str) -> Iterator[StreamEvent]:
+        """
+        Parse one content delta; yield text / thinking events and update buffers.
+        """
+        buf = self._pending + delta
+        self._pending = ""
+
+        while buf:
+            if self.inside:
+                buf = yield from self._feed_inside(buf)
+            else:
+                buf = yield from self._feed_outside(buf)
+
+    def _feed_outside(self, buf: str) -> Generator[StreamEvent, None, str]:
+        """
+        Emit text until a start marker; return leftover buffer.
+        """
+        idx = buf.find(self.marker.start)
+        if idx != -1:
+            before = buf[:idx]
+            if before:
+                self._append_text(before)
+                yield TextEvent(content=before)
+            self.inside = True
+            return buf[idx + len(self.marker.start) :]
+
+        partial = self._partial_tag_suffix(buf, self.marker.start)
+        if partial:
+            emit = buf[: -len(partial)]
+            if emit:
+                self._append_text(emit)
+                yield TextEvent(content=emit)
+            self._pending = partial
+        else:
+            self._append_text(buf)
+            yield TextEvent(content=buf)
+        return ""
+
+    def _feed_inside(self, buf: str) -> Generator[StreamEvent, None, str]:
+        """
+        Emit thinking until an end marker; return leftover buffer.
+        """
+        idx = buf.find(self.marker.end)
+        if idx != -1:
+            thought = buf[:idx]
+            if thought:
+                self._append_think(thought)
+                yield ThinkingEvent(content=thought)
+            self.inside = False
+            return buf[idx + len(self.marker.end) :].lstrip()
+
+        partial = self._partial_tag_suffix(buf, self.marker.end)
+        if partial:
+            emit = buf[: -len(partial)]
+            if emit:
+                self._append_think(emit)
+                yield ThinkingEvent(content=emit)
+            self._pending = partial
+        else:
+            self._append_think(buf)
+            yield ThinkingEvent(content=buf)
+        return ""
+
+    def _append_text(self, new: str) -> None:
+        self.text_buffer += new
+        self._complete_buffer += new
+
+    def _append_think(self, new: str) -> None:
+        self.think_buffer += new
+        self._complete_buffer += new
+
+    @staticmethod
+    def _partial_tag_suffix(text: str, tag: str) -> str:
+        """
+        Longest non-empty suffix of ``text`` that is a prefix of ``tag``.
+        """
+        for i in range(min(len(tag) - 1, len(text)), 0, -1):
+            if text.endswith(tag[:i]):
+                return text[-i:]
+        return ""

corekit/llm/tools/__init__.py

@@ -0,0 +1,19 @@
+"""
+Tool vocabulary, registry, and detectors.
+
+``ToolLoop`` lives in ``corekit.llm.tools.loop`` and is re-exported from
+``corekit.llm`` — importing it here would cycle through ``events``.
+"""
+
+from corekit.llm.tools.base import Tool, ToolCall, ToolResult
+from corekit.llm.tools.detection import MAX_TOOL_ITERATIONS, ToolCallLoopDetector
+from corekit.llm.tools.registry import ToolRegistry
+
+__all__ = [
+    "MAX_TOOL_ITERATIONS",
+    "Tool",
+    "ToolCall",
+    "ToolCallLoopDetector",
+    "ToolRegistry",
+    "ToolResult",
+]

corekit/llm/tools/base.py

@@ -0,0 +1,118 @@
+"""
+Tool call / result vocabulary and the Tool contract.
+"""
+
+import abc
+import json
+from typing import Any
+
+from pydantic import BaseModel
+
+from corekit.llm.enum import Role, ToolKind, WireField
+
+__all__ = ["Tool", "ToolCall", "ToolResult"]
+
+
+class ToolCall(BaseModel):
+    """
+    A model-requested tool invocation.
+    """
+
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+    def arguments_json(self) -> str:
+        """
+        Serialize arguments for an OpenAI ``function.arguments`` field.
+        """
+        return json.dumps(self.arguments, default=str)
+
+    def signature(self) -> str:
+        """
+        Stable identity for loop detection (name + sorted args).
+        """
+        return f"{self.name}:{json.dumps(self.arguments, sort_keys=True, default=str)}"
+
+
+class ToolResult(BaseModel):
+    """
+    Outcome of executing a tool call, ready to feed back as a tool message.
+    """
+
+    tool_call_id: str
+    name: str
+    content: Any  # JSON-serializable; rendered back to the model as a string
+    is_error: bool = False
+
+    def to_message_content(self) -> str:
+        """
+        Serialize content for the OpenAI ``role=tool`` message body.
+        """
+        if isinstance(self.content, str):
+            return self.content
+        return json.dumps(self.content, default=str)
+
+    def to_openai_message(self) -> dict[str, Any]:
+        """
+        OpenAI ``role=tool`` chat-message dict for the next completion turn.
+        """
+        return {
+            WireField.ROLE: Role.TOOL,
+            WireField.TOOL_CALL_ID: self.tool_call_id,
+            WireField.CONTENT: self.to_message_content(),
+        }
+
+    def with_call_id(self, tool_call_id: str) -> "ToolResult":
+        """
+        Return a copy with ``tool_call_id`` set (tools often leave it blank).
+        """
+        return self.model_copy(update={"tool_call_id": tool_call_id})
+
+
+class Tool(abc.ABC):
+    """
+    One callable tool. Subclass, set class attributes, implement ``execute``.
+
+    The agentic loop always awaits ``aexecute``. Override ``aexecute`` for
+    I/O-bound tools; the default implementation runs sync ``execute``.
+    """
+
+    name: str
+    description: str
+    arguments_model: type[BaseModel]
+    kind: ToolKind = ToolKind.DATA
+
+    @abc.abstractmethod
+    def execute(self, arguments: dict[str, Any]) -> ToolResult:
+        """
+        Run the tool with validated (or raw) arguments.
+        """
+        raise NotImplementedError
+
+    async def aexecute(self, arguments: dict[str, Any]) -> ToolResult:
+        """
+        Async entry point. Default delegates to ``execute``.
+        """
+        return self.execute(arguments)
+
+    def parse_arguments(self, arguments: dict[str, Any]) -> BaseModel:
+        """
+        Validate arguments against ``arguments_model``.
+        """
+        return self.arguments_model.model_validate(arguments)
+
+    def to_openai_schema(self) -> dict[str, Any]:
+        """
+        OpenAI function-tool schema for chat completions.
+        """
+        schema = self.arguments_model.model_json_schema()
+        schema.pop("title", None)
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": schema,
+            },
+        }

corekit/llm/tools/detection.py

@@ -0,0 +1,99 @@
+"""
+Detectors for stuck thinking / tool-call loops.
+"""
+
+from corekit.llm.tools.base import ToolCall
+from corekit.observability import Loggable
+
+__all__ = [
+    "MAX_TOOL_ITERATIONS",
+    "ToolCallLoopDetector",
+]
+
+# Hard circuit breaker — stuck loops are usually caught earlier by
+# ToolCallLoopDetector. Exists so a non-repeating pathological sequence cannot
+# run forever.
+MAX_TOOL_ITERATIONS = 50
+
+_REPEAT_WINDOW = 400
+_REPEAT_LOOKBACK = 2_000
+_REPEAT_CALL_THRESHOLD = 3
+
+
+class ToolCallLoopDetector(Loggable):
+    """
+    Catches a model stuck making no progress — identical calls or repeated
+    errors — without capping a long sequence of distinct productive calls.
+
+    ``Loggable`` rather than ``Benchmarkable``: this object decides and warns;
+    wall-clock timing of model turns belongs on ``ToolLoop``.
+    """
+
+    def __init__(self, threshold: int = _REPEAT_CALL_THRESHOLD) -> None:
+        super().__init__()
+        self._threshold = threshold
+        self._last_call_signature: str | None = None
+        self._repeat_call_count = 0
+        self._last_error_signature: str | None = None
+        self._repeat_error_count = 0
+
+    @staticmethod
+    def is_thinking_looping(
+        buffer: str,
+        *,
+        window: int = _REPEAT_WINDOW,
+        lookback: int = _REPEAT_LOOKBACK,
+    ) -> bool:
+        """
+        True if the tail of the thinking buffer is a repeated substring.
+        """
+        if len(buffer) < window * 2:
+            return False
+        tail = buffer[-window:]
+        prior = buffer[-(min(len(buffer), window + lookback)) : -window]
+        return tail in prior
+
+    def check_thinking(self, buffer: str) -> bool:
+        """
+        Return True (and log) when the thinking buffer is looping.
+        """
+        if not self.is_thinking_looping(buffer):
+            return False
+        self.warning(f"Thinking loop detected at {len(buffer)} chars")
+        return True
+
+    def record_call(self, call: ToolCall) -> bool:
+        """
+        Record a tool call; return True (and log) if it matches the last call
+        enough times.
+        """
+        signature = call.signature()
+        if signature == self._last_call_signature:
+            self._repeat_call_count += 1
+        else:
+            self._last_call_signature = signature
+            self._repeat_call_count = 1
+        if self._repeat_call_count < self._threshold:
+            return False
+        self.warning(f"Repeated identical call detected — {call.name!r} (x{self._repeat_call_count})")
+        return True
+
+    def record_result(self, call: ToolCall, is_error: bool) -> bool:
+        """
+        Record a tool result; return True (and log) if the same call keeps failing.
+        """
+        if not is_error:
+            self._last_error_signature = None
+            self._repeat_error_count = 0
+            return False
+
+        signature = call.signature()
+        if signature == self._last_error_signature:
+            self._repeat_error_count += 1
+        else:
+            self._last_error_signature = signature
+            self._repeat_error_count = 1
+        if self._repeat_error_count < self._threshold:
+            return False
+        self.warning(f"Repeated failing call detected — {call.name!r} (x{self._repeat_error_count})")
+        return True

corekit/llm/tools/loop.py

@@ -0,0 +1,255 @@
+"""
+Agentic tool-calling loop over an injected chat client.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Sequence
+from typing import Any
+
+from corekit.llm.enum import StopReason, ToolKind
+from corekit.llm.events import (
+    StopEvent,
+    StreamEvent,
+    ThinkingEvent,
+    ToolCallEvent,
+    ToolResultEvent,
+    UIComponentEvent,
+)
+from corekit.llm.messages import AssistantMessage, ChatTurn, Message
+from corekit.llm.protocols import ChatClient, MetricsRecorder
+from corekit.llm.streaming import ModelStreamingState
+from corekit.llm.tools.base import ToolCall
+from corekit.llm.tools.detection import MAX_TOOL_ITERATIONS, ToolCallLoopDetector
+from corekit.llm.tools.registry import ToolRegistry
+from corekit.llm.wire import ChoiceDelta, CompletionTurn, UsageInfo
+from corekit.observability import Benchmarkable
+from corekit.utils.coercion import safe_dict
+
+__all__ = ["ToolLoop"]
+
+
+class ToolLoop(Benchmarkable):
+    """
+    Streamed tool-calling loop against an OpenAI-compatible chat client.
+
+    DATA tools loop; UI tools halt. Plain-text turns stream live and exit.
+    The client is injected — this never imports a provider SDK.
+
+    Tool execution always goes through ``ToolRegistry.aexecute`` so I/O-bound
+    tools can override ``Tool.aexecute`` without blocking the event loop.
+
+    Per-turn timing uses ``Benchmarkable``: ``reset_timing()`` at turn start,
+    ``timing("first_token")`` on the first chunk, ``timing("turn_done")`` at
+    the end — so TTFT and total duration come from splits rather than ad-hoc
+    clocks.
+    """
+
+    def __init__(
+        self,
+        client: ChatClient,
+        registry: ToolRegistry,
+        *,
+        max_iterations: int = MAX_TOOL_ITERATIONS,
+        metrics: MetricsRecorder | None = None,
+        model: str | None = None,
+    ) -> None:
+        super().__init__()
+        self._client = client
+        self._registry = registry
+        self._max_iterations = max_iterations
+        self._metrics = metrics
+        self._model = model
+        self._loop_detector = ToolCallLoopDetector()
+
+    async def stream(
+        self,
+        messages: Sequence[ChatTurn],
+        *,
+        model: str | None = None,
+        internal: bool = False,
+        **complete_kwargs: Any,
+    ) -> AsyncIterator[StreamEvent]:
+        """
+        Run the loop starting from ``messages``. Yields stream events until
+        text completes, a UI tool halts, or a stop condition fires.
+
+        ``messages`` may be ``Message`` instances or raw OpenAI message dicts.
+        """
+        working: list[dict[str, Any]] = Message.as_wire(messages, internal=internal)
+        tools_schema = self._registry.openai_schemas() or None
+        resolved_model = model or self._model
+
+        self.info(f"[ToolLoop] Starting — tools: {[t.name for t in self._registry.list_tools()]}")
+
+        for iteration in range(self._max_iterations):
+            self.info(f"[ToolLoop] Iteration {iteration + 1}/{self._max_iterations}")
+            self.reset_timing()
+            state = ModelStreamingState()
+            turn = CompletionTurn()
+
+            async for event in self._stream_model_turn(
+                working,
+                state=state,
+                turn=turn,
+                tools_schema=tools_schema,
+                model=resolved_model,
+                **complete_kwargs,
+            ):
+                yield event
+                if isinstance(event, StopEvent):
+                    return
+
+            turn_done = self.timing("turn_done")
+            self._record_metrics(turn, model=resolved_model, total_duration_ms=turn_done.total_ms)
+
+            if turn.usage is not None:
+                yield turn.usage.to_event()
+
+            self.info(
+                f"[ToolLoop] Turn done — finish_reason: {turn.finish_reason!r}, tool_calls: {len(turn.tool_calls)}"
+            )
+            self.info(f"[ToolLoop] Final State: {state.format()}")
+
+            if not turn.tool_calls:
+                return
+
+            tool_calls = turn.parsed_tool_calls()
+            working.append(AssistantMessage.with_tools(state.text_buffer.strip(), tool_calls).format())
+
+            async for event in self._dispatch_tool_calls(tool_calls, working):
+                yield event
+                if isinstance(event, (StopEvent, UIComponentEvent)):
+                    return
+
+        self.warning(f"[ToolLoop] Hit max_iterations={self._max_iterations}")
+        yield StopEvent(reason=StopReason.MAX_ITERATIONS)
+
+    async def _stream_model_turn(
+        self,
+        working: list[dict[str, Any]],
+        *,
+        state: ModelStreamingState,
+        turn: CompletionTurn,
+        tools_schema: list[dict[str, Any]] | None,
+        model: str | None,
+        **complete_kwargs: Any,
+    ) -> AsyncIterator[StreamEvent]:
+        """
+        Stream one ``complete()`` call into events and update ``state`` / ``turn``.
+        """
+        self.info(f"[ToolLoop] Streaming — messages: {len(working)}, tools: {tools_schema is not None}")
+
+        # ToolLoop is a streaming consumer; never let callers flip stream off.
+        stream_kwargs = {k: v for k, v in complete_kwargs.items() if k != "stream"}
+        async for chunk in self._client.complete(
+            working,
+            tools=tools_schema,
+            model=model,
+            stream=True,
+            **stream_kwargs,
+        ):
+            if turn.first_token_ms is None:
+                turn.first_token_ms = self.timing("first_token").total_ms
+
+            turn.absorb_usage(UsageInfo.from_chunk(chunk))
+            delta = ChoiceDelta.from_chunk(chunk)
+            if delta is None:
+                continue
+
+            turn.absorb_delta(delta)
+
+            if delta.reasoning_content:
+                async for event in self._emit_native_reasoning(state, delta.reasoning_content):
+                    yield event
+                    if isinstance(event, StopEvent):
+                        return
+
+            if delta.content:
+                for event in state.feed(delta.content):
+                    yield event
+
+    async def _emit_native_reasoning(
+        self,
+        state: ModelStreamingState,
+        reasoning: str,
+    ) -> AsyncIterator[StreamEvent]:
+        """
+        Handle ``delta.reasoning_content`` (llama.cpp native thinking field).
+        """
+        state.think_buffer += reasoning
+        if self._loop_detector.check_thinking(state.think_buffer):
+            yield StopEvent(reason=StopReason.LOOP_DETECTED)
+            return
+        yield ThinkingEvent(content=reasoning)
+
+    async def _dispatch_tool_calls(
+        self,
+        tool_calls: list[ToolCall],
+        working: list[dict[str, Any]],
+    ) -> AsyncIterator[StreamEvent]:
+        """
+        Execute parsed tool calls; append DATA results onto ``working``.
+        """
+        self.info(f"[ToolLoop] Parsed calls: {[c.name for c in tool_calls]}")
+
+        for call in tool_calls:
+            self.info(f"[ToolLoop] Executing — {call.name!r}")
+            yield ToolCallEvent(tool_call=call)
+
+            if self._loop_detector.record_call(call):
+                yield StopEvent(reason=StopReason.LOOP_DETECTED)
+                return
+
+            kind = self._registry.get_tool_kind(call.name)
+            self.info(f"[ToolLoop] Kind: {kind!r}")
+
+            if kind == ToolKind.UI:
+                self.info("[ToolLoop] UI tool — halting")
+                yield await self._execute_ui_tool(call)
+                return
+
+            if kind is None:
+                self.warning(f"[ToolLoop] Unknown tool {call.name!r}")
+
+            result = await self._registry.aexecute(call.name, call.id, call.arguments)
+            self.info(f"[ToolLoop] DATA result — is_error: {result.is_error}")
+            yield ToolResultEvent(tool_result=result)
+            working.append(result.to_openai_message())
+
+            if self._loop_detector.record_result(call, result.is_error):
+                yield StopEvent(reason=StopReason.LOOP_DETECTED)
+                return
+
+    async def _execute_ui_tool(self, call: ToolCall) -> UIComponentEvent:
+        result = await self._registry.aexecute(call.name, call.id, call.arguments)
+        self.info(f"[ToolLoop] UI result — is_error: {result.is_error}")
+        if result.is_error:
+            self.warning(f"[ToolLoop] UI tool failed: {result.content}")
+        return UIComponentEvent(
+            tool_call_id=call.id,
+            name=call.name,
+            component=safe_dict(result.content),
+        )
+
+    def _record_metrics(
+        self,
+        turn: CompletionTurn,
+        *,
+        model: str | None,
+        total_duration_ms: float,
+        succeeded: bool = True,
+    ) -> None:
+        if self._metrics is None:
+            return
+        usage = turn.usage
+        self._metrics.record(
+            model=model,
+            streamed=True,
+            succeeded=succeeded,
+            prompt_tokens=usage.prompt_tokens if usage else None,
+            completion_tokens=usage.completion_tokens if usage else None,
+            total_tokens=usage.total_tokens if usage else None,
+            time_to_first_token_ms=turn.first_token_ms,
+            total_duration_ms=total_duration_ms,
+        )