chatlas 0.12.0__py3-none-any.whl → 0.13.1__py3-none-any.whl

chatlas/_provider_openai.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
 
 import base64
+import json
+import os
+import re
+import tempfile
+import warnings
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Literal, Optional, cast, overload
 
 import orjson
 from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
+from openai.types.batch import Batch
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
 from pydantic import BaseModel
 
 from ._chat import Chat
@@ -24,18 +31,20 @@ from ._content import (
 )
 from ._logging import log_model_default
 from ._merge import merge_dicts
-from ._provider import ModelInfo, Provider, StandardModelParamNames, StandardModelParams
+from ._provider import (
+    BatchStatus,
+    ModelInfo,
+    Provider,
+    StandardModelParamNames,
+    StandardModelParams,
+)
 from ._tokens import get_token_pricing, tokens_log
 from ._tools import Tool, basemodel_to_param_schema
 from ._turn import Turn, user_turn
 from ._utils import MISSING, MISSING_TYPE, is_testing, split_http_client_kwargs
 
 if TYPE_CHECKING:
-    from openai.types.chat import (
-        ChatCompletion,
-        ChatCompletionChunk,
-        ChatCompletionMessageParam,
-    )
+    from openai.types.chat import ChatCompletionMessageParam
     from openai.types.chat.chat_completion_assistant_message_param import (
         ContentArrayOfContentPart,
     )
@@ -45,10 +54,6 @@ if TYPE_CHECKING:
     from openai.types.chat_model import ChatModel
 
     from .types.openai import ChatAzureClientArgs, ChatClientArgs, SubmitInputArgs
-else:
-    ChatCompletion = object
-    ChatCompletionChunk = object
-
 
 # The dictionary form of ChatCompletion (TODO: stronger typing)?
 ChatCompletionDict = dict[str, Any]
@@ -171,6 +176,21 @@ def ChatOpenAI(
     )
 
 
+# Seems there is no native typing support for `files.content()` results
+# so mock them based on the docs here
+# https://platform.openai.com/docs/guides/batch#5-retrieve-the-results
+class BatchResult(BaseModel):
+    id: str
+    custom_id: str
+    response: BatchResultResponse
+
+
+class BatchResultResponse(BaseModel):
+    status_code: int
+    request_id: str
+    body: ChatCompletionDict
+
+
 class OpenAIProvider(
     Provider[ChatCompletion, ChatCompletionChunk, ChatCompletionDict, "SubmitInputArgs"]
 ):
@@ -353,8 +373,6 @@ class OpenAIProvider(
         return merge_dicts(completion, chunkd)
 
     def stream_turn(self, completion, has_data_model) -> Turn:
-        from openai.types.chat import ChatCompletion
-
         delta = completion["choices"][0].pop("delta")  # type: ignore
         completion["choices"][0]["message"] = delta  # type: ignore
         completion = ChatCompletion.construct(**completion)
@@ -662,6 +680,119 @@ class OpenAIProvider(
             "stop_sequences",
         }
 
+    def has_batch_support(self) -> bool:
+        return True
+
+    def batch_submit(
+        self,
+        conversations: list[list[Turn]],
+        data_model: Optional[type[BaseModel]] = None,
+    ):
+        # First put the requests in a file
+        # https://platform.openai.com/docs/api-reference/batch/request-input
+        # https://platform.openai.com/docs/api-reference/batch
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
+            temp_path = f.name
+
+            for i, turns in enumerate(conversations):
+                kwargs = self._chat_perform_args(
+                    stream=False,
+                    turns=turns,
+                    tools={},
+                    data_model=data_model,
+                )
+
+                body = {
+                    "messages": kwargs.get("messages", []),
+                    "model": self.model,
+                }
+
+                if "response_format" in kwargs:
+                    body["response_format"] = kwargs["response_format"]
+
+                request = {
+                    "custom_id": f"request-{i}",
+                    "method": "POST",
+                    "url": "/v1/chat/completions",
+                    "body": body,
+                }
+
+                f.write(orjson.dumps(request).decode() + "\n")
+
+        try:
+            with open(temp_path, "rb") as f:
+                file_response = self._client.files.create(file=f, purpose="batch")
+
+            batch = self._client.batches.create(
+                input_file_id=file_response.id,
+                endpoint="/v1/chat/completions",
+                completion_window="24h",
+            )
+
+            return batch.model_dump()
+        finally:
+            os.unlink(temp_path)
+
+    def batch_poll(self, batch):
+        batch = Batch.model_validate(batch)
+        b = self._client.batches.retrieve(batch.id)
+        return b.model_dump()
+
+    def batch_status(self, batch):
+        batch = Batch.model_validate(batch)
+        counts = batch.request_counts
+        total, completed, failed = 0, 0, 0
+        if counts is not None:
+            total = counts.total
+            completed = counts.completed
+            failed = counts.failed
+
+        return BatchStatus(
+            working=batch.status not in ["completed", "failed", "cancelled"],
+            n_processing=total - completed - failed,
+            n_succeeded=completed,
+            n_failed=failed,
+        )
+
+    def batch_retrieve(self, batch):
+        batch = Batch.model_validate(batch)
+        if batch.output_file_id is None:
+            raise ValueError("Batch has no output file")
+
+        # Download and parse JSONL results
+        response = self._client.files.content(batch.output_file_id)
+        results: list[dict[str, Any]] = []
+        for line in response.text.splitlines():
+            results.append(json.loads(line))
+
+        # Sort by custom_id to maintain order
+        def extract_id(x: str):
+            match = re.search(r"-(\d+)$", x)
+            return int(match.group(1)) if match else 0
+
+        results.sort(key=lambda x: int(extract_id(x.get("custom_id", ""))))
+
+        return results
+
+    def batch_result_turn(
+        self,
+        result,
+        has_data_model: bool = False,
+    ) -> Turn | None:
+        response = BatchResult.model_validate(result).response
+        if response.status_code != 200:
+            # TODO: offer advice on what to do?
+            warnings.warn(f"Batch request failed: {response.body}")
+            return None
+
+        completion = ChatCompletion.construct(**response.body)
+        return self._as_turn(completion, has_data_model)
+
+
+# -------------------------------------------------------------------------------------
+# Azure OpenAI Chat
+# -------------------------------------------------------------------------------------
+
 
 def ChatAzureOpenAI(
     *,

chatlas/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.12.0'
-__version_tuple__ = version_tuple = (0, 12, 0)
+__version__ = version = '0.13.1'
+__version_tuple__ = version_tuple = (0, 13, 1)
 
 __commit_id__ = commit_id = None

{chatlas-0.12.0.dist-info → chatlas-0.13.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chatlas
-Version: 0.12.0
+Version: 0.13.1
 Summary: A simple and consistent interface for chatting with LLMs
 Project-URL: Homepage, https://posit-dev.github.io/chatlas
 Project-URL: Documentation, https://posit-dev.github.io/chatlas
@@ -44,6 +44,7 @@ Requires-Dist: pillow; extra == 'dev'
 Requires-Dist: python-dotenv; extra == 'dev'
 Requires-Dist: ruff>=0.6.5; extra == 'dev'
 Requires-Dist: shiny; extra == 'dev'
+Requires-Dist: shinychat; extra == 'dev'
 Requires-Dist: snowflake-ml-python>=1.8.4; extra == 'dev'
 Requires-Dist: tenacity; extra == 'dev'
 Requires-Dist: tiktoken; extra == 'dev'

{chatlas-0.12.0.dist-info → chatlas-0.13.1.dist-info}/RECORD

@@ -1,8 +1,10 @@
-chatlas/__init__.py,sha256=CyViGMiz50clcVu3vpZgOq_qP4hmoYGOlcHKlRPcLJo,2416
+chatlas/__init__.py,sha256=M3zK10LguXW6bybDqatRV9SXd0y6axYu5QzENjbOVd0,2633
 chatlas/_auto.py,sha256=aeMN2_EM-xK-Yx5JaCuwYRZZ29eqn_0oM7QR5zayrec,8912
+chatlas/_batch_chat.py,sha256=1KkHENB-l7VmhCizhdvbJO5WQmRntQS6EvcSJ6VLgvM,5546
+chatlas/_batch_job.py,sha256=2__JIOo_JpcQyAAzO07r6eS4urpAxEc9m7_zsjFieQw,7359
 chatlas/_callbacks.py,sha256=3RpPaOQonTqScjXbaShgKJ1Rc-YxzWerxKRBjVssFnc,1838
-chatlas/_chat.py,sha256=pZOmlg0rqrzn7bi5QECAraG3ZA1MJPS0Ovk96wvgAGg,85712
-chatlas/_content.py,sha256=xi00PQgYi0hC2mTROOGVZiKfTkvlG9dud213iRQcwnY,22761
+chatlas/_chat.py,sha256=oESXNVzDCJ1DpV_fBRgwK6N_fs_EoJGrlez5dJjqx5c,90664
+chatlas/_content.py,sha256=BdJQ5G5onT9Cf1tNFeXsCWWTD2zSIjWz50FYIk6_DDI,22767
 chatlas/_content_image.py,sha256=EUK6wAint-JatLsiwvaPDu4D3W-NcIsDCkzABkXgfDg,8304
 chatlas/_content_pdf.py,sha256=cffeuJxzhUDukQ-Srkmpy62M8X12skYpU_FVq-Wvya4,2420
 chatlas/_display.py,sha256=wyQzSc6z1VqrJfkTLkw1wQcti9s1Pr4qT8UxFJESn4U,4664
@@ -11,18 +13,18 @@ chatlas/_live_render.py,sha256=UMZltE35LxziDKPMEeDwQ9meZ95SeqwhJi7j-y9pcro,4004
 chatlas/_logging.py,sha256=weKvXZDIZ88X7X61ruXM_S0AAhQ5mgiW9dR-km8x7Mg,3324
 chatlas/_mcp_manager.py,sha256=smMXeKZzP90MrlCdnTHMyo7AWHwl7J2jkU8dKSlnEsQ,10237
 chatlas/_merge.py,sha256=SGj_BetgA7gaOqSBKOhYmW3CYeQKTEehFrXvx3y4OYE,3924
-chatlas/_provider.py,sha256=k0rJ2uzGDacXVJZZVoLlySNCSFOjYbOC7k_VUT7j_Ms,6453
-chatlas/_provider_anthropic.py,sha256=cOBkAEj6gyl0NGdLk3QvI6pruZZ2fUsXJZAN2i3_j3k,27394
+chatlas/_provider.py,sha256=-5Oyq8tehHJtbBWQUyFUvdqTqZNUcOq2pO5qfAw5oQo,9057
+chatlas/_provider_anthropic.py,sha256=sPPEaDObGuY7JDqU533wlUDA-HaX3sumYWaD3kdG4nE,30964
 chatlas/_provider_cloudflare.py,sha256=vFbqgQPmosopJa9qsVxTkjPn4vYC_wOlgqa6_QmwTho,5227
 chatlas/_provider_databricks.py,sha256=JIOTm0HMe0qVAt8eS0WgGKugBwBdmL80JHLFH59ongU,4850
 chatlas/_provider_deepseek.py,sha256=6nPtPSo-Po6sD4i8PZJHuI5T2oATpLi5djXFGdlserk,4906
-chatlas/_provider_github.py,sha256=mhzoZGCb9d5DJvBll8_Ule6EljNJTNMGV2yS5KaZiY4,5622
+chatlas/_provider_github.py,sha256=5h7xwgUB0UdFpqr3UmUuFcC9hfn6kJvZDab0uAURIek,5642
 chatlas/_provider_google.py,sha256=0vAFSSk8SVBEtaDmy7Tl0XlHDqR4qwxWJR8HeSgh79E,21060
 chatlas/_provider_groq.py,sha256=XB2JDyuF95CcSbNkgk7JHcuy9KCW7hxTVaONDSjK8U8,3671
 chatlas/_provider_huggingface.py,sha256=feJ416X0UdtyoeHZbkgolFf62D7zxNwM7i_X3NYsQQw,4669
 chatlas/_provider_mistral.py,sha256=-p4rut0KCn-PrwnOlvr6lK8-K-OXvc5H9vTX-rCzUkk,5309
 chatlas/_provider_ollama.py,sha256=jFAu4v0NLUwdG_W_nKagBHOah0VKl_auTsgcYinP9rI,4119
-chatlas/_provider_openai.py,sha256=SwzEBwA491HOL6YvEI5soDQIVXnSLpwMmU-0DC8k7QA,26422
+chatlas/_provider_openai.py,sha256=f8ijdXMMHy17VcuA2ImMpaYvaiKjHzcnTdR8LH1xE40,30654
 chatlas/_provider_openrouter.py,sha256=9sCXvROVIiUdwfEbkVA-15_kc6ouFUP2uV2MmUe2rFk,4385
 chatlas/_provider_perplexity.py,sha256=5q_LsUCJQ5w-jRveLDMPvZTX-GU2TVURp65mUMyDh10,4248
 chatlas/_provider_portkey.py,sha256=6wKrLZmKVxOqyO6P3HBgWqPe7y1N8une_1wp0aJq7pU,4087
@@ -33,7 +35,7 @@ chatlas/_tools.py,sha256=8rhGOsEviBJXk5Qb-a1RRb_C-DE2T3DOeN6IhblkxqI,12408
 chatlas/_turn.py,sha256=yK7alUxeP8d2iBc7amyz20BtEqcpvX6BCwWZsnlQ5R4,4515
 chatlas/_typing_extensions.py,sha256=BXmbhywjm5ssmyVLGwyP_5TWZMAobzrrgZLYkB6_etE,1076
 chatlas/_utils.py,sha256=Kku2fa1mvTYCr5D28VxE6-fwfy2e2doCi-eKQkLEg4Y,4686
-chatlas/_version.py,sha256=VDAgmYWykomGcTucvPJWNS0ePk26QojGHvaE9chtgGc,706
+chatlas/_version.py,sha256=RBWKvLuPH5uJN_RJXTOaV5SvvBYz1PUZvU4D8m3gAvo,706
 chatlas/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chatlas/data/prices.json,sha256=X6qALp-dWc4nfus9lIqHoKzk3PZDPHTLoxxcN2m6fXc,62645
 chatlas/types/__init__.py,sha256=1n0xrJ7TRIKsZ2z06FLFgGqfKMFtXSIxxPvJ2j0hvPw,850
@@ -48,7 +50,7 @@ chatlas/types/openai/__init__.py,sha256=Q2RAr1bSH1nHsxICK05nAmKmxdhKmhbBkWD_XHiV
 chatlas/types/openai/_client.py,sha256=mAoQftcJIp0ssIhS8q3TIW9u6zTRNtYDmpZJO8L0mC0,849
 chatlas/types/openai/_client_azure.py,sha256=Tf_PFRl0QAj4Nk5CD0ZNIO-SRsT39bVkEJlUTry1fb8,960
 chatlas/types/openai/_submit.py,sha256=EDtIUFcNIJ5QAt0wVyBXvUshK8FA9e86wcZDQ_HUOYs,7829
-chatlas-0.12.0.dist-info/METADATA,sha256=Xd9L5THh_2C6ckT6Khvn2B6js3JhT8WyCjlAp1JzOOc,5594
-chatlas-0.12.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-chatlas-0.12.0.dist-info/licenses/LICENSE,sha256=zyuGzPOC7CcbOaBHsQ3UEyKYRO56KDUkor0OA4LqqDg,1081
-chatlas-0.12.0.dist-info/RECORD,,
+chatlas-0.13.1.dist-info/METADATA,sha256=gG-XEXaUc9LFlqoyqHa5nNzH_Dz-QHeL6HqxECWHAgM,5635
+chatlas-0.13.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+chatlas-0.13.1.dist-info/licenses/LICENSE,sha256=zyuGzPOC7CcbOaBHsQ3UEyKYRO56KDUkor0OA4LqqDg,1081
+chatlas-0.13.1.dist-info/RECORD,,