نسخهٔ وب مستقل — پلتفرم نیروی کار هوش مصنوعی
مخزن اختصاصی برنامهٔ وب؛ محتوای پوشهٔ web در ریشهٔ مخزن قرار گرفت. تاریخچهٔ پروژهٔ ربات (agy-bot) در این مخزن وجود ندارد و هیچ وابستگی runtime به آن نیست.
This commit is contained in:
mjalmoudy94
@@ -0,0 +1,276 @@
|
||||
"""Repositories for the AI Control Plane.
|
||||
|
||||
Platform-owner-only data. The application layer is responsible for never
|
||||
serialising any of it to a customer-facing response; these repositories just
|
||||
move rows.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.ai_control.models import (
|
||||
AiModel,
|
||||
ModelCredential,
|
||||
Provider,
|
||||
ProviderType,
|
||||
RateCard,
|
||||
RateUnit,
|
||||
RoutingMode,
|
||||
RoutingPolicy,
|
||||
)
|
||||
from app.infrastructure.db.models import (
|
||||
AiModelCredentialRow,
|
||||
AiModelRow,
|
||||
AiProviderRow,
|
||||
RateCardRow,
|
||||
RoutingPolicyRow,
|
||||
)
|
||||
|
||||
|
||||
def _provider_to_domain(row: AiProviderRow) -> Provider:
|
||||
return Provider(
|
||||
id=row.id, name=row.name, type=ProviderType(row.type),
|
||||
base_url=row.base_url, is_active=row.is_active, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _credential_to_domain(row: AiModelCredentialRow) -> ModelCredential:
|
||||
return ModelCredential(
|
||||
id=row.id, provider_id=row.provider_id, name=row.name,
|
||||
encrypted_api_key=row.encrypted_api_key, is_active=row.is_active,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _model_to_domain(row: AiModelRow) -> AiModel:
|
||||
return AiModel(
|
||||
id=row.id, provider_id=row.provider_id, internal_name=row.internal_name,
|
||||
display_name_for_admin=row.display_name_for_admin,
|
||||
capabilities=list(row.capabilities or []),
|
||||
input_price_per_1m=row.input_price_per_1m,
|
||||
output_price_per_1m=row.output_price_per_1m,
|
||||
is_active=row.is_active, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _policy_to_domain(row: RoutingPolicyRow) -> RoutingPolicy:
|
||||
return RoutingPolicy(
|
||||
id=row.id, name=row.name, description=row.description,
|
||||
mode=RoutingMode(row.mode),
|
||||
fallback_policy_json=dict(row.fallback_policy_json or {}),
|
||||
budget_limits_json=dict(row.budget_limits_json or {}),
|
||||
is_default=row.is_default, is_active=row.is_active, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _rate_card_to_domain(row: RateCardRow) -> RateCard:
|
||||
return RateCard(
|
||||
id=row.id, name=row.name, unit=RateUnit(row.unit),
|
||||
customer_price=row.customer_price,
|
||||
internal_cost_estimate=row.internal_cost_estimate,
|
||||
is_active=row.is_active, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class ProviderRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list(self) -> list[Provider]:
|
||||
stmt = select(AiProviderRow).order_by(AiProviderRow.created_at)
|
||||
return [_provider_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def get(self, provider_id: str) -> Provider | None:
|
||||
row = await self.session.get(AiProviderRow, provider_id)
|
||||
return _provider_to_domain(row) if row else None
|
||||
|
||||
async def get_by_name(self, name: str) -> Provider | None:
|
||||
stmt = select(AiProviderRow).where(AiProviderRow.name == name)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _provider_to_domain(row) if row else None
|
||||
|
||||
async def create(self, provider: Provider) -> Provider:
|
||||
row = AiProviderRow(
|
||||
id=provider.id, name=provider.name, type=provider.type.value,
|
||||
base_url=provider.base_url, is_active=provider.is_active,
|
||||
created_at=provider.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _provider_to_domain(row)
|
||||
|
||||
async def set_active(self, provider_id: str, active: bool) -> None:
|
||||
row = await self.session.get(AiProviderRow, provider_id)
|
||||
if row:
|
||||
row.is_active = active
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
class ModelCredentialRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_provider(self, provider_id: str) -> list[ModelCredential]:
|
||||
stmt = select(AiModelCredentialRow).where(
|
||||
AiModelCredentialRow.provider_id == provider_id
|
||||
)
|
||||
return [_credential_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def get(self, credential_id: str) -> ModelCredential | None:
|
||||
row = await self.session.get(AiModelCredentialRow, credential_id)
|
||||
return _credential_to_domain(row) if row else None
|
||||
|
||||
async def get_active_for_provider(self, provider_id: str) -> ModelCredential | None:
|
||||
stmt = (
|
||||
select(AiModelCredentialRow)
|
||||
.where(
|
||||
AiModelCredentialRow.provider_id == provider_id,
|
||||
AiModelCredentialRow.is_active.is_(True),
|
||||
)
|
||||
.order_by(AiModelCredentialRow.created_at.desc())
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalars().first()
|
||||
return _credential_to_domain(row) if row else None
|
||||
|
||||
async def create(self, credential: ModelCredential) -> ModelCredential:
|
||||
row = AiModelCredentialRow(
|
||||
id=credential.id, provider_id=credential.provider_id, name=credential.name,
|
||||
encrypted_api_key=credential.encrypted_api_key,
|
||||
is_active=credential.is_active, created_at=credential.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _credential_to_domain(row)
|
||||
|
||||
|
||||
class AiModelRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list(self, *, provider_id: str | None = None) -> list[AiModel]:
|
||||
stmt = select(AiModelRow)
|
||||
if provider_id:
|
||||
stmt = stmt.where(AiModelRow.provider_id == provider_id)
|
||||
stmt = stmt.order_by(AiModelRow.created_at)
|
||||
return [_model_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def get(self, model_id: str) -> AiModel | None:
|
||||
row = await self.session.get(AiModelRow, model_id)
|
||||
return _model_to_domain(row) if row else None
|
||||
|
||||
async def create(self, model: AiModel) -> AiModel:
|
||||
row = AiModelRow(
|
||||
id=model.id, provider_id=model.provider_id,
|
||||
internal_name=model.internal_name,
|
||||
display_name_for_admin=model.display_name_for_admin,
|
||||
capabilities=list(model.capabilities),
|
||||
input_price_per_1m=model.input_price_per_1m,
|
||||
output_price_per_1m=model.output_price_per_1m,
|
||||
is_active=model.is_active, created_at=model.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _model_to_domain(row)
|
||||
|
||||
async def pick_for_mode(self, mode: RoutingMode) -> AiModel | None:
|
||||
"""Choose a model for a routing mode.
|
||||
|
||||
Deterministic and explainable: filter by required capability, then take
|
||||
the cheapest active model. A real implementation would add health and
|
||||
latency signals; keeping it simple here makes the behaviour testable.
|
||||
"""
|
||||
required: set[str] = set()
|
||||
if mode == RoutingMode.VISION:
|
||||
required.add("vision")
|
||||
if mode == RoutingMode.CODING:
|
||||
required.add("code")
|
||||
if mode == RoutingMode.REASONING:
|
||||
required.add("reasoning")
|
||||
|
||||
candidates = [m for m in await self.list() if m.is_active]
|
||||
if required:
|
||||
filtered = [m for m in candidates if required.issubset(set(m.capabilities))]
|
||||
if filtered:
|
||||
candidates = filtered
|
||||
if not candidates:
|
||||
return None
|
||||
return min(candidates, key=lambda m: (m.input_price_per_1m, m.output_price_per_1m))
|
||||
|
||||
|
||||
class RoutingPolicyRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list(self) -> list[RoutingPolicy]:
|
||||
stmt = select(RoutingPolicyRow).order_by(RoutingPolicyRow.created_at)
|
||||
return [_policy_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def get(self, policy_id: str) -> RoutingPolicy | None:
|
||||
row = await self.session.get(RoutingPolicyRow, policy_id)
|
||||
return _policy_to_domain(row) if row else None
|
||||
|
||||
async def get_default(self) -> RoutingPolicy | None:
|
||||
stmt = select(RoutingPolicyRow).where(RoutingPolicyRow.is_default.is_(True))
|
||||
row = (await self.session.execute(stmt)).scalars().first()
|
||||
return _policy_to_domain(row) if row else None
|
||||
|
||||
async def create(self, policy: RoutingPolicy) -> RoutingPolicy:
|
||||
if policy.is_default:
|
||||
await self._clear_defaults()
|
||||
row = RoutingPolicyRow(
|
||||
id=policy.id, name=policy.name, description=policy.description,
|
||||
mode=policy.mode.value,
|
||||
fallback_policy_json=policy.fallback_policy_json,
|
||||
budget_limits_json=policy.budget_limits_json,
|
||||
is_default=policy.is_default, is_active=policy.is_active,
|
||||
created_at=policy.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _policy_to_domain(row)
|
||||
|
||||
async def _clear_defaults(self) -> None:
|
||||
rows = (await self.session.execute(select(RoutingPolicyRow))).scalars().all()
|
||||
for row in rows:
|
||||
row.is_default = False
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
class RateCardRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list(self) -> list[RateCard]:
|
||||
stmt = select(RateCardRow).order_by(RateCardRow.unit)
|
||||
return [_rate_card_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def get_by_unit(self, unit: RateUnit) -> RateCard | None:
|
||||
stmt = select(RateCardRow).where(RateCardRow.unit == unit.value)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _rate_card_to_domain(row) if row else None
|
||||
|
||||
async def upsert(self, card: RateCard) -> RateCard:
|
||||
existing = await self.get_by_unit(card.unit)
|
||||
if existing:
|
||||
row = await self.session.get(RateCardRow, existing.id)
|
||||
assert row is not None
|
||||
row.customer_price = card.customer_price
|
||||
row.internal_cost_estimate = card.internal_cost_estimate
|
||||
row.name = card.name
|
||||
row.is_active = card.is_active
|
||||
await self.session.flush()
|
||||
return _rate_card_to_domain(row)
|
||||
return await self.create(card)
|
||||
|
||||
async def create(self, card: RateCard) -> RateCard:
|
||||
row = RateCardRow(
|
||||
id=card.id, name=card.name, unit=card.unit.value,
|
||||
customer_price=card.customer_price,
|
||||
internal_cost_estimate=card.internal_cost_estimate,
|
||||
is_active=card.is_active, created_at=card.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _rate_card_to_domain(row)
|
||||
@@ -0,0 +1,496 @@
|
||||
"""Repositories for Audit & Observability and the Usage Ledger."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import case, func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.application.ports.audit import UsageAggregate, UsageBreakdown, UsageRecord
|
||||
from app.domain.ai_control.models import RateUnit
|
||||
from app.domain.audit.models import AuditEvent, ErrorEvent, QualityMetric, Span, Trace
|
||||
from app.infrastructure.db.models import (
|
||||
AuditEventRow,
|
||||
ErrorEventRow,
|
||||
QualityMetricRow,
|
||||
SpanRow,
|
||||
TraceRow,
|
||||
UsageEventRow,
|
||||
)
|
||||
|
||||
|
||||
def _audit_to_domain(row: AuditEventRow) -> AuditEvent:
|
||||
return AuditEvent(
|
||||
id=row.id,
|
||||
organization_id=row.organization_id,
|
||||
actor_user_id=row.actor_user_id,
|
||||
action=row.action,
|
||||
target_type=row.target_type,
|
||||
target_id=row.target_id,
|
||||
metadata_json=row.metadata_json or {},
|
||||
ip_address=row.ip_address,
|
||||
user_agent=row.user_agent,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _error_to_domain(row: ErrorEventRow) -> ErrorEvent:
|
||||
return ErrorEvent(
|
||||
id=row.id,
|
||||
organization_id=row.organization_id,
|
||||
run_id=row.run_id,
|
||||
kind=row.kind,
|
||||
message=row.message,
|
||||
stack=row.stack,
|
||||
metadata_json=row.metadata_json or {},
|
||||
source=row.source,
|
||||
severity=row.severity,
|
||||
error_type=row.error_type,
|
||||
fingerprint=row.fingerprint,
|
||||
status_code=row.status_code,
|
||||
method=row.method,
|
||||
path=row.path,
|
||||
user_id=row.user_id,
|
||||
request_id=row.request_id,
|
||||
ip=row.ip,
|
||||
user_agent=row.user_agent,
|
||||
resolved_at=row.resolved_at,
|
||||
resolved_note=row.resolved_note,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class AuditRepository:
|
||||
"""Concrete ``AuditPort``. Append-only."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def record(self, *, event: AuditEvent) -> AuditEvent:
|
||||
row = AuditEventRow(
|
||||
id=event.id,
|
||||
organization_id=event.organization_id,
|
||||
actor_user_id=event.actor_user_id,
|
||||
action=event.action,
|
||||
target_type=event.target_type,
|
||||
target_id=event.target_id,
|
||||
metadata_json=event.metadata_json,
|
||||
ip_address=event.ip_address,
|
||||
user_agent=event.user_agent,
|
||||
created_at=event.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _audit_to_domain(row)
|
||||
|
||||
async def list_for_organization(
|
||||
self, *, organization_id: str, limit: int = 100, offset: int = 0
|
||||
) -> list[AuditEvent]:
|
||||
stmt = (
|
||||
select(AuditEventRow)
|
||||
.where(AuditEventRow.organization_id == organization_id)
|
||||
.order_by(AuditEventRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return [_audit_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_platform(self, *, limit: int = 100, offset: int = 0) -> list[AuditEvent]:
|
||||
stmt = (
|
||||
select(AuditEventRow)
|
||||
.order_by(AuditEventRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return [_audit_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
|
||||
class UsageLedgerRepository:
|
||||
"""Concrete ``UsageLedgerPort``. The ledger is the billing source of truth."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def record(self, *, usage: UsageRecord) -> None:
|
||||
total = usage.total_price or (usage.quantity * usage.unit_price)
|
||||
row = UsageEventRow(
|
||||
organization_id=usage.organization_id,
|
||||
project_id=usage.project_id,
|
||||
user_id=usage.user_id,
|
||||
employee_id=usage.employee_id,
|
||||
run_id=usage.run_id,
|
||||
unit=usage.unit.value,
|
||||
quantity=usage.quantity,
|
||||
unit_price=usage.unit_price,
|
||||
internal_cost_estimate=usage.internal_cost_estimate,
|
||||
total_price=total,
|
||||
metadata_json=usage.metadata_json,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
|
||||
async def breakdown(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
since: datetime | None = None,
|
||||
until: datetime | None = None,
|
||||
) -> UsageBreakdown:
|
||||
filters = [UsageEventRow.organization_id == organization_id]
|
||||
if since is not None:
|
||||
filters.append(UsageEventRow.created_at >= since)
|
||||
if until is not None:
|
||||
filters.append(UsageEventRow.created_at <= until)
|
||||
|
||||
by_unit_stmt = (
|
||||
select(
|
||||
UsageEventRow.unit,
|
||||
func.coalesce(func.sum(UsageEventRow.quantity), 0.0),
|
||||
func.coalesce(func.sum(UsageEventRow.total_price), 0.0),
|
||||
func.coalesce(func.sum(UsageEventRow.internal_cost_estimate), 0.0),
|
||||
)
|
||||
.where(*filters)
|
||||
.group_by(UsageEventRow.unit)
|
||||
)
|
||||
by_unit = [
|
||||
UsageAggregate(
|
||||
unit=RateUnit(unit),
|
||||
quantity=float(qty),
|
||||
total_price=float(price),
|
||||
internal_cost_estimate=float(cost),
|
||||
)
|
||||
for unit, qty, price, cost in (await self.session.execute(by_unit_stmt)).all()
|
||||
]
|
||||
|
||||
emp_stmt = (
|
||||
select(
|
||||
UsageEventRow.employee_id,
|
||||
func.coalesce(func.sum(UsageEventRow.total_price), 0.0),
|
||||
)
|
||||
.where(*filters, UsageEventRow.employee_id.is_not(None))
|
||||
.group_by(UsageEventRow.employee_id)
|
||||
)
|
||||
by_employee = {
|
||||
str(emp): float(total)
|
||||
for emp, total in (await self.session.execute(emp_stmt)).all()
|
||||
}
|
||||
|
||||
proj_stmt = (
|
||||
select(
|
||||
UsageEventRow.project_id,
|
||||
func.coalesce(func.sum(UsageEventRow.total_price), 0.0),
|
||||
)
|
||||
.where(*filters, UsageEventRow.project_id.is_not(None))
|
||||
.group_by(UsageEventRow.project_id)
|
||||
)
|
||||
by_project = {
|
||||
str(proj): float(total)
|
||||
for proj, total in (await self.session.execute(proj_stmt)).all()
|
||||
}
|
||||
|
||||
return UsageBreakdown(
|
||||
by_unit=by_unit,
|
||||
by_employee=by_employee,
|
||||
by_project=by_project,
|
||||
total_price=sum(a.total_price for a in by_unit),
|
||||
total_cost=sum(a.internal_cost_estimate for a in by_unit),
|
||||
)
|
||||
|
||||
async def list_events(
|
||||
self, *, organization_id: str, limit: int = 100, offset: int = 0
|
||||
) -> list[UsageEventRow]:
|
||||
stmt = (
|
||||
select(UsageEventRow)
|
||||
.where(UsageEventRow.organization_id == organization_id)
|
||||
.order_by(UsageEventRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return list((await self.session.execute(stmt)).scalars())
|
||||
|
||||
async def platform_totals(self) -> list[tuple[str, float, float]]:
|
||||
"""``(organization_id, quantity, total_price)`` across every tenant.
|
||||
|
||||
Platform-owner console only.
|
||||
"""
|
||||
stmt = (
|
||||
select(
|
||||
UsageEventRow.organization_id,
|
||||
func.coalesce(func.sum(UsageEventRow.quantity), 0.0),
|
||||
func.coalesce(func.sum(UsageEventRow.total_price), 0.0),
|
||||
)
|
||||
.group_by(UsageEventRow.organization_id)
|
||||
.order_by(func.sum(UsageEventRow.total_price).desc())
|
||||
)
|
||||
rows = (await self.session.execute(stmt)).all()
|
||||
return [(str(org), float(qty), float(price)) for org, qty, price in rows]
|
||||
|
||||
|
||||
class ObservabilityRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create_trace(self, trace: Trace) -> Trace:
|
||||
self.session.add(
|
||||
TraceRow(
|
||||
id=trace.id, organization_id=trace.organization_id,
|
||||
run_id=trace.run_id, name=trace.name, created_at=trace.created_at,
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
return trace
|
||||
|
||||
async def add_span(self, span: Span) -> Span:
|
||||
self.session.add(
|
||||
SpanRow(
|
||||
id=span.id, trace_id=span.trace_id, parent_span_id=span.parent_span_id,
|
||||
name=span.name, kind=span.kind, attributes_json=span.attributes_json,
|
||||
started_at=span.started_at, ended_at=span.ended_at,
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
return span
|
||||
|
||||
async def record_error(self, event: ErrorEvent) -> None:
|
||||
self.session.add(
|
||||
ErrorEventRow(
|
||||
id=event.id,
|
||||
organization_id=event.organization_id,
|
||||
run_id=event.run_id,
|
||||
kind=event.kind,
|
||||
message=event.message,
|
||||
stack=event.stack,
|
||||
metadata_json=event.metadata_json,
|
||||
source=event.source,
|
||||
severity=event.severity,
|
||||
error_type=event.error_type,
|
||||
fingerprint=event.fingerprint,
|
||||
status_code=event.status_code,
|
||||
method=event.method,
|
||||
path=event.path,
|
||||
user_id=event.user_id,
|
||||
request_id=event.request_id,
|
||||
ip=event.ip,
|
||||
user_agent=event.user_agent,
|
||||
resolved_at=event.resolved_at,
|
||||
resolved_note=event.resolved_note,
|
||||
created_at=event.created_at,
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
|
||||
async def record_metric(self, metric: QualityMetric) -> None:
|
||||
self.session.add(
|
||||
QualityMetricRow(
|
||||
id=metric.id, organization_id=metric.organization_id, name=metric.name,
|
||||
value=metric.value, dimensions_json=metric.dimensions_json,
|
||||
created_at=metric.created_at,
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
|
||||
async def list_spans(self, *, trace_id: str) -> list[Span]:
|
||||
stmt = select(SpanRow).where(SpanRow.trace_id == trace_id).order_by(SpanRow.started_at)
|
||||
return [
|
||||
Span(
|
||||
id=r.id, trace_id=r.trace_id, parent_span_id=r.parent_span_id, name=r.name,
|
||||
kind=r.kind, attributes_json=r.attributes_json or {},
|
||||
started_at=r.started_at, ended_at=r.ended_at,
|
||||
)
|
||||
for r in (await self.session.execute(stmt)).scalars()
|
||||
]
|
||||
|
||||
# -- error console --------------------------------------------------------
|
||||
|
||||
def _error_filters(
|
||||
self,
|
||||
*,
|
||||
severity: str | None = None,
|
||||
source: str | None = None,
|
||||
status: str | None = None,
|
||||
kind: str | None = None,
|
||||
organization_id: str | None = None,
|
||||
fingerprint: str | None = None,
|
||||
since: datetime | None = None,
|
||||
search: str | None = None,
|
||||
) -> list:
|
||||
"""Build the WHERE clauses shared by the list and count queries."""
|
||||
clauses: list = []
|
||||
if severity:
|
||||
clauses.append(ErrorEventRow.severity == severity)
|
||||
if source:
|
||||
clauses.append(ErrorEventRow.source == source)
|
||||
if status == "open":
|
||||
clauses.append(ErrorEventRow.resolved_at.is_(None))
|
||||
elif status == "resolved":
|
||||
clauses.append(ErrorEventRow.resolved_at.is_not(None))
|
||||
if kind:
|
||||
clauses.append(ErrorEventRow.kind == kind)
|
||||
if organization_id:
|
||||
clauses.append(ErrorEventRow.organization_id == organization_id)
|
||||
if fingerprint:
|
||||
clauses.append(ErrorEventRow.fingerprint == fingerprint)
|
||||
if since is not None:
|
||||
clauses.append(ErrorEventRow.created_at >= since)
|
||||
if search:
|
||||
clauses.append(ErrorEventRow.message.ilike(f"%{search}%"))
|
||||
return clauses
|
||||
|
||||
async def list_errors(
|
||||
self,
|
||||
*,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
**filters,
|
||||
) -> list[ErrorEvent]:
|
||||
stmt = (
|
||||
select(ErrorEventRow)
|
||||
.where(*self._error_filters(**filters))
|
||||
# Newest first: the failure someone is reacting to is almost always
|
||||
# the most recent one.
|
||||
.order_by(ErrorEventRow.created_at.desc())
|
||||
.limit(limit)
|
||||
.offset(offset)
|
||||
)
|
||||
return [_error_to_domain(row) for row in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def count_errors(self, **filters) -> int:
|
||||
stmt = select(func.count()).select_from(ErrorEventRow).where(
|
||||
*self._error_filters(**filters)
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
async def get_error(self, error_id: str) -> ErrorEvent | None:
|
||||
row = (
|
||||
await self.session.execute(
|
||||
select(ErrorEventRow).where(ErrorEventRow.id == error_id)
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
return _error_to_domain(row) if row is not None else None
|
||||
|
||||
async def list_error_groups(
|
||||
self,
|
||||
*,
|
||||
limit: int = 50,
|
||||
since: datetime | None = None,
|
||||
status: str | None = None,
|
||||
severity: str | None = None,
|
||||
source: str | None = None,
|
||||
organization_id: str | None = None,
|
||||
) -> list[dict]:
|
||||
"""One row per distinct bug, newest activity first.
|
||||
|
||||
This is the view that makes the console usable: a bug that fired 4 000
|
||||
times is one line with a count, not 4 000 lines. ``open_count`` is what
|
||||
tells you whether anyone has actually dealt with it.
|
||||
"""
|
||||
clauses: list = [ErrorEventRow.fingerprint != ""]
|
||||
if since is not None:
|
||||
clauses.append(ErrorEventRow.created_at >= since)
|
||||
if severity:
|
||||
clauses.append(ErrorEventRow.severity == severity)
|
||||
if source:
|
||||
clauses.append(ErrorEventRow.source == source)
|
||||
if organization_id:
|
||||
clauses.append(ErrorEventRow.organization_id == organization_id)
|
||||
|
||||
open_count = func.sum(
|
||||
case((ErrorEventRow.resolved_at.is_(None), 1), else_=0)
|
||||
).label("open_count")
|
||||
stmt = (
|
||||
select(
|
||||
ErrorEventRow.fingerprint,
|
||||
func.count().label("occurrences"),
|
||||
open_count,
|
||||
func.min(ErrorEventRow.created_at).label("first_seen"),
|
||||
func.max(ErrorEventRow.created_at).label("last_seen"),
|
||||
func.max(ErrorEventRow.kind).label("kind"),
|
||||
func.max(ErrorEventRow.error_type).label("error_type"),
|
||||
func.max(ErrorEventRow.severity).label("severity"),
|
||||
func.max(ErrorEventRow.source).label("source"),
|
||||
func.max(ErrorEventRow.path).label("path"),
|
||||
func.max(ErrorEventRow.message).label("sample_message"),
|
||||
)
|
||||
.where(*clauses)
|
||||
.group_by(ErrorEventRow.fingerprint)
|
||||
.order_by(func.max(ErrorEventRow.created_at).desc())
|
||||
.limit(limit)
|
||||
)
|
||||
rows = (await self.session.execute(stmt)).all()
|
||||
groups = [
|
||||
{
|
||||
"fingerprint": row.fingerprint,
|
||||
"occurrences": int(row.occurrences),
|
||||
"open_count": int(row.open_count or 0),
|
||||
"first_seen": row.first_seen,
|
||||
"last_seen": row.last_seen,
|
||||
"kind": row.kind or "",
|
||||
"error_type": row.error_type or "",
|
||||
"severity": row.severity or "",
|
||||
"source": row.source or "",
|
||||
"path": row.path or "",
|
||||
"sample_message": row.sample_message or "",
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
if status == "open":
|
||||
return [g for g in groups if g["open_count"] > 0]
|
||||
if status == "resolved":
|
||||
return [g for g in groups if g["open_count"] == 0]
|
||||
return groups
|
||||
|
||||
async def resolve_error(
|
||||
self, error_id: str, *, note: str, resolved_at: datetime
|
||||
) -> ErrorEvent | None:
|
||||
row = (
|
||||
await self.session.execute(
|
||||
select(ErrorEventRow).where(ErrorEventRow.id == error_id)
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
if row is None:
|
||||
return None
|
||||
row.resolved_at = resolved_at
|
||||
row.resolved_note = note
|
||||
await self.session.flush()
|
||||
return _error_to_domain(row)
|
||||
|
||||
async def resolve_error_group(
|
||||
self, fingerprint: str, *, note: str, resolved_at: datetime
|
||||
) -> int:
|
||||
"""Resolve every open occurrence of one bug. Returns how many changed."""
|
||||
stmt = select(ErrorEventRow).where(
|
||||
ErrorEventRow.fingerprint == fingerprint,
|
||||
ErrorEventRow.resolved_at.is_(None),
|
||||
)
|
||||
rows = list((await self.session.execute(stmt)).scalars())
|
||||
for row in rows:
|
||||
row.resolved_at = resolved_at
|
||||
row.resolved_note = note
|
||||
await self.session.flush()
|
||||
return len(rows)
|
||||
|
||||
async def reopen_error_group(self, fingerprint: str) -> int:
|
||||
"""Undo a resolution — for when the bug comes back."""
|
||||
stmt = select(ErrorEventRow).where(
|
||||
ErrorEventRow.fingerprint == fingerprint,
|
||||
ErrorEventRow.resolved_at.is_not(None),
|
||||
)
|
||||
rows = list((await self.session.execute(stmt)).scalars())
|
||||
for row in rows:
|
||||
row.resolved_at = None
|
||||
row.resolved_note = ""
|
||||
await self.session.flush()
|
||||
return len(rows)
|
||||
|
||||
async def count_open_by_severity(self) -> dict[str, int]:
|
||||
"""Unresolved counts per severity, for the console's header tiles."""
|
||||
stmt = (
|
||||
select(ErrorEventRow.severity, func.count())
|
||||
.where(ErrorEventRow.resolved_at.is_(None))
|
||||
.group_by(ErrorEventRow.severity)
|
||||
)
|
||||
return {
|
||||
str(severity): int(count)
|
||||
for severity, count in (await self.session.execute(stmt)).all()
|
||||
}
|
||||
@@ -0,0 +1,212 @@
|
||||
"""Repositories for Billing: plans, subscriptions, wallets, invoices, payments.
|
||||
|
||||
The usage ledger lives in ``repositories/audit.py`` because it is shared with the
|
||||
observability layer; billing reads it but never writes to it directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.billing.models import (
|
||||
CreditWallet,
|
||||
Invoice,
|
||||
InvoiceStatus,
|
||||
Payment,
|
||||
PaymentStatus,
|
||||
Plan,
|
||||
Subscription,
|
||||
SubscriptionStatus,
|
||||
)
|
||||
from app.infrastructure.db.models import (
|
||||
CreditWalletRow,
|
||||
InvoiceRow,
|
||||
PaymentRow,
|
||||
PlanRow,
|
||||
SubscriptionRow,
|
||||
)
|
||||
|
||||
|
||||
def _plan(row: PlanRow) -> Plan:
|
||||
return Plan(
|
||||
id=row.id, name=row.name, monthly_price=row.monthly_price,
|
||||
included_credit=row.included_credit, features_json=dict(row.features_json or {}),
|
||||
is_active=row.is_active, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _subscription(row: SubscriptionRow) -> Subscription:
|
||||
return Subscription(
|
||||
id=row.id, organization_id=row.organization_id, plan_id=row.plan_id,
|
||||
status=SubscriptionStatus(row.status),
|
||||
external_subscription_id=row.external_subscription_id,
|
||||
current_period_end=row.current_period_end, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _wallet(row: CreditWalletRow) -> CreditWallet:
|
||||
return CreditWallet(
|
||||
id=row.id, organization_id=row.organization_id, balance=row.balance,
|
||||
currency=row.currency, low_balance_threshold=row.low_balance_threshold,
|
||||
updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
|
||||
def _invoice(row: InvoiceRow) -> Invoice:
|
||||
return Invoice(
|
||||
id=row.id, organization_id=row.organization_id, amount=row.amount,
|
||||
currency=row.currency, status=InvoiceStatus(row.status),
|
||||
external_invoice_id=row.external_invoice_id, period_start=row.period_start,
|
||||
period_end=row.period_end, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _payment(row: PaymentRow) -> Payment:
|
||||
return Payment(
|
||||
id=row.id, organization_id=row.organization_id, amount=row.amount,
|
||||
status=PaymentStatus(row.status), provider=row.provider,
|
||||
external_payment_id=row.external_payment_id, invoice_id=row.invoice_id,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class PlanRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list(self) -> list[Plan]:
|
||||
stmt = select(PlanRow).order_by(PlanRow.monthly_price)
|
||||
return [_plan(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def get(self, plan_id: str) -> Plan | None:
|
||||
row = await self.session.get(PlanRow, plan_id)
|
||||
return _plan(row) if row else None
|
||||
|
||||
|
||||
class SubscriptionRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get_for_org(self, organization_id: str) -> Subscription | None:
|
||||
stmt = select(SubscriptionRow).where(
|
||||
SubscriptionRow.organization_id == organization_id
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalars().first()
|
||||
return _subscription(row) if row else None
|
||||
|
||||
async def upsert(self, subscription: Subscription) -> Subscription:
|
||||
existing = await self.get_for_org(subscription.organization_id)
|
||||
if existing:
|
||||
row = await self.session.get(SubscriptionRow, existing.id)
|
||||
assert row is not None
|
||||
row.plan_id = subscription.plan_id
|
||||
row.status = str(subscription.status)
|
||||
row.external_subscription_id = subscription.external_subscription_id
|
||||
row.current_period_end = subscription.current_period_end
|
||||
await self.session.flush()
|
||||
return _subscription(row)
|
||||
|
||||
row = SubscriptionRow(
|
||||
id=subscription.id, organization_id=subscription.organization_id,
|
||||
plan_id=subscription.plan_id, status=str(subscription.status),
|
||||
external_subscription_id=subscription.external_subscription_id,
|
||||
current_period_end=subscription.current_period_end,
|
||||
created_at=subscription.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _subscription(row)
|
||||
|
||||
|
||||
class CreditWalletRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get_for_org(self, organization_id: str) -> CreditWallet | None:
|
||||
stmt = select(CreditWalletRow).where(
|
||||
CreditWalletRow.organization_id == organization_id
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _wallet(row) if row else None
|
||||
|
||||
async def ensure(self, organization_id: str) -> CreditWallet:
|
||||
"""Return the org's wallet, creating a zero-balance one if absent."""
|
||||
existing = await self.get_for_org(organization_id)
|
||||
if existing:
|
||||
return existing
|
||||
row = CreditWalletRow(organization_id=organization_id, balance=0.0, currency="USD")
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _wallet(row)
|
||||
|
||||
async def adjust(self, *, organization_id: str, delta: float) -> CreditWallet:
|
||||
"""Add (or subtract) credit. Callers must validate the delta's sign."""
|
||||
wallet = await self.ensure(organization_id)
|
||||
row = await self.session.get(CreditWalletRow, wallet.id)
|
||||
assert row is not None
|
||||
row.balance = float(row.balance) + float(delta)
|
||||
await self.session.flush()
|
||||
return _wallet(row)
|
||||
|
||||
|
||||
class InvoiceRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_org(self, organization_id: str) -> list[Invoice]:
|
||||
stmt = (
|
||||
select(InvoiceRow)
|
||||
.where(InvoiceRow.organization_id == organization_id)
|
||||
.order_by(InvoiceRow.created_at.desc())
|
||||
)
|
||||
return [_invoice(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, invoice: Invoice) -> Invoice:
|
||||
row = InvoiceRow(
|
||||
id=invoice.id, organization_id=invoice.organization_id,
|
||||
amount=invoice.amount, currency=invoice.currency,
|
||||
status=str(invoice.status),
|
||||
external_invoice_id=invoice.external_invoice_id,
|
||||
period_start=invoice.period_start, period_end=invoice.period_end,
|
||||
created_at=invoice.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _invoice(row)
|
||||
|
||||
|
||||
class PaymentRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_org(self, organization_id: str) -> list[Payment]:
|
||||
stmt = (
|
||||
select(PaymentRow)
|
||||
.where(PaymentRow.organization_id == organization_id)
|
||||
.order_by(PaymentRow.created_at.desc())
|
||||
)
|
||||
return [_payment(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, payment: Payment) -> Payment:
|
||||
row = PaymentRow(
|
||||
id=payment.id, organization_id=payment.organization_id,
|
||||
amount=payment.amount, status=str(payment.status),
|
||||
provider=payment.provider, external_payment_id=payment.external_payment_id,
|
||||
invoice_id=payment.invoice_id, created_at=payment.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _payment(row)
|
||||
|
||||
async def set_status(
|
||||
self, payment_id: str, status: PaymentStatus, *, external_id: str | None = None
|
||||
) -> Payment | None:
|
||||
row = await self.session.get(PaymentRow, payment_id)
|
||||
if row is None:
|
||||
return None
|
||||
row.status = str(status)
|
||||
if external_id is not None:
|
||||
row.external_payment_id = external_id
|
||||
await self.session.flush()
|
||||
return _payment(row)
|
||||
@@ -0,0 +1,269 @@
|
||||
"""Repositories for Identity & Access.
|
||||
|
||||
Each repository maps between the ORM rows and the pure domain dataclasses, so
|
||||
the domain layer never sees SQLAlchemy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.identity.models import (
|
||||
Invite,
|
||||
Membership,
|
||||
MembershipStatus,
|
||||
Organization,
|
||||
Role,
|
||||
User,
|
||||
)
|
||||
from app.infrastructure.db.models import (
|
||||
InviteRow,
|
||||
MembershipRow,
|
||||
OrganizationRow,
|
||||
UserRow,
|
||||
)
|
||||
|
||||
|
||||
def _user_to_domain(row: UserRow) -> User:
|
||||
return User(
|
||||
id=row.id,
|
||||
email=row.email,
|
||||
password_hash=row.password_hash,
|
||||
name=row.name,
|
||||
avatar_url=row.avatar_url,
|
||||
is_platform_owner=row.is_platform_owner,
|
||||
is_active=row.is_active,
|
||||
created_at=row.created_at,
|
||||
updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
|
||||
def _org_to_domain(row: OrganizationRow) -> Organization:
|
||||
return Organization(
|
||||
id=row.id, name=row.name, slug=row.slug,
|
||||
owner_user_id=row.owner_user_id, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _membership_to_domain(row: MembershipRow) -> Membership:
|
||||
return Membership(
|
||||
id=row.id,
|
||||
organization_id=row.organization_id,
|
||||
user_id=row.user_id,
|
||||
role=Role(row.role),
|
||||
status=MembershipStatus(row.status),
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _invite_to_domain(row: InviteRow) -> Invite:
|
||||
return Invite(
|
||||
id=row.id,
|
||||
organization_id=row.organization_id,
|
||||
email=row.email,
|
||||
role=Role(row.role),
|
||||
token_hash=row.token_hash,
|
||||
invited_by_user_id=row.invited_by_user_id,
|
||||
expires_at=row.expires_at,
|
||||
accepted_at=row.accepted_at,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class UserRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, user_id: str) -> User | None:
|
||||
row = await self.session.get(UserRow, user_id)
|
||||
return _user_to_domain(row) if row else None
|
||||
|
||||
async def get_by_email(self, email: str) -> User | None:
|
||||
stmt = select(UserRow).where(func.lower(UserRow.email) == email.strip().lower())
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _user_to_domain(row) if row else None
|
||||
|
||||
async def create(self, user: User) -> User:
|
||||
row = UserRow(
|
||||
id=user.id,
|
||||
email=user.normalized_email(),
|
||||
password_hash=user.password_hash,
|
||||
name=user.name,
|
||||
avatar_url=user.avatar_url,
|
||||
is_platform_owner=user.is_platform_owner,
|
||||
is_active=user.is_active,
|
||||
created_at=user.created_at,
|
||||
updated_at=user.updated_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _user_to_domain(row)
|
||||
|
||||
async def update_password(self, user_id: str, password_hash: str) -> None:
|
||||
row = await self.session.get(UserRow, user_id)
|
||||
if row:
|
||||
row.password_hash = password_hash
|
||||
await self.session.flush()
|
||||
|
||||
async def count(self) -> int:
|
||||
return int((await self.session.execute(select(func.count(UserRow.id)))).scalar_one())
|
||||
|
||||
|
||||
class OrganizationRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, org_id: str) -> Organization | None:
|
||||
row = await self.session.get(OrganizationRow, org_id)
|
||||
return _org_to_domain(row) if row else None
|
||||
|
||||
async def get_by_slug(self, slug: str) -> Organization | None:
|
||||
stmt = select(OrganizationRow).where(OrganizationRow.slug == slug)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _org_to_domain(row) if row else None
|
||||
|
||||
async def list_all_slugs(self) -> set[str]:
|
||||
rows = (await self.session.execute(select(OrganizationRow.slug))).scalars().all()
|
||||
return set(rows)
|
||||
|
||||
async def create(self, org: Organization) -> Organization:
|
||||
row = OrganizationRow(
|
||||
id=org.id, name=org.name, slug=org.slug,
|
||||
owner_user_id=org.owner_user_id, created_at=org.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _org_to_domain(row)
|
||||
|
||||
async def list_for_user(self, user_id: str) -> list[tuple[Organization, Role]]:
|
||||
stmt = (
|
||||
select(OrganizationRow, MembershipRow.role)
|
||||
.join(MembershipRow, MembershipRow.organization_id == OrganizationRow.id)
|
||||
.where(MembershipRow.user_id == user_id, MembershipRow.status == "active")
|
||||
.order_by(OrganizationRow.created_at)
|
||||
)
|
||||
rows = (await self.session.execute(stmt)).all()
|
||||
return [(_org_to_domain(o), Role(r)) for o, r in rows]
|
||||
|
||||
async def list_all(self) -> list[Organization]:
|
||||
"""Every organization. Platform-owner console only — never expose this
|
||||
through a tenant-scoped endpoint."""
|
||||
stmt = select(OrganizationRow).order_by(OrganizationRow.created_at)
|
||||
return [_org_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def delete(self, org_id: str) -> None:
|
||||
row = await self.session.get(OrganizationRow, org_id)
|
||||
if row:
|
||||
await self.session.delete(row)
|
||||
await self.session.flush()
|
||||
|
||||
async def count(self) -> int:
|
||||
return int(
|
||||
(await self.session.execute(select(func.count(OrganizationRow.id)))).scalar_one()
|
||||
)
|
||||
|
||||
|
||||
class MembershipRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, org_id: str, user_id: str) -> Membership | None:
|
||||
stmt = select(MembershipRow).where(
|
||||
MembershipRow.organization_id == org_id, MembershipRow.user_id == user_id
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _membership_to_domain(row) if row else None
|
||||
|
||||
async def get_by_id(self, membership_id: str) -> Membership | None:
|
||||
row = await self.session.get(MembershipRow, membership_id)
|
||||
return _membership_to_domain(row) if row else None
|
||||
|
||||
async def list_for_org(self, org_id: str) -> list[Membership]:
|
||||
stmt = (
|
||||
select(MembershipRow)
|
||||
.where(MembershipRow.organization_id == org_id)
|
||||
.order_by(MembershipRow.created_at)
|
||||
)
|
||||
return [_membership_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def add(self, membership: Membership) -> Membership:
|
||||
row = MembershipRow(
|
||||
id=membership.id,
|
||||
organization_id=membership.organization_id,
|
||||
user_id=membership.user_id,
|
||||
role=membership.role.value,
|
||||
status=membership.status.value,
|
||||
created_at=membership.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _membership_to_domain(row)
|
||||
|
||||
async def set_role(self, membership_id: str, role: Role) -> None:
|
||||
row = await self.session.get(MembershipRow, membership_id)
|
||||
if row:
|
||||
row.role = role.value
|
||||
await self.session.flush()
|
||||
|
||||
async def set_status(self, membership_id: str, status: MembershipStatus) -> None:
|
||||
row = await self.session.get(MembershipRow, membership_id)
|
||||
if row:
|
||||
row.status = status.value
|
||||
await self.session.flush()
|
||||
|
||||
async def remove(self, membership_id: str) -> None:
|
||||
row = await self.session.get(MembershipRow, membership_id)
|
||||
if row:
|
||||
await self.session.delete(row)
|
||||
await self.session.flush()
|
||||
|
||||
async def count_owners(self, org_id: str) -> int:
|
||||
stmt = select(func.count(MembershipRow.id)).where(
|
||||
MembershipRow.organization_id == org_id,
|
||||
MembershipRow.role == Role.OWNER.value,
|
||||
MembershipRow.status == MembershipStatus.ACTIVE.value,
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
|
||||
class InviteRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get_by_token_hash(self, token_hash: str) -> Invite | None:
|
||||
stmt = select(InviteRow).where(InviteRow.token_hash == token_hash)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _invite_to_domain(row) if row else None
|
||||
|
||||
async def list_pending(self, org_id: str) -> list[Invite]:
|
||||
stmt = (
|
||||
select(InviteRow)
|
||||
.where(InviteRow.organization_id == org_id, InviteRow.accepted_at.is_(None))
|
||||
.order_by(InviteRow.created_at.desc())
|
||||
)
|
||||
return [_invite_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, invite: Invite) -> Invite:
|
||||
row = InviteRow(
|
||||
id=invite.id,
|
||||
organization_id=invite.organization_id,
|
||||
email=invite.email.strip().lower(),
|
||||
role=invite.role.value,
|
||||
token_hash=invite.token_hash,
|
||||
invited_by_user_id=invite.invited_by_user_id,
|
||||
expires_at=invite.expires_at,
|
||||
accepted_at=invite.accepted_at,
|
||||
created_at=invite.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _invite_to_domain(row)
|
||||
|
||||
async def mark_accepted(self, invite_id: str, when: datetime) -> None:
|
||||
row = await self.session.get(InviteRow, invite_id)
|
||||
if row:
|
||||
row.accepted_at = when
|
||||
await self.session.flush()
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Repositories for the Knowledge Base.
|
||||
|
||||
Retrieval is tenant-scoped by construction: every search joins
|
||||
``document_chunks -> documents -> knowledge_sources`` and filters on
|
||||
``knowledge_sources.organization_id``. There is no search entry point that takes
|
||||
a chunk id directly, so a caller cannot reach across tenants.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import func, or_, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.knowledge.models import (
|
||||
Document,
|
||||
DocumentChunk,
|
||||
EmbeddingJob,
|
||||
EmbeddingJobStatus,
|
||||
KnowledgeSource,
|
||||
RetrievalPolicy,
|
||||
SourceKind,
|
||||
)
|
||||
from app.infrastructure.db.models import (
|
||||
DocumentChunkRow,
|
||||
DocumentRow,
|
||||
EmbeddingJobRow,
|
||||
KnowledgeSourceRow,
|
||||
RetrievalPolicyRow,
|
||||
)
|
||||
|
||||
|
||||
def _source(row: KnowledgeSourceRow) -> KnowledgeSource:
|
||||
return KnowledgeSource(
|
||||
id=row.id, organization_id=row.organization_id, project_id=row.project_id,
|
||||
name=row.name, kind=SourceKind(row.kind), uri=row.uri, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _document(row: DocumentRow) -> Document:
|
||||
return Document(
|
||||
id=row.id, source_id=row.source_id, title=row.title, content=row.content,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class KnowledgeSourceRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, source_id: str) -> KnowledgeSource | None:
|
||||
row = await self.session.get(KnowledgeSourceRow, source_id)
|
||||
return _source(row) if row else None
|
||||
|
||||
async def get_for_org(self, organization_id: str, source_id: str) -> KnowledgeSource | None:
|
||||
stmt = select(KnowledgeSourceRow).where(
|
||||
KnowledgeSourceRow.id == source_id,
|
||||
KnowledgeSourceRow.organization_id == organization_id,
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _source(row) if row else None
|
||||
|
||||
async def list_for_org(
|
||||
self, organization_id: str, *, limit: int = 100, offset: int = 0
|
||||
) -> list[KnowledgeSource]:
|
||||
stmt = (
|
||||
select(KnowledgeSourceRow)
|
||||
.where(KnowledgeSourceRow.organization_id == organization_id)
|
||||
.order_by(KnowledgeSourceRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return [_source(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, source: KnowledgeSource) -> KnowledgeSource:
|
||||
row = KnowledgeSourceRow(
|
||||
id=source.id, organization_id=source.organization_id,
|
||||
project_id=source.project_id, name=source.name, kind=str(source.kind),
|
||||
uri=source.uri, created_at=source.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _source(row)
|
||||
|
||||
async def delete(self, source_id: str) -> None:
|
||||
row = await self.session.get(KnowledgeSourceRow, source_id)
|
||||
if row:
|
||||
await self.session.delete(row)
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
class DocumentRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_source(self, source_id: str) -> list[Document]:
|
||||
stmt = (
|
||||
select(DocumentRow)
|
||||
.where(DocumentRow.source_id == source_id)
|
||||
.order_by(DocumentRow.created_at)
|
||||
)
|
||||
return [_document(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, document: Document) -> Document:
|
||||
row = DocumentRow(
|
||||
id=document.id, source_id=document.source_id, title=document.title,
|
||||
content=document.content, created_at=document.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _document(row)
|
||||
|
||||
|
||||
class DocumentChunkRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create_many(self, chunks: list[DocumentChunk]) -> int:
|
||||
for chunk in chunks:
|
||||
self.session.add(
|
||||
DocumentChunkRow(
|
||||
id=chunk.id, document_id=chunk.document_id, index=chunk.index,
|
||||
content=chunk.content, embedding_json=list(chunk.embedding_json),
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
return len(chunks)
|
||||
|
||||
async def count_for_source(self, source_id: str) -> int:
|
||||
stmt = (
|
||||
select(func.count())
|
||||
.select_from(DocumentChunkRow)
|
||||
.join(DocumentRow, DocumentRow.id == DocumentChunkRow.document_id)
|
||||
.where(DocumentRow.source_id == source_id)
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
async def count_for_sources(self, source_ids: list[str]) -> dict[str, int]:
|
||||
"""Chunk counts for many sources in one grouped query.
|
||||
|
||||
Listing sources needs a count per row; doing that one source at a time
|
||||
is an N+1. ``GROUP BY`` collapses it into a single round trip, and
|
||||
sources with no chunks are simply absent from the result.
|
||||
"""
|
||||
if not source_ids:
|
||||
return {}
|
||||
stmt = (
|
||||
select(DocumentRow.source_id, func.count(DocumentChunkRow.id))
|
||||
.select_from(DocumentChunkRow)
|
||||
.join(DocumentRow, DocumentRow.id == DocumentChunkRow.document_id)
|
||||
.where(DocumentRow.source_id.in_(source_ids))
|
||||
.group_by(DocumentRow.source_id)
|
||||
)
|
||||
rows = (await self.session.execute(stmt)).all()
|
||||
return {str(source_id): int(count) for source_id, count in rows}
|
||||
|
||||
async def keyword_search(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
query: str,
|
||||
project_id: str | None = None,
|
||||
limit: int = 5,
|
||||
) -> list[dict]:
|
||||
"""MVP retrieval: case-insensitive substring match over chunk text.
|
||||
|
||||
A real deployment swaps this for a vector search through
|
||||
``VectorStorePort``; the tenant scoping stays identical.
|
||||
"""
|
||||
needle = f"%{query.strip().lower()}%"
|
||||
if not query.strip():
|
||||
return []
|
||||
|
||||
stmt = (
|
||||
select(
|
||||
DocumentChunkRow.id,
|
||||
DocumentChunkRow.content,
|
||||
DocumentChunkRow.index,
|
||||
DocumentRow.title,
|
||||
KnowledgeSourceRow.id,
|
||||
KnowledgeSourceRow.name,
|
||||
KnowledgeSourceRow.project_id,
|
||||
)
|
||||
.join(DocumentRow, DocumentRow.id == DocumentChunkRow.document_id)
|
||||
.join(KnowledgeSourceRow, KnowledgeSourceRow.id == DocumentRow.source_id)
|
||||
.where(
|
||||
KnowledgeSourceRow.organization_id == organization_id,
|
||||
func.lower(DocumentChunkRow.content).like(needle),
|
||||
)
|
||||
.order_by(DocumentChunkRow.document_id, DocumentChunkRow.index)
|
||||
.limit(min(limit, 50))
|
||||
)
|
||||
if project_id is not None:
|
||||
# An org-wide source (project_id is NULL) is visible to every project.
|
||||
stmt = stmt.where(
|
||||
or_(
|
||||
KnowledgeSourceRow.project_id.is_(None),
|
||||
KnowledgeSourceRow.project_id == project_id,
|
||||
)
|
||||
)
|
||||
|
||||
rows = (await self.session.execute(stmt)).all()
|
||||
return [
|
||||
{
|
||||
"chunk_id": chunk_id,
|
||||
"content": content,
|
||||
"index": index,
|
||||
"document_title": title,
|
||||
"source_id": source_id,
|
||||
"source_name": source_name,
|
||||
"project_id": source_project_id,
|
||||
}
|
||||
for chunk_id, content, index, title, source_id, source_name, source_project_id in rows
|
||||
]
|
||||
|
||||
|
||||
class EmbeddingJobRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create(self, job: EmbeddingJob) -> EmbeddingJob:
|
||||
row = EmbeddingJobRow(
|
||||
id=job.id, source_id=job.source_id, status=str(job.status),
|
||||
error=job.error, created_at=job.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return EmbeddingJob(
|
||||
id=row.id, source_id=row.source_id,
|
||||
status=EmbeddingJobStatus(row.status), error=row.error,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
async def latest_for_source(self, source_id: str) -> EmbeddingJob | None:
|
||||
stmt = (
|
||||
select(EmbeddingJobRow)
|
||||
.where(EmbeddingJobRow.source_id == source_id)
|
||||
.order_by(EmbeddingJobRow.created_at.desc())
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalars().first()
|
||||
if row is None:
|
||||
return None
|
||||
return EmbeddingJob(
|
||||
id=row.id, source_id=row.source_id,
|
||||
status=EmbeddingJobStatus(row.status), error=row.error,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
async def set_status(
|
||||
self, job_id: str, status: EmbeddingJobStatus, *, error: str | None = None
|
||||
) -> None:
|
||||
row = await self.session.get(EmbeddingJobRow, job_id)
|
||||
if row:
|
||||
row.status = str(status)
|
||||
if error is not None:
|
||||
row.error = error
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
class RetrievalPolicyRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_org(self, organization_id: str) -> list[RetrievalPolicy]:
|
||||
stmt = select(RetrievalPolicyRow).where(
|
||||
RetrievalPolicyRow.organization_id == organization_id
|
||||
)
|
||||
return [
|
||||
RetrievalPolicy(
|
||||
id=r.id, organization_id=r.organization_id, name=r.name,
|
||||
scope_json=dict(r.scope_json or {}), created_at=r.created_at,
|
||||
)
|
||||
for r in (await self.session.execute(stmt)).scalars()
|
||||
]
|
||||
|
||||
async def create(self, policy: RetrievalPolicy) -> RetrievalPolicy:
|
||||
row = RetrievalPolicyRow(
|
||||
id=policy.id, organization_id=policy.organization_id, name=policy.name,
|
||||
scope_json=dict(policy.scope_json), created_at=policy.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return RetrievalPolicy(
|
||||
id=row.id, organization_id=row.organization_id, name=row.name,
|
||||
scope_json=dict(row.scope_json or {}), created_at=row.created_at,
|
||||
)
|
||||
@@ -0,0 +1,253 @@
|
||||
"""Repositories for Projects, their Git repositories, credentials and preview
|
||||
domains."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.projects.models import (
|
||||
Project,
|
||||
ProjectGitCredential,
|
||||
ProjectGitRepo,
|
||||
)
|
||||
from app.infrastructure.db.models import (
|
||||
PreviewDomainRow,
|
||||
PreviewPortRow,
|
||||
ProjectGitCredentialRow,
|
||||
ProjectRepositoryRow,
|
||||
ProjectRow,
|
||||
)
|
||||
|
||||
|
||||
def _project_to_domain(row: ProjectRow) -> Project:
|
||||
return Project(
|
||||
id=row.id, organization_id=row.organization_id, name=row.name, slug=row.slug,
|
||||
description=row.description, created_by_user_id=row.created_by_user_id,
|
||||
default_branch=row.default_branch, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _repo_to_domain(row: ProjectRepositoryRow) -> ProjectGitRepo:
|
||||
return ProjectGitRepo(
|
||||
id=row.id, project_id=row.project_id, provider=row.provider, owner=row.owner,
|
||||
repo_name=row.repo_name, clone_url=row.clone_url, web_url=row.web_url,
|
||||
default_branch=row.default_branch, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class ProjectGitRepoRepository:
|
||||
"""Persistence for the Git repository attached to a project."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get_for_project(self, project_id: str) -> ProjectGitRepo | None:
|
||||
stmt = select(ProjectRepositoryRow).where(
|
||||
ProjectRepositoryRow.project_id == project_id
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _repo_to_domain(row) if row else None
|
||||
|
||||
async def upsert(self, repo: ProjectGitRepo) -> ProjectGitRepo:
|
||||
existing = await self.get_for_project(repo.project_id)
|
||||
if existing:
|
||||
row = await self.session.get(ProjectRepositoryRow, existing.id)
|
||||
assert row is not None
|
||||
row.provider = repo.provider
|
||||
row.owner = repo.owner
|
||||
row.repo_name = repo.repo_name
|
||||
row.clone_url = repo.clone_url
|
||||
row.web_url = repo.web_url
|
||||
row.default_branch = repo.default_branch
|
||||
await self.session.flush()
|
||||
return _repo_to_domain(row)
|
||||
|
||||
row = ProjectRepositoryRow(
|
||||
id=repo.id, project_id=repo.project_id, provider=repo.provider,
|
||||
owner=repo.owner, repo_name=repo.repo_name, clone_url=repo.clone_url,
|
||||
web_url=repo.web_url, default_branch=repo.default_branch,
|
||||
created_at=repo.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _repo_to_domain(row)
|
||||
|
||||
|
||||
class ProjectGitCredentialRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get_for_project(self, project_id: str) -> ProjectGitCredential | None:
|
||||
stmt = select(ProjectGitCredentialRow).where(
|
||||
ProjectGitCredentialRow.project_id == project_id
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
return None
|
||||
return ProjectGitCredential(
|
||||
id=row.id, project_id=row.project_id, username=row.username,
|
||||
encrypted_password=row.encrypted_password, created_at=row.created_at,
|
||||
)
|
||||
|
||||
async def upsert(self, credential: ProjectGitCredential) -> ProjectGitCredential:
|
||||
existing = await self.get_for_project(credential.project_id)
|
||||
if existing:
|
||||
row = await self.session.get(ProjectGitCredentialRow, existing.id)
|
||||
assert row is not None
|
||||
row.username = credential.username
|
||||
row.encrypted_password = credential.encrypted_password
|
||||
await self.session.flush()
|
||||
return ProjectGitCredential(
|
||||
id=row.id, project_id=row.project_id, username=row.username,
|
||||
encrypted_password=row.encrypted_password, created_at=row.created_at,
|
||||
)
|
||||
|
||||
row = ProjectGitCredentialRow(
|
||||
id=credential.id, project_id=credential.project_id,
|
||||
username=credential.username,
|
||||
encrypted_password=credential.encrypted_password,
|
||||
created_at=credential.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return credential
|
||||
|
||||
|
||||
class ProjectRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, project_id: str) -> Project | None:
|
||||
row = await self.session.get(ProjectRow, project_id)
|
||||
return _project_to_domain(row) if row else None
|
||||
|
||||
async def list_for_org(
|
||||
self, organization_id: str, *, limit: int = 100, offset: int = 0
|
||||
) -> list[Project]:
|
||||
stmt = (
|
||||
select(ProjectRow)
|
||||
.where(ProjectRow.organization_id == organization_id)
|
||||
.order_by(ProjectRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return [_project_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_slugs(self, organization_id: str) -> set[str]:
|
||||
stmt = select(ProjectRow.slug).where(ProjectRow.organization_id == organization_id)
|
||||
return set((await self.session.execute(stmt)).scalars().all())
|
||||
|
||||
async def create(self, project: Project) -> Project:
|
||||
row = ProjectRow(
|
||||
id=project.id, organization_id=project.organization_id, name=project.name,
|
||||
slug=project.slug, description=project.description,
|
||||
created_by_user_id=project.created_by_user_id,
|
||||
default_branch=project.default_branch, created_at=project.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _project_to_domain(row)
|
||||
|
||||
async def delete(self, project_id: str) -> None:
|
||||
row = await self.session.get(ProjectRow, project_id)
|
||||
if row:
|
||||
await self.session.delete(row)
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
class PreviewDomainRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
# -- ports ---------------------------------------------------------------
|
||||
|
||||
async def get_port(self, project_id: str, service_name: str) -> int | None:
|
||||
stmt = select(PreviewPortRow.port).where(
|
||||
PreviewPortRow.project_id == project_id,
|
||||
PreviewPortRow.service_name == service_name,
|
||||
)
|
||||
return (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
|
||||
async def list_used_ports(self, *, start: int, end: int) -> set[int]:
|
||||
stmt = select(PreviewPortRow.port).where(
|
||||
PreviewPortRow.port.between(start, end)
|
||||
)
|
||||
return set((await self.session.execute(stmt)).scalars().all())
|
||||
|
||||
async def add_port(
|
||||
self, *, organization_id: str, project_id: str, service_name: str, port: int
|
||||
) -> None:
|
||||
self.session.add(
|
||||
PreviewPortRow(
|
||||
organization_id=organization_id, project_id=project_id,
|
||||
service_name=service_name, port=port, status="active",
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
|
||||
# -- domains -------------------------------------------------------------
|
||||
|
||||
async def get_by_host(self, host: str) -> dict | None:
|
||||
stmt = select(PreviewDomainRow).where(PreviewDomainRow.host == host)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return self._row_to_dict(row) if row else None
|
||||
|
||||
async def get(self, domain_id: str) -> dict | None:
|
||||
row = await self.session.get(PreviewDomainRow, domain_id)
|
||||
return self._row_to_dict(row) if row else None
|
||||
|
||||
async def list_for_project(self, project_id: str) -> list[dict]:
|
||||
stmt = (
|
||||
select(PreviewDomainRow)
|
||||
.where(PreviewDomainRow.project_id == project_id)
|
||||
.order_by(PreviewDomainRow.created_at)
|
||||
)
|
||||
return [self._row_to_dict(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def count_for_project(self, project_id: str) -> int:
|
||||
return len(await self.list_for_project(project_id))
|
||||
|
||||
async def create(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
project_id: str,
|
||||
host: str,
|
||||
port: int,
|
||||
created_by_user_id: str | None,
|
||||
) -> dict:
|
||||
row = PreviewDomainRow(
|
||||
organization_id=organization_id, project_id=project_id, host=host,
|
||||
port=port, status="active", protected=False,
|
||||
created_by_user_id=created_by_user_id,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return self._row_to_dict(row)
|
||||
|
||||
async def set_status(self, domain_id: str, status: str) -> None:
|
||||
row = await self.session.get(PreviewDomainRow, domain_id)
|
||||
if row:
|
||||
row.status = status
|
||||
await self.session.flush()
|
||||
|
||||
async def delete(self, domain_id: str) -> None:
|
||||
row = await self.session.get(PreviewDomainRow, domain_id)
|
||||
if row:
|
||||
await self.session.delete(row)
|
||||
await self.session.flush()
|
||||
|
||||
@staticmethod
|
||||
def _row_to_dict(row: PreviewDomainRow) -> dict:
|
||||
return {
|
||||
"id": row.id,
|
||||
"organization_id": row.organization_id,
|
||||
"project_id": row.project_id,
|
||||
"host": row.host,
|
||||
"port": row.port,
|
||||
"status": row.status,
|
||||
"protected": row.protected,
|
||||
"created_by_user_id": row.created_by_user_id,
|
||||
"created_at": row.created_at,
|
||||
}
|
||||
@@ -0,0 +1,362 @@
|
||||
"""Repositories for the Sandbox Gateway.
|
||||
|
||||
Every tenant-scoped read joins back to ``workspaces.organization_id`` so a
|
||||
workspace or session id from another tenant resolves to ``None``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.base import utcnow
|
||||
from app.domain.sandbox.models import (
|
||||
CommandRun,
|
||||
NetworkPolicy,
|
||||
PreviewUrl,
|
||||
SandboxProviderKind,
|
||||
SandboxResourceLimit,
|
||||
SandboxSession,
|
||||
SandboxTemplate,
|
||||
SessionStatus,
|
||||
Snapshot,
|
||||
Workspace,
|
||||
WorkspaceStatus,
|
||||
)
|
||||
from app.infrastructure.db.models import (
|
||||
CommandRunRow,
|
||||
PreviewUrlRow,
|
||||
SandboxResourceLimitRow,
|
||||
SandboxSessionRow,
|
||||
SandboxTemplateRow,
|
||||
SnapshotRow,
|
||||
WorkspaceRow,
|
||||
)
|
||||
|
||||
|
||||
def _template(row: SandboxTemplateRow) -> SandboxTemplate:
|
||||
return SandboxTemplate(
|
||||
id=row.id, name=row.name, provider=row.provider, image=row.image,
|
||||
description=row.description, default_cpu=row.default_cpu,
|
||||
default_memory_mb=row.default_memory_mb, default_disk_mb=row.default_disk_mb,
|
||||
is_system=row.is_system, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _workspace(row: WorkspaceRow) -> Workspace:
|
||||
return Workspace(
|
||||
id=row.id, organization_id=row.organization_id, project_id=row.project_id,
|
||||
template_id=row.template_id, provider=row.provider,
|
||||
external_workspace_id=row.external_workspace_id,
|
||||
status=WorkspaceStatus(row.status), created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _session(row: SandboxSessionRow) -> SandboxSession:
|
||||
return SandboxSession(
|
||||
id=row.id, workspace_id=row.workspace_id,
|
||||
external_session_id=row.external_session_id,
|
||||
status=SessionStatus(row.status), started_at=row.started_at,
|
||||
stopped_at=row.stopped_at, last_active_at=row.last_active_at,
|
||||
)
|
||||
|
||||
|
||||
class SandboxTemplateRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list(self) -> list[SandboxTemplate]:
|
||||
stmt = select(SandboxTemplateRow).order_by(SandboxTemplateRow.name)
|
||||
return [_template(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def get(self, template_id: str) -> SandboxTemplate | None:
|
||||
row = await self.session.get(SandboxTemplateRow, template_id)
|
||||
return _template(row) if row else None
|
||||
|
||||
async def get_by_name(self, name: str) -> SandboxTemplate | None:
|
||||
stmt = select(SandboxTemplateRow).where(SandboxTemplateRow.name == name)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _template(row) if row else None
|
||||
|
||||
async def create(self, template: SandboxTemplate) -> SandboxTemplate:
|
||||
row = SandboxTemplateRow(
|
||||
id=template.id, name=template.name, provider=str(template.provider),
|
||||
image=template.image, description=template.description,
|
||||
default_cpu=template.default_cpu,
|
||||
default_memory_mb=template.default_memory_mb,
|
||||
default_disk_mb=template.default_disk_mb, is_system=template.is_system,
|
||||
created_at=template.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _template(row)
|
||||
|
||||
|
||||
class WorkspaceRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, workspace_id: str) -> Workspace | None:
|
||||
"""Unscoped lookup. Platform components only — tenant paths use
|
||||
:meth:`get_for_org` so a foreign id resolves to ``None``."""
|
||||
row = await self.session.get(WorkspaceRow, workspace_id)
|
||||
return _workspace(row) if row else None
|
||||
|
||||
async def get_for_org(self, organization_id: str, workspace_id: str) -> Workspace | None:
|
||||
stmt = select(WorkspaceRow).where(
|
||||
WorkspaceRow.id == workspace_id,
|
||||
WorkspaceRow.organization_id == organization_id,
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _workspace(row) if row else None
|
||||
|
||||
async def list_for_project(
|
||||
self, project_id: str, *, limit: int = 100, offset: int = 0
|
||||
) -> list[Workspace]:
|
||||
stmt = (
|
||||
select(WorkspaceRow)
|
||||
.where(WorkspaceRow.project_id == project_id)
|
||||
.order_by(WorkspaceRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return [_workspace(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, workspace: Workspace) -> Workspace:
|
||||
row = WorkspaceRow(
|
||||
id=workspace.id, organization_id=workspace.organization_id,
|
||||
project_id=workspace.project_id, template_id=workspace.template_id,
|
||||
provider=str(workspace.provider),
|
||||
external_workspace_id=workspace.external_workspace_id,
|
||||
status=str(workspace.status), created_at=workspace.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _workspace(row)
|
||||
|
||||
async def set_status(self, workspace_id: str, status: WorkspaceStatus) -> None:
|
||||
row = await self.session.get(WorkspaceRow, workspace_id)
|
||||
if row:
|
||||
row.status = str(status)
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
class SandboxSessionRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, session_id: str) -> SandboxSession | None:
|
||||
row = await self.session.get(SandboxSessionRow, session_id)
|
||||
return _session(row) if row else None
|
||||
|
||||
async def get_for_org(self, organization_id: str, session_id: str) -> SandboxSession | None:
|
||||
"""Resolve a session through its workspace, which carries the tenant."""
|
||||
stmt = (
|
||||
select(SandboxSessionRow)
|
||||
.join(WorkspaceRow, WorkspaceRow.id == SandboxSessionRow.workspace_id)
|
||||
.where(
|
||||
SandboxSessionRow.id == session_id,
|
||||
WorkspaceRow.organization_id == organization_id,
|
||||
)
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _session(row) if row else None
|
||||
|
||||
async def list_for_workspace(self, workspace_id: str) -> list[SandboxSession]:
|
||||
stmt = (
|
||||
select(SandboxSessionRow)
|
||||
.where(SandboxSessionRow.workspace_id == workspace_id)
|
||||
.order_by(SandboxSessionRow.started_at.desc())
|
||||
)
|
||||
return [_session(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_stale_active(
|
||||
self, *, idle_before: datetime, limit: int = 200
|
||||
) -> list[SandboxSession]:
|
||||
"""Active sessions untouched since ``idle_before``.
|
||||
|
||||
The gateway reaper's query. Only ``starting``/``running``/``paused``
|
||||
count as active — a stopped session is already gone and must not be
|
||||
stopped twice.
|
||||
"""
|
||||
stmt = (
|
||||
select(SandboxSessionRow)
|
||||
.where(
|
||||
SandboxSessionRow.status.in_(
|
||||
[
|
||||
str(SessionStatus.STARTING),
|
||||
str(SessionStatus.RUNNING),
|
||||
str(SessionStatus.PAUSED),
|
||||
]
|
||||
),
|
||||
SandboxSessionRow.last_active_at < idle_before,
|
||||
)
|
||||
.order_by(SandboxSessionRow.last_active_at.asc())
|
||||
.limit(min(limit, 500))
|
||||
)
|
||||
return [_session(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, sandbox_session: SandboxSession) -> SandboxSession:
|
||||
row = SandboxSessionRow(
|
||||
id=sandbox_session.id, workspace_id=sandbox_session.workspace_id,
|
||||
external_session_id=sandbox_session.external_session_id,
|
||||
status=str(sandbox_session.status), started_at=sandbox_session.started_at,
|
||||
stopped_at=sandbox_session.stopped_at,
|
||||
last_active_at=sandbox_session.last_active_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _session(row)
|
||||
|
||||
async def set_status(
|
||||
self, session_id: str, status: SessionStatus, *, stopped: bool = False
|
||||
) -> None:
|
||||
row = await self.session.get(SandboxSessionRow, session_id)
|
||||
if row:
|
||||
row.status = str(status)
|
||||
row.last_active_at = utcnow()
|
||||
if stopped:
|
||||
row.stopped_at = utcnow()
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
class CommandRunRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_session(self, session_id: str) -> list[CommandRun]:
|
||||
stmt = (
|
||||
select(CommandRunRow)
|
||||
.where(CommandRunRow.session_id == session_id)
|
||||
.order_by(CommandRunRow.created_at)
|
||||
)
|
||||
return [
|
||||
CommandRun(
|
||||
id=r.id, session_id=r.session_id, command=r.command,
|
||||
exit_code=r.exit_code, stdout_ref=r.stdout_ref, stderr_ref=r.stderr_ref,
|
||||
duration_ms=r.duration_ms, truncated=r.truncated, created_at=r.created_at,
|
||||
)
|
||||
for r in (await self.session.execute(stmt)).scalars()
|
||||
]
|
||||
|
||||
async def create(self, run: CommandRun) -> CommandRun:
|
||||
row = CommandRunRow(
|
||||
id=run.id, session_id=run.session_id, command=run.command,
|
||||
exit_code=run.exit_code, stdout_ref=run.stdout_ref, stderr_ref=run.stderr_ref,
|
||||
duration_ms=run.duration_ms, truncated=run.truncated, created_at=run.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return run
|
||||
|
||||
|
||||
class SnapshotRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_workspace(self, workspace_id: str) -> list[Snapshot]:
|
||||
stmt = (
|
||||
select(SnapshotRow)
|
||||
.where(SnapshotRow.workspace_id == workspace_id)
|
||||
.order_by(SnapshotRow.created_at.desc())
|
||||
)
|
||||
return [
|
||||
Snapshot(
|
||||
id=r.id, workspace_id=r.workspace_id, name=r.name,
|
||||
external_snapshot_id=r.external_snapshot_id, created_at=r.created_at,
|
||||
)
|
||||
for r in (await self.session.execute(stmt)).scalars()
|
||||
]
|
||||
|
||||
async def create(self, snapshot: Snapshot) -> Snapshot:
|
||||
row = SnapshotRow(
|
||||
id=snapshot.id, workspace_id=snapshot.workspace_id, name=snapshot.name,
|
||||
external_snapshot_id=snapshot.external_snapshot_id,
|
||||
created_at=snapshot.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return snapshot
|
||||
|
||||
|
||||
class PreviewUrlRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_workspace(self, workspace_id: str) -> list[PreviewUrl]:
|
||||
stmt = select(PreviewUrlRow).where(PreviewUrlRow.workspace_id == workspace_id)
|
||||
return [
|
||||
PreviewUrl(
|
||||
id=r.id, workspace_id=r.workspace_id, port=r.port, url=r.url,
|
||||
created_at=r.created_at,
|
||||
)
|
||||
for r in (await self.session.execute(stmt)).scalars()
|
||||
]
|
||||
|
||||
async def create(self, preview: PreviewUrl) -> PreviewUrl:
|
||||
row = PreviewUrlRow(
|
||||
id=preview.id, workspace_id=preview.workspace_id, port=preview.port,
|
||||
url=preview.url, created_at=preview.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return preview
|
||||
|
||||
|
||||
class ResourceLimitRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get_for_scope(self, scope: str, scope_id: str | None) -> SandboxResourceLimit | None:
|
||||
stmt = select(SandboxResourceLimitRow).where(
|
||||
SandboxResourceLimitRow.scope == scope,
|
||||
SandboxResourceLimitRow.scope_id == scope_id,
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
return None
|
||||
return SandboxResourceLimit(
|
||||
id=row.id, scope=row.scope, scope_id=row.scope_id, cpu=row.cpu,
|
||||
memory_mb=row.memory_mb, disk_mb=row.disk_mb,
|
||||
timeout_seconds=row.timeout_seconds, max_output_bytes=row.max_output_bytes,
|
||||
network_policy=NetworkPolicy(row.network_policy),
|
||||
)
|
||||
|
||||
async def upsert(self, limit: SandboxResourceLimit) -> SandboxResourceLimit:
|
||||
stmt = select(SandboxResourceLimitRow).where(
|
||||
SandboxResourceLimitRow.scope == limit.scope,
|
||||
SandboxResourceLimitRow.scope_id == limit.scope_id,
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
row = SandboxResourceLimitRow(
|
||||
id=limit.id, scope=limit.scope, scope_id=limit.scope_id, cpu=limit.cpu,
|
||||
memory_mb=limit.memory_mb, disk_mb=limit.disk_mb,
|
||||
timeout_seconds=limit.timeout_seconds,
|
||||
max_output_bytes=limit.max_output_bytes,
|
||||
network_policy=str(limit.network_policy),
|
||||
)
|
||||
self.session.add(row)
|
||||
else:
|
||||
row.cpu = limit.cpu
|
||||
row.memory_mb = limit.memory_mb
|
||||
row.disk_mb = limit.disk_mb
|
||||
row.timeout_seconds = limit.timeout_seconds
|
||||
row.max_output_bytes = limit.max_output_bytes
|
||||
row.network_policy = str(limit.network_policy)
|
||||
await self.session.flush()
|
||||
return limit
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CommandRunRepository",
|
||||
"PreviewUrlRepository",
|
||||
"ResourceLimitRepository",
|
||||
"SandboxProviderKind",
|
||||
"SandboxSessionRepository",
|
||||
"SandboxTemplateRepository",
|
||||
"SnapshotRepository",
|
||||
"WorkspaceRepository",
|
||||
]
|
||||
@@ -0,0 +1,684 @@
|
||||
"""Repositories for the AI Workforce bounded context.
|
||||
|
||||
Covers agent templates, employees, prompt versions, tool permissions, and the
|
||||
whole Manager-Chat chain: threads -> messages -> tasks -> runs -> steps ->
|
||||
tool calls -> artifacts.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.workforce.models import (
|
||||
AgentTemplate,
|
||||
Artifact,
|
||||
Employee,
|
||||
EmployeeStatus,
|
||||
EmployeeToolPermission,
|
||||
Message,
|
||||
MessageRole,
|
||||
PromptVersion,
|
||||
RoleKey,
|
||||
Run,
|
||||
RunStatus,
|
||||
RunStep,
|
||||
Task,
|
||||
TaskStatus,
|
||||
Thread,
|
||||
ToolCall,
|
||||
ToolCallStatus,
|
||||
)
|
||||
from app.infrastructure.db.models import (
|
||||
AgentTemplateRow,
|
||||
ArtifactRow,
|
||||
EmployeeRow,
|
||||
EmployeeToolPermissionRow,
|
||||
MessageRow,
|
||||
PromptVersionRow,
|
||||
RunRow,
|
||||
RunStepRow,
|
||||
TaskRow,
|
||||
ThreadRow,
|
||||
ToolCallRow,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mappers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _template(row: AgentTemplateRow) -> AgentTemplate:
|
||||
return AgentTemplate(
|
||||
id=row.id, name=row.name, role_key=RoleKey(row.role_key),
|
||||
description=row.description, default_prompt=row.default_prompt,
|
||||
default_tools_json=list(row.default_tools_json or []),
|
||||
default_routing_policy_id=row.default_routing_policy_id,
|
||||
is_system=row.is_system, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _employee(row: EmployeeRow) -> Employee:
|
||||
return Employee(
|
||||
id=row.id, organization_id=row.organization_id, project_id=row.project_id,
|
||||
template_id=row.template_id, name=row.name, role_key=RoleKey(row.role_key),
|
||||
status=EmployeeStatus(row.status), routing_policy_id=row.routing_policy_id,
|
||||
budget_limit_json=dict(row.budget_limit_json or {}), created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _thread(row: ThreadRow) -> Thread:
|
||||
return Thread(
|
||||
id=row.id, organization_id=row.organization_id, project_id=row.project_id,
|
||||
title=row.title, created_by_user_id=row.created_by_user_id,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _message(row: MessageRow) -> Message:
|
||||
return Message(
|
||||
id=row.id, thread_id=row.thread_id, role=MessageRole(row.role),
|
||||
content=row.content, employee_id=row.employee_id, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _task(row: TaskRow) -> Task:
|
||||
return Task(
|
||||
id=row.id, organization_id=row.organization_id, project_id=row.project_id,
|
||||
thread_id=row.thread_id, title=row.title, description=row.description,
|
||||
status=TaskStatus(row.status),
|
||||
assigned_employee_id=row.assigned_employee_id,
|
||||
created_by_user_id=row.created_by_user_id,
|
||||
created_at=row.created_at, updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
|
||||
def _run(row: RunRow) -> Run:
|
||||
return Run(
|
||||
id=row.id, task_id=row.task_id, employee_id=row.employee_id,
|
||||
status=RunStatus(row.status), error=row.error, started_at=row.started_at,
|
||||
finished_at=row.finished_at, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _run_step(row: RunStepRow) -> RunStep:
|
||||
return RunStep(
|
||||
id=row.id, run_id=row.run_id, index=row.index, kind=row.kind,
|
||||
title=row.title, payload_json=dict(row.payload_json or {}),
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _tool_call(row: ToolCallRow) -> ToolCall:
|
||||
return ToolCall(
|
||||
id=row.id, run_id=row.run_id, tool_key=row.tool_key,
|
||||
arguments_json=dict(row.arguments_json or {}),
|
||||
result_json=dict(row.result_json or {}),
|
||||
status=ToolCallStatus(row.status), requires_approval=row.requires_approval,
|
||||
approved_by_user_id=row.approved_by_user_id, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _artifact(row: ArtifactRow) -> Artifact:
|
||||
return Artifact(
|
||||
id=row.id, organization_id=row.organization_id, project_id=row.project_id,
|
||||
run_id=row.run_id, name=row.name, kind=row.kind, content_ref=row.content_ref,
|
||||
size_bytes=row.size_bytes, created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Templates & employees
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AgentTemplateRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list(self) -> list[AgentTemplate]:
|
||||
stmt = select(AgentTemplateRow).order_by(AgentTemplateRow.role_key)
|
||||
return [_template(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def get(self, template_id: str) -> AgentTemplate | None:
|
||||
row = await self.session.get(AgentTemplateRow, template_id)
|
||||
return _template(row) if row else None
|
||||
|
||||
async def get_by_name(self, name: str) -> AgentTemplate | None:
|
||||
stmt = select(AgentTemplateRow).where(AgentTemplateRow.name == name)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _template(row) if row else None
|
||||
|
||||
async def create(self, template: AgentTemplate) -> AgentTemplate:
|
||||
row = AgentTemplateRow(
|
||||
id=template.id, name=template.name, role_key=str(template.role_key),
|
||||
description=template.description, default_prompt=template.default_prompt,
|
||||
default_tools_json=list(template.default_tools_json),
|
||||
default_routing_policy_id=template.default_routing_policy_id,
|
||||
is_system=template.is_system, created_at=template.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _template(row)
|
||||
|
||||
|
||||
class EmployeeRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, employee_id: str) -> Employee | None:
|
||||
row = await self.session.get(EmployeeRow, employee_id)
|
||||
return _employee(row) if row else None
|
||||
|
||||
async def get_for_org(self, organization_id: str, employee_id: str) -> Employee | None:
|
||||
"""Tenant-scoped lookup: an id from another org resolves to ``None``."""
|
||||
stmt = select(EmployeeRow).where(
|
||||
EmployeeRow.id == employee_id,
|
||||
EmployeeRow.organization_id == organization_id,
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _employee(row) if row else None
|
||||
|
||||
async def list_for_org(
|
||||
self, organization_id: str, *, limit: int = 100, offset: int = 0
|
||||
) -> list[Employee]:
|
||||
stmt = (
|
||||
select(EmployeeRow)
|
||||
.where(EmployeeRow.organization_id == organization_id)
|
||||
.order_by(EmployeeRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return [_employee(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_for_project(self, project_id: str) -> list[Employee]:
|
||||
stmt = (
|
||||
select(EmployeeRow)
|
||||
.where(EmployeeRow.project_id == project_id)
|
||||
.order_by(EmployeeRow.created_at.desc())
|
||||
)
|
||||
return [_employee(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def count_for_org(self, organization_id: str) -> int:
|
||||
stmt = select(func.count()).select_from(EmployeeRow).where(
|
||||
EmployeeRow.organization_id == organization_id
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
async def create(self, employee: Employee) -> Employee:
|
||||
row = EmployeeRow(
|
||||
id=employee.id, organization_id=employee.organization_id,
|
||||
project_id=employee.project_id, template_id=employee.template_id,
|
||||
name=employee.name, role_key=str(employee.role_key),
|
||||
status=str(employee.status),
|
||||
routing_policy_id=employee.routing_policy_id,
|
||||
budget_limit_json=dict(employee.budget_limit_json),
|
||||
created_at=employee.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _employee(row)
|
||||
|
||||
async def update(self, employee: Employee) -> Employee:
|
||||
row = await self.session.get(EmployeeRow, employee.id)
|
||||
if row is None:
|
||||
raise LookupError("employee not found")
|
||||
row.name = employee.name
|
||||
row.role_key = str(employee.role_key)
|
||||
row.status = str(employee.status)
|
||||
row.project_id = employee.project_id
|
||||
row.routing_policy_id = employee.routing_policy_id
|
||||
row.budget_limit_json = dict(employee.budget_limit_json)
|
||||
await self.session.flush()
|
||||
return _employee(row)
|
||||
|
||||
|
||||
class PromptVersionRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_employee(self, employee_id: str) -> list[PromptVersion]:
|
||||
stmt = (
|
||||
select(PromptVersionRow)
|
||||
.where(PromptVersionRow.employee_id == employee_id)
|
||||
.order_by(PromptVersionRow.version.desc())
|
||||
)
|
||||
return [
|
||||
PromptVersion(
|
||||
id=r.id, employee_id=r.employee_id, version=r.version,
|
||||
content=r.content, created_by_user_id=r.created_by_user_id,
|
||||
created_at=r.created_at,
|
||||
)
|
||||
for r in (await self.session.execute(stmt)).scalars()
|
||||
]
|
||||
|
||||
async def latest(self, employee_id: str) -> PromptVersion | None:
|
||||
versions = await self.list_for_employee(employee_id)
|
||||
return versions[0] if versions else None
|
||||
|
||||
async def next_version_number(self, employee_id: str) -> int:
|
||||
stmt = select(func.max(PromptVersionRow.version)).where(
|
||||
PromptVersionRow.employee_id == employee_id
|
||||
)
|
||||
current = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return int(current or 0) + 1
|
||||
|
||||
async def create(self, version: PromptVersion) -> PromptVersion:
|
||||
row = PromptVersionRow(
|
||||
id=version.id, employee_id=version.employee_id, version=version.version,
|
||||
content=version.content, created_by_user_id=version.created_by_user_id,
|
||||
created_at=version.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return PromptVersion(
|
||||
id=row.id, employee_id=row.employee_id, version=row.version,
|
||||
content=row.content, created_by_user_id=row.created_by_user_id,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class EmployeeToolPermissionRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_employee(self, employee_id: str) -> list[EmployeeToolPermission]:
|
||||
stmt = select(EmployeeToolPermissionRow).where(
|
||||
EmployeeToolPermissionRow.employee_id == employee_id
|
||||
)
|
||||
return [
|
||||
EmployeeToolPermission(
|
||||
id=r.id, employee_id=r.employee_id, tool_key=r.tool_key,
|
||||
allowed=r.allowed, policy_json=dict(r.policy_json or {}),
|
||||
)
|
||||
for r in (await self.session.execute(stmt)).scalars()
|
||||
]
|
||||
|
||||
async def is_allowed(self, employee_id: str, tool_key: str) -> bool:
|
||||
stmt = select(EmployeeToolPermissionRow.allowed).where(
|
||||
EmployeeToolPermissionRow.employee_id == employee_id,
|
||||
EmployeeToolPermissionRow.tool_key == tool_key,
|
||||
)
|
||||
value = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
# Deny by default: an absent row means "not permitted".
|
||||
return bool(value)
|
||||
|
||||
async def set_permission(
|
||||
self, *, employee_id: str, tool_key: str, allowed: bool, policy: dict | None = None
|
||||
) -> EmployeeToolPermission:
|
||||
stmt = select(EmployeeToolPermissionRow).where(
|
||||
EmployeeToolPermissionRow.employee_id == employee_id,
|
||||
EmployeeToolPermissionRow.tool_key == tool_key,
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
row = EmployeeToolPermissionRow(
|
||||
employee_id=employee_id, tool_key=tool_key, allowed=allowed,
|
||||
policy_json=dict(policy or {}),
|
||||
)
|
||||
self.session.add(row)
|
||||
else:
|
||||
row.allowed = allowed
|
||||
row.policy_json = dict(policy or {})
|
||||
await self.session.flush()
|
||||
return EmployeeToolPermission(
|
||||
id=row.id, employee_id=row.employee_id, tool_key=row.tool_key,
|
||||
allowed=row.allowed, policy_json=dict(row.policy_json or {}),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Chat: threads & messages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ThreadRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, thread_id: str) -> Thread | None:
|
||||
row = await self.session.get(ThreadRow, thread_id)
|
||||
return _thread(row) if row else None
|
||||
|
||||
async def get_for_org(self, organization_id: str, thread_id: str) -> Thread | None:
|
||||
stmt = select(ThreadRow).where(
|
||||
ThreadRow.id == thread_id, ThreadRow.organization_id == organization_id
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _thread(row) if row else None
|
||||
|
||||
async def list_for_org(
|
||||
self, organization_id: str, *, limit: int = 100, offset: int = 0
|
||||
) -> list[Thread]:
|
||||
stmt = (
|
||||
select(ThreadRow)
|
||||
.where(ThreadRow.organization_id == organization_id)
|
||||
.order_by(ThreadRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return [_thread(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, thread: Thread) -> Thread:
|
||||
row = ThreadRow(
|
||||
id=thread.id, organization_id=thread.organization_id,
|
||||
project_id=thread.project_id, title=thread.title,
|
||||
created_by_user_id=thread.created_by_user_id, created_at=thread.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _thread(row)
|
||||
|
||||
|
||||
class MessageRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_thread(self, thread_id: str) -> list[Message]:
|
||||
stmt = (
|
||||
select(MessageRow)
|
||||
.where(MessageRow.thread_id == thread_id)
|
||||
.order_by(MessageRow.created_at)
|
||||
)
|
||||
return [_message(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def count_for_thread(self, thread_id: str) -> int:
|
||||
stmt = select(func.count()).select_from(MessageRow).where(
|
||||
MessageRow.thread_id == thread_id
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
async def create(self, message: Message) -> Message:
|
||||
row = MessageRow(
|
||||
id=message.id, thread_id=message.thread_id, role=str(message.role),
|
||||
content=message.content, employee_id=message.employee_id,
|
||||
created_at=message.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _message(row)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Work: tasks, runs, steps, tool calls, artifacts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TaskRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, task_id: str) -> Task | None:
|
||||
row = await self.session.get(TaskRow, task_id)
|
||||
return _task(row) if row else None
|
||||
|
||||
async def get_for_org(self, organization_id: str, task_id: str) -> Task | None:
|
||||
stmt = select(TaskRow).where(
|
||||
TaskRow.id == task_id, TaskRow.organization_id == organization_id
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _task(row) if row else None
|
||||
|
||||
async def list_for_org(
|
||||
self,
|
||||
organization_id: str,
|
||||
*,
|
||||
status: TaskStatus | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> list[Task]:
|
||||
stmt = select(TaskRow).where(TaskRow.organization_id == organization_id)
|
||||
if status is not None:
|
||||
stmt = stmt.where(TaskRow.status == str(status))
|
||||
stmt = stmt.order_by(TaskRow.created_at.desc()).limit(min(limit, 500)).offset(offset)
|
||||
return [_task(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_for_employee(self, employee_id: str) -> list[Task]:
|
||||
stmt = (
|
||||
select(TaskRow)
|
||||
.where(TaskRow.assigned_employee_id == employee_id)
|
||||
.order_by(TaskRow.created_at.desc())
|
||||
)
|
||||
return [_task(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_runnable(self, *, limit: int = 10) -> list[Task]:
|
||||
"""Tasks the runtime worker should pick up, oldest first.
|
||||
|
||||
Cross-tenant by design: the worker is a platform component, not a
|
||||
tenant. Only ``new``/``triaged`` qualify — anything else is either
|
||||
already executing or closed, and re-running it would double-bill.
|
||||
"""
|
||||
stmt = (
|
||||
select(TaskRow)
|
||||
.where(TaskRow.status.in_([str(TaskStatus.NEW), str(TaskStatus.TRIAGED)]))
|
||||
.order_by(TaskRow.created_at.asc())
|
||||
.limit(min(limit, 100))
|
||||
)
|
||||
return [_task(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def count_open_for_employee(self, employee_id: str) -> int:
|
||||
stmt = (
|
||||
select(func.count())
|
||||
.select_from(TaskRow)
|
||||
.where(
|
||||
TaskRow.assigned_employee_id == employee_id,
|
||||
TaskRow.status.notin_(
|
||||
[str(TaskStatus.DONE), str(TaskStatus.FAILED), str(TaskStatus.CANCELLED)]
|
||||
),
|
||||
)
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
async def create(self, task: Task) -> Task:
|
||||
row = TaskRow(
|
||||
id=task.id, organization_id=task.organization_id,
|
||||
project_id=task.project_id, thread_id=task.thread_id,
|
||||
title=task.title, description=task.description, status=str(task.status),
|
||||
assigned_employee_id=task.assigned_employee_id,
|
||||
created_by_user_id=task.created_by_user_id,
|
||||
created_at=task.created_at, updated_at=task.updated_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _task(row)
|
||||
|
||||
async def set_status(self, task_id: str, status: TaskStatus) -> Task | None:
|
||||
row = await self.session.get(TaskRow, task_id)
|
||||
if row is None:
|
||||
return None
|
||||
row.status = str(status)
|
||||
await self.session.flush()
|
||||
return _task(row)
|
||||
|
||||
|
||||
class RunRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, run_id: str) -> Run | None:
|
||||
row = await self.session.get(RunRow, run_id)
|
||||
return _run(row) if row else None
|
||||
|
||||
async def list_for_task(self, task_id: str) -> list[Run]:
|
||||
stmt = (
|
||||
select(RunRow)
|
||||
.where(RunRow.task_id == task_id)
|
||||
.order_by(RunRow.created_at.desc())
|
||||
)
|
||||
return [_run(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, run: Run) -> Run:
|
||||
row = RunRow(
|
||||
id=run.id, task_id=run.task_id, employee_id=run.employee_id,
|
||||
status=str(run.status), error=run.error, started_at=run.started_at,
|
||||
finished_at=run.finished_at, created_at=run.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _run(row)
|
||||
|
||||
async def set_status(
|
||||
self,
|
||||
run_id: str,
|
||||
status: RunStatus,
|
||||
*,
|
||||
error: str | None = None,
|
||||
started_at=None,
|
||||
finished_at=None,
|
||||
) -> Run | None:
|
||||
row = await self.session.get(RunRow, run_id)
|
||||
if row is None:
|
||||
return None
|
||||
row.status = str(status)
|
||||
if error is not None:
|
||||
row.error = error
|
||||
if started_at is not None:
|
||||
row.started_at = started_at
|
||||
if finished_at is not None:
|
||||
row.finished_at = finished_at
|
||||
await self.session.flush()
|
||||
return _run(row)
|
||||
|
||||
|
||||
class RunStepRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_run(self, run_id: str) -> list[RunStep]:
|
||||
stmt = (
|
||||
select(RunStepRow)
|
||||
.where(RunStepRow.run_id == run_id)
|
||||
.order_by(RunStepRow.index)
|
||||
)
|
||||
return [_run_step(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def next_index(self, run_id: str) -> int:
|
||||
stmt = select(func.max(RunStepRow.index)).where(RunStepRow.run_id == run_id)
|
||||
current = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return int(current if current is not None else -1) + 1
|
||||
|
||||
async def append(
|
||||
self, *, run_id: str, kind: str, title: str = "", payload: dict | None = None
|
||||
) -> RunStep:
|
||||
step = RunStep(
|
||||
run_id=run_id, index=await self.next_index(run_id), kind=kind, title=title,
|
||||
payload_json=dict(payload or {}),
|
||||
)
|
||||
row = RunStepRow(
|
||||
id=step.id, run_id=step.run_id, index=step.index, kind=step.kind,
|
||||
title=step.title, payload_json=step.payload_json, created_at=step.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _run_step(row)
|
||||
|
||||
|
||||
class ToolCallRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, call_id: str) -> ToolCall | None:
|
||||
row = await self.session.get(ToolCallRow, call_id)
|
||||
return _tool_call(row) if row else None
|
||||
|
||||
async def get_for_org(self, organization_id: str, call_id: str) -> ToolCall | None:
|
||||
"""Tenant-scoped lookup, resolved through run -> task -> organization."""
|
||||
stmt = (
|
||||
select(ToolCallRow)
|
||||
.join(RunRow, RunRow.id == ToolCallRow.run_id)
|
||||
.join(TaskRow, TaskRow.id == RunRow.task_id)
|
||||
.where(ToolCallRow.id == call_id, TaskRow.organization_id == organization_id)
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _tool_call(row) if row else None
|
||||
|
||||
async def list_for_run(self, run_id: str) -> list[ToolCall]:
|
||||
stmt = (
|
||||
select(ToolCallRow)
|
||||
.where(ToolCallRow.run_id == run_id)
|
||||
.order_by(ToolCallRow.created_at)
|
||||
)
|
||||
return [_tool_call(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_pending_approvals_for_org(
|
||||
self, organization_id: str
|
||||
) -> list[ToolCall]:
|
||||
"""Tool calls awaiting approval, scoped through the owning run -> task -> org."""
|
||||
stmt = (
|
||||
select(ToolCallRow)
|
||||
.join(RunRow, RunRow.id == ToolCallRow.run_id)
|
||||
.join(TaskRow, TaskRow.id == RunRow.task_id)
|
||||
.where(
|
||||
TaskRow.organization_id == organization_id,
|
||||
ToolCallRow.status == str(ToolCallStatus.AWAITING_APPROVAL),
|
||||
)
|
||||
.order_by(ToolCallRow.created_at)
|
||||
)
|
||||
return [_tool_call(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, call: ToolCall) -> ToolCall:
|
||||
row = ToolCallRow(
|
||||
id=call.id, run_id=call.run_id, tool_key=call.tool_key,
|
||||
arguments_json=dict(call.arguments_json),
|
||||
result_json=dict(call.result_json), status=str(call.status),
|
||||
requires_approval=call.requires_approval,
|
||||
approved_by_user_id=call.approved_by_user_id, created_at=call.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _tool_call(row)
|
||||
|
||||
async def set_status(
|
||||
self,
|
||||
call_id: str,
|
||||
status: ToolCallStatus,
|
||||
*,
|
||||
result: dict | None = None,
|
||||
approved_by_user_id: str | None = None,
|
||||
) -> ToolCall | None:
|
||||
row = await self.session.get(ToolCallRow, call_id)
|
||||
if row is None:
|
||||
return None
|
||||
row.status = str(status)
|
||||
if result is not None:
|
||||
row.result_json = dict(result)
|
||||
if approved_by_user_id is not None:
|
||||
row.approved_by_user_id = approved_by_user_id
|
||||
await self.session.flush()
|
||||
return _tool_call(row)
|
||||
|
||||
|
||||
class ArtifactRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_for_org(
|
||||
self, organization_id: str, *, limit: int = 100, offset: int = 0
|
||||
) -> list[Artifact]:
|
||||
stmt = (
|
||||
select(ArtifactRow)
|
||||
.where(ArtifactRow.organization_id == organization_id)
|
||||
.order_by(ArtifactRow.created_at.desc())
|
||||
.limit(min(limit, 500))
|
||||
.offset(offset)
|
||||
)
|
||||
return [_artifact(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_for_run(self, run_id: str) -> list[Artifact]:
|
||||
stmt = (
|
||||
select(ArtifactRow)
|
||||
.where(ArtifactRow.run_id == run_id)
|
||||
.order_by(ArtifactRow.created_at)
|
||||
)
|
||||
return [_artifact(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def create(self, artifact: Artifact) -> Artifact:
|
||||
row = ArtifactRow(
|
||||
id=artifact.id, organization_id=artifact.organization_id,
|
||||
project_id=artifact.project_id, run_id=artifact.run_id,
|
||||
name=artifact.name, kind=artifact.kind, content_ref=artifact.content_ref,
|
||||
size_bytes=artifact.size_bytes, created_at=artifact.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _artifact(row)
|
||||
Reference in New Issue
Block a user