crewai_agents.py

from crewai import Agent, Task, Crew, Process
from typing import Dict, Any, List
import json
import os
from llm_client import LLMClient
from soil_analyzer import SoilLayerAnalyzer
from pydantic import BaseModel, Field
from config import LLM_PROVIDERS, get_default_provider_and_model

class CrewAIGeotechSystem:
    def __init__(self, model=None, api_key=None):
        # Handle API key - if explicitly passed as empty string, use that (for mock mode)
        if api_key == "":
            self.api_key = ""
        else:
            self.api_key = api_key or ""
        
        _, default_model = get_default_provider_and_model()
        self.model = model or default_model
        
        # Check if we have a valid API key
        self.has_api_key = bool(self.api_key and self.api_key.strip())
        
        if self.has_api_key:
            # Initialize our working LLMClient for actual LLM calls
            self.llm_client = LLMClient(model=self.model, api_key=self.api_key)
            
            # We'll use direct LLM calls instead of CrewAI agents due to compatibility issues
            self.geotech_agent = "Geotechnical Engineer (Direct LLM)"
            self.senior_geotech_agent = "Senior Geotechnical Engineer (Direct LLM)"
        else:
            # No API key available - set to None to trigger mock mode
            self.llm_client = None
            self.geotech_agent = None
            self.senior_geotech_agent = None
    
    def _run_geotech_analysis_direct(self, soil_data: Dict[str, Any]) -> str:
        """Run geotechnical analysis using direct LLM calls"""
        
        analysis_prompt = f"""
        As an experienced geotechnical engineer, analyze the following soil boring log data and provide a comprehensive geotechnical assessment:

        DATA:
        {json.dumps(soil_data, indent=2)}

        CRITICAL UNIT CONVERSION REQUIREMENTS:
        1. **Su (Undrained Shear Strength) Unit Conversions:**
           - t/m² → kPa: MULTIPLY BY 9.81 (NOT 10!)
           - ksc (kg/cm²) → kPa: multiply by 98.0
           - psi → kPa: multiply by 6.895
           - MPa → kPa: multiply by 1000
           - tsf (tons/ft²) → kPa: multiply by 95.76

        2. **Common Unit Errors to Check:**
           - If Su values seem unusually high (>500 kPa for soft clay), check if units are incorrectly converted
           - If Su values seem unusually low (<10 kPa for stiff clay), check if conversion factor was applied
           - t/m² is commonly misconverted using factor of 10 instead of 9.81 - BE CAREFUL!

        CRITICAL LAYER ANALYSIS REQUIREMENTS:
        3. **Su Value Consistency Within Layers:**
           - **EXAMINE each layer for Su value variations**
           - If multiple Su values exist within a single layer, check for consistency
           - **LAYER SPLITTING CRITERIA:**
             * If Su values within a layer vary by >30%, consider splitting the layer
             * If Su values show clear trend (increasing/decreasing), split at transition points
             * If one Su value is >2x another Su value in same layer, MUST split the layer
           
        4. **Layer Splitting Protocol:**
           - Identify depth ranges where Su values are similar (within 20-30% variation)
           - Create new layer boundaries at points where Su values change significantly
           - Each new sub-layer should have consistent Su values (average or representative value)
           - Update soil descriptions to reflect the new layer characteristics

        5. **Su Value Assignment for Layers:**
           - If Su values are consistent within layer: use average value
           - If Su values vary significantly: split layer and assign representative values
           - Document the original Su readings and how they were processed

        Your responsibilities:
        1. **CAREFULLY validate ALL unit conversions** - pay special attention to Su values
        2. Check if any Su values are in t/m² and need conversion to kPa using factor 9.81
        3. **ANALYZE Su value consistency within each layer**
        4. **SPLIT layers where Su values vary significantly (>30% variation or >2x difference)**
        5. Validate all geotechnical parameters for consistency and reasonableness
        6. Check layer classifications and transitions
        7. Verify strength parameter correlations (Su vs water content, SPT vs consistency)
        8. Calculate layer statistics and identify any outliers
        9. Ensure ALL unit conversions are correct (kPa, degrees, etc.)
        10. Check depth continuity and layer boundaries

        Focus on:
        - **UNIT CONVERSION ACCURACY** (especially Su values from t/m² to kPa)
        - **Su VALUE CONSISTENCY within layers** - split if values vary significantly
        - Parameter consistency (e.g., soft clay should have low Su, high water content)
        - Reasonable strength ranges for each soil type
        - Proper calculation methodology
        - Layer transition logic

        LAYER SPLITTING DECISION TREE:
        ✓ Check if multiple Su values exist in each layer
        ✓ Calculate variation between Su values (max-min)/average
        ✓ If variation >30% OR max/min ratio >2.0: SPLIT the layer
        ✓ Create new layers with consistent Su values
        ✓ Assign average Su value to each new sub-layer
        ✓ Update layer descriptions and boundaries

        UNIT CONVERSION VERIFICATION CHECKLIST:
        ✓ Check if any Su values are in t/m² - if yes, multiply by 9.81 to get kPa
        ✓ Verify Su values are reasonable for soil consistency (soft clay: 10-30 kPa, stiff clay: 50-100 kPa)
        ✓ Check SPT N-values are reasonable for soil description
        ✓ Ensure depth units are in meters

        Provide detailed analysis with any concerns noted, especially unit conversion issues and layer splitting recommendations.
        """
        
        try:
            # Use our working LLMClient
            response = self.llm_client.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are an experienced geotechnical engineer with expertise in soil mechanics and foundation design. You are particularly careful about unit conversions and have seen many errors caused by incorrect unit conversion factors."},
                    {"role": "user", "content": analysis_prompt}
                ],
                temperature=0.1,
                max_tokens=2000
            )
            
            return response.choices[0].message.content
                
        except Exception as e:
            return f"Analysis failed: {str(e)}"
    
    def _run_senior_review_direct(self, analysis_result: str) -> str:
        """Run senior engineer review using direct LLM calls"""
        
        review_prompt = f"""
        As a senior geotechnical engineer with 20+ years of experience, review the following geotechnical analysis for consistency, accuracy, and engineering reasonableness:

        ANALYSIS TO REVIEW:
        {analysis_result}

        CRITICAL REVIEW FOCUS - UNIT CONVERSIONS:
        Your PRIMARY responsibility is to catch unit conversion errors that can lead to catastrophic design failures.

        1. **Su (Undrained Shear Strength) CRITICAL CHECKS:**
           - Are any Su values still in t/m²? If yes, they MUST be converted to kPa using factor 9.81
           - Are Su values reasonable for the soil consistency described?
           - Soft clay: 10-30 kPa | Medium clay: 30-60 kPa | Stiff clay: 60-120 kPa | Very stiff clay: >120 kPa
           - If Su > 500 kPa for soft clay → MAJOR RED FLAG - likely unit error
           - If Su < 10 kPa for stiff clay → MAJOR RED FLAG - likely unit error

        2. **COMMON UNIT CONVERSION ERRORS TO IDENTIFY:**
           - Using factor 10 instead of 9.81 for t/m² → kPa conversion
           - Missing unit conversions (values still in original units)
           - Wrong conversion factors applied

        CRITICAL LAYER SPLITTING REVIEW:
        3. **Su Value Consistency Within Layers:**
           - **EXAMINE if layers with varying Su values were properly split**
           - Check if any single layer contains Su values that vary by >30%
           - **LAYER SPLITTING VALIDATION:**
             * Were layers with Su variation >30% properly split?
             * Were layers with Su ratio >2.0 (max/min) properly split?
             * Are the new layer boundaries logical and well-defined?
             * Does each sub-layer have consistent Su values?

        4. **Layer Splitting Quality Control:**
           - Verify that split layers have representative Su values (average or appropriate)
           - Check that layer descriptions match the Su values assigned
           - Ensure depth boundaries are clearly defined for split layers
           - Validate that soil consistency matches the Su values in each sub-layer

        Your senior engineer review responsibilities:

        1. **UNIT CONVERSION VALIDATION (HIGHEST PRIORITY):**
           - Su vs water content relationships for clay soils
           - SPT N-values vs soil consistency correlations  
           - Strength ranges within expected bounds for each soil type
           - Unit conversion accuracy (kPa, degrees, m) - ESPECIALLY t/m² to kPa using 9.81

        2. **LAYER SPLITTING VALIDATION (HIGH PRIORITY):**
           - Check if layers with varying Su values were appropriately split
           - Verify consistency of Su values within each layer
           - Validate that layer boundaries make engineering sense

        3. **PARAMETER CONSISTENCY CHECKS:**
           - Layer boundaries and transitions are logical
           - Classification consistency across depth
           - Parameter ranges match soil descriptions

        4. **RED FLAGS TO IDENTIFY:**
           - **CRITICAL:** Su values in wrong units (t/m² not converted to kPa)
           - **CRITICAL:** Su values unreasonable for consistency (high Su with soft clay, low Su with stiff clay)
           - **CRITICAL:** Single layer with highly variable Su values (>30% variation) not split
           - High water content with very high Su (unusual for clay)
           - Low water content with very low Su
           - Soft consistency with high SPT N-values
           - Hard consistency with low SPT N-values
           - Strength values that appear incorrectly converted

        5. **DECISION CRITERIA:**
           - If you find ANY unit conversion errors: **REJECT and require re-investigation**
           - If Su values are inconsistent with soil consistency: **REJECT and require re-investigation**
           - If layers with varying Su values were not properly split: **REJECT and require re-investigation**
           - If parameters look reasonable: Approve with confidence assessment
           - If minor concerns exist: Approve with notes

        LAYER SPLITTING REVIEW CHECKLIST:
        ✓ Are there any layers with multiple different Su values?
        ✓ Were layers with Su variation >30% properly split?
        ✓ Were layers with Su ratio >2.0 properly split?
        ✓ Do split layers have consistent Su values within each sub-layer?
        ✓ Are layer descriptions updated to match the split layers?

        UNIT CONVERSION VERIFICATION CHECKLIST FOR REVIEW:
        ✓ Are ALL Su values properly converted to kPa?
        ✓ Are Su values reasonable for the described soil consistency?
        ✓ Has the correct factor (9.81) been used for t/m² → kPa conversion?
        ✓ Are SPT N-values consistent with soil descriptions?

        Provide your professional judgment on whether this analysis is acceptable or requires revision.
        Be EXTREMELY specific about any unit conversion issues and layer splitting issues found and provide clear guidance for correction.
        
        REMEMBER: Unit conversion errors and improper layer definition can lead to foundation failures. Be thorough.
        """
        
        try:
            # Use our working LLMClient
            response = self.llm_client.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are a senior geotechnical engineer with extensive experience in complex foundation projects and rigorous quality control. You have seen foundation failures caused by unit conversion errors and are extremely vigilant about this issue."},
                    {"role": "user", "content": review_prompt}
                ],
                temperature=0.1,
                max_tokens=2000
            )
            
            return response.choices[0].message.content
            
        except Exception as e:
            return f"Review failed: {str(e)}"
    
    def _run_reinvestigation_direct(self, original_analysis: str, review_feedback: str) -> str:
        """Run re-investigation using direct LLM calls"""
        
        reinvestigation_prompt = f"""
        Based on the senior engineer's review, re-investigate and address the following issues:

        ORIGINAL ANALYSIS: 
        {original_analysis}

        SENIOR REVIEW FEEDBACK: 
        {review_feedback}

        RE-INVESTIGATION REQUIREMENTS:
        
        **PRIORITY 1: UNIT CONVERSION CORRECTIONS**
        If the senior engineer identified unit conversion issues:
        1. **Su (Undrained Shear Strength) Corrections:**
           - If values are in t/m², convert to kPa by multiplying by 9.81
           - Double-check ALL conversion factors used
           - Verify final Su values are reasonable for soil consistency
           
        2. **Verify Conversion Factors:**
           - t/m² → kPa: multiply by 9.81 (NOT 10)
           - ksc → kPa: multiply by 98.0
           - psi → kPa: multiply by 6.895
           - MPa → kPa: multiply by 1000

        **PRIORITY 2: LAYER SPLITTING CORRECTIONS**
        If the senior engineer identified layer splitting issues:
        1. **Su Value Variation Analysis:**
           - Re-examine each layer for Su value consistency
           - Calculate variation: (max Su - min Su) / average Su
           - Calculate ratio: max Su / min Su
           
        2. **Layer Splitting Protocol:**
           - If Su variation >30% OR ratio >2.0: SPLIT the layer
           - Create new layer boundaries at points where Su values change significantly
           - Assign consistent Su values to each new sub-layer (use average within range)
           - Update layer descriptions to reflect new boundaries
           
        3. **New Layer Definition:**
           - Each sub-layer should have Su values within 20-30% variation
           - Update depth ranges for each new sub-layer
           - Ensure soil consistency descriptions match Su values
           - Verify layer transitions are logical

        **GENERAL RE-INVESTIGATION:**
        1. Address each specific concern raised by the senior engineer
        2. Re-examine parameter relationships and correlations
        3. Double-check ALL unit conversions and calculations
        4. Provide revised analysis with explanations for changes
        5. Justify any assumptions or interpretations made

        **LAYER SPLITTING EXAMPLE:**
        Original Layer: 2.0-6.0m, Clay, Su values: 15, 25, 45 kPa
        Analysis: Variation = (45-15)/28 = 107% > 30%, Ratio = 45/15 = 3.0 > 2.0
        Action: SPLIT INTO:
        - Layer 2a: 2.0-4.0m, Soft Clay, Su = 20 kPa (average of 15, 25)
        - Layer 2b: 4.0-6.0m, Medium Clay, Su = 45 kPa

        **UNIT CONVERSION RE-CHECK PROTOCOL:**
        ✓ Identify original units for each parameter
        ✓ Apply correct conversion factors
        ✓ Verify converted values are reasonable
        ✓ Check consistency with soil descriptions

        **LAYER SPLITTING RE-CHECK PROTOCOL:**
        ✓ Check Su value variation within each layer
        ✓ Split layers with >30% variation or >2.0 ratio
        ✓ Assign representative Su values to sub-layers
        ✓ Update layer descriptions and boundaries
        ✓ Verify consistency between Su and soil consistency
        
        Focus specifically on the issues identified in the review.
        Provide a comprehensive revised analysis that addresses ALL concerns.
        
        **Show your work:** For any corrections, clearly state:
        - Original layer configuration
        - Su values and their variation/ratio
        - Splitting decision and rationale
        - New layer boundaries and Su assignments
        - Unit conversion details (original value, factor, final value)
        - Verification that results are reasonable
        """
        
        try:
            # Use our working LLMClient
            response = self.llm_client.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are an experienced geotechnical engineer conducting a thorough re-investigation based on senior engineer feedback. You are particularly focused on correcting any unit conversion errors identified."},
                    {"role": "user", "content": reinvestigation_prompt}
                ],
                temperature=0.1,
                max_tokens=2000
            )
            
            return response.choices[0].message.content
            
        except Exception as e:
            return f"Re-investigation failed: {str(e)}"
    
    def _run_final_review_direct(self, revised_analysis: str, previous_concerns: str) -> str:
        """Run final review using direct LLM calls"""
        
        final_review_prompt = f"""
        Conduct final review of the re-investigated analysis:

        REVISED ANALYSIS:
        {revised_analysis}

        PREVIOUS CONCERNS:
        {previous_concerns}

        Final validation requirements:
        1. Confirm all previous concerns have been adequately addressed
        2. Verify parameter consistency and engineering reasonableness
        3. Check that explanations are technically sound
        4. Provide final approval or additional guidance if still needed

        Make final determination: APPROVED or REQUIRES FURTHER WORK
        """
        
        try:
            # Use our working LLMClient
            response = self.llm_client.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are a senior geotechnical engineer conducting final validation with authority to approve or reject the analysis."},
                    {"role": "user", "content": final_review_prompt}
                ],
                temperature=0.1,
                max_tokens=2000
            )
            
            return response.choices[0].message.content
            
        except Exception as e:
            return f"Final review failed: {str(e)}"
    
    def run_geotechnical_analysis(self, soil_data: Dict[str, Any]) -> Dict[str, Any]:
        """Run the complete two-agent geotechnical analysis workflow using direct LLM calls"""
        
        # Handle case where no LLM is available (testing mode)
        if not self.has_api_key:
            return self._mock_analysis_for_testing(soil_data)
        
        try:
            # Run initial analysis using direct LLM call
            analysis_result = self._run_geotech_analysis_direct(soil_data)
            
            if "failed:" in analysis_result:
                return {
                    "error": f"Initial analysis failed: {analysis_result}",
                    "status": "error",
                    "workflow": "failed"
                }
            
            # Run review using direct LLM call
            review_result = self._run_senior_review_direct(analysis_result)
            
            if "failed:" in review_result:
                return {
                    "error": f"Review failed: {review_result}",
                    "status": "error",
                    "workflow": "failed"
                }
            
            # Check if re-investigation is needed based on review content
            review_text = review_result.lower()
            needs_reinvestigation = any(keyword in review_text for keyword in [
                "re-investigate", "reject", "inconsistent", "unusual", "verify", "additional testing",
                "requires revision", "not acceptable", "re-examination", "concerning", "requires further work"
            ])
            
            if needs_reinvestigation:
                # Run re-investigation
                final_analysis = self._run_reinvestigation_direct(analysis_result, review_result)
                
                if "failed:" in final_analysis:
                    return {
                        "error": f"Re-investigation failed: {final_analysis}",
                        "status": "error",
                        "workflow": "failed"
                    }
                
                # Final review of re-investigation
                final_review = self._run_final_review_direct(final_analysis, review_result)
                
                if "failed:" in final_review:
                    return {
                        "error": f"Final review failed: {final_review}",
                        "status": "error",
                        "workflow": "failed"
                    }
                
                return {
                    "initial_analysis": analysis_result,
                    "initial_review": review_result,
                    "reinvestigation": final_analysis,
                    "final_review": final_review,
                    "status": "completed_with_revision",
                    "workflow": "two_stage_review_direct_llm"
                }
            
            else:
                return {
                    "analysis": analysis_result,
                    "review": review_result,
                    "status": "approved",
                    "workflow": "single_stage_review_direct_llm"
                }
        
        except Exception as e:
            return {
                "error": f"CrewAI analysis failed: {str(e)}",
                "status": "error",
                "workflow": "failed"
            }
    
    def _mock_analysis_for_testing(self, soil_data: Dict[str, Any]) -> Dict[str, Any]:
        """Provides mock analysis for testing when no API key is available"""
        
        # Simulate comprehensive analysis based on input data
        project_name = soil_data.get('project_info', {}).get('project_name', 'Unknown Project')
        boring_id = soil_data.get('project_info', {}).get('boring_id', 'Unknown Boring')
        soil_layers = soil_data.get('soil_layers', [])
        
        # Generate realistic mock analysis
        mock_analysis = f"""
        GEOTECHNICAL ANALYSIS REPORT - {project_name} ({boring_id})
        
        EXECUTIVE SUMMARY:
        Analyzed {len(soil_layers)} soil layers from the boring log data.
        Overall soil conditions appear consistent with typical soil behavior patterns.
        
        LAYER ANALYSIS:
        """
        
        for i, layer in enumerate(soil_layers, 1):
            soil_type = layer.get('soil_type', 'Unknown')
            consistency = layer.get('consistency', 'Unknown')
            depth_from = layer.get('depth_from', 0)
            depth_to = layer.get('depth_to', 0)
            strength_value = layer.get('strength_value', 'N/A')
            strength_unit = layer.get('strength_unit', '')
            
            mock_analysis += f"""
        Layer {i} ({depth_from}-{depth_to}m): {soil_type.title()}
        - Consistency: {consistency}
        - Strength: {strength_value} {strength_unit}
        - Classification appears reasonable for {consistency} {soil_type}
        - Depth continuity validated ✓
            """
        
        mock_analysis += """
        
        VALIDATION CHECKS:
        ✓ Layer depth continuity confirmed
        ✓ Strength parameters within expected ranges
        ✓ Soil classification consistency verified
        ✓ Unit conversions validated
        
        RECOMMENDATIONS:
        - Data appears consistent and suitable for preliminary design
        - Standard geotechnical correlations apply
        - Consider additional testing for final design if required
        """
        
        mock_review = f"""
        SENIOR ENGINEER REVIEW - {project_name}
        
        I have reviewed the geotechnical analysis and findings:
        
        TECHNICAL REVIEW:
        - Analysis methodology is sound and follows standard practices
        - Parameter correlations are reasonable and well-documented  
        - Soil classification is consistent with strength parameters
        - Depth boundaries and layer transitions are appropriate
        
        QUALITY ASSURANCE:
        - All calculations have been verified
        - Unit conversions are correct
        - Data consistency checks passed
        - Engineering correlations within acceptable ranges
        
        APPROVAL STATUS: ✅ APPROVED
        
        The analysis meets professional standards and is suitable for use in geotechnical design.
        No re-investigation required at this time.
        
        Senior Geotechnical Engineer Review Complete.
        """
        
        return {
            "status": "approved",
            "workflow": "two_agent_review",
            "analysis": mock_analysis.strip(),
            "review": mock_review.strip(),
            "summary": f"Mock analysis completed for {len(soil_layers)} soil layers. All parameters validated.",
            "timestamp": "2024-06-26T10:00:00Z",
            "system": "CrewAI Mock System (No API Key)",
            "reinvestigation_required": False
            }

# Example usage function
def analyze_soil_with_crewai(soil_data: Dict[str, Any]) -> Dict[str, Any]:
    """Main function to run CrewAI-based geotechnical analysis"""
    system = CrewAIGeotechSystem()
    return system.run_geotechnical_analysis(soil_data)