from pydantic import BaseModel, model_validator
from typing import Dict, List, Optional, Any
from datetime import datetime, timezone

class Message(BaseModel):
    prompt: str
    preamble: str = ""
    content: str = ""
    response: str = ""        
    metadata: dict[str, Any] = {
        "rag": { "documents": [] },
        "tools": [],
        "eval_count": 0,
        "eval_duration": 0,
        "prompt_eval_count": 0,
        "prompt_eval_duration": 0,
    }
    actions: List[str] = []
    timestamp: datetime = datetime.now(timezone.utc)

    def add_action(self, action: str | list[str]) -> None:
        """Add a actions(s) to the message."""
        if isinstance(action, str):
          self.actions.append(action)
        else:
          self.actions.extend(action)
        
    def get_summary(self) -> str:
        """Return a summary of the message."""
        response_summary = (
            f"Response: {self.response} (Actions: {', '.join(self.actions)})"
            if self.response else "No response yet"
        )
        return (
            f"Message at {self.timestamp}:\n"
            f"Query: {self.preamble}{self.content}\n"
            f"{response_summary}"
        )