Coverage for src / osiris_cli / neural_tools.py: 0%

1#!/usr/bin/env python3 

2""" 

3Neural & Cognitive Tools for Osiris CLI 

4Handles scratchpad management, model state dumping, and behavioral cloning. 

5""" 

6 

7import os 

8import json 

9import time 

10from pathlib import Path 

11from typing import Dict, Any, Optional, List 

12from datetime import datetime 

13 

14class NeuralTools: 

15 """Specialized tools for neural state management and behavioral alignment.""" 

16 

17 def __init__(self, project_root: Optional[Path] = None): 

18 self.root = project_root or Path.cwd() 

19 self.scratchpad_path = self.root / "SCRATCHPAD.md" 

20 self.state_dir = self.root / ".osiris" / "neural_state" 

21 self.state_dir.mkdir(parents=True, exist_ok=True) 

22 self.cloning_dir = self.root / ".osiris" / "behavioral_cloning" 

23 self.cloning_dir.mkdir(parents=True, exist_ok=True) 

24 

25 def scratchpad(self, action: str, content: Optional[str] = None) -> str: 

26 """ 

27 Manage a persistent project-level scratchpad for planning and temporary notes. 

28  

29 Args: 

30 action: "read", "write", or "append" 

31 content: Content to write or append 

32 """ 

33 if action == "read": 

34 if not self.scratchpad_path.exists(): 

35 return "Scratchpad is empty. Use 'write' to initialize it." 

36 return self.scratchpad_path.read_text(encoding='utf-8') 

37 

38 elif action == "write": 

39 self.scratchpad_path.write_text(content or "", encoding='utf-8') 

40 return f"Scratchpad updated at {datetime.now().strftime('%H:%M:%S')}" 

41 

42 elif action == "append": 

43 with open(self.scratchpad_path, "a", encoding='utf-8') as f: 

44 f.write(f"\n\n--- Added {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n{content}") 

45 return "Content appended to scratchpad." 

46 

47 return "Invalid action. Use 'read', 'write', or 'append'." 

48 

49 def dump_state(self, state_obj: Dict[str, Any], label: str = "latest") -> str: 

50 """ 

51 Persist neural model state (optimized for Flax/JAX structures). 

52  

53 Args: 

54 state_obj: The state dictionary to persist 

55 label: Unique label for this state dump 

56 """ 

57 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") 

58 filename = f"state_{label}_{timestamp}.json" 

59 target_path = self.state_dir / filename 

60 

61 try: 

62 # Handle JAX/Flax specific types if present (simplified for now) 

63 # In a real JAX env, we would use orbax or msgpack 

64 serializable_state = self._make_serializable(state_obj) 

65 

66 with open(target_path, "w", encoding='utf-8') as f: 

67 json.dump(serializable_state, f, indent=2) 

68 

69 return f"Neural state dumped to {target_path}" 

70 except Exception as e: 

71 return f"Failed to dump state: {str(e)}" 

72 

73 def behavioral_cloning(self, mission_id: str, success: bool = True) -> str: 

74 """ 

75 Export the current interaction sequence for behavioral cloning and alignment. 

76  

77 Args: 

78 mission_id: Identifier for the current task/mission 

79 success: Whether the mission was accomplished successfully 

80 """ 

81 from .session import session 

82 

83 history = session.messages 

84 

85 payload = { 

86 "mission_id": mission_id, 

87 "timestamp": datetime.now().isoformat(), 

88 "status": "success" if success else "failed", 

89 "interactions": history, 

90 "environment": { 

91 "cwd": str(self.root), 

92 "os": os.name 

93 } 

94 } 

95 

96 filename = f"clone_{mission_id}_{datetime.now().strftime('%H%M%S')}.json" 

97 target_path = self.cloning_dir / filename 

98 

99 with open(target_path, "w", encoding='utf-8') as f: 

100 json.dump(payload, f, indent=2) 

101 

102 return f"Behavioral trace archived for alignment: {target_path}" 

103 

104 def pad_outputs(self, outputs: List[Any], max_length: int, pad_value: Any = 0) -> List[Any]: 

105 """ 

106 Standardized padding for neural sequence outputs (tensors or token lists). 

107  

108 Args: 

109 outputs: List of sequences to pad 

110 max_length: Target length for all sequences 

111 pad_value: Value to use for padding 

112 """ 

113 padded = [] 

114 for seq in outputs: 

115 if len(seq) > max_length: 

116 padded.append(seq[:max_length]) 

117 else: 

118 padded.append(list(seq) + [pad_value] * (max_length - len(seq))) 

119 return padded 

120 

121 def _make_serializable(self, obj: Any) -> Any: 

122 """Recursively convert objects to JSON-serializable formats.""" 

123 if isinstance(obj, (str, int, float, bool, type(None))): 

124 return obj 

125 elif isinstance(obj, dict): 

126 return {k: self._make_serializable(v) for k, v in obj.items()} 

127 elif isinstance(obj, (list, tuple)): 

128 return [self._make_serializable(item) for item in obj] 

129 elif hasattr(obj, "tolist"): # Numpy/JAX arrays 

130 return obj.tolist() 

131 else: 

132 return str(obj) 

133 

134# Singleton instance 

135neural_tools = NeuralTools()