38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
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}"
|
|
) |