agent/consultation.py

"""
Interactive Network Consultation
Multi-turn conversation to gather complete requirements
"""

import json
import logging
from typing import Dict, List, Tuple, Optional
from agent.llm_client import LLMClient, LLMMessage

logger = logging.getLogger(__name__)


CONSULTATION_SYSTEM_PROMPT = """You are an expert network consultant helping a client design their network infrastructure.

Your job is to:
1. Understand their business needs and technical requirements
2. Ask clarifying questions to fill in gaps
3. Probe for important details they may have forgotten (security, compliance, scalability, budget)
4. Extract structured information: devices needed, VLANs, subnets, bandwidth, redundancy needs

When you have enough information, respond with JSON in this format:
```json
{
  "consultation_complete": true,
  "network_intent": {
    "description": "full description",
    "locations": [...],
    "business_requirements": [...],
    "constraints": [...],
    "timeline": "...",
    "budget": "...",
    "vendor_preference": "...",
    "compliance_requirements": [...],
    "bandwidth_requirements": {...},
    "redundancy_requirements": {...}
  }
}
```

If you need more information, respond with:
```json
{
  "consultation_complete": false,
  "questions": ["question 1", "question 2", ...],
  "summary_so_far": "what we know so far"
}
```

Be professional, thorough, and ask smart questions that a real network consultant would ask.
"""


class NetworkConsultant:
    """
    Interactive consultation agent
    Gathers complete requirements through conversation
    """
    
    def __init__(self):
        self.llm = LLMClient()
        self.conversation_history: List[LLMMessage] = []
        self.intent_data: Optional[Dict] = None
    
    def start_consultation(self, initial_description: str) -> Tuple[bool, str, Optional[Dict]]:
        """
        Start consultation process
        Returns: (is_complete, next_questions_or_summary, intent_dict)
        """
        # Initialize conversation
        self.conversation_history = [
            LLMMessage(role="system", content=CONSULTATION_SYSTEM_PROMPT),
            LLMMessage(role="user", content=f"Initial request: {initial_description}")
        ]
        
        # Get LLM response
        response = self.llm.chat(self.conversation_history, temperature=0.3)
        
        # Parse response
        try:
            # Extract JSON from response
            json_start = response.find('{')
            json_end = response.rfind('}') + 1
            if json_start >= 0 and json_end > json_start:
                json_str = response[json_start:json_end]
                result = json.loads(json_str)
                
                if result.get("consultation_complete"):
                    self.intent_data = result.get("network_intent")
                    summary = self._format_intent_summary(self.intent_data)
                    return True, summary, self.intent_data
                else:
                    questions_text = self._format_questions(result.get("questions", []))
                    return False, questions_text, None
            else:
                # Fallback if no JSON
                return False, response, None
                
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse LLM response as JSON: {e}")
            return False, response, None
    
    def continue_consultation(self, user_response: str) -> Tuple[bool, str, Optional[Dict]]:
        """
        Continue multi-turn consultation
        Returns: (is_complete, next_questions_or_summary, intent_dict)
        """
        # Add user response to history
        self.conversation_history.append(
            LLMMessage(role="user", content=user_response)
        )
        
        # Get LLM response
        response = self.llm.chat(self.conversation_history, temperature=0.3)
        self.conversation_history.append(
            LLMMessage(role="assistant", content=response)
        )
        
        # Parse response
        try:
            json_start = response.find('{')
            json_end = response.rfind('}') + 1
            if json_start >= 0 and json_end > json_start:
                json_str = response[json_start:json_end]
                result = json.loads(json_str)
                
                if result.get("consultation_complete"):
                    self.intent_data = result.get("network_intent")
                    summary = self._format_intent_summary(self.intent_data)
                    return True, summary, self.intent_data
                else:
                    questions_text = self._format_questions(result.get("questions", []))
                    summary = result.get("summary_so_far", "")
                    output = f"**Progress Summary:**\n{summary}\n\n**Additional Questions:**\n{questions_text}"
                    return False, output, None
            else:
                return False, response, None
                
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse LLM response: {e}")
            return False, response, None
    
    def _format_questions(self, questions: List[str]) -> str:
        """Format questions as numbered list"""
        return "\n".join(f"{i+1}. {q}" for i, q in enumerate(questions))
    
    def _format_intent_summary(self, intent: Dict) -> str:
        """Format final intent as readable summary"""
        lines = ["## 📋 Consultation Complete!\n"]
        
        lines.append(f"**Description:** {intent.get('description', 'N/A')}\n")
        
        if intent.get('locations'):
            lines.append(f"**Locations:** {len(intent['locations'])} sites")
            for loc in intent['locations']:
                lines.append(f"  - {loc}")
            lines.append("")
        
        if intent.get('business_requirements'):
            lines.append("**Business Requirements:**")
            for req in intent['business_requirements']:
                lines.append(f"  - {req}")
            lines.append("")
        
        if intent.get('budget'):
            lines.append(f"**Budget:** {intent['budget']}\n")
        
        if intent.get('timeline'):
            lines.append(f"**Timeline:** {intent['timeline']}\n")
        
        if intent.get('vendor_preference'):
            lines.append(f"**Preferred Vendors:** {intent['vendor_preference']}\n")
        
        if intent.get('compliance_requirements'):
            lines.append("**Compliance:**")
            for req in intent['compliance_requirements']:
                lines.append(f"  - {req}")
            lines.append("")
        
        return "\n".join(lines)