Spaces:
Sleeping
Sleeping
| from typing import Optional, Any, Dict | |
| from pydantic import BaseModel, ConfigDict | |
| class AgentSession(BaseModel): | |
| """Agent Session that is stored in the database""" | |
| # Session UUID | |
| session_id: str | |
| # ID of the agent that this session is associated with | |
| agent_id: Optional[str] = None | |
| # ID of the user interacting with this agent | |
| user_id: Optional[str] = None | |
| # Agent Memory | |
| memory: Optional[Dict[str, Any]] = None | |
| # Agent Metadata | |
| agent_data: Optional[Dict[str, Any]] = None | |
| # User Metadata | |
| user_data: Optional[Dict[str, Any]] = None | |
| # Session Metadata | |
| session_data: Optional[Dict[str, Any]] = None | |
| # The Unix timestamp when this session was created | |
| created_at: Optional[int] = None | |
| # The Unix timestamp when this session was last updated | |
| updated_at: Optional[int] = None | |
| model_config = ConfigDict(from_attributes=True) | |
| def monitoring_data(self) -> Dict[str, Any]: | |
| # Google Gemini adds a "parts" field to the messages, which is not serializable | |
| # If the provider is Google, remove the "parts" from the messages | |
| if self.agent_data is not None: | |
| if self.agent_data.get("model", {}).get("provider") == "Google" and self.memory is not None: | |
| # Remove parts from runs' response messages | |
| if "runs" in self.memory: | |
| for _run in self.memory["runs"]: | |
| if "response" in _run and "messages" in _run["response"]: | |
| for m in _run["response"]["messages"]: | |
| if isinstance(m, dict): | |
| m.pop("parts", None) | |
| # Remove parts from top-level memory messages | |
| if "messages" in self.memory: | |
| for m in self.memory["messages"]: | |
| if isinstance(m, dict): | |
| m.pop("parts", None) | |
| monitoring_data = self.model_dump() | |
| return monitoring_data | |
| def telemetry_data(self) -> Dict[str, Any]: | |
| return self.model_dump(include={"model", "created_at", "updated_at"}) | |