survey_generator.py

"""
Survey Generation Module - Generate AI-powered surveys from outlines
"""
import json
from typing import List, Dict, Optional
from llm_backend import LLMBackend


class SurveyGenerator:
    """
    Generates professional surveys from user outlines using AI.
    Follows industry best practices for qualitative research.
    """

    def __init__(self, llm_backend: LLMBackend):
        self.llm = llm_backend

    def generate_survey(self,
                       outline: str,
                       survey_type: str = "qualitative",
                       num_questions: int = 10,
                       target_audience: str = "general") -> Dict:
        """
        Generate a complete survey from an outline.

        Args:
            outline: User's outline or topic description
            survey_type: Type of survey (qualitative, quantitative, mixed)
            num_questions: Target number of questions
            target_audience: Description of target respondents

        Returns:
            Dict containing survey metadata and questions
        """
        prompt = self._build_generation_prompt(outline, survey_type, num_questions, target_audience)

        messages = [
            {"role": "system", "content": self._get_system_prompt()},
            {"role": "user", "content": prompt}
        ]

        try:
            response = self.llm.generate(messages, max_tokens=2000, temperature=0.7)
            survey_data = self._parse_survey_response(response)

            # Add metadata
            survey_data["metadata"] = {
                "outline": outline,
                "survey_type": survey_type,
                "target_audience": target_audience,
                "generated_question_count": len(survey_data.get("questions", []))
            }

            return survey_data

        except Exception as e:
            raise Exception(f"Survey generation failed: {str(e)}")

    def _get_system_prompt(self) -> str:
        """System prompt for survey generation"""
        return """You are an expert survey designer and qualitative researcher with deep knowledge of:
- Industry best practices for survey design
- Question formulation techniques (open-ended, closed-ended, Likert scales)
- Avoiding bias and leading questions
- Survey flow and respondent experience
- Research methodologies (interviews, focus groups, ethnographic studies)

Your task is to generate professional, well-structured surveys that will yield high-quality research data.
Follow these principles:
1. Use clear, unambiguous language
2. Avoid double-barreled questions
3. Include a logical flow from general to specific
4. Balance open-ended and structured questions appropriately
5. Consider the respondent's cognitive load
6. Include screening questions when relevant
7. Add instructions and context where helpful

Always respond with valid JSON containing the survey structure."""

    def _build_generation_prompt(self, outline, survey_type, num_questions, target_audience) -> str:
        """Build the user prompt for survey generation"""
        return f"""Generate a professional {survey_type} survey based on the following outline:

OUTLINE:
{outline}

REQUIREMENTS:
- Target number of questions: {num_questions}
- Target audience: {target_audience}
- Survey type: {survey_type}

Please generate a complete survey with:
1. A clear title
2. An introduction/welcome message
3. Well-crafted questions following best practices
4. Appropriate question types for the research goals
5. A thank you/closing message

Respond with a JSON object in this exact format:
{{
  "title": "Survey Title",
  "introduction": "Welcome message and instructions",
  "questions": [
    {{
      "id": 1,
      "question_text": "The question to ask",
      "question_type": "open_ended|multiple_choice|likert_scale|yes_no|rating",
      "options": ["option1", "option2"],
      "required": true|false,
      "help_text": "Optional clarification"
    }}
  ],
  "closing": "Thank you message"
}}

For open-ended questions, omit the "options" field.
For multiple choice and Likert questions, include appropriate options.
Ensure questions follow best practices and are unbiased."""

    def _parse_survey_response(self, response: str) -> Dict:
        """Parse LLM response into survey structure"""
        # Try to extract JSON from response
        response = response.strip()

        # Handle code blocks
        if "```json" in response:
            start = response.find("```json") + 7
            end = response.find("```", start)
            response = response[start:end].strip()
        elif "```" in response:
            start = response.find("```") + 3
            end = response.find("```", start)
            response = response[start:end].strip()

        try:
            survey_data = json.loads(response)

            # Validate required fields
            required_fields = ["title", "introduction", "questions", "closing"]
            for field in required_fields:
                if field not in survey_data:
                    raise ValueError(f"Missing required field: {field}")

            # Validate questions
            if not isinstance(survey_data["questions"], list) or len(survey_data["questions"]) == 0:
                raise ValueError("Survey must contain at least one question")

            return survey_data

        except json.JSONDecodeError as e:
            raise Exception(f"Failed to parse survey JSON: {str(e)}\nResponse: {response}")

    def refine_question(self, question: str, improvement_type: str = "clarity") -> str:
        """
        Refine a single survey question.

        Args:
            question: The question to improve
            improvement_type: Type of improvement (clarity, neutrality, specificity)

        Returns:
            Improved question text
        """
        prompt = f"""Improve the following survey question for better {improvement_type}:

Original Question: {question}

Provide an improved version that:
- {"Is clearer and easier to understand" if improvement_type == "clarity" else ""}
- {"Removes bias and leading language" if improvement_type == "neutrality" else ""}
- {"Is more specific and actionable" if improvement_type == "specificity" else ""}

Respond with only the improved question text, no explanation."""

        messages = [
            {"role": "system", "content": "You are an expert survey question designer."},
            {"role": "user", "content": prompt}
        ]

        return self.llm.generate(messages, max_tokens=150, temperature=0.5).strip()

    def add_follow_up_questions(self, base_question: str, num_follow_ups: int = 3) -> List[str]:
        """
        Generate follow-up questions for deeper exploration.

        Args:
            base_question: The main question
            num_follow_ups: Number of follow-up questions to generate

        Returns:
            List of follow-up question texts
        """
        prompt = f"""Generate {num_follow_ups} follow-up questions for this main question:

Main Question: {base_question}

The follow-up questions should:
1. Probe deeper into the topic
2. Explore different aspects or dimensions
3. Encourage detailed responses
4. Follow a logical progression

Respond with a JSON array of question strings."""

        messages = [
            {"role": "system", "content": "You are an expert in qualitative research interviews."},
            {"role": "user", "content": prompt}
        ]

        response = self.llm.generate(messages, max_tokens=500, temperature=0.7)

        try:
            # Extract JSON array
            if "[" in response:
                start = response.find("[")
                end = response.rfind("]") + 1
                follow_ups = json.loads(response[start:end])
                return follow_ups[:num_follow_ups]
        except:
            pass

        # Fallback: split by newlines
        lines = [line.strip() for line in response.split("\n") if line.strip()]
        return [line.lstrip("0123456789.-) ") for line in lines if "?" in line][:num_follow_ups]