tests/test_v6_components.py

"""Tests for MORPH-AI v6 components (no model required)."""

import sys
import torch
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))

from architecture import (
    MorphConfig, MorphModel, MixtureOfDepths, MemoryEfficientAttention,
    DynamicMoEBlock, MultimodalFusion, ToolUseModule, DocumentModule,
    VideoModule, CodeSandbox, QuantizedMemoryModule, build_code_features,
)
from audio import AudioModule, AudioFacts
from video import VideoModule as VM
from document import DocumentModule as DM
from tools import ToolRegistry, ToolCall


def test_mod_forward():
    cfg = MorphConfig(mod_hidden=64, mod_keep_prob=0.8)
    mod = MixtureOfDepths(128, cfg.mod_hidden, cfg.mod_keep_prob)
    x = torch.randn(2, 16, 128)
    out, probs = mod(x)
    assert out.shape == x.shape
    assert probs.shape == (2, 16, 1)
    assert (probs >= 0).all() and (probs <= 1).all()


def test_memory_efficient_attention():
    attn = MemoryEfficientAttention(64, 4)
    x = torch.randn(2, 8, 64)
    out = attn(x)
    assert out.shape == x.shape


def test_dynamic_moe_block():
    moe = DynamicMoEBlock(128, 4, 256, 2, prune_threshold=0.01)
    x = torch.randn(2, 8, 128)
    out, aux = moe(x)
    assert out.shape == x.shape
    assert aux.item() >= 0


def test_multimodal_fusion():
    cfg = MorphConfig()
    fusion = MultimodalFusion(cfg, 128)
    text = torch.randn(2, 8, 128)
    vision = torch.randn(2, 8, 768)
    out = fusion(text, vision)
    assert out.shape == text.shape


def test_tool_use_module():
    cfg = MorphConfig()
    tool = ToolUseModule(cfg, 128)
    x = torch.randn(2, 8, 128)
    tool_emb, args = tool(x)
    assert tool_emb.shape[0] == 2
    assert args.shape == (2, 128)


def test_document_module():
    cfg = MorphConfig()
    doc = DocumentModule(cfg, 128)
    x = torch.randn(2, 8, 128)
    out = doc(x)
    assert out.shape == x.shape


def test_video_module():
    cfg = MorphConfig()
    video = VideoModule(cfg, 128)
    x = torch.randn(2, 8, 128)
    frames = torch.randn(4, 128)
    out = video(x, frames.unsqueeze(0))
    assert out.shape == x.shape


def test_code_sandbox():
    sandbox = CodeSandbox(timeout=2.0)
    result = sandbox.execute("1 + 1")
    assert result["success"] is True
    assert "2" in result["output"]

    result = sandbox.execute("import os")
    assert result["success"] is False
    assert "not allowed" in result["error"]


def test_quantized_memory():
    mem = QuantizedMemoryModule(256, 128, 256, quantize=True, quant_bits=8)
    x = torch.randn(2, 8, 256)
    mem.write(x)
    out = mem.read(x)
    assert out.shape == x.shape


def test_tool_registry():
    registry = ToolRegistry()
    assert "calculator" in registry._tools
    assert "search" in registry._tools
    assert "code_exec" in registry._tools
    schema = registry.get_tool_schema()
    assert "calculator" in schema


def test_audio_module_init():
    cfg = MorphConfig()
    audio = AudioModule(cfg, 768)
    assert audio is not None


def test_build_code_features_v7():
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct", trust_remote_code=True)
    ids = torch.tensor([[1, 2, 3, 4, 5]])
    feats = build_code_features(tokenizer, ids)
    assert feats.shape[-1] == 4


def test_video_analyze():
    from video import VideoModule as VM
    vm = VM(MorphConfig(), 128)
    facts = vm.analyze("nonexistent.mp4")
    assert isinstance(facts.to_text(), str)


def test_document_extract():
    cfg = MorphConfig()
    dm = DocumentModule(cfg, 128)
    result = dm.extract_text("nonexistent.pdf")
    assert isinstance(result, str)


if __name__ == "__main__":
    test_mod_forward()
    test_memory_efficient_attention()
    test_dynamic_moe_block()
    test_multimodal_fusion()
    test_tool_use_module()
    test_document_module()
    test_video_module()
    test_code_sandbox()
    test_quantized_memory()
    test_tool_registry()
    test_audio_module_init()
    test_build_code_features_v7()
    test_video_analyze()
    test_document_extract()
    print("All v6 component tests passed!")