pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/functions/date.py

@@ -0,0 +1,188 @@
+"""
+Pixeltable UDFs for `DateType`.
+
+Usage example:
+```python
+import pixeltable as pxt
+
+t = pxt.get_table(...)
+t.select(t.date_col.year, t.date_col.weekday()).collect()
+```
+"""
+
+from datetime import date, timedelta
+
+import sqlalchemy as sql
+
+import pixeltable as pxt
+from pixeltable.utils.code import local_public_names
+
+_SQL_ZERO = sql.literal(0)
+
+# NOT YET SUPPORTED date +/- integer
+# NOT YET SUPPORTED date1 - date2 -> integer
+# NOT YET SUPPORTED timestamp(date)
+# NOT YET SUPPORTED date(timestamp)
+
+
+@pxt.udf(is_property=True)
+def year(self: date) -> int:
+    """
+    Between 1 and 9999 inclusive.
+
+    (Between [`MINYEAR`](https://docs.python.org/3/library/datetime.html#datetime.MINYEAR) and
+    [`MAXYEAR`](https://docs.python.org/3/library/datetime.html#datetime.MAXYEAR) as defined by the Python `datetime`
+    library).
+
+    Equivalent to [`date.year`](https://docs.python.org/3/library/datetime.html#datetime.date.year).
+    """
+    return self.year
+
+
+@year.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('year', self)
+
+
+@pxt.udf(is_property=True)
+def month(self: date) -> int:
+    """
+    Between 1 and 12 inclusive.
+
+    Equivalent to [`date.month`](https://docs.python.org/3/library/datetime.html#datetime.date.month).
+    """
+    return self.month
+
+
+@month.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('month', self)
+
+
+@pxt.udf(is_property=True)
+def day(self: date) -> int:
+    """
+    Between 1 and the number of days in the given month of the given year.
+
+    Equivalent to [`date.day`](https://docs.python.org/3/library/datetime.html#datetime.date.day).
+    """
+    return self.day
+
+
+@day.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('day', self)
+
+
+@pxt.udf(is_method=True)
+def make_date(year: int, month: int, day: int) -> date:
+    """
+    Create a date.
+
+    Equivalent to [`datetime()`](https://docs.python.org/3/library/datetime.html#datetime.date).
+    """
+    return date(year, month, day)
+
+
+@make_date.to_sql
+def _(year: sql.ColumnElement, month: sql.ColumnElement, day: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.make_date(year.cast(sql.Integer), month.cast(sql.Integer), day.cast(sql.Integer))
+
+
+@pxt.udf(is_method=True)
+def weekday(self: date) -> int:
+    """
+    Between 0 (Monday) and 6 (Sunday) inclusive.
+
+    Equivalent to [`date.weekday()`](https://docs.python.org/3/library/datetime.html#datetime.date.weekday).
+    """
+    return self.weekday()
+
+
+@weekday.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('isodow', self) - 1
+
+
+@pxt.udf(is_method=True)
+def isoweekday(self: date) -> int:
+    """
+    Return the day of the week as an integer, where Monday is 1 and Sunday is 7.
+
+    Equivalent to [`date.isoweekday()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoweekday).
+    """
+    return self.isoweekday()
+
+
+@isoweekday.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('isodow', self)
+
+
+@pxt.udf(is_method=True)
+def isocalendar(self: date) -> dict:
+    """
+    Return a dictionary with three entries: `'year'`, `'week'`, and `'weekday'`.
+
+    Equivalent to
+    [`date.isocalendar()`](https://docs.python.org/3/library/datetime.html#datetime.date.isocalendar).
+    """
+    iso_year, iso_week, iso_weekday = self.isocalendar()
+    return {'year': iso_year, 'week': iso_week, 'weekday': iso_weekday}
+
+
+@pxt.udf(is_method=True)
+def isoformat(self: date, sep: str = 'T', timespec: str = 'auto') -> str:
+    """
+    Return a string representing the date and time in ISO 8601 format.
+
+    Equivalent to [`date.isoformat()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoformat).
+
+    Args:
+        sep: Separator between date and time.
+        timespec: The number of additional terms in the output. See the
+            [`date.isoformat()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoformat)
+            documentation for more details.
+    """
+    return self.isoformat()
+
+
+@pxt.udf(is_method=True)
+def toordinal(self: date) -> int:
+    """
+    Return the proleptic Gregorian ordinal of the date, where January 1 of year 1 has ordinal 1.
+
+    Equivalent to [`date.toordinal()`](https://docs.python.org/3/library/datetime.html#datetime.date.toordinal).
+    """
+    return self.toordinal()
+
+
+@pxt.udf(is_method=True)
+def strftime(self: date, format: str) -> str:
+    """
+    Return a string representing the date and time, controlled by an explicit format string.
+
+    Equivalent to [`date.strftime()`](https://docs.python.org/3/library/datetime.html#datetime.date.strftime).
+
+    Args:
+        format: The format string to control the output. For a complete list of formatting directives, see
+            [`strftime()` and `strptime()` Behavior](https://docs.python.org/3/library/datetime.html#strftime-and-strptime-behavior).
+    """
+    return self.strftime(format)
+
+
+@pxt.udf(is_method=True)
+def add_days(self: date, n: int) -> date:
+    """
+    Add `n` days to the date.
+
+    Equivalent to [`date + timedelta(days=n)`](https://docs.python.org/3/library/datetime.html#datetime.timedelta).
+    """
+    return self + timedelta(days=n)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/deepseek.py

@@ -0,0 +1,113 @@
+"""
+Pixeltable UDFs for Deepseek AI models.
+
+Provides integration with Deepseek's language models for chat completions
+and other AI capabilities.
+"""
+
+import json
+from typing import TYPE_CHECKING, Any
+
+import httpx
+
+import pixeltable as pxt
+from pixeltable import env
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import openai
+
+
+@env.register_client('deepseek')
+def _(api_key: str) -> 'openai.AsyncOpenAI':
+    import openai
+
+    return openai.AsyncOpenAI(
+        api_key=api_key,
+        base_url='https://api.deepseek.com',
+        http_client=httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=100, max_connections=500)),
+    )
+
+
+def _deepseek_client() -> 'openai.AsyncOpenAI':
+    return env.Env.get().get_client('deepseek')
+
+
+@pxt.udf(resource_pool='request-rate:deepseek')
+async def chat_completions(
+    messages: list,
+    *,
+    model: str,
+    model_kwargs: dict[str, Any] | None = None,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: dict[str, Any] | None = None,
+) -> dict:
+    """
+    Creates a model response for the given chat conversation.
+
+    Equivalent to the Deepseek `chat/completions` API endpoint.
+    For additional details, see: <https://api-docs.deepseek.com/api/create-chat-completion>
+
+    Deepseek uses the OpenAI SDK, so you will need to install the `openai` package to use this UDF.
+
+    Request throttling:
+    Applies the rate limit set in the config (section `deepseek`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install openai`
+
+    Args:
+        messages: A list of messages to use for chat completion, as described in the Deepseek API documentation.
+        model: The model to use for chat completion.
+        model_kwargs: Additional keyword args for the Deepseek `chat/completions` API.
+            For details on the available parameters, see: <https://api-docs.deepseek.com/api/create-chat-completion>
+        tools: An optional list of Pixeltable tools to use for the request.
+        tool_choice: An optional tool choice configuration.
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `deepseek-chat` to an existing Pixeltable column `tbl.prompt`
+        of the table `tbl`:
+
+        >>> messages = [
+        ...     {'role': 'system', 'content': 'You are a helpful assistant.'},
+        ...     {'role': 'user', 'content': tbl.prompt}
+        ... ]
+        >>> tbl.add_computed_column(response=chat_completions(messages, model='deepseek-chat'))
+    """
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    if tools is not None:
+        model_kwargs['tools'] = [{'type': 'function', 'function': tool} for tool in tools]
+
+    if tool_choice is not None:
+        if tool_choice['auto']:
+            model_kwargs['tool_choice'] = 'auto'
+        elif tool_choice['required']:
+            model_kwargs['tool_choice'] = 'required'
+        else:
+            assert tool_choice['tool'] is not None
+            model_kwargs['tool_choice'] = {'type': 'function', 'function': {'name': tool_choice['tool']}}
+
+    if tool_choice is not None and not tool_choice['parallel_tool_calls']:
+        if 'extra_body' not in model_kwargs:
+            model_kwargs['extra_body'] = {}
+        model_kwargs['extra_body']['parallel_tool_calls'] = False
+
+    result = await _deepseek_client().chat.completions.with_raw_response.create(
+        messages=messages, model=model, **model_kwargs
+    )
+
+    return json.loads(result.text)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/document.py

@@ -0,0 +1,81 @@
+"""
+Pixeltable UDFs for `DocumentType`.
+"""
+
+from typing import Any, Literal
+
+import pixeltable as pxt
+
+
+def document_splitter(
+    document: Any,
+    separators: str,
+    *,
+    elements: list[Literal['text', 'image']] | None = None,
+    limit: int | None = None,
+    overlap: int | None = None,
+    metadata: str = '',
+    skip_tags: list[str] | None = None,
+    tiktoken_encoding: str | None = 'cl100k_base',
+    tiktoken_target_model: str | None = None,
+    image_dpi: int = 300,
+    image_format: str = 'png',
+) -> tuple[type[pxt.iterators.ComponentIterator], dict[str, Any]]:
+    """Iterator over chunks of a document. The document is chunked according to the specified `separators`.
+
+    The iterator yields a `text` field containing the text of the chunk, and it may also
+    include additional metadata fields if specified in the `metadata` parameter, as explained below.
+
+    Chunked text will be cleaned with `ftfy.fix_text` to fix up common problems with unicode sequences.
+
+    Args:
+        separators: separators to use to chunk the document. Options are:
+             `'heading'`, `'paragraph'`, `'sentence'`, `'token_limit'`, `'char_limit'`, `'page'`.
+             This may be a comma-separated string, e.g., `'heading,token_limit'`.
+        elements: list of elements to extract from the document. Options are:
+            `'text'`, `'image'`. Defaults to `['text']` if not specified. The `'image'` element is only supported
+            for the `'page'` separator on PDF documents.
+        limit: the maximum number of tokens or characters in each chunk, if `'token_limit'`
+             or `'char_limit'` is specified.
+        metadata: additional metadata fields to include in the output. Options are:
+             `'title'`, `'heading'` (HTML and Markdown), `'sourceline'` (HTML), `'page'` (PDF), `'bounding_box'`
+             (PDF). The input may be a comma-separated string, e.g., `'title,heading,sourceline'`.
+        image_dpi: DPI to use when extracting images from PDFs. Defaults to 300.
+        image_format: format to use when extracting images from PDFs. Defaults to 'png'.
+
+    Examples:
+        All these examples assume an existing table `tbl` with a column `doc` of type `pxt.Document`.
+
+        Create a view that splits all documents into chunks of up to 300 tokens:
+
+        >>> pxt.create_view('chunks', tbl, iterator=document_splitter(tbl.doc, separators='token_limit', limit=300))
+
+        Create a view that splits all documents along sentence boundaries, including title and heading metadata:
+
+        >>> pxt.create_view(
+        ...     'sentence_chunks',
+        ...     tbl,
+        ...     iterator=document_splitter(tbl.doc, separators='sentence', metadata='title,heading')
+        ... )
+    """
+
+    kwargs: dict[str, Any] = {}
+    if elements is not None:
+        kwargs['elements'] = elements
+    if limit is not None:
+        kwargs['limit'] = limit
+    if overlap is not None:
+        kwargs['overlap'] = overlap
+    if metadata != '':
+        kwargs['metadata'] = metadata
+    if skip_tags is not None:
+        kwargs['skip_tags'] = skip_tags
+    if tiktoken_encoding != 'cl100k_base':
+        kwargs['tiktoken_encoding'] = tiktoken_encoding
+    if tiktoken_target_model is not None:
+        kwargs['tiktoken_target_model'] = tiktoken_target_model
+    if image_dpi != 300:
+        kwargs['image_dpi'] = image_dpi
+    if image_format != 'png':
+        kwargs['image_format'] = image_format
+    return pxt.iterators.document.DocumentSplitter._create(document=document, separators=separators, **kwargs)

pixeltable/functions/fal.py

@@ -0,0 +1,76 @@
+"""
+Pixeltable UDFs
+that wrap various endpoints from the fal.ai API. In order to use them, you must
+first `pip install fal-client` and configure your fal.ai credentials, as described in
+the [Working with fal.ai](https://docs.pixeltable.com/notebooks/integrations/working-with-fal) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import fal_client
+
+
+@register_client('fal')
+def _(api_key: str) -> 'fal_client.AsyncClient':
+    import fal_client
+
+    return fal_client.AsyncClient(key=api_key)
+
+
+def _fal_client() -> 'fal_client.AsyncClient':
+    return Env.get().get_client('fal')
+
+
+@pxt.udf(resource_pool='request-rate:fal')
+async def run(input: dict[str, Any], *, app: str) -> pxt.Json:
+    """
+    Run a model on fal.ai.
+
+    Uses fal's queue-based subscribe mechanism for reliable execution.
+    For additional details, see: <https://fal.ai/docs>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `fal`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install fal-client`
+
+    Args:
+        input: The input parameters for the model.
+        app: The name or ID of the fal.ai application to run (e.g., 'fal-ai/flux/schnell').
+
+    Returns:
+        The output of the model as a JSON object.
+
+    Examples:
+        Add a computed column that applies the model `fal-ai/flux/schnell`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> input = {'prompt': tbl.prompt}
+        ... tbl.add_computed_column(response=run(input, app='fal-ai/flux/schnell'))
+
+        Add a computed column that uses the model `fal-ai/fast-sdxl`
+        to generate images from an existing Pixeltable column `tbl.prompt`:
+
+        >>> input = {'prompt': tbl.prompt, 'image_size': 'square', 'num_inference_steps': 25}
+        ... tbl.add_computed_column(response=run(input, app='fal-ai/fast-sdxl'))
+        ... tbl.add_computed_column(image=tbl.response['images'][0]['url'].astype(pxt.Image))
+    """
+    Env.get().require_package('fal_client')
+    client = _fal_client()
+    result = await client.subscribe(app, arguments=input)
+    return result
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/fireworks.py

@@ -1,14 +1,15 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+Pixeltable UDFs
 that wrap various endpoints from the Fireworks AI API. In order to use them, you must
 first `pip install fireworks-ai` and configure your Fireworks AI credentials, as described in
-the [Working with Fireworks](https://pixeltable.readme.io/docs/working-with-fireworks) tutorial.
+the [Working with Fireworks](https://docs.pixeltable.com/notebooks/integrations/working-with-fireworks) tutorial.
 """
 
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import pixeltable as pxt
 from pixeltable import env
+from pixeltable.config import Config
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
@@ -26,21 +27,19 @@ def _fireworks_client() -> 'fireworks.client.Fireworks':
     return env.Env.get().get_client('fireworks')
 
 
-@pxt.udf
-def chat_completions(
-    messages: list[dict[str, str]],
-    *,
-    model: str,
-    max_tokens: Optional[int] = None,
-    top_k: Optional[int] = None,
-    top_p: Optional[float] = None,
-    temperature: Optional[float] = None,
+@pxt.udf(resource_pool='request-rate:fireworks')
+async def chat_completions(
+    messages: list[dict[str, str]], *, model: str, model_kwargs: dict[str, Any] | None = None
 ) -> dict:
     """
     Creates a model response for the given chat conversation.
 
     Equivalent to the Fireworks AI `chat/completions` API endpoint.
-    For additional details, see: [https://docs.fireworks.ai/api-reference/post-chatcompletions](https://docs.fireworks.ai/api-reference/post-chatcompletions)
+    For additional details, see: <https://docs.fireworks.ai/api-reference/post-chatcompletions>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `fireworks`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
 
     __Requirements:__
 
@@ -49,8 +48,8 @@ def chat_completions(
     Args:
         messages: A list of messages comprising the conversation so far.
         model: The name of the model to use.
-
-    For details on the other parameters, see: [https://docs.fireworks.ai/api-reference/post-chatcompletions](https://docs.fireworks.ai/api-reference/post-chatcompletions)
+        model_kwargs: Additional keyword args for the Fireworks `chat_completions` API. For details on the available
+            parameters, see: <https://docs.fireworks.ai/api-reference/post-chatcompletions>
 
     Returns:
         A dictionary containing the response and other metadata.
@@ -60,15 +59,68 @@ def chat_completions(
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
         >>> messages = [{'role': 'user', 'content': tbl.prompt}]
-        ... tbl['response'] = chat_completions(messages, model='accounts/fireworks/models/mixtral-8x22b-instruct')
+        ... tbl.add_computed_column(
+        ...     response=chat_completions(messages, model='accounts/fireworks/models/mixtral-8x22b-instruct')
+        ... )
     """
-    kwargs = {'max_tokens': max_tokens, 'top_k': top_k, 'top_p': top_p, 'temperature': temperature}
-    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).dict()
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    # for debugging purposes:
+    # res_sync = _fireworks_client().chat.completions.create(model=model, messages=messages, **kwargs_not_none)
+    # res_sync_dict = res_sync.dict()
+
+    if 'request_timeout' not in model_kwargs:
+        model_kwargs['request_timeout'] = Config.get().get_int_value('timeout', section='fireworks') or 600
+    # TODO: this timeout doesn't really work, I think it only applies to returning the stream, but not to the timing
+    # of the chunks; addressing this would require a timeout for the task running this udf
+    stream = _fireworks_client().chat.completions.acreate(model=model, messages=messages, **model_kwargs)
+    chunks = []
+    async for chunk in stream:
+        chunks.append(chunk)
+
+    res = {
+        'id': chunks[0].id,
+        'object': 'chat.completion',
+        'created': chunks[0].created,
+        'model': chunks[0].model,
+        'choices': [
+            {
+                'index': 0,
+                'message': {
+                    'role': None,
+                    'content': '',
+                    'tool_calls': None,
+                    'tool_call_id': None,
+                    'function': None,
+                    'name': None,
+                },
+                'finish_reason': None,
+                'logprobs': None,
+                'raw_output': None,
+            }
+        ],
+        'usage': {},
+    }
+    for chunk in chunks:
+        d = chunk.dict()
+        if 'usage' in d and d['usage'] is not None:
+            res['usage'] = d['usage']
+        if chunk.choices[0].finish_reason is not None:
+            res['choices'][0]['finish_reason'] = chunk.choices[0].finish_reason
+        if chunk.choices[0].delta.role is not None:
+            res['choices'][0]['message']['role'] = chunk.choices[0].delta.role
+        if chunk.choices[0].delta.content is not None:
+            res['choices'][0]['message']['content'] += chunk.choices[0].delta.content
+        if chunk.choices[0].delta.tool_calls is not None:
+            res['choices'][0]['message']['tool_calls'] = chunk.choices[0].delta.tool_calls
+        if chunk.choices[0].delta.function is not None:
+            res['choices'][0]['message']['function'] = chunk.choices[0].delta.function
+    return res
 
 
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__