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

chatlas/__init__.py

@@ -1,5 +1,11 @@
 from . import types
 from ._auto import ChatAuto
+from ._batch_chat import (
+    batch_chat,
+    batch_chat_completed,
+    batch_chat_structured,
+    batch_chat_text,
+)
 from ._chat import Chat
 from ._content import (
     ContentToolRequest,
@@ -36,6 +42,10 @@ except ImportError:  # pragma: no cover
     __version__ = "0.0.0"  # stub value for docs
 
 __all__ = (
+    "batch_chat",
+    "batch_chat_completed",
+    "batch_chat_structured",
+    "batch_chat_text",
     "ChatAnthropic",
     "ChatAuto",
     "ChatBedrockAnthropic",

chatlas/_batch_chat.py

@@ -0,0 +1,211 @@
+"""
+Batch chat processing for submitting multiple requests simultaneously.
+
+This module provides functionality for submitting multiple chat requests
+in batches to providers that support it (currently OpenAI and Anthropic).
+Batch processing can take up to 24 hours but offers significant cost savings
+(up to 50% less than regular requests).
+"""
+
+from __future__ import annotations
+
+import copy
+from pathlib import Path
+from typing import TypeVar, Union
+
+from pydantic import BaseModel
+
+from ._batch_job import BatchJob, ContentT
+from ._chat import Chat
+
+ChatT = TypeVar("ChatT", bound=Chat)
+BaseModelT = TypeVar("BaseModelT", bound=BaseModel)
+
+
+def batch_chat(
+    chat: ChatT,
+    prompts: list[ContentT] | list[list[ContentT]],
+    path: Union[str, Path],
+    wait: bool = True,
+) -> list[ChatT | None]:
+    """
+    Submit multiple chat requests in a batch.
+
+    This function allows you to submit multiple chat requests simultaneously
+    using provider batch APIs (currently OpenAI and Anthropic). Batch processing
+    can take up to 24 hours but offers significant cost savings.
+
+    Parameters
+    ----------
+    chat
+        Chat instance to use for the batch
+    prompts
+        List of prompts to process. Each can be a string or list of strings.
+    path
+        Path to file (with .json extension) to store batch state
+    wait
+        If True, wait for batch to complete. If False, return None if incomplete.
+
+    Returns
+    -------
+    List of Chat objects (one per prompt) if complete, None if wait=False and incomplete.
+    Individual Chat objects may be None if their request failed.
+
+    Example
+    -------
+
+    ```python
+    from chatlas import ChatOpenAI
+
+    chat = ChatOpenAI()
+    prompts = [
+        "What's the capital of France?",
+        "What's the capital of Germany?",
+        "What's the capital of Italy?",
+    ]
+
+    chats = batch_chat(chat, prompts, "capitals.json")
+    for i, result_chat in enumerate(chats):
+        if result_chat:
+            print(f"Prompt {i + 1}: {result_chat.get_last_turn().text}")
+    ```
+    """
+    job = BatchJob(chat, prompts, path, wait=wait)
+    job.step_until_done()
+
+    chats = []
+    assistant_turns = job.result_turns()
+    for user, assistant in zip(job.user_turns, assistant_turns):
+        if assistant is not None:
+            new_chat = copy.deepcopy(chat)
+            new_chat.add_turn(user)
+            new_chat.add_turn(assistant)
+            chats.append(new_chat)
+        else:
+            chats.append(None)
+
+    return chats
+
+
+def batch_chat_text(
+    chat: Chat,
+    prompts: list[ContentT] | list[list[ContentT]],
+    path: Union[str, Path],
+    wait: bool = True,
+) -> list[str | None]:
+    """
+    Submit multiple chat requests in a batch and return text responses.
+
+    This is a convenience function that returns just the text of the responses
+    rather than full Chat objects.
+
+    Parameters
+    ----------
+    chat
+        Chat instance to use for the batch
+    prompts
+        List of prompts to process
+    path
+        Path to file (with .json extension) to store batch state
+    wait
+        If True, wait for batch to complete
+
+    Return
+    ------
+    List of text responses (or None for failed requests)
+    """
+    chats = batch_chat(chat, prompts, path, wait=wait)
+
+    texts = []
+    for x in chats:
+        if x is None:
+            texts.append(None)
+            continue
+        last_turn = x.get_last_turn()
+        if last_turn is None:
+            texts.append(None)
+            continue
+        texts.append(last_turn.text)
+
+    return texts
+
+
+def batch_chat_structured(
+    chat: Chat,
+    prompts: list[ContentT] | list[list[ContentT]],
+    path: Union[str, Path],
+    data_model: type[BaseModelT],
+    wait: bool = True,
+) -> list[BaseModelT | None]:
+    """
+    Submit multiple structured data requests in a batch.
+
+    Parameters
+    ----------
+    chat
+        Chat instance to use for the batch
+    prompts
+        List of prompts to process
+    path
+        Path to file (with .json extension) to store batch state
+    data_model
+        Pydantic model class for structured responses
+    wait
+        If True, wait for batch to complete
+
+    Return
+    ------
+    List of structured data objects (or None for failed requests)
+    """
+    job = BatchJob(chat, prompts, path, data_model=data_model, wait=wait)
+    result = job.step_until_done()
+
+    if result is None:
+        return []
+
+    res: list[BaseModelT | None] = []
+    assistant_turns = job.result_turns()
+    for turn in assistant_turns:
+        if turn is None:
+            res.append(None)
+        else:
+            json = chat._extract_turn_json(turn)
+            model = data_model.model_validate(json)
+            res.append(model)
+
+    return res
+
+
+def batch_chat_completed(
+    chat: Chat,
+    prompts: list[ContentT] | list[list[ContentT]],
+    path: Union[str, Path],
+) -> bool:
+    """
+    Check if a batch job is completed without waiting.
+
+    Parameters
+    ----------
+    chat
+        Chat instance used for the batch
+    prompts
+        List of prompts used for the batch
+    path
+        Path to batch state file
+
+    Returns
+    -------
+    True if batch is complete, False otherwise
+    """
+    job = BatchJob(chat, prompts, path, wait=False)
+    stage = job.stage
+
+    if stage == "submitting":
+        return False
+    elif stage == "waiting":
+        status = job._poll()
+        return not status.working
+    elif stage == "retrieving" or stage == "done":
+        return True
+    else:
+        raise ValueError(f"Unknown batch stage: {stage}")

chatlas/_batch_job.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import time
+from datetime import timedelta
+from pathlib import Path
+from typing import Any, Literal, Optional, TypeVar, Union
+
+from pydantic import BaseModel
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
+from ._chat import Chat
+from ._content import Content
+from ._provider import BatchStatus
+from ._turn import Turn, user_turn
+from ._typing_extensions import TypedDict
+
+BatchStage = Literal["submitting", "waiting", "retrieving", "done"]
+
+
+class BatchStateHash(TypedDict):
+    provider: str
+    model: str
+    prompts: str
+    user_turns: str
+
+
+class BatchState(BaseModel):
+    version: int
+    stage: BatchStage
+    batch: dict[str, Any]
+    results: list[dict[str, Any]]
+    started_at: int
+    hash: BatchStateHash
+
+
+ContentT = TypeVar("ContentT", bound=Union[str, Content])
+
+
+class BatchJob:
+    """
+    Manages the lifecycle of a batch processing job.
+
+    A batch job goes through several stages:
+    1. "submitting" - Initial submission to the provider
+    2. "waiting" - Waiting for processing to complete
+    3. "retrieving" - Downloading results
+    4. "done" - Processing complete
+    """
+
+    def __init__(
+        self,
+        chat: Chat,
+        prompts: list[ContentT] | list[list[ContentT]],
+        path: Union[str, Path],
+        data_model: Optional[type[BaseModel]] = None,
+        wait: bool = True,
+    ):
+        if not chat.provider.has_batch_support():
+            raise ValueError("Batch requests are not supported by this provider")
+
+        self.chat = chat
+        self.prompts = prompts
+        self.path = Path(path)
+        self.data_model = data_model
+        self.should_wait = wait
+
+        # Convert prompts to user turns
+        self.user_turns: list[Turn] = []
+        for prompt in prompts:
+            if not isinstance(prompt, (str, Content)):
+                turn = user_turn(*prompt)
+            else:
+                turn = user_turn(prompt)
+            self.user_turns.append(turn)
+
+        # Job state management
+        self.provider = chat.provider
+        self.stage: BatchStage = "submitting"
+        self.batch: dict[str, Any] = {}
+        self.results: list[dict[str, Any]] = []
+
+        # Load existing state if file exists and is not empty
+        if self.path.exists() and self.path.stat().st_size > 0:
+            self._load_state()
+        else:
+            self.started_at = time.time()
+
+    def _load_state(self) -> None:
+        with open(self.path, "r") as f:
+            state = BatchState.model_validate_json(f.read())
+
+        self.stage = state.stage
+        self.batch = state.batch
+        self.results = state.results
+        self.started_at = state.started_at
+
+        # Verify hash to ensure consistency
+        stored_hash = state.hash
+        current_hash = self._compute_hash()
+
+        for key, value in current_hash.items():
+            if stored_hash.get(key) != value:
+                raise ValueError(
+                    f"Batch state mismatch: {key} doesn't match stored value. "
+                    f"Do you need to pick a different path?"
+                )
+
+    def _save_state(self) -> None:
+        state = BatchState(
+            version=1,
+            stage=self.stage,
+            batch=self.batch,
+            results=self.results,
+            started_at=int(self.started_at) if self.started_at else 0,
+            hash=self._compute_hash(),
+        )
+
+        with open(self.path, "w") as f:
+            f.write(state.model_dump_json(indent=2))
+
+    def _compute_hash(self) -> BatchStateHash:
+        turns = self.chat.get_turns(include_system_prompt=True)
+        return {
+            "provider": self.provider.name,
+            "model": self.provider.model,
+            "prompts": self._hash([str(p) for p in self.prompts]),
+            "user_turns": self._hash([str(turn) for turn in turns]),
+        }
+
+    @staticmethod
+    def _hash(x: Any) -> str:
+        return hashlib.md5(json.dumps(x, sort_keys=True).encode()).hexdigest()
+
+    def step(self) -> bool:
+        if self.stage == "submitting":
+            return self._submit()
+        elif self.stage == "waiting":
+            return self._wait()
+        elif self.stage == "retrieving":
+            return self._retrieve()
+        else:
+            raise ValueError(f"Unknown stage: {self.stage}")
+
+    def step_until_done(self) -> Optional["BatchJob"]:
+        while self.stage != "done":
+            if not self.step():
+                return None
+        return self
+
+    def _submit(self) -> bool:
+        existing_turns = self.chat.get_turns(include_system_prompt=True)
+
+        conversations = []
+        for turn in self.user_turns:
+            conversation = existing_turns + [turn]
+            conversations.append(conversation)
+
+        self.batch = self.provider.batch_submit(conversations, self.data_model)
+        self.stage = "waiting"
+        self._save_state()
+        return True
+
+    def _wait(self) -> bool:
+        # Always poll once, even when wait=False
+        status = self._poll()
+
+        if self.should_wait:
+            console = Console()
+
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("Processing..."),
+                TextColumn("[{task.fields[elapsed]}]"),
+                TextColumn("{task.fields[n_processing]} pending |"),
+                TextColumn("[green]{task.fields[n_succeeded]}[/green] done |"),
+                TextColumn("[red]{task.fields[n_failed]}[/red] failed"),
+                console=console,
+            ) as progress:
+                task = progress.add_task(
+                    "processing",
+                    elapsed=self._elapsed(),
+                    n_processing=status.n_processing,
+                    n_succeeded=status.n_succeeded,
+                    n_failed=status.n_failed,
+                )
+
+                while status.working:
+                    time.sleep(0.5)
+                    status = self._poll()
+                    progress.update(
+                        task,
+                        elapsed=self._elapsed(),
+                        n_processing=status.n_processing,
+                        n_succeeded=status.n_succeeded,
+                        n_failed=status.n_failed,
+                    )
+
+        if not status.working:
+            self.stage = "retrieving"
+            self._save_state()
+            return True
+        else:
+            return False
+
+    def _poll(self) -> "BatchStatus":
+        if not self.batch:
+            raise ValueError("No batch to poll")
+        self.batch = self.provider.batch_poll(self.batch)
+        self._save_state()
+        return self.provider.batch_status(self.batch)
+
+    def _elapsed(self) -> str:
+        return str(timedelta(seconds=int(time.time()) - int(self.started_at)))
+
+    def _retrieve(self) -> bool:
+        if not self.batch:
+            raise ValueError("No batch to retrieve")
+        self.results = self.provider.batch_retrieve(self.batch)
+        self.stage = "done"
+        self._save_state()
+        return True
+
+    def result_turns(self) -> list[Turn | None]:
+        turns = []
+        for result in self.results:
+            turn = self.provider.batch_result_turn(
+                result, has_data_model=self.data_model is not None
+            )
+            turns.append(turn)
+
+        return turns

chatlas/_chat.py

@@ -78,6 +78,7 @@ CompletionT = TypeVar("CompletionT")
 EchoOptions = Literal["output", "all", "none", "text"]
 
 T = TypeVar("T")
+BaseModelT = TypeVar("BaseModelT", bound=BaseModel)
 
 
 def is_present(value: T | None | MISSING_TYPE) -> TypeGuard[T]:
@@ -209,6 +210,7 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         self,
         *,
         include_system_prompt: bool = False,
+        tool_result_role: Literal["assistant", "user"] = "user",
     ) -> list[Turn[CompletionT]]:
         """
         Get all the turns (i.e., message contents) in the chat.
@@ -217,14 +219,50 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         ----------
         include_system_prompt
             Whether to include the system prompt in the turns.
+        tool_result_role
+            The role to assign to turns containing tool results. By default,
+            tool results are assigned a role of "user" since they represent
+            information provided to the assistant. If set to "assistant" tool
+            result content (plus the surrounding assistant turn contents) is
+            collected into a single assistant turn. This is convenient for
+            display purposes and more generally if you want the tool calling
+            loop to be contained in a single turn.
         """
 
         if not self._turns:
             return self._turns
 
         if not include_system_prompt and self._turns[0].role == "system":
-            return self._turns[1:]
-        return self._turns
+            turns = self._turns[1:]
+        else:
+            turns = self._turns
+
+        if tool_result_role == "user":
+            return turns
+
+        if tool_result_role != "assistant":
+            raise ValueError(
+                f"Expected `tool_result_role` to be one of 'user' or 'assistant', not '{tool_result_role}'"
+            )
+
+        # If a turn is purely a tool result, change its role
+        turns2 = copy.deepcopy(turns)
+        for turn in turns2:
+            if all(isinstance(c, ContentToolResult) for c in turn.contents):
+                turn.role = tool_result_role
+
+        # If two consecutive turns have the same role (i.e., assistant), collapse them into one
+        final_turns: list[Turn[CompletionT]] = []
+        for x in turns2:
+            if not final_turns:
+                final_turns.append(x)
+                continue
+            if x.role != final_turns[-1].role:
+                final_turns.append(x)
+            else:
+                final_turns[-1].contents.extend(x.contents)
+
+        return final_turns
 
     def get_last_turn(
         self,
@@ -531,7 +569,7 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         args
             The input to get a token count for.
         data_model
-            If the input is meant for data extraction (i.e., `.extract_data()`), then
+            If the input is meant for data extraction (i.e., `.chat_structured()`), then
             this should be the Pydantic model that describes the structure of the data to
             extract.
 
@@ -585,7 +623,7 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         args
             The input to get a token count for.
         data_model
-            If this input is meant for data extraction (i.e., `.extract_data_async()`),
+            If this input is meant for data extraction (i.e., `.chat_structured_async()`),
             then this should be the Pydantic model that describes the structure of the data
             to extract.
 
@@ -608,6 +646,7 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         port: int = 0,
         host: str = "127.0.0.1",
         launch_browser: bool = True,
+        bookmark_store: Literal["url", "server", "disable"] = "url",
         bg_thread: Optional[bool] = None,
         echo: Optional[EchoOptions] = None,
         content: Literal["text", "all"] = "all",
@@ -626,6 +665,12 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
             The host to run the app on (the default is "127.0.0.1").
         launch_browser
             Whether to launch a browser window.
+        bookmark_store
+            One of the following (default is "url"):
+              - `"url"`: Store bookmarks in the URL (default).
+              - `"server"`: Store bookmarks on the server (requires a server-side
+                storage backend).
+              - `"disable"`: Disable bookmarking.
         bg_thread
             Whether to run the app in a background thread. If `None`, the app will
             run in a background thread if the current environment is a notebook.
@@ -647,24 +692,37 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
             from shiny import App, run_app, ui
         except ImportError:
             raise ImportError(
-                "The `shiny` package is required for the `browser` method. "
+                "The `shiny` package is required for the `app()` method. "
                 "Install it with `pip install shiny`."
             )
 
-        app_ui = ui.page_fillable(
-            ui.chat_ui("chat"),
-            fillable_mobile=True,
-        )
+        try:
+            from shinychat import (
+                Chat,
+                chat_ui,
+                message_content,  # pyright: ignore[reportAttributeAccessIssue]
+            )
+        except ImportError:
+            raise ImportError(
+                "The `shinychat` package is required for the `app()` method. "
+                "Install it with `pip install shinychat`."
+            )
 
-        def server(input):  # noqa: A002
-            chat = ui.Chat(
-                "chat",
-                messages=[
-                    {"role": turn.role, "content": turn.text}
-                    for turn in self.get_turns()
-                ],
+        messages = [
+            message_content(x) for x in self.get_turns(tool_result_role="assistant")
+        ]
+
+        def app_ui(x):
+            return ui.page_fillable(
+                chat_ui("chat", messages=messages),
+                fillable_mobile=True,
             )
 
+        def server(input):  # noqa: A002
+            chat = Chat("chat")
+
+            chat.enable_bookmarking(self)
+
             @chat.on_user_submit
             async def _(user_input: str):
                 if stream:
@@ -688,7 +746,7 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
                         )
                     )
 
-        app = App(app_ui, server)
+        app = App(app_ui, server, bookmark_store=bookmark_store)
 
         def _run_app():
             run_app(app, launch_browser=launch_browser, port=port, host=host)
@@ -997,20 +1055,22 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
 
         return wrapper()
 
-    def extract_data(
+    def chat_structured(
         self,
         *args: Content | str,
-        data_model: type[BaseModel],
+        data_model: type[BaseModelT],
         echo: EchoOptions = "none",
         stream: bool = False,
-    ) -> dict[str, Any]:
+    ) -> BaseModelT:
         """
-        Extract structured data from the given input.
+        Extract structured data.
 
         Parameters
         ----------
         args
-            The input to extract data from.
+            The input to send to the chatbot. This is typically the text you
+            want to extract data from, but it can be omitted if the data is
+            obvious from the existing conversation.
         data_model
             A Pydantic model describing the structure of the data to extract.
         echo
@@ -1024,10 +1084,47 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
 
         Returns
         -------
-        dict[str, Any]
-            The extracted data.
+        BaseModelT
+            An instance of the provided `data_model` containing the extracted data.
         """
+        dat = self._submit_and_extract_data(
+            *args,
+            data_model=data_model,
+            echo=echo,
+            stream=stream,
+        )
+        return data_model.model_validate(dat)
 
+    def extract_data(
+        self,
+        *args: Content | str,
+        data_model: type[BaseModel],
+        echo: EchoOptions = "none",
+        stream: bool = False,
+    ) -> dict[str, Any]:
+        """
+        Deprecated: use `.chat_structured()` instead.
+        """
+        warnings.warn(
+            "The `extract_data()` method is deprecated and will be removed in a future release. "
+            "Use the `chat_structured()` method instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self._submit_and_extract_data(
+            *args,
+            data_model=data_model,
+            echo=echo,
+            stream=stream,
+        )
+
+    def _submit_and_extract_data(
+        self,
+        *args: Content | str,
+        data_model: type[BaseModel],
+        echo: EchoOptions = "none",
+        stream: bool = False,
+    ) -> dict[str, Any]:
         display = self._markdown_display(echo=echo)
 
         response = ChatResponse(
@@ -1046,33 +1143,24 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         turn = self.get_last_turn()
         assert turn is not None
 
-        res: list[ContentJson] = []
-        for x in turn.contents:
-            if isinstance(x, ContentJson):
-                res.append(x)
-
-        if len(res) != 1:
-            raise ValueError(
-                f"Data extraction failed: {len(res)} data results received."
-            )
-
-        json = res[0]
-        return json.value
+        return Chat._extract_turn_json(turn)
 
-    async def extract_data_async(
+    async def chat_structured_async(
         self,
         *args: Content | str,
-        data_model: type[BaseModel],
+        data_model: type[BaseModelT],
         echo: EchoOptions = "none",
         stream: bool = False,
-    ) -> dict[str, Any]:
+    ) -> BaseModelT:
         """
         Extract structured data from the given input asynchronously.
 
         Parameters
         ----------
         args
-            The input to extract data from.
+            The input to send to the chatbot. This is typically the text you
+            want to extract data from, but it can be omitted if the data is
+            obvious from the existing conversation.
         data_model
             A Pydantic model describing the structure of the data to extract.
         echo
@@ -1087,10 +1175,47 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
 
         Returns
         -------
-        dict[str, Any]
-            The extracted data.
+        BaseModelT
+            An instance of the provided `data_model` containing the extracted data.
         """
+        dat = await self._submit_and_extract_data_async(
+            *args,
+            data_model=data_model,
+            echo=echo,
+            stream=stream,
+        )
+        return data_model.model_validate(dat)
 
+    async def extract_data_async(
+        self,
+        *args: Content | str,
+        data_model: type[BaseModel],
+        echo: EchoOptions = "none",
+        stream: bool = False,
+    ) -> dict[str, Any]:
+        """
+        Deprecated: use `.chat_structured_async()` instead.
+        """
+        warnings.warn(
+            "The `extract_data_async()` method is deprecated and will be removed in a future release. "
+            "Use the `chat_structured_async()` method instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return await self._submit_and_extract_data_async(
+            *args,
+            data_model=data_model,
+            echo=echo,
+            stream=stream,
+        )
+
+    async def _submit_and_extract_data_async(
+        self,
+        *args: Content | str,
+        data_model: type[BaseModel],
+        echo: EchoOptions = "none",
+        stream: bool = False,
+    ) -> dict[str, Any]:
         display = self._markdown_display(echo=echo)
 
         response = ChatResponseAsync(
@@ -1109,6 +1234,10 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         turn = self.get_last_turn()
         assert turn is not None
 
+        return Chat._extract_turn_json(turn)
+
+    @staticmethod
+    def _extract_turn_json(turn: Turn) -> dict[str, Any]:
         res: list[ContentJson] = []
         for x in turn.contents:
             if isinstance(x, ContentJson):

chatlas/_content.py

@@ -603,7 +603,7 @@ class ContentJson(Content):
     JSON content
 
     This content type primarily exists to signal structured data extraction
-    (i.e., data extracted via [](`~chatlas.Chat`)'s `.extract_data()` method)
+    (i.e., data extracted via [](`~chatlas.Chat`)'s `.chat_structured()` method)
 
     Parameters
     ----------
@@ -630,7 +630,7 @@ class ContentPDF(Content):
     PDF content
 
     This content type primarily exists to signal PDF data extraction
-    (i.e., data extracted via [](`~chatlas.Chat`)'s `.extract_data()` method)
+    (i.e., data extracted via [](`~chatlas.Chat`)'s `.chat_structured()` method)
 
     Parameters
     ----------

chatlas/_provider.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from datetime import date
 from typing import (
+    Any,
     AsyncIterable,
     Generic,
     Iterable,
@@ -100,6 +101,16 @@ StandardModelParamNames = Literal[
 ]
 
 
+# Provider-agnostic batch status info
+class BatchStatus(BaseModel):
+    """Status information for a batch job."""
+
+    working: bool
+    n_processing: int
+    n_succeeded: int
+    n_failed: int
+
+
 class Provider(
     ABC,
     Generic[
@@ -261,3 +272,80 @@ class Provider(
 
     @abstractmethod
     def supported_model_params(self) -> set[StandardModelParamNames]: ...
+
+    def has_batch_support(self) -> bool:
+        """
+        Returns whether this provider supports batch processing.
+        Override this method to return True for providers that implement batch methods.
+        """
+        return False
+
+    def batch_submit(
+        self,
+        conversations: list[list[Turn]],
+        data_model: Optional[type[BaseModel]] = None,
+    ) -> dict[str, Any]:
+        """
+        Submit a batch of conversations for processing.
+
+        Args:
+            conversations: List of conversation histories (each is a list of Turns)
+            data_model: Optional structured data model for responses
+
+        Returns:
+            BatchInfo containing batch job information
+        """
+        raise NotImplementedError("This provider does not support batch processing")
+
+    def batch_poll(self, batch: dict[str, Any]) -> dict[str, Any]:
+        """
+        Poll the status of a submitted batch.
+
+        Args:
+            batch: Batch information returned from batch_submit
+
+        Returns:
+            Updated batch information
+        """
+        raise NotImplementedError("This provider does not support batch processing")
+
+    def batch_status(self, batch: dict[str, Any]) -> BatchStatus:
+        """
+        Get the status of a batch.
+
+        Args:
+            batch: Batch information
+
+        Returns:
+            BatchStatus with processing status information
+        """
+        raise NotImplementedError("This provider does not support batch processing")
+
+    def batch_retrieve(self, batch: dict[str, Any]) -> list[dict[str, Any]]:
+        """
+        Retrieve results from a completed batch.
+
+        Args:
+            batch: Batch information
+
+        Returns:
+            List of BatchResult objects, one for each request in the batch
+        """
+        raise NotImplementedError("This provider does not support batch processing")
+
+    def batch_result_turn(
+        self,
+        result: dict[str, Any],
+        has_data_model: bool = False,
+    ) -> Turn | None:
+        """
+        Convert a batch result to a Turn.
+
+        Args:
+            result: Individual BatchResult from batch_retrieve
+            has_data_model: Whether the request used a structured data model
+
+        Returns:
+            Turn object or None if the result was an error
+        """
+        raise NotImplementedError("This provider does not support batch processing")

chatlas/_provider_anthropic.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 import base64
+import re
 import warnings
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast, overload
 
 import orjson
+from openai.types.chat import ChatCompletionToolParam
 from pydantic import BaseModel
 
 from ._chat import Chat
@@ -21,7 +23,13 @@ from ._content import (
     ContentToolResultResource,
 )
 from ._logging import log_model_default
-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
@@ -38,11 +46,12 @@ if TYPE_CHECKING:
     )
     from anthropic.types.document_block_param import DocumentBlockParam
     from anthropic.types.image_block_param import ImageBlockParam
+    from anthropic.types.message_create_params import MessageCreateParamsNonStreaming
+    from anthropic.types.messages.batch_create_params import Request as BatchRequest
     from anthropic.types.model_param import ModelParam
     from anthropic.types.text_block_param import TextBlockParam
     from anthropic.types.tool_result_block_param import ToolResultBlockParam
     from anthropic.types.tool_use_block_param import ToolUseBlockParam
-    from openai.types.chat import ChatCompletionToolParam
 
     from .types.anthropic import ChatBedrockClientArgs, ChatClientArgs, SubmitInputArgs
 
@@ -631,6 +640,101 @@ class AnthropicProvider(
             completion=completion,
         )
 
+    def has_batch_support(self) -> bool:
+        return True
+
+    def batch_submit(
+        self,
+        conversations: list[list[Turn]],
+        data_model: Optional[type[BaseModel]] = None,
+    ):
+        from anthropic import NotGiven
+
+        requests: list["BatchRequest"] = []
+
+        for i, turns in enumerate(conversations):
+            kwargs = self._chat_perform_args(
+                stream=False,
+                turns=turns,
+                tools={},
+                data_model=data_model,
+            )
+
+            params: "MessageCreateParamsNonStreaming" = {
+                "messages": kwargs.get("messages", {}),
+                "model": self.model,
+                "max_tokens": kwargs.get("max_tokens", 4096),
+            }
+
+            # If data_model, tools/tool_choice should be present
+            tools = kwargs.get("tools")
+            tool_choice = kwargs.get("tool_choice")
+            if tools and not isinstance(tools, NotGiven):
+                params["tools"] = tools
+            if tool_choice and not isinstance(tool_choice, NotGiven):
+                params["tool_choice"] = tool_choice
+
+            requests.append({"custom_id": f"request-{i}", "params": params})
+
+        batch = self._client.messages.batches.create(requests=requests)
+        return batch.model_dump()
+
+    def batch_poll(self, batch):
+        from anthropic.types.messages import MessageBatch
+
+        batch = MessageBatch.model_validate(batch)
+        b = self._client.messages.batches.retrieve(batch.id)
+        return b.model_dump()
+
+    def batch_status(self, batch) -> "BatchStatus":
+        from anthropic.types.messages import MessageBatch
+
+        batch = MessageBatch.model_validate(batch)
+        status = batch.processing_status
+        counts = batch.request_counts
+
+        return BatchStatus(
+            working=status != "ended",
+            n_processing=counts.processing,
+            n_succeeded=counts.succeeded,
+            n_failed=counts.errored + counts.canceled + counts.expired,
+        )
+
+    # https://docs.anthropic.com/en/api/retrieving-message-batch-results
+    def batch_retrieve(self, batch):
+        from anthropic.types.messages import MessageBatch
+
+        batch = MessageBatch.model_validate(batch)
+        if batch.results_url is None:
+            raise ValueError("Batch has no results URL")
+
+        results: list[dict[str, Any]] = []
+        for res in self._client.messages.batches.results(batch.id):
+            results.append(res.model_dump())
+
+        # 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: extract_id(x.get("custom_id", "")))
+
+        return results
+
+    def batch_result_turn(self, result, has_data_model: bool = False) -> Turn | None:
+        from anthropic.types.messages.message_batch_individual_response import (
+            MessageBatchIndividualResponse,
+        )
+
+        result = MessageBatchIndividualResponse.model_validate(result)
+        if result.result.type != "succeeded":
+            # TODO: offer advice on what to do?
+            warnings.warn(f"Batch request didn't succeed: {result.result}")
+            return None
+
+        message = result.result.message
+        return self._as_turn(message, has_data_model)
+
 
 def ChatBedrockAnthropic(
     *,

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.0'
+__version_tuple__ = version_tuple = (0, 13, 0)
 
 __commit_id__ = commit_id = None

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chatlas
-Version: 0.12.0
+Version: 0.13.0
 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.0.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,8 +13,8 @@ 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
@@ -22,7 +24,7 @@ chatlas/_provider_groq.py,sha256=XB2JDyuF95CcSbNkgk7JHcuy9KCW7hxTVaONDSjK8U8,367
 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=2OwVfU6dYXS16K64p8xtyqfPM0bDW3r_7GeF3HzefhA,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.0.dist-info/METADATA,sha256=bpXFGxJZFx1hIkGuIWfC-MoHLhT3cDzsotmuhtFyktY,5635
+chatlas-0.13.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+chatlas-0.13.0.dist-info/licenses/LICENSE,sha256=zyuGzPOC7CcbOaBHsQ3UEyKYRO56KDUkor0OA4LqqDg,1081
+chatlas-0.13.0.dist-info/RECORD,,