agentic-mobile-control/skills_learning/models.py

from __future__ import annotations

from dataclasses import dataclass, field, replace
from datetime import datetime
from typing import Any, Literal
from uuid import uuid4

from core.models import utc_now

LOCAL_SYNTHESIS_SOURCE = "local-synthesis"
SkillKind = Literal["knowledge", "flow_template"]


@dataclass(frozen=True)
class SkillMetadata:
    id: str = field(default_factory=lambda: uuid4().hex)
    name: str = ""
    description: str = ""
    kind: SkillKind = "flow_template"
    tags: list[str] = field(default_factory=list)
    source: str = LOCAL_SYNTHESIS_SOURCE
    version: int = 1
    parent_version_id: str | None = None
    originating_goal: str | None = None
    created_at: datetime = field(default_factory=utc_now)
    updated_at: datetime = field(default_factory=utc_now)

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "kind": self.kind,
            "tags": list(self.tags),
            "source": self.source,
            "version": self.version,
            "parent_version_id": self.parent_version_id,
            "originating_goal": self.originating_goal,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "SkillMetadata":
        return cls(
            id=str(data.get("id") or uuid4().hex),
            name=str(data.get("name") or ""),
            description=str(data.get("description") or ""),
            kind=data.get("kind") or "flow_template",
            tags=[str(tag) for tag in data.get("tags", [])],
            source=str(data.get("source") or LOCAL_SYNTHESIS_SOURCE),
            version=int(data.get("version") or 1),
            parent_version_id=data.get("parent_version_id"),
            originating_goal=data.get("originating_goal"),
            created_at=_parse_datetime(data.get("created_at")),
            updated_at=_parse_datetime(data.get("updated_at")),
        )


@dataclass(frozen=True)
class FlowStep:
    tool_name: str
    args: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "tool_name": self.tool_name,
            "args": dict(self.args),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "FlowStep":
        return cls(
            tool_name=str(data.get("tool_name") or data.get("action") or ""),
            args=dict(data.get("args") or {}),
        )


@dataclass(frozen=True)
class Skill:
    metadata: SkillMetadata

    @property
    def id(self) -> str:
        return self.metadata.id

    @property
    def name(self) -> str:
        return self.metadata.name

    @property
    def description(self) -> str:
        return self.metadata.description

    @property
    def source(self) -> str:
        return self.metadata.source

    @property
    def version(self) -> int:
        return self.metadata.version

    @property
    def parent_version_id(self) -> str | None:
        return self.metadata.parent_version_id

    @property
    def originating_goal(self) -> str | None:
        return self.metadata.originating_goal


@dataclass(frozen=True)
class FlowTemplateSkill(Skill):
    steps: list[FlowStep] = field(default_factory=list)
    parameters: dict[str, dict[str, Any]] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            **self.metadata.to_dict(),
            "steps": [step.to_dict() for step in self.steps],
            "parameters": {
                name: dict(schema)
                for name, schema in self.parameters.items()
            },
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "FlowTemplateSkill":
        return cls(
            metadata=SkillMetadata.from_dict(data),
            steps=[
                FlowStep.from_dict(step)
                for step in data.get("steps", [])
            ],
            parameters={
                str(name): dict(schema)
                for name, schema in (data.get("parameters") or {}).items()
            },
        )

    def with_metadata(self, **changes: Any) -> "FlowTemplateSkill":
        return replace(self, metadata=replace(self.metadata, **changes))

    def with_updates(
        self,
        *,
        steps: list[FlowStep] | None = None,
        parameters: dict[str, dict[str, Any]] | None = None,
        **metadata_changes: Any,
    ) -> "FlowTemplateSkill":
        metadata = replace(
            self.metadata,
            updated_at=utc_now(),
            **metadata_changes,
        )
        return replace(
            self,
            metadata=metadata,
            steps=list(steps) if steps is not None else list(self.steps),
            parameters={
                name: dict(schema)
                for name, schema in (
                    parameters if parameters is not None else self.parameters
                ).items()
            },
        )


@dataclass(frozen=True)
class KnowledgeSkill(Skill):
    content: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            **self.metadata.to_dict(),
            "content": self.content,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "KnowledgeSkill":
        return cls(
            metadata=SkillMetadata.from_dict(data),
            content=str(data.get("content") or ""),
        )

    def with_metadata(self, **changes: Any) -> "KnowledgeSkill":
        return replace(self, metadata=replace(self.metadata, **changes))


def skill_embedding_text(skill: FlowTemplateSkill) -> str:
    goal = skill.originating_goal or ""
    return f"{skill.name}: {skill.description}\nOriginal goal: {goal}"


def clone_skill(skill: FlowTemplateSkill) -> FlowTemplateSkill:
    return FlowTemplateSkill.from_dict(skill.to_dict())


def _parse_datetime(value: Any) -> datetime:
    if isinstance(value, datetime):
        return value
    if isinstance(value, str):
        try:
            return datetime.fromisoformat(value)
        except ValueError:
            pass
    return utc_now()