118 lines
4.8 KiB
Python
118 lines
4.8 KiB
Python
"""Keyless stub agent for workflow testing.
|
|
|
|
Connects to the orchestrator over WebSocket, registers, and returns canned results WITHOUT
|
|
calling any LLM (no OPENAI_API_KEY needed). To make the master review loop deterministic, the
|
|
"testing" task returns a conflicting framework (unittest) on its FIRST execution and an aligned
|
|
one (pytest) on retry — so the heuristic critic rejects exactly once and then accepts.
|
|
|
|
Run standalone against a manually started orchestrator:
|
|
set "ORCHESTRATOR_URL=ws://localhost:8000"
|
|
set "AGENT_ID=stub-1"
|
|
set "AGENT_CAPABILITIES=python,code_generation,testing,pytest,technical-writing,general"
|
|
python scripts/stub_agent.py
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import time
|
|
|
|
import websockets
|
|
|
|
|
|
class StubAgent:
|
|
def __init__(self, orchestrator_url: str, agent_id: str, capabilities: list[str]):
|
|
self.orchestrator_url = orchestrator_url.rstrip("/")
|
|
self.agent_id = agent_id
|
|
self.capabilities = capabilities
|
|
self.ws = None
|
|
self.runs: dict[str, int] = {} # task_id -> execution count
|
|
self.running = True
|
|
|
|
def _result_for(self, task_id: str, description: str) -> dict:
|
|
count = self.runs.get(task_id, 0) + 1
|
|
self.runs[task_id] = count
|
|
tid = task_id.lower()
|
|
if "implementation" in tid:
|
|
summary = "implemented add(a, b); tests should use pytest"
|
|
elif "testing" in tid:
|
|
# First attempt conflicts (unittest); retry aligns with implementation (pytest).
|
|
summary = "wrote tests with unittest" if count == 1 else "wrote tests with pytest"
|
|
elif "documentation" in tid:
|
|
summary = "documented usage of add(a, b)"
|
|
else:
|
|
summary = f"completed: {description[:40]}"
|
|
return {
|
|
"success": True,
|
|
"task_id": task_id,
|
|
"subtasks": [{"status": "completed", "summary": summary, "changes": summary, "files": []}],
|
|
"awaiting_handoff": False,
|
|
"agent_id": self.agent_id,
|
|
"usage": {
|
|
"model_id": "stub", "model_tokens": 0, "prompt_tokens": 0,
|
|
"completion_tokens": 0, "model_cost_usd": 0.0, "runtime_seconds": 0.0,
|
|
},
|
|
}
|
|
|
|
async def send(self, payload: dict):
|
|
await self.ws.send(json.dumps(payload))
|
|
|
|
async def handle(self, msg: dict):
|
|
msg_type = msg.get("type")
|
|
if msg_type == "task_assignment":
|
|
task_id = msg["task_id"]
|
|
await self.send({"type": "task_accepted", "task_id": task_id, "available_slots": 4})
|
|
await self.send({"type": "task_start", "agent_id": self.agent_id, "task_id": task_id, "timestamp": time.time()})
|
|
await asyncio.sleep(0.2) # simulate work
|
|
result = self._result_for(task_id, msg.get("description", ""))
|
|
await self.send({
|
|
"type": "task_complete", "agent_id": self.agent_id,
|
|
"task_id": task_id, "result": result, "timestamp": time.time(),
|
|
})
|
|
elif msg_type == "peer_message" and not msg.get("is_reply"):
|
|
# Answer an inbound peer query so collaboration round-trips complete.
|
|
await self.send({
|
|
"type": "peer_message", "agent_id": self.agent_id,
|
|
"target_agent_id": msg.get("from_agent_id"), "task_id": msg.get("task_id"),
|
|
"content": f"stub guidance from {self.agent_id}",
|
|
"correlation_id": msg.get("correlation_id"), "is_reply": True, "timestamp": time.time(),
|
|
})
|
|
|
|
async def heartbeat(self):
|
|
while self.running:
|
|
try:
|
|
await self.send({
|
|
"type": "heartbeat", "agent_id": self.agent_id,
|
|
"timestamp": time.time(), "active_tasks": 0, "available_slots": 4,
|
|
})
|
|
except Exception:
|
|
return
|
|
await asyncio.sleep(5)
|
|
|
|
async def run(self):
|
|
async with websockets.connect(f"{self.orchestrator_url}/ws/{self.agent_id}") as ws:
|
|
self.ws = ws
|
|
await self.send({
|
|
"type": "register", "agent_id": self.agent_id,
|
|
"capabilities": self.capabilities, "available_slots": 4, "active_task_ids": [],
|
|
})
|
|
hb = asyncio.create_task(self.heartbeat())
|
|
try:
|
|
async for raw in ws:
|
|
await self.handle(json.loads(raw))
|
|
finally:
|
|
self.running = False
|
|
hb.cancel()
|
|
|
|
|
|
async def _main():
|
|
agent = StubAgent(
|
|
os.getenv("ORCHESTRATOR_URL", "ws://localhost:8000"),
|
|
os.getenv("AGENT_ID", "stub-1"),
|
|
os.getenv("AGENT_CAPABILITIES", "python,code_generation,testing,pytest,technical-writing,general").split(","),
|
|
)
|
|
await agent.run()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(_main())
|