diff --git a/src/hawk/client.py b/src/hawk/client.py index 357046e..6397513 100644 --- a/src/hawk/client.py +++ b/src/hawk/client.py @@ -20,6 +20,7 @@ SessionDetail, SessionSummary, StatsResponse, + ToolResult, ) DEFAULT_BASE_URL = "http://127.0.0.1:4590" @@ -124,6 +125,8 @@ def chat( autonomy: str | None = None, cwd: str | None = None, agent: str | None = None, + tools: list[dict[str, Any]] | None = None, + tool_results: list[ToolResult] | None = None, ) -> ChatResponse: """Send a prompt and return the complete response.""" request = ChatRequest( @@ -134,6 +137,8 @@ def chat( autonomy=autonomy, cwd=cwd, agent=agent, + tools=tools, + tool_results=tool_results, ) def _do() -> ChatResponse: @@ -156,6 +161,8 @@ def chat_stream( autonomy: str | None = None, cwd: str | None = None, agent: str | None = None, + tools: list[dict[str, Any]] | None = None, + tool_results: list[ToolResult] | None = None, ) -> StreamReader: """Send a prompt and stream the response via SSE.""" request = ChatRequest( @@ -166,6 +173,8 @@ def chat_stream( autonomy=autonomy, cwd=cwd, agent=agent, + tools=tools, + tool_results=tool_results, ) def _do() -> StreamReader: @@ -190,12 +199,18 @@ def _do() -> StreamReader: def create_session( self, name: str | None = None, + model: str | None = None, + cwd: str | None = None, metadata: dict[str, Any] | None = None, ) -> SessionDetail: """Create a new session.""" body: dict[str, Any] = {} if name is not None: body["name"] = name + if model is not None: + body["model"] = model + if cwd is not None: + body["cwd"] = cwd if metadata is not None: body["metadata"] = metadata @@ -333,6 +348,8 @@ async def chat( autonomy: str | None = None, cwd: str | None = None, agent: str | None = None, + tools: list[dict[str, Any]] | None = None, + tool_results: list[ToolResult] | None = None, ) -> ChatResponse: """Send a prompt and return the complete response.""" request = ChatRequest( @@ -343,6 +360,8 @@ async def chat( autonomy=autonomy, cwd=cwd, agent=agent, + tools=tools, + tool_results=tool_results, ) async def _do() -> ChatResponse: @@ -365,6 +384,8 @@ async def chat_stream( autonomy: str | None = None, cwd: str | None = None, agent: str | None = None, + tools: list[dict[str, Any]] | None = None, + tool_results: list[ToolResult] | None = None, ) -> AsyncStreamReader: """Send a prompt and stream the response via SSE.""" request = ChatRequest( @@ -375,6 +396,8 @@ async def chat_stream( autonomy=autonomy, cwd=cwd, agent=agent, + tools=tools, + tool_results=tool_results, ) async def _do() -> AsyncStreamReader: @@ -399,12 +422,18 @@ async def _do() -> AsyncStreamReader: async def create_session( self, name: str | None = None, + model: str | None = None, + cwd: str | None = None, metadata: dict[str, Any] | None = None, ) -> SessionDetail: """Create a new session.""" body: dict[str, Any] = {} if name is not None: body["name"] = name + if model is not None: + body["model"] = model + if cwd is not None: + body["cwd"] = cwd if metadata is not None: body["metadata"] = metadata diff --git a/src/hawk/sessions.py b/src/hawk/sessions.py new file mode 100644 index 0000000..ea331f5 --- /dev/null +++ b/src/hawk/sessions.py @@ -0,0 +1,174 @@ +"""Phase-attributed session state for multi-phase agent pipelines. + +Mirrors hawk-core-contracts/sessions/sessions.go so Python agents can track +token spend and context snapshots with the same structure as Go agents. + +Phase attribution data shows that code review alone consumes 59.4% of tokens +(Tokenomics paper, arXiv 2601.14470). These dataclasses let callers record +per-phase usage and surface it in dashboards without coupling to the Go runtime. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from enum import Enum + + +class Phase(str, Enum): + """Phase identifies a step in the localize → repair → validate pipeline.""" + + LOCALIZE = "localize" + REPAIR = "repair" + VALIDATE = "validate" + REVIEW = "review" + PLANNING = "planning" + UNKNOWN = "" + + @classmethod + def parse(cls, s: str) -> Phase: + """Return the Phase matching s, or Phase.UNKNOWN if unrecognised.""" + for member in cls: + if member.value == s: + return member + return cls.UNKNOWN + + +@dataclass +class PhaseUsage: + """Token usage totals for a single pipeline phase.""" + + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + cost_usd: float = 0.0 + + +@dataclass +class ContextSnapshot: + """A point-in-time snapshot of context window state within a phase.""" + + session_id: str + phase: Phase + token_count: int + compressed_at: float = field(default_factory=time.time) # Unix timestamp + summary: str = "" + anchors: list[str] = field(default_factory=list) + + +@dataclass +class ToolCallRecord: + """A record of a single tool call with phase attribution and token counts.""" + + session_id: str + phase: Phase + tool_name: str + input_tokens: int = 0 + output_tokens: int = 0 + duration_ms: int = 0 + error: str = "" + timestamp: float = field(default_factory=time.time) # Unix timestamp + + +@dataclass +class CostAccumulator: + """Accumulates per-phase token usage for a session. + + Usage:: + + acc = CostAccumulator(session_id="sess-abc") + acc.add(Phase.LOCALIZE, input_tokens=500, output_tokens=100, cost_usd=0.002) + acc.add(Phase.REVIEW, input_tokens=3000, output_tokens=600, cost_usd=0.012) + print(acc.phase_share(Phase.REVIEW)) # ~0.857 + """ + + session_id: str + by_phase: dict[Phase, PhaseUsage] = field(default_factory=dict) + total_tokens: int = 0 + total_cost_usd: float = 0.0 + + def add( + self, + phase: Phase, + input_tokens: int, + output_tokens: int, + cost_usd: float = 0.0, + ) -> None: + """Record token usage for phase.""" + if phase not in self.by_phase: + self.by_phase[phase] = PhaseUsage() + pu = self.by_phase[phase] + pu.input_tokens += input_tokens + pu.output_tokens += output_tokens + pu.total_tokens += input_tokens + output_tokens + pu.cost_usd += cost_usd + self.total_tokens += input_tokens + output_tokens + self.total_cost_usd += cost_usd + + def phase_share(self, phase: Phase) -> float: + """Return the fraction of total tokens consumed by phase (0.0-1.0).""" + if self.total_tokens == 0: + return 0.0 + usage = self.by_phase.get(phase) + if usage is None: + return 0.0 + return usage.total_tokens / self.total_tokens + + def __str__(self) -> str: + parts = [ + f"session={self.session_id} total={self.total_tokens}tok ${self.total_cost_usd:.4f}" + ] + for phase, pu in sorted( + self.by_phase.items(), key=lambda kv: kv[1].total_tokens, reverse=True + ): + share = self.phase_share(phase) + parts.append( + f" {phase.value or 'unknown'}: {pu.total_tokens}tok ({share:.1%}) ${pu.cost_usd:.4f}" + ) + return "\n".join(parts) + + +@dataclass +class ResolutionRequest: + """SDK-level request to run a multi-phase resolution against a repository.""" + + session_id: str = "" + root_dir: str = "" + query: str = "" + max_files: int = 0 + max_symbols: int = 0 + language: str = "" + + +@dataclass +class PatchCandidate: + """A single proposed code change returned by the repair phase.""" + + file_path: str = "" + symbol: str = "" + original_body: str = "" + patched_body: str = "" + confidence: float = 0.0 + + +@dataclass +class PhaseMetrics: + """Token spend and timing for a single pipeline phase.""" + + phase: Phase = Phase.UNKNOWN + input_tokens: int = 0 + output_tokens: int = 0 + elapsed_ms: int = 0 + + +@dataclass +class ResolutionResult: + """Complete output of a multi-phase resolution run.""" + + session_id: str = "" + candidates: list[PatchCandidate] = field(default_factory=list) + validation_passed: bool = False + validation_failures: list[str] = field(default_factory=list) + phase_metrics: list[PhaseMetrics] = field(default_factory=list) + total_input_tokens: int = 0 + total_output_tokens: int = 0 diff --git a/src/hawk/tools.py b/src/hawk/tools.py index 760d306..615f695 100644 --- a/src/hawk/tools.py +++ b/src/hawk/tools.py @@ -10,6 +10,8 @@ if TYPE_CHECKING: from .types import ChatResponse +from .types import ToolResult + @dataclass class Tool: @@ -135,7 +137,7 @@ def chat_with_tools( break # Execute each tool call and collect results. - tool_results = [] + tool_results: list[ToolResult] = [] for tc in response.tool_calls: tool_name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", None) arguments = ( @@ -146,12 +148,11 @@ def chat_with_tools( else: result = json.dumps({"error": f"Unknown tool: {tool_name}"}) tool_results.append( - { - "tool_use_id": tc.get("id", "") - if isinstance(tc, dict) - else getattr(tc, "id", ""), - "content": result, - } + ToolResult( + tool_use_id=tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", ""), + content=result, + is_error=bool(tool_name) and tool_name not in tool_map, + ) ) # Send tool results back to continue the conversation. @@ -207,7 +208,7 @@ async def chat_with_tools_async( if not hasattr(response, "tool_calls") or not response.tool_calls: break - tool_results = [] + tool_results: list[ToolResult] = [] for tc in response.tool_calls: tool_name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", None) arguments = ( @@ -218,12 +219,11 @@ async def chat_with_tools_async( else: result = json.dumps({"error": f"Unknown tool: {tool_name}"}) tool_results.append( - { - "tool_use_id": tc.get("id", "") - if isinstance(tc, dict) - else getattr(tc, "id", ""), - "content": result, - } + ToolResult( + tool_use_id=tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", ""), + content=result, + is_error=bool(tool_name) and tool_name not in tool_map, + ) ) response = await client.chat( diff --git a/src/hawk/types.py b/src/hawk/types.py index 0401004..c664b57 100644 --- a/src/hawk/types.py +++ b/src/hawk/types.py @@ -19,6 +19,8 @@ class ChatRequest(BaseModel): autonomy: str | None = None cwd: str | None = None agent: str | None = None + tools: list[dict[str, Any]] | None = None + tool_results: list[ToolResult] | None = Field(None, alias="tool_results") model_config = {"populate_by_name": True} @@ -81,8 +83,8 @@ class Message(BaseModel): role: str content: str | None = None - tool_use: Any | None = Field(None, alias="tool_use") - tool_result: Any | None = Field(None, alias="tool_result") + tool_use: list[ToolCall] | None = Field(None, alias="tool_use") + tool_result: list[ToolResult] | None = Field(None, alias="tool_result") model_config = {"populate_by_name": True} @@ -95,6 +97,16 @@ class ToolCall(BaseModel): arguments: dict[str, Any] = Field(default_factory=dict) +class ToolResult(BaseModel): + """Result of executing a tool call.""" + + tool_use_id: str = Field(alias="tool_use_id") + content: str + is_error: bool = Field(False, alias="is_error") + + model_config = {"populate_by_name": True} + + class Usage(BaseModel): """Token usage information.""" diff --git a/tests/test_sessions.py b/tests/test_sessions.py new file mode 100644 index 0000000..dd4ea77 --- /dev/null +++ b/tests/test_sessions.py @@ -0,0 +1,96 @@ +"""Tests for hawk.sessions phase-attributed session state.""" + +from hawk.sessions import ContextSnapshot, CostAccumulator, Phase, ToolCallRecord + + +def test_phase_parse_known(): + assert Phase.parse("localize") == Phase.LOCALIZE + assert Phase.parse("repair") == Phase.REPAIR + assert Phase.parse("validate") == Phase.VALIDATE + assert Phase.parse("review") == Phase.REVIEW + assert Phase.parse("planning") == Phase.PLANNING + + +def test_phase_parse_unknown(): + assert Phase.parse("nonexistent") == Phase.UNKNOWN + assert Phase.parse("") == Phase.UNKNOWN + + +def test_phase_is_string(): + # Phase values must be plain strings to ensure JSON serialisation parity + # with the Go contracts package. + assert Phase.LOCALIZE == "localize" + assert Phase.REVIEW == "review" + + +def test_cost_accumulator_add_single_phase(): + acc = CostAccumulator(session_id="test-session") + acc.add(Phase.REVIEW, input_tokens=3000, output_tokens=600, cost_usd=0.012) + + assert acc.total_tokens == 3600 + assert abs(acc.total_cost_usd - 0.012) < 1e-9 + assert Phase.REVIEW in acc.by_phase + pu = acc.by_phase[Phase.REVIEW] + assert pu.input_tokens == 3000 + assert pu.output_tokens == 600 + + +def test_cost_accumulator_add_multiple_phases(): + acc = CostAccumulator(session_id="test-session") + acc.add(Phase.LOCALIZE, input_tokens=500, output_tokens=100, cost_usd=0.002) + acc.add(Phase.REPAIR, input_tokens=1000, output_tokens=200, cost_usd=0.004) + acc.add(Phase.REVIEW, input_tokens=3000, output_tokens=600, cost_usd=0.012) + + assert acc.total_tokens == 500 + 100 + 1000 + 200 + 3000 + 600 + assert len(acc.by_phase) == 3 + + +def test_cost_accumulator_phase_share_dominant(): + acc = CostAccumulator(session_id="test-session") + acc.add(Phase.REVIEW, input_tokens=5940, output_tokens=0) + acc.add(Phase.LOCALIZE, input_tokens=1000, output_tokens=0) + + share = acc.phase_share(Phase.REVIEW) + assert share > 0.8, f"review share expected > 0.8, got {share}" + + +def test_cost_accumulator_phase_share_zero_total(): + acc = CostAccumulator(session_id="empty") + assert acc.phase_share(Phase.REVIEW) == 0.0 + + +def test_cost_accumulator_phase_share_missing_phase(): + acc = CostAccumulator(session_id="test") + acc.add(Phase.REPAIR, input_tokens=100, output_tokens=20) + assert acc.phase_share(Phase.REVIEW) == 0.0 + + +def test_cost_accumulator_str(): + acc = CostAccumulator(session_id="sess-abc") + acc.add(Phase.REVIEW, input_tokens=3000, output_tokens=600, cost_usd=0.012) + s = str(acc) + assert "sess-abc" in s + assert "review" in s + + +def test_context_snapshot_defaults(): + snap = ContextSnapshot(session_id="s1", phase=Phase.LOCALIZE, token_count=1024) + assert snap.session_id == "s1" + assert snap.phase == Phase.LOCALIZE + assert snap.token_count == 1024 + assert snap.anchors == [] + assert snap.compressed_at > 0 + + +def test_tool_call_record_fields(): + rec = ToolCallRecord( + session_id="s1", + phase=Phase.REPAIR, + tool_name="SearchBySymbol", + input_tokens=50, + output_tokens=200, + duration_ms=42, + ) + assert rec.tool_name == "SearchBySymbol" + assert rec.phase == Phase.REPAIR + assert rec.duration_ms == 42