agent.py

import os
import asyncio
from dotenv import load_dotenv
from openai import AsyncOpenAI
from agents import Agent, Runner, OpenAIChatCompletionsModel

# ---------------------------
# 1. Load environment
# ---------------------------
load_dotenv()

OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
if not OPENROUTER_API_KEY:
    raise ValueError("OPENROUTER_API_KEY is missing in .env file")

# ---------------------------
# 2. OpenRouter + Gemini model
# ---------------------------
BASE_URL = "https://openrouter.ai/api/v1"

# ✅ Same model that worked in your last project
MODEL_NAME = "google/gemini-2.0-flash-001"

# Create OpenAI-compatible client that talks to OpenRouter
client = AsyncOpenAI(
    api_key=OPENROUTER_API_KEY,
    base_url=BASE_URL,
)

# Shared model adapter for Agents SDK
study_model = OpenAIChatCompletionsModel(
    model=MODEL_NAME,
    openai_client=client,
)

# ---------------------------
# 3. Fix event loop issue (Streamlit + asyncio)
# ---------------------------
def _ensure_event_loop() -> None:
    """
    Runner.run_sync() needs an asyncio event loop.
    On Python 3.12, there may be no default loop in the Streamlit thread,
    so we create and set one if missing.
    """
    try:
        asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

# ---------------------------
# 4. Summary Agent
# ---------------------------
summary_agent = Agent(
    name="StudyNotesSummarizer",
    instructions=(
        "You are a helpful study assistant. "
        "Given raw text from a PDF, create a clear, exam-focused summary "
        "using headings and bullet points. Keep it concise but meaningful."
    ),
    model=study_model,
)

# --------- QUIZ AGENT (UPGRADED) ---------
quiz_agent = Agent(
    name="StudyQuizGenerator",
    instructions=(
        "You are an exam paper setter. Using the provided study material:\n"
        "1. Generate EXACTLY 5 MCQs.\n"
        "2. EACH MCQ MUST have options A, B, C, D.\n"
        "3. Options must be realistic and NOT empty.\n"
        "4. After MCQs, generate EXACTLY 5 short questions.\n"
        "5. Do NOT include answers unless user asks.\n"
        "6. Format MCQs like:\n"
        "Q1. Question text?\n"
        "A) ...\nB) ...\nC) ...\nD) ...\n\n"
        "Short questions like:\n"
        "Q1. ......................?\n"
    ),
    model=study_model,
)

# ---------------------------
# 6. Public helper functions
# ---------------------------
def generate_summary(pdf_text: str) -> str:
    """Call the summary agent to summarize the PDF text."""
    _ensure_event_loop()

    prompt = (
        "Summarize the following study notes for a student preparing for exams. "
        "Use clear headings and bullet points.\n\n"
        f"{pdf_text}"
    )

    result = Runner.run_sync(summary_agent, prompt)
    return result.final_output


# --------- UPDATED STRICT QUIZ GENERATOR ---------
def generate_quiz(pdf_text: str) -> str:
    _ensure_event_loop()

    prompt = (
        "Create an exam-style quiz from the following material.\n\n"
        "You MUST follow this structure:\n\n"
        "=== MCQs (5) ===\n"
        "Q1. <question>\n"
        "A) <option>\nB) <option>\nC) <option>\nD) <option>\n\n"
        "Repeat for Q2–Q5.\n\n"
        "=== Short Questions (5) ===\n"
        "Q1. <short question>\n"
        "Repeat for Q2–Q5.\n\n"
        "DO NOT give answers.\n"
        "DO NOT skip options.\n"
        "DO NOT create less than 5 questions.\n\n"
        f"STUDY MATERIAL:\n{pdf_text}"
    )

    result = Runner.run_sync(quiz_agent, prompt)
    return result.final_output