-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsessions.py
More file actions
174 lines (137 loc) · 4.89 KB
/
Copy pathsessions.py
File metadata and controls
174 lines (137 loc) · 4.89 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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