Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions src/hawk/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
SessionDetail,
SessionSummary,
StatsResponse,
ToolResult,
)

DEFAULT_BASE_URL = "http://127.0.0.1:4590"
Expand Down Expand Up @@ -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(
Expand All @@ -134,6 +137,8 @@ def chat(
autonomy=autonomy,
cwd=cwd,
agent=agent,
tools=tools,
tool_results=tool_results,
)

def _do() -> ChatResponse:
Expand All @@ -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(
Expand All @@ -166,6 +173,8 @@ def chat_stream(
autonomy=autonomy,
cwd=cwd,
agent=agent,
tools=tools,
tool_results=tool_results,
)

def _do() -> StreamReader:
Expand All @@ -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

Expand Down Expand Up @@ -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(
Expand All @@ -343,6 +360,8 @@ async def chat(
autonomy=autonomy,
cwd=cwd,
agent=agent,
tools=tools,
tool_results=tool_results,
)

async def _do() -> ChatResponse:
Expand All @@ -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(
Expand All @@ -375,6 +396,8 @@ async def chat_stream(
autonomy=autonomy,
cwd=cwd,
agent=agent,
tools=tools,
tool_results=tool_results,
)

async def _do() -> AsyncStreamReader:
Expand All @@ -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

Expand Down
174 changes: 174 additions & 0 deletions src/hawk/sessions.py
Original file line number Diff line number Diff line change
@@ -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
28 changes: 14 additions & 14 deletions src/hawk/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
if TYPE_CHECKING:
from .types import ChatResponse

from .types import ToolResult


@dataclass
class Tool:
Expand Down Expand Up @@ -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 = (
Expand All @@ -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.
Expand Down Expand Up @@ -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 = (
Expand All @@ -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(
Expand Down
16 changes: 14 additions & 2 deletions src/hawk/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}

Expand Down Expand Up @@ -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}

Expand All @@ -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."""

Expand Down
Loading
Loading