Merge pull request #343 from alex-dixon/studio-refactor

Separate storage, studio dependencies

docs/src/installation.rst

@@ -10,20 +10,87 @@ Installing ell
 
    .. code-block:: bash
 
-      pip install -U ell-ai
+      pip install -U ell-ai[all]
 
-   By default, this installs only the OpenAI client SDK. If you want to include the Anthropic client SDK, use the "anthropic" extra like so:
-
-   .. code-block:: bash
-
-      pip install -U 'ell-ai[anthropic]'
+   This installs ``ell``, ``ell-studio``, versioning and tracing with SQLite, and the default provider clients.
 
 2. Verify installation:
 
    .. code-block:: bash
 
       python -c "import ell; print(ell.__version__)"
 
+Custom Installation
+-------------------
+
+You can create a custom ``ell`` installation with the following options.
+
+Install ``ell`` without storage or ``ell-studio`` and with the default OpenAI client:
+
+.. code-block:: bash
+
+   pip install -U ell-ai
+
+Supported options:
+
+``anthropic``
+~~~~~~~~~~~~~
+Adds the Anthropic client.
+
+.. code-block:: bash
+
+   pip install -U ell-ai[anthropic]
+
+
+``groq``
+~~~~~~~~
+Adds the Groq client.
+
+.. code-block:: bash
+
+   pip install -U ell-ai[groq]
+
+
+``studio``
+~~~~~~~~~~
+Adds ``ell-studio``.
+
+.. code-block:: bash
+
+   pip install -U ell-ai[studio]
+
+
+``sqlite``
+~~~~~~~~~~
+SQLite storage for versioning and tracing.
+
+.. code-block:: bash
+
+   pip install -U ell-ai[sqlite]
+
+
+``postgres``
+~~~~~~~~~~~~
+Postgres storage for versioning and tracing.
+
+Include this option if you'd like to use ``ell-studio`` with Postgres.
+
+.. code-block:: bash
+
+   pip install -U ell-ai[postgres]
+
+Combining options
+~~~~~~~~~~~~~~~~~
+
+All options are additive and can be combined as needed.
+
+Example: Install ``ell`` with ``ell-studio``, Postgres, and the Anthropic client:
+
+.. code-block:: bash
+
+   pip install -U ell-ai[studio, postgres, anthropic]
+
+
 API Key Setup
 -------------
 

pyproject.toml

@@ -27,21 +27,25 @@ include = [
 
 [tool.poetry.dependencies]
 python = ">=3.9"
-fastapi = "^0.111.1"
 numpy = ">=1.26.0"
 dill = "^0.3.8"
 colorama = "^0.4.6"
 cattrs = "^23.2.3"
 openai = "^1.51.0"
-anthropic = { version = "^0.34.2", optional = true }
-groq = { version = "^0.11.0", optional = true }
-sqlmodel = ">=0.0.21, <0.1.0"
-uvicorn = "^0.30.3"
 requests = "^2.32.3"
 typing-extensions = "^4.12.2"
 black = "^24.8.0"
 pillow = "^10.4.0"
 psutil = "^5.9.0"
+# Providers
+anthropic = { version = "^0.34.2", optional = true }
+groq = { version = "^0.11.0", optional = true }
+# Storage
+psycopg2 = { version = ">=2.7", optional = true }
+sqlmodel = { version = ">=0.0.21, <0.1.0", optional = true }
+# Studio
+fastapi = { version = "^0.111.1", optional = true }
+uvicorn = { version = "^0.30.3", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.2"
@@ -53,7 +57,18 @@ sphinx-rtd-theme = "^2.0.0"
 # causes poetry to mark it as optional = true (even if explicitly specified optional = false).
 anthropic = ["anthropic"]
 groq = ["groq"]
-all = ["anthropic", "groq"]
+sqlite = [ 'sqlmodel' ]
+postgres = ['sqlmodel', 'psycopg2']
+studio = ['fastapi', 'uvicorn', 'sqlmodel']
+all = [
+    "anthropic",
+    "groq",
+    # default storage dependencies
+    'sqlmodel',
+    # allow running stduio by default
+    'fastapi',
+    'uvicorn',
+]
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

src/ell/configurator.py

@@ -1,14 +1,19 @@
 from functools import lru_cache, wraps
-from typing import Dict, Any, Optional, Tuple, Union, Type
+from typing import Dict, Any, Optional, Tuple, Union, Type, TYPE_CHECKING
 import openai
 import logging
 from contextlib import contextmanager
 import threading
 from pydantic import BaseModel, ConfigDict, Field
-from ell.store import Store
 from ell.provider import Provider
 from dataclasses import dataclass, field
 
+if TYPE_CHECKING:
+    from ell.stores import Store
+else:
+    Store = None
+
+
 _config_logger = logging.getLogger(__name__)
 
 @dataclass(frozen=True)
@@ -175,8 +180,11 @@ def init(
     config.lazy_versioning = lazy_versioning
 
     if isinstance(store, str):
-        from ell.stores.sql import SQLiteStore
-        config.store = SQLiteStore(store)
+        try:
+            from ell.stores.sql import SQLiteStore
+            config.store = SQLiteStore(store)
+        except ImportError:
+            raise ImportError("Failed importing SQLiteStore. Install with `pip install -U ell-ai[all]`. More info: https://docs.ell.so/installation")
     else:
         config.store = store
     config.autocommit = autocommit or config.autocommit

src/ell/lmp/_track.py

@@ -1,25 +1,26 @@
 import json
 import logging
 import threading
-from ell.types import SerializedLMP, Invocation, InvocationTrace, InvocationContents
-from ell.types.studio import LMPType, utc_now
+from ell.types.lmp import LMPType
 from ell.util._warnings import _autocommit_warning
 import ell.util.closure
 from ell.configurator import config
-from ell.types._lstr import _lstr
 
 import inspect
 
 import secrets
-import time
-from datetime import datetime
 from functools import wraps
-from typing import Any, Callable, Dict, Iterable, Optional, OrderedDict, Tuple
+from typing import Any, Callable, Dict, Optional
 
 from ell.util.serialization import get_immutable_vars
 from ell.util.serialization import compute_state_cache_key
 from ell.util.serialization import prepare_invocation_params
 
+try:
+    from ell.stores.studio import SerializedLMP, Invocation, InvocationContents, utc_now
+except ImportError:
+    SerializedLMP = Invocation =  InvocationContents = utc_now = None
+
 logger = logging.getLogger(__name__)
 
 # Thread-local storage for the invocation stack

src/ell/lmp/complex.py

@@ -3,10 +3,10 @@
 from ell.provider import EllCallParams
 from ell.types._lstr import _lstr
 from ell.types import Message, ContentBlock
-from ell.types.message import LMP, InvocableLM, LMPParams, MessageOrDict, _lstr_generic
-from ell.types.studio import LMPType
+from ell.types.message import LMP, MessageOrDict
+from ell.types.lmp import LMPType
 from ell.util._warnings import _no_api_key_warning, _warnings
-from ell.util.verbosity import compute_color, model_usage_logger_pre
+from ell.util.verbosity import model_usage_logger_pre
 
 from ell.util.verbosity import model_usage_logger_post_end, model_usage_logger_post_intermediate, model_usage_logger_post_start
 

src/ell/lmp/tool.py

@@ -1,6 +1,6 @@
 from functools import wraps
 import json
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 
 from pydantic import Field, create_model
 from pydantic.fields import FieldInfo
@@ -9,11 +9,10 @@
 # from ell.util.verbosity import compute_color, tool_usage_logger_pre
 from ell.configurator import config
 from ell.types._lstr import _lstr
-from ell.types.studio import LMPType
+from ell.types.lmp import LMPType
 import inspect
 
-from ell.types.message import ContentBlock, InvocableTool, ToolResult, to_content_blocks
-
+from ell.types.message import ContentBlock, InvocableTool, ToolResult
 
 
 def tool(*, exempt_from_tracking: bool = False, **tool_kwargs):

src/ell/stores/__init__.py

@@ -0,0 +1,4 @@
+try:
+    import sqlmodel
+except ImportError:
+    raise ImportError("ell.stores has missing dependencies. Install them with `pip install -U ell-ai[sqlite]` or `pip install -U ell-ai[postgres]`. More info: https://docs.ell.so/installation/custom-installation")

src/ell/stores/sql.py

@@ -1,24 +1,17 @@
 from datetime import datetime, timedelta
-import json
 import os
-from typing import Any, Optional, Dict, List, Set, Union
-from pydantic import BaseModel
+from typing import Any, Optional, Dict, List, Set
 from sqlmodel import Session, SQLModel, create_engine, select
-import ell.store
-import cattrs
-import numpy as np
+import ell.stores.store
 from sqlalchemy.sql import text
-from ell.types import InvocationTrace, SerializedLMP, Invocation, InvocationContents
-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.stores.studio import InvocationTrace, SerializedLMP, Invocation
+from sqlalchemy import func, and_
 from ell.util.serialization import pydantic_ltype_aware_cattr
 import gzip
 import json
 
-class SQLStore(ell.store.Store):
-    def __init__(self, db_uri: str, blob_store: Optional[ell.store.BlobStore] = None):
+class SQLStore(ell.stores.store.Store):
+    def __init__(self, db_uri: str, blob_store: Optional[ell.stores.store.BlobStore] = None):
         self.engine = create_engine(db_uri,
                                     json_serializer=lambda obj: json.dumps(pydantic_ltype_aware_cattr.unstructure(obj), 
                                      sort_keys=True, default=repr, ensure_ascii=False))
@@ -218,7 +211,7 @@ def __init__(self, db_dir: str):
         blob_store = SQLBlobStore(db_dir)
         super().__init__(f'sqlite:///{db_path}', blob_store=blob_store)
 
-class SQLBlobStore(ell.store.BlobStore):
+class SQLBlobStore(ell.stores.store.BlobStore):
     def __init__(self, db_dir: str):
         self.db_dir = db_dir
 

src/ell/store.py → src/ell/stores/store.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from typing import Any, Optional, Dict, List, Set, Union
 from ell.types._lstr import _lstr
-from ell.types import SerializedLMP, Invocation
+from ell.stores.studio import SerializedLMP, Invocation
 from ell.types.message import InvocableLM
 
 class BlobStore(ABC):

src/ell/types/studio.py → src/ell/stores/studio.py

@@ -4,9 +4,9 @@
 
 import sqlalchemy.types as types
 
+from ell.types.lmp import LMPType
 from ell.types.message import Any, Any, Field, Message, Optional
 
-from sqlmodel import Column, Field, SQLModel
 from typing import Optional
 from dataclasses import dataclass
 from typing import Dict, List, Literal, Union, Any, Optional
@@ -51,14 +51,6 @@ def UTCTimestampField(index:bool=False, **kwargs:Any):
         sa_column=Column(UTCTimestamp(timezone=True), index=index, **kwargs))
 
 
-class LMPType(str, enum.Enum):
-    LM = "LM"
-    TOOL = "TOOL"
-    MULTIMODAL = "MULTIMODAL"
-    OTHER = "OTHER"
-
-
-
 class SerializedLMPBase(SQLModel):
     lmp_id: Optional[str] = Field(default=None, primary_key=True)
     name: str = Field(index=True)

src/ell/studio/__init__.py

@@ -0,0 +1,5 @@
+try:
+    import fastapi
+    import ell.stores
+except ImportError:
+    raise ImportError("ell.studio is missing dependencies. Install them with `pip install -U ell-ai[studio]. More info: https://docs.ell.so/installation/custom-installation")

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 import SerializedLMPBase, InvocationBase, InvocationContentsBase
+from ell.stores.studio import SerializedLMPBase, InvocationBase, InvocationContentsBase
 
 
 class SerializedLMPWithUses(SerializedLMPBase):

src/ell/studio/server.py

@@ -11,7 +11,7 @@
 from ell.studio.connection_manager import ConnectionManager
 from ell.studio.datamodels import InvocationPublicWithConsumes, SerializedLMPWithUses
 
-from ell.types import SerializedLMP
+from ell.stores.studio import SerializedLMP
 from datetime import datetime, timedelta
 from sqlmodel import select
 

src/ell/types/__init__.py

@@ -17,19 +17,6 @@
     _content_to_text_only,
     _content_to_text,
 )
-from ell.types.studio import (
-    utc_now,
-    SerializedLMPUses,
-    UTCTimestampField,
-    LMPType,
-    SerializedLMPBase,
-    SerializedLMP,
-    InvocationTrace,
-    InvocationBase,
-    InvocationContentsBase,
-    InvocationContents,
-    Invocation,
-)
 from ell.types._lstr import _lstr
 
 __all__ = [
@@ -46,16 +33,5 @@
     "assistant",
     "_content_to_text_only",
     "_content_to_text",
-    "utc_now",
-    "SerializedLMPUses",
-    "UTCTimestampField",
-    "LMPType",
-    "SerializedLMPBase",
-    "SerializedLMP",
-    "InvocationTrace",
-    "InvocationBase",
-    "InvocationContentsBase",
-    "InvocationContents",
-    "Invocation",
     "_lstr",
 ]