mjalmoudy94
581f703910
Phase 3 of the AI-workforce engine (TODO-agent-2): - engine/manager.py: ManagerTurns.chat (CREATE_TASK/PAUSE/RESUME/CANCEL, HIRE/UPDATE_EMPLOYEE/FIRE, KNOWLEDGE), triage (ASSIGN/REJECT/ASK_USER), final_review (COMPLETE/REVISE +cycle), error_review (RESOLVE_ERROR/CANCEL), bootstrap_team; post_user_message now delegates to it. - One billing path: call_model/BillingScope in runner.py is the single model entry point (fixes the agybot chat-billing bug). - Team changes capped at 3/turn, employee cap, last manager cannot be fired, hires pick a granted model of the requested tier strictly. - container_engine.py wires the real WorkforceBootstrapPort. - Fixes: REVISE is a body tag (prompt and parser now agree); ASSIGN employee="" means pick-for-me; the runner's protocol-error retry is capped by engine_max_protocol_retries (was unbounded until the turn timeout). - 20 new manager tests + regression tests; report in docs/plan/reports/agent-2-phase-3.md.
260 lines
8.4 KiB
Python
260 lines
8.4 KiB
Python
"""Fakes for the engine tests: a scripted model and a scripted sandbox.
|
|
|
|
Two rules the engine tests insist on: the model never answers by calling a
|
|
real provider, and nothing ever executes on the host. Both fakes also *record*
|
|
what they were given, because half of what the engine must guarantee is about
|
|
prompts (no secret ever reaches one) and commands (git never runs).
|
|
|
|
``ScriptedModel.complete`` returns the queued responses in order and raises
|
|
the queued exceptions, so a test can script "read a file, run the tests, report"
|
|
exactly the way a real agent would behave — including the failures.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from app.application.ports import ChatMessage, CommandResult, ModelHealth, ModelResponse
|
|
from app.application.ports.model_grants import GrantedModel
|
|
from app.application.ports.sandbox_provider import (
|
|
SessionInfo,
|
|
SnapshotInfo,
|
|
WorkspaceInfo,
|
|
)
|
|
|
|
|
|
@dataclass(slots=True, kw_only=True)
|
|
class ScriptedReply:
|
|
"""One queued model answer: text, or an exception to raise instead."""
|
|
|
|
content: str = ""
|
|
raises: BaseException | None = None
|
|
input_tokens: int = 10
|
|
output_tokens: int = 5
|
|
|
|
|
|
class ScriptedModel:
|
|
"""A ``ModelProviderPort`` that replays a script.
|
|
|
|
``seen`` holds every call it received (system prompt + messages), so a test
|
|
can assert what the model was actually told — the no-secrets rule is only
|
|
testable this way.
|
|
"""
|
|
|
|
def __init__(self, replies: list[ScriptedReply] | None = None) -> None:
|
|
self.replies: list[ScriptedReply] = list(replies or [])
|
|
self.seen: list[dict] = []
|
|
self.calls = 0
|
|
|
|
def queue(self, content: str, **kwargs) -> None:
|
|
self.replies.append(ScriptedReply(content=content, **kwargs))
|
|
|
|
async def complete(
|
|
self,
|
|
*,
|
|
model,
|
|
messages: list[ChatMessage],
|
|
system: str = "",
|
|
temperature: float = 0.2,
|
|
max_tokens: int = 2048,
|
|
credential=None,
|
|
credentials=None,
|
|
) -> ModelResponse:
|
|
self.calls += 1
|
|
self.seen.append(
|
|
{
|
|
"model": model,
|
|
"system": system,
|
|
"messages": list(messages),
|
|
"credential": credential,
|
|
}
|
|
)
|
|
if not self.replies:
|
|
return ModelResponse(
|
|
content="[[WORK_REPORT]]پاسخ پیشفرض[[/WORK_REPORT]]",
|
|
model_internal_name=getattr(model, "internal_name", ""),
|
|
input_tokens=1,
|
|
output_tokens=1,
|
|
)
|
|
reply = self.replies.pop(0)
|
|
if reply.raises is not None:
|
|
raise reply.raises
|
|
return ModelResponse(
|
|
content=reply.content,
|
|
model_internal_name=getattr(model, "internal_name", ""),
|
|
input_tokens=reply.input_tokens,
|
|
output_tokens=reply.output_tokens,
|
|
)
|
|
|
|
async def health(self) -> ModelHealth:
|
|
return ModelHealth(ok=True, detail="scripted")
|
|
|
|
@property
|
|
def prompts(self) -> list[str]:
|
|
"""Every system prompt the model was given, in order."""
|
|
return [call["system"] for call in self.seen]
|
|
|
|
def all_text(self) -> str:
|
|
"""Everything the model saw, for the "no secret" assertion."""
|
|
parts = [call["system"] for call in self.seen]
|
|
for call in self.seen:
|
|
parts.extend(message.content for message in call["messages"])
|
|
return "\n".join(parts)
|
|
|
|
|
|
class FakeSandbox:
|
|
"""A ``SandboxProviderPort`` over a dict filesystem and scripted commands.
|
|
|
|
``commands`` maps the command *prefix* to a queue of results, so a test can
|
|
make the tests red once and green afterwards: the engine must notice.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
files: dict[str, str] | None = None,
|
|
commands: dict[str, list[CommandResult]] | None = None,
|
|
) -> None:
|
|
self.files: dict[str, str] = dict(files or {})
|
|
self.commands: dict[str, list[CommandResult]] = {
|
|
key: list(value) for key, value in (commands or {}).items()
|
|
}
|
|
self.executed: list[dict] = []
|
|
self.writes: list[dict] = []
|
|
self.reads: list[str] = []
|
|
|
|
# -- SandboxProviderPort ------------------------------------------------
|
|
|
|
async def create_workspace(
|
|
self, *, organization_id: str, project_id: str, template_id: str, image: str
|
|
) -> WorkspaceInfo:
|
|
return WorkspaceInfo(external_id="ws-1", status="ready")
|
|
|
|
async def start_session(self, *, workspace_external_id: str) -> SessionInfo:
|
|
return SessionInfo(external_id="sess-1", status="running")
|
|
|
|
async def run_command(
|
|
self,
|
|
*,
|
|
session_external_id: str,
|
|
command: str,
|
|
timeout_seconds: int,
|
|
max_output_bytes: int,
|
|
env: dict[str, str] | None = None,
|
|
) -> CommandResult:
|
|
self.executed.append(
|
|
{
|
|
"command": command,
|
|
"timeout_seconds": timeout_seconds,
|
|
"max_output_bytes": max_output_bytes,
|
|
"session": session_external_id,
|
|
}
|
|
)
|
|
for prefix, queue in self.commands.items():
|
|
if command.startswith(prefix) and queue:
|
|
return queue.pop(0)
|
|
return CommandResult(exit_code=0, stdout="ok")
|
|
|
|
async def read_file(self, *, session_external_id: str, path: str) -> str:
|
|
self.reads.append(path)
|
|
if path not in self.files:
|
|
raise FileNotFoundError(path)
|
|
return self.files[path]
|
|
|
|
async def write_file(self, *, session_external_id: str, path: str, content: str) -> None:
|
|
self.files[path] = content
|
|
self.writes.append({"path": path, "content": content})
|
|
|
|
async def snapshot(self, *, session_external_id: str, name: str) -> SnapshotInfo:
|
|
return SnapshotInfo(external_id="snap-1", name=name)
|
|
|
|
async def pause(self, *, session_external_id: str) -> None:
|
|
return None
|
|
|
|
async def resume(self, *, session_external_id: str) -> SessionInfo:
|
|
return SessionInfo(external_id="sess-1", status="running")
|
|
|
|
async def destroy(self, *, session_external_id: str) -> None:
|
|
return None
|
|
|
|
async def get_preview_url(self, *, session_external_id: str, port: int) -> str:
|
|
return f"http://127.0.0.1:{port}"
|
|
|
|
|
|
class FakeGrants:
|
|
"""A ``ModelGrantPort`` with an explicit allow-list.
|
|
|
|
``granted`` is either a list of model ids (all ``balanced``) or a mapping
|
|
of id → neutral tier, so a test can have a "strong" and a "fast" model and
|
|
assert the engine picked the right one.
|
|
"""
|
|
|
|
def __init__(self, granted: list[str] | dict[str, str] | None = None) -> None:
|
|
if granted is None:
|
|
self._tiers: dict[str, str] = {
|
|
"model-strong": "strong",
|
|
"model-fast": "fast",
|
|
}
|
|
elif isinstance(granted, dict):
|
|
self._tiers = dict(granted)
|
|
else:
|
|
self._tiers = dict.fromkeys(granted, "balanced")
|
|
self.checks: list[str] = []
|
|
|
|
@property
|
|
def granted(self) -> list[str]:
|
|
return list(self._tiers)
|
|
|
|
async def granted_models(self, *, organization_id: str) -> list[GrantedModel]:
|
|
return [
|
|
GrantedModel(ai_model_id=model_id, tier=tier)
|
|
for model_id, tier in self._tiers.items()
|
|
]
|
|
|
|
async def is_granted(self, *, organization_id: str, ai_model_id: str) -> bool:
|
|
self.checks.append(ai_model_id)
|
|
return ai_model_id in self._tiers
|
|
|
|
def grant(self, model_id: str, tier: str = "balanced") -> None:
|
|
self._tiers[model_id] = tier
|
|
|
|
|
|
class RecordingMemory:
|
|
"""A ``MemoryPort`` that records its queries and returns fixed context."""
|
|
|
|
def __init__(self, context: str = "") -> None:
|
|
self.context = context
|
|
self.queries: list[dict] = []
|
|
|
|
async def build_context(
|
|
self, *, organization_id: str, role: str, query: str, budget_chars: int = 2500
|
|
) -> str:
|
|
self.queries.append(
|
|
{"organization_id": organization_id, "role": role, "query": query}
|
|
)
|
|
return self.context
|
|
|
|
async def record(
|
|
self,
|
|
*,
|
|
organization_id: str,
|
|
kind: str,
|
|
content: str,
|
|
source: str,
|
|
evidence_task_id: str | None = None,
|
|
key: str | None = None,
|
|
) -> str | None:
|
|
return None
|
|
|
|
async def on_task_finished(self, *, organization_id: str, task) -> None:
|
|
return None
|
|
|
|
|
|
__all__ = [
|
|
"FakeGrants",
|
|
"FakeSandbox",
|
|
"RecordingMemory",
|
|
"ScriptedModel",
|
|
"ScriptedReply",
|
|
]
|