data-harness 0.2.0__tar.gz → 0.3.0__tar.gz

{data_harness-0.2.0 → data_harness-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-harness
-Version: 0.2.0
+Version: 0.3.0
 Summary: Data agent with Python-native tools (no bash)
 Project-URL: Homepage, https://github.com/maxkskhor/data-harness
 Project-URL: Repository, https://github.com/maxkskhor/data-harness

{data_harness-0.2.0 → data_harness-0.3.0}/data_harness/__init__.py

@@ -12,6 +12,19 @@ from data_harness.providers.base import (
     StopReason,
 )
 from data_harness.result import CacheStorageInfo, RunResult, Usage
+from data_harness.streaming import (
+    ContentBlockDeltaEvent,
+    ContentBlockStartEvent,
+    ContentBlockStopEvent,
+    ContentDelta,
+    InputJSONDelta,
+    MessageDeltaEvent,
+    MessageStartEvent,
+    MessageStopEvent,
+    StreamEvent,
+    TextDelta,
+    ToolResultEvent,
+)
 from data_harness.types import (
     ContentBlock,
     Message,
@@ -31,17 +44,28 @@ __all__ = [
     "AsyncProviderAdapter",
     "CacheStorageInfo",
     "ContentBlock",
+    "ContentBlockDeltaEvent",
+    "ContentBlockStartEvent",
+    "ContentBlockStopEvent",
+    "ContentDelta",
+    "InputJSONDelta",
     "MaxTurnsExceeded",
     "Message",
+    "MessageDeltaEvent",
+    "MessageStartEvent",
+    "MessageStopEvent",
     "NormalizedResponse",
     "ProviderAdapter",
     "RunResult",
     "StopReason",
+    "StreamEvent",
     "SubagentRecursionError",
     "TextBlock",
+    "TextDelta",
     "ToolAnnotations",
     "ToolNotFoundError",
     "ToolResultBlock",
+    "ToolResultEvent",
     "ToolSpec",
     "ToolUseBlock",
     "Usage",

{data_harness-0.2.0 → data_harness-0.3.0}/data_harness/agent.py

@@ -20,6 +20,7 @@ from data_harness.loop import AsyncHarness, Harness
 from data_harness.providers.base import AsyncProviderAdapter, ProviderAdapter
 from data_harness.result import RunResult
 from data_harness.schema import infer_input_schema
+from data_harness.streaming import StreamEvent
 from data_harness.tools.connectors import ConnectorRegistry
 from data_harness.tools.interpreter import PythonInterpreter
 from data_harness.tools.planner import Planner
@@ -385,18 +386,24 @@ class AsyncAgent:
             raise RuntimeError(result.error or "unknown error")
         return result.text
 
-    async def run_stream(self, user_message: str) -> AsyncGenerator[str, None]:
-        """Stream assistant tokens for a one-shot run.
+    async def run_stream(self, user_message: str) -> AsyncGenerator[StreamEvent, None]:
+        """Stream events for a one-shot run.
+
+        Yields StreamEvent objects (message_start, content_block_*, message_delta,
+        message_stop, tool_result) following the Claude Agent SDK protocol.
 
         Usage::
 
-            async for chunk in agent.run_stream("hello"):
-                print(chunk, end="", flush=True)
+            async for event in agent.run_stream("hello"):
+                if event.type == "content_block_delta":
+                    from data_harness.streaming import TextDelta
+                    if isinstance(event.delta, TextDelta):
+                        print(event.delta.text, end="", flush=True)
         """
         harness = self._make_harness()
         self._last_harness = harness
-        async for chunk in harness.run_stream(user_message):
-            yield chunk
+        async for event in harness.run_stream(user_message):
+            yield event
         self._last_run_file = harness.run_file
 
     def _build_tools(
@@ -507,10 +514,10 @@ class AsyncAgentSession:
             raise RuntimeError(result.error or "unknown error")
         return result.text
 
-    async def ask_stream(self, user_message: str) -> AsyncGenerator[str, None]:
-        """Stream assistant tokens for a follow-up turn."""
-        async for chunk in self._harness.ask_stream(user_message):
-            yield chunk
+    async def ask_stream(self, user_message: str) -> AsyncGenerator[StreamEvent, None]:
+        """Stream events for a follow-up turn."""
+        async for event in self._harness.ask_stream(user_message):
+            yield event
         self._agent._last_harness = self._harness
         self._agent._last_run_file = self._harness.run_file
 

{data_harness-0.2.0 → data_harness-0.3.0}/data_harness/loop.py

@@ -18,6 +18,11 @@ from data_harness.providers.base import (
     StopReason,
 )
 from data_harness.result import CacheStorageInfo, RunResult, Usage
+from data_harness.streaming import (
+    StreamEvent,
+    ToolResultEvent,
+    accumulate_stream_events,
+)
 from data_harness.types import (
     Message,
     TextBlock,
@@ -355,27 +360,29 @@ class AsyncHarness:
             raise RuntimeError(result.error or "unknown error")
         return result.text
 
-    async def run_stream(self, user_message: str) -> AsyncGenerator[str, None]:
-        """Stream assistant tokens for a one-shot run.
+    async def run_stream(self, user_message: str) -> AsyncGenerator[StreamEvent, None]:
+        """Stream events for a one-shot run.
 
-        Yields text chunks as they arrive. Tool-use turns are dispatched
-        internally; only final-answer text reaches the caller.
-        The JSONL logger records fully assembled messages, not individual chunks.
+        Yields StreamEvent objects following the same protocol as the Claude
+        Agent SDK.  Each provider turn emits message_start,
+        content_block_start/delta/stop, message_delta, and message_stop events.
+        After the harness dispatches a tool call a ToolResultEvent is emitted.
+        The JSONL logger records fully assembled messages, not individual events.
         """
         self._run_file = setup_logger(self._run_dir)
         self._messages = [Message(role="user", content=[TextBlock(text=user_message)])]
-        async for chunk in self._run_loop_stream():
-            yield chunk
+        async for event in self._run_loop_stream():
+            yield event
 
-    async def ask_stream(self, user_message: str) -> AsyncGenerator[str, None]:
-        """Stream assistant tokens for a follow-up turn in a session."""
+    async def ask_stream(self, user_message: str) -> AsyncGenerator[StreamEvent, None]:
+        """Stream events for a follow-up turn in a session."""
         if self._run_file is None:
             self._run_file = setup_logger(self._run_dir)
         self._messages.append(
             Message(role="user", content=[TextBlock(text=user_message)])
         )
-        async for chunk in self._run_loop_stream():
-            yield chunk
+        async for event in self._run_loop_stream():
+            yield event
 
     async def _run_loop_result(self) -> RunResult:
         if self._run_file is None:
@@ -471,7 +478,7 @@ class AsyncHarness:
                     cache_storage=self._build_cache_storage(),
                 )
 
-    async def _run_loop_stream(self) -> AsyncGenerator[str, None]:
+    async def _run_loop_stream(self) -> AsyncGenerator[StreamEvent, None]:
         if self._run_file is None:
             raise RuntimeError("run_file must be initialised before running the loop")
 
@@ -481,35 +488,17 @@ class AsyncHarness:
             self._apply_reminders(turn)
             visible_tools = [t for t in self._tools if t.visible]
 
-            chunk_queue: asyncio.Queue[str | None] = asyncio.Queue()
-
-            async def on_chunk(text: str) -> None:
-                await chunk_queue.put(text)
+            events_this_turn: list[StreamEvent] = []
 
-            async def do_stream() -> NormalizedResponse:
+            with time_block() as tb:
                 try:
-                    return await self._adapter.stream(
+                    async for evt in self._adapter.stream_events(
                         system=self._system,
                         messages=self._messages,
                         tools=visible_tools,
-                        on_chunk=on_chunk,
-                    )
-                finally:
-                    await chunk_queue.put(None)  # sentinel: stream is done
-
-            with time_block() as tb:
-                stream_task: asyncio.Task[NormalizedResponse] = asyncio.create_task(
-                    do_stream()
-                )
-
-                while True:
-                    chunk = await chunk_queue.get()
-                    if chunk is None:
-                        break
-                    yield chunk
-
-                try:
-                    response = await stream_task
+                    ):
+                        events_this_turn.append(evt)
+                        yield evt
                 except Exception as exc:
                     log_error_turn(
                         turn=turn,
@@ -521,6 +510,7 @@ class AsyncHarness:
                     return
 
             latency = tb.elapsed_ms
+            response = accumulate_stream_events(events_this_turn)
 
             total_usage = total_usage + Usage(
                 input_tokens=response.input_tokens,
@@ -535,6 +525,20 @@ class AsyncHarness:
 
             if response.stop_reason == StopReason.TOOL_USE:
                 tool_results = await self._dispatch_tools(response.content)
+
+                tool_name_map = {
+                    b.tool_use_id: b.tool_name
+                    for b in response.content
+                    if isinstance(b, ToolUseBlock)
+                }
+                for result in tool_results:
+                    yield ToolResultEvent(
+                        tool_use_id=result.tool_use_id,
+                        tool_name=tool_name_map.get(result.tool_use_id, ""),
+                        content=result.content,
+                        is_error=result.is_error,
+                    )
+
                 self._messages.append(Message(role="user", content=list(tool_results)))
 
             tool_error_count = sum(1 for r in tool_results if r.is_error)

{data_harness-0.2.0 → data_harness-0.3.0}/data_harness/providers/anthropic.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
 
 import anthropic
 
@@ -19,6 +20,9 @@ from data_harness.types import (
     ToolUseBlock,
 )
 
+if TYPE_CHECKING:
+    from data_harness.streaming import StreamEvent
+
 _STOP_REASON_MAP = {
     "end_turn": StopReason.END_TURN,
     "tool_use": StopReason.TOOL_USE,
@@ -156,14 +160,24 @@ class AsyncAnthropicAdapter(_AnthropicHelpers, AsyncProviderAdapter):
         )
         return self._normalize_response(resp)
 
-    async def stream(
+    async def stream_events(
         self,
         system: str,
         messages: list[Message],
         tools: list[ToolSpec],
-        *,
-        on_chunk: Callable[[str], Awaitable[None]],
-    ) -> NormalizedResponse:
+    ) -> AsyncGenerator[StreamEvent, None]:
+        """Yield StreamEvents by mapping raw Anthropic SSE events."""
+        from data_harness.streaming import (
+            ContentBlockDeltaEvent,
+            ContentBlockStartEvent,
+            ContentBlockStopEvent,
+            InputJSONDelta,
+            MessageDeltaEvent,
+            MessageStartEvent,
+            MessageStopEvent,
+            TextDelta,
+        )
+
         api_system = self._build_system(system)
         api_messages = self._build_messages(messages)
         api_tools = self._build_tools(tools)
@@ -175,8 +189,60 @@ class AsyncAnthropicAdapter(_AnthropicHelpers, AsyncProviderAdapter):
             messages=api_messages,
             tools=api_tools or anthropic.NOT_GIVEN,
         ) as stream:
-            async for text in stream.text_stream:
-                await on_chunk(text)
-            final = await stream.get_final_message()
+            async for event in stream:
+                if event.type == "message_start":
+                    yield MessageStartEvent()
+
+                elif event.type == "content_block_start":
+                    block = event.content_block
+                    if block.type == "text":
+                        yield ContentBlockStartEvent(
+                            index=event.index,
+                            content_block=TextBlock(text=""),
+                        )
+                    elif block.type == "tool_use":
+                        yield ContentBlockStartEvent(
+                            index=event.index,
+                            content_block=ToolUseBlock(
+                                tool_use_id=block.id,
+                                tool_name=block.name,
+                                tool_input={},
+                            ),
+                        )
+
+                elif event.type == "content_block_delta":
+                    delta = event.delta
+                    if delta.type == "text_delta":
+                        yield ContentBlockDeltaEvent(
+                            index=event.index,
+                            delta=TextDelta(text=delta.text),
+                        )
+                    elif delta.type == "input_json_delta":
+                        yield ContentBlockDeltaEvent(
+                            index=event.index,
+                            delta=InputJSONDelta(partial_json=delta.partial_json),
+                        )
+
+                elif event.type == "content_block_stop":
+                    yield ContentBlockStopEvent(index=event.index)
+
+                elif event.type == "message_delta":
+                    stop_reason = _STOP_REASON_MAP.get(
+                        event.delta.stop_reason, StopReason.END_TURN
+                    )
+                    yield MessageDeltaEvent(
+                        stop_reason=stop_reason,
+                        input_tokens=event.usage.input_tokens,
+                        output_tokens=event.usage.output_tokens,
+                        cache_read_tokens=getattr(
+                            event.usage, "cache_read_input_tokens", 0
+                        )
+                        or 0,
+                        cache_write_tokens=getattr(
+                            event.usage, "cache_creation_input_tokens", 0
+                        )
+                        or 0,
+                    )
 
-        return self._normalize_response(final)
+                elif event.type == "message_stop":
+                    yield MessageStopEvent()

data_harness-0.3.0/data_harness/providers/base.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, Awaitable, Callable
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from data_harness.types import ContentBlock, Message, TextBlock, ToolSpec, ToolUseBlock
+
+if TYPE_CHECKING:
+    from data_harness.streaming import StreamEvent
+
+
+class StopReason(Enum):
+    END_TURN = "end_turn"
+    TOOL_USE = "tool_use"
+    MAX_TOKENS = "max_tokens"
+    STOP_SEQUENCE = "stop_sequence"
+
+
+@dataclass
+class NormalizedResponse:
+    stop_reason: StopReason
+    content: list[ContentBlock]
+    input_tokens: int
+    output_tokens: int
+    cache_read_tokens: int
+    cache_write_tokens: int
+
+
+class ProviderAdapter(ABC):
+    @abstractmethod
+    def chat(
+        self,
+        system: str,
+        messages: list[Message],
+        tools: list[ToolSpec],
+    ) -> NormalizedResponse: ...
+
+    @abstractmethod
+    def format_cache_control(self, obj: dict) -> dict: ...
+
+
+class AsyncProviderAdapter(ABC):
+    @abstractmethod
+    async def chat(
+        self,
+        system: str,
+        messages: list[Message],
+        tools: list[ToolSpec],
+    ) -> NormalizedResponse: ...
+
+    @abstractmethod
+    def format_cache_control(self, obj: dict) -> dict: ...
+
+    async def stream_events(
+        self,
+        system: str,
+        messages: list[Message],
+        tools: list[ToolSpec],
+    ) -> AsyncGenerator[StreamEvent, None]:
+        """Yield stream events for one provider turn.
+
+        The default implementation calls chat() and synthesises the six
+        standard event types from the assembled response.  Override in
+        provider subclasses to emit real token-level events.
+        """
+        from data_harness.streaming import (
+            ContentBlockDeltaEvent,
+            ContentBlockStartEvent,
+            ContentBlockStopEvent,
+            InputJSONDelta,
+            MessageDeltaEvent,
+            MessageStartEvent,
+            MessageStopEvent,
+            TextDelta,
+        )
+
+        response = await self.chat(system, messages, tools)
+        yield MessageStartEvent()
+        for i, block in enumerate(response.content):
+            if isinstance(block, TextBlock):
+                yield ContentBlockStartEvent(index=i, content_block=TextBlock(text=""))
+                yield ContentBlockDeltaEvent(index=i, delta=TextDelta(text=block.text))
+                yield ContentBlockStopEvent(index=i)
+            elif isinstance(block, ToolUseBlock):
+                yield ContentBlockStartEvent(
+                    index=i,
+                    content_block=ToolUseBlock(
+                        tool_use_id=block.tool_use_id,
+                        tool_name=block.tool_name,
+                        tool_input={},
+                    ),
+                )
+                yield ContentBlockDeltaEvent(
+                    index=i,
+                    delta=InputJSONDelta(partial_json=json.dumps(block.tool_input)),
+                )
+                yield ContentBlockStopEvent(index=i)
+        yield MessageDeltaEvent(
+            stop_reason=response.stop_reason,
+            input_tokens=response.input_tokens,
+            output_tokens=response.output_tokens,
+            cache_read_tokens=response.cache_read_tokens,
+            cache_write_tokens=response.cache_write_tokens,
+        )
+        yield MessageStopEvent()
+
+    async def stream(
+        self,
+        system: str,
+        messages: list[Message],
+        tools: list[ToolSpec],
+        *,
+        on_chunk: Callable[[str], Awaitable[None]],
+    ) -> NormalizedResponse:
+        """Backward-compat text-only streaming; calls stream_events() internally."""
+        from data_harness.streaming import (
+            ContentBlockDeltaEvent,
+            TextDelta,
+            accumulate_stream_events,
+        )
+
+        events = []
+        async for evt in self.stream_events(system, messages, tools):
+            events.append(evt)
+            if isinstance(evt, ContentBlockDeltaEvent) and isinstance(
+                evt.delta, TextDelta
+            ):
+                await on_chunk(evt.delta.text)
+        return accumulate_stream_events(events)

data_harness-0.3.0/data_harness/streaming.py

@@ -0,0 +1,186 @@
+"""Stream event types for AsyncHarness.run_stream() and AsyncAgent.run_stream().
+
+The event protocol mirrors the Claude Agent SDK's raw SSE event shape so that
+callers get the same discriminated-union stream whether they use data-harness or
+the Anthropic SDK directly.  An additional ToolResultEvent is emitted after the
+harness dispatches a tool call — it has no equivalent in the raw provider stream.
+
+Typical iteration pattern::
+
+    async for event in agent.run_stream("..."):
+        match event.type:
+            case "content_block_delta":
+                if isinstance(event.delta, TextDelta):
+                    print(event.delta.text, end="", flush=True)
+            case "tool_result":
+                print(f"[{event.tool_name}] → {event.content[:80]}")
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Literal
+
+from data_harness.providers.base import NormalizedResponse, StopReason
+from data_harness.types import TextBlock, ToolUseBlock
+
+# ---------------------------------------------------------------------------
+# Delta types (carried by ContentBlockDeltaEvent)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TextDelta:
+    text: str
+    type: Literal["text_delta"] = field(default="text_delta", init=False)
+
+
+@dataclass
+class InputJSONDelta:
+    partial_json: str
+    type: Literal["input_json_delta"] = field(default="input_json_delta", init=False)
+
+
+ContentDelta = TextDelta | InputJSONDelta
+
+
+# ---------------------------------------------------------------------------
+# Stream event types (same discriminator strings as the Anthropic SDK)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class MessageStartEvent:
+    type: Literal["message_start"] = field(default="message_start", init=False)
+
+
+@dataclass
+class ContentBlockStartEvent:
+    index: int
+    content_block: TextBlock | ToolUseBlock
+    type: Literal["content_block_start"] = field(
+        default="content_block_start", init=False
+    )
+
+
+@dataclass
+class ContentBlockDeltaEvent:
+    index: int
+    delta: TextDelta | InputJSONDelta
+    type: Literal["content_block_delta"] = field(
+        default="content_block_delta", init=False
+    )
+
+
+@dataclass
+class ContentBlockStopEvent:
+    index: int
+    type: Literal["content_block_stop"] = field(
+        default="content_block_stop", init=False
+    )
+
+
+@dataclass
+class MessageDeltaEvent:
+    stop_reason: StopReason
+    input_tokens: int
+    output_tokens: int
+    cache_read_tokens: int
+    cache_write_tokens: int
+    type: Literal["message_delta"] = field(default="message_delta", init=False)
+
+
+@dataclass
+class MessageStopEvent:
+    type: Literal["message_stop"] = field(default="message_stop", init=False)
+
+
+@dataclass
+class ToolResultEvent:
+    """Emitted by the harness after a tool call is dispatched and returns.
+
+    This event has no equivalent in the raw provider stream.  It signals that
+    the harness has finished executing the tool and the next turn is starting.
+    """
+
+    tool_use_id: str
+    tool_name: str
+    content: str
+    is_error: bool
+    type: Literal["tool_result"] = field(default="tool_result", init=False)
+
+
+StreamEvent = (
+    MessageStartEvent
+    | ContentBlockStartEvent
+    | ContentBlockDeltaEvent
+    | ContentBlockStopEvent
+    | MessageDeltaEvent
+    | MessageStopEvent
+    | ToolResultEvent
+)
+
+
+# ---------------------------------------------------------------------------
+# Helper: reconstruct a NormalizedResponse from accumulated stream events
+# ---------------------------------------------------------------------------
+
+
+def accumulate_stream_events(events: list[StreamEvent]) -> NormalizedResponse:
+    """Build a NormalizedResponse by replaying a list of StreamEvents."""
+    block_by_index: dict[int, TextBlock | ToolUseBlock] = {}
+    json_by_index: dict[int, str] = {}
+    stop_reason = StopReason.END_TURN
+    input_tokens = output_tokens = cache_read = cache_write = 0
+
+    for evt in events:
+        if isinstance(evt, ContentBlockStartEvent):
+            if isinstance(evt.content_block, TextBlock):
+                block_by_index[evt.index] = TextBlock(text="")
+            elif isinstance(evt.content_block, ToolUseBlock):
+                cb = evt.content_block
+                block_by_index[evt.index] = ToolUseBlock(
+                    tool_use_id=cb.tool_use_id,
+                    tool_name=cb.tool_name,
+                    tool_input={},
+                )
+                json_by_index[evt.index] = ""
+
+        elif isinstance(evt, ContentBlockDeltaEvent):
+            if isinstance(evt.delta, TextDelta):
+                b = block_by_index.get(evt.index)
+                if isinstance(b, TextBlock):
+                    b.text += evt.delta.text
+            elif isinstance(evt.delta, InputJSONDelta):
+                json_by_index[evt.index] = (
+                    json_by_index.get(evt.index, "") + evt.delta.partial_json
+                )
+
+        elif isinstance(evt, ContentBlockStopEvent):
+            if evt.index in json_by_index:
+                b = block_by_index.get(evt.index)
+                if isinstance(b, ToolUseBlock):
+                    raw = json_by_index[evt.index]
+                    try:
+                        b.tool_input = json.loads(raw) if raw else {}
+                    except json.JSONDecodeError:
+                        b.tool_input = {}
+
+        elif isinstance(evt, MessageDeltaEvent):
+            stop_reason = evt.stop_reason
+            input_tokens = evt.input_tokens
+            output_tokens = evt.output_tokens
+            cache_read = evt.cache_read_tokens
+            cache_write = evt.cache_write_tokens
+
+    content = [block_by_index[i] for i in sorted(block_by_index.keys())]
+
+    return NormalizedResponse(
+        stop_reason=stop_reason,
+        content=content,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        cache_read_tokens=cache_read,
+        cache_write_tokens=cache_write,
+    )

{data_harness-0.2.0 → data_harness-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "data-harness"
-version = "0.2.0"
+version = "0.3.0"
 description = "Data agent with Python-native tools (no bash)"
 readme = "README.md"
 license = { text = "MIT" }