first
This commit is contained in:
valenti
@@ -0,0 +1,317 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user