pixeltable 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

pixeltable/env.py

@@ -5,20 +5,20 @@ import glob
 import http.server
 import importlib
 import importlib.util
+import inspect
 import logging
 import os
-import socketserver
 import sys
 import threading
 import uuid
 import warnings
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Optional, Dict, Any, List
 
 import pgserver
 import sqlalchemy as sql
 import yaml
-from sqlalchemy_utils.functions import database_exists, create_database, drop_database
 from tqdm import TqdmWarning
 
 import pixeltable.exceptions as excs
@@ -59,12 +59,11 @@ class Env:
         # info about installed packages that are utilized by some parts of the code;
         # package name -> version; version == []: package is installed, but we haven't determined the version yet
         self._installed_packages: Dict[str, Optional[List[int]]] = {}
-        self._nos_client: Optional[Any] = None
         self._spacy_nlp: Optional[Any] = None  # spacy.Language
-        self._httpd: Optional[http.server.ThreadingHTTPServer] = None
+        self._httpd: Optional[http.server.HTTPServer] = None
         self._http_address: Optional[str] = None
 
-        self._registered_clients: dict[str, Any] = {}
+        self._registered_clients: dict[str, ApiClient] = {}
 
         # logging-related state
         self._logger = logging.getLogger('pixeltable')
@@ -120,8 +119,8 @@ class Env:
         if level is not None:
             self.set_log_level(level)
         if add is not None:
-            for module, level in [t.split(':') for t in add.split(',')]:
-                self.set_module_log_level(module, int(level))
+            for module, level_str in [t.split(':') for t in add.split(',')]:
+                self.set_module_log_level(module, int(level_str))
         if remove is not None:
             for module in remove.split(','):
                 self.set_module_log_level(module, None)
@@ -263,24 +262,19 @@ class Env:
         self._db_url = self._db_server.get_uri(database=self._db_name)
 
         if reinit_db:
-            if database_exists(self.db_url):
-                drop_database(self.db_url)
+            if self._store_db_exists():
+                self._drop_store_db()
 
-        if not database_exists(self.db_url):
+        if not self._store_db_exists():
             self._logger.info(f'creating database at {self.db_url}')
-            create_database(self.db_url)
-            self._sa_engine = sql.create_engine(self.db_url, echo=echo, future=True)
+            self._create_store_db()
+            self._create_engine(echo=echo)
             from pixeltable.metadata import schema
-
             schema.Base.metadata.create_all(self._sa_engine)
             metadata.create_system_info(self._sa_engine)
-            # enable pgvector
-            with self._sa_engine.begin() as conn:
-                conn.execute(sql.text('CREATE EXTENSION vector'))
         else:
             self._logger.info(f'found database {self.db_url}')
-            if self._sa_engine is None:
-                self._sa_engine = sql.create_engine(self.db_url, echo=echo, future=True)
+            self._create_engine(echo=echo)
 
         print(f'Connected to Pixeltable database at: {self.db_url}')
 
@@ -288,57 +282,110 @@ class Env:
         self._set_up_runtime()
         self.log_to_stdout(False)
 
-    def _upgrade_metadata(self) -> None:
-        metadata.upgrade_md(self._sa_engine)
-
-    def _create_nos_client(self) -> None:
-        import nos
+    def _create_engine(self, echo: bool = False) -> None:
+        self._sa_engine = sql.create_engine(self.db_url, echo=echo, future=True, isolation_level='AUTOCOMMIT')
 
-        self._logger.info('connecting to NOS')
-        nos.init(logging_level=logging.DEBUG)
-        self._nos_client = nos.client.InferenceClient()
-        self._logger.info('waiting for NOS')
-        self._nos_client.WaitForServer()
-
-        # now that we have a client, we can create the module
-        import importlib
+    def _store_db_exists(self) -> bool:
+        assert self._db_name is not None
+        # don't try to connect to self.db_name, it may not exist
+        db_url = self._db_server.get_uri(database='postgres')
+        engine = sql.create_engine(db_url, future=True)
+        try:
+            with engine.begin() as conn:
+                stmt = f"SELECT COUNT(*) FROM pg_database WHERE datname = '{self._db_name}'"
+                result = conn.scalar(sql.text(stmt))
+                assert result <= 1
+                return result == 1
+        finally:
+            engine.dispose()
+
+
+    def _create_store_db(self) -> None:
+        assert self._db_name is not None
+        # create the db
+        pg_db_url = self._db_server.get_uri(database='postgres')
+        engine = sql.create_engine(pg_db_url, future=True, isolation_level='AUTOCOMMIT')
+        preparer = engine.dialect.identifier_preparer
+        try:
+            with engine.begin() as conn:
+                # use C collation to get standard C/Python-style sorting
+                stmt = (
+                    f"CREATE DATABASE {preparer.quote(self._db_name)} "
+                    "ENCODING 'utf-8' LC_COLLATE 'C' LC_CTYPE 'C' TEMPLATE template0"
+                )
+                conn.execute(sql.text(stmt))
+        finally:
+            engine.dispose()
 
+        # enable pgvector
+        store_db_url = self._db_server.get_uri(database=self._db_name)
+        engine = sql.create_engine(store_db_url, future=True, isolation_level='AUTOCOMMIT')
         try:
-            importlib.import_module('pixeltable.functions.nos')
-            # it's already been created
-            return
-        except ImportError:
-            pass
-        from pixeltable.functions.util import create_nos_modules
+            with engine.begin() as conn:
+                conn.execute(sql.text('CREATE EXTENSION vector'))
+        finally:
+            engine.dispose()
+
+    def _drop_store_db(self) -> None:
+        assert self._db_name is not None
+        db_url = self._db_server.get_uri(database='postgres')
+        engine = sql.create_engine(db_url, future=True, isolation_level='AUTOCOMMIT')
+        preparer = engine.dialect.identifier_preparer
+        try:
+            with engine.begin() as conn:
+                # terminate active connections
+                stmt = (f"""
+                    SELECT pg_terminate_backend(pg_stat_activity.pid)
+                    FROM pg_stat_activity
+                    WHERE pg_stat_activity.datname = '{self._db_name}'
+                    AND pid <> pg_backend_pid()
+                """)
+                conn.execute(sql.text(stmt))
+                # drop db
+                stmt = f'DROP DATABASE {preparer.quote(self._db_name)}'
+                conn.execute(sql.text(stmt))
+        finally:
+            engine.dispose()
 
-        _ = create_nos_modules()
+    def _upgrade_metadata(self) -> None:
+        metadata.upgrade_md(self._sa_engine)
 
-    def get_client(self, name: str, init: Callable, environ: Optional[str] = None) -> Any:
-        """
-        Gets the client with the specified name, using `init` to construct one if necessary.
+    def _register_client(self, name: str, init_fn: Callable) -> None:
+        sig = inspect.signature(init_fn)
+        param_names = list(sig.parameters.keys())
+        self._registered_clients[name] = ApiClient(init_fn=init_fn, param_names=param_names)
 
-        - name: The name of the client
-        - init: A `Callable` with signature `fn(api_key: str) -> Any` that constructs a client object
-        - environ: The name of the environment variable to use for the API key, if no API key is found in config
-            (defaults to f'{name.upper()}_API_KEY')
+    def get_client(self, name: str) -> Any:
         """
-        if name in self._registered_clients:
-            return self._registered_clients[name]
+        Gets the client with the specified name, initializing it if necessary.
 
-        if environ is None:
-            environ = f'{name.upper()}_API_KEY'
-
-        if name in self._config and 'api_key' in self._config[name]:
-            api_key = self._config[name]['api_key']
-        else:
-            api_key = os.environ.get(environ)
-        if api_key is None or api_key == '':
-            raise excs.Error(f'`{name}` client not initialized (no API key configured).')
+        Args:
+            - name: The name of the client
+        """
+        cl = self._registered_clients[name]
+        if cl.client_obj is not None:
+            return cl.client_obj  # Already initialized
+
+        # Construct a client. For each client parameter, first check if the parameter is in the environment;
+        # if not, look in Pixeltable config from `config.yaml`.
+
+        init_kwargs: dict[str, str] = {}
+        for param in cl.param_names:
+            environ = f'{name.upper()}_{param.upper()}'
+            if environ in os.environ:
+                init_kwargs[param] = os.environ[environ]
+            elif name.lower() in self._config and param in self._config[name.lower()]:
+                init_kwargs[param] = self._config[name.lower()][param.lower()]
+            if param not in init_kwargs or init_kwargs[param] == '':
+                raise excs.Error(
+                    f'`{name}` client not initialized: parameter `{param}` is not configured.\n'
+                    f'To fix this, specify the `{environ}` environment variable, or put `{param.lower()}` in '
+                    f'the `{name.lower()}` section of $PIXELTABLE_HOME/config.yaml.'
+                )
 
-        client = init(api_key)
-        self._registered_clients[name] = client
+        cl.client_obj = cl.init_fn(**init_kwargs)
         self._logger.info(f'Initialized `{name}` client.')
-        return client
+        return cl.client_obj
 
     def _start_web_server(self) -> None:
         """
@@ -380,6 +427,7 @@ class Env:
         check('transformers')
         check('sentence_transformers')
         check('yolox')
+        check('whisperx')
         check('boto3')
         check('fitz')  # pymupdf
         check('pyarrow')
@@ -392,9 +440,7 @@ class Env:
         check('openai')
         check('together')
         check('fireworks')
-        check('nos')
-        if self.is_installed_package('nos'):
-            self._create_nos_client()
+        check('label_studio_sdk')
         check('openpyxl')
 
     def require_package(self, package: str, min_version: Optional[List[int]] = None) -> None:
@@ -405,7 +451,7 @@ class Env:
             return
 
         # check whether we have a version >= the required one
-        if self._installed_packages[package] == []:
+        if not self._installed_packages[package]:
             m = importlib.import_module(package)
             module_version = [int(x) for x in m.__version__.split('.')]
             self._installed_packages[package] = module_version
@@ -415,8 +461,8 @@ class Env:
         if any([a < b for a, b in zip(installed_version, normalized_min_version)]):
             raise excs.Error(
                 (
-                    f'The installed version of package {package} is {".".join([str[v] for v in installed_version])}, '
-                    f'but version  >={".".join([str[v] for v in min_version])} is required'
+                    f'The installed version of package {package} is {".".join(str(v) for v in installed_version)}, '
+                    f'but version  >={".".join(str(v) for v in min_version)} is required'
                 )
             )
 
@@ -456,11 +502,44 @@ class Env:
         assert self._sa_engine is not None
         return self._sa_engine
 
-    @property
-    def nos_client(self) -> Any:
-        return self._nos_client
-
     @property
     def spacy_nlp(self) -> Any:
         assert self._spacy_nlp is not None
         return self._spacy_nlp
+
+
+def register_client(name: str) -> Callable:
+    """Decorator that registers a third-party API client for use by Pixeltable.
+
+    The decorated function is an initialization wrapper for the client, and can have
+    any number of string parameters, with a signature such as:
+
+    ```
+    def my_client(api_key: str, url: str) -> my_client_sdk.Client:
+        return my_client_sdk.Client(api_key=api_key, url=url)
+    ```
+
+    The initialization wrapper will not be called immediately; initialization will
+    be deferred until the first time the client is used. At initialization time,
+    Pixeltable will attempt to load the client parameters from config. For each
+    config parameter:
+    - If an environment variable named MY_CLIENT_API_KEY (for example) is set, use it;
+    - Otherwise, look for 'api_key' in the 'my_client' section of config.yaml.
+
+    If all config parameters are found, Pixeltable calls the initialization function;
+    otherwise it throws an exception.
+
+    Args:
+        - name (str): The name of the API client (e.g., 'openai' or 'label-studio').
+    """
+    def decorator(fn: Callable) -> None:
+        Env.get()._register_client(name, fn)
+
+    return decorator
+
+
+@dataclass
+class ApiClient:
+    init_fn: Callable
+    param_names: list[str]
+    client_obj: Optional[Any] = None

pixeltable/exprs/column_ref.py

@@ -108,7 +108,7 @@ class ColumnRef(Expr):
     def _from_dict(cls, d: Dict, components: List[Expr]) -> Expr:
         tbl_id, version, col_id = UUID(d['tbl_id']), d['tbl_version'], d['col_id']
         tbl_version = catalog.Catalog.get().tbl_versions[(tbl_id, version)]
-        assert col_id in tbl_version.cols_by_id
-        col = tbl_version.cols_by_id[col_id]
+        # don't use tbl_version.cols_by_id here, this might be a snapshot reference to a column that was then dropped
+        col = next(col for col in tbl_version.cols if col.id == col_id)
         return cls(col)
 

pixeltable/exprs/comparison.py

@@ -4,18 +4,44 @@ from typing import Optional, List, Any, Dict, Tuple
 
 import sqlalchemy as sql
 
+from .column_ref import ColumnRef
 from .data_row import DataRow
 from .expr import Expr
 from .globals import ComparisonOperator
+from .literal import Literal
 from .predicate import Predicate
 from .row_builder import RowBuilder
+import pixeltable.exceptions as excs
+import pixeltable.index as index
 
 
 class Comparison(Predicate):
     def __init__(self, operator: ComparisonOperator, op1: Expr, op2: Expr):
         super().__init__()
         self.operator = operator
-        self.components = [op1, op2]
+
+        # if this is a comparison of a column to a literal (ie, could be used as a search argument in an index lookup),
+        # normalize it to <column> <operator> <literal>.
+        if isinstance(op1, ColumnRef) and isinstance(op2, Literal):
+            self.is_search_arg_comparison = True
+            self.components = [op1, op2]
+        elif isinstance(op1, Literal) and isinstance(op2, ColumnRef):
+            self.is_search_arg_comparison = True
+            self.components = [op2, op1]
+            self.operator = self.operator.reverse()
+        else:
+            self.is_search_arg_comparison = False
+            self.components = [op1, op2]
+
+        import pixeltable.index as index
+        if self.is_search_arg_comparison and self._op2.col_type.is_string_type() \
+                and len(self._op2.val) >= index.BtreeIndex.MAX_STRING_LEN:
+            # we can't use an index for this after all
+            raise excs.Error(
+                f'String literal too long for comparison against indexed column {self._op1.col.name!r} '
+                f'(max length is {index.BtreeIndex.MAX_STRING_LEN - 1})'
+            )
+
         self.id = self._create_id()
 
     def __str__(self) -> str:
@@ -37,6 +63,18 @@ class Comparison(Predicate):
 
     def sql_expr(self) -> Optional[sql.ClauseElement]:
         left = self._op1.sql_expr()
+        if self.is_search_arg_comparison:
+            # reference the index value column if there is an index and this is not a snapshot
+            # (indices don't apply to snapshots)
+            tbl = self._op1.col.tbl
+            idx_info = [
+                info for info in self._op1.col.get_idx_info().values() if isinstance(info.idx, index.BtreeIndex)
+            ]
+            if len(idx_info) > 0 and not tbl.is_snapshot:
+                # there shouldn't be multiple B-tree indices on a column
+                assert len(idx_info) == 1
+                left = idx_info[0].val_col.sa_col
+
         right = self._op2.sql_expr()
         if left is None or right is None:
             return None

pixeltable/exprs/data_row.py

@@ -197,3 +197,10 @@ class DataRow:
             pass
         self.vals[index] = None
 
+    @property
+    def rowid(self) -> Tuple[int]:
+        return self.pk[:-1]
+
+    @property
+    def v_min(self) -> int:
+        return self.pk[-1]

pixeltable/exprs/expr.py

@@ -169,16 +169,22 @@ class Expr(abc.ABC):
         memo[id(self)] = result
         return result
 
-    def substitute(self, old: Expr, new: Expr) -> Expr:
+    def substitute(self, spec: dict[Expr, Expr]) -> Expr:
         """
         Replace 'old' with 'new' recursively.
         """
-        if self.equals(old):
-            return new.copy()
+        for old, new in spec.items():
+            if self.equals(old):
+                return new.copy()
         for i in range(len(self.components)):
-            self.components[i] = self.components[i].substitute(old, new)
+            self.components[i] = self.components[i].substitute(spec)
         return self
 
+    @classmethod
+    def list_substitute(cls, expr_list: List[Expr], spec: dict[Expr, Expr]) -> None:
+        for i in range(len(expr_list)):
+            expr_list[i] = expr_list[i].substitute(spec)
+
     def resolve_computed_cols(self, resolve_cols: Optional[Set[catalog.Column]] = None) -> Expr:
         """
         Recursively replace ColRefs to unstored computed columns with their value exprs.
@@ -196,9 +202,7 @@ class Expr(abc.ABC):
             ])
             if len(target_col_refs) == 0:
                 return result
-            for ref in target_col_refs:
-                assert ref.col.value_expr is not None
-                result = result.substitute(ref, ref.col.value_expr)
+            result = result.substitute({ref: ref.col.value_expr for ref in target_col_refs})
 
     def is_bound_by(self, tbl: catalog.TableVersionPath) -> bool:
         """Returns True if this expr can be evaluated in the context of tbl."""
@@ -225,11 +229,6 @@ class Expr(abc.ABC):
             self.components[i] = self.components[i]._retarget(tbl_versions)
         return self
 
-    @classmethod
-    def list_substitute(cls, expr_list: List[Expr], old: Expr, new: Expr) -> None:
-        for i in range(len(expr_list)):
-            expr_list[i] = expr_list[i].substitute(old, new)
-
     @abc.abstractmethod
     def __str__(self) -> str:
         pass

pixeltable/exprs/function_call.py

@@ -174,9 +174,6 @@ class FunctionCall(Expr):
                     f'Parameter {param_name}: argument type {arg.col_type} does not match parameter type '
                     f'{param_type}')
 
-    def is_nos_call(self) -> bool:
-        return isinstance(self.fn, func.NOSFunction)
-
     def _equals(self, other: FunctionCall) -> bool:
         if self.fn != other.fn:
             return False

pixeltable/exprs/globals.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import datetime
-from typing import Union
 import enum
-
+from typing import Union
 
 # Python types corresponding to our literal types
 LiteralPythonTypes = Union[str, int, float, bool, datetime.datetime, datetime.date]
@@ -33,6 +34,17 @@ class ComparisonOperator(enum.Enum):
         if self == self.GE:
             return '>='
 
+    def reverse(self) -> ComparisonOperator:
+        if self == self.LT:
+            return self.GT
+        if self == self.LE:
+            return self.GE
+        if self == self.GT:
+            return self.LT
+        if self == self.GE:
+            return self.LE
+        return self
+
 
 class LogicalOperator(enum.Enum):
     AND = 0

pixeltable/exprs/similarity_expr.py

@@ -23,13 +23,15 @@ class SimilarityExpr(Expr):
 
         # determine index to use
         idx_info = col_ref.col.get_idx_info()
-        if len(idx_info) == 0:
+        import pixeltable.index as index
+        embedding_idx_info = [info for info in idx_info.values() if isinstance(info.idx, index.EmbeddingIndex)]
+        if len(embedding_idx_info) == 0:
             raise excs.Error(f'No index found for column {col_ref.col}')
-        if len(idx_info) > 1:
+        if len(embedding_idx_info) > 1:
             raise excs.Error(
                 f'Column {col_ref.col.name} has multiple indices; use the index name to disambiguate, '
                 f'e.g., `{col_ref.col.name}.<index-name>.similarity(...)`')
-        self.idx_info = next(iter(idx_info.values()))
+        self.idx_info = embedding_idx_info[0]
         idx = self.idx_info.idx
 
         if item.col_type.is_string_type() and idx.txt_embed is None:

pixeltable/ext/functions/whisperx.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+import torch
+import whisperx
+from whisperx.asr import FasterWhisperPipeline
+
+import pixeltable as pxt
+
+
+@pxt.udf(param_types=[pxt.AudioType(), pxt.StringType(), pxt.StringType(), pxt.StringType(), pxt.IntType()])
+def transcribe(
+    audio: str, *, model: str, compute_type: Optional[str] = None, language: Optional[str] = None, chunk_size: int = 30
+) -> dict:
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    compute_type = compute_type or ('float16' if device == 'cuda' else 'int8')
+    model = _lookup_model(model, device, compute_type)
+    audio_array = whisperx.load_audio(audio)
+    result = model.transcribe(audio_array, batch_size=16, language=language, chunk_size=chunk_size)
+    return result
+
+
+def _lookup_model(model_id: str, device: str, compute_type: str) -> FasterWhisperPipeline:
+    key = (model_id, device, compute_type)
+    if key not in _model_cache:
+        model = whisperx.load_model(model_id, device, compute_type=compute_type)
+        _model_cache[key] = model
+    return _model_cache[key]
+
+
+_model_cache = {}

pixeltable/ext/functions/yolox.py

@@ -56,6 +56,22 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
     return results
 
 
+@pxt.udf
+def yolo_to_coco(detections: dict) -> list:
+    bboxes, labels = detections['bboxes'], detections['labels']
+    num_annotations = len(detections['bboxes'])
+    assert num_annotations == len(detections['labels'])
+    result = []
+    for i in range(num_annotations):
+        bbox = bboxes[i]
+        ann = {
+            'bbox': [round(bbox[0]), round(bbox[1]), round(bbox[2] - bbox[0]), round(bbox[3] - bbox[1])],
+            'category': labels[i],
+        }
+        result.append(ann)
+    return result
+
+
 def _images_to_tensors(images: Iterable[PIL.Image.Image], exp: Exp) -> Iterator[torch.Tensor]:
     for image in images:
         image_transform, _ = _val_transform(np.array(image), None, exp.test_size)

pixeltable/func/aggregate_function.py

@@ -140,7 +140,7 @@ def uda(
         update_types: List[ts.ColumnType],
         init_types: Optional[List[ts.ColumnType]] = None,
         requires_order_by: bool = False, allows_std_agg: bool = True, allows_window: bool = False,
-) -> Callable:
+) -> Callable[[Type[Aggregator]], AggregateFunction]:
     """Decorator for user-defined aggregate functions.
 
     The decorated class must inherit from Aggregator and implement the following methods:
@@ -162,7 +162,7 @@ def uda(
     if init_types is None:
         init_types = []
 
-    def decorator(cls: Type[Aggregator]) -> Type[Function]:
+    def decorator(cls: Type[Aggregator]) -> AggregateFunction:
         # validate type parameters
         num_init_params = len(inspect.signature(cls.__init__).parameters) - 1
         if num_init_params > 0:

pixeltable/func/expr_template_function.py

@@ -50,6 +50,7 @@ class ExprTemplateFunction(Function):
             {param_name: default for param_name, default in self.defaults.items() if param_name not in bound_args})
         result = self.expr.copy()
         import pixeltable.exprs as exprs
+        arg_exprs: dict[exprs.Expr, exprs.Expr] = {}
         for param_name, arg in bound_args.items():
             param_expr = self.param_exprs_by_name[param_name]
             if not isinstance(arg, exprs.Expr):
@@ -59,7 +60,8 @@ class ExprTemplateFunction(Function):
                     raise excs.Error(f'{self.self_name}(): cannot convert argument {arg} to a Pixeltable expression')
             else:
                 arg_expr = arg
-            result = result.substitute(param_expr, arg_expr)
+            arg_exprs[param_expr] = arg_expr
+        result = result.substitute(arg_exprs)
         import pixeltable.exprs as exprs
         assert not result.contains(exprs.Variable)
         return result

pixeltable/func/udf.py

@@ -28,7 +28,7 @@ def udf(
         batch_size: Optional[int] = None,
         substitute_fn: Optional[Callable] = None,
         _force_stored: bool = False
-) -> Callable: ...
+) -> Callable[[Callable], Function]: ...
 
 
 def udf(*args, **kwargs):
@@ -131,7 +131,7 @@ def make_function(
 def expr_udf(py_fn: Callable) -> ExprTemplateFunction: ...
 
 @overload
-def expr_udf(*, param_types: Optional[List[ts.ColumnType]] = None) -> Callable: ...
+def expr_udf(*, param_types: Optional[List[ts.ColumnType]] = None) -> Callable[[Callable], ExprTemplateFunction]: ...
 
 def expr_udf(*args: Any, **kwargs: Any) -> Any:
     def decorator(py_fn: Callable, param_types: Optional[List[ts.ColumnType]]) -> ExprTemplateFunction:

pixeltable/functions/fireworks.py

@@ -6,8 +6,13 @@ import pixeltable as pxt
 from pixeltable import env
 
 
-def fireworks_client() -> fireworks.client.Fireworks:
-    return env.Env.get().get_client('fireworks', lambda api_key: fireworks.client.Fireworks(api_key=api_key))
+@env.register_client('fireworks')
+def _(api_key: str) -> fireworks.client.Fireworks:
+    return fireworks.client.Fireworks(api_key=api_key)
+
+
+def _fireworks_client() -> fireworks.client.Fireworks:
+    return env.Env.get().get_client('fireworks')
 
 
 @pxt.udf
@@ -26,8 +31,8 @@ def chat_completions(
         'top_p': top_p,
         'temperature': temperature
     }
-    kwargs_not_none = dict(filter(lambda x: x[1] is not None, kwargs.items()))
-    return fireworks_client().chat.completions.create(
+    kwargs_not_none = {k: v for k, v in kwargs.items() if v is not None}
+    return _fireworks_client().chat.completions.create(
         model=model,
         messages=messages,
         **kwargs_not_none