97 lines
3.7 KiB
Python
97 lines
3.7 KiB
Python
"""SQLAlchemy-backed model-capability store."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.model_capabilities.base import ModelCapabilityStore
|
|
from deerflow.persistence.model_capabilities.model import ModelCapabilityRow
|
|
|
|
|
|
def _to_dict(row: ModelCapabilityRow) -> dict[str, Any]:
|
|
return {
|
|
"model_name": row.model_name,
|
|
"supports_reasoning_effort": bool(row.supports_reasoning_effort),
|
|
"default_reasoning_effort": row.default_reasoning_effort or "xhigh",
|
|
"updated_by": row.updated_by,
|
|
"updated_at": row.updated_at.isoformat() if isinstance(row.updated_at, datetime) else row.updated_at,
|
|
}
|
|
|
|
|
|
class ModelCapabilityRepository(ModelCapabilityStore):
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
async def list_capabilities(self) -> list[dict[str, Any]]:
|
|
async with self._sf() as session:
|
|
result = await session.execute(select(ModelCapabilityRow).order_by(ModelCapabilityRow.model_name.asc()))
|
|
return [_to_dict(row) for row in result.scalars()]
|
|
|
|
async def get_capability(self, model_name: str) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
row = await session.get(ModelCapabilityRow, model_name)
|
|
return _to_dict(row) if row is not None else None
|
|
|
|
async def set_reasoning_effort_support(
|
|
self,
|
|
model_name: str,
|
|
*,
|
|
supports_reasoning_effort: bool,
|
|
default_reasoning_effort: str | None = None,
|
|
updated_by: str | None = None,
|
|
) -> dict[str, Any]:
|
|
async with self._sf() as session:
|
|
row = await session.get(ModelCapabilityRow, model_name)
|
|
if row is None:
|
|
row = ModelCapabilityRow(model_name=model_name)
|
|
session.add(row)
|
|
row.supports_reasoning_effort = bool(supports_reasoning_effort)
|
|
if default_reasoning_effort is not None:
|
|
row.default_reasoning_effort = default_reasoning_effort
|
|
row.updated_by = updated_by
|
|
row.updated_at = datetime.now(UTC)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return _to_dict(row)
|
|
|
|
async def seed_reasoning_effort_support(
|
|
self,
|
|
model_name: str,
|
|
*,
|
|
updated_by: str | None = None,
|
|
) -> bool:
|
|
"""Insert a legacy ``true`` without ever overwriting an existing row.
|
|
|
|
The primary-key conflict handles concurrent Gateway startup safely: if
|
|
another worker or an administrator creates the row after our read, the
|
|
losing seed transaction rolls back and preserves that authoritative
|
|
value, including ``false``.
|
|
"""
|
|
|
|
async with self._sf() as session:
|
|
if await session.get(ModelCapabilityRow, model_name) is not None:
|
|
return False
|
|
session.add(
|
|
ModelCapabilityRow(
|
|
model_name=model_name,
|
|
supports_reasoning_effort=True,
|
|
updated_by=updated_by,
|
|
)
|
|
)
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError:
|
|
await session.rollback()
|
|
return False
|
|
return True
|
|
|
|
async def delete_capability(self, model_name: str) -> None:
|
|
async with self._sf() as session:
|
|
await session.execute(delete(ModelCapabilityRow).where(ModelCapabilityRow.model_name == model_name))
|
|
await session.commit()
|