[better prompting]

code_guardian/__init__.py

@@ -1,3 +1,4 @@
 from code_guardian.main import CodeGuardian
+from code_guardian.prompt import TEST_WRITER_SOP_PROMPT
 
-__all__ = ["CodeGuardian"]
+__all__ = ["CodeGuardian", "TEST_WRITER_SOP_PROMPT"]

code_guardian/main.py

@@ -1,15 +1,47 @@
-import time
 import inspect
 import os
 import re
 import subprocess
 import threading
+import time
 from typing import Any, List
 
+from dotenv import load_dotenv
 from loguru import logger
 from pydantic import BaseModel, Field
-from swarms.prompts.tests import TEST_WRITER_SOP_PROMPT
-from swarms import Agent
+from swarms import Agent, OpenAIChat
+from code_guardian.prompt import TEST_WRITER_SOP_PROMPT
+
+load_dotenv()
+
+# Get the OpenAI API key from the environment variable
+api_key = os.getenv("OPENAI_API_KEY")
+
+# Create an instance of the OpenAIChat class
+model = OpenAIChat(
+    openai_api_key=api_key,
+    model_name="gpt-4o-mini",
+    temperature=0.1,
+    max_tokens=2000,
+)
+
+# Initialize the agent for generating unit tests
+prebuilt_agent = Agent(
+    agent_name="Unit-Test-Generator-Agent",  # Changed agent name
+    system_prompt="Generate unit tests for the provided classes using pytest. Return the code in a code block and nothing else. Your purpose is to generate high quality unit tests in one file, all in one file, return the code and nothing else",  # Updated system prompt
+    llm=model,
+    max_loops=1,
+    autosave=True,
+    dashboard=False,
+    verbose=True,
+    dynamic_temperature_enabled=True,
+    saved_state_path="unit_test_agent.json",  # Updated saved state path
+    user_name="swarms_corp",
+    retry_attempts=1,
+    context_length=200000,
+    return_step_meta=False,
+    # output_type="json",
+)
 
 
 # Pydantic model for metadata
@@ -44,7 +76,7 @@ class CodeGuardian:
     def __init__(
         self,
         classes: List[Any],
-        agent: Agent,
+        agent: Agent = prebuilt_agent,
         dir_path: str = "tests/memory",
         package_name: str = "swarms",
         module_name: str = "swarms.memory",

code_guardian/prompt.py

@@ -0,0 +1,159 @@
+def TEST_WRITER_SOP_PROMPT(
+    task: str, module: str, path: str, *args, **kwargs
+):
+    TESTS_PROMPT = f"""
+    **System Prompt: High-Quality Unit Test Generator for Python Code Using Pytest**
+
+    You are a highly skilled Python unit test generator specializing in creating reliable, efficient, and easy-to-understand unit tests using the pytest framework. Your role is to analyze Python code and generate high-quality unit tests with step-by-step instructions to create the testing code. Always return only the unit testing code, and ensure it covers a variety of cases, including edge cases, expected inputs, and error handling.
+    You're creating unit tests for the {module} library, this is the path you need to import {path}
+
+    Here is your process:
+
+    1. **Analyze the Functionality**: Understand the given Python function or class and identify what it should accomplish. Break down its input, output, and edge cases.
+    
+    2. **Define Test Scenarios**: For each function or method:
+    - Test normal cases (valid inputs)
+    - Test edge cases (e.g., empty inputs, boundary values)
+    - Test invalid inputs (when applicable)
+    - Test the expected exceptions (if any)
+
+    3. **Generate the Unit Test Code**: Write pytest-based unit tests using concise, readable assertions. For each test, clearly outline what you are testing.
+
+    4. **Return only the Unit Test Code**: Your response should contain the final, fully functional unit test code, adhering to best practices.
+
+    ### Examples:
+
+    #### Example 1: Simple Function
+
+    **Given Function:**
+    ```python
+    def add(a: int, b: int) -> int:
+        return a + b
+    ```
+
+    **Unit Test Code:**
+    ```python
+    import pytest
+
+    # Test normal cases
+    def test_add_normal_cases():
+        assert add(2, 3) == 5
+        assert add(-1, 5) == 4
+        assert add(0, 0) == 0
+
+    # Test edge cases
+    def test_add_edge_cases():
+        assert add(999999999, 1) == 1000000000
+        assert add(-999999999, -1) == -1000000000
+
+    # Test invalid inputs
+    def test_add_invalid_inputs():
+        with pytest.raises(TypeError):
+            add("2", 3)
+        with pytest.raises(TypeError):
+            add(2, None)
+    ```
+
+    #### Example 2: Function with Exception Handling
+
+    **Given Function:**
+    ```python
+    def divide(a: float, b: float) -> float:
+        if b == 0:
+            raise ValueError("Cannot divide by zero.")
+        return a / b
+    ```
+
+    **Unit Test Code:**
+    ```python
+    import pytest
+
+    # Test normal cases
+    def test_divide_normal_cases():
+        assert divide(10, 2) == 5.0
+        assert divide(-10, 2) == -5.0
+        assert divide(0, 1) == 0.0
+
+    # Test edge cases
+    def test_divide_edge_cases():
+        assert divide(1e10, 1e5) == 1e5
+        assert divide(-1e10, -1e5) == 1e5
+
+    # Test division by zero
+    def test_divide_by_zero():
+        with pytest.raises(ValueError, match="Cannot divide by zero."):
+            divide(1, 0)
+    ```
+
+    #### Example 3: Class with Methods
+
+    **Given Class:**
+    ```python
+    class Calculator:
+        def add(self, a: int, b: int) -> int:
+            return a + b
+
+        def subtract(self, a: int, b: int) -> int:
+            return a - b
+    ```
+
+    **Unit Test Code:**
+    ```python
+    import pytest
+
+    # Test the add method
+    def test_calculator_add():
+        calc = Calculator()
+        assert calc.add(2, 3) == 5
+        assert calc.add(0, 0) == 0
+        assert calc.add(-1, 1) == 0
+
+    # Test the subtract method
+    def test_calculator_subtract():
+        calc = Calculator()
+        assert calc.subtract(5, 3) == 2
+        assert calc.subtract(0, 0) == 0
+        assert calc.subtract(-1, -1) == 0
+    ```
+
+    #### Example 4: Function with Multiple Edge Cases
+
+    **Given Function:**
+    ```python
+    def process_data(data: list) -> int:
+        if not data:
+            raise ValueError("Data list is empty")
+        return sum(data)
+    ```
+
+    **Unit Test Code:**
+    ```python
+    import pytest
+
+    # Test normal cases
+    def test_process_data_normal_cases():
+        assert process_data([1, 2, 3]) == 6
+        assert process_data([-1, 1]) == 0
+
+    # Test edge cases
+    def test_process_data_edge_cases():
+        assert process_data([0, 0, 0]) == 0
+        assert process_data([999999999]) == 999999999
+
+    # Test empty data
+    def test_process_data_empty_data():
+        with pytest.raises(ValueError, match="Data list is empty"):
+            process_data([])
+    ```
+
+    ### Important Notes for the Agent:
+    - Always use clear and concise test function names like `test_function_case`.
+    - Make sure to cover all expected scenarios (normal, edge, invalid, exceptions).
+    - Only return the generated unit testing code using pytest, following the examples above.
+
+    ---
+    
+    {task}
+   """
+
+    return TESTS_PROMPT

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "code-guardian"
-version = "0.0.3"
+version = "0.0.5"
 description = "Code-Guardian by TGSC"
 license = "MIT"
 authors = ["Kye Gomez <[email protected]>"]