fixed migrations and module seperation

src/ell/evaluation/evaluation.py

@@ -14,13 +14,13 @@
 from ell.evaluation.results import _ResultDatapoint, EvaluationResults
 from ell.evaluation.serialization import write_evaluation, write_evaluation_run_end, write_evaluation_run_intermediate, write_evaluation_run_start
 from ell.evaluation.util import get_lmp_output
-from ell.types.studio import LMPType
+from ell.stores.models import LMPType
 
 from ell.evaluation.util import validate_callable_dict
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from ell.types.message import LMP
-from ell.types.studio.evaluations import EvaluationLabelerType
+from ell.stores.models.evaluations import EvaluationLabelerType
 from ell.util.tqdm import tqdm
 import inspect
 

src/ell/evaluation/results.py

@@ -3,7 +3,7 @@
 import numpy as np
 from dataclasses import dataclass, field
 
-from ell.types.studio.evaluations import EvaluationLabelerType
+from ell.stores.models.evaluations import EvaluationLabelerType
 
 Datapoint = Dict[str, Any]
 Dataset = List[Dict[str, Any]]

src/ell/evaluation/serialization.py

@@ -17,7 +17,7 @@
 
 import itertools
 
-from ell.types.studio.evaluations import (
+from ell.stores.models.evaluations import (
     EvaluationLabel,
     SerializedEvaluation as SerializedEvaluation,
     EvaluationLabeler,

src/ell/lmp/_track.py

@@ -17,7 +17,7 @@
 from ell.util.serialization import prepare_invocation_params
 
 try:
-    from ell.types.studio.core import SerializedLMP, Invocation, InvocationContents, utc_now
+    from ell.stores.models.core import SerializedLMP, Invocation, InvocationContents, utc_now
 except ImportError:
     SerializedLMP = Invocation =  InvocationContents = utc_now = None
 

src/ell/lmp/function.py

@@ -2,7 +2,7 @@
 from typing import Any, Callable
 from ell.configurator import config
 from ell.lmp._track import _track
-from ell.types.studio import LMPType
+from ell.stores.models import LMPType
 from ell.util.verbosity import model_usage_logger_pre
 
 def function(*, exempt_from_tracking: bool = False, _exempt_from_logging: bool = False, type = LMPType.FUNCTION, **function_kwargs):

src/ell/stores/migrations/script.py.mako

@@ -10,7 +10,7 @@ from typing import Sequence, Union
 from alembic import op
 import sqlalchemy as sa
 import sqlmodel
-import ell.types.studio.core
+import ell.stores.models.core
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.

src/ell/stores/migrations/versions/b3aff868f213_evaluations.py → src/ell/stores/migrations/versions/f6528d04bbbd_evaluations.py

@@ -1,20 +1,20 @@
 """evaluations
 
-Revision ID: b3aff868f213
+Revision ID: f6528d04bbbd
 Revises: 4524fb60d23e
-Create Date: 2024-11-19 18:51:57.273837+00:00
+Create Date: 2024-11-19 19:31:38.381105+00:00
 
 """
 from typing import Sequence, Union
 
 from alembic import op
 import sqlalchemy as sa
 import sqlmodel
-import ell.types.studio.core
+import ell.stores.models.core
 
 
 # revision identifiers, used by Alembic.
-revision: str = 'b3aff868f213'
+revision: str = 'f6528d04bbbd'
 down_revision: Union[str, None] = '4524fb60d23e'
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -25,7 +25,7 @@ def upgrade() -> None:
     op.create_table('serializedevaluation',
     sa.Column('id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
     sa.Column('name', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
-    sa.Column('created_at', ell.types.studio.core.UTCTimestamp(timezone=True), nullable=False),
+    sa.Column('created_at', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=False),
     sa.Column('dataset_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
     sa.Column('n_evals', sa.Integer(), nullable=False),
     sa.Column('version_number', sa.Integer(), nullable=False),
@@ -49,8 +49,8 @@ def upgrade() -> None:
     sa.Column('evaluation_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
     sa.Column('evaluated_lmp_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
     sa.Column('api_params', sa.JSON(), nullable=True),
-    sa.Column('start_time', ell.types.studio.core.UTCTimestamp(timezone=True), nullable=True),
-    sa.Column('end_time', ell.types.studio.core.UTCTimestamp(timezone=True), nullable=True),
+    sa.Column('start_time', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
+    sa.Column('end_time', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
     sa.Column('success', sa.Boolean(), nullable=True),
     sa.Column('error', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
     sa.ForeignKeyConstraint(['evaluated_lmp_id'], ['serializedlmp.lmp_id'], ),
@@ -70,9 +70,9 @@ def upgrade() -> None:
     op.create_table('evaluationrunlabelersummary',
     sa.Column('evaluation_labeler_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
     sa.Column('evaluation_run_id', sa.Integer(), nullable=False),
-    sa.Column('created_at', ell.types.studio.core.UTCTimestamp(timezone=True), nullable=True),
-    sa.Column('updated_at', ell.types.studio.core.UTCTimestamp(timezone=True), nullable=True),
-    sa.Column('finalized_at', ell.types.studio.core.UTCTimestamp(timezone=True), nullable=True),
+    sa.Column('created_at', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
+    sa.Column('updated_at', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
+    sa.Column('finalized_at', ell.stores.models.core.UTCTimestamp(timezone=True), nullable=True),
     sa.Column('is_scalar', sa.Boolean(), nullable=False),
     sa.Column('data', sa.JSON(), nullable=True),
     sa.Column('count', sa.Integer(), nullable=False),

src/ell/stores/sql.py

@@ -12,15 +12,15 @@
 from ell.types._lstr import _lstr
 from sqlalchemy import or_, func, and_, extract, FromClause
 from sqlalchemy.types import TypeDecorator, VARCHAR
-from ell.types.studio import SerializedLMPUses, utc_now
-from ell.types.studio.evaluations import (
+from ell.stores.models import SerializedLMPUses, utc_now
+from ell.stores.models.evaluations import (
     EvaluationLabeler,
     EvaluationResultDatapoint,
     EvaluationRunLabelerSummary,
     SerializedEvaluation,
     SerializedEvaluationRun,
 )
-from ell.types.studio.core import InvocationTrace, SerializedLMP, Invocation, InvocationContents
+from ell.stores.models.core import InvocationTrace, SerializedLMP, Invocation, InvocationContents
 from sqlalchemy import func, and_
 from ell.util.serialization import pydantic_ltype_aware_cattr
 import gzip

src/ell/stores/store.py

@@ -3,9 +3,9 @@
 from datetime import datetime
 from typing import Any, Optional, Dict, List, Set, Union
 from ell.types._lstr import _lstr
-from ell.types.studio.core import SerializedLMP, Invocation
+from ell.stores.models.core import SerializedLMP, Invocation
 from ell.types.message import InvocableLM
-from ell.types.studio.evaluations import EvaluationResultDatapoint, EvaluationRunLabelerSummary, SerializedEvaluation, SerializedEvaluationRun
+from ell.stores.models.evaluations import EvaluationResultDatapoint, EvaluationRunLabelerSummary, SerializedEvaluation, SerializedEvaluationRun
 # from ell.types.studio import SerializedEvaluation, SerializedEvaluationRun
 
 class BlobStore(ABC):

src/ell/studio/datamodels.py

@@ -1,15 +1,15 @@
 from datetime import datetime
 from typing import List, Optional, Dict, Any
 from sqlmodel import SQLModel
-from ell.types.studio.evaluations import (
+from ell.stores.models.evaluations import (
     EvaluationLabelBase,
     EvaluationLabelerBase,
     SerializedEvaluationBase,
     SerializedEvaluationRunBase,
     EvaluationRunLabelerSummaryBase,
     EvaluationResultDatapointBase,
 )
-from ell.types.studio.core import SerializedLMPBase, InvocationBase, InvocationContentsBase
+from ell.stores.models.core import SerializedLMPBase, InvocationBase, InvocationContentsBase
 
 
 class SerializedLMPWithUses(SerializedLMPBase):

src/ell/studio/server.py

@@ -11,10 +11,10 @@
 from ell.studio.connection_manager import ConnectionManager
 from ell.studio.datamodels import EvaluationResultDatapointPublic, InvocationPublicWithConsumes, SerializedLMPWithUses, EvaluationPublic, SpecificEvaluationRunPublic
 
-from ell.types.studio.core import SerializedLMP
+from ell.stores.models.core import SerializedLMP
 from datetime import datetime, timedelta
 from sqlmodel import select
-from ell.types.studio.evaluations import SerializedEvaluation
+from ell.stores.models.evaluations import SerializedEvaluation
 
 
 logger = logging.getLogger(__name__)

tests/test_migrations.py

@@ -116,7 +116,7 @@ def test_existing_tables_no_alembic(temp_db_url):
         result = conn.execute(text("SELECT version_num FROM ell_alembic_version"))
         version = result.scalar()
         # Get current head version from alembic config
-        assert version == "b3aff868f213"
+        assert version == "f6528d04bbbd"
 
 def test_multiple_migrations(temp_db_url):
     """Test running multiple migrations in sequence"""

tests/test_results.py

@@ -1,5 +1,5 @@
 from ell.evaluation.results import _ResultDatapoint, EvaluationResults, Label
-from ell.types.studio.evaluations import EvaluationLabelerType
+from ell.stores.models.evaluations import EvaluationLabelerType
 import numpy as np
 
 def test_evaluation_results_from_rowar_results():

tests/test_sql_store.py

@@ -4,7 +4,7 @@
 from sqlalchemy import Engine, create_engine, func
 
 from ell.types.lmp import LMPType
-from ell.types.studio.core import utc_now
+from ell.stores.models.core import utc_now
 
 
 @pytest.fixture