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

1import ast 

2from pathlib import Path 

3from typing import List, Dict, Any 

4from rich.console import Console 

5from .config import settings 

6 

7console = Console() 

8 

9class ContextManager: 

10 

11 def __init__(self): 

12 self.loaded_files: Dict[str, str] = {} 

13 self.image_files: List[str] = [] # Vision-Ready tracking 

14 self.file_structures: Dict[str, Dict[str, Any]] = {} 

15 

16 def _estimate_tokens(self, text: str) -> int: 

17 """Rough estimation: 4 chars per token.""" 

18 return len(text) // 4 

19 

20 def get_total_tokens(self) -> int: 

21 return sum(self._estimate_tokens(content) for content in self.loaded_files.values()) 

22 

23 def add_file(self, file_path: str) -> bool: 

24 path = Path(file_path) 

25 if not path.exists(): 

26 return False 

27 

28 if not path.is_file(): 

29 return False 

30 

31 try: 

32 content = path.read_text(encoding='utf-8', errors='replace') 

33 

34 # Smart Context Truncation 

35 limit = settings.max_context_chars 

36 if len(content) > limit: 

37 lines = content.splitlines() 

38 total_lines = len(lines) 

39 # Keep head and tail 

40 head = lines[:50] 

41 tail = lines[-20:] 

42 

43 truncated_content = "\n".join(head) 

44 truncated_content += f"\n\n... [TRUNCATED: {total_lines - 70} lines hidden] ...\n" 

45 truncated_content += f"... [File too large ({len(content)} chars > {limit}). Use read_file(path, start_line, end_line) to view missing parts.] ...\n\n" 

46 truncated_content += "\n".join(tail) 

47 

48 self.loaded_files[str(path)] = truncated_content 

49 else: 

50 self.loaded_files[str(path)] = content 

51 # Analyze structure immediately 

52 self.file_structures[str(path)] = self.get_file_structure(str(path), content) 

53 return True 

54 except Exception as e: 

55 return False 

56 

57 def remove_file(self, file_path: str): 

58 if file_path in self.loaded_files: 

59 del self.loaded_files[file_path] 

60 if file_path in self.file_structures: 

61 del self.file_structures[file_path] 

62 

63 def clear(self): 

64 self.loaded_files = {} 

65 self.file_structures = {} 

66 

67 def get_file_structure(self, path: str, content: str) -> Dict[str, Any]: 

68 """Parse file structure (AST for Python, basic stats for others).""" 

69 struct = {"classes": {}, "functions": [], "type": "unknown"} 

70 

71 if path.endswith(".py"): 

72 struct["type"] = "python" 

73 try: 

74 tree = ast.parse(content) 

75 for node in ast.walk(tree): 

76 if isinstance(node, ast.ClassDef): 

77 methods = [n.name for n in node.body if isinstance(n, ast.FunctionDef)] 

78 struct["classes"][node.name] = methods 

79 elif isinstance(node, ast.FunctionDef): 

80 # Only top-level functions (approximate check: if parent is Module?) 

81 # ast.walk is flat. We can iterate body for top-level. 

82 pass 

83 

84 # Better top-level iteration 

85 for node in tree.body: 

86 if isinstance(node, ast.FunctionDef): 

87 struct["functions"].append(node.name) 

88 elif isinstance(node, ast.AsyncFunctionDef): 

89 struct["functions"].append(node.name) 

90 except Exception: 

91 pass 

92 return struct 

93 

94 def get_system_prompt_addition(self) -> str: 

95 if not self.loaded_files: 

96 return "" 

97 

98 text = "\n\n# CURRENT FILE CONTEXT\n" 

99 for name, content in self.loaded_files.items(): 

100 text += f"## File: {name}\n```\n{content}\n```\n" 

101 return text 

102 

103context = ContextManager()