devdox-ai-locust 0.1.2__tar.gz → 0.1.3.post1__tar.gz

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: devdox_ai_locust
-Version: 0.1.2
+Version: 0.1.3.post1
 Summary: AI-powered Locust load test generator from API documentation
 Author-email: Hayat Bourji <hayat.bourgi@montyholding.com>
 Maintainer-email: Hayat Bourji <hayat.bourgi@montyholding.com>
@@ -72,6 +72,21 @@ Dynamic: license-file
 
 DevDox AI Locust automatically generates comprehensive Locust load testing scripts from your API documentation (OpenAPI/Swagger specs). Using advanced AI capabilities, it creates realistic test scenarios, handles complex authentication flows, and generates production-ready performance tests.
 
+
+## 🆕 What's New in 0.1.3.post1
+
+### Performance & Reliability Improvements
+
+- **🚀 Asynchronous API Calls**: Migrated from `Together` to `AsyncTogether` for non-blocking API interactions, significantly improving generation speed and responsiveness
+- **⚡ Enhanced Timeout Handling**: Implemented robust timeout logic with configurable retry mechanisms for better error resilience
+- **🔧 Improved Code Extraction**: Enhanced `<code>` block parsing with detailed validation, multiple fallback scenarios, and better error messages
+- **🛡️ Better Error Management**: Comprehensive error handling throughout the AI generation pipeline with graceful degradation
+
+### Bug Fixes & Stability
+
+- Fixed edge cases in code block extraction where malformed responses could cause generation failures
+- Improved retry logic to handle transient API errors without interrupting the generation process
+
 ## ✨ Features
 
 - 🤖 **AI-Enhanced Generation**: Uses Together AI to create intelligent, realistic load test scenarios
@@ -82,6 +97,7 @@ DevDox AI Locust automatically generates comprehensive Locust load testing scrip
 - 🛠️ **Template-Based**: Highly customizable Jinja2 templates for different testing needs
 - 🔄 **Hybrid Approach**: Combines rule-based generation with AI enhancement
 - 📈 **Comprehensive Coverage**: Handles various HTTP methods, content types, and response scenarios
+- ⚡ **Asynchronous Processing**: Fast, non-blocking test generation with async/await
 
 ## 🚀 Quick Start
 
@@ -121,12 +137,9 @@ echo "API_KEY=your_together_ai_api_key_here" > .env
 # Generate from OpenAPI URL
 devdox_ai_locust generate --openapi-url https://api.example.com/openapi.json --output ./tests
 
-# Generate from local file
-dal generate --openapi-file ./api-spec.yaml --output ./load-tests
-
 # Generate with custom configuration
 devdox_ai_locust generate \
-   https://petstore.swagger.io/v2/swagger.json \
+   https://petstore.swagger.io/v3/swagger.json \
   --output ./petstore-tests \
   --together-api-key your_api_key \
   

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/README.md

@@ -8,6 +8,21 @@
 
 DevDox AI Locust automatically generates comprehensive Locust load testing scripts from your API documentation (OpenAPI/Swagger specs). Using advanced AI capabilities, it creates realistic test scenarios, handles complex authentication flows, and generates production-ready performance tests.
 
+
+## 🆕 What's New in 0.1.3.post1
+
+### Performance & Reliability Improvements
+
+- **🚀 Asynchronous API Calls**: Migrated from `Together` to `AsyncTogether` for non-blocking API interactions, significantly improving generation speed and responsiveness
+- **⚡ Enhanced Timeout Handling**: Implemented robust timeout logic with configurable retry mechanisms for better error resilience
+- **🔧 Improved Code Extraction**: Enhanced `<code>` block parsing with detailed validation, multiple fallback scenarios, and better error messages
+- **🛡️ Better Error Management**: Comprehensive error handling throughout the AI generation pipeline with graceful degradation
+
+### Bug Fixes & Stability
+
+- Fixed edge cases in code block extraction where malformed responses could cause generation failures
+- Improved retry logic to handle transient API errors without interrupting the generation process
+
 ## ✨ Features
 
 - 🤖 **AI-Enhanced Generation**: Uses Together AI to create intelligent, realistic load test scenarios
@@ -18,6 +33,7 @@ DevDox AI Locust automatically generates comprehensive Locust load testing scrip
 - 🛠️ **Template-Based**: Highly customizable Jinja2 templates for different testing needs
 - 🔄 **Hybrid Approach**: Combines rule-based generation with AI enhancement
 - 📈 **Comprehensive Coverage**: Handles various HTTP methods, content types, and response scenarios
+- ⚡ **Asynchronous Processing**: Fast, non-blocking test generation with async/await
 
 ## 🚀 Quick Start
 
@@ -57,12 +73,9 @@ echo "API_KEY=your_together_ai_api_key_here" > .env
 # Generate from OpenAPI URL
 devdox_ai_locust generate --openapi-url https://api.example.com/openapi.json --output ./tests
 
-# Generate from local file
-dal generate --openapi-file ./api-spec.yaml --output ./load-tests
-
 # Generate with custom configuration
 devdox_ai_locust generate \
-   https://petstore.swagger.io/v2/swagger.json \
+   https://petstore.swagger.io/v3/swagger.json \
   --output ./petstore-tests \
   --together-api-key your_api_key \
   

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "devdox_ai_locust"
-version = "0.1.2"
+version = "0.1.3.post1"
 description = "AI-powered Locust load test generator from API documentation"
 readme = "README.md"
 license = {text = "Apache-2.0" }

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/src/devdox_ai_locust/cli.py

@@ -6,7 +6,7 @@ from datetime import datetime, timezone
 from typing import Optional, Tuple, Union, List, Dict, Any
 from rich.console import Console
 from rich.table import Table
-from together import Together
+from together import AsyncTogether
 
 from .hybrid_loctus_generator import HybridLocustGenerator
 from .config import Settings
@@ -199,7 +199,7 @@ async def _generate_and_create_tests(
     auth: bool = False,
 ) -> List[Dict[Any, Any]]:
     """Generate tests using AI and create test files"""
-    together_client = Together(api_key=api_key)
+    together_client = AsyncTogether(api_key=api_key)
 
     with console.status("[bold green]Generating Locust tests with AI..."):
         generator = HybridLocustGenerator(ai_client=together_client)

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/src/devdox_ai_locust/config.py

@@ -8,7 +8,7 @@ from pydantic_settings import BaseSettings
 class Settings(BaseSettings):
     """Application settings."""
 
-    VERSION: str = "0.1.2"
+    VERSION: str = "0.1.3.post1"
 
     API_KEY: str = ""  # Fallback for backward compatibility
 

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/src/devdox_ai_locust/hybrid_loctus_generator.py

@@ -19,7 +19,7 @@ import shutil
 from devdox_ai_locust.utils.open_ai_parser import Endpoint
 from devdox_ai_locust.utils.file_creation import FileCreationConfig, SafeFileCreator
 from devdox_ai_locust.locust_generator import LocustTestGenerator, TestDataConfig
-from together import Together
+from together import AsyncTogether
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +27,15 @@ logger = logging.getLogger(__name__)
 test_data_file_path = "test_data.py"
 
 
+@dataclass
+class ErrorClassification:
+    """Classification of an error for retry logic"""
+
+    is_retryable: bool
+    backoff_seconds: float
+    error_type: str
+
+
 @dataclass
 class AIEnhancementConfig:
     """Configuration for AI enhancement"""
@@ -202,7 +211,7 @@ class HybridLocustGenerator:
 
     def __init__(
         self,
-        ai_client: Together,
+        ai_client: AsyncTogether,
         ai_config: Optional[AIEnhancementConfig] = None,
         test_config: Optional[TestDataConfig] = None,
         prompt_dir: str = "prompt",
@@ -211,7 +220,12 @@ class HybridLocustGenerator:
         self.ai_config = ai_config or AIEnhancementConfig()
         self.template_generator = LocustTestGenerator(test_config)
         self.prompt_dir = self._find_project_root() / prompt_dir
+        self._api_semaphore = asyncio.Semaphore(5)
         self._setup_jinja_env()
+        self.MAX_RETRIES = 3
+        self.RATE_LIMIT_BACKOFF = 10
+        self.NON_RETRYABLE_CODES = ["401", "403", "unauthorized", "forbidden"]
+        self.RATE_LIMIT_INDICATORS = ["429", "rate limit"]
 
     def _find_project_root(self) -> Path:
         """Find the project root by looking for setup.py, pyproject.toml, or .git"""
@@ -229,6 +243,45 @@ class HybridLocustGenerator:
             autoescape=False,
         )
 
+    def _classify_error(self, error: Exception, attempt: int) -> ErrorClassification:
+        """
+        Classify an error to determine retry behavior.
+
+        Args:
+            error: The exception that occurred
+            attempt: Current attempt number (0-indexed)
+
+        Returns:
+            ErrorClassification with retry decision and backoff time
+        """
+        error_str = str(error).lower()
+
+        # Non-retryable errors (auth/permission)
+        if any(code in error_str for code in self.NON_RETRYABLE_CODES):
+            logger.error(f"Authentication error, not retrying: {error}")
+            return ErrorClassification(
+                is_retryable=False, backoff_seconds=0, error_type="auth"
+            )
+
+        # Rate limit errors (retryable with longer backoff)
+        if any(indicator in error_str for indicator in self.RATE_LIMIT_INDICATORS):
+            logger.warning(f"Rate limit hit on attempt {attempt + 1}")
+            return ErrorClassification(
+                is_retryable=True,
+                backoff_seconds=self.RATE_LIMIT_BACKOFF,
+                error_type="rate_limit",
+            )
+
+        # Other retryable errors (exponential backoff)
+        logger.warning(
+            f"Retryable error on attempt {attempt + 1}: {type(error).__name__}"
+        )
+        return ErrorClassification(
+            is_retryable=True,
+            backoff_seconds=2**attempt,  # Exponential: 1s, 2s, 4s
+            error_type="retryable",
+        )
+
     async def generate_from_endpoints(
         self,
         endpoints: List[Endpoint],
@@ -573,10 +626,8 @@ class HybridLocustGenerator:
 
         return ""
 
-    async def _call_ai_service(self, prompt: str) -> Optional[str]:
-        """Call AI service with retry logic and validation"""
-
-        messages = [
+    def _build_messages(self, prompt: str) -> list[dict]:
+        return [
             {
                 "role": "system",
                 "content": "You are an expert Python developer specializing in Locust load testing. Generate clean, production-ready code with proper error handling. "
@@ -586,30 +637,42 @@ class HybridLocustGenerator:
             {"role": "user", "content": prompt},
         ]
 
-        for attempt in range(3):  # Retry logic
-            try:
-                response = await asyncio.wait_for(
-                    asyncio.to_thread(
-                        self.ai_client.chat.completions.create,
-                        model=self.ai_config.model,
-                        messages=messages,
-                        max_tokens=self.ai_config.max_tokens,
-                        temperature=self.ai_config.temperature,
-                        top_p=0.9,
-                        top_k=40,
-                        repetition_penalty=1.1,
-                    ),
-                    timeout=self.ai_config.timeout,
+    async def _make_api_call(self, messages: list[dict]) -> Optional[str]:
+        """Make API call - ONE job"""
+        async with self._api_semaphore:
+            api_call = self.ai_client.chat.completions.create(
+                model=self.ai_config.model,
+                messages=messages,
+                max_tokens=self.ai_config.max_tokens,
+                temperature=self.ai_config.temperature,
+                top_p=0.9,
+                top_k=40,
+                repetition_penalty=1.1,
+            )
+
+            # Wait for the API call with timeout
+            response = await asyncio.wait_for(
+                api_call,
+                timeout=self.ai_config.timeout,
+            )
+            if response.choices and response.choices[0].message:
+                content = response.choices[0].message.content.strip()
+                # Clean up the response
+                content = self._clean_ai_response(
+                    self.extract_code_from_response(content)
                 )
+                return content
 
-                if response.choices and response.choices[0].message:
-                    content = response.choices[0].message.content.strip()
+        return None
 
-                    # Clean up the response
-                    content = self._clean_ai_response(
-                        self.extract_code_from_response(content)
-                    )
+    async def _call_ai_service(self, prompt: str) -> Optional[str]:
+        """Call AI service with retry logic and validation"""
+        messages = self._build_messages(prompt)
 
+        for attempt in range(self.MAX_RETRIES):  # Retry logic
+            try:
+                async with self._api_semaphore:
+                    content = await self._make_api_call(messages)
                     if content:
                         return content
 
@@ -617,24 +680,41 @@ class HybridLocustGenerator:
                 logger.warning(f"AI service timeout on attempt {attempt + 1}")
 
             except Exception as e:
-                logger.warning(f"AI service error on attempt {attempt + 1}: {e}")
+                classification = self._classify_error(e, attempt)  # Helper 3
+
+                if not classification.is_retryable:
+                    return ""
+
+                if attempt < self.MAX_RETRIES - 1:
+                    await asyncio.sleep(classification.backoff_seconds)
+
+                    continue
 
-            if attempt < 2:  # Wait before retry
+            if attempt < self.MAX_RETRIES - 1:
                 await asyncio.sleep(2**attempt)
 
         return ""
 
     def extract_code_from_response(self, response_text: str) -> str:
         # Extract content between <code> tags
+        pattern = r"<code>(.*?)</code>"
+        matches = re.findall(pattern, response_text, re.DOTALL)
 
-        code_match = re.search(r"<code>(.*?)</code>", response_text, re.DOTALL)
-        if code_match:
-            content = code_match.group(1).strip()
-            # Additional validation - ensure we got actual content
-            if content and len(content) > 0:
-                return content
+        if not matches:
+            logger.warning("No <code> tags found, using full response")
+            return response_text.strip()
+
+        content = max(matches, key=len).strip()
+
+        # Content too short - use full response
+        if not content or len(content) <= 10:
+            logger.warning(
+                f"Code in tags too short ({len(content)} chars), using full response"
+            )
+            return response_text.strip()
 
-        return response_text.strip()
+        logger.debug(f"Extracted {len(content)} chars from <code> tags")
+        return str(content)
 
     def _clean_ai_response(self, content: str) -> str:
         """Clean and validate AI response"""

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/src/devdox_ai_locust.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: devdox_ai_locust
-Version: 0.1.2
+Version: 0.1.3.post1
 Summary: AI-powered Locust load test generator from API documentation
 Author-email: Hayat Bourji <hayat.bourgi@montyholding.com>
 Maintainer-email: Hayat Bourji <hayat.bourgi@montyholding.com>
@@ -72,6 +72,21 @@ Dynamic: license-file
 
 DevDox AI Locust automatically generates comprehensive Locust load testing scripts from your API documentation (OpenAPI/Swagger specs). Using advanced AI capabilities, it creates realistic test scenarios, handles complex authentication flows, and generates production-ready performance tests.
 
+
+## 🆕 What's New in 0.1.3.post1
+
+### Performance & Reliability Improvements
+
+- **🚀 Asynchronous API Calls**: Migrated from `Together` to `AsyncTogether` for non-blocking API interactions, significantly improving generation speed and responsiveness
+- **⚡ Enhanced Timeout Handling**: Implemented robust timeout logic with configurable retry mechanisms for better error resilience
+- **🔧 Improved Code Extraction**: Enhanced `<code>` block parsing with detailed validation, multiple fallback scenarios, and better error messages
+- **🛡️ Better Error Management**: Comprehensive error handling throughout the AI generation pipeline with graceful degradation
+
+### Bug Fixes & Stability
+
+- Fixed edge cases in code block extraction where malformed responses could cause generation failures
+- Improved retry logic to handle transient API errors without interrupting the generation process
+
 ## ✨ Features
 
 - 🤖 **AI-Enhanced Generation**: Uses Together AI to create intelligent, realistic load test scenarios
@@ -82,6 +97,7 @@ DevDox AI Locust automatically generates comprehensive Locust load testing scrip
 - 🛠️ **Template-Based**: Highly customizable Jinja2 templates for different testing needs
 - 🔄 **Hybrid Approach**: Combines rule-based generation with AI enhancement
 - 📈 **Comprehensive Coverage**: Handles various HTTP methods, content types, and response scenarios
+- ⚡ **Asynchronous Processing**: Fast, non-blocking test generation with async/await
 
 ## 🚀 Quick Start
 
@@ -121,12 +137,9 @@ echo "API_KEY=your_together_ai_api_key_here" > .env
 # Generate from OpenAPI URL
 devdox_ai_locust generate --openapi-url https://api.example.com/openapi.json --output ./tests
 
-# Generate from local file
-dal generate --openapi-file ./api-spec.yaml --output ./load-tests
-
 # Generate with custom configuration
 devdox_ai_locust generate \
-   https://petstore.swagger.io/v2/swagger.json \
+   https://petstore.swagger.io/v3/swagger.json \
   --output ./petstore-tests \
   --together-api-key your_api_key \
   

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/tests/test_cli.py

@@ -438,7 +438,7 @@ class TestGenerateAndCreateTests:
     """Test test generation and creation functionality."""
 
     @pytest.mark.asyncio
-    @patch("devdox_ai_locust.cli.Together")
+    @patch("devdox_ai_locust.cli.AsyncTogether")
     @patch("devdox_ai_locust.cli.HybridLocustGenerator")
     async def test_generate_and_create_tests_success(
         self,
@@ -450,7 +450,7 @@ class TestGenerateAndCreateTests:
     ):
         """Test successful test generation and creation."""
         # Mock Together client
-        mock_client = Mock()
+        mock_client = AsyncMock()
         mock_together_class.return_value = mock_client
 
         # Mock generator

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/tests/test_config.py

@@ -15,7 +15,7 @@ class TestSettings:
         """Test default settings values."""
         test_settings = Settings(_env_file=".env.example")
 
-        assert test_settings.VERSION == "0.1.2"
+        assert test_settings.VERSION == "0.1.3.post1"
         assert test_settings.API_KEY == ""
 
     def test_settings_with_env_vars(self):
@@ -108,7 +108,7 @@ class TestSettings:
     def test_global_settings_instance(self):
         """Test the global settings instance."""
         assert isinstance(settings, Settings)
-        assert settings.VERSION == "0.1.2"
+        assert settings.VERSION == "0.1.3.post1"
 
 
 class TestSettingsMethods:
@@ -120,7 +120,7 @@ class TestSettingsMethods:
         original_version = test_settings.VERSION
 
         # VERSION should be a class-level constant
-        assert original_version == "0.1.2"
+        assert original_version == "0.1.3.post1"
 
         # Even if we try to change it, it should remain the same
         # (depending on Pydantic implementation)
@@ -131,7 +131,7 @@ class TestSettingsMethods:
 
         # Create new instance to verify class-level value
         new_settings = Settings()
-        assert new_settings.VERSION == "0.1.2"
+        assert new_settings.VERSION == "0.1.3.post1"
 
     def test_settings_field_types(self):
         """Test that settings fields have correct types."""

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.3.post1}/tests/test_hybrid_loctus_generator.py

@@ -12,6 +12,7 @@ from devdox_ai_locust.hybrid_loctus_generator import (
     AIEnhancementConfig,
     EnhancementResult,
     EnhancementProcessor,
+    ErrorClassification,
 )
 from devdox_ai_locust.locust_generator import TestDataConfig
 
@@ -375,21 +376,31 @@ class TestHybridLocustGeneratorAsync:
         assert "import locust" in result
 
     @pytest.mark.asyncio
-    async def test_call_ai_service_timeout(self, mock_together_client):
-        """Test AI service call with timeout."""
-        generator = HybridLocustGenerator(ai_client=mock_together_client)
-        generator.ai_config.timeout = 0.1  # Very short timeout
+    async def test_ai_call_with_timeout(mock_together_client):
+        """Test AI service call that times out"""
+
+        async def mock_timeout(*args, **kwargs):
+            """Simulate a timeout by sleeping longer than expected"""
+            await asyncio.sleep(10)  # Long enough to trigger timeout
+            raise asyncio.TimeoutError("Simulated timeout")
+
+        mock_together_client.chat = Mock()
+        mock_together_client.chat.completions = Mock()
+        mock_together_client.chat.completions.create = AsyncMock(
+            side_effect=mock_timeout
+        )
 
-        # Mock a slow response
-        async def slow_response(*args, **kwargs):
-            await asyncio.sleep(1)
-            return mock_together_client.chat.completions.create(*args, **kwargs)
+        generator = HybridLocustGenerator(
+            ai_client=mock_together_client,
+            ai_config=AIEnhancementConfig(timeout=1),  # Short timeout
+        )
 
-        with patch("asyncio.to_thread", side_effect=slow_response):
-            result = await generator._call_ai_service("Test prompt")
+        # Call should timeout and return empty string after retries
+        result = await generator._call_ai_service("test prompt")
 
-            # Should return empty string on timeout
-            assert result == ""
+        assert result == ""
+        # Should have tried 3 times
+        assert mock_together_client.chat.completions.create.call_count == 3
 
     @pytest.mark.asyncio
     async def test_call_ai_service_with_retry(self, mock_together_client):
@@ -894,3 +905,186 @@ class TestHybridLocustGeneratorEdgeCases:
         for result in results:
             if not isinstance(result, Exception):
                 assert isinstance(result, str)
+
+
+class TestErrorClassification:
+    """Test ErrorClassification dataclass"""
+
+    def test_error_classification_creation(self):
+        """Test creating ErrorClassification"""
+        classification = ErrorClassification(
+            is_retryable=True, backoff_seconds=2.0, error_type="rate_limit"
+        )
+
+        assert classification.is_retryable is True
+        assert classification.backoff_seconds == 2.0
+        assert classification.error_type == "rate_limit"
+
+    def test_non_retryable_classification(self):
+        """Test non-retryable error classification"""
+        classification = ErrorClassification(
+            is_retryable=False, backoff_seconds=0, error_type="auth"
+        )
+
+        assert classification.is_retryable is False
+        assert classification.backoff_seconds == 0
+
+
+class TestBuildMessages:
+    """Test _build_messages method"""
+
+    def test_build_messages_structure(self, mock_together_client):
+        """Test message structure"""
+        generator = HybridLocustGenerator(ai_client=mock_together_client)
+
+        messages = generator._build_messages("test prompt")
+
+        assert len(messages) == 2
+        assert messages[0]["role"] == "system"
+        assert messages[1]["role"] == "user"
+        assert messages[1]["content"] == "test prompt"
+
+    def test_build_messages_system_prompt(self, mock_together_client):
+        """Test system prompt content"""
+        generator = HybridLocustGenerator(ai_client=mock_together_client)
+
+        messages = generator._build_messages("test")
+
+        system_content = messages[0]["content"]
+        assert "Locust load testing" in system_content
+        assert "<code>" in system_content
+        assert "DO NOT TRUNCATE" in system_content
+
+
+class TestMakeApiCall:
+    """Test _make_api_call method"""
+
+    @pytest.mark.asyncio
+    async def test_make_api_call_success(self, mock_together_client):
+        """Test successful API call"""
+        generator = HybridLocustGenerator(ai_client=mock_together_client)
+
+        messages = [{"role": "user", "content": "test"}]
+        result = await generator._make_api_call(messages)
+
+        assert result is not None
+        assert "import locust" in result
+
+    @pytest.mark.asyncio
+    async def test_make_api_call_empty_response(self, mock_together_client):
+        """Test API call with empty response"""
+        # Mock empty response
+        mock_response = Mock()
+        mock_response.choices = []
+
+        async def mock_create(*args, **kwargs):
+            await asyncio.sleep(0.01)
+            return mock_response
+
+        mock_together_client.chat = Mock()
+        mock_together_client.chat.completions = Mock()
+        mock_together_client.chat.completions.create = AsyncMock(
+            side_effect=mock_create
+        )
+
+        generator = HybridLocustGenerator(ai_client=mock_together_client)
+
+        messages = [{"role": "user", "content": "test"}]
+        result = await generator._make_api_call(messages)
+
+        assert result is None
+
+
+class TestCallAIService:
+    """Test _call_ai_service method with refactored code"""
+
+    @pytest.mark.asyncio
+    async def test_call_ai_service_success(self, mock_together_client):
+        """Test successful AI service call"""
+        generator = HybridLocustGenerator(ai_client=mock_together_client)
+
+        result = await generator._call_ai_service("Test prompt")
+
+        assert result is not None
+        assert "import locust" in result
+        assert mock_together_client.chat.completions.create.call_count == 1
+
+    @pytest.mark.asyncio
+    async def test_call_ai_service_with_timeout(self, mock_together_client):
+        """Test AI service call that times out"""
+
+        async def mock_timeout(*args, **kwargs):
+            await asyncio.sleep(10)
+            raise asyncio.TimeoutError("Simulated timeout")
+
+        mock_together_client.chat = Mock()
+        mock_together_client.chat.completions = Mock()
+        mock_together_client.chat.completions.create = AsyncMock(
+            side_effect=mock_timeout
+        )
+
+        generator = HybridLocustGenerator(
+            ai_client=mock_together_client,
+            ai_config=AIEnhancementConfig(timeout=1),
+        )
+
+        result = await generator._call_ai_service("test prompt")
+
+        assert result == ""
+        assert mock_together_client.chat.completions.create.call_count == 3
+
+    @pytest.mark.asyncio
+    async def test_call_ai_service_auth_error_no_retry(self, mock_together_client):
+        """Test that auth errors are not retried"""
+
+        async def mock_auth_error(*args, **kwargs):
+            await asyncio.sleep(0.01)
+            raise Exception("401 Unauthorized")
+
+        mock_together_client.chat = Mock()
+        mock_together_client.chat.completions = Mock()
+        mock_together_client.chat.completions.create = AsyncMock(
+            side_effect=mock_auth_error
+        )
+
+        generator = HybridLocustGenerator(ai_client=mock_together_client)
+
+        result = await generator._call_ai_service("test prompt")
+
+        assert result == ""
+        # Should only try once for auth errors
+        assert mock_together_client.chat.completions.create.call_count == 1
+
+    @pytest.mark.asyncio
+    async def test_call_ai_service_rate_limit_retry(self, mock_together_client):
+        """Test rate limit handling with retries"""
+        call_count = {"count": 0}
+
+        # Create successful response for 3rd attempt
+        mock_message = Mock()
+        mock_message.content = "<code>success_code</code>"
+        mock_choice = Mock()
+        mock_choice.message = mock_message
+        mock_response = Mock()
+        mock_response.choices = [mock_choice]
+
+        async def mock_rate_limit(*args, **kwargs):
+            call_count["count"] += 1
+            await asyncio.sleep(0.01)
+
+            if call_count["count"] < 3:
+                raise Exception("429 Rate limit exceeded")
+            return mock_response
+
+        mock_together_client.chat = Mock()
+        mock_together_client.chat.completions = Mock()
+        mock_together_client.chat.completions.create = AsyncMock(
+            side_effect=mock_rate_limit
+        )
+
+        generator = HybridLocustGenerator(ai_client=mock_together_client)
+
+        result = await generator._call_ai_service("test prompt")
+
+        assert result == "success_code"
+        assert call_count["count"] == 3