app/services/classification_service.py

import json
import logging
import os

import httpx
from google import genai
from google.genai import errors, types
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

import app.config  # Loads Backend/.env
from app.models.department import Department

logger = logging.getLogger(__name__)


class ClassificationResult(BaseModel):
    department_id: int | None
    confidence: float = Field(ge=0, le=1)


def classify_query(
    db: Session,
    message: str,
) -> tuple[Department, float] | None:
    """Return a validated Gemini result, or None to use keyword routing."""
    gemini_keys = [
    os.getenv("GEMINI_API_KEY_1"),
    os.getenv("GEMINI_API_KEY_2"),
    os.getenv("GEMINI_API_KEY_3"),
    os.getenv("GEMINI_API_KEY_4"),
    os.getenv("GEMINI_API_KEY_5"),
    os.getenv("GEMINI_API_KEY_6"),
    os.getenv("GEMINI_API_KEY_7"),
    os.getenv("GEMINI_API_KEY_8"),
    os.getenv("GEMINI_API_KEY_9"),
]


    gemini_keys = [
    key for key in gemini_keys if key
]


    api_key = gemini_keys[0] if gemini_keys else None
    if not api_key:
        logger.warning("Gemini key missing; using keyword routing.")
        return None

    departments = (
        db.query(Department)
        .filter(Department.is_active.is_(True))
        .all()
    )
    if not departments:
        return None

    department_data = [
        {
            "id": department.id,
            "name": department.name,
            "code": department.code,
            "description": department.description or "",
            "keywords": department.keywords or "",
            "example_queries": department.query or "",
        }
        for department in departments
    ]

    try:
        with genai.Client(
            api_key=api_key,
            http_options=types.HttpOptions(
                timeout=15000,
                retry_options=types.HttpRetryOptions(attempts=1),
            ),
        ) as client:
            response = client.models.generate_content(
                  model="gemini-3.6-flash",
                contents=json.dumps({
                    "departments": department_data,
                    "student_message": message,
                }),
                config=types.GenerateContentConfig(
                    system_instruction=(
                     "You are an AI query routing assistant for  a university. "
                      "Classify student queries and select the most suitable department. "
                      "Use only the supplied department information. "
                      "Return only the department_id and confidence score. "
                      "Do not answer the student query."
),
                    response_mime_type="application/json",
                    response_schema=ClassificationResult,
                    automatic_function_calling=(
                        types.AutomaticFunctionCallingConfig(disable=True)
                    ),
                ),
            )

        result = ClassificationResult.model_validate_json(
            response.text or ""
        )

        department = next(
            (
                department
                for department in departments
                if department.id == result.department_id
            ),
            None,
        )

        if department is None:
            logger.warning(
                "Gemini returned no valid department; using keyword routing."
            )
            return None

        return department, result.confidence

    except (errors.APIError, httpx.HTTPError, ValueError) as exc:
        # Log the error type without exposing credentials or query content.
        logger.warning(
            "Gemini classification failed (%s); using keyword routing.",
            type(exc).__name__,
        )
        return None