app.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import openai
import os
import json
import re
from typing import Dict, List, Optional, Tuple, Any

app = FastAPI(title="TestCreationAgent", 
              description="An API for collecting test creation parameters through conversation")

# Add CORS middleware to allow requests from frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Define subject chapters mapping
SUBJECT_CHAPTERS = {
    "Mathematics": [
        "Number Systems", "Polynomials", "Coordinate Geometry", "Linear Equations in Two Variables",
        "Introduction to Euclid's Geometry", "Lines and Angles", "Triangles", "Quadrilaterals",
        "Areas of Parallelograms and Triangles", "Circles", "Constructions", "Heron's Formula",
        "Surface Areas and Volumes", "Statistics", "Probability", "Real Numbers",
        "Pair of Linear Equations in Two Variables", "Quadratic Equations", "Arithmetic Progressions",
        "Introduction to Trigonometry", "Some Applications of Trigonometry", "Areas Related to Circles",
        "Sets", "Relations and Functions", "Trigonometric Functions", "Principle of Mathematical Induction",
        "Complex Numbers and Quadratic Equations", "Linear Inequalities", "Permutations and Combinations",
        "Binomial Theorem", "Sequences and Series", "Straight Lines", "Conic Sections",
        "Introduction to Three Dimensional Geometry", "Limits and Derivatives",
        "Inverse Trigonometric Functions", "Matrices", "Determinants",
        "Continuity and Differentiability", "Application of Derivatives", "Integrals",
        "Application of Integrals", "Differential Equations", "Vector Algebra",
        "Three Dimensional Geometry", "Linear Programming"
    ],
    "Physics": [
        "Motion", "Force and Laws of Motion", "Gravitation", "Work and Energy", "Sound",
        "Light: Reflection and Refraction", "Human Eye and Colourful World", "Electricity",
        "Magnetic Effects of Electric Current", "Physical World and Measurement", "Kinematics",
        "Laws of Motion", "Work, Energy and Power", "Motion of System of Particles and Rigid Body",
        "Properties of Bulk Matter", "Thermodynamics", "Behaviour of Perfect Gases and Kinetic Theory",
        "Oscillations and Waves", "Electrostatics", "Current Electricity",
        "Magnetic Effects of Current and Magnetism", "Electromagnetic Induction and Alternating Currents",
        "Electromagnetic Waves", "Optics", "Dual Nature of Radiation and Matter", "Atoms", "Nuclei",
        "Semiconductor Electronics: Materials, Devices and Simple Circuits", "Vectors"
    ],
    "Chemistry": [
        "Matter in Our Surroundings", "Is Matter Around Us Pure?", "Atoms and Molecules",
        "Structure of the Atom", "Chemical Reactions and Equations", "Acids, Bases and Salts",
        "Metals and Non-metals", "Carbon and Its Compounds", "Periodic Classification of Elements",
        "Some Basic Concepts of Chemistry", "Structure of Atom",
        "Classification of Elements and Periodicity in Properties",
        "Chemical Bonding and Molecular Structure", "States of Matter: Gases and Liquids",
        "Thermodynamics", "Equilibrium", "Redox Reactions",
        "Organic Chemistry: Some Basic Principles and Techniques", "Hydrocarbons",
        "Environmental Chemistry", "Solid State", "Solutions", "Electrochemistry",
        "Chemical Kinetics", "Surface Chemistry", "General Principles and Processes of Isolation of Elements",
        "p-Block Elements", "d- and f-Block Elements", "Coordination Compounds",
        "Haloalkanes and Haloarenes", "Alcohols, Phenols and Ethers",
        "Aldehydes, Ketones and Carboxylic Acids", "Amines", "Biomolecules", "Polymers",
        "Chemistry in Everyday Life"
    ],
    "Organic Chemistry": [
        "Organic Chemistry: Some Basic Principles and Techniques", "Hydrocarbons",
        "Haloalkanes and Haloarenes", "Alcohols, Phenols and Ethers",
        "Aldehydes, Ketones and Carboxylic Acids", "Amines", "Biomolecules",
        "Polymers", "Chemistry in Everyday Life"
    ],
    "Inorganic Chemistry": [
        "Classification of Elements and Periodicity in Properties",
        "Chemical Bonding and Molecular Structure", "Redox Reactions",
        "p-Block Elements", "d- and f-Block Elements", "Coordination Compounds"
    ]
}

# Create a flat mapping of misspelled/approximate chapter names to correct ones
CHAPTER_MAPPING = {}
for subject, chapters in SUBJECT_CHAPTERS.items():
    for chapter in chapters:
        # Add the correct chapter name
        CHAPTER_MAPPING[chapter.lower()] = (subject, chapter)

        # Add common misspellings/variations
        if chapter.lower() == "thermodynamics":
            CHAPTER_MAPPING["termodyanamics"] = (subject, chapter)
            CHAPTER_MAPPING["termodyn"] = (subject, chapter)
            CHAPTER_MAPPING["thermo"] = (subject, chapter)
            CHAPTER_MAPPING["thermodynamic"] = (subject, chapter)


class UserInput(BaseModel):
    message: str
    session_id: str


class SessionState(BaseModel):
    params: Dict[str, str] = {
        "chapters_of_the_test": "",
        "questions_per_chapter": "",
        "difficulty_distribution": "",
        "test_duration": "",
        "test_date": "",
        "test_time": ""
    }
    completed: bool = False
    attempt_count: int = 0


# In-memory session storage
sessions = {}


def normalize_chapter_name(chapter_input: str) -> Optional[Tuple[str, str]]:
    """
    Maps user input to standardized chapter names from the curriculum.
    Returns tuple of (subject, correct_chapter_name) or None if no match.
    """
    if not chapter_input:
        return None

    # Direct mapping for exact matches or known misspellings
    norm_input = chapter_input.lower().strip()
    if norm_input in CHAPTER_MAPPING:
        return CHAPTER_MAPPING[norm_input]

    # Try fuzzy matching if no direct match
    # Look for partial matches
    for chapter_key, (subject, correct_name) in CHAPTER_MAPPING.items():
        if norm_input in chapter_key or chapter_key in norm_input:
            return (subject, correct_name)

    # No match found
    return None


async def llm_extractParams(user_input: str, current_params: Dict[str, str]) -> Dict[str, str]:
    """
    Extracts structured test parameters from natural language input
    and updates the provided params dictionary.
    """
    system_prompt = """
You are an expert educational test creation assistant that extracts test setup parameters from user input.
Extract ONLY the parameters explicitly mentioned in the user's message.

Return a JSON object with all the following keys:
- chapters_of_the_test (string: list of chapters or topics)
- questions_per_chapter (string or number: how many questions per chapter)
- difficulty_distribution (string: e.g., "easy:40%, medium:40%, hard:20%" or any format specified)
- test_duration (string or number: time in minutes)
- test_date (string: in any reasonable date format)
- test_time (string: time of day)

Important rules:
- Do NOT make assumptions - if information isn't provided, leave as empty string ("")
- Only fill in values explicitly mentioned by the user
- For difficulty_distribution:
  * Convert numeric sequences like "30 40 30" to "easy:30%, medium:40%, hard:30%" if they appear to be distributions
  * Convert descriptions like "mostly hard" to approximate percentages (e.g., "easy:20%, medium:20%, hard:60%")
  * Accept formats like "60 easy, 20 medium, 20 hard" and convert to percentages
- Return valid JSON with all keys, even if empty
"""
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_input}
    ]

    try:
        response = openai.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            temperature=0.2
        )

        extracted_json = response.choices[0].message.content.strip()

        # Handle potential JSON formatting issues by extracting JSON from response
        if not extracted_json.startswith('{'):
            # Find JSON object in text if it's not a clean JSON response
            start_idx = extracted_json.find('{')
            end_idx = extracted_json.rfind('}') + 1
            if start_idx >= 0 and end_idx > start_idx:
                extracted_json = extracted_json[start_idx:end_idx]
            else:
                raise ValueError("Unable to extract valid JSON from response")

        # Parse and update the current_params safely
        extracted_dict = json.loads(extracted_json)
        updated_params = current_params.copy()
        
        for key in updated_params:
            if key.lower() in extracted_dict and extracted_dict[key.lower()]:
                updated_params[key] = extracted_dict[key.lower()]
            elif key in extracted_dict and extracted_dict[key]:
                updated_params[key] = extracted_dict[key]

        # Apply chapter mapping if chapters were specified
        if updated_params["chapters_of_the_test"] and updated_params["chapters_of_the_test"] != current_params["chapters_of_the_test"]:
            chapters_input = updated_params["chapters_of_the_test"]
            # Split multiple chapters if comma-separated
            chapter_list = [ch.strip() for ch in re.split(r',|;', chapters_input)]

            mapped_chapters = []
            for chapter in chapter_list:
                result = normalize_chapter_name(chapter)
                if result:
                    subject, correct_name = result
                    mapped_chapters.append(f"{correct_name} ({subject})")
                else:
                    mapped_chapters.append(chapter)  # Keep as-is if no mapping found

            updated_params["chapters_of_the_test"] = ", ".join(mapped_chapters)

        return updated_params

    except json.JSONDecodeError as e:
        print(f"Error: Could not parse response as JSON: {e}")
        return current_params
    except Exception as e:
        print(f"Error during parameter extraction: {e}")
        return current_params


def gate(params: Dict[str, str]) -> List[str]:
    """
    Checks which fields are still empty in the params.
    Returns a list of missing parameter keys.
    """
    return [key for key, val in params.items() if not val]


async def llm_getMissingParams(missing_keys: List[str]) -> str:
    """
    Generates a human-readable prompt to ask user for missing fields.
    """
    # Create context-aware prompts for specific missing fields
    context_details = {
        "chapters_of_the_test": "such as Math, Science, History, etc.",
        "questions_per_chapter": "the number of questions for each chapter",
        "difficulty_distribution": "as percentages or numbers (easy, medium, hard)",
        "test_duration": "in minutes",
        "test_date": "when the test will be given",
        "test_time": "the time of day for the test"
    }

    # Create a more specific prompt based on what's missing
    if len(missing_keys) == 1:
        key = missing_keys[0]
        prompt = f"Please provide the {key.replace('_', ' ')} {context_details.get(key, '')}."
    else:
        formatted_missing = [f"{key.replace('_', ' ')} ({context_details.get(key, '')})" for key in missing_keys]
        prompt = f"The following test details are still needed: {', '.join(formatted_missing)}."

    messages = [
        {"role": "system", "content": "You are a helpful assistant who creates clear, concise questions to collect missing test setup information. Keep your response under 2 sentences and focus only on what's missing."},
        {"role": "user", "content": prompt}
    ]

    try:
        response = openai.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            temperature=0.3
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"Error generating prompt for missing values: {e}")
        return f"Please provide the following missing information: {', '.join(missing_keys)}."


@app.on_event("startup")
async def startup_event():
    # Set up OpenAI API key from environment variable
    openai.api_key = os.getenv("OPENAI_API_KEY")
    if not openai.api_key:
        print("⚠️ WARNING: OPENAI_API_KEY environment variable not set.")


@app.get("/")
async def root():
    return {"message": "Test Creation Agent API is running"}


@app.post("/chat")
async def chat(user_input: UserInput):
    session_id = user_input.session_id
    
    # Initialize session if it doesn't exist
    if session_id not in sessions:
        sessions[session_id] = SessionState()
    
    session = sessions[session_id]
    
    # If this is the first message, send a welcome message
    if session.attempt_count == 0:
        session.attempt_count += 1
        return {
            "response": "👋 Welcome! Please provide the test setup details. I need: chapters, questions per chapter, difficulty distribution, test duration, date, and time.",
            "session_state": {
                "params": session.params,
                "completed": False
            }
        }
    
    # Process user input to extract parameters
    session.params = await llm_extractParams(user_input.message, session.params)
    session.attempt_count += 1
    
    # Check if we have all required parameters
    missing = gate(session.params)
    
    # If we have all parameters or exceeded max attempts, return completion
    max_attempts = 10
    if not missing or session.attempt_count > max_attempts:
        session.completed = True
        if not missing:
            result = "✅ All test parameters are now complete:"
        else:
            result = "⚠️ Some parameters could not be filled after multiple attempts:"
        
        # Format the parameters as a readable string
        for k, v in session.params.items():
            result += f"\n- {k.replace('_', ' ').title()}: {v or 'Not provided'}"
        
        return {
            "response": result,
            "session_state": {
                "params": session.params,
                "completed": True
            }
        }
    
    # Otherwise, ask for missing parameters
    follow_up_prompt = await llm_getMissingParams(missing)
    
    return {
        "response": follow_up_prompt,
        "session_state": {
            "params": session.params,
            "completed": False
        }
    }


@app.get("/session/{session_id}")
async def get_session(session_id: str):
    if session_id not in sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    
    session = sessions[session_id]
    return {
        "params": session.params,
        "completed": session.completed,
        "attempt_count": session.attempt_count
    }


@app.delete("/session/{session_id}")
async def delete_session(session_id: str):
    if session_id in sessions:
        del sessions[session_id]
    return {"message": "Session deleted successfully"}


@app.post("/reset")
async def reset_session(user_input: UserInput):
    session_id = user_input.session_id
    sessions[session_id] = SessionState()
    
    return {
        "response": "Session reset. 👋 Welcome! Please provide the test setup details. I need: chapters, questions per chapter, difficulty distribution, test duration, date, and time.",
        "session_state": {
            "params": sessions[session_id].params,
            "completed": False
        }
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=int(os.getenv("PORT", 8000)), reload=True)