feat(org): model grants, org memory, MCP gateway (v2 phases A1-3..5)
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
mjalmoudy94
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
c93b25741e
commit
dee2f72c46
@@ -24,7 +24,17 @@ from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import DateTime, ForeignKey, Index, String, Text
|
||||
from sqlalchemy import (
|
||||
Boolean,
|
||||
DateTime,
|
||||
Float,
|
||||
ForeignKey,
|
||||
Index,
|
||||
Integer,
|
||||
String,
|
||||
Text,
|
||||
UniqueConstraint,
|
||||
)
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from app.core.database import Base
|
||||
@@ -103,4 +113,139 @@ class OrgSetupDraftRow(Base):
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["OrgSetupDraftRow", "PlatformInviteRow"]
|
||||
class DefaultModelGrantRow(Base):
|
||||
"""The platform owner's default grant set, copied into every new organization.
|
||||
|
||||
Kept in a table of its own (never a column on ``ai_models``) so the model
|
||||
catalogue stays a pure, owner-managed list.
|
||||
"""
|
||||
|
||||
__tablename__ = "default_model_grants"
|
||||
__table_args__ = (UniqueConstraint("ai_model_id", name="uq_default_model_grant"),)
|
||||
|
||||
id: Mapped[str] = pk()
|
||||
ai_model_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("ai_models.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
daily_credit_cap: Mapped[float | None] = mapped_column(Float, nullable=True)
|
||||
is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
||||
created_by_user_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("users.id"), nullable=True
|
||||
)
|
||||
created_at: Mapped[datetime] = created()
|
||||
|
||||
|
||||
class OrgPolicyRow(Base):
|
||||
"""Per-organization policy the owner sets: limits and automation switches."""
|
||||
|
||||
__tablename__ = "org_policies"
|
||||
|
||||
organization_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("organizations.id", ondelete="CASCADE"), primary_key=True
|
||||
)
|
||||
max_employees: Mapped[int] = mapped_column(Integer, default=5, nullable=False)
|
||||
#: ``None`` means "no explicit budget" (an empty/unlimited wallet view).
|
||||
monthly_credit_budget: Mapped[float | None] = mapped_column(Float, nullable=True)
|
||||
automation_enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
||||
#: MCP access for this organization (Phase 5); the owner may switch it off.
|
||||
mcp_enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=utcnow, onupdate=utcnow, nullable=False
|
||||
)
|
||||
|
||||
|
||||
class OrgMemoryRow(Base):
|
||||
"""One thing the organization remembers about itself.
|
||||
|
||||
A memory is *small* by design (≤ 1,500 characters): longer material is a
|
||||
document and belongs in the Knowledge Base. ``content_hash`` is what makes
|
||||
"teach the same thing twice" a correction instead of a duplicate, and
|
||||
``key`` gives a rule a stable address so the engine can update it.
|
||||
"""
|
||||
|
||||
__tablename__ = "org_memories"
|
||||
__table_args__ = (
|
||||
Index("ix_org_memories_org_kind", "organization_id", "kind"),
|
||||
Index("ix_org_memories_org_key", "organization_id", "key"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = pk()
|
||||
organization_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False, index=True
|
||||
)
|
||||
kind: Mapped[str] = mapped_column(String(16), nullable=False)
|
||||
#: Optional stable handle, unique in spirit within one org.
|
||||
key: Mapped[str | None] = mapped_column(String(80), nullable=True)
|
||||
content: Mapped[str] = mapped_column(String(1500), nullable=False)
|
||||
#: A digest of the normalised content — the dedupe key the domain defines.
|
||||
content_hash: Mapped[str] = mapped_column(String(32), nullable=False, default="")
|
||||
source: Mapped[str] = mapped_column(String(16), default="user", nullable=False)
|
||||
evidence_task_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
pinned: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
|
||||
confidence: Mapped[float] = mapped_column(Float, default=1.0, nullable=False)
|
||||
expires_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
hits: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
last_used_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
created_by_user_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("users.id"), nullable=True
|
||||
)
|
||||
created_at: Mapped[datetime] = created()
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=utcnow, onupdate=utcnow, nullable=False
|
||||
)
|
||||
#: Soft delete: a memory an admin removed stays in the audit story.
|
||||
deleted_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
|
||||
|
||||
class McpTokenRow(Base):
|
||||
"""A scoped token an external AI agent uses to call the MCP gateway.
|
||||
|
||||
The token itself is shown exactly once, at creation; only its hash and a
|
||||
six-character prefix (for display) are stored. Every call is attributed to
|
||||
that prefix in the audit trail, so a leaked token can be identified and
|
||||
revoked long before anyone can reuse it.
|
||||
"""
|
||||
|
||||
__tablename__ = "mcp_tokens"
|
||||
|
||||
id: Mapped[str] = pk()
|
||||
organization_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False, index=True
|
||||
)
|
||||
name: Mapped[str] = mapped_column(String(120), nullable=False)
|
||||
token_hash: Mapped[str] = mapped_column(String(128), unique=True, nullable=False)
|
||||
#: First six characters, for display only. Never enough to authenticate.
|
||||
prefix: Mapped[str] = mapped_column(String(16), nullable=False)
|
||||
#: Comma-separated scopes; an empty string means "no scopes", never "all".
|
||||
scopes: Mapped[str] = mapped_column(String(255), default="", nullable=False)
|
||||
created_by_user_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("users.id"), nullable=True
|
||||
)
|
||||
expires_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
last_used_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
revoked_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
created_at: Mapped[datetime] = created()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DefaultModelGrantRow",
|
||||
"McpTokenRow",
|
||||
"OrgMemoryRow",
|
||||
"OrgPolicyRow",
|
||||
"OrgSetupDraftRow",
|
||||
"PlatformInviteRow",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,281 @@
|
||||
"""Repositories for org memory (Agent 1).
|
||||
|
||||
One table, three access shapes:
|
||||
|
||||
* reads are always scoped to one ``organization_id`` — there is no "list all
|
||||
memories" call, because a memory from another tenant must be unreachable by
|
||||
construction rather than by a filter someone remembers to add;
|
||||
* dedupe asks the database for the entry it suspects is the same one (by hash,
|
||||
by key), so "teach it twice" becomes an update;
|
||||
* eviction asks for the oldest unpinned lessons, which is the only thing this
|
||||
table is allowed to throw away.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import String, cast, func, or_, select, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.base import utcnow
|
||||
from app.domain.memory.models import Memory, MemoryKind, MemorySource
|
||||
from app.infrastructure.db.models_org import OrgMemoryRow
|
||||
|
||||
|
||||
def _to_domain(row: OrgMemoryRow) -> Memory:
|
||||
return Memory(
|
||||
id=row.id,
|
||||
organization_id=row.organization_id,
|
||||
kind=MemoryKind(row.kind),
|
||||
key=row.key,
|
||||
content=row.content,
|
||||
source=MemorySource(row.source),
|
||||
evidence_task_id=row.evidence_task_id,
|
||||
pinned=bool(row.pinned),
|
||||
confidence=float(row.confidence or 0.0),
|
||||
expires_at=row.expires_at,
|
||||
hits=int(row.hits or 0),
|
||||
last_used_at=row.last_used_at,
|
||||
created_by_user_id=row.created_by_user_id,
|
||||
created_at=row.created_at,
|
||||
updated_at=row.updated_at,
|
||||
deleted_at=row.deleted_at,
|
||||
)
|
||||
|
||||
|
||||
class MemoryRepository:
|
||||
"""Org memory persistence. Every method is tenant-scoped or keyed by id."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
# -- reads ---------------------------------------------------------------
|
||||
|
||||
async def list_for_org(
|
||||
self,
|
||||
organization_id: str,
|
||||
*,
|
||||
kind: MemoryKind | None = None,
|
||||
search: str | None = None,
|
||||
include_deleted: bool = False,
|
||||
limit: int = 200,
|
||||
offset: int = 0,
|
||||
) -> list[Memory]:
|
||||
stmt = select(OrgMemoryRow).where(
|
||||
OrgMemoryRow.organization_id == organization_id
|
||||
)
|
||||
if not include_deleted:
|
||||
stmt = stmt.where(OrgMemoryRow.deleted_at.is_(None))
|
||||
if kind is not None:
|
||||
stmt = stmt.where(OrgMemoryRow.kind == str(kind))
|
||||
if search:
|
||||
# Literal LIKE, not user-written pattern: a memory search must not
|
||||
# be able to inject SQL. Parameterised by SQLAlchemy either way.
|
||||
needle = f"%{search.strip()}%"
|
||||
stmt = stmt.where(
|
||||
or_(
|
||||
OrgMemoryRow.content.like(needle),
|
||||
cast(OrgMemoryRow.key, String).like(needle),
|
||||
)
|
||||
)
|
||||
stmt = stmt.order_by(
|
||||
OrgMemoryRow.pinned.desc(), OrgMemoryRow.updated_at.desc()
|
||||
).limit(limit).offset(offset)
|
||||
rows = (await self.session.execute(stmt)).scalars().all()
|
||||
return [_to_domain(row) for row in rows]
|
||||
|
||||
async def get(self, organization_id: str, memory_id: str) -> Memory | None:
|
||||
row = await self.session.scalar(
|
||||
select(OrgMemoryRow).where(
|
||||
OrgMemoryRow.organization_id == organization_id,
|
||||
OrgMemoryRow.id == memory_id,
|
||||
)
|
||||
)
|
||||
return _to_domain(row) if row else None
|
||||
|
||||
async def find_duplicate(
|
||||
self, organization_id: str, *, kind: MemoryKind, key: str | None, digest: str
|
||||
) -> Memory | None:
|
||||
"""The entry this write would duplicate: same key, or same content.
|
||||
|
||||
Matching on the key first is what turns "our port is 8000" into a
|
||||
correction of one memory instead of a second, similar memory.
|
||||
"""
|
||||
if key:
|
||||
row = await self.session.scalar(
|
||||
select(OrgMemoryRow).where(
|
||||
OrgMemoryRow.organization_id == organization_id,
|
||||
OrgMemoryRow.key == key,
|
||||
OrgMemoryRow.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
if row is not None:
|
||||
return _to_domain(row)
|
||||
row = await self.session.scalar(
|
||||
select(OrgMemoryRow).where(
|
||||
OrgMemoryRow.organization_id == organization_id,
|
||||
OrgMemoryRow.kind == str(kind),
|
||||
OrgMemoryRow.content_hash == digest,
|
||||
OrgMemoryRow.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
return _to_domain(row) if row else None
|
||||
|
||||
async def count_for_task(self, organization_id: str, evidence_task_id: str) -> int:
|
||||
return int(
|
||||
await self.session.scalar(
|
||||
select(func.count())
|
||||
.select_from(OrgMemoryRow)
|
||||
.where(
|
||||
OrgMemoryRow.organization_id == organization_id,
|
||||
OrgMemoryRow.evidence_task_id == evidence_task_id,
|
||||
OrgMemoryRow.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
or 0
|
||||
)
|
||||
|
||||
async def total_chars(self, organization_id: str) -> int:
|
||||
"""Characters the org currently spends on live memory."""
|
||||
rows = (
|
||||
await self.session.execute(
|
||||
select(func.length(OrgMemoryRow.content)).where(
|
||||
OrgMemoryRow.organization_id == organization_id,
|
||||
OrgMemoryRow.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
).scalars()
|
||||
return int(sum(int(value or 0) for value in rows))
|
||||
|
||||
# -- writes --------------------------------------------------------------
|
||||
|
||||
async def create(self, memory: Memory) -> Memory:
|
||||
row = OrgMemoryRow(
|
||||
id=memory.id,
|
||||
organization_id=memory.organization_id,
|
||||
kind=str(memory.kind),
|
||||
key=memory.key,
|
||||
content=memory.content,
|
||||
content_hash=memory.content_hash,
|
||||
source=str(memory.source),
|
||||
evidence_task_id=memory.evidence_task_id,
|
||||
pinned=memory.pinned,
|
||||
confidence=memory.confidence,
|
||||
expires_at=memory.expires_at,
|
||||
hits=memory.hits,
|
||||
last_used_at=memory.last_used_at,
|
||||
created_by_user_id=memory.created_by_user_id,
|
||||
created_at=memory.created_at,
|
||||
updated_at=memory.updated_at,
|
||||
deleted_at=memory.deleted_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return memory
|
||||
|
||||
async def save(self, memory: Memory) -> Memory:
|
||||
"""Persist a mutated in-memory record, hash and timestamp included."""
|
||||
row = await self.session.scalar(
|
||||
select(OrgMemoryRow).where(
|
||||
OrgMemoryRow.organization_id == memory.organization_id,
|
||||
OrgMemoryRow.id == memory.id,
|
||||
)
|
||||
)
|
||||
if row is None: # pragma: no cover - the service always loads first
|
||||
return await self.create(memory)
|
||||
row.kind = str(memory.kind)
|
||||
row.key = memory.key
|
||||
row.content = memory.content
|
||||
row.content_hash = memory.content_hash
|
||||
row.source = str(memory.source)
|
||||
row.evidence_task_id = memory.evidence_task_id
|
||||
row.pinned = memory.pinned
|
||||
row.confidence = memory.confidence
|
||||
row.expires_at = memory.expires_at
|
||||
row.hits = memory.hits
|
||||
row.last_used_at = memory.last_used_at
|
||||
row.deleted_at = memory.deleted_at
|
||||
row.updated_at = utcnow()
|
||||
await self.session.flush()
|
||||
return memory
|
||||
|
||||
async def mark_used(self, ids: list[str]) -> None:
|
||||
"""Count a ``build_context`` hit. One UPDATE, never one per memory."""
|
||||
if not ids:
|
||||
return
|
||||
await self.session.execute(
|
||||
update(OrgMemoryRow)
|
||||
.where(OrgMemoryRow.id.in_(ids))
|
||||
.values(hits=OrgMemoryRow.hits + 1, last_used_at=utcnow())
|
||||
)
|
||||
|
||||
async def soft_delete(self, organization_id: str, memory_id: str) -> bool:
|
||||
row = await self.session.scalar(
|
||||
select(OrgMemoryRow).where(
|
||||
OrgMemoryRow.organization_id == organization_id,
|
||||
OrgMemoryRow.id == memory_id,
|
||||
OrgMemoryRow.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
if row is None:
|
||||
return False
|
||||
row.deleted_at = utcnow()
|
||||
await self.session.flush()
|
||||
return True
|
||||
|
||||
# -- maintenance ---------------------------------------------------------
|
||||
|
||||
async def expire(self, organization_id: str | None = None) -> int:
|
||||
"""Soft-delete expired lessons. Pinned entries never expire.
|
||||
|
||||
Called from ``expire_memories`` — which the plan wants exposed as a
|
||||
plain function because the worker that owns the schedule is not mine.
|
||||
"""
|
||||
stmt = select(OrgMemoryRow).where(
|
||||
OrgMemoryRow.deleted_at.is_(None),
|
||||
OrgMemoryRow.pinned.is_(False),
|
||||
OrgMemoryRow.expires_at.is_not(None),
|
||||
OrgMemoryRow.expires_at < utcnow(),
|
||||
)
|
||||
if organization_id:
|
||||
stmt = stmt.where(OrgMemoryRow.organization_id == organization_id)
|
||||
rows = (await self.session.execute(stmt)).scalars().all()
|
||||
now = utcnow()
|
||||
for row in rows:
|
||||
row.deleted_at = now
|
||||
await self.session.flush()
|
||||
return len(rows)
|
||||
|
||||
async def evictable_lessons(
|
||||
self, organization_id: str, *, limit: int
|
||||
) -> list[Memory]:
|
||||
"""The oldest unpinned lessons — the only entries this table may drop."""
|
||||
rows = (
|
||||
await self.session.execute(
|
||||
select(OrgMemoryRow)
|
||||
.where(
|
||||
OrgMemoryRow.organization_id == organization_id,
|
||||
OrgMemoryRow.deleted_at.is_(None),
|
||||
OrgMemoryRow.pinned.is_(False),
|
||||
OrgMemoryRow.kind == str(MemoryKind.LESSON),
|
||||
)
|
||||
.order_by(OrgMemoryRow.created_at.asc())
|
||||
.limit(limit)
|
||||
)
|
||||
).scalars().all()
|
||||
return [_to_domain(row) for row in rows]
|
||||
|
||||
async def hard_delete(self, organization_id: str, memory_id: str) -> bool:
|
||||
row = await self.session.scalar(
|
||||
select(OrgMemoryRow).where(
|
||||
OrgMemoryRow.organization_id == organization_id,
|
||||
OrgMemoryRow.id == memory_id,
|
||||
)
|
||||
)
|
||||
if row is None:
|
||||
return False
|
||||
await self.session.delete(row)
|
||||
await self.session.flush()
|
||||
return True
|
||||
|
||||
|
||||
__all__ = ["MemoryRepository"]
|
||||
@@ -0,0 +1,237 @@
|
||||
"""Repositories for owner-granted models and per-organization policy (Agent 1).
|
||||
|
||||
``org_model_grants`` is foundation-owned but has no repository elsewhere; these
|
||||
live in a file no other workstream edits (the local shim the plan allows).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.base import new_id, utcnow
|
||||
from app.infrastructure.db.models import OrgModelGrantRow
|
||||
from app.infrastructure.db.models_org import DefaultModelGrantRow, OrgPolicyRow
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class OrgGrantRecord:
|
||||
organization_id: str
|
||||
ai_model_id: str
|
||||
daily_credit_cap: float | None = None
|
||||
is_active: bool = True
|
||||
created_by_user_id: str | None = None
|
||||
id: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class DefaultGrantRecord:
|
||||
ai_model_id: str
|
||||
daily_credit_cap: float | None = None
|
||||
is_active: bool = True
|
||||
created_by_user_id: str | None = None
|
||||
id: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class OrgPolicyRecord:
|
||||
organization_id: str
|
||||
max_employees: int = 5
|
||||
monthly_credit_budget: float | None = None
|
||||
automation_enabled: bool = True
|
||||
mcp_enabled: bool = True
|
||||
updated_at: datetime | None = None
|
||||
|
||||
|
||||
class OrgModelGrantRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
@staticmethod
|
||||
def _to_record(row: OrgModelGrantRow) -> OrgGrantRecord:
|
||||
return OrgGrantRecord(
|
||||
id=row.id,
|
||||
organization_id=row.organization_id,
|
||||
ai_model_id=row.ai_model_id,
|
||||
daily_credit_cap=row.daily_credit_cap,
|
||||
is_active=row.is_active,
|
||||
created_by_user_id=row.created_by_user_id,
|
||||
)
|
||||
|
||||
async def get(self, organization_id: str, ai_model_id: str) -> OrgGrantRecord | None:
|
||||
stmt = select(OrgModelGrantRow).where(
|
||||
OrgModelGrantRow.organization_id == organization_id,
|
||||
OrgModelGrantRow.ai_model_id == ai_model_id,
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return self._to_record(row) if row else None
|
||||
|
||||
async def list_for_org(
|
||||
self, organization_id: str, *, active_only: bool = False
|
||||
) -> list[OrgGrantRecord]:
|
||||
stmt = select(OrgModelGrantRow).where(
|
||||
OrgModelGrantRow.organization_id == organization_id
|
||||
)
|
||||
if active_only:
|
||||
stmt = stmt.where(OrgModelGrantRow.is_active)
|
||||
rows = (await self.session.execute(stmt)).scalars().all()
|
||||
return [self._to_record(r) for r in rows]
|
||||
|
||||
async def upsert(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
ai_model_id: str,
|
||||
daily_credit_cap: float | None,
|
||||
is_active: bool,
|
||||
created_by_user_id: str | None,
|
||||
) -> OrgGrantRecord:
|
||||
row = await self.session.scalar(
|
||||
select(OrgModelGrantRow).where(
|
||||
OrgModelGrantRow.organization_id == organization_id,
|
||||
OrgModelGrantRow.ai_model_id == ai_model_id,
|
||||
)
|
||||
)
|
||||
if row is None:
|
||||
row = OrgModelGrantRow(
|
||||
id=new_id(),
|
||||
organization_id=organization_id,
|
||||
ai_model_id=ai_model_id,
|
||||
daily_credit_cap=daily_credit_cap,
|
||||
is_active=is_active,
|
||||
created_by_user_id=created_by_user_id,
|
||||
)
|
||||
self.session.add(row)
|
||||
else:
|
||||
row.daily_credit_cap = daily_credit_cap
|
||||
row.is_active = is_active
|
||||
await self.session.flush()
|
||||
return self._to_record(row)
|
||||
|
||||
async def delete(self, organization_id: str, ai_model_id: str) -> bool:
|
||||
row = await self.session.scalar(
|
||||
select(OrgModelGrantRow).where(
|
||||
OrgModelGrantRow.organization_id == organization_id,
|
||||
OrgModelGrantRow.ai_model_id == ai_model_id,
|
||||
)
|
||||
)
|
||||
if row is None:
|
||||
return False
|
||||
await self.session.delete(row)
|
||||
await self.session.flush()
|
||||
return True
|
||||
|
||||
|
||||
class DefaultModelGrantRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
@staticmethod
|
||||
def _to_record(row: DefaultModelGrantRow) -> DefaultGrantRecord:
|
||||
return DefaultGrantRecord(
|
||||
id=row.id,
|
||||
ai_model_id=row.ai_model_id,
|
||||
daily_credit_cap=row.daily_credit_cap,
|
||||
is_active=row.is_active,
|
||||
created_by_user_id=row.created_by_user_id,
|
||||
)
|
||||
|
||||
async def list_all(self) -> list[DefaultGrantRecord]:
|
||||
rows = (await self.session.execute(select(DefaultModelGrantRow))).scalars().all()
|
||||
return [self._to_record(r) for r in rows]
|
||||
|
||||
async def list_active(self) -> list[DefaultGrantRecord]:
|
||||
stmt = select(DefaultModelGrantRow).where(DefaultModelGrantRow.is_active)
|
||||
return [self._to_record(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def upsert(
|
||||
self,
|
||||
*,
|
||||
ai_model_id: str,
|
||||
daily_credit_cap: float | None,
|
||||
is_active: bool,
|
||||
created_by_user_id: str | None,
|
||||
) -> DefaultGrantRecord:
|
||||
row = await self.session.scalar(
|
||||
select(DefaultModelGrantRow).where(
|
||||
DefaultModelGrantRow.ai_model_id == ai_model_id
|
||||
)
|
||||
)
|
||||
if row is None:
|
||||
row = DefaultModelGrantRow(
|
||||
id=new_id(),
|
||||
ai_model_id=ai_model_id,
|
||||
daily_credit_cap=daily_credit_cap,
|
||||
is_active=is_active,
|
||||
created_by_user_id=created_by_user_id,
|
||||
)
|
||||
self.session.add(row)
|
||||
else:
|
||||
row.daily_credit_cap = daily_credit_cap
|
||||
row.is_active = is_active
|
||||
await self.session.flush()
|
||||
return self._to_record(row)
|
||||
|
||||
async def delete(self, ai_model_id: str) -> bool:
|
||||
row = await self.session.scalar(
|
||||
select(DefaultModelGrantRow).where(
|
||||
DefaultModelGrantRow.ai_model_id == ai_model_id
|
||||
)
|
||||
)
|
||||
if row is None:
|
||||
return False
|
||||
await self.session.delete(row)
|
||||
await self.session.flush()
|
||||
return True
|
||||
|
||||
|
||||
class OrgPolicyRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
@staticmethod
|
||||
def _to_record(row: OrgPolicyRow) -> OrgPolicyRecord:
|
||||
return OrgPolicyRecord(
|
||||
organization_id=row.organization_id,
|
||||
max_employees=row.max_employees,
|
||||
monthly_credit_budget=row.monthly_credit_budget,
|
||||
automation_enabled=row.automation_enabled,
|
||||
mcp_enabled=row.mcp_enabled,
|
||||
updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
async def get(self, organization_id: str) -> OrgPolicyRecord | None:
|
||||
row = await self.session.get(OrgPolicyRow, organization_id)
|
||||
return self._to_record(row) if row else None
|
||||
|
||||
async def upsert(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
max_employees: int,
|
||||
monthly_credit_budget: float | None,
|
||||
automation_enabled: bool,
|
||||
mcp_enabled: bool | None = None,
|
||||
) -> OrgPolicyRecord:
|
||||
row = await self.session.get(OrgPolicyRow, organization_id)
|
||||
if row is None:
|
||||
row = OrgPolicyRow(
|
||||
organization_id=organization_id,
|
||||
max_employees=max_employees,
|
||||
monthly_credit_budget=monthly_credit_budget,
|
||||
automation_enabled=automation_enabled,
|
||||
mcp_enabled=True if mcp_enabled is None else mcp_enabled,
|
||||
)
|
||||
self.session.add(row)
|
||||
else:
|
||||
row.max_employees = max_employees
|
||||
row.monthly_credit_budget = monthly_credit_budget
|
||||
row.automation_enabled = automation_enabled
|
||||
if mcp_enabled is not None:
|
||||
row.mcp_enabled = mcp_enabled
|
||||
row.updated_at = utcnow()
|
||||
await self.session.flush()
|
||||
return self._to_record(row)
|
||||
Reference in New Issue
Block a user