diff --git a/src/ell/evaluation/evaluation.py b/src/ell/evaluation/evaluation.py index e6d806b7..7bcbe968 100644 --- a/src/ell/evaluation/evaluation.py +++ b/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 diff --git a/src/ell/evaluation/results.py b/src/ell/evaluation/results.py index 14c1c46d..71da991e 100644 --- a/src/ell/evaluation/results.py +++ b/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]] diff --git a/src/ell/evaluation/serialization.py b/src/ell/evaluation/serialization.py index cb460f13..e946628d 100644 --- a/src/ell/evaluation/serialization.py +++ b/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, diff --git a/src/ell/lmp/_track.py b/src/ell/lmp/_track.py index 7503672f..bb5deb59 100644 --- a/src/ell/lmp/_track.py +++ b/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 diff --git a/src/ell/lmp/function.py b/src/ell/lmp/function.py index 06226348..636dd9c1 100644 --- a/src/ell/lmp/function.py +++ b/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): diff --git a/src/ell/stores/migrations/script.py.mako b/src/ell/stores/migrations/script.py.mako index 40e30f47..6b825de1 100644 --- a/src/ell/stores/migrations/script.py.mako +++ b/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. diff --git a/src/ell/stores/migrations/versions/b3aff868f213_evaluations.py b/src/ell/stores/migrations/versions/f6528d04bbbd_evaluations.py similarity index 87% rename from src/ell/stores/migrations/versions/b3aff868f213_evaluations.py rename to src/ell/stores/migrations/versions/f6528d04bbbd_evaluations.py index 17cc7628..60c731b6 100644 --- a/src/ell/stores/migrations/versions/b3aff868f213_evaluations.py +++ b/src/ell/stores/migrations/versions/f6528d04bbbd_evaluations.py @@ -1,8 +1,8 @@ """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 @@ -10,11 +10,11 @@ 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), diff --git a/src/ell/types/studio/__init__.py b/src/ell/stores/models/__init__.py similarity index 100% rename from src/ell/types/studio/__init__.py rename to src/ell/stores/models/__init__.py diff --git a/src/ell/types/studio/core.py b/src/ell/stores/models/core.py similarity index 100% rename from src/ell/types/studio/core.py rename to src/ell/stores/models/core.py diff --git a/src/ell/types/studio/evaluations.py b/src/ell/stores/models/evaluations.py similarity index 100% rename from src/ell/types/studio/evaluations.py rename to src/ell/stores/models/evaluations.py diff --git a/src/ell/stores/sql.py b/src/ell/stores/sql.py index 9efbbd0d..4ed66634 100644 --- a/src/ell/stores/sql.py +++ b/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 diff --git a/src/ell/stores/store.py b/src/ell/stores/store.py index 601861c0..73774397 100644 --- a/src/ell/stores/store.py +++ b/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): diff --git a/src/ell/studio/datamodels.py b/src/ell/studio/datamodels.py index 6269df5d..7b67f5f6 100644 --- a/src/ell/studio/datamodels.py +++ b/src/ell/studio/datamodels.py @@ -1,7 +1,7 @@ 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, @@ -9,7 +9,7 @@ EvaluationRunLabelerSummaryBase, EvaluationResultDatapointBase, ) -from ell.types.studio.core import SerializedLMPBase, InvocationBase, InvocationContentsBase +from ell.stores.models.core import SerializedLMPBase, InvocationBase, InvocationContentsBase class SerializedLMPWithUses(SerializedLMPBase): diff --git a/src/ell/studio/server.py b/src/ell/studio/server.py index 05729a8c..40ea6a80 100644 --- a/src/ell/studio/server.py +++ b/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__) diff --git a/tests/test_migrations.py b/tests/test_migrations.py index 14778010..5e0faf13 100644 --- a/tests/test_migrations.py +++ b/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""" diff --git a/tests/test_results.py b/tests/test_results.py index b4908645..5796b6a2 100644 --- a/tests/test_results.py +++ b/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(): diff --git a/tests/test_sql_store.py b/tests/test_sql_store.py index 165350a0..06f68fa6 100644 --- a/tests/test_sql_store.py +++ b/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