Update services/client.py

services/client.py

@@ -1,6 +1,4 @@
 from concurrent.futures import ThreadPoolExecutor
-from multiprocessing.pool import ThreadPool
-from pydoc import cli
 from openai import OpenAI
 import os
 from dotenv import load_dotenv
@@ -13,87 +11,86 @@ BASE_URL = "https://api.novita.ai/v3/openai"
 MODEL_NAME = "meta-llama/llama-3.2-1b-instruct"
 
 client = OpenAI(
-    base_url= BASE_URL,
+    base_url=BASE_URL,
     api_key=os.environ.get("NOVITA_API_KEY"),
 )
 
-def generate_commented_code(code:str, max_token:int):
+def generate_commented_code(code: str, max_token: int):
+    system_prompt = (
+        "You are a senior software engineer skilled in writing precise comments only. "
+        "Add clear, concise, and professional comments to the provided code. "
+        "Include docstrings for all functions/classes and inline comments for complex logic. "
+        "Strictly do not add or provide alternative code."
+    )
     try:
         chat_responses = client.chat.completions.create(
-        model= MODEL_NAME,
-        temperature=0.7,
-        stream= False,
-        max_tokens= max_token,
-        messages = [
-                {
-                    "role": "system",
-                    "content": (
-                        "You are a senior software engineer skilled in writing precise comments only. Follow these instructions strictly:\n\n"
-                        "Identify the programming language used and then do following:\n"
-                        "Add **clear, concise, and professional comments** to the provided code. Include:\n"
-                        "Docstrings for all functions and classes\n"
-                        "Inline comments explaining non-obvious or critical logic\n\n"
-                        "Strictly do not do anything else than writing comments. Stricltly do not add or provide alternative code"
-
-                        "Example response:"
-                        '''
-                        def add(a, b):
-                            """
-                            Returns the sum of two numbers.
-
-                            Args:
-                                a (int): The first number.
-                                b (int): The second number.
-
-                            Returns:
-                                int: The sum of a and b.
-
-                            Raises:
-                                TypeError: If either a or b is not an integer.
-                                ValueError: If either a or b is negative.
-                            """
-                            # Check if both inputs are integers
-                            if not isinstance(a, int) or not isinstance(b, int):
-                                raise TypeError("Both inputs must be integers.")
-                            # Check if both inputs are non-negative
-                            if a < 0 or b < 0:
-                                raise ValueError("Both inputs must be non-negative.")
+            model=MODEL_NAME,
+            temperature=0.7,
+            stream=False,
+            max_tokens=max_token,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": code}
+            ]
+        )
+        return chat_responses.choices[0].message.content
+    except Exception as ex:
+        return f"Error generating comments: {ex}"
 
-                            # Use the built-in addition operator to perform the calculation
-                            # This is more efficient and readable than using a loop or recursive function
-                            return a + b
-                        '''
-                    )
-                },
+def generate_documentation(code: str, max_token: int):
+    system_prompt = (
+        "You are a senior software engineer skilled in writing structured documentation only. "
+        "Produce a strictly formatted JSON object with the following keys: "
+        "'name', 'description', 'parameters', 'returns', 'exceptions', and 'examples'."
+    )
+    try:
+        chat_responses = client.chat.completions.create(
+            model=MODEL_NAME,
+            temperature=0.7,
+            stream=False,
+            max_tokens=max_token,
+            messages=[
+                {"role": "system", "content": system_prompt},
                 {"role": "user", "content": code}
             ]
         )
-        
-        result = chat_responses.choices[0].message.content
-        return result
+        return chat_responses.choices[0].message.content
     except Exception as ex:
-        return f"Error generating comments {ex}"
+        return f"Error generating documentation: {ex}"
 
-def generate_documentation(code:str, max_token:int):
+def generate_improved_code(code: str, max_token: int):
+    system_prompt = (
+        "You are a senior software engineer. Analyze the code for possible improvements. "
+        "If optimization is possible, provide a cleaner or more efficient version. "
+        "Otherwise, write 'Already optimal'. "
+        "Always report the time complexity of both the original and improved versions."
+    )
     try:
         chat_responses = client.chat.completions.create(
-            model= MODEL_NAME,
+            model=MODEL_NAME,
             temperature=0.7,
-            stream= False,
-            max_tokens= max_token,
-            messages = [
-                    {
-                        "role": "system",
-                        "content": (
-                            "You are a senior software engineer skilled in writing structured documentation only. Follow these instructions strictly:\n\n"
-                            "Identify the programming language used and then do following:\n"
-                            "Produce a **strictly formatted JSON** object with:\n"
-                            "name: Function or class name\n"
-                            "description: Brief summary of what it does\n"
-                            "parameters: Dictionary with param name, type, and description\n"
-                            "returns: Return name, type, and description\n"
-                            "exceptions (optional): Raised exceptions with reasons\n"
-                            "examples: Code usage examples with expected output\n\n"
+            stream=False,
+            max_tokens=max_token,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": code}
+            ]
+        )
+        return chat_responses.choices[0].message.content
+    except Exception as ex:
+        return f"Error generating improved code: {ex}"
+
+def generate_code_comments_and_docs(code: str, max_token: int = 1024) -> List:
+    try:
+        with ThreadPoolExecutor() as executor:
+            future1 = executor.submit(generate_commented_code, code, max_token)
+            future2 = executor.submit(generate_documentation, code, max_token)
+            future3 = executor.submit(generate_improved_code, code, max_token)
 
-                            "Example response: "
-                            """
+            commented = future1.result()
+            documentation = future2.result()
+            improved = future3.result()
+
+        return [commented, documentation, improved]
+    except Exception as ex:
+        return [f"Error generating comments: {ex}", "{}", f"Error: {ex}"]