shared_model/systems.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch

from .shared_model import SharedBackboneAdapter, get_shared_adapter, register_shared_model


@dataclass
class SharedSystems:
    adapter: SharedBackboneAdapter
    consciousness: Optional[Any] = None
    agi_memory: Optional[Any] = None
    errors: Optional[Dict[str, str]] = None


def initialize_shared_systems(
    model,
    tokenizer=None,
    *,
    enable_consciousness: bool = True,
    enable_agi_memory: bool = True,
) -> SharedSystems:
    adapter = register_shared_model(model, tokenizer)
    errors: Dict[str, str] = {}
    consciousness = None
    agi_memory = None

    if enable_consciousness:
        try:
            from ConsciousnessSystem.config import ConsciousnessSystemConfig
            from ConsciousnessSystem.consciousness_system import ConsciousnessSystem

            cfg = ConsciousnessSystemConfig(
                device=str(adapter.device),
                hidden_dim=adapter.hidden_size,
                use_fp16=adapter.dtype in (torch.float16, torch.bfloat16),
            )
            consciousness = ConsciousnessSystem(cfg, qwen_model=model)
        except Exception as exc:
            errors["consciousness"] = str(exc)

    if enable_agi_memory:
        try:
            from agi_memory.core import AGIMemoryConfig, AGIMemorySystem

            cfg = AGIMemoryConfig(
                device=adapter.device,
                use_fp16=adapter.dtype in (torch.float16, torch.bfloat16),
                disable_vae_compression=True,
                enable_rag=False,
                enable_dynamic_action=False,
                enable_neuro_fusion=False,
                enable_memory_optimization=False,
            )
            agi_memory = AGIMemorySystem(cfg)
        except Exception as exc:
            errors["agi_memory"] = str(exc)

    return SharedSystems(adapter=adapter, consciousness=consciousness, agi_memory=agi_memory, errors=errors)