chatlas 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

chatlas/_google.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import base64
 import json
-from typing import TYPE_CHECKING, Any, Literal, Optional, overload
+from typing import TYPE_CHECKING, Any, Literal, Optional, cast, overload
 
 from pydantic import BaseModel
 
@@ -16,20 +17,19 @@ from ._content import (
     ContentToolResult,
 )
 from ._logging import log_model_default
+from ._merge import merge_dicts
 from ._provider import Provider
-from ._tools import Tool, basemodel_to_param_schema
-from ._turn import Turn, normalize_turns
+from ._tokens import tokens_log
+from ._tools import Tool
+from ._turn import Turn, normalize_turns, user_turn
 
 if TYPE_CHECKING:
-    from google.generativeai.types.content_types import (
-        ContentDict,
-        FunctionDeclaration,
-        PartType,
-    )
-    from google.generativeai.types.generation_types import (
-        AsyncGenerateContentResponse,
+    from google.genai.types import Content as GoogleContent
+    from google.genai.types import (
         GenerateContentResponse,
-        GenerationConfig,
+        GenerateContentResponseDict,
+        Part,
+        PartDict,
     )
 
     from .types.google import ChatClientArgs, SubmitInputArgs
@@ -61,8 +61,8 @@ def ChatGoogle(
     ::: {.callout-note}
     ## Python requirements
 
-    `ChatGoogle` requires the `google-generativeai` package
-    (e.g., `pip install google-generativeai`).
+    `ChatGoogle` requires the `google-genai` package
+    (e.g., `pip install google-genai`).
     :::
 
     Examples
@@ -95,17 +95,13 @@ def ChatGoogle(
         The API key to use for authentication. You generally should not supply
         this directly, but instead set the `GOOGLE_API_KEY` environment variable.
     kwargs
-        Additional arguments to pass to the `genai.GenerativeModel` constructor.
+        Additional arguments to pass to the `genai.Client` constructor.
 
     Returns
     -------
     Chat
         A Chat object.
 
-    Limitations
-    -----------
-    `ChatGoogle` currently doesn't work with streaming tools.
-
     Note
     ----
     Pasting an API key into a chat constructor (e.g., `ChatGoogle(api_key="...")`)
@@ -144,63 +140,49 @@ def ChatGoogle(
     """
 
     if model is None:
-        model = log_model_default("gemini-1.5-flash")
-
-    turns = normalize_turns(
-        turns or [],
-        system_prompt=system_prompt,
-    )
+        model = log_model_default("gemini-2.0-flash")
 
     return Chat(
         provider=GoogleProvider(
-            turns=turns,
             model=model,
             api_key=api_key,
             kwargs=kwargs,
         ),
-        turns=turns,
+        turns=normalize_turns(
+            turns or [],
+            system_prompt=system_prompt,
+        ),
     )
 
 
-# The dictionary form of ChatCompletion (TODO: stronger typing)?
-GenerateContentDict = dict[str, Any]
-
-
 class GoogleProvider(
-    Provider[GenerateContentResponse, GenerateContentResponse, GenerateContentDict]
+    Provider[
+        GenerateContentResponse, GenerateContentResponse, "GenerateContentResponseDict"
+    ]
 ):
     def __init__(
         self,
         *,
-        turns: list[Turn],
         model: str,
         api_key: str | None,
         kwargs: Optional["ChatClientArgs"],
     ):
         try:
-            from google.generativeai import GenerativeModel
+            from google import genai
         except ImportError:
             raise ImportError(
-                f"The {self.__class__.__name__} class requires the `google-generativeai` package. "
-                "Install it with `pip install google-generativeai`."
+                f"The {self.__class__.__name__} class requires the `google-genai` package. "
+                "Install it with `pip install google-genai`."
             )
 
-        if api_key is not None:
-            import google.generativeai as genai
-
-            genai.configure(api_key=api_key)
-
-        system_prompt = None
-        if len(turns) > 0 and turns[0].role == "system":
-            system_prompt = turns[0].text
+        self._model = model
 
         kwargs_full: "ChatClientArgs" = {
-            "model_name": model,
-            "system_instruction": system_prompt,
+            "api_key": api_key,
             **(kwargs or {}),
         }
 
-        self._client = GenerativeModel(**kwargs_full)
+        self._client = genai.Client(**kwargs_full)
 
     @overload
     def chat_perform(
@@ -232,8 +214,11 @@ class GoogleProvider(
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ):
-        kwargs = self._chat_perform_args(stream, turns, tools, data_model, kwargs)
-        return self._client.generate_content(**kwargs)
+        kwargs = self._chat_perform_args(turns, tools, data_model, kwargs)
+        if stream:
+            return self._client.models.generate_content_stream(**kwargs)
+        else:
+            return self._client.models.generate_content(**kwargs)
 
     @overload
     async def chat_perform_async(
@@ -265,101 +250,160 @@ class GoogleProvider(
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ):
-        kwargs = self._chat_perform_args(stream, turns, tools, data_model, kwargs)
-        return await self._client.generate_content_async(**kwargs)
+        kwargs = self._chat_perform_args(turns, tools, data_model, kwargs)
+        if stream:
+            return await self._client.aio.models.generate_content_stream(**kwargs)
+        else:
+            return await self._client.aio.models.generate_content(**kwargs)
 
     def _chat_perform_args(
         self,
-        stream: bool,
         turns: list[Turn],
         tools: dict[str, Tool],
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ) -> "SubmitInputArgs":
+        from google.genai.types import FunctionDeclaration, GenerateContentConfig
+        from google.genai.types import Tool as GoogleTool
+
         kwargs_full: "SubmitInputArgs" = {
-            "contents": self._google_contents(turns),
-            "stream": stream,
-            "tools": self._gemini_tools(list(tools.values())) if tools else None,
+            "model": self._model,
+            "contents": cast("GoogleContent", self._google_contents(turns)),
             **(kwargs or {}),
         }
 
-        if data_model:
-            config = kwargs_full.get("generation_config", {})
-            params = basemodel_to_param_schema(data_model)
+        config = kwargs_full.get("config")
+        if config is None:
+            config = GenerateContentConfig()
+        if isinstance(config, dict):
+            config = GenerateContentConfig.model_construct(**config)
 
-            if "additionalProperties" in params:
-                del params["additionalProperties"]
+        if config.system_instruction is None:
+            if len(turns) > 0 and turns[0].role == "system":
+                config.system_instruction = turns[0].text
 
-            mime_type = "application/json"
-            if isinstance(config, dict):
-                config["response_schema"] = params
-                config["response_mime_type"] = mime_type
-            elif isinstance(config, GenerationConfig):
-                config.response_schema = params
-                config.response_mime_type = mime_type
+        if data_model:
+            config.response_schema = data_model
+            config.response_mime_type = "application/json"
+
+        if tools:
+            config.tools = [
+                GoogleTool(
+                    function_declarations=[
+                        FunctionDeclaration.from_callable(
+                            client=self._client, callable=tool.func
+                        )
+                        for tool in tools.values()
+                    ]
+                )
+            ]
 
-            kwargs_full["generation_config"] = config
+        kwargs_full["config"] = config
 
         return kwargs_full
 
     def stream_text(self, chunk) -> Optional[str]:
-        if chunk.parts:
+        try:
+            # Errors if there is no text (e.g., tool request)
             return chunk.text
-        return None
+        except Exception:
+            return None
 
     def stream_merge_chunks(self, completion, chunk):
-        # The .resolve() in .stream_turn() does the merging for us
-        return {}
-
-    def stream_turn(
-        self, completion, has_data_model, stream: GenerateContentResponse
-    ) -> Turn:
-        stream.resolve()
-        return self._as_turn(
-            stream,
-            has_data_model,
+        chunkd = chunk.model_dump()
+        if completion is None:
+            return cast("GenerateContentResponseDict", chunkd)
+        return cast(
+            "GenerateContentResponseDict",
+            merge_dicts(completion, chunkd),  # type: ignore
         )
 
-    async def stream_turn_async(
-        self, completion, has_data_model, stream: AsyncGenerateContentResponse
-    ) -> Turn:
-        await stream.resolve()
+    def stream_turn(self, completion, has_data_model) -> Turn:
         return self._as_turn(
-            stream,
+            completion,
             has_data_model,
         )
 
     def value_turn(self, completion, has_data_model) -> Turn:
+        completion = cast("GenerateContentResponseDict", completion.model_dump())
         return self._as_turn(completion, has_data_model)
 
-    def _google_contents(self, turns: list[Turn]) -> list["ContentDict"]:
-        contents: list["ContentDict"] = []
+    def token_count(
+        self,
+        *args: Content | str,
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+    ):
+        kwargs = self._token_count_args(
+            *args,
+            tools=tools,
+            data_model=data_model,
+        )
+
+        res = self._client.models.count_tokens(**kwargs)
+        return res.total_tokens or 0
+
+    async def token_count_async(
+        self,
+        *args: Content | str,
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+    ):
+        kwargs = self._token_count_args(
+            *args,
+            tools=tools,
+            data_model=data_model,
+        )
+
+        res = await self._client.aio.models.count_tokens(**kwargs)
+        return res.total_tokens or 0
+
+    def _token_count_args(
+        self,
+        *args: Content | str,
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+    ) -> dict[str, Any]:
+        turn = user_turn(*args)
+
+        kwargs = self._chat_perform_args(
+            turns=[turn],
+            tools=tools,
+            data_model=data_model,
+        )
+
+        args_to_keep = ["model", "contents", "tools"]
+
+        return {arg: kwargs[arg] for arg in args_to_keep if arg in kwargs}
+
+    def _google_contents(self, turns: list[Turn]) -> list["GoogleContent"]:
+        from google.genai.types import Content as GoogleContent
+
+        contents: list["GoogleContent"] = []
         for turn in turns:
             if turn.role == "system":
                 continue  # System messages are handled separately
             elif turn.role == "user":
                 parts = [self._as_part_type(c) for c in turn.contents]
-                contents.append({"role": turn.role, "parts": parts})
+                contents.append(GoogleContent(role=turn.role, parts=parts))
             elif turn.role == "assistant":
                 parts = [self._as_part_type(c) for c in turn.contents]
-                contents.append({"role": "model", "parts": parts})
+                contents.append(GoogleContent(role="model", parts=parts))
             else:
                 raise ValueError(f"Unknown role {turn.role}")
         return contents
 
-    def _as_part_type(self, content: Content) -> "PartType":
-        from google.generativeai.types.content_types import protos
+    def _as_part_type(self, content: Content) -> "Part":
+        from google.genai.types import FunctionCall, FunctionResponse, Part
 
         if isinstance(content, ContentText):
-            return protos.Part(text=content.text)
+            return Part.from_text(text=content.text)
         elif isinstance(content, ContentJson):
-            return protos.Part(text="<structured data/>")
-        elif isinstance(content, ContentImageInline):
-            return protos.Part(
-                inline_data={
-                    "mime_type": content.content_type,
-                    "data": content.data,
-                }
+            return Part.from_text(text="<structured data/>")
+        elif isinstance(content, ContentImageInline) and content.data:
+            return Part.from_bytes(
+                data=base64.b64decode(content.data),
+                mime_type=content.content_type,
             )
         elif isinstance(content, ContentImageRemote):
             raise NotImplementedError(
@@ -367,90 +411,197 @@ class GoogleProvider(
                 "Consider downloading the image and using content_image_file() instead."
             )
         elif isinstance(content, ContentToolRequest):
-            return protos.Part(
-                function_call={
-                    "name": content.id,
-                    "args": content.arguments,
-                }
+            return Part(
+                function_call=FunctionCall(
+                    id=content.id,
+                    name=content.name,
+                    # Goes in a dict, so should come out as a dict
+                    args=cast(dict[str, Any], content.arguments),
+                )
             )
         elif isinstance(content, ContentToolResult):
-            return protos.Part(
-                function_response={
-                    "name": content.id,
-                    "response": {"value": content.get_final_value()},
-                }
+            if content.error:
+                resp = {"error": content.error}
+            else:
+                resp = {"result": str(content.value)}
+            return Part(
+                # TODO: seems function response parts might need role='tool'???
+                # https://github.com/googleapis/python-genai/blame/c8cfef85c/README.md#L344
+                function_response=FunctionResponse(
+                    id=content.id,
+                    name=content.name,
+                    response=resp,
+                )
             )
         raise ValueError(f"Unknown content type: {type(content)}")
 
     def _as_turn(
         self,
-        message: "GenerateContentResponse | AsyncGenerateContentResponse",
+        message: "GenerateContentResponseDict",
         has_data_model: bool,
     ) -> Turn:
-        contents = []
-
-        msg = message.candidates[0].content
-
-        for part in msg.parts:
-            if part.text:
+        from google.genai.types import FinishReason
+
+        candidates = message.get("candidates")
+        if not candidates:
+            return Turn("assistant", "")
+
+        parts: list["PartDict"] = []
+        finish_reason = None
+        for candidate in candidates:
+            content = candidate.get("content")
+            if content:
+                parts.extend(content.get("parts") or {})
+            finish = candidate.get("finish_reason")
+            if finish:
+                finish_reason = finish
+
+        contents: list[Content] = []
+        for part in parts:
+            text = part.get("text")
+            if text:
                 if has_data_model:
-                    contents.append(ContentJson(json.loads(part.text)))
+                    contents.append(ContentJson(json.loads(text)))
                 else:
-                    contents.append(ContentText(part.text))
-            if part.function_call:
-                func = part.function_call
-                contents.append(
-                    ContentToolRequest(
-                        func.name,
-                        name=func.name,
-                        arguments=dict(func.args),
+                    contents.append(ContentText(text))
+            function_call = part.get("function_call")
+            if function_call:
+                # Seems name is required but id is optional?
+                name = function_call.get("name")
+                if name:
+                    contents.append(
+                        ContentToolRequest(
+                            id=function_call.get("id") or name,
+                            name=name,
+                            arguments=function_call.get("args"),
+                        )
                     )
-                )
-            if part.function_response:
-                func = part.function_response
-                contents.append(
-                    ContentToolResult(
-                        func.name,
-                        value=func.response,
+            function_response = part.get("function_response")
+            if function_response:
+                # Seems name is required but id is optional?
+                name = function_response.get("name")
+                if name:
+                    contents.append(
+                        ContentToolResult(
+                            id=function_response.get("id") or name,
+                            value=function_response.get("response"),
+                            name=name,
+                        )
                     )
-                )
 
-        usage = message.usage_metadata
-        tokens = (
-            usage.prompt_token_count,
-            usage.candidates_token_count,
-        )
+        usage = message.get("usage_metadata")
+        tokens = (0, 0)
+        if usage:
+            tokens = (
+                usage.get("prompt_token_count") or 0,
+                usage.get("candidates_token_count") or 0,
+            )
 
-        finish = message.candidates[0].finish_reason
+        tokens_log(self, tokens)
+
+        if isinstance(finish_reason, FinishReason):
+            finish_reason = finish_reason.name
 
         return Turn(
             "assistant",
             contents,
             tokens=tokens,
-            finish_reason=finish.name,
+            finish_reason=finish_reason,
             completion=message,
         )
 
-    def _gemini_tools(self, tools: list[Tool]) -> list["FunctionDeclaration"]:
-        from google.generativeai.types.content_types import FunctionDeclaration
-
-        res: list["FunctionDeclaration"] = []
-        for tool in tools:
-            fn = tool.schema["function"]
-            params = None
-            if "parameters" in fn and fn["parameters"]["properties"]:
-                params = {
-                    "type": "object",
-                    "properties": fn["parameters"]["properties"],
-                    "required": fn["parameters"]["required"],
-                }
-
-            res.append(
-                FunctionDeclaration(
-                    name=fn["name"],
-                    description=fn.get("description", ""),
-                    parameters=params,
-                )
-            )
 
-        return res
+def ChatVertex(
+    *,
+    model: Optional[str] = None,
+    project: Optional[str] = None,
+    location: Optional[str] = None,
+    api_key: Optional[str] = None,
+    system_prompt: Optional[str] = None,
+    turns: Optional[list[Turn]] = None,
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", GenerateContentResponse]:
+    """
+    Chat with a Google Vertex AI model.
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## Python requirements
+
+    `ChatGoogle` requires the `google-genai` package
+    (e.g., `pip install google-genai`).
+    :::
+
+    ::: {.callout-note}
+    ## Credentials
+
+    To use Google's models (i.e., Vertex AI), you'll need to sign up for an account
+    with [Vertex AI](https://cloud.google.com/vertex-ai), then specify the appropriate
+    model, project, and location.
+    :::
+
+    Parameters
+    ----------
+    model
+        The model to use for the chat. The default, None, will pick a reasonable
+        default, and warn you about it. We strongly recommend explicitly choosing
+        a model for all but the most casual use.
+    project
+        The Google Cloud project ID (e.g., "your-project-id"). If not provided, the
+        GOOGLE_CLOUD_PROJECT environment variable will be used.
+    location
+        The Google Cloud location (e.g., "us-central1"). If not provided, the
+        GOOGLE_CLOUD_LOCATION environment variable will be used.
+    system_prompt
+        A system prompt to set the behavior of the assistant.
+    turns
+        A list of turns to start the chat with (i.e., continuing a previous
+        conversation). If not provided, the conversation begins from scratch.
+        Do not provide non-`None` values for both `turns` and `system_prompt`.
+        Each message in the list should be a dictionary with at least `role`
+        (usually `system`, `user`, or `assistant`, but `tool` is also possible).
+        Normally there is also a `content` field, which is a string.
+
+    Returns
+    -------
+    Chat
+        A Chat object.
+
+    Examples
+    --------
+
+    ```python
+    import os
+    from chatlas import ChatVertex
+
+    chat = ChatVertex(
+        project="your-project-id",
+        location="us-central1",
+    )
+    chat.chat("What is the capital of France?")
+    ```
+    """
+
+    if kwargs is None:
+        kwargs = {}
+
+    kwargs["vertexai"] = True
+    kwargs["project"] = project
+    kwargs["location"] = location
+
+    if model is None:
+        model = log_model_default("gemini-2.0-flash")
+
+    return Chat(
+        provider=GoogleProvider(
+            model=model,
+            api_key=api_key,
+            kwargs=kwargs,
+        ),
+        turns=normalize_turns(
+            turns or [],
+            system_prompt=system_prompt,
+        ),
+    )

chatlas/_merge.py

@@ -1,5 +1,5 @@
 # Adapted from https://github.com/langchain-ai/langchain/blob/master/libs/core/langchain_core/utils/_merge.py
-# Also tweaked to more closely match https://github.com/hadley/elmer/blob/main/R/utils-merge.R
+# Also tweaked to more closely match https://github.com/hadley/ellmer/blob/main/R/utils-merge.R
 
 from __future__ import annotations
 

chatlas/_ollama.py

@@ -48,6 +48,13 @@ def ChatOllama(
     (e.g. `ollama pull llama3.2`).
     :::
 
+    ::: {.callout-note}
+    ## Python requirements
+
+    `ChatOllama` requires the `openai` package (e.g., `pip install openai`).
+    :::
+
+
     Examples
     --------
 
@@ -103,6 +110,7 @@ def ChatOllama(
 
     return ChatOpenAI(
         system_prompt=system_prompt,
+        api_key="ollama",  # ignored
         turns=turns,
         base_url=f"{base_url}/v1",
         model=model,