Setting the file. One moment. — line 107
This file
- Number
- 7.7
- Position
- 7 of 18
- Type
- Python
- Size
- 10 KB
- Lines
- 283
references/multi-turn-simulation/harbor_example.py
Python·283 lines·10 KB
12from harbor.environments.base import BaseEnvironment
13from harbor.models.agent.context import AgentContext
14from harbor.models.trajectories import Agent, FinalMetrics, Step, Trajectory
15
16from .model_user import ModelUser
17from .runner import (
18 ConversationResult,
19 ConversationRunError,
20 run_llm_user_conversation,
21 run_scripted_conversation,
22)
23
24
25MAX_EVIDENCE_CHARS = 1_000_000
26
27
28def _timestamp() -> str:
29 return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
30
31
32@dataclass(frozen=True)
33class HarnessReply:
34 """One visible response and directly observed evidence for a Harness turn."""
35
36 message: str
37 evidence: Mapping[str, object]
38
39
40class HarnessSession(Protocol):
41 session_id: str
42 public_config: Mapping[str, object]
43
44 async def send(self, user_message: str) -> HarnessReply: ...
45
46
47class RecordingSession:
48 """Adapt a repository Harness session to the conversation runner."""
49
50 def __init__(self, harness: HarnessSession) -> None:
51 self.harness = harness
52 self.events: list[dict[str, object]] = []
53 self.exchanges: list[dict[str, object]] = []
54
55 async def send(self, user_message: str) -> str:
56 self.events.append(
57 {"role": "user", "content": user_message, "timestamp": _timestamp()}
58 )
59 try:
60 reply = await self.harness.send(user_message)
61 if not isinstance(reply, HarnessReply):
62 raise TypeError("Harness session must return HarnessReply")
63 except Exception as error:
64 self.exchanges.append(
65 {"user_message": user_message, "error": type(error).__name__}
66 )
67 raise
68 encoded_evidence = json.dumps(dict(reply.evidence))
69 if len(encoded_evidence) > MAX_EVIDENCE_CHARS:
70 raise ValueError("Harness evidence is too large")
71 evidence = json.loads(encoded_evidence)
72 self.events.append(
73 {
74 "role": "assistant",
75 "content": reply.message,
76 "evidence": evidence,
77 "timestamp": _timestamp(),
78 }
79 )
80 self.exchanges.append(
81 {
82 "user_message": user_message,
83 "assistant_message": reply.message,
84 "evidence": evidence,
85 }
86 )
87 return reply.message
88
89
90class MultiTurnHarborAgent(BaseAgent):
91 """Subclass this adapter and implement the repository-specific bindings."""
92
93 SUPPORTS_ATIF = True
94 MAX_TURNS = 8
95 SIMULATOR_MODEL: Optional[str] = None
96
97 @staticmethod
98 def name() -> str:
99 return "multi-turn-harness"
100
101 def version(self) -> str:
102 return "1.0.0"
103
104 async def setup(self, environment: BaseEnvironment) -> None:
105 return None
106
107 async def create_harness_session(
108 self, environment: BaseEnvironment, context: AgentContext
109 ) -> HarnessSession:
110 raise NotImplementedError
111
112 def scripted_followups(self) -> Optional[Tuple[str, ...]]:
113 return None
114
115 def user_contract(self) -> Optional[str]:
116 return None
117
118 async def call_user_model(self, system: str, payload: str) -> str:
119 raise NotImplementedError
120
121 async def read_user_observation(
122 self, environment: BaseEnvironment
123 ) -> Mapping[str, object]:
124 return {}
125
126 async def run(
127 self, instruction: str, environment: BaseEnvironment, context: AgentContext
128 ) -> None:
129 result: Optional[ConversationResult] = None
130 model_user: Optional[ModelUser] = None
131 session: Optional[RecordingSession] = None
132 error_type: Optional[str] = None
133 contract: Optional[str] = None
134 run_error: Optional[BaseException] = None
135
136 try:
137 session = RecordingSession(
138 await self.create_harness_session(environment, context)
139 )
140 followups = self.scripted_followups()
141 contract = self.user_contract()
142 if (followups is None) == (contract is None):
143 raise ValueError("choose exactly one of scripted_followups or user_contract")
144 if followups is not None:
145 result = await run_scripted_conversation(
146 first_message=instruction,
147 followups=followups,
148 session=session,
149 )
150 else:
151 if not self.SIMULATOR_MODEL:
152 raise ValueError("set SIMULATOR_MODEL for an LLM user")
153 model_user = ModelUser(
154 contract=contract or "",
155 call_model=self.call_user_model,
156 read_observation=lambda: self.read_user_observation(environment),
157 )
158 result = await run_llm_user_conversation(
159 first_message=instruction,
160 session=session,
161 simulated_user=model_user,
162 max_turns=self.MAX_TURNS,
163 )
164 except ConversationRunError as error:
165 result = error.result
166 error_type = type(error).__name__
167 run_error = error
168 except asyncio.CancelledError as error:
169 error_type = type(error).__name__
170 run_error = error
171 except Exception as error:
172 error_type = type(error).__name__
173 run_error = error
174
175 interaction = self._interaction(
176 instruction, session, result, model_user, contract, error_type
177 )
178 artifact_errors: list[dict[str, str]] = []
179 interaction_path = self.logs_dir / "interaction.json"
180 trajectory_path = self.logs_dir / "trajectory.json"
181
182 try:
183 self.logs_dir.mkdir(parents=True, exist_ok=True)
184 interaction_path.write_text(
185 json.dumps(interaction, indent=2), encoding="utf-8"
186 )
187 except Exception as error:
188 artifact_errors.append(
189 {"operation": "write interaction", "error": type(error).__name__}
190 )
191
192 try:
193 trajectory_path.write_text(
194 json.dumps(self._trajectory(interaction).to_json_dict(), indent=2),
195 encoding="utf-8",
196 )
197 except Exception as error:
198 artifact_errors.append(
199 {"operation": "write trajectory", "error": type(error).__name__}
200 )
201
202 if artifact_errors and interaction_path.is_file():
203 interaction["artifact_errors"] = artifact_errors
204 try:
205 interaction_path.write_text(
206 json.dumps(interaction, indent=2), encoding="utf-8"
207 )
208 except Exception:
209 pass
210
211 if run_error is not None:
212 raise run_error.with_traceback(run_error.__traceback__)
213 if artifact_errors:
214 operations = ", ".join(item["operation"] for item in artifact_errors)
215 raise RuntimeError(f"artifact handling failed: {operations}")
216
217 def _interaction(
218 self,
219 instruction: str,
220 session: Optional[RecordingSession],
221 result: Optional[ConversationResult],
222 model_user: Optional[ModelUser],
223 contract: Optional[str],
224 error_type: Optional[str],
225 ) -> dict[str, object]:
226 harness = session.harness if session else None
227 turns = [asdict(turn) for turn in result.turns] if result else []
228 if turns and session:
229 exchanges = iter(session.exchanges)
230 for turn in turns:
231 if turn["role"] != "assistant":
232 continue
233 exchange = next(exchanges, {})
234 if "evidence" in exchange:
235 turn["evidence"] = exchange["evidence"]
236 return {
237 "instruction": instruction,
238 "session_id": harness.session_id if harness else None,
239 "harness_config": json.loads(json.dumps(dict(harness.public_config)))
240 if harness
241 else None,
242 "turns": turns or (session.events if session else []),
243 "harness_exchanges": session.exchanges if session else [],
244 "simulation": {
245 "mode": "llm_user" if contract is not None else "scripted",
246 "model": self.SIMULATOR_MODEL if contract is not None else None,
247 "user_contract": contract,
248 },
249 "simulator_decisions": model_user.records if model_user else [],
250 "termination": result.termination if result else "error",
251 "error": {"type": error_type} if error_type else None,
252 }
253
254 def _trajectory(self, interaction: Mapping[str, object]) -> Trajectory:
255 raw_turns = interaction["turns"]
256 turns = raw_turns if isinstance(raw_turns, list) else []
257 steps: list[Step] = []
258 for index, turn in enumerate(turns, 1):
259 role = turn["role"]
260 extra = {
261 key: value
262 for key, value in turn.items()
263 if key not in {"role", "content", "timestamp"} and value is not None
264 }
265 values = {
266 "step_id": index,
267 "source": "agent" if role == "assistant" else "user",
268 "message": turn["content"],
269 "extra": extra or None,
270 }
271 if turn.get("timestamp"):
272 values["timestamp"] = turn["timestamp"]
273 steps.append(Step(**values))
274 if not steps:
275 steps.append(Step(step_id=1, source="user", message=interaction["instruction"]))
276 return Trajectory(
277 schema_version="ATIF-v1.7",
278 session_id=interaction.get("session_id"),
279 agent=Agent(name=self.name(), version=self.version()),
280 steps=steps,
281 final_metrics=FinalMetrics(total_steps=len(steps)),
282 extra={"termination": interaction["termination"], "error": interaction["error"]},
283 )