MAIA/backend/modules/nlp/service.py

from sqlalchemy.orm import Session
from sqlalchemy import desc  # Import desc for ordering
from google import genai
import json
from datetime import datetime, timezone
from typing import List  # Import List

# Import the new model and Enum
from .models import ChatMessage, MessageSender
from core.config import settings

client = genai.Client(api_key=settings.GOOGLE_API_KEY)

### Base prompt for MAIA, used for inital user requests
SYSTEM_PROMPT = (
    """
You are MAIA - My AI Assistant. Your job is to parse user requests into structured JSON commands and generate a user-facing response text.

Available functions/intents:
1. ask_ai(request: str): Use for simple questions (e.g., weather, facts). Forward the user's request.
2. get_calendar_events(start: Optional[datetime], end: Optional[datetime]): Retrieve calendar events.
3. add_calendar_event(title: str, description: str, start: datetime, end: Optional[datetime], location: str): Add a new event.
4. update_calendar_event(event_id: int, title: Optional[str], description: Optional[str], start: Optional[datetime], end: Optional[datetime], location: Optional[str]): Update an existing event. Requires event_id.
5. delete_calendar_event(event_id: int): Delete an event. Requires event_id.
6. get_todos(): Retrieve the user's TODO list.
7. add_todo(task: str, date: Optional[datetime], remind: Optional[bool]): Add a new task to the user's TODO list.
8. update_todo(todo_id: int, task: Optional[str], date: Optional[datetime], remind: Optional[bool], complete: Optional[bool]): Update an existing TODO item. Requires todo_id.
9. delete_todo(todo_id: int): Delete a TODO item. Requires todo_id.
10. clarification_needed(request: str): Use this if the user's request is ambiguous or lacks necessary information (like event_id or todo_id for update/delete). The original user request should be passed in the 'request' parameter.

**IMPORTANT:** Respond ONLY with JSON containing BOTH "intent" and "params", AND a "response_text" field.
- "response_text" should be a friendly, user-facing message confirming the action taken, providing the answer, asking for clarification OR can be empty if the query does not require a response to the user.

Examples:

User: Add a meeting tomorrow at 3pm about project X
MAIA:
{
  "intent": "add_calendar_event",
  "params": {
    "title": "Meeting",
    "description": "Project X",
    "start": "2025-04-22 15:00:00.000000+00:00",
    "end": null,
    "location": null
  },
  "response_text": "Okay, I've added a meeting about Project X to your calendar for tomorrow at 3 PM."
}

User: What's the weather like?
MAIA:
{
  "intent": "ask_ai",
  "params": {
    "request": "What's the weather like?"
  },
  "response_text": "Let me check the weather for you."
}

User: Delete the team sync event.
MAIA:
{
  "intent": "clarification_needed",
  "params": {
      "request": "Delete the team sync event."
  },
  "response_text": "Okay, I can help with that. Could you please provide the ID or more specific details about the 'team sync' event you want me to delete?"
}

User: Add 'Buy groceries' to my todo list
MAIA:
{
  "intent": "add_todo",
  "params": {
    "task": "Buy groceries",
    "date": null,
    "remind": false
  },
  "response_text": "I've added 'Buy groceries' to your TODO list."
}

User: Show me my todos
MAIA:
{
  "intent": "get_todos",
  "params": {},
  "response_text": "Okay, fetching your TODO list now."
}

User: Mark task 15 as complete
MAIA:
{
  "intent": "update_todo",
  "params": {
    "todo_id": 15,
    "complete": true
  },
  "response_text": "Got it, I've marked task 15 as complete."
}

User: Delete task 2
MAIA:
{
  "intent": "delete_todo",
  "params": {
    "todo_id": 2
  },
  "response_text": "Okay, I've deleted task 2 from your list."
}

The datetime right now is """
    + str(datetime.now(timezone.utc))
    + """.
"""
)

### Prompt for MAIA to forward user request to AI
SYSTEM_FORWARD_PROMPT = f"""
You are MAIA - My AI Assistant. Your job is to answer user simple user requests.
Here is some context for you:
 - The datetime right now is {str(datetime.now(timezone.utc))}.
Here is the user request:

"""

# --- Chat History Service Functions ---


def save_chat_message(db: Session, user_id: int, sender: MessageSender, text: str):
    """Saves a chat message to the database."""
    db_message = ChatMessage(user_id=user_id, sender=sender, text=text)
    db.add(db_message)
    db.commit()
    db.refresh(db_message)
    return db_message


def get_chat_history(db: Session, user_id: int, limit: int = 50) -> List[ChatMessage]:
    """Retrieves the last 'limit' chat messages for a user."""
    return (
        db.query(ChatMessage)
        .filter(ChatMessage.user_id == user_id)
        .order_by(desc(ChatMessage.timestamp))
        .limit(limit)
        .all()[::-1]
    )  # Reverse to get oldest first for display order


# --- Existing NLP Service Functions ---


def process_request(request: str):
    """
    Process the user request using the Google GenAI API.
    Expects a JSON response with intent, params, and response_text.
    """
    response = client.models.generate_content(
        model="gemini-2.0-flash",
        contents=SYSTEM_PROMPT + f"\n\nUser: {request}\nMAIA:",
        config={
            "temperature": 0.3,  # Less creativity, more factual
            "response_mime_type": "application/json",
        },
    )

    # Parse the JSON response
    try:
        parsed_response = json.loads(response.text)
        # Validate required fields
        if not all(k in parsed_response for k in ("intent", "params", "response_text")):
            raise ValueError(
                "AI response missing required fields (intent, params, response_text)"
            )
        return parsed_response
    except (json.JSONDecodeError, ValueError) as e:
        print(f"Error parsing AI response: {e}")
        print(f"Raw AI response: {response.text}")
        # Return a structured error that the API layer can handle
        return {
            "intent": "error",
            "params": {},
            "response_text": "Sorry, I had trouble understanding that request or formulating a response. Could you please try rephrasing?",
        }


def ask_ai(request: str):
    """
    Ask the AI a question.
    This is only called by MAIA when the intent is a simple question.
    """
    response = client.models.generate_content(
        model="gemini-2.0-flash",
        contents=SYSTEM_FORWARD_PROMPT + request,
    )
    return response.text