Coverage for app \ engines \ base.py: 30%
99 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-11-30 09:36 +0100
1"""
2Base class for all SLM engines
4Engines are neural network-based components that use
5Small Language Models for code understanding and generation.
6"""
7from abc import ABC, abstractmethod
8from typing import Dict, Any, Optional
9import logging
11from app.models.schemas import TaskType, Language
13logger = logging.getLogger(__name__)
16class BaseEngine(ABC):
17 """Base class for all SLM engines"""
19 def __init__(self, name: str, model_path: Optional[str] = None):
20 self.name = name
21 self.model_path = model_path
22 self.model = None
23 self.initialized = False
24 logger.info(f"Creating engine: {name}")
26 @abstractmethod
27 async def initialize(self):
28 """
29 Initialize the engine and load model
31 Should set self.initialized = True when done
32 """
33 pass
35 @abstractmethod
36 async def process(
37 self,
38 task: TaskType,
39 code: str,
40 language: Language,
41 context: Optional[str] = None,
42 trace: Optional[str] = None,
43 **kwargs
44 ) -> Dict[str, Any]:
45 """
46 Process a task using the SLM
48 Args:
49 task: Type of task
50 code: Source code
51 language: Programming language
52 context: Additional context
53 trace: Error trace (if applicable)
54 **kwargs: Additional parameters
56 Returns:
57 Dict with:
58 - success: bool
59 - result: str (generated/fixed code or explanation)
60 - explanation: Optional[str]
61 - suggestions: Optional[List[str]]
62 """
63 pass
65 @abstractmethod
66 async def shutdown(self):
67 """Cleanup and free resources"""
68 pass
70 def build_prompt(self, task: TaskType, code: str, context: Optional[str] = None) -> str:
71 """
72 Build prompt for SLM based on task type
74 Args:
75 task: Type of task
76 code: Source code
77 context: Additional context
79 Returns:
80 Formatted prompt string
81 """
82 if task == TaskType.FIX:
83 return self._build_fix_prompt(code, context)
84 elif task == TaskType.EXPLAIN:
85 return self._build_explain_prompt(code, context)
86 elif task == TaskType.REFACTOR:
87 return self._build_refactor_prompt(code, context)
88 elif task == TaskType.TEST:
89 return self._build_test_prompt(code, context)
90 elif task == TaskType.BOILERPLATE:
91 return self._build_boilerplate_prompt(context)
92 elif task == TaskType.TRANSLATE:
93 return self._build_translate_prompt(code, context)
94 else:
95 return f"Code:\n{code}\n\nTask: {task}"
97 def get_stop_tokens(self, task: TaskType) -> list:
98 """
99 Get stop tokens to prevent over-generation
101 Args:
102 task: Type of task
104 Returns:
105 List of stop tokens
106 """
107 # Common stop tokens
108 common_stops = ["```", "\n\n\n", "# Example", "# Test"]
110 # Task-specific stops
111 task_stops = {
112 TaskType.FIX: ["# Fixed code:", "# Original:"],
113 TaskType.BOILERPLATE: ["# Usage:", "# Example usage:"],
114 TaskType.TEST: ["# Run tests:"],
115 TaskType.EXPLAIN: ["# Code:", "# Summary:"],
116 }
118 return common_stops + task_stops.get(task, [])
120 def _build_fix_prompt(self, code: str, context: Optional[str]) -> str:
121 """Build prompt for code fixing"""
122 prompt = """Fix the following Python code. Return ONLY the corrected code without explanation.
124Rules:
125- Fix syntax errors
126- Fix logic errors
127- Maintain original functionality
128- Keep the same structure
129- Do not add comments unless necessary
131"""
132 if context:
133 prompt += f"Additional context: {context}\n\n"
135 prompt += f"Code to fix:\n```python\n{code}\n```\n\n"
136 prompt += "Fixed code:\n```python\n"
137 return prompt
139 def _build_explain_prompt(self, code: str, context: Optional[str]) -> str:
140 """Build prompt for explaining code"""
141 prompt = """Explain the following Python code in detail.
143Rules:
144- Provide a high-level summary.
145- Break down the code into logical sections and explain each.
146- Highlight key concepts and potential improvements.
147- Use clear and concise language.
149"""
150 if context:
151 prompt += f"Focus on: {context}\n\n"
153 prompt += f"Code to explain:\n```python\n{code}\n```\n\n"
154 prompt += "Explanation:\n"
155 return prompt
157 def _build_refactor_prompt(self, code: str, context: Optional[str]) -> str:
158 """Build prompt for refactoring code"""
159 prompt = """Refactor the following Python code to improve readability, maintainability, and performance. Return ONLY the refactored code without explanation.
161Rules:
162- Maintain original functionality.
163- Apply Python best practices and idioms.
164- Improve variable names, function structure, and overall design.
165- Do not add comments unless necessary.
167"""
168 if context:
169 prompt += f"Refactoring requirements: {context}\n\n"
171 prompt += f"Original code:\n```python\n{code}\n```\n\n"
172 prompt += "Refactored code:\n```python\n"
173 return prompt
175 def _build_test_prompt(self, code: str, context: Optional[str]) -> str:
176 """Build prompt for generating unit tests"""
177 prompt = """Generate comprehensive unit tests for the following Python code using `pytest`. Return ONLY the test code without explanation.
179Rules:
180- Cover normal cases, edge cases, and error cases.
181- Use descriptive test function names.
182- Include assertions for expected behavior.
183- Use `pytest.fixture` for setup if needed.
185"""
186 if context:
187 prompt += f"Test requirements: {context}\n\n"
189 prompt += f"Code to test:\n```python\n{code}\n```\n\n"
190 prompt += "Test code:\n```python\n"
191 return prompt
193 def _build_translate_prompt(self, code: str, context: Optional[str]) -> str:
194 """Build prompt for translating code"""
195 prompt = """Translate the following code. Return ONLY the translated code without explanation.
197"""
198 if context:
199 prompt += f"{context}\n\n"
201 prompt += f"Code to translate:\n```\n{code}\n```\n\n"
202 prompt += "Translated code:\n```\n"
203 return prompt
205 def _build_boilerplate_prompt(self, context: Optional[str]) -> str:
206 """Build prompt for boilerplate generation"""
207 prompt = """Generate clean, well-structured Python code based on the description below.
209Requirements:
210- Follow PEP 8 style guide
211- Include docstrings
212- Handle edge cases
213- Use type hints where appropriate
214- Keep it simple and readable
216"""
217 if context:
218 prompt += f"Description: {context}\n\n"
219 else:
220 prompt += "Description: Create a basic implementation\n\n"
222 prompt += "Code:\n```python\n"
223 return prompt
225 def _default_prompt(
226 self,
227 code: str,
228 language: Language,
229 context: Optional[str],
230 trace: Optional[str]
231 ) -> str:
232 """Default prompt"""
233 return f"Process the following {language} code:\n\n```{language}\n{code}\n```"
235 def _format_result(
236 self,
237 success: bool,
238 result: Optional[str] = None,
239 explanation: Optional[str] = None,
240 suggestions: Optional[list] = None
241 ) -> Dict[str, Any]:
242 """Helper to format results consistently"""
243 return {
244 "success": success,
245 "result": result,
246 "explanation": explanation,
247 "suggestions": suggestions or []
248 }
250 def _extract_code_from_response(self, response: str) -> str:
251 """Extract code from model response (handles markdown code blocks)"""
252 import re
254 # Look for ```language\ncode\n``` pattern
255 pattern = r'```(?:\w+)?\n(.*?)\n?```'
256 matches = re.findall(pattern, response, re.DOTALL)
258 if matches:
259 # Return first code block
260 return matches[0].strip()
262 # If no code blocks, check if response looks like code already
263 # (happens when prompt ends with ``` and model just generates code)
264 lines = response.strip().split('\n')
266 # Skip empty lines at start/end
267 while lines and not lines[0].strip():
268 lines.pop(0)
269 while lines and not lines[-1].strip():
270 lines.pop()
272 # If we have content, return it
273 if lines:
274 return '\n'.join(lines)
276 # Empty response
277 return ""