Skip to content

Commit aa01451

Browse files
Patel230Patel230
authored
Merge pull request #21 from GrayCodeAI/chore/sync-quality-hardening-and-sessions
chore(sync): quality-hardening + SDK sessions module, client/tools/types updates
2 parents 48e8d5e + a8beb20 commit aa01451

5 files changed

Lines changed: 327 additions & 16 deletions

File tree

‎src/hawk/client.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
SessionDetail,
2121
SessionSummary,
2222
StatsResponse,
23+
ToolResult,
2324
)
2425

2526
DEFAULT_BASE_URL = "http://127.0.0.1:4590"
@@ -124,6 +125,8 @@ def chat(
124125
autonomy: str | None = None,
125126
cwd: str | None = None,
126127
agent: str | None = None,
128+
tools: list[dict[str, Any]] | None = None,
129+
tool_results: list[ToolResult] | None = None,
127130
) -> ChatResponse:
128131
"""Send a prompt and return the complete response."""
129132
request = ChatRequest(
@@ -134,6 +137,8 @@ def chat(
134137
autonomy=autonomy,
135138
cwd=cwd,
136139
agent=agent,
140+
tools=tools,
141+
tool_results=tool_results,
137142
)
138143

139144
def _do() -> ChatResponse:
@@ -156,6 +161,8 @@ def chat_stream(
156161
autonomy: str | None = None,
157162
cwd: str | None = None,
158163
agent: str | None = None,
164+
tools: list[dict[str, Any]] | None = None,
165+
tool_results: list[ToolResult] | None = None,
159166
) -> StreamReader:
160167
"""Send a prompt and stream the response via SSE."""
161168
request = ChatRequest(
@@ -166,6 +173,8 @@ def chat_stream(
166173
autonomy=autonomy,
167174
cwd=cwd,
168175
agent=agent,
176+
tools=tools,
177+
tool_results=tool_results,
169178
)
170179

171180
def _do() -> StreamReader:
@@ -190,12 +199,18 @@ def _do() -> StreamReader:
190199
def create_session(
191200
self,
192201
name: str | None = None,
202+
model: str | None = None,
203+
cwd: str | None = None,
193204
metadata: dict[str, Any] | None = None,
194205
) -> SessionDetail:
195206
"""Create a new session."""
196207
body: dict[str, Any] = {}
197208
if name is not None:
198209
body["name"] = name
210+
if model is not None:
211+
body["model"] = model
212+
if cwd is not None:
213+
body["cwd"] = cwd
199214
if metadata is not None:
200215
body["metadata"] = metadata
201216

@@ -333,6 +348,8 @@ async def chat(
333348
autonomy: str | None = None,
334349
cwd: str | None = None,
335350
agent: str | None = None,
351+
tools: list[dict[str, Any]] | None = None,
352+
tool_results: list[ToolResult] | None = None,
336353
) -> ChatResponse:
337354
"""Send a prompt and return the complete response."""
338355
request = ChatRequest(
@@ -343,6 +360,8 @@ async def chat(
343360
autonomy=autonomy,
344361
cwd=cwd,
345362
agent=agent,
363+
tools=tools,
364+
tool_results=tool_results,
346365
)
347366

348367
async def _do() -> ChatResponse:
@@ -365,6 +384,8 @@ async def chat_stream(
365384
autonomy: str | None = None,
366385
cwd: str | None = None,
367386
agent: str | None = None,
387+
tools: list[dict[str, Any]] | None = None,
388+
tool_results: list[ToolResult] | None = None,
368389
) -> AsyncStreamReader:
369390
"""Send a prompt and stream the response via SSE."""
370391
request = ChatRequest(
@@ -375,6 +396,8 @@ async def chat_stream(
375396
autonomy=autonomy,
376397
cwd=cwd,
377398
agent=agent,
399+
tools=tools,
400+
tool_results=tool_results,
378401
)
379402

380403
async def _do() -> AsyncStreamReader:
@@ -399,12 +422,18 @@ async def _do() -> AsyncStreamReader:
399422
async def create_session(
400423
self,
401424
name: str | None = None,
425+
model: str | None = None,
426+
cwd: str | None = None,
402427
metadata: dict[str, Any] | None = None,
403428
) -> SessionDetail:
404429
"""Create a new session."""
405430
body: dict[str, Any] = {}
406431
if name is not None:
407432
body["name"] = name
433+
if model is not None:
434+
body["model"] = model
435+
if cwd is not None:
436+
body["cwd"] = cwd
408437
if metadata is not None:
409438
body["metadata"] = metadata
410439

‎src/hawk/sessions.py‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
"""Phase-attributed session state for multi-phase agent pipelines.
2+
3+
Mirrors hawk-core-contracts/sessions/sessions.go so Python agents can track
4+
token spend and context snapshots with the same structure as Go agents.
5+
6+
Phase attribution data shows that code review alone consumes 59.4% of tokens
7+
(Tokenomics paper, arXiv 2601.14470). These dataclasses let callers record
8+
per-phase usage and surface it in dashboards without coupling to the Go runtime.
9+
"""
10+
11+
from __future__ import annotations
12+
13+
import time
14+
from dataclasses import dataclass, field
15+
from enum import Enum
16+
17+
18+
class Phase(str, Enum):
19+
"""Phase identifies a step in the localize → repair → validate pipeline."""
20+
21+
LOCALIZE = "localize"
22+
REPAIR = "repair"
23+
VALIDATE = "validate"
24+
REVIEW = "review"
25+
PLANNING = "planning"
26+
UNKNOWN = ""
27+
28+
@classmethod
29+
def parse(cls, s: str) -> Phase:
30+
"""Return the Phase matching s, or Phase.UNKNOWN if unrecognised."""
31+
for member in cls:
32+
if member.value == s:
33+
return member
34+
return cls.UNKNOWN
35+
36+
37+
@dataclass
38+
class PhaseUsage:
39+
"""Token usage totals for a single pipeline phase."""
40+
41+
input_tokens: int = 0
42+
output_tokens: int = 0
43+
total_tokens: int = 0
44+
cost_usd: float = 0.0
45+
46+
47+
@dataclass
48+
class ContextSnapshot:
49+
"""A point-in-time snapshot of context window state within a phase."""
50+
51+
session_id: str
52+
phase: Phase
53+
token_count: int
54+
compressed_at: float = field(default_factory=time.time) # Unix timestamp
55+
summary: str = ""
56+
anchors: list[str] = field(default_factory=list)
57+
58+
59+
@dataclass
60+
class ToolCallRecord:
61+
"""A record of a single tool call with phase attribution and token counts."""
62+
63+
session_id: str
64+
phase: Phase
65+
tool_name: str
66+
input_tokens: int = 0
67+
output_tokens: int = 0
68+
duration_ms: int = 0
69+
error: str = ""
70+
timestamp: float = field(default_factory=time.time) # Unix timestamp
71+
72+
73+
@dataclass
74+
class CostAccumulator:
75+
"""Accumulates per-phase token usage for a session.
76+
77+
Usage::
78+
79+
acc = CostAccumulator(session_id="sess-abc")
80+
acc.add(Phase.LOCALIZE, input_tokens=500, output_tokens=100, cost_usd=0.002)
81+
acc.add(Phase.REVIEW, input_tokens=3000, output_tokens=600, cost_usd=0.012)
82+
print(acc.phase_share(Phase.REVIEW)) # ~0.857
83+
"""
84+
85+
session_id: str
86+
by_phase: dict[Phase, PhaseUsage] = field(default_factory=dict)
87+
total_tokens: int = 0
88+
total_cost_usd: float = 0.0
89+
90+
def add(
91+
self,
92+
phase: Phase,
93+
input_tokens: int,
94+
output_tokens: int,
95+
cost_usd: float = 0.0,
96+
) -> None:
97+
"""Record token usage for phase."""
98+
if phase not in self.by_phase:
99+
self.by_phase[phase] = PhaseUsage()
100+
pu = self.by_phase[phase]
101+
pu.input_tokens += input_tokens
102+
pu.output_tokens += output_tokens
103+
pu.total_tokens += input_tokens + output_tokens
104+
pu.cost_usd += cost_usd
105+
self.total_tokens += input_tokens + output_tokens
106+
self.total_cost_usd += cost_usd
107+
108+
def phase_share(self, phase: Phase) -> float:
109+
"""Return the fraction of total tokens consumed by phase (0.0-1.0)."""
110+
if self.total_tokens == 0:
111+
return 0.0
112+
usage = self.by_phase.get(phase)
113+
if usage is None:
114+
return 0.0
115+
return usage.total_tokens / self.total_tokens
116+
117+
def __str__(self) -> str:
118+
parts = [
119+
f"session={self.session_id} total={self.total_tokens}tok ${self.total_cost_usd:.4f}"
120+
]
121+
for phase, pu in sorted(
122+
self.by_phase.items(), key=lambda kv: kv[1].total_tokens, reverse=True
123+
):
124+
share = self.phase_share(phase)
125+
parts.append(
126+
f" {phase.value or 'unknown'}: {pu.total_tokens}tok ({share:.1%}) ${pu.cost_usd:.4f}"
127+
)
128+
return "\n".join(parts)
129+
130+
131+
@dataclass
132+
class ResolutionRequest:
133+
"""SDK-level request to run a multi-phase resolution against a repository."""
134+
135+
session_id: str = ""
136+
root_dir: str = ""
137+
query: str = ""
138+
max_files: int = 0
139+
max_symbols: int = 0
140+
language: str = ""
141+
142+
143+
@dataclass
144+
class PatchCandidate:
145+
"""A single proposed code change returned by the repair phase."""
146+
147+
file_path: str = ""
148+
symbol: str = ""
149+
original_body: str = ""
150+
patched_body: str = ""
151+
confidence: float = 0.0
152+
153+
154+
@dataclass
155+
class PhaseMetrics:
156+
"""Token spend and timing for a single pipeline phase."""
157+
158+
phase: Phase = Phase.UNKNOWN
159+
input_tokens: int = 0
160+
output_tokens: int = 0
161+
elapsed_ms: int = 0
162+
163+
164+
@dataclass
165+
class ResolutionResult:
166+
"""Complete output of a multi-phase resolution run."""
167+
168+
session_id: str = ""
169+
candidates: list[PatchCandidate] = field(default_factory=list)
170+
validation_passed: bool = False
171+
validation_failures: list[str] = field(default_factory=list)
172+
phase_metrics: list[PhaseMetrics] = field(default_factory=list)
173+
total_input_tokens: int = 0
174+
total_output_tokens: int = 0

‎src/hawk/tools.py‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
if TYPE_CHECKING:
1111
from .types import ChatResponse
1212

13+
from .types import ToolResult
14+
1315

1416
@dataclass
1517
class Tool:
@@ -135,7 +137,7 @@ def chat_with_tools(
135137
break
136138

137139
# Execute each tool call and collect results.
138-
tool_results = []
140+
tool_results: list[ToolResult] = []
139141
for tc in response.tool_calls:
140142
tool_name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", None)
141143
arguments = (
@@ -146,12 +148,11 @@ def chat_with_tools(
146148
else:
147149
result = json.dumps({"error": f"Unknown tool: {tool_name}"})
148150
tool_results.append(
149-
{
150-
"tool_use_id": tc.get("id", "")
151-
if isinstance(tc, dict)
152-
else getattr(tc, "id", ""),
153-
"content": result,
154-
}
151+
ToolResult(
152+
tool_use_id=tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", ""),
153+
content=result,
154+
is_error=bool(tool_name) and tool_name not in tool_map,
155+
)
155156
)
156157

157158
# Send tool results back to continue the conversation.
@@ -207,7 +208,7 @@ async def chat_with_tools_async(
207208
if not hasattr(response, "tool_calls") or not response.tool_calls:
208209
break
209210

210-
tool_results = []
211+
tool_results: list[ToolResult] = []
211212
for tc in response.tool_calls:
212213
tool_name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", None)
213214
arguments = (
@@ -218,12 +219,11 @@ async def chat_with_tools_async(
218219
else:
219220
result = json.dumps({"error": f"Unknown tool: {tool_name}"})
220221
tool_results.append(
221-
{
222-
"tool_use_id": tc.get("id", "")
223-
if isinstance(tc, dict)
224-
else getattr(tc, "id", ""),
225-
"content": result,
226-
}
222+
ToolResult(
223+
tool_use_id=tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", ""),
224+
content=result,
225+
is_error=bool(tool_name) and tool_name not in tool_map,
226+
)
227227
)
228228

229229
response = await client.chat(

‎src/hawk/types.py‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ class ChatRequest(BaseModel):
1919
autonomy: str | None = None
2020
cwd: str | None = None
2121
agent: str | None = None
22+
tools: list[dict[str, Any]] | None = None
23+
tool_results: list[ToolResult] | None = Field(None, alias="tool_results")
2224

2325
model_config = {"populate_by_name": True}
2426

@@ -81,8 +83,8 @@ class Message(BaseModel):
8183

8284
role: str
8385
content: str | None = None
84-
tool_use: Any | None = Field(None, alias="tool_use")
85-
tool_result: Any | None = Field(None, alias="tool_result")
86+
tool_use: list[ToolCall] | None = Field(None, alias="tool_use")
87+
tool_result: list[ToolResult] | None = Field(None, alias="tool_result")
8688

8789
model_config = {"populate_by_name": True}
8890

@@ -95,6 +97,16 @@ class ToolCall(BaseModel):
9597
arguments: dict[str, Any] = Field(default_factory=dict)
9698

9799

100+
class ToolResult(BaseModel):
101+
"""Result of executing a tool call."""
102+
103+
tool_use_id: str = Field(alias="tool_use_id")
104+
content: str
105+
is_error: bool = Field(False, alias="is_error")
106+
107+
model_config = {"populate_by_name": True}
108+
109+
98110
class Usage(BaseModel):
99111
"""Token usage information."""
100112

0 commit comments

Comments
 (0)