From c08d35560efbc80d07c0a4eee8a4e8738352ee85 Mon Sep 17 00:00:00 2001
From: Alex Dixon <alex.dixon@embeddedinsurance.com>
Date: Mon, 5 Aug 2024 22:17:42 -0700
Subject: [PATCH] cleanup

---
 src/ell/store.py      |  5 +++--
 src/ell/stores/sql.py |  4 ++--
 src/ell/types.py      | 10 ++--------
 3 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/src/ell/store.py b/src/ell/store.py
index 02a46ffc..11348a35 100644
--- a/src/ell/store.py
+++ b/src/ell/store.py
@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
+from datetime import datetime
 from typing import Any, Optional, Dict, List, Set, Union
 from ell.lstr import lstr
 from ell.types import InvocableLM
@@ -15,7 +16,7 @@ def write_lmp(self, lmp_id: str, name: str, source: str, dependencies: List[str]
                   version_number: int,
                   uses: Dict[str, Any], 
                   commit_message: Optional[str] = None,
-                  created_at: Optional[float]=None) -> Optional[Any]:
+                  created_at: Optional[datetime]=None) -> Optional[Any]:
         """
         Write an LMP (Language Model Package) to the storage.
 
@@ -33,7 +34,7 @@ def write_lmp(self, lmp_id: str, name: str, source: str, dependencies: List[str]
 
     @abstractmethod
     def write_invocation(self, id: str, lmp_id: str, args: str, kwargs: str, result: Union[lstr, List[lstr]], invocation_kwargs: Dict[str, Any], 
-                         created_at: Optional[float], consumes: Set[str], prompt_tokens: Optional[int] = None,
+                         created_at: Optional[datetime], consumes: Set[str], prompt_tokens: Optional[int] = None,
                          completion_tokens: Optional[int] = None, latency_ms: Optional[float] = None,
                          state_cache_key: Optional[str] = None,
                          cost_estimate: Optional[float] = None) -> Optional[Any]:
diff --git a/src/ell/stores/sql.py b/src/ell/stores/sql.py
index 36b2a7c0..7e32e18b 100644
--- a/src/ell/stores/sql.py
+++ b/src/ell/stores/sql.py
@@ -1,4 +1,4 @@
-import datetime
+from datetime import datetime
 import json
 import os
 from typing import Any, Optional, Dict, List, Set, Union
@@ -26,7 +26,7 @@ def write_lmp(self, lmp_id: str, name: str, source: str, dependencies: List[str]
                   global_vars: Dict[str, Any],
                   free_vars: Dict[str, Any],
                   commit_message: Optional[str] = None,
-                  created_at: Optional[float]=None) -> Optional[Any]:
+                  created_at: Optional[datetime]=None) -> Optional[Any]:
         with Session(self.engine) as session:
             lmp = session.query(SerializedLMP).filter(SerializedLMP.lmp_id == lmp_id).first()
             
diff --git a/src/ell/types.py b/src/ell/types.py
index 7514386a..eb5a4ead 100644
--- a/src/ell/types.py
+++ b/src/ell/types.py
@@ -1,15 +1,14 @@
 # Let's define the core types.
 from dataclasses import dataclass
-from typing import Callable, Dict, List, Union
+from typing import Callable, Dict, List, Union, Any, Optional
 
-from typing import Any
 from ell.lstr import lstr
 from ell.util.dict_sync_meta import DictSyncMeta
 
 from datetime import datetime, timezone
 from typing import Any, List, Optional
 from sqlmodel import Field, SQLModel, Relationship, JSON, Column
-from sqlalchemy import TIMESTAMP, func
+from sqlalchemy import func
 import sqlalchemy.types as types
 
 _lstr_generic = Union[lstr, str]
@@ -44,10 +43,6 @@ class Message(dict, metaclass=DictSyncMeta):
 LMP = Union[OneTurn, MultiTurnLMP, ChatLMP]
 InvocableLM = Callable[..., _lstr_generic]
 
-from datetime import timezone
-from sqlmodel import Field
-from typing import Optional
-
 
 def utc_now() -> datetime:
     """
@@ -92,7 +87,6 @@ class SerializedLMP(SQLModel, table=True):
     # Timestamp of when the LMP was created
     created_at: datetime = UTCTimestampField(
         index=True,
-        default=func.now(),
         nullable=False
     )
     is_lm: bool  # Boolean indicating if it is an LM (Language Model) or an LMP