pixeltable 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

pixeltable/functions/openai.py

@@ -1,9 +1,14 @@
 import base64
 import io
-from typing import Optional
+import pathlib
+import uuid
+from typing import Optional, TypeVar, Union, Callable
 
 import PIL.Image
 import numpy as np
+import openai
+import tenacity
+from openai._types import NOT_GIVEN, NotGiven
 
 import pixeltable as pxt
 import pixeltable.type_system as ts
@@ -11,43 +16,148 @@ from pixeltable import env
 from pixeltable.func import Batch
 
 
+def openai_client() -> openai.OpenAI:
+    return env.Env.get().get_client('openai', lambda api_key: openai.OpenAI(api_key=api_key))
+
+
+# Exponential backoff decorator using tenacity.
+# TODO(aaron-siegel): Right now this hardwires random exponential backoff with defaults suggested
+# by OpenAI. Should we investigate making this more customizable in the future?
+def _retry(fn: Callable) -> Callable:
+    return tenacity.retry(
+        retry=tenacity.retry_if_exception_type(openai.RateLimitError),
+        wait=tenacity.wait_random_exponential(multiplier=3, max=180),
+        stop=tenacity.stop_after_attempt(20)
+    )(fn)
+
+
+#####################################
+# Audio Endpoints
+
+@pxt.udf(return_type=ts.AudioType())
+@_retry
+def speech(
+        input: str,
+        *,
+        model: str,
+        voice: str,
+        response_format: Optional[str] = None,
+        speed: Optional[float] = None
+) -> str:
+    content = openai_client().audio.speech.create(
+        input=input,
+        model=model,
+        voice=voice,
+        response_format=_opt(response_format),
+        speed=_opt(speed)
+    )
+    ext = response_format or 'mp3'
+    output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
+    content.stream_to_file(output_filename, chunk_size=1 << 20)
+    return output_filename
+
+
+@pxt.udf(
+    param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True),
+                 ts.StringType(nullable=True), ts.FloatType(nullable=True)]
+)
+@_retry
+def transcriptions(
+        audio: str,
+        *,
+        model: str,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    file = pathlib.Path(audio)
+    transcription = openai_client().audio.transcriptions.create(
+        file=file,
+        model=model,
+        language=_opt(language),
+        prompt=_opt(prompt),
+        temperature=_opt(temperature)
+    )
+    return transcription.dict()
+
+
+@pxt.udf(
+    param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True), ts.FloatType(nullable=True)]
+)
+@_retry
+def translations(
+        audio: str,
+        *,
+        model: str,
+        prompt: Optional[str] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    file = pathlib.Path(audio)
+    translation = openai_client().audio.translations.create(
+        file=file,
+        model=model,
+        prompt=_opt(prompt),
+        temperature=_opt(temperature)
+    )
+    return translation.dict()
+
+
+#####################################
+# Chat Endpoints
+
 @pxt.udf
+@_retry
 def chat_completions(
         messages: list,
+        *,
         model: str,
         frequency_penalty: Optional[float] = None,
-        logit_bias: Optional[dict] = None,
+        logit_bias: Optional[dict[str, int]] = None,
+        logprobs: Optional[bool] = None,
+        top_logprobs: Optional[int] = None,
         max_tokens: Optional[int] = None,
         n: Optional[int] = None,
         presence_penalty: Optional[float] = None,
         response_format: Optional[dict] = None,
         seed: Optional[int] = None,
+        stop: Optional[list[str]] = None,
+        temperature: Optional[float] = None,
         top_p: Optional[float] = None,
-        temperature: Optional[float] = None
+        tools: Optional[list[dict]] = None,
+        tool_choice: Optional[dict] = None,
+        user: Optional[str] = None
 ) -> dict:
-    from openai._types import NOT_GIVEN
-    result = env.Env.get().openai_client.chat.completions.create(
+    result = openai_client().chat.completions.create(
         messages=messages,
         model=model,
-        frequency_penalty=frequency_penalty if frequency_penalty is not None else NOT_GIVEN,
-        logit_bias=logit_bias if logit_bias is not None else NOT_GIVEN,
-        max_tokens=max_tokens if max_tokens is not None else NOT_GIVEN,
-        n=n if n is not None else NOT_GIVEN,
-        presence_penalty=presence_penalty if presence_penalty is not None else NOT_GIVEN,
-        response_format=response_format if response_format is not None else NOT_GIVEN,
-        seed=seed if seed is not None else NOT_GIVEN,
-        top_p=top_p if top_p is not None else NOT_GIVEN,
-        temperature=temperature if temperature is not None else NOT_GIVEN
+        frequency_penalty=_opt(frequency_penalty),
+        logit_bias=_opt(logit_bias),
+        logprobs=_opt(logprobs),
+        top_logprobs=_opt(top_logprobs),
+        max_tokens=_opt(max_tokens),
+        n=_opt(n),
+        presence_penalty=_opt(presence_penalty),
+        response_format=_opt(response_format),
+        seed=_opt(seed),
+        stop=_opt(stop),
+        temperature=_opt(temperature),
+        top_p=_opt(top_p),
+        tools=_opt(tools),
+        tool_choice=_opt(tool_choice),
+        user=_opt(user)
     )
     return result.dict()
 
 
 @pxt.udf
+@_retry
 def vision(
         prompt: str,
         image: PIL.Image.Image,
+        *,
         model: str = 'gpt-4-vision-preview'
 ) -> str:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
     bytes_arr = io.BytesIO()
     image.save(bytes_arr, format='png')
     b64_bytes = base64.b64encode(bytes_arr.getvalue())
@@ -61,28 +171,86 @@ def vision(
              }}
          ]}
     ]
-    result = env.Env.get().openai_client.chat.completions.create(
+    result = openai_client().chat.completions.create(
         messages=messages,
         model=model
     )
     return result.choices[0].message.content
 
 
-@pxt.udf
-def moderations(input: str, model: Optional[str] = None) -> dict:
-    result = env.Env().get().openai_client.moderations.create(input=input, model=model)
-    return result.dict()
-
+#####################################
+# Embeddings Endpoints
 
 @pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
-def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
-    result = env.Env().get().openai_client.embeddings.create(
+@_retry
+def embeddings(
+        input: Batch[str],
+        *,
+        model: str,
+        user: Optional[str] = None
+) -> Batch[np.ndarray]:
+    result = openai_client().embeddings.create(
         input=input,
         model=model,
+        user=_opt(user),
         encoding_format='float'
     )
-    embeddings = [
+    return [
         np.array(data.embedding, dtype=np.float64)
         for data in result.data
     ]
-    return embeddings
+
+
+#####################################
+# Images Endpoints
+
+@pxt.udf
+@_retry
+def image_generations(
+        prompt: str,
+        *,
+        model: Optional[str] = None,
+        quality: Optional[str] = None,
+        size: Optional[str] = None,
+        style: Optional[str] = None,
+        user: Optional[str] = None
+) -> PIL.Image.Image:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
+    result = openai_client().images.generate(
+        prompt=prompt,
+        model=_opt(model),
+        quality=_opt(quality),
+        size=_opt(size),
+        style=_opt(style),
+        user=_opt(user),
+        response_format="b64_json"
+    )
+    b64_str = result.data[0].b64_json
+    b64_bytes = base64.b64decode(b64_str)
+    img = PIL.Image.open(io.BytesIO(b64_bytes))
+    img.load()
+    return img
+
+
+#####################################
+# Moderations Endpoints
+
+@pxt.udf
+@_retry
+def moderations(
+        input: str,
+        *,
+        model: Optional[str] = None
+) -> dict:
+    result = openai_client().moderations.create(
+        input=input,
+        model=_opt(model)
+    )
+    return result.dict()
+
+
+_T = TypeVar('_T')
+
+
+def _opt(arg: _T) -> Union[_T, NotGiven]:
+    return arg if arg is not None else NOT_GIVEN

pixeltable/functions/together.py

@@ -1,27 +1,122 @@
+import base64
+import io
 from typing import Optional
 
+import PIL.Image
+import numpy as np
+import together
+
 import pixeltable as pxt
+from pixeltable import env
+from pixeltable.func import Batch
+
+
+def together_client() -> together.Together:
+    return env.Env.get().get_client('together', lambda api_key: together.Together(api_key=api_key))
 
 
 @pxt.udf
 def completions(
         prompt: str,
+        *,
         model: str,
         max_tokens: Optional[int] = None,
-        repetition_penalty: Optional[float] = None,
         stop: Optional[list] = None,
-        top_k: Optional[int] = None,
+        temperature: Optional[float] = None,
         top_p: Optional[float] = None,
-        temperature: Optional[float] = None
+        top_k: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        logprobs: Optional[int] = None,
+        echo: Optional[bool] = None,
+        n: Optional[int] = None,
+        safety_model: Optional[str] = None
 ) -> dict:
-    import together
-    return together.Complete.create(
-        prompt,
-        model,
+    return together_client().completions.create(
+        prompt=prompt,
+        model=model,
         max_tokens=max_tokens,
-        repetition_penalty=repetition_penalty,
         stop=stop,
+        temperature=temperature,
+        top_p=top_p,
         top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        logprobs=logprobs,
+        echo=echo,
+        n=n,
+        safety_model=safety_model
+    ).dict()
+
+
+@pxt.udf
+def chat_completions(
+        messages: list[dict[str, str]],
+        *,
+        model: str,
+        max_tokens: Optional[int] = None,
+        stop: Optional[list[str]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        logprobs: Optional[int] = None,
+        echo: Optional[bool] = None,
+        n: Optional[int] = None,
+        safety_model: Optional[str] = None,
+        response_format: Optional[dict] = None,
+        tools: Optional[dict] = None,
+        tool_choice: Optional[dict] = None
+) -> dict:
+    return together_client().chat.completions.create(
+        messages=messages,
+        model=model,
+        max_tokens=max_tokens,
+        stop=stop,
+        temperature=temperature,
         top_p=top_p,
-        temperature=temperature
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        logprobs=logprobs,
+        echo=echo,
+        n=n,
+        safety_model=safety_model,
+        response_format=response_format,
+        tools=tools,
+        tool_choice=tool_choice
+    ).dict()
+
+
+@pxt.udf(batch_size=32, return_type=pxt.ArrayType((None,), dtype=pxt.FloatType()))
+def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
+    result = together_client().embeddings.create(input=input, model=model)
+    return [
+        np.array(data.embedding, dtype=np.float64)
+        for data in result.data
+    ]
+
+
+@pxt.udf
+def image_generations(
+        prompt: str,
+        *,
+        model: str,
+        steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        negative_prompt: Optional[str] = None,
+) -> PIL.Image.Image:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
+    result = together_client().images.generate(
+        prompt=prompt,
+        model=model,
+        steps=steps,
+        seed=seed,
+        height=height,
+        width=width,
+        negative_prompt=negative_prompt
     )
+    b64_str = result.data[0].b64_json
+    b64_bytes = base64.b64decode(b64_str)
+    img = PIL.Image.open(io.BytesIO(b64_bytes))
+    img.load()
+    return img

pixeltable/functions/util.py

@@ -39,3 +39,14 @@ def create_nos_modules() -> List[types.ModuleType]:
         setattr(sub_module, model_id, pt_func)
 
     return new_modules
+
+
+def resolve_torch_device(device: str) -> str:
+    import torch
+    if device == 'auto':
+        if torch.cuda.is_available():
+            return 'cuda'
+        if torch.backends.mps.is_available():
+            return 'mps'
+        return 'cpu'
+    return device

pixeltable/index/__init__.py

@@ -0,0 +1,2 @@
+from .base import IndexBase
+from .embedding_index import EmbeddingIndex

pixeltable/index/base.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import abc
+from typing import Any
+
+import sqlalchemy as sql
+
+import pixeltable.catalog as catalog
+
+
+class IndexBase(abc.ABC):
+    """
+    Internal interface used by the catalog and runtime system to interact with indices:
+    - types and expressions needed to create and populate the index value column
+    - creating/dropping the index
+    - TODO: translating queries into sqlalchemy predicates
+    """
+    @abc.abstractmethod
+    def __init__(self, c: catalog.Column, **kwargs: Any):
+        pass
+
+    @abc.abstractmethod
+    def index_value_expr(self) -> 'pixeltable.exprs.Expr':
+        """Return expression that computes the value that goes into the index"""
+        pass
+
+    @abc.abstractmethod
+    def index_sa_type(self) -> sql.sqltypes.TypeEngine:
+        """Return the sqlalchemy type of the index value column"""
+        pass
+
+    @abc.abstractmethod
+    def create_index(self, index_name: str, index_value_col: catalog.Column, conn: sql.engine.Connection) -> None:
+        """Create the index on the index value column"""
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def display_name(cls) -> str:
+        pass
+
+    @abc.abstractmethod
+    def as_dict(self) -> dict:
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def from_dict(cls, c: catalog.Column, d: dict) -> IndexBase:
+        pass

pixeltable/index/embedding_index.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import pgvector.sqlalchemy
+import sqlalchemy as sql
+
+import pixeltable.catalog as catalog
+import pixeltable.exceptions as excs
+import pixeltable.func as func
+import pixeltable.type_system as ts
+from .base import IndexBase
+
+
+class EmbeddingIndex(IndexBase):
+    """
+    Internal interface used by the catalog and runtime system to interact with (embedding) indices:
+    - types and expressions needed to create and populate the index value column
+    - creating/dropping the index
+    - translating 'matches' queries into sqlalchemy predicates
+    """
+
+    def __init__(
+            self, c: catalog.Column, text_embed: Optional[func.Function] = None,
+            img_embed: Optional[func.Function] = None):
+        if not c.col_type.is_string_type() and not c.col_type.is_image_type():
+            raise excs.Error(f'Embedding index requires string or image column')
+        if c.col_type.is_string_type() and text_embed is None:
+                raise excs.Error(f'Text embedding function is required for column {c.name} (parameter `txt_embed`)')
+        if c.col_type.is_image_type() and img_embed is None:
+            raise excs.Error(f'Image embedding function is required for column {c.name} (parameter `img_embed`)')
+        if text_embed is not None:
+            # verify signature
+            self._validate_embedding_fn(text_embed, 'txt_embed', ts.ColumnType.Type.STRING)
+        if img_embed is not None:
+            # verify signature
+            self._validate_embedding_fn(img_embed, 'img_embed', ts.ColumnType.Type.IMAGE)
+
+        from pixeltable.exprs import ColumnRef
+        self.value_expr = text_embed(ColumnRef(c)) if c.col_type.is_string_type() else img_embed(ColumnRef(c))
+        assert self.value_expr.col_type.is_array_type()
+        self.txt_embed = text_embed
+        self.img_embed = img_embed
+        vector_size = self.value_expr.col_type.shape[0]
+        assert vector_size is not None
+        self.index_col_type = pgvector.sqlalchemy.Vector(vector_size)
+
+    def index_value_expr(self) -> 'pixeltable.exprs.Expr':
+        """Return expression that computes the value that goes into the index"""
+        return self.value_expr
+
+    def index_sa_type(self) -> sql.sqltypes.TypeEngine:
+        """Return the sqlalchemy type of the index value column"""
+        return self.index_col_type
+
+    def create_index(self, index_name: str, index_value_col: catalog.Column, conn: sql.engine.Connection) -> None:
+        """Create the index on the index value column"""
+        idx = sql.Index(
+            index_name, index_value_col.sa_col,
+            postgresql_using='hnsw',
+            postgresql_with={'m': 16, 'ef_construction': 64},
+            postgresql_ops={index_value_col.sa_col.name: 'vector_cosine_ops'}
+        )
+        idx.create(bind=conn)
+
+    @classmethod
+    def display_name(cls) -> str:
+        return 'embedding'
+
+    @classmethod
+    def _validate_embedding_fn(cls, embed_fn: func.Function, name: str, expected_type: ts.ColumnType.Type) -> None:
+        """Validate the signature"""
+        assert isinstance(embed_fn, func.Function)
+        sig = embed_fn.signature
+        if not sig.return_type.is_array_type():
+            raise excs.Error(f'{name} must return an array, but returns {sig.return_type}')
+        else:
+            shape = sig.return_type.shape
+            if len(shape) != 1 or shape[0] == None:
+                raise excs.Error(f'{name} must return a 1D array of a specific length, but returns {sig.return_type}')
+        if len(sig.parameters) != 1 or sig.parameters_by_pos[0].col_type.type_enum != expected_type:
+            raise excs.Error(
+                f'{name} must take a single {expected_type.name.lower()} parameter, but has signature {sig}')
+
+    def as_dict(self) -> dict:
+        return {
+            'txt_embed': None if self.txt_embed is None else self.txt_embed.as_dict(),
+            'img_embed': None if self.img_embed is None else self.img_embed.as_dict()
+        }
+
+    @classmethod
+    def from_dict(cls, c: catalog.Column, d: dict) -> EmbeddingIndex:
+        txt_embed = func.Function.from_dict(d['txt_embed']) if d['txt_embed'] is not None else None
+        img_embed = func.Function.from_dict(d['img_embed']) if d['img_embed'] is not None else None
+        return cls(c, text_embed=txt_embed, img_embed=img_embed)

pixeltable/metadata/schema.py

@@ -1,4 +1,4 @@
-from typing import Optional, List, Dict, get_type_hints, Type, Any, TypeVar, Tuple, Union
+from typing import Optional, List, get_type_hints, Type, Any, TypeVar, Tuple, Union
 import platform
 import uuid
 import dataclasses
@@ -71,16 +71,43 @@ class Dir(Base):
 
 
 @dataclasses.dataclass
-class ColumnHistory:
+class ColumnMd:
     """
-    Records when a column was added/dropped, which is needed to GC unreachable storage columns
-    (a column that was added after table snapshot n and dropped before table snapshot n+1 can be removed
-    from the stored table).
-    One record per column (across all schema versions).
+    Records the non-versioned metadata of a column.
+    - immutable attributes: type, primary key, etc.
+    - when a column was added/dropped, which is needed to GC unreachable storage columns
+      (a column that was added after table snapshot n and dropped before table snapshot n+1 can be removed
+      from the stored table).
      """
-    col_id: int
+    id: int
     schema_version_add: int
     schema_version_drop: Optional[int]
+    col_type: dict
+
+    # if True, is part of the primary key
+    is_pk: bool
+
+    # if set, this is a computed column
+    value_expr: Optional[dict]
+
+    # if True, the column is present in the stored table
+    stored: Optional[bool]
+
+
+@dataclasses.dataclass
+class IndexMd:
+    """
+    Metadata needed to instantiate an EmbeddingIndex
+    """
+    id: int
+    name: str
+    indexed_col_id: int  # column being indexed
+    index_val_col_id: int  # column holding the values to be indexed
+    index_val_undo_col_id: int  # column holding index values for deleted rows
+    schema_version_add: int
+    schema_version_drop: Optional[int]
+    class_fqn: str
+    init_args: dict[str, Any]
 
 
 @dataclasses.dataclass
@@ -91,13 +118,13 @@ class ViewMd:
     base_versions: List[Tuple[str, Optional[int]]]
 
     # filter predicate applied to the base table; view-only
-    predicate: Optional[Dict[str, Any]]
+    predicate: Optional[dict[str, Any]]
 
     # ComponentIterator subclass; only for component views
     iterator_class_fqn: Optional[str]
 
     # args to pass to the iterator class constructor; only for component views
-    iterator_args: Optional[Dict[str, Any]]
+    iterator_args: Optional[dict[str, Any]]
 
 
 @dataclasses.dataclass
@@ -109,15 +136,15 @@ class TableMd:
     # each version has a corresponding schema version (current_version >= current_schema_version)
     current_schema_version: int
 
-    # used to assign Column.id
-    next_col_id: int
+    next_col_id: int  # used to assign Column.id
+    next_idx_id: int  # used to assign IndexMd.id
 
     # - used to assign the rowid column in the storage table
     # - every row is assigned a unique and immutable rowid on insertion
     next_row_id: int
 
-    column_history: Dict[int, ColumnHistory]  # col_id -> ColumnHistory
-
+    column_md: dict[int, ColumnMd]  # col_id -> ColumnMd
+    index_md: dict[int, IndexMd]  # index_id -> IndexMd
     view_md: Optional[ViewMd]
 
 
@@ -155,24 +182,20 @@ class TableVersion(Base):
 @dataclasses.dataclass
 class SchemaColumn:
     """
-    Records the logical (user-visible) schema of a table.
-    Contains the full set of columns for each new schema version: one record per (column x schema version).
+    Records the versioned metadata of a column.
     """
     pos: int
     name: str
-    col_type: dict
-    is_pk: bool
-    value_expr: Optional[dict]
-    stored: Optional[bool]
-    # if True, creates vector index for this column
-    is_indexed: bool
 
 
 @dataclasses.dataclass
 class TableSchemaVersionMd:
+    """
+    Records all versioned table metadata.
+    """
     schema_version: int
     preceding_schema_version: Optional[int]
-    columns: Dict[int, SchemaColumn]  # col_id -> SchemaColumn
+    columns: dict[int, SchemaColumn]  # col_id -> SchemaColumn
     num_retained_versions: int
     comment: str