dataact 0.1.0__py3-none-any.whl

dataact/format.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from dataact.cache import SessionCache
+
+_INLINE_STR_MAX = 500
+_INLINE_JSON_MAX = 2000
+_INLINE_COLLECTION_MAX = 20
+
+
+def format_tool_output(
+    value: Any,
+    cache: "SessionCache | None" = None,
+    preferred_name: str | None = None,
+) -> str:
+    """Decide whether to inline output or cache it; return a string for the message."""
+    if isinstance(value, Exception):
+        return f"Error: {type(value).__name__}: {value}"
+
+    # DataFrame / ndarray → always cache
+    if _is_dataframe(value) or _is_ndarray(value):
+        return _cache_value(value, cache, preferred_name or _default_name(value))
+
+    # Short string → inline
+    if isinstance(value, str):
+        if len(value) <= _INLINE_STR_MAX:
+            return value
+        return _cache_value(value, cache, preferred_name or "text_result")
+
+    # Scalar (int, float, bool, None)
+    if isinstance(value, (int, float, bool)) or value is None:
+        return str(value)
+
+    # Short dict/list → inline JSON repr
+    if isinstance(value, (dict, list)):
+        try:
+            serialized = json.dumps(value, default=repr)
+        except Exception:
+            serialized = repr(value)
+        if (
+            len(serialized) <= _INLINE_JSON_MAX
+            and _collection_size(value) <= _INLINE_COLLECTION_MAX
+        ):
+            return serialized
+        return _cache_value(value, cache, preferred_name or _default_name(value))
+
+    # Unknown object with preferred_name → cache
+    if preferred_name is not None and cache is not None:
+        return _cache_value(value, cache, preferred_name)
+
+    # Unknown object → repr truncated
+    r = repr(value)
+    if len(r) > _INLINE_STR_MAX:
+        return r[:_INLINE_STR_MAX] + "..."
+    return r
+
+
+def _cache_value(value: Any, cache: "SessionCache | None", name: str) -> str:
+    if cache is None:
+        # No cache available, fall back to repr
+        r = repr(value)
+        if len(r) > _INLINE_STR_MAX:
+            return r[:_INLINE_STR_MAX] + "..."
+        return r
+    resolved = cache.put(name, value)
+    snapshot = cache.snapshot(resolved)
+    return f"Saved as `{resolved}`\nSnapshot: {snapshot}"
+
+
+def _default_name(value: Any) -> str:
+    if _is_dataframe(value):
+        return "dataframe"
+    if _is_ndarray(value):
+        return "array"
+    if isinstance(value, dict):
+        return "result_dict"
+    if isinstance(value, list):
+        return "result_list"
+    return "result"
+
+
+def _collection_size(value: Any) -> int:
+    if isinstance(value, dict):
+        return len(value)
+    if isinstance(value, list):
+        return len(value)
+    return 0
+
+
+def _is_dataframe(value: Any) -> bool:
+    try:
+        import pandas as pd
+
+        return isinstance(value, pd.DataFrame)
+    except ImportError:
+        return False
+
+
+def _is_ndarray(value: Any) -> bool:
+    try:
+        import numpy as np
+
+        return isinstance(value, np.ndarray)
+    except ImportError:
+        return False

dataact/logger.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+
+from dataact.providers.base import NormalizedResponse
+from dataact.serialize import to_jsonable
+from dataact.types import Message, ToolResultBlock
+
+
+def setup_logger(run_dir: str = "./runs") -> str:
+    """Create a timestamped JSONL file and configure loguru. Returns the path."""
+    Path(run_dir).mkdir(parents=True, exist_ok=True)
+    ts = datetime.now(tz=timezone.utc).strftime("%Y%m%dT%H%M%S")
+    run_file = str(Path(run_dir) / f"{ts}.jsonl")
+    # Ensure the file exists
+    Path(run_file).touch()
+    logger.remove()
+    logger.add(
+        lambda msg: None, level="INFO"
+    )  # suppress default stderr; callers can add sinks
+    return run_file
+
+
+def log_turn(
+    turn: int,
+    system: str,
+    messages: list[Message],
+    response: NormalizedResponse,
+    tool_results: list[ToolResultBlock],
+    latency_ms: float,
+    run_file: str,
+    cache_storage: dict[str, dict[str, str]] | None = None,
+) -> None:
+    """Append one JSON line to the run JSONL file."""
+    system_hash = hashlib.sha256(system.encode()).hexdigest()
+
+    record: dict[str, Any] = {
+        "turn": turn,
+        "timestamp": datetime.now(tz=timezone.utc).isoformat(),
+        "system_hash": system_hash,
+        "messages": to_jsonable(messages),
+        "response_content": to_jsonable(response.content),
+        "stop_reason": response.stop_reason.value,
+        "tool_results": to_jsonable(tool_results),
+        "metrics": {
+            "input_tokens": response.input_tokens,
+            "output_tokens": response.output_tokens,
+            "cache_read_tokens": response.cache_read_tokens,
+            "cache_write_tokens": response.cache_write_tokens,
+            "latency_ms": latency_ms,
+        },
+    }
+
+    if turn == 1:
+        record["system"] = system
+    if cache_storage is not None:
+        record["cache_storage"] = cache_storage
+
+    with open(run_file, "a") as f:
+        f.write(json.dumps(record) + "\n")

dataact/loop.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+from typing import Callable
+
+from dataact.cache import SessionCache
+from dataact.exceptions import MaxTurnsExceeded
+from dataact.format import format_tool_output
+from dataact.logger import log_turn, setup_logger
+from dataact.observe import time_block
+from dataact.providers.base import NormalizedResponse, ProviderAdapter, StopReason
+from dataact.types import Message, TextBlock, ToolResultBlock, ToolSpec, ToolUseBlock
+
+_MAX_TURN_REMINDER = (
+    "This is the final turn. You MUST produce your complete final output now. "
+    "Do not use any more tools. Respond with your answer directly."
+)
+
+
+class Harness:
+    def __init__(
+        self,
+        adapter: ProviderAdapter,
+        system: str,
+        tools: list[ToolSpec],
+        max_turns: int = 25,
+        run_dir: str = "./runs",
+        cache: SessionCache | None = None,
+    ) -> None:
+        self._adapter = adapter
+        self._system = system
+        self._tools = list(tools)
+        self._max_turns = max_turns
+        self._run_dir = run_dir
+        self._cache = cache if cache is not None else SessionCache()
+        self._messages: list[Message] = []
+        self._reminders: list[Callable[[int, int], str | None]] = []
+
+    def register_reminder(self, hook: Callable[[int, int], str | None]) -> None:
+        self._reminders.append(hook)
+
+    def run(self, user_message: str) -> str:
+        run_file = setup_logger(self._run_dir)
+        self._messages = [Message(role="user", content=[TextBlock(text=user_message)])]
+        last_response: NormalizedResponse | None = None
+
+        for turn in range(1, self._max_turns + 1):
+            self._apply_reminders(turn)
+            visible_tools = [t for t in self._tools if t.visible]
+
+            with time_block() as tb:
+                response = self._adapter.chat(
+                    system=self._system,
+                    messages=self._messages,
+                    tools=visible_tools,
+                )
+            last_response = response
+            latency = tb.elapsed_ms
+
+            # Append assistant message
+            self._messages.append(Message(role="assistant", content=response.content))
+
+            tool_results: list[ToolResultBlock] = []
+
+            if response.stop_reason == StopReason.TOOL_USE:
+                tool_results = self._dispatch_tools(response.content)
+                user_msg = Message(role="user", content=list(tool_results))
+                self._messages.append(user_msg)
+
+            log_turn(
+                turn=turn,
+                system=self._system,
+                messages=self._messages,
+                response=response,
+                tool_results=tool_results,
+                latency_ms=latency,
+                run_file=run_file,
+                cache_storage=self._cache.storage_metadata(),
+            )
+
+            if response.stop_reason == StopReason.END_TURN:
+                return self._extract_text(response)
+
+            if turn == self._max_turns:
+                raise MaxTurnsExceeded(turn, last_response)
+
+        raise MaxTurnsExceeded(self._max_turns, last_response)
+
+    def _apply_reminders(self, turn: int) -> None:
+        reminder_texts: list[str] = []
+
+        for hook in self._reminders:
+            text = hook(turn, self._max_turns)
+            if text:
+                reminder_texts.append(text)
+
+        # Built-in max-turn reminder
+        if turn == self._max_turns - 1:
+            reminder_texts.append(_MAX_TURN_REMINDER)
+
+        if not reminder_texts:
+            return
+
+        combined = "\n\n".join(reminder_texts)
+        reminder_block = TextBlock(text=combined)
+
+        # Append to existing user message or create a new one
+        if self._messages and self._messages[-1].role == "user":
+            self._messages[-1].content.append(reminder_block)
+        else:
+            self._messages.append(Message(role="user", content=[reminder_block]))
+
+    def _dispatch_tools(self, content: list) -> list[ToolResultBlock]:
+        tool_uses = [b for b in content if isinstance(b, ToolUseBlock)]
+        results = []
+        tool_map = {t.name: t for t in self._tools}
+
+        for tub in tool_uses:
+            spec = tool_map.get(tub.tool_name)
+            if spec is None or spec.handler is None:
+                results.append(
+                    ToolResultBlock(
+                        tool_use_id=tub.tool_use_id,
+                        content=f"Tool not found: {tub.tool_name!r}",
+                        is_error=True,
+                    )
+                )
+                continue
+            try:
+                raw = spec.handler(**tub.tool_input)
+                output = format_tool_output(raw, cache=self._cache)
+            except Exception as exc:
+                output = repr(exc)
+                results.append(
+                    ToolResultBlock(
+                        tool_use_id=tub.tool_use_id,
+                        content=output,
+                        is_error=True,
+                    )
+                )
+                continue
+            results.append(
+                ToolResultBlock(
+                    tool_use_id=tub.tool_use_id,
+                    content=output,
+                    is_error=False,
+                )
+            )
+
+        return results
+
+    def _extract_text(self, response: NormalizedResponse) -> str:
+        texts = [b.text for b in response.content if isinstance(b, TextBlock)]
+        return "\n".join(texts)

dataact/observe.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Generator
+
+
+@dataclass
+class TurnMetrics:
+    turn: int
+    input_tokens: int
+    output_tokens: int
+    cache_read_tokens: int
+    cache_write_tokens: int
+    latency_ms: float
+
+
+class _TimeResult:
+    def __init__(self) -> None:
+        self.elapsed_ms: float = 0.0
+
+
+@contextmanager
+def time_block() -> Generator[_TimeResult, None, None]:
+    result = _TimeResult()
+    start = time.monotonic()
+    try:
+        yield result
+    finally:
+        result.elapsed_ms = (time.monotonic() - start) * 1000

dataact/providers/anthropic.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import copy
+
+import anthropic
+
+from dataact.providers.base import NormalizedResponse, ProviderAdapter, StopReason
+from dataact.types import Message, TextBlock, ToolResultBlock, ToolSpec, ToolUseBlock
+
+_STOP_REASON_MAP = {
+    "end_turn": StopReason.END_TURN,
+    "tool_use": StopReason.TOOL_USE,
+    "max_tokens": StopReason.MAX_TOKENS,
+    "stop_sequence": StopReason.STOP_SEQUENCE,
+}
+
+
+class AnthropicAdapter(ProviderAdapter):
+    def __init__(
+        self, model: str = "claude-sonnet-4-6", max_tokens: int = 8096
+    ) -> None:
+        self._model = model
+        self._max_tokens = max_tokens
+        self._client = anthropic.Anthropic()
+
+    def format_cache_control(self, obj: dict) -> dict:
+        result = copy.copy(obj)
+        result["cache_control"] = {"type": "ephemeral"}
+        return result
+
+    def chat(
+        self,
+        system: str,
+        messages: list[Message],
+        tools: list[ToolSpec],
+    ) -> NormalizedResponse:
+        # Deep-copy inputs so we never mutate harness state
+        api_system = self._build_system(system)
+        api_messages = self._build_messages(messages)
+        api_tools = self._build_tools(tools)
+
+        resp = self._client.messages.create(
+            model=self._model,
+            max_tokens=self._max_tokens,
+            system=api_system,
+            messages=api_messages,
+            tools=api_tools or anthropic.NOT_GIVEN,
+        )
+
+        stop_reason = _STOP_REASON_MAP.get(resp.stop_reason, StopReason.END_TURN)
+        content = self._normalize_content(resp.content)
+
+        return NormalizedResponse(
+            stop_reason=stop_reason,
+            content=content,
+            input_tokens=resp.usage.input_tokens,
+            output_tokens=resp.usage.output_tokens,
+            cache_read_tokens=getattr(resp.usage, "cache_read_input_tokens", 0) or 0,
+            cache_write_tokens=getattr(resp.usage, "cache_creation_input_tokens", 0)
+            or 0,
+        )
+
+    def _build_system(self, system: str) -> list[dict]:
+        return [self.format_cache_control({"type": "text", "text": system})]
+
+    def _build_messages(self, messages: list[Message]) -> list[dict]:
+        result = []
+        for i, msg in enumerate(messages):
+            blocks = [self._block_to_dict(b) for b in msg.content]
+            # Apply cache_control to last user message
+            if i == len(messages) - 1 and msg.role == "user" and blocks:
+                last_block = blocks[-1]
+                blocks[-1] = self.format_cache_control(last_block)
+            result.append({"role": msg.role, "content": blocks})
+        return result
+
+    def _block_to_dict(self, block) -> dict:
+        if isinstance(block, TextBlock):
+            return {"type": "text", "text": block.text}
+        if isinstance(block, ToolUseBlock):
+            return {
+                "type": "tool_use",
+                "id": block.tool_use_id,
+                "name": block.tool_name,
+                "input": block.tool_input,
+            }
+        if isinstance(block, ToolResultBlock):
+            return {
+                "type": "tool_result",
+                "tool_use_id": block.tool_use_id,
+                "content": block.content,
+                "is_error": block.is_error,
+            }
+        raise ValueError(f"Unknown block type: {type(block)}")
+
+    def _build_tools(self, tools: list[ToolSpec]) -> list[dict]:
+        return [t.to_provider_dict() for t in tools]
+
+    def _normalize_content(self, content) -> list:
+        blocks = []
+        for block in content:
+            if block.type == "text":
+                blocks.append(TextBlock(text=block.text))
+            elif block.type == "tool_use":
+                blocks.append(
+                    ToolUseBlock(
+                        tool_use_id=block.id,
+                        tool_name=block.name,
+                        tool_input=dict(block.input),
+                    )
+                )
+        return blocks

dataact/providers/base.py

@@ -0,0 +1,35 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+
+from dataact.types import ContentBlock, Message, ToolSpec
+
+
+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: ...

dataact/providers/openai.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import copy
+import json
+
+import openai
+
+from dataact.providers.base import NormalizedResponse, ProviderAdapter, StopReason
+from dataact.types import Message, TextBlock, ToolResultBlock, ToolSpec, ToolUseBlock
+
+_STOP_REASON_MAP = {
+    "stop": StopReason.END_TURN,
+    "tool_calls": StopReason.TOOL_USE,
+    "length": StopReason.MAX_TOKENS,
+    # TODO: expose content-filter handling separately if the core stop model grows.
+    "content_filter": StopReason.END_TURN,
+    "function_call": StopReason.TOOL_USE,
+}
+
+
+class OpenAIAdapter(ProviderAdapter):
+    def __init__(self, model: str = "gpt-4o-mini", max_tokens: int = 4096) -> None:
+        self._model = model
+        self._max_tokens = max_tokens
+        self._client = openai.OpenAI()
+
+    def format_cache_control(self, obj: dict) -> dict:
+        return copy.copy(obj)
+
+    def chat(
+        self,
+        system: str,
+        messages: list[Message],
+        tools: list[ToolSpec],
+    ) -> NormalizedResponse:
+        api_messages = self._build_messages(system, messages)
+        api_tools = self._build_tools(tools)
+
+        response = self._client.chat.completions.create(
+            model=self._model,
+            max_tokens=self._max_tokens,
+            messages=api_messages,
+            tools=api_tools or openai.NOT_GIVEN,
+        )
+
+        choice = response.choices[0]
+        stop_reason = _STOP_REASON_MAP.get(choice.finish_reason, StopReason.END_TURN)
+
+        return NormalizedResponse(
+            stop_reason=stop_reason,
+            content=self._normalize_message(choice.message),
+            input_tokens=getattr(response.usage, "prompt_tokens", 0) or 0,
+            output_tokens=getattr(response.usage, "completion_tokens", 0) or 0,
+            cache_read_tokens=0,
+            cache_write_tokens=0,
+        )
+
+    def _build_messages(self, system: str, messages: list[Message]) -> list[dict]:
+        result = [{"role": "system", "content": system}]
+        for message in messages:
+            text_blocks = [b for b in message.content if isinstance(b, TextBlock)]
+            tool_uses = [b for b in message.content if isinstance(b, ToolUseBlock)]
+            tool_results = [
+                b for b in message.content if isinstance(b, ToolResultBlock)
+            ]
+
+            if text_blocks or tool_uses:
+                api_message: dict = {
+                    "role": message.role,
+                    "content": "\n".join(b.text for b in text_blocks)
+                    if text_blocks
+                    else None,
+                }
+                if tool_uses:
+                    api_message["tool_calls"] = [
+                        {
+                            "id": block.tool_use_id,
+                            "type": "function",
+                            "function": {
+                                "name": block.tool_name,
+                                "arguments": json.dumps(block.tool_input),
+                            },
+                        }
+                        for block in tool_uses
+                    ]
+                result.append(api_message)
+
+            for block in tool_results:
+                result.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": block.tool_use_id,
+                        "content": block.content,
+                    }
+                )
+        return result
+
+    def _build_tools(self, tools: list[ToolSpec]) -> list[dict]:
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": tool.input_schema,
+                },
+            }
+            for tool in tools
+        ]
+
+    def _normalize_message(self, message) -> list:
+        blocks = []
+        content = getattr(message, "content", None)
+        if content:
+            blocks.append(TextBlock(text=content))
+
+        for call in getattr(message, "tool_calls", None) or []:
+            blocks.append(
+                ToolUseBlock(
+                    tool_use_id=call.id,
+                    tool_name=call.function.name,
+                    tool_input=json.loads(call.function.arguments or "{}"),
+                )
+            )
+        return blocks