forked from sagnik/Project_Velocity
254 lines
8.3 KiB
Python
254 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
import asyncpg
|
|
|
|
|
|
_JSON_KEYS = {
|
|
"context_refs",
|
|
"requested_outputs",
|
|
"payload",
|
|
"input",
|
|
"output",
|
|
"citations",
|
|
"before_state",
|
|
"after_state",
|
|
"detail",
|
|
}
|
|
|
|
|
|
def _row_dict(row: asyncpg.Record | None) -> dict[str, Any] | None:
|
|
if row is None:
|
|
return None
|
|
data = dict(row)
|
|
for key in _JSON_KEYS:
|
|
if isinstance(data.get(key), str):
|
|
data[key] = json.loads(data[key])
|
|
return data
|
|
|
|
|
|
class ColonyRepository:
|
|
def __init__(self, pool: asyncpg.Pool) -> None:
|
|
self.pool = pool
|
|
|
|
async def create_mission(self, mission: dict[str, Any]) -> dict[str, Any]:
|
|
async with self.pool.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
INSERT INTO colony_missions (
|
|
mission_id, tenant_id, mission_type, origin_surface, actor_id,
|
|
actor_role, risk_level, sensitivity_class, status, review_status,
|
|
time_budget_ms, token_budget, user_goal, normalized_goal,
|
|
context_refs, requested_outputs, payload, created_at, updated_at
|
|
) VALUES (
|
|
$1::uuid, $2, $3, $4, $5,
|
|
$6, $7, $8, 'pending', NULL,
|
|
$9, $10, $11, $12,
|
|
$13::jsonb, $14::jsonb, $15::jsonb, NOW(), NOW()
|
|
)
|
|
RETURNING *
|
|
""",
|
|
mission["mission_id"],
|
|
mission["tenant_id"],
|
|
mission["mission_type"],
|
|
mission["origin_surface"],
|
|
mission["actor_id"],
|
|
mission.get("actor_role"),
|
|
mission["risk_level"],
|
|
mission["sensitivity_class"],
|
|
mission["time_budget_ms"],
|
|
mission["token_budget"],
|
|
mission["user_goal"],
|
|
mission["normalized_goal"],
|
|
json.dumps(mission["context_refs"]),
|
|
json.dumps(mission["requested_outputs"]),
|
|
json.dumps(mission["payload"]),
|
|
)
|
|
return _row_dict(row)
|
|
|
|
async def get_mission(self, mission_id: str, tenant_id: str) -> dict[str, Any] | None:
|
|
async with self.pool.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
SELECT *
|
|
FROM colony_missions
|
|
WHERE mission_id = $1::uuid
|
|
AND tenant_id = $2
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
)
|
|
return _row_dict(row)
|
|
|
|
async def update_status(
|
|
self,
|
|
mission_id: str,
|
|
tenant_id: str,
|
|
status: str,
|
|
*,
|
|
review_status: str | None = None,
|
|
) -> dict[str, Any] | None:
|
|
async with self.pool.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
UPDATE colony_missions
|
|
SET status = $3,
|
|
review_status = COALESCE($4, review_status),
|
|
updated_at = NOW(),
|
|
completed_at = CASE
|
|
WHEN $3 IN ('completed', 'failed', 'dispatch_failed') THEN NOW()
|
|
ELSE completed_at
|
|
END
|
|
WHERE mission_id = $1::uuid
|
|
AND tenant_id = $2
|
|
RETURNING *
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
status,
|
|
review_status,
|
|
)
|
|
return _row_dict(row)
|
|
|
|
async def list_missions(self, tenant_id: str, *, limit: int, offset: int) -> list[dict[str, Any]]:
|
|
async with self.pool.acquire() as conn:
|
|
rows = await conn.fetch(
|
|
"""
|
|
SELECT *
|
|
FROM colony_missions
|
|
WHERE tenant_id = $1
|
|
ORDER BY created_at DESC
|
|
LIMIT $2 OFFSET $3
|
|
""",
|
|
tenant_id,
|
|
limit,
|
|
offset,
|
|
)
|
|
return [_row_dict(row) for row in rows]
|
|
|
|
async def artifacts(self, mission_id: str, tenant_id: str) -> dict[str, list[dict[str, Any]]]:
|
|
async with self.pool.acquire() as conn:
|
|
mission = await conn.fetchrow(
|
|
"SELECT mission_id FROM colony_missions WHERE mission_id = $1::uuid AND tenant_id = $2",
|
|
mission_id,
|
|
tenant_id,
|
|
)
|
|
if mission is None:
|
|
return {}
|
|
tasks = await conn.fetch(
|
|
"""
|
|
SELECT *
|
|
FROM colony_tasks
|
|
WHERE mission_id = $1::uuid
|
|
AND tenant_id = $2
|
|
ORDER BY created_at ASC
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
)
|
|
results = await conn.fetch(
|
|
"""
|
|
SELECT *
|
|
FROM colony_worker_results
|
|
WHERE mission_id = $1::uuid
|
|
AND tenant_id = $2
|
|
ORDER BY created_at ASC
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
)
|
|
proposals = await conn.fetch(
|
|
"""
|
|
SELECT *
|
|
FROM colony_writeback_proposals
|
|
WHERE mission_id = $1::uuid
|
|
AND tenant_id = $2
|
|
ORDER BY created_at DESC
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
)
|
|
return {
|
|
"tasks": [_row_dict(row) for row in tasks],
|
|
"worker_results": [_row_dict(row) for row in results],
|
|
"writeback_proposals": [_row_dict(row) for row in proposals],
|
|
}
|
|
|
|
async def pending_writeback_proposals(self, mission_id: str, tenant_id: str) -> list[dict[str, Any]]:
|
|
async with self.pool.acquire() as conn:
|
|
rows = await conn.fetch(
|
|
"""
|
|
SELECT *
|
|
FROM colony_writeback_proposals
|
|
WHERE mission_id = $1::uuid
|
|
AND tenant_id = $2
|
|
AND approval_status = 'pending'
|
|
ORDER BY created_at ASC
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
)
|
|
return [_row_dict(row) for row in rows]
|
|
|
|
async def approve_pending_writebacks(self, mission_id: str, tenant_id: str, actor_id: str) -> int:
|
|
async with self.pool.acquire() as conn:
|
|
result = await conn.execute(
|
|
"""
|
|
UPDATE colony_writeback_proposals
|
|
SET approval_status = 'approved',
|
|
approved_by = $3,
|
|
approved_at = NOW()
|
|
WHERE mission_id = $1::uuid
|
|
AND tenant_id = $2
|
|
AND approval_status = 'pending'
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
actor_id,
|
|
)
|
|
return int(result.rsplit(" ", 1)[-1])
|
|
|
|
async def reject_pending_writebacks(self, mission_id: str, tenant_id: str, actor_id: str, reason: str) -> int:
|
|
async with self.pool.acquire() as conn:
|
|
result = await conn.execute(
|
|
"""
|
|
UPDATE colony_writeback_proposals
|
|
SET approval_status = 'rejected',
|
|
rejected_by = $3,
|
|
rejected_at = NOW(),
|
|
rejection_reason = $4
|
|
WHERE mission_id = $1::uuid
|
|
AND tenant_id = $2
|
|
AND approval_status = 'pending'
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
actor_id,
|
|
reason,
|
|
)
|
|
return int(result.rsplit(" ", 1)[-1])
|
|
|
|
async def log_event(
|
|
self,
|
|
*,
|
|
mission_id: str | None,
|
|
tenant_id: str,
|
|
event_type: str,
|
|
actor: str | None,
|
|
detail: dict[str, Any] | None = None,
|
|
) -> None:
|
|
async with self.pool.acquire() as conn:
|
|
await conn.execute(
|
|
"""
|
|
INSERT INTO colony_event_log (mission_id, tenant_id, event_type, actor, detail, created_at)
|
|
VALUES ($1::uuid, $2, $3, $4, $5::jsonb, NOW())
|
|
""",
|
|
mission_id,
|
|
tenant_id,
|
|
event_type,
|
|
actor,
|
|
json.dumps(detail or {}),
|
|
)
|