Coverage for src/ollamapy/analysis_engine.py: 20%
79 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-01 12:29 -0400
1"""AI analysis engine for action selection and parameter extraction."""
3import re
4from typing import List, Dict, Tuple, Any
5from .ollama_client import OllamaClient
6from .skills import get_available_actions, SKILL_REGISTRY
7from .parameter_utils import extract_parameter_from_response
10class AnalysisEngine:
11 """Handles AI-based action selection and parameter extraction."""
13 def __init__(self, analysis_model: str, client: OllamaClient):
14 """Initialize the analysis engine.
16 Args:
17 analysis_model: The model to use for analysis
18 client: The OllamaClient instance
19 """
20 self.analysis_model = analysis_model
21 self.client = client
22 self.actions = get_available_actions()
24 def remove_thinking_blocks(self, text: str) -> str:
25 """Remove <think></think> blocks from AI output.
27 This allows models with thinking steps to be used without interference.
29 Args:
30 text: The text to clean
32 Returns:
33 The text with thinking blocks removed
34 """
35 # Remove <think>...</think> blocks (including nested content)
36 cleaned = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
37 return cleaned.strip()
39 def get_cleaned_response(
40 self,
41 prompt: str,
42 system_message: str = "You are a decision assistant. Answer only 'yes' or 'no' to questions.",
43 show_context: bool = True,
44 ) -> str:
45 """Get a cleaned response from the analysis model.
47 Args:
48 prompt: The prompt to send
49 system_message: The system message to use
50 show_context: Whether to show context usage
52 Returns:
53 The cleaned response content
54 """
55 response_content = ""
56 try:
57 # Show context usage if requested
58 if show_context:
59 self.client.print_context_usage(
60 self.analysis_model, prompt, system_message
61 )
63 # Stream the response from the analysis model
64 for chunk in self.client.chat_stream(
65 model=self.analysis_model,
66 messages=[{"role": "user", "content": prompt}],
67 system=system_message,
68 ):
69 response_content += chunk
71 # Remove thinking blocks if present
72 return self.remove_thinking_blocks(response_content)
74 except Exception as e:
75 print(f"\n❌ Error getting response: {e}")
76 return ""
78 def ask_yes_no_question(self, prompt: str, show_context: bool = True) -> bool:
79 """Ask the analysis model a yes/no question and parse the response.
81 This is the core of our simplified analysis. We ask a clear yes/no question
82 and parse the response to determine if the answer is yes.
84 Args:
85 prompt: The yes/no question to ask
87 Returns:
88 True if the model answered yes, False otherwise
89 """
90 cleaned_response = self.get_cleaned_response(prompt, show_context=show_context)
92 # Convert to lowercase for easier parsing
93 response_lower = cleaned_response.lower().strip()
95 # Check for yes indicators at the beginning of the response
96 # This handles "yes", "yes.", "yes,", "yes!", etc.
97 if response_lower.startswith("yes"):
98 return True
100 # Also check if just "yes" appears alone in the first few characters
101 if response_lower[:10].strip() == "yes":
102 return True
104 # If we see "no" at the start, definitely return False
105 if response_lower.startswith("no"):
106 return False
108 # Default to False if unclear
109 return False
111 def extract_single_parameter(
112 self,
113 user_input: str,
114 action_name: str,
115 param_name: str,
116 param_spec: Dict[str, Any],
117 ) -> Any:
118 """Extract a single parameter value from user input.
120 This asks the AI to extract just one parameter value, making it simple and reliable.
122 Args:
123 user_input: The original user input
124 action_name: The name of the action being executed
125 param_name: The name of the parameter to extract
126 param_spec: The specification for this parameter (type, description, required)
128 Returns:
129 The extracted parameter value, or None if not found
130 """
131 param_type = param_spec.get("type", "string")
132 param_desc = param_spec.get("description", "")
134 # Build a simple, focused prompt for parameter extraction
135 prompt = f"""From this user input: "{user_input}"
137Extract the value for the parameter '{param_name}' which is described as: {param_desc}
139The parameter type is: {param_type}
141Respond with ONLY the parameter value, nothing else.
142If the parameter value cannot be found in the user input, respond with only: NOT_FOUND
144Examples for {param_type} type:
145- If type is number and user says "square root of 16", respond: 16
146- If type is string and user says "weather in Paris", respond: Paris
147- If type is string and user says "calculate 5+3", respond: 5+3
148"""
150 system_message = "You are a parameter extractor. Respond only with the extracted value or NOT_FOUND."
151 cleaned_response = self.get_cleaned_response(
152 prompt, system_message, show_context=True
153 ).strip()
155 return extract_parameter_from_response(cleaned_response, param_type)
157 def select_all_applicable_actions(
158 self, user_input: str
159 ) -> List[Tuple[str, Dict[str, Any]]]:
160 """Select ALL applicable actions and extract their parameters.
162 This evaluates EVERY action and returns a list of all that apply.
164 Args:
165 user_input: The user's input to analyze
167 Returns:
168 List of tuples containing (action_name, parameters_dict)
169 """
170 print(f"🔍 Analyzing user input with {self.analysis_model}...")
172 selected_actions = []
174 # Iterate through EVERY action and check if it's applicable
175 for action_name, action_info in self.actions.items():
176 # Build a comprehensive prompt for this specific action
177 description = action_info["description"]
178 vibe_phrases = action_info.get("vibe_test_phrases", [])
179 parameters = action_info.get("parameters", {})
181 # Create the yes/no prompt for this action
182 prompt = f"""Consider this user input: "{user_input}"
184Should the '{action_name}' action be used?
186Action description: {description}
188Example phrases that would trigger this action:
189{chr(10).join(f'- "{phrase}"' for phrase in vibe_phrases[:5]) if vibe_phrases else '- No examples available'}
191{f"This action requires parameters: {', '.join(parameters.keys())}" if parameters else "This action requires no parameters"}
193Answer only 'yes' if this action should be used for the user's input, or 'no' if it should not.
194"""
196 # Ask if this action is applicable
197 print(f" Checking {action_name}... ", end="", flush=True)
199 if self.ask_yes_no_question(prompt):
200 print("✓ Selected!", end="")
202 # Extract parameters if needed
203 extracted_params = {}
204 if parameters:
205 print(" Extracting parameters:", end="")
207 for param_name, param_spec in parameters.items():
208 value = self.extract_single_parameter(
209 user_input, action_name, param_name, param_spec
210 )
212 if value is not None:
213 extracted_params[param_name] = value
214 print(f" {param_name}✓", end="")
215 else:
216 if param_spec.get("required", False):
217 print(f" {param_name}✗(required)", end="")
218 # Still add the action, but note the missing parameter
219 else:
220 print(f" {param_name}✗", end="")
222 selected_actions.append((action_name, extracted_params))
223 print() # New line after this action
224 else:
225 print("✗")
227 if selected_actions:
228 print(
229 f"🎯 Selected {len(selected_actions)} action(s): {', '.join([a[0] for a in selected_actions])}"
230 )
231 else:
232 print("🎯 No specific actions needed for this query")
234 return selected_actions
236 def generate_custom_python_script(self, user_input: str) -> str:
237 """Generate a custom Python script based on user input.
239 Args:
240 user_input: The user's request
242 Returns:
243 Generated Python script
244 """
245 prompt = f"""Generate a Python script to help with this request: {user_input}
247Requirements:
248- Create a standalone Python script that solves the user's request
249- Use print() statements for all output
250- Include error handling where appropriate
251- Do not use any external libraries unless absolutely necessary
252- Make the script clear and well-commented
254Output ONLY the Python code, no explanations or markdown:"""
256 system_message = (
257 "You are a Python code generator. Output only valid Python code."
258 )
259 script = self.get_cleaned_response(prompt, system_message)
261 # Clean up any markdown formatting if present
262 if "```python" in script:
263 script = script.split("```python")[1].split("```")[0]
264 elif "```" in script:
265 script = script.split("```")[1].split("```")[0]
267 return script.strip()