devdox-ai-locust 0.1.2__tar.gz → 0.1.7__tar.gz

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: devdox_ai_locust
-Version: 0.1.2
+Version: 0.1.7
 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>
@@ -67,11 +67,20 @@ Dynamic: license-file
 [![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Python 3.12+](https://img.shields.io/badge/python-3.12+-blue.svg)](https://www.python.org/downloads/)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Quality Gate Status](https://sonarcloud.io/api/project_badges/measure?project=montymobile1_devdox-ai-locust&metric=alert_status)](https://sonarcloud.io/dashboard?id=montymobile1_devdox-ai-locust)
 
 > **AI-powered Locust load test generator from API documentation**
 
 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.
 
+## [0.1.7] - 2025-11-11
+## 🆕 What's New in 0.1.7
+
+**Improved Naming Convention:** Generated class names now automatically capitalize group names to ensure consistent and professional naming in generated Python code.
+
+**MongoDB Environment Setup:** Added example MongoDB variables (MONGODB_URI, MONGODB_DB_NAME, etc.) in the .env.example file to make it easier to configure MongoDB connections.
+
+
 ## ✨ Features
 
 - 🤖 **AI-Enhanced Generation**: Uses Together AI to create intelligent, realistic load test scenarios
@@ -82,6 +91,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,18 +131,35 @@ 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://petstore3.swagger.io/api/v3/openapi.json \
   --output ./petstore-tests \
   --together-api-key your_api_key \
   
+ # Generate with db integration 
+devdox_ai_locust generate \
+  https://petstore3.swagger.io/api/v3/openapi.json \
+  --output ./petstore-tests \
+  --db-type mongo \
 ```
 
 
+## 🚀 Installation with Inputs
+
+Add this step to your GitHub Actions workflow:
+
+```yaml
+- name: DevDox Locust Test Generator
+  uses: montymobile1/devdox-ai-locust@v0.1.6
+  with:
+    swagger_url: "https://portal-api.devdox.ai/openapi.json"
+    output: "generated_tests"
+    users: "15"
+    spawn_rate: "3"
+    run_time: "10m"
+    together_api_key: ${{ secrets.TOGETHER_API_KEY }}
+
 
 ## 📖 Documentation
 

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.7}/README.md

@@ -3,11 +3,20 @@
 [![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Python 3.12+](https://img.shields.io/badge/python-3.12+-blue.svg)](https://www.python.org/downloads/)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Quality Gate Status](https://sonarcloud.io/api/project_badges/measure?project=montymobile1_devdox-ai-locust&metric=alert_status)](https://sonarcloud.io/dashboard?id=montymobile1_devdox-ai-locust)
 
 > **AI-powered Locust load test generator from API documentation**
 
 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.
 
+## [0.1.7] - 2025-11-11
+## 🆕 What's New in 0.1.7
+
+**Improved Naming Convention:** Generated class names now automatically capitalize group names to ensure consistent and professional naming in generated Python code.
+
+**MongoDB Environment Setup:** Added example MongoDB variables (MONGODB_URI, MONGODB_DB_NAME, etc.) in the .env.example file to make it easier to configure MongoDB connections.
+
+
 ## ✨ Features
 
 - 🤖 **AI-Enhanced Generation**: Uses Together AI to create intelligent, realistic load test scenarios
@@ -18,6 +27,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,18 +67,35 @@ 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://petstore3.swagger.io/api/v3/openapi.json \
   --output ./petstore-tests \
   --together-api-key your_api_key \
   
+ # Generate with db integration 
+devdox_ai_locust generate \
+  https://petstore3.swagger.io/api/v3/openapi.json \
+  --output ./petstore-tests \
+  --db-type mongo \
 ```
 
 
+## 🚀 Installation with Inputs
+
+Add this step to your GitHub Actions workflow:
+
+```yaml
+- name: DevDox Locust Test Generator
+  uses: montymobile1/devdox-ai-locust@v0.1.6
+  with:
+    swagger_url: "https://portal-api.devdox.ai/openapi.json"
+    output: "generated_tests"
+    users: "15"
+    spawn_rate: "3"
+    run_time: "10m"
+    together_api_key: ${{ secrets.TOGETHER_API_KEY }}
+
 
 ## 📖 Documentation
 

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.7}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "devdox_ai_locust"
-version = "0.1.2"
+version = "0.1.7"
 description = "AI-powered Locust load test generator from API documentation"
 readme = "README.md"
 license = {text = "Apache-2.0" }
@@ -101,7 +101,7 @@ exclude = ["tests*"]
 
 
 [tool.setuptools.package-data]
-devdox_ai_locust = ["schemas/*.json", "templates/*.j2","prompt/*.j2", "py.typed","*.j2"]
+devdox_ai_locust = ["schemas/*.json", "templates/*.j2","templates/mongo/*","prompt/*.j2", "py.typed","*.j2"]
 
 # flake8 configuration (legacy - remove if using ruff)
 [tool.flake8]

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.7}/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
@@ -197,19 +197,20 @@ async def _generate_and_create_tests(
     custom_requirement: Optional[str] = "",
     host: Optional[str] = "0.0.0.0",
     auth: bool = False,
+    db_type: str = "",
 ) -> 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)
-
         test_files, test_directories = await generator.generate_from_endpoints(
             endpoints=endpoints,
             api_info=api_info,
             custom_requirement=custom_requirement,
             target_host=host,
             include_auth=auth,
+            db_type=db_type,
         )
 
     # Create test files
@@ -271,6 +272,12 @@ def cli(ctx: click.Context, verbose: bool) -> None:
 )
 @click.option("--host", "-H", type=str, help="Target host URL")
 @click.option("--auth/--no-auth", default=True, help="Include authentication in tests")
+@click.option(
+    "--db-type",
+    type=click.Choice(["", "mongo", "postgresql"], case_sensitive=False),
+    default="",
+    help="Database type for testing (empty for no database, mongo, or postgresql)",
+)
 @click.option("--dry-run", is_flag=True, help="Generate tests without running them")
 @click.option(
     "--custom-requirement", type=str, help="Custom requirements for test generation"
@@ -291,6 +298,7 @@ def generate(
     run_time: str,
     host: Optional[str],
     auth: bool,
+    db_type: str,
     dry_run: bool,
     custom_requirement: Optional[str],
     together_api_key: Optional[str],
@@ -309,6 +317,7 @@ def generate(
                 run_time,
                 host,
                 auth,
+                db_type,
                 dry_run,
                 custom_requirement,
                 together_api_key,
@@ -332,6 +341,7 @@ async def _async_generate(
     run_time: str,
     host: Optional[str],
     auth: bool,
+    db_type: str,
     dry_run: bool,
     custom_requirement: Optional[str],
     together_api_key: Optional[str],
@@ -343,7 +353,6 @@ async def _async_generate(
     try:
         _, api_key = _initialize_config(together_api_key)
         output_dir = _setup_output_directory(output)
-
         # Display configuration
         if ctx.obj["verbose"]:
             _display_configuration(
@@ -363,7 +372,14 @@ async def _async_generate(
         )
 
         created_files = await _generate_and_create_tests(
-            api_key, endpoints, api_info, output_dir, custom_requirement, host, auth
+            api_key,
+            endpoints,
+            api_info,
+            output_dir,
+            custom_requirement,
+            host,
+            auth,
+            db_type,
         )
 
         # Show results

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.7}/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.7"
 
     API_KEY: str = ""  # Fallback for backward compatibility
 

{devdox_ai_locust-0.1.2 → devdox_ai_locust-0.1.7}/src/devdox_ai_locust/hybrid_loctus_generator.py

@@ -19,12 +19,22 @@ 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__)
 
 
 test_data_file_path = "test_data.py"
+data_provider_path = "data_provider.py"
+
+
+@dataclass
+class ErrorClassification:
+    """Classification of an error for retry logic"""
+
+    is_retryable: bool
+    backoff_seconds: float
+    error_type: str
 
 
 @dataclass
@@ -109,6 +119,7 @@ class EnhancementProcessor:
         base_files: Dict[str, str],
         directory_files: List[Dict[str, Any]],
         grouped_endpoints: Dict[str, List[Endpoint]],
+        db_type: str = "",
     ) -> Tuple[List[Dict[str, Any]], List[str]]:
         """Process workflow enhancements"""
         enhanced_directory_files: List[Dict[str, Any]] = []
@@ -125,10 +136,13 @@ class EnhancementProcessor:
             first_workflow = base_workflow_files[0]
             # Get the content from the dictionary - adjust key name as needed
             base_workflow_content = first_workflow.get("base_workflow.py", "")
-
         for workflow_item in directory_files:
             enhanced_workflow_item = await self._enhance_single_workflow(
-                workflow_item, base_files, base_workflow_content, grouped_endpoints
+                workflow_item,
+                base_files,
+                base_workflow_content,
+                grouped_endpoints,
+                db_type,
             )
             if enhanced_workflow_item:
                 enhanced_directory_files.append(enhanced_workflow_item["files"])
@@ -142,6 +156,7 @@ class EnhancementProcessor:
         base_files: Dict[str, str],
         base_workflow_files: str,
         grouped_endpoints: Dict[str, List[Endpoint]],
+        db_type: str = "",
     ) -> Dict[str, Any] | None:
         """Enhance a single workflow file"""
         for key, value in workflow_item.items():
@@ -155,6 +170,7 @@ class EnhancementProcessor:
                 base_workflow=base_workflow_files,
                 grouped_enpoints=workflow_endpoints_dict,
                 auth_endpoints=auth_endpoints,
+                db_type=db_type,
             )
             if enhanced_workflow:
                 return {
@@ -165,14 +181,19 @@ class EnhancementProcessor:
         return None
 
     async def process_test_data_enhancement(
-        self, base_files: Dict[str, str], endpoints: List[Endpoint]
+        self, base_files: Dict[str, str], endpoints: List[Endpoint], db_type: str = ""
     ) -> Tuple[Dict[str, str], List[str]]:
         """Process test data enhancement"""
         enhanced_files = {}
         enhancements = []
         if self.ai_config and self.ai_config.enhance_test_data:
             enhanced_test_data = await self.locust_generator.enhance_test_data_file(
-                base_files.get(test_data_file_path, ""), endpoints
+                base_files.get(test_data_file_path, ""),
+                endpoints,
+                db_type,
+                base_files.get(data_provider_path, ""),
+                base_files.get("db_config.py", ""),
+                data_provider_path,
             )
             if enhanced_test_data:
                 enhanced_files[test_data_file_path] = enhanced_test_data
@@ -202,7 +223,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 +232,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 +255,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],
@@ -236,6 +301,7 @@ class HybridLocustGenerator:
         custom_requirement: Optional[str] = None,
         target_host: Optional[str] = None,
         include_auth: bool = True,
+        db_type: str = "",
     ) -> Tuple[Dict[str, str], List[Dict[str, Any]]]:
         """
         Generate Locust tests using hybrid approach
@@ -255,6 +321,7 @@ class HybridLocustGenerator:
                     api_info,
                     include_auth=include_auth,
                     target_host=target_host,
+                    db_type=db_type,
                 )
             )
 
@@ -269,6 +336,7 @@ class HybridLocustGenerator:
                     directory_files,
                     grouped_enpoints,
                     custom_requirement,
+                    db_type,
                 )
                 if enhancement_result.success:
                     logger.info(
@@ -369,6 +437,7 @@ class HybridLocustGenerator:
         directory_files: List[Dict[str, Any]],
         grouped_endpoints: Dict[str, List[Endpoint]],
         custom_requirement: Optional[str] = None,
+        db_type: str = "",
     ) -> EnhancementResult:
         """Enhance base files with AI - Refactored for reduced cognitive complexity"""
         start_time = asyncio.get_event_loop().time()
@@ -381,6 +450,7 @@ class HybridLocustGenerator:
                 directory_files,
                 grouped_endpoints,
                 custom_requirement,
+                db_type,
             )
 
             processing_time = asyncio.get_event_loop().time() - start_time
@@ -409,6 +479,7 @@ class HybridLocustGenerator:
         directory_files: List[Dict[str, Any]],
         grouped_endpoints: Dict[str, List[Endpoint]],
         custom_requirement: Optional[str] = None,
+        db_type: str = "",
     ) -> EnhancementResult:
         """Process all enhancements using the enhancement processor"""
         processor = EnhancementProcessor(self.ai_config, self)
@@ -423,7 +494,7 @@ class HybridLocustGenerator:
             processor.process_domain_flows_enhancement(
                 endpoints, api_info, custom_requirement
             ),
-            processor.process_test_data_enhancement(base_files, endpoints),
+            processor.process_test_data_enhancement(base_files, endpoints, db_type),
             processor.process_validation_enhancement(base_files, endpoints),
         ]
 
@@ -448,7 +519,7 @@ class HybridLocustGenerator:
                 workflow_files,
                 workflow_enhancements,
             ) = await processor.process_workflow_enhancements(
-                base_files, directory_files, grouped_endpoints
+                base_files, directory_files, grouped_endpoints, db_type
             )
             enhanced_directory_files.extend(workflow_files)
             enhancements_applied.extend(workflow_enhancements)
@@ -504,6 +575,7 @@ class HybridLocustGenerator:
         base_workflow: str,
         grouped_enpoints: Dict[str, List[Endpoint]],
         auth_endpoints: List[Endpoint],
+        db_type: str = "",
     ) -> Optional[str]:
         try:
             template = self.jinja_env.get_template("workflow.j2")
@@ -515,6 +587,7 @@ class HybridLocustGenerator:
                 base_workflow=base_workflow,
                 auth_endpoints=auth_endpoints,
                 base_content=base_content,
+                db_type=db_type,
             )
             enhanced_content = await self._call_ai_service(prompt)
             return enhanced_content
@@ -524,7 +597,13 @@ class HybridLocustGenerator:
         return ""
 
     async def enhance_test_data_file(
-        self, base_content: str, endpoints: List[Endpoint]
+        self,
+        base_content: str,
+        endpoints: List[Endpoint],
+        db_type: str = "",
+        data_provider: str = "",
+        db_config: str = "",
+        data_provider_path: str = "",
     ) -> Optional[str]:
         """Enhance test data generation with domain knowledge"""
 
@@ -539,11 +618,14 @@ class HybridLocustGenerator:
                 "base_content": base_content,
                 "schemas_info": schemas_info,
                 "endpoints": endpoints,
+                "db_type": db_type,
+                "data_provider_content": data_provider,
+                "db_config": db_config,
+                "data_provider_path": data_provider_path,
             }
 
             # Render enhanced content
             prompt = template.render(**context)
-
             enhanced_content = await self._call_ai_service(prompt)
             if enhanced_content and self._validate_python_code(enhanced_content):
                 return enhanced_content
@@ -573,10 +655,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 +666,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 +709,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)
 
-            if attempt < 2:  # Wait before retry
+                    continue
+
+            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"""