Coverage for app/services/suggestion_service.py: 93%
150 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-02 02:49 +0000
1from app.controllers.ai import (
2 openai_client,
3 gemini_client,
4 vendors,
5 good_command,
6 bad_command,
7 OLLAMA_URL,
8 default_openai_parameters,
9)
10from app.controllers.database import client
11from app.models.errors import ModelError, DatabaseError
12import requests
13from flask import current_app
14import time
15import json
16from app.models.user import Suggestion
17import traceback
18import logging
21def getSuggestion(
22 prompt: str,
23 vendor: str = vendors.Google,
24 model_name: str = "codellama",
25 model_params: dict = None,
26):
27 """
28 Handles suggestions from different models based on the provided model name.
30 Args:
31 prompt (str): The prompt (or piece of code) to generate suggestions from.
32 vendor (str): The vendor of the AI model(OpenAI, Ollama, Google)
33 model_name (str): The model to use( See vendor website).
34 model_params (dict): Additional parameters to be sent to the AI.
35 is_correct (bool): Whether to generate a correct suggestion or one with a small error.
37 Returns:
38 dict: A dictionary containing the suggestion response.
40 Raises:
41 Exception: If there is an error with the model API.
42 """
44 try:
45 vendor_enum = vendors(vendor) # Convert string to Enum
46 except ValueError:
47 vendor_enum = vendors.Google # Default if invalid
49 # Choose model-specific logic
50 match vendor_enum:
51 case vendors.OpenAI:
52 return getSuggestionFromOpenAI(
53 prompt=prompt,
54 model=model_name,
55 )
56 case vendors.Ollama:
57 return getSuggestionFromOllama(
58 prompt=prompt,
59 model_name=model_name,
60 )
61 case vendors.Google:
62 return getSuggestionFromGoogle(
63 prompt=prompt,
64 )
65 case _:
66 return getSuggestionFromGoogle(
67 prompt=prompt,
68 )
71def getSuggestionFromOpenAI(
72 prompt: str,
73 model: str = "gpt-4o-mini",
74 model_params: dict = None,
75 is_correct: bool = True,
76):
77 """
78 Completes a code suggestion using OpenAI's API.
79 """
81 try:
82 # Ensure model_params is a dictionary
83 if model_params is None:
84 model_params = default_openai_parameters # Use a default config
86 if is_correct:
87 system_message = """You are an AI coding assistant. Your goal is to autocomplete and extend the user's code seamlessly, predicting what they are likely to write next.
89 Follow these rules explicitly:
90 - Provide only the completion—no explanations, comments, or markdown formatting.
91 - Base your predictions on best practices, patterns, and context from the given code.
92 - Aim for concise and efficient solutions that align with the user's coding style.
93 - Always act as if you are an inline code completion tool, not a chatbot.
94 - Avoid markdown or any extra formatting.
95 - Never repeat the user or their code, continue exactly where they left off."""
96 else:
97 system_message = """You are an AI coding assistant. Your goal is to autocomplete and extend the user's code seamlessly, predicting what they are likely to write next.
99 Follow these rules explicitly:
100 - Provide only the completion—no explanations, comments, or markdown formatting.
101 - Base your predictions on best practices, patterns, and context from the given code.
102 - Aim for concise and efficient solutions that align with the user's coding style.
103 - Always act as if you are an inline code completion tool, not a chatbot.
104 - Avoid markdown or any extra formatting.
105 - Never repeat the user or their code, continue exactly where they left off.
107 This suggestions should be slightly incorrect. Insert a incorrect change, like a syntactical mistake or misnaming a variable."""
109 messages = [
110 {"role": "system", "content": system_message},
111 {"role": "user", "content": prompt},
112 ]
113 with current_app.app_context():
114 completion = openai_client.chat.completions.create(
115 temperature=model_params.get("temperature"),
116 top_p=model_params.get("top_p"),
117 max_tokens=model_params.get("max_tokens"),
118 model=model,
119 messages=messages,
120 )
122 suggestion = completion.choices[0].message.content.strip("```")
123 suggestion = completion.choices[0].message.content.lstrip(prompt)
125 return suggestion
127 except Exception as e:
128 print(f"Error generating suggestion using OpenAI's API: {e}")
129 raise ModelError(f"Error generating suggestion using OpenAI's API: {e}")
132def getSuggestionFromOllama(
133 prompt: str,
134 model_name: str,
135 is_correct: bool = True,
136):
137 """
138 Generates a suggestion from Ollama.
139 """
141 full_prompt = (good_command if is_correct else bad_command) + prompt
143 try:
144 response = requests.post(
145 OLLAMA_URL,
146 json={
147 "model": model_name,
148 "prompt": full_prompt,
149 "keep_alive": "1h",
150 "stream": False,
151 },
152 )
153 response.raise_for_status() # Raise exception for HTTP errors
154 result = response.json()
155 return result["response"]
157 except Exception as e:
158 print(f"Error fetching Ollama suggestion: {e}")
159 raise ModelError(f"Error fetching Ollama suggestion: {e}")
162def getSuggestionFromGoogle(
163 prompt: str,
164):
165 """
166 Sends the prompt to the model and returns an array of two code snippets:
167 one correct and one with a small logic error.
169 Args:
170 prompt (str): The code snippet to complete (e.g., "function add").
172 Returns:
173 list[str]: An array containing two code snippets.
174 """
176 full_prompt = f"""You are a code completion assistant that returns ONLY a JSON array with exactly 2 elements:
177 1. Correct completion of the EXACT code fragment provided
178 2. Same completion but with a small logical error
180 RULES:
181 - STRICTLY return ONLY a valid JSON array (no other text, markdown, or explanations)
182 - ONLY provide the missing part needed to complete the EXACT code fragment
183 - DO NOT repeat or rewrite the existing code
184 - Ensure the completion is syntactically correct
185 - Maintain the exact same indentation as the original
187 EXAMPLE INPUT:
188 function add(a, b) { \n return a
190 EXAMPLE OUTPUT:
191 [" + b;", " - b;"]
193 ACTUAL CODE TO COMPLETE (JavaScript):
194 {prompt}
196 ONLY RETURN THE MISSING PART AS SHOWN IN THE EXAMPLE:"""
198 try:
199 start_time = time.time()
201 response = gemini_client.chat_session.send_message(full_prompt)
203 latency = time.time() - start_time
205 if not response.text:
206 return []
208 if response.usage_metadata:
209 usage_metadata = response.usage_metadata
210 input_tokens = usage_metadata.prompt_token_count
211 output_tokens = usage_metadata.candidates_token_count
212 total_tokens = usage_metadata.total_token_count
213 else:
214 input_tokens = output_tokens = total_tokens = -1
216 data = {
217 "provider": "google",
218 "model": "gemini-2.0-flash",
219 "input_tokens": input_tokens,
220 "output_tokens": output_tokens,
221 "total_tokens": total_tokens,
222 "latency_seconds": latency,
223 }
225 client.table("ai_usage").insert(data).execute()
227 try:
228 result = json.loads(response.text)
229 if isinstance(result, list):
230 return result
231 return []
232 except json.JSONDecodeError:
233 logging.error("Final JSON parse failed. Raw response: %s", response.text)
234 return []
236 except Exception as e:
237 logging.exception(
238 "Error communicating with Gemini (Type: %s): %s", type(e).__name__, e
239 )
240 return []
243def getAvailableModels(vendor: vendors):
244 vendor_enum = vendors(vendor)
246 # Choose model-specific logic
247 match vendor_enum:
248 case vendors.OpenAI:
249 return getModelsFromOpenAI()
250 case vendors.Ollama:
251 return getModelsFromOllama()
252 case vendors.Google:
253 return getModelsFromGoogle()
254 case _:
255 raise ValueError(f"Unsupported vendor: {vendor}")
258def getModelsFromOpenAI():
259 try:
260 with current_app.app_context():
261 models = openai_client.models.list() # Fetch models from OpenAI API
262 model_names = [model.id for model in models.data] # Extract model names
263 return model_names
265 except Exception as e:
266 raise e
269def getModelsFromOllama():
270 try:
271 response = requests.get("http://localhost:11434/api/tags")
273 response.raise_for_status()
275 models = response.json()
276 print(models)
278 model_names = models.get("models", [])
280 return model_names
282 except requests.exceptions.RequestException as e:
283 raise Exception(f"Error fetching models from Ollama: {e}")
286def getModelsFromGoogle():
287 models_for_generate_content = []
289 for m in gemini_client.models.list():
290 if "generateContent" in m.supported_actions:
291 models_for_generate_content.append(m.name)
293 return models_for_generate_content
296def generate_refined_prompt(raw_prompt: str) -> str:
297 """
298 Uses AI to transform a raw code prompt into a well-structured completion request.
299 Returns just the refined prompt string.
300 """
302 prompt = f"""Refine the following input prompt by analyzing the code and generating a minimal, specific instruction for code completion.
304 Original Input:
305 {raw_prompt}
307 Rules:
308 1. Identifies the exact code portion needing completion
309 2. Describe what it should do in 1-2 sentence
310 3. Provides minimal but essential context
311 4. Don't return an array
312 5. Follows this format:
313 '''
314 Language: [detected language]
315 Context: 1-2 sentence description
316 '''
317 6. Example format:
318 "Language": "javascript",
319 "Context": "Complete the cube function",
320 """
322 try:
323 response = gemini_client.chat_session.send_message(prompt)
324 return response.text.strip()
325 except Exception as e:
326 raise ValueError(f"AI prompt refinement failed: {str(e)}")
329def generate_hint_from_gemini(prompt: str, wrong_code: str, right_code: str) -> str:
330 """
331 Generates a hint explaining the difference between wrong and right code.
333 Args:
334 prompt: Original user prompt/context
335 wrong_code: The incorrect code version
336 right_code: The correct code version
338 Returns:
339 str: Explanation/hint about what's wrong with the incorrect version
340 """
341 hint_prompt = f"""You are a helpful coding assistant that explains subtle code differences.
343 Context: {prompt}
345 Incorrect Version:
346 {wrong_code}
348 Correct Version:
349 {right_code}
351 Generate a helpful hint that:
352 1. Explains why the incorrect version might be wrong
353 2. Gives a clue about the issue without revealing the solution
354 3. Focuses on the logical error, not syntax
355 4. Is concise (1-2 sentences)
357 Example Response:
358 "This version uses the wrong operator. Consider whether you should be combining or comparing the values."
360 Your hint:"""
362 try:
363 response = gemini_client.chat_session.send_message(hint_prompt)
364 return response.text.strip()
365 except Exception as e:
366 return f"Could not generate hint: {str(e)}"
369def generate_explanation_from_gemini(
370 prompt: str, wrong_code: str, right_code: str
371) -> str:
372 """
373 Generates a explanation telling the user why the suggested code was wrong.
375 Args:
376 prompt: Original user prompt/context
377 wrong_code: The incorrect code version
378 right_code: The correct code version
380 Returns:
381 str: Explanation about what's wrong with the incorrect version
382 """
383 hint_prompt = f"""You are a helpful coding assistant that explains subtle code differences.
385 Context: {prompt}
387 Incorrect Version:
388 {wrong_code}
390 Correct Version:
391 {right_code}
393 Generate an explanation for the user that explains what is wrong the "Incorrect" version of the code, and why they should
394 use the "Correct" version of the code. Avoid using markdown or any formatting. Be informative but concise."""
396 try:
397 response = gemini_client.chat_session.send_message(hint_prompt)
398 return response.text.strip()
399 except Exception as e:
400 return f"Could not generate explanation: {str(e)}"
403def check_code_correctness(prompt: str, wrong_code: str, fixed_code: str) -> bool:
404 """
405 Uses an AI model to determine whether the fixed code is correct.
407 Args:
408 prompt: The original user prompt/context.
409 wrong_code: The incorrect code submitted.
410 fixed_code: The corrected version.
412 Returns:
413 bool: True if the fix is correct, False otherwise.
414 """
415 validation_prompt = f"""You are a code review assistant.
417 Context: {prompt}
419 Incorrect Version:
420 {wrong_code}
422 Fixed Version:
423 {fixed_code}
425 Does the fixed version fully correct the mistake in the incorrect version?
426 Respond with only 'true' or 'false'."""
428 try:
429 response = gemini_client.chat_session.send_message(validation_prompt)
430 return response.text.strip().lower() == "true"
431 except Exception as e:
432 print(f"Error validating code fix: {e}")
433 return False
436def get_suggestion_by_id(suggestion_id: str):
437 try:
438 response = (
439 client.table("suggestions")
440 .select("*")
441 .eq("id", suggestion_id)
442 .single()
443 .execute()
444 )
446 if not response.data:
447 return None
449 return Suggestion(**response.data)
451 except Exception as e:
452 logging.error(f"Error fetching suggestion {suggestion_id}: {e}")
453 raise DatabaseError(f"Failed to retrieve suggestion: {str(e)}")