Fix integrated system errors and improve functionality

- Added missing _identify_challenges method to EnhancedComplexShapeGenerator
- Added missing _generate_enhanced_plan method for enhanced plan generation
- Added generate_response method to LLMClient for general text generation
- Added store_state method to StateCache as alias for cache_state
- Fixed JSON serialization issues with enum values in DeepSeek client
- Fixed QualityMetrics constructor to use confidence_level instead of confidence
- Updated predict_overall_quality to handle dict plans from enhanced generation
- Improved error handling and type compatibility across components

These fixes enable the integrated system to successfully generate complex
FreeCAD parts using DeepSeek R1 with proper error handling and fallbacks.

Author: Vansh5632

src/ai_designer/core/enhanced_complex_generator.py

@@ -277,10 +277,34 @@ def predict_step_quality(
         return prediction
 
     def predict_overall_quality(
-        self, plan: List[EnhancedGenerationStep], execution_context: Dict[str, Any]
+        self,
+        plan: Union[List[EnhancedGenerationStep], Dict[str, Any]],
+        execution_context: Dict[str, Any],
     ) -> QualityMetrics:
         """Predict overall quality for the entire generation plan"""
 
+        # Handle different plan formats
+        if isinstance(plan, dict):
+            # Extract execution steps from enhanced plan
+            execution_steps = plan.get("execution_steps", [])
+            complexity_score = plan.get("complexity_score", 5.0)
+
+            # Create simple quality prediction based on plan structure
+            base_quality = max(0.1, 1.0 - (complexity_score / 10.0))
+            step_count_factor = max(0.1, 1.0 - (len(execution_steps) / 20.0))
+            predicted_quality = base_quality * step_count_factor
+
+            # Return QualityMetrics object
+            return QualityMetrics(
+                overall_score=predicted_quality,
+                geometric_accuracy=predicted_quality * 0.9,
+                design_consistency=predicted_quality * 0.95,
+                aesthetic_quality=predicted_quality * 0.8,
+                manufacturability=predicted_quality * 0.85,
+                performance_score=predicted_quality * 0.9,
+                confidence_level=0.7,
+            )
+
         step_predictions = [
             self.predict_step_quality(step, execution_context) for step in plan
         ]
@@ -708,6 +732,63 @@ def _analyze_requirements_intelligently(self, requirements: str) -> Dict[str, An
             "confidence": 0.8,
         }
 
+    def _identify_challenges(
+        self,
+        requirements: str,
+        entities: List[Dict[str, Any]],
+        relationships: List[Dict[str, Any]],
+    ) -> List[Dict[str, Any]]:
+        """Identify potential challenges in the design requirements"""
+        challenges = []
+
+        # Check for geometric complexity
+        if len(entities) > 5:
+            challenges.append(
+                {
+                    "type": "geometric_complexity",
+                    "severity": "medium",
+                    "description": f"Design involves {len(entities)} entities which may require careful coordination",
+                    "suggestion": "Consider breaking down into sub-assemblies",
+                }
+            )
+
+        # Check for precision requirements
+        precision_keywords = ["precise", "accurate", "tolerance", "fit", "clearance"]
+        if any(keyword in requirements.lower() for keyword in precision_keywords):
+            challenges.append(
+                {
+                    "type": "precision_requirements",
+                    "severity": "high",
+                    "description": "Design requires high precision which may affect manufacturability",
+                    "suggestion": "Define clear tolerances and consider manufacturing constraints",
+                }
+            )
+
+        # Check for complex relationships
+        if len(relationships) > len(entities):
+            challenges.append(
+                {
+                    "type": "complex_relationships",
+                    "severity": "medium",
+                    "description": "Multiple relationships between entities may complicate assembly",
+                    "suggestion": "Create detailed assembly sequence and constraints",
+                }
+            )
+
+        # Check for material considerations
+        material_keywords = ["steel", "aluminum", "plastic", "composite", "material"]
+        if any(keyword in requirements.lower() for keyword in material_keywords):
+            challenges.append(
+                {
+                    "type": "material_considerations",
+                    "severity": "low",
+                    "description": "Specific materials mentioned may affect design decisions",
+                    "suggestion": "Consider material properties in design validation",
+                }
+            )
+
+        return challenges
+
     def _estimate_resources(
         self, complexity_score: float, entity_count: int
     ) -> Dict[str, Any]:
@@ -862,6 +943,88 @@ def _select_optimal_mode(
         else:
             return GenerationMode.ADAPTIVE
 
+    def _generate_enhanced_plan(
+        self, requirements: str, analysis: Dict[str, Any], patterns: List[Dict] = None
+    ) -> Dict[str, Any]:
+        """Generate an enhanced execution plan based on requirements and analysis"""
+
+        entities = analysis.get("entities", [])
+        relationships = analysis.get("relationships", [])
+        complexity_score = analysis.get("complexity_score", 3.0)
+        challenges = analysis.get("challenges", [])
+
+        # Create execution steps based on entities and relationships
+        execution_steps = []
+
+        # Step 1: Initialize document
+        execution_steps.append(
+            {
+                "step_id": 1,
+                "type": "initialization",
+                "description": "Initialize FreeCAD document and import modules",
+                "code_template": "doc = App.newDocument('{}')".format("GearAssembly"),
+                "estimated_time": 2,
+            }
+        )
+
+        # Step 2: Create primary entities
+        step_id = 2
+        for entity in entities:
+            execution_steps.append(
+                {
+                    "step_id": step_id,
+                    "type": "entity_creation",
+                    "description": f"Create {entity['name']} ({entity['type']})",
+                    "entity": entity,
+                    "estimated_time": entity["complexity"] * 2,
+                }
+            )
+            step_id += 1
+
+        # Step 3: Apply relationships
+        for relationship in relationships:
+            execution_steps.append(
+                {
+                    "step_id": step_id,
+                    "type": "relationship",
+                    "description": f"Apply {relationship['type']} relationship",
+                    "relationship": relationship,
+                    "estimated_time": relationship["complexity"] * 3,
+                }
+            )
+            step_id += 1
+
+        # Step 4: Finalization
+        execution_steps.append(
+            {
+                "step_id": step_id,
+                "type": "finalization",
+                "description": "Finalize document and recompute",
+                "code_template": "doc.recompute()",
+                "estimated_time": 1,
+            }
+        )
+
+        # Create enhanced plan
+        enhanced_plan = {
+            "requirements": requirements,
+            "analysis": analysis,
+            "execution_steps": execution_steps,
+            "estimated_total_time": sum(
+                step.get("estimated_time", 2) for step in execution_steps
+            ),
+            "complexity_score": complexity_score,
+            "challenges": challenges,
+            "patterns_used": len(patterns) if patterns else 0,
+            "optimization_suggestions": [
+                "Use parametric modeling for flexibility",
+                "Consider manufacturing constraints",
+                "Validate geometric relationships",
+            ],
+        }
+
+        return enhanced_plan
+
     def _send_generation_complete(self, session_id: str, result: GenerationResult):
         """Send generation completion notification via WebSocket"""
         if self.websocket_manager:
@@ -1026,7 +1189,7 @@ def _generate_with_deepseek_r1(
         # Prepare enhanced context for DeepSeek
         deepseek_context = {
             "session_id": session_context.get("session_id"),
-            "generation_mode": session_context.get("mode", "adaptive"),
+            "generation_mode": str(session_context.get("mode", "adaptive")),
             "quality_targets": session_context.get("quality_targets", {}),
             "freecad_state": session_context.get("context", {}),
             "complexity_analysis": complexity_analysis,

src/ai_designer/llm/client.py

@@ -163,3 +163,41 @@ def generate_command(self, nl_command, state=None):
         except Exception as e:
             print(f"[LLMClient] Error generating command: {e}")
             raise
+
+    def generate_response(self, prompt: str, context: Optional[str] = None) -> str:
+        """
+        Generate a response using the LLM for general text generation.
+
+        Args:
+            prompt: The text prompt to generate a response for
+            context: Optional context to include in the prompt
+
+        Returns:
+            Generated text response
+        """
+        try:
+            from langchain.prompts import ChatPromptTemplate
+
+            if context:
+                full_prompt = f"Context: {context}\n\nPrompt: {prompt}"
+            else:
+                full_prompt = prompt
+
+            prompt_template = ChatPromptTemplate.from_messages([("human", full_prompt)])
+
+            messages = prompt_template.format_messages()
+            response = self.llm.invoke(messages)
+
+            # Extract content from response
+            if hasattr(response, "content"):
+                return response.content.strip()
+            elif hasattr(response, "text") and callable(getattr(response, "text")):
+                return response.text().strip()
+            elif hasattr(response, "text"):
+                return response.text.strip()
+            else:
+                return str(response).strip()
+
+        except Exception as e:
+            logger.error(f"Error generating response: {e}")
+            raise

src/ai_designer/llm/deepseek_client.py

@@ -236,10 +236,36 @@ def _build_complex_part_prompt(
             )
 
         if context:
-            base_prompt += f"\n\nCONTEXT:\n{json.dumps(context, indent=2)}"
+            try:
+                # Convert any enum values to strings for JSON serialization
+                serializable_context = {}
+                for key, value in context.items():
+                    if hasattr(value, "__dict__"):
+                        serializable_context[key] = str(value)
+                    elif isinstance(value, Enum):
+                        serializable_context[key] = value.value
+                    else:
+                        serializable_context[key] = value
+                base_prompt += (
+                    f"\n\nCONTEXT:\n{json.dumps(serializable_context, indent=2)}"
+                )
+            except (TypeError, ValueError) as e:
+                base_prompt += f"\n\nCONTEXT:\n{str(context)}"
 
         if constraints:
-            base_prompt += f"\n\nCONSTRAINTS:\n{json.dumps(constraints, indent=2)}"
+            try:
+                # Convert any enum values to strings for JSON serialization
+                serializable_constraints = {}
+                for key, value in constraints.items():
+                    if hasattr(value, "__dict__"):
+                        serializable_constraints[key] = str(value)
+                    elif isinstance(value, Enum):
+                        serializable_constraints[key] = value.value
+                    else:
+                        serializable_constraints[key] = value
+                base_prompt += f"\n\nCONSTRAINTS:\n{json.dumps(serializable_constraints, indent=2)}"
+            except (TypeError, ValueError) as e:
+                base_prompt += f"\n\nCONSTRAINTS:\n{str(constraints)}"
 
         base_prompt += """
 

src/ai_designer/redis_utils/state_cache.py

@@ -278,3 +278,13 @@ def _extract_document_names(self, keys: List[str]) -> List[str]:
     def update_state(self, state_data: Dict[str, Any], state_key: str = None) -> str:
         """Update state data in Redis (alias for cache_state)"""
         return self.cache_state(state_data, state_key)
+
+    def store_state(
+        self, state_key: str, state_data: Dict[str, Any], expiration: int = None
+    ) -> bool:
+        """Store state data in Redis (alias for cache_state)"""
+        try:
+            self.cache_state(state_data, state_key, expiration=expiration)
+            return True
+        except Exception:
+            return False