add pyro (de)serializer for datetime

This commit is contained in:
Reinier van der Leer
2025-02-12 17:11:43 +01:00
parent aa21bf7b2d
commit 1dfc840907
3 changed files with 24 additions and 8 deletions

View File

@@ -366,7 +366,7 @@ async def get_execution_results(graph_exec_id: str) -> list[ExecutionResult]:
async def get_executions_in_timerange(
user_id: str, start_time: str, end_time: str
user_id: str, start_time: datetime, end_time: datetime
) -> list[ExecutionResult]:
try:
executions = await AgentGraphExecution.prisma().find_many(
@@ -374,8 +374,8 @@ async def get_executions_in_timerange(
"AND": [
{
"startedAt": {
"gte": datetime.fromisoformat(start_time),
"lte": datetime.fromisoformat(end_time),
"gte": start_time,
"lte": end_time,
}
},
{"userId": user_id},

View File

@@ -144,15 +144,17 @@ async def migrate_and_encrypt_user_integrations():
)
async def get_active_user_ids_in_timerange(start_time: str, end_time: str) -> list[str]:
async def get_active_user_ids_in_timerange(
start_time: datetime, end_time: datetime
) -> list[str]:
try:
users = await User.prisma().find_many(
where={
"AgentGraphExecutions": {
"some": {
"createdAt": {
"gte": datetime.fromisoformat(start_time),
"lte": datetime.fromisoformat(end_time),
"gte": start_time,
"lte": end_time,
}
}
}
@@ -168,8 +170,8 @@ async def get_active_user_ids_in_timerange(start_time: str, end_time: str) -> li
async def get_active_users_ids() -> list[str]:
user_ids = await get_active_user_ids_in_timerange(
(datetime.now() - timedelta(days=30)).isoformat(),
datetime.now().isoformat(),
datetime.now() - timedelta(days=30),
datetime.now(),
)
return user_ids

View File

@@ -6,6 +6,7 @@ import threading
import time
import typing
from abc import ABC, abstractmethod
from datetime import datetime
from enum import Enum
from types import NoneType, UnionType
from typing import (
@@ -307,3 +308,16 @@ def _pydantic_models_from_type_annotation(annotation) -> Iterator[type[BaseModel
yield annotype
elif annotype not in builtin_types and not issubclass(annotype, Enum):
raise TypeError(f"Unsupported type encountered: {annotype}")
# Register (de)serializers for datetime objects
pyro.register_class_to_dict(
datetime,
lambda dt: {
"__class__": datetime.__qualname__,
"iso": cast(datetime, dt).isoformat(),
},
)
pyro.register_dict_to_class(
datetime.__qualname__, lambda _, dict: datetime.fromisoformat(dict["iso"])
)