318 lines
10 KiB
Python
318 lines
10 KiB
Python
"""Job handler dispatch + default handlers (AI + logging).
|
|
|
|
Handlers share the following signature::
|
|
|
|
async fn(
|
|
envelope: JobRequestEnvelope,
|
|
*,
|
|
deps: WorkerDependencies,
|
|
) -> HandlerResult
|
|
|
|
WorkerDependencies wraps Mongo collections, the AI Gateway, and logging.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Awaitable, Callable, Protocol
|
|
from uuid import UUID
|
|
|
|
import structlog
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
|
|
from asf_events.base import JobRequestEnvelope
|
|
from asf_domain import JobStatus, JobUpdate, JobUsage
|
|
|
|
from app.config import WorkerSettings # type: ignore[attr-defined]
|
|
|
|
log = structlog.get_logger(__name__)
|
|
|
|
|
|
JobHandlerFn = Callable[["JobRequestEnvelope", "WorkerDependencies"], Awaitable["HandlerResult"]]
|
|
|
|
|
|
@dataclass
|
|
class HandlerResult:
|
|
status: JobStatus
|
|
result: dict[str, Any] = field(default_factory=dict)
|
|
error: str | None = None
|
|
error_detail: dict[str, Any] | None = None
|
|
usage: JobUsage | None = None
|
|
|
|
|
|
@dataclass
|
|
class WorkerDependencies:
|
|
db: AsyncIOMotorDatabase
|
|
jobs: AsyncIOMotorCollection
|
|
ai_executions: AsyncIOMotorCollection
|
|
artifacts: AsyncIOMotorCollection
|
|
settings: WorkerSettings
|
|
ai_gateway: Any | None = None # will be set if enable_ai_handlers=True
|
|
worker_id: str = "asf-worker"
|
|
|
|
|
|
# ==========================================================================
|
|
# Default handlers
|
|
# ==========================================================================
|
|
|
|
|
|
async def _mark_job_done(
|
|
deps: WorkerDependencies,
|
|
job_id: UUID,
|
|
result: HandlerResult,
|
|
) -> None:
|
|
"""Persist the final handler outcome back to MongoDB jobs collection.
|
|
|
|
This is the ONLY place worker-state leaves process memory and hits the
|
|
durable source of truth (MongoDB).
|
|
"""
|
|
updates: dict[str, Any] = {
|
|
"status": result.status.value
|
|
if isinstance(result.status, JobStatus)
|
|
else str(result.status),
|
|
"completed_at": __import__("datetime").datetime.now(
|
|
__import__("datetime").timezone.utc
|
|
),
|
|
"updated_at": __import__("datetime").datetime.now(
|
|
__import__("datetime").timezone.utc
|
|
),
|
|
"worker": deps.worker_id,
|
|
}
|
|
if result.result:
|
|
updates["result"] = result.result
|
|
if result.error:
|
|
updates["error"] = result.error
|
|
if result.error_detail:
|
|
updates["error_detail"] = result.error_detail
|
|
if result.usage is not None:
|
|
updates["usage"] = (
|
|
result.usage.model_dump()
|
|
if hasattr(result.usage, "model_dump")
|
|
else result.usage
|
|
)
|
|
await deps.jobs.update_one({"_id": job_id}, {"$set": updates})
|
|
|
|
|
|
async def noop_handler(
|
|
env: JobRequestEnvelope, deps: WorkerDependencies
|
|
) -> HandlerResult:
|
|
"""Catch-all handler: acknowledges the job and leaves a note.
|
|
|
|
Useful when a queue is not yet staffed with a real handler — ensures we
|
|
never leave work stuck "running" with no outcome.
|
|
"""
|
|
log.warn(
|
|
"no_handler_registered",
|
|
job_id=str(env.job_id),
|
|
queue=env.type,
|
|
payload_keys=list(env.payload.keys()),
|
|
)
|
|
return HandlerResult(
|
|
status=JobStatus.FAILED,
|
|
error=f"No worker handler registered for job_type={env.type!r}",
|
|
error_detail={"hint": "register a handler in app.worker.HANDLERS_BY_TYPE"},
|
|
)
|
|
|
|
|
|
async def echo_handler(
|
|
env: JobRequestEnvelope, deps: WorkerDependencies
|
|
) -> HandlerResult:
|
|
"""Simple echo / dev handler: stores the payload back as result."""
|
|
return HandlerResult(
|
|
status=JobStatus.COMPLETED,
|
|
result={"echo": env.payload, "handled_by": deps.worker_id},
|
|
usage=JobUsage(cpu_seconds=0.001),
|
|
)
|
|
|
|
|
|
async def ai_handler(
|
|
env: JobRequestEnvelope, deps: WorkerDependencies
|
|
) -> HandlerResult:
|
|
"""AI-interactive + AI-background handler.
|
|
|
|
Dispatches to the ASF AI Gateway (MockProvider in bootstrap) for:
|
|
- "ask_asf" kind (chat completion)
|
|
- "analyze_code"
|
|
- "generate_tests"
|
|
- "decompose_requirement"
|
|
|
|
In the bootstrap phase we safely fall back to echo if no gateway is
|
|
attached (e.g. test runs with AI disabled).
|
|
"""
|
|
|
|
kind = env.payload.get("kind", "unknown")
|
|
log.info(
|
|
"ai_handler_start",
|
|
job_id=str(env.job_id),
|
|
kind=kind,
|
|
provider=deps.settings.ai_gateway_provider,
|
|
)
|
|
|
|
if deps.ai_gateway is None:
|
|
return HandlerResult(
|
|
status=JobStatus.COMPLETED,
|
|
result={
|
|
"note": (
|
|
"ai_gateway not attached to this worker. "
|
|
"Install services/asf-ai and wire it in worker."
|
|
),
|
|
"payload": env.payload,
|
|
},
|
|
)
|
|
|
|
from asf_ai import ( # type: ignore[import-not-found]
|
|
AIGateway,
|
|
ChatRequest,
|
|
ChatMessage,
|
|
TraceContext,
|
|
CodeAnalysisRequest,
|
|
TestGenerationRequest,
|
|
RequirementDecomposeRequest,
|
|
)
|
|
|
|
gw: AIGateway = deps.ai_gateway
|
|
trace = TraceContext(
|
|
tenant_id=env.tenant_id,
|
|
project_id=env.project_id,
|
|
job_id=env.job_id,
|
|
correlation_id=env.correlation_id,
|
|
)
|
|
|
|
try:
|
|
if kind == "ask_asf":
|
|
question = env.payload.get("question", "")
|
|
req = ChatRequest(
|
|
trace=trace,
|
|
messages=[ChatMessage(role="user", content=question)],
|
|
context_summary=env.payload.get("context_summary"),
|
|
)
|
|
resp = await gw.chat(req)
|
|
return HandlerResult(
|
|
status=JobStatus.COMPLETED,
|
|
result={
|
|
"kind": "ask_asf",
|
|
"question": question,
|
|
"answer": resp.reply,
|
|
"references": resp.references,
|
|
},
|
|
usage=JobUsage(
|
|
input_tokens=resp.input_tokens,
|
|
output_tokens=resp.output_tokens,
|
|
estimated_cost=resp.estimated_cost,
|
|
),
|
|
)
|
|
|
|
if kind == "analyze_code":
|
|
req = CodeAnalysisRequest(
|
|
trace=trace,
|
|
code_snippet=env.payload.get("code_snippet", ""),
|
|
language=env.payload.get("language"),
|
|
file_path=env.payload.get("file_path"),
|
|
)
|
|
resp = await gw.analyze_code(req)
|
|
return HandlerResult(
|
|
status=JobStatus.COMPLETED,
|
|
result={
|
|
"kind": "analyze_code",
|
|
"summary": resp.summary,
|
|
"findings": [f.model_dump(mode="json") for f in resp.findings],
|
|
},
|
|
usage=JobUsage(
|
|
input_tokens=resp.input_tokens,
|
|
output_tokens=resp.output_tokens,
|
|
estimated_cost=resp.estimated_cost,
|
|
),
|
|
)
|
|
|
|
if kind == "generate_tests":
|
|
req = TestGenerationRequest(
|
|
trace=trace,
|
|
code_snippet=env.payload.get("code_snippet", ""),
|
|
language=env.payload.get("language", "python"),
|
|
test_framework=env.payload.get("framework", "pytest"),
|
|
goal=env.payload.get("goal", "Generate unit tests"),
|
|
)
|
|
resp = await gw.generate_tests(req)
|
|
return HandlerResult(
|
|
status=JobStatus.COMPLETED,
|
|
result={
|
|
"kind": "generate_tests",
|
|
"test_cases": resp.test_cases,
|
|
"fixtures_needed": resp.fixtures_needed,
|
|
"coverage_notes": resp.coverage_notes,
|
|
},
|
|
usage=JobUsage(
|
|
input_tokens=resp.input_tokens,
|
|
output_tokens=resp.output_tokens,
|
|
estimated_cost=resp.estimated_cost,
|
|
),
|
|
)
|
|
|
|
if kind == "decompose_requirement":
|
|
req = RequirementDecomposeRequest(
|
|
trace=trace,
|
|
requirement_text=env.payload.get("requirement_text", ""),
|
|
context=env.payload.get("context"),
|
|
max_subtasks=int(env.payload.get("max_subtasks", 20)),
|
|
)
|
|
resp = await gw.decompose_requirement(req)
|
|
return HandlerResult(
|
|
status=JobStatus.COMPLETED,
|
|
result={
|
|
"kind": "decompose_requirement",
|
|
"subtasks": [s.model_dump(mode="json") for s in resp.subtasks],
|
|
"architecture_notes": resp.architecture_notes,
|
|
"open_questions": resp.open_questions,
|
|
},
|
|
usage=JobUsage(
|
|
input_tokens=resp.input_tokens,
|
|
output_tokens=resp.output_tokens,
|
|
estimated_cost=resp.estimated_cost,
|
|
),
|
|
)
|
|
|
|
return await echo_handler(env, deps)
|
|
|
|
except Exception as exc: # noqa: BLE001
|
|
log.exception("ai_handler_failed", exc=str(exc))
|
|
return HandlerResult(
|
|
status=JobStatus.FAILED,
|
|
error=f"AI handler raised: {type(exc).__name__}: {exc}",
|
|
error_detail={"payload": env.payload},
|
|
)
|
|
|
|
|
|
async def indexing_handler(
|
|
env: JobRequestEnvelope, deps: WorkerDependencies
|
|
) -> HandlerResult:
|
|
"""Stub handler for repository / doc indexing jobs."""
|
|
log.info(
|
|
"indexing_handler_stub",
|
|
job_id=str(env.job_id),
|
|
payload=env.payload,
|
|
)
|
|
return HandlerResult(
|
|
status=JobStatus.COMPLETED,
|
|
result={"stub": True, "indexed_items": 0, "source": env.payload.get("source")},
|
|
usage=JobUsage(cpu_seconds=0.005),
|
|
)
|
|
|
|
|
|
async def notification_handler(
|
|
env: JobRequestEnvelope, deps: WorkerDependencies
|
|
) -> HandlerResult:
|
|
"""Stub handler for email / webhook / in-app notifications."""
|
|
log.info(
|
|
"notification_stub",
|
|
job_id=str(env.job_id),
|
|
payload_keys=list(env.payload.keys()),
|
|
)
|
|
return HandlerResult(
|
|
status=JobStatus.COMPLETED,
|
|
result={"note": "notifications are a stub in bootstrap"},
|
|
)
|
|
|
|
|
|
# Final status persistence step, exported for reuse in the main worker loop.
|
|
mark_job_done = _mark_job_done
|