import unittest from pathlib import Path import tempfile import torch from jspaceai import ( ActionEvent, ActionPolicy, CharTokenizer, ConsensusSnapshot, InMemoryVectorMemoryStore, JSpaceConfig, JSpaceLanguageModel, JSpaceModel, LanguageConfig, LanguageTrainingConfig, LanguageTrainingSession, MultimodalConfig, MultimodalJSpaceModel, OnlineLanguageLearner, WorkspaceEvent, WorkspaceRuntime, compose_action_params, ) from main_chat import generate_response class SmokeTests(unittest.TestCase): def setUp(self): torch.manual_seed(0) def test_core_forward_shapes(self): config = JSpaceConfig( input_dim=8, workspace_dim=16, expert_dim=8, num_experts=3, ode_steps=2, noise_std=0.0, ) model = JSpaceModel(config) xs = torch.randn(2, 5, 8) preds, info = model(xs, record_trajectory=True) self.assertEqual(tuple(preds.shape), (2, 5, 8)) self.assertEqual(tuple(info["alpha"].shape), (2, 5, 3)) self.assertEqual(tuple(info["w_trajectory"].shape), (2, 5, 2, 16)) def test_language_fast_generation_preserves_eval_and_rk4(self): tokenizer = CharTokenizer.from_text("学而时习之") config = LanguageConfig( vocab_size=tokenizer.vocab_size, embed_dim=8, input_dim=8, workspace_dim=16, expert_dim=8, num_experts=3, ode_steps=1, noise_std=0.0, ) model = JSpaceLanguageModel(config) model.eval() response = generate_response( model, tokenizer, "学", n_new=3, temperature=1.0, top_k=2, fast=True, ) self.assertEqual(len(response), 3) self.assertFalse(model.training) self.assertTrue(all(expert.use_rk4 for expert in model.experts)) def test_language_generate_preserves_training_mode(self): tokenizer = CharTokenizer.from_text("abcabc") config = LanguageConfig( vocab_size=tokenizer.vocab_size, embed_dim=8, input_dim=8, workspace_dim=16, expert_dim=8, num_experts=3, ode_steps=1, noise_std=0.0, ) model = JSpaceLanguageModel(config) model.eval() model.generate([1], n_new=1) self.assertFalse(model.training) model.train() model.generate([1], n_new=1) self.assertTrue(model.training) def test_multimodal_single_step_records_trajectory(self): config = MultimodalConfig( vocab_size=20, embed_dim=8, input_dim=8, workspace_dim=16, expert_dim=8, num_experts=4, ode_steps=2, noise_std=0.0, audio_frame_size=128, ) model = MultimodalJSpaceModel(config) token = torch.tensor([1]) outputs, info = model.forward_multimodal( "text", token, record_trajectory=True, ) self.assertEqual(tuple(outputs["w"].shape), (1, 16)) self.assertIn("alpha", info) self.assertEqual(len(info["w_trajectory"]), 2) def test_tokenizer_unknown_maps_to_zero(self): tokenizer = CharTokenizer.from_text("abc") self.assertEqual(tokenizer.encode("?"), [0]) self.assertEqual(tokenizer.decode([0]), "") def test_consensus_snapshot_extracts_primary_slot(self): w = torch.ones(1, 4) alpha = torch.tensor([[0.1, 0.7, 0.2]]) snapshot = ConsensusSnapshot.from_workspace( w, alpha, modality="text", labels=["vision", "language", "memory"], ) self.assertEqual(snapshot.primary_slot().label, "language") self.assertGreater(snapshot.confidence, 0.0) def test_online_language_learner_tracks_session_state(self): tokenizer = CharTokenizer.from_text("学而时习之学而时习之") config = LanguageConfig( vocab_size=tokenizer.vocab_size, embed_dim=8, input_dim=8, workspace_dim=16, expert_dim=8, num_experts=3, ode_steps=1, noise_std=0.0, ) model = JSpaceLanguageModel(config) learner = OnlineLanguageLearner( model, config, tokenizer, seq_len=8, replay_batch_size=1, consolidate_every=2, ) first = learner.learn_text("学而时习之") second = learner.learn_text("学而时习之") self.assertEqual(first["step"], 1) self.assertEqual(second["step"], 2) self.assertGreaterEqual(second["replay_loss"], 0.0) def test_language_training_session_saves_checkpoint(self): tokenizer = CharTokenizer.from_text("学而时习之学而时习之") config = LanguageConfig( vocab_size=tokenizer.vocab_size, embed_dim=8, input_dim=8, workspace_dim=16, expert_dim=8, num_experts=3, ode_steps=1, noise_std=0.0, ) model = JSpaceLanguageModel(config) train_config = LanguageTrainingConfig( seq_len=8, batch_size=2, lr=1e-2, replay_batch_size=1, validate_every=1, validate_batches=1, save_every=1, consolidate_every=0, ) session = LanguageTrainingSession( model, config, tokenizer, train_config, device="cpu", ) with tempfile.TemporaryDirectory() as tmpdir: ckpt_path = Path(tmpdir) / "model.pt" history = session.fit_text( "学而时习之学而时习之", max_steps=2, checkpoint_path=ckpt_path, ) ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) self.assertEqual(len(history), 2) self.assertIn("trainer", ckpt) self.assertEqual(ckpt["trainer"]["global_step"], 2) def test_compose_action_params_respects_discrete_mode(self): raw = torch.tensor([0.5, -0.25, 0.3, -0.4, 0.8]).numpy() move = compose_action_params(1, raw) left_click = compose_action_params(2, raw) self.assertEqual(move.tolist(), [0.5, -0.25, 0.0, 0.0, 0.0]) self.assertEqual(left_click.tolist(), [0.0, 0.0, 1.0, 0.0, 0.0]) def test_action_policy_decides_from_workspace(self): policy = ActionPolicy(workspace_dim=4, exploration=0.0) policy.value_model.action_weights[4, 0] = 2.0 w = torch.tensor([[1.0, 0.0, 0.0, 0.0]]) decision = policy.decide(w) self.assertEqual(decision.action_idx, 4) self.assertEqual(decision.action_name, "scroll") def test_workspace_event_and_memory_store_round_trip(self): event = WorkspaceEvent.from_tensor( torch.tensor([[1.0, 0.0, 0.0, 0.0]]), modality="text", step=7, context={"source": "test"}, ) action = ActionEvent.from_action_info({ "action_idx": 1, "action_name": "mouse_move", "action_params": [0.1, 0.0, 0.0, 0.0, 0.0], "executed": True, "risk": 0.0, }, step=7) memory = InMemoryVectorMemoryStore(capacity=4, workspace_dim=4) memory.put(event) results = memory.query(torch.tensor([[0.9, 0.0, 0.0, 0.0]]), top_k=1) self.assertEqual(event.to_dict()["step"], 7) self.assertEqual(action.to_dict()["action_name"], "mouse_move") self.assertEqual(results[0].event.modality, "text") self.assertGreater(results[0].similarity, 0.9) def test_workspace_runtime_runs_minimal_loop(self): class DummySenses: def start(self): return None def stop(self): return None class DummyAudio: def stop(self): return None class DummyMemory: def __init__(self): self.items = [] def store(self, w, context=None): self.items.append((w, context or {})) def size(self): return len(self.items) class DummyConfig: workspace_dim = 4 class DummyAgent: def __init__(self): self.config = DummyConfig() self.state = { "w": torch.zeros(1, 4), "m": [torch.zeros(1, 2)], } self.senses = DummySenses() self.audio_actuator = DummyAudio() self.hippocampus = DummyMemory() self.last_consensus = None self.learned = [] def perceive(self): return {} def think(self, sensory_data): del sensory_data self.state["w"] = self.state["w"] + 0.25 return self.state["w"], "idle" def decide_and_act(self, w, modality): del w, modality return { "action_idx": 1, "action_name": "mouse_move", "action_params": [0.0, 0.0, 0.0, 0.0, 0.0], "raw_action_params": [0.0, 0.0, 0.0, 0.0, 0.0], "executed": False, "action_strength": 0.0, "risk": 0.0, } def remember(self, w, context): self.hippocampus.store(w[0].numpy(), context) def learn(self, w, action_idx, reward=0.0): del w self.learned.append((action_idx, reward)) agent = DummyAgent() with tempfile.TemporaryDirectory() as tmpdir: runtime = WorkspaceRuntime(agent, save_dir=Path(tmpdir), device="cpu") info = runtime.step() self.assertEqual(info["step"], 1) self.assertEqual(info["modality"], "idle") self.assertEqual(info["memory_count"], 1) self.assertEqual(agent.learned[0][0], 1) self.assertEqual(info["workspace_event"]["step"], 1) self.assertEqual(info["action_event"]["action_name"], "mouse_move") if __name__ == "__main__": unittest.main()