agent-platform/services/tool-service/app/domain/repositories.py

from core_shared import JSONValue
from sqlalchemy import func, select
from sqlalchemy.orm import Session

from app.db.models import ToolBinding, ToolCredential, ToolDefinition, ToolVersion


class ToolDefinitionRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        plugin_id: str | None,
        code: str,
        name: str,
        tool_type: str,
        description: str | None) -> ToolDefinition:
        entity = ToolDefinition(
            plugin_id=plugin_id,
            code=code,
            name=name,
            tool_type=tool_type,
            description=description)
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_all(self) -> list[ToolDefinition]:
        stmt = (
            select(ToolDefinition)
            .order_by(ToolDefinition.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, tool_id: str) -> ToolDefinition | None:
        stmt = (
            select(ToolDefinition)
            .where(ToolDefinition.id == tool_id)
        )
        return self.db.scalar(stmt)

    def save(self, entity: ToolDefinition) -> ToolDefinition:
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def delete(self, entity: ToolDefinition) -> None:
        self.db.delete(entity)
        self.db.commit()


class ToolVersionRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        tool_id: str,
        input_schema_json: dict[str, JSONValue] | None,
        output_schema_json: dict[str, JSONValue] | None,
        invoke_config_json: dict[str, JSONValue] | None,
        timeout_ms: int | None,
        retry_policy_json: dict[str, JSONValue] | None) -> ToolVersion:
        entity = ToolVersion(
            tool_id=tool_id,
            version_no=self._next_version_no(tool_id),
            input_schema_json=input_schema_json,
            output_schema_json=output_schema_json,
            invoke_config_json=invoke_config_json,
            timeout_ms=timeout_ms,
            retry_policy_json=retry_policy_json)
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_by_tool(self, *, tool_id: str) -> list[ToolVersion]:
        stmt = (
            select(ToolVersion)
            .where(ToolVersion.tool_id == tool_id)
            .order_by(ToolVersion.version_no.desc())
        )
        return list(self.db.scalars(stmt))

    def list_all(self) -> list[ToolVersion]:
        stmt = (
            select(ToolVersion)
            .order_by(ToolVersion.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, tool_version_id: str) -> ToolVersion | None:
        stmt = (
            select(ToolVersion)
            .where(ToolVersion.id == tool_version_id)
        )
        return self.db.scalar(stmt)

    def save(self, entity: ToolVersion) -> ToolVersion:
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def delete(self, entity: ToolVersion) -> None:
        self.db.delete(entity)
        self.db.commit()

    def _next_version_no(self, tool_id: str) -> int:
        stmt = select(func.max(ToolVersion.version_no)).where(ToolVersion.tool_id == tool_id)
        current_max = self.db.scalar(stmt)
        return (current_max or 0) + 1


class ToolBindingRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        app_id: str,
        tool_version_id: str,
        credential_id: str | None,
        binding_scope: str,
        enabled: bool,
        config_json: dict[str, JSONValue] | None) -> ToolBinding:
        entity = ToolBinding(
            app_id=app_id,
            tool_version_id=tool_version_id,
            credential_id=credential_id,
            binding_scope=binding_scope,
            enabled=enabled,
            config_json=config_json)
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_by_scope(self, *, app_id: str | None = None) -> list[ToolBinding]:
        stmt = select(ToolBinding)
        if app_id is not None:
            stmt = stmt.where(ToolBinding.app_id == app_id)
        stmt = stmt.order_by(ToolBinding.created_time.desc())
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, binding_id: str) -> ToolBinding | None:
        stmt = (
            select(ToolBinding)
            .where(ToolBinding.id == binding_id)
        )
        return self.db.scalar(stmt)

    def save(self, entity: ToolBinding) -> ToolBinding:
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def delete(self, entity: ToolBinding) -> None:
        self.db.delete(entity)
        self.db.commit()


class ToolCredentialRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        name: str,
        credential_type: str,
        encrypted_payload_text: str,
        secret_fingerprint: str,
        encryption_algorithm: str,
        metadata_json: dict[str, JSONValue]) -> ToolCredential:
        entity = ToolCredential(
            name=name,
            credential_type=credential_type,
            encrypted_payload_text=encrypted_payload_text,
            secret_fingerprint=secret_fingerprint,
            encryption_algorithm=encryption_algorithm,
            metadata_json=metadata_json)
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_all(self) -> list[ToolCredential]:
        stmt = (
            select(ToolCredential)
            .order_by(ToolCredential.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, credential_id: str) -> ToolCredential | None:
        stmt = (
            select(ToolCredential)
            .where(ToolCredential.id == credential_id)
        )
        return self.db.scalar(stmt)

    def save(self, entity: ToolCredential) -> ToolCredential:
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def delete(self, entity: ToolCredential) -> None:
        self.db.delete(entity)
        self.db.commit()