JARVIS/backend/app/agents/registry/indexes.py

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType

from app.agents.registry.loader import RegistryBundle
from app.agents.registry.models import (
    AgentManifest,
    CapabilityManifest,
    SpecialistTemplateManifest,
    SubCommanderManifest,
)


@dataclass(frozen=True)
class RegistryIndexes:
    agent_by_id: Mapping[str, AgentManifest]
    agent_by_role_value: Mapping[str, AgentManifest]
    sub_commander_by_id: Mapping[str, SubCommanderManifest]
    capability_by_id: Mapping[str, CapabilityManifest]
    specialist_template_by_id: Mapping[str, SpecialistTemplateManifest]
    agent_prompt_key_by_id: Mapping[str, str]
    sub_commander_prompt_key_by_id: Mapping[str, str]
    skill_context_key_by_agent_id: Mapping[str, str]
    capability_id_by_tool_name: Mapping[str, str]
    capability_ids_by_sub_commander_id: Mapping[str, tuple[str, ...]]
    spawnable_role_values_by_agent_id: Mapping[str, tuple[str, ...]]


def summarize_registry_indexes(indexes: RegistryIndexes) -> dict[str, int]:
    return {
        "agent_count": len(indexes.agent_by_id),
        "sub_commander_count": len(indexes.sub_commander_by_id),
        "capability_count": len(indexes.capability_by_id),
        "specialist_template_count": len(indexes.specialist_template_by_id),
    }


def build_registry_indexes(bundle: RegistryBundle) -> RegistryIndexes:
    agent_by_id = {agent.agent_id: agent for agent in bundle.agents}
    sub_commander_by_id = {
        sub_commander.sub_commander_id: sub_commander
        for sub_commander in bundle.sub_commanders
    }
    capability_by_id = {
        capability.capability_id: capability for capability in bundle.capabilities
    }
    specialist_template_by_id = {
        template.template_id: template for template in bundle.specialist_templates
    }

    return RegistryIndexes(
        agent_by_id=MappingProxyType(agent_by_id),
        agent_by_role_value=MappingProxyType({
            agent.role_value: agent for agent in bundle.agents
        }),
        sub_commander_by_id=MappingProxyType(sub_commander_by_id),
        capability_by_id=MappingProxyType(capability_by_id),
        specialist_template_by_id=MappingProxyType(specialist_template_by_id),
        agent_prompt_key_by_id=MappingProxyType({
            agent.agent_id: agent.system_prompt_key for agent in bundle.agents
        }),
        sub_commander_prompt_key_by_id=MappingProxyType({
            sub_commander.sub_commander_id: sub_commander.sub_commander_id
            for sub_commander in bundle.sub_commanders
        }),
        skill_context_key_by_agent_id=MappingProxyType({
            agent.agent_id: agent.skill_context_key
            for agent in bundle.agents
            if agent.skill_context_key is not None
        }),
        capability_id_by_tool_name=MappingProxyType({
            capability.tool_name: capability.capability_id
            for capability in bundle.capabilities
        }),
        capability_ids_by_sub_commander_id=MappingProxyType({
            sub_commander.sub_commander_id: tuple(sub_commander.capability_ids)
            for sub_commander in bundle.sub_commanders
        }),
        spawnable_role_values_by_agent_id=MappingProxyType({
            agent.agent_id: tuple(agent.allowed_spawn_role_values)
            for agent in bundle.agents
            if agent.can_spawn_children and agent.allowed_spawn_role_values
        }),
    )