not-again-ai 0.10.3__tar.gz → 0.12.0__tar.gz

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: not-again-ai
-Version: 0.10.3
+Version: 0.12.0
 Summary: Designed to once and for all collect all the little things that come up over and over again in AI projects and put them in one place.
 Home-page: https://github.com/DaveCoDev/not-again-ai
 License: MIT
@@ -21,18 +21,21 @@ Provides-Extra: llm
 Provides-Extra: local-llm
 Provides-Extra: statistics
 Provides-Extra: viz
+Requires-Dist: azure-ai-inference (==1.0.0b3) ; extra == "llm"
+Requires-Dist: azure-identity (>=1.17,<2.0) ; extra == "llm"
 Requires-Dist: jinja2 (>=3.1,<4.0) ; extra == "local-llm"
 Requires-Dist: loguru (==0.7.2)
 Requires-Dist: numpy (>=1.26,<2.0) ; extra == "statistics" or extra == "viz"
-Requires-Dist: ollama (>=0.2,<0.3) ; extra == "local-llm"
-Requires-Dist: openai (>=1.36,<2.0) ; extra == "llm"
+Requires-Dist: ollama (>=0.3,<0.4) ; extra == "local-llm"
+Requires-Dist: openai (>=1.40,<2.0) ; extra == "llm"
 Requires-Dist: pandas (>=2.2,<3.0) ; extra == "viz"
+Requires-Dist: pydantic (>=2.8,<3.0) ; extra == "llm"
 Requires-Dist: python-liquid (>=1.12,<2.0) ; extra == "llm"
 Requires-Dist: scikit-learn (>=1.5,<2.0) ; extra == "statistics"
 Requires-Dist: scipy (>=1.14,<2.0) ; extra == "statistics"
 Requires-Dist: seaborn (>=0.13,<0.14) ; extra == "viz"
 Requires-Dist: tiktoken (>=0.7,<0.8) ; extra == "llm"
-Requires-Dist: transformers (>=4.42,<5.0) ; extra == "local-llm"
+Requires-Dist: transformers (>=4.43,<5.0) ; extra == "local-llm"
 Project-URL: Documentation, https://github.com/DaveCoDev/not-again-ai
 Project-URL: Repository, https://github.com/DaveCoDev/not-again-ai
 Description-Content-Type: text/markdown
@@ -71,23 +74,26 @@ Note that local LLM requires separate installations and will not work out of the
 The package is split into subpackages, so you can install only the parts you need.
 * **Base only**: `pip install not_again_ai`
 * **LLM**: `pip install not_again_ai[llm]`
-    1. If you wish to use OpenAI
+    1. OpenAI API
         1. Go to https://platform.openai.com/settings/profile?tab=api-keys to get your API key.
         1. (Optional) Set the `OPENAI_API_KEY` and the `OPENAI_ORG_ID` environment variables.
+    1. GitHub Models
+       1. Get a Personal Access Token from https://github.com/settings/tokens and set the `GITHUB_TOKEN` environment variable. The token does not need any permissions.
+       1. Check the [Github Marketplace](https://github.com/marketplace/models) to see which models are available.
 * **Local LLM**: `pip install not_again_ai[llm,local_llm]`
     1. Some HuggingFace transformers tokenizers are gated behind access requests. If you wish to use these, you will need to request access from HuggingFace on the model card.
        1. Then set the `HF_TOKEN` environment variable to your HuggingFace API token which can be found here: https://huggingface.co/settings/tokens
-    1. If you wish to use Ollama:
+    2. If you wish to use Ollama:
         1. Follow the instructions at https://github.com/ollama/ollama to install Ollama for your system. 
-        1. (Optional) [Add Ollama as a startup service (recommended)](https://github.com/ollama/ollama/blob/main/docs/linux.md#adding-ollama-as-a-startup-service-recommended)
-        1. (Optional) To make the Ollama service accessible on your local network from a Linux server, add the following to the `/etc/systemd/system/ollama.service` file which will make Ollama available at `http://<local_address>:11434`:
+        2. (Optional) [Add Ollama as a startup service (recommended)](https://github.com/ollama/ollama/blob/main/docs/linux.md#adding-ollama-as-a-startup-service-recommended)
+        3. (Optional) To make the Ollama service accessible on your local network from a Linux server, add the following to the `/etc/systemd/system/ollama.service` file which will make Ollama available at `http://<local_address>:11434`:
             ```bash
             [Service]
             ...
             Environment="OLLAMA_HOST=0.0.0.0"
             ```
-        1. It is recommended to always have the latest version of Ollama. To update Ollama check the [docs](https://github.com/ollama/ollama/blob/main/docs/). The command for Linux is: `curl -fsSL https://ollama.com/install.sh | sh`
-    2. HuggingFace transformers and other requirements are hardware dependent so for providers other than Ollama, this only installs some generic dependencies. Check the [notebooks](notebooks/local_llm/) for more details on what is available and how to install it.
+        4. It is recommended to always have the latest version of Ollama. To update Ollama check the [docs](https://github.com/ollama/ollama/blob/main/docs/). The command for Linux is: `curl -fsSL https://ollama.com/install.sh | sh`
+    3. HuggingFace transformers and other requirements are hardware dependent so for providers other than Ollama, this only installs some generic dependencies. Check the [notebooks](notebooks/local_llm/) for more details on what is available and how to install it.
 * **Statistics**: `pip install not_again_ai[statistics]`
 * **Visualization**: `pip install not_again_ai[viz]`
 
@@ -295,7 +301,7 @@ installed package to indicate that inline type annotations should be checked.
 
 ## Typos
 
-Checks for typos using [typos](https://github.com/crate-ci/typos)
+Check for typos using [typos](https://github.com/crate-ci/typos)
 
 ```bash
 (.venv) $ nox -s typos

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/README.md

@@ -32,23 +32,26 @@ Note that local LLM requires separate installations and will not work out of the
 The package is split into subpackages, so you can install only the parts you need.
 * **Base only**: `pip install not_again_ai`
 * **LLM**: `pip install not_again_ai[llm]`
-    1. If you wish to use OpenAI
+    1. OpenAI API
         1. Go to https://platform.openai.com/settings/profile?tab=api-keys to get your API key.
         1. (Optional) Set the `OPENAI_API_KEY` and the `OPENAI_ORG_ID` environment variables.
+    1. GitHub Models
+       1. Get a Personal Access Token from https://github.com/settings/tokens and set the `GITHUB_TOKEN` environment variable. The token does not need any permissions.
+       1. Check the [Github Marketplace](https://github.com/marketplace/models) to see which models are available.
 * **Local LLM**: `pip install not_again_ai[llm,local_llm]`
     1. Some HuggingFace transformers tokenizers are gated behind access requests. If you wish to use these, you will need to request access from HuggingFace on the model card.
        1. Then set the `HF_TOKEN` environment variable to your HuggingFace API token which can be found here: https://huggingface.co/settings/tokens
-    1. If you wish to use Ollama:
+    2. If you wish to use Ollama:
         1. Follow the instructions at https://github.com/ollama/ollama to install Ollama for your system. 
-        1. (Optional) [Add Ollama as a startup service (recommended)](https://github.com/ollama/ollama/blob/main/docs/linux.md#adding-ollama-as-a-startup-service-recommended)
-        1. (Optional) To make the Ollama service accessible on your local network from a Linux server, add the following to the `/etc/systemd/system/ollama.service` file which will make Ollama available at `http://<local_address>:11434`:
+        2. (Optional) [Add Ollama as a startup service (recommended)](https://github.com/ollama/ollama/blob/main/docs/linux.md#adding-ollama-as-a-startup-service-recommended)
+        3. (Optional) To make the Ollama service accessible on your local network from a Linux server, add the following to the `/etc/systemd/system/ollama.service` file which will make Ollama available at `http://<local_address>:11434`:
             ```bash
             [Service]
             ...
             Environment="OLLAMA_HOST=0.0.0.0"
             ```
-        1. It is recommended to always have the latest version of Ollama. To update Ollama check the [docs](https://github.com/ollama/ollama/blob/main/docs/). The command for Linux is: `curl -fsSL https://ollama.com/install.sh | sh`
-    2. HuggingFace transformers and other requirements are hardware dependent so for providers other than Ollama, this only installs some generic dependencies. Check the [notebooks](notebooks/local_llm/) for more details on what is available and how to install it.
+        4. It is recommended to always have the latest version of Ollama. To update Ollama check the [docs](https://github.com/ollama/ollama/blob/main/docs/). The command for Linux is: `curl -fsSL https://ollama.com/install.sh | sh`
+    3. HuggingFace transformers and other requirements are hardware dependent so for providers other than Ollama, this only installs some generic dependencies. Check the [notebooks](notebooks/local_llm/) for more details on what is available and how to install it.
 * **Statistics**: `pip install not_again_ai[statistics]`
 * **Visualization**: `pip install not_again_ai[viz]`
 
@@ -256,7 +259,7 @@ installed package to indicate that inline type annotations should be checked.
 
 ## Typos
 
-Checks for typos using [typos](https://github.com/crate-ci/typos)
+Check for typos using [typos](https://github.com/crate-ci/typos)
 
 ```bash
 (.venv) $ nox -s typos

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "not-again-ai"
-version = "0.10.3"
+version = "0.12.0"
 description = "Designed to once and for all collect all the little things that come up over and over again in AI projects and put them in one place."
 authors = ["DaveCoDev <dave.co.dev@gmail.com>"]
 license = "MIT"
@@ -29,20 +29,23 @@ python = "^3.11 || ^3.12"
 loguru = { version = "==0.7.2" }
 
 # Optional dependencies are defined here, and groupings are defined below.
+azure-ai-inference = { version = "==1.0.0b3", optional = true }
+azure-identity = { version = "^1.17", optional = true }
 jinja2 = { version = "^3.1", optional = true }
 numpy = { version = "^1.26", optional = true }
-ollama = { version = "^0.2", optional = true }
-openai = { version = "^1.36", optional = true }
+ollama = { version = "^0.3", optional = true }
+openai = { version = "^1.40", optional = true }
 pandas = { version = "^2.2", optional = true }
+pydantic = { version = "^2.8", optional = true }
 python-liquid = { version = "^1.12", optional = true }
 scipy = { version = "^1.14", optional = true }
 scikit-learn = { version = "^1.5", optional = true }
 seaborn = { version = "^0.13", optional = true }
 tiktoken = { version = "^0.7", optional = true }
-transformers = { version = "^4.42", optional = true }
+transformers = { version = "^4.43", optional = true }
 
 [tool.poetry.extras]
-llm = ["openai", "python-liquid", "tiktoken"]
+llm = ["azure-ai-inference", "azure-identity", "openai", "pydantic", "python-liquid", "tiktoken"]
 local_llm = ["jinja2", "ollama", "transformers"]
 statistics = ["numpy", "scikit-learn", "scipy"]
 viz = ["numpy", "pandas", "seaborn"]

not_again_ai-0.12.0/src/not_again_ai/llm/gh_models/azure_ai_client.py

@@ -0,0 +1,20 @@
+import os
+
+from azure.ai.inference import ChatCompletionsClient
+from azure.core.credentials import AzureKeyCredential
+
+
+def azure_ai_client(
+    token: str | None = None,
+    endpoint: str = "https://models.inference.ai.azure.com",
+) -> ChatCompletionsClient:
+    if not token:
+        token = os.getenv("GITHUB_TOKEN")
+        if not token:
+            raise ValueError("Token must be provided or GITHUB_TOKEN environment variable must be set")
+
+    client = ChatCompletionsClient(
+        endpoint=endpoint,
+        credential=AzureKeyCredential(token),
+    )
+    return client

not_again_ai-0.12.0/src/not_again_ai/llm/gh_models/chat_completion.py

@@ -0,0 +1,81 @@
+import contextlib
+import json
+import time
+from typing import Any
+
+from azure.ai.inference import ChatCompletionsClient
+from azure.ai.inference.models import ChatCompletionsToolDefinition, ChatRequestMessage
+
+
+def chat_completion(
+    messages: list[ChatRequestMessage],
+    model: str,
+    client: ChatCompletionsClient,
+    tools: list[ChatCompletionsToolDefinition] | None = None,
+    max_tokens: int | None = None,
+    temperature: float | None = None,
+    json_mode: bool = False,
+    seed: int | None = None,
+) -> dict[str, Any]:
+    """Gets a response from GitHub Models using the Azure AI Inference SDK.
+    See the available models at https://github.com/marketplace/models
+    Full documentation of the SDK is at: https://learn.microsoft.com/en-us/azure/ai-studio/reference/reference-model-inference-chat-completions
+    And samples at: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ai/azure-ai-inference/samples
+
+    Returns:
+        dict[str, Any]: A dictionary with the following keys
+            message (str | dict): The content of the generated assistant message.
+                If json_mode is True, this will be a dictionary.
+            tool_names (list[str], optional): The names of the tools called by the model.
+                If the model does not support tools, a ResponseError is raised.
+            tool_args_list (list[dict], optional): The arguments of the tools called by the model.
+            prompt_tokens (int): The number of tokens in the messages sent to the model.
+            completion_tokens (int): The number of tokens used by the model to generate the completion.
+            response_duration (float): The time, in seconds, taken to generate the response by using the model.
+            system_fingerprint (str, optional): If seed is set, a unique identifier for the model used to generate the response.
+    """
+    response_format = {"type": "json_object"} if json_mode else None
+    start_time = time.time()
+    response = client.complete(  # type: ignore
+        messages=messages,
+        model=model,
+        response_format=response_format,  # type: ignore
+        max_tokens=max_tokens,
+        temperature=temperature,
+        tools=tools,
+        seed=seed,
+    )
+    end_time = time.time()
+    response_duration = end_time - start_time
+
+    response_data = {}
+    finish_reason = response.choices[0].finish_reason
+    response_data["finish_reason"] = finish_reason.value  # type: ignore
+
+    message = response.choices[0].message.content
+    if message and json_mode:
+        with contextlib.suppress(json.JSONDecodeError):
+            message = json.loads(message)
+    response_data["message"] = message
+
+    # Check for tool calls because even if the finish_reason is stop, the model may have called a tool
+    tool_calls = response.choices[0].message.tool_calls
+    if tool_calls:
+        tool_names = []
+        tool_args_list = []
+        for tool_call in tool_calls:
+            tool_names.append(tool_call.function.name)  # type: ignore
+            tool_args_list.append(json.loads(tool_call.function.arguments))  # type: ignore
+        response_data["tool_names"] = tool_names
+        response_data["tool_args_list"] = tool_args_list
+
+    if seed is not None and hasattr(response, "system_fingerprint"):
+        response_data["system_fingerprint"] = response.system_fingerprint
+
+    usage = response.usage
+    if usage is not None:
+        response_data["completion_tokens"] = usage.completion_tokens
+        response_data["prompt_tokens"] = usage.prompt_tokens
+    response_data["response_duration"] = round(response_duration, 4)
+
+    return response_data

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/src/not_again_ai/llm/openai_api/chat_completion.py

@@ -4,6 +4,7 @@ import time
 from typing import Any
 
 from openai import OpenAI
+from pydantic import BaseModel
 
 
 def chat_completion(
@@ -15,6 +16,7 @@ def chat_completion(
     max_tokens: int | None = None,
     temperature: float = 0.7,
     json_mode: bool = False,
+    json_schema: dict[str, Any] | None = None,
     seed: int | None = None,
     logprobs: tuple[bool, int | None] | None = None,
     n: int = 1,
@@ -32,8 +34,8 @@ def chat_completion(
             https://platform.openai.com/docs/models/model-endpoint-compatibility
             for details on which models work with the Chat API.
         client (OpenAI): An instance of the OpenAI client.
-        tools (list[dict[str, Any]], optional): A list of tools the model may generate JSON inputs for.
-            Defaults to None.
+        tools (list[dict[str, Any]], optional):A list of tools the model may call.
+            Use this to provide a list of functions the model may generate JSON inputs for. Defaults to None.
         tool_choice (str, optional): The tool choice to use. Can be "auto", "required", "none", or a specific function name.
             Note the function name cannot be any of "auto", "required", or "none". Defaults to "auto".
         max_tokens (int, optional): The maximum number of tokens to generate in the chat completion.
@@ -44,6 +46,9 @@ def chat_completion(
         json_mode (bool, optional): When JSON mode is enabled, the model is constrained to only
             generate strings that parse into valid JSON object and will return a dictionary.
             See https://platform.openai.com/docs/guides/text-generation/json-mode
+        json_schema (dict, optional): Enables Structured Outputs which ensures the model will
+            always generate responses that adhere to your supplied JSON Schema.
+            See https://platform.openai.com/docs/guides/structured-outputs/structured-outputs
         seed (int, optional): If specified, OpenAI will make a best effort to sample deterministically,
             such that repeated requests with the same `seed` and parameters should return the same result.
             Determinism is not guaranteed, and you should refer to the `system_fingerprint` response
@@ -58,23 +63,35 @@ def chat_completion(
 
     Returns:
         dict[str, Any]: A dictionary with the following keys:
-            'finish_reason' (str): The reason the model stopped generating further tokens.
+            finish_reason (str): The reason the model stopped generating further tokens.
                 Can be 'stop', 'length', or 'tool_calls'.
-            'tool_names' (list[str], optional): The names of the tools called by the model.
-            'tool_args_list' (list[dict], optional): The arguments of the tools called by the model.
-            'message' (str | dict): The content of the generated assistant message.
+            tool_names (list[str], optional): The names of the tools called by the model.
+            tool_args_list (list[dict], optional): The arguments of the tools called by the model.
+            message (str | dict): The content of the generated assistant message.
                 If json_mode is True, this will be a dictionary.
-            'logprobs' (list[dict[str, Any] | list[dict[str, Any]]]): If logprobs[1] is between 1 and 5, each element in the list
+            logprobs (list[dict[str, Any] | list[dict[str, Any]]]): If logprobs[1] is between 1 and 5, each element in the list
                 will be a list of dictionaries containing the token, logprob, and bytes for the top `logprobs[1]` logprobs. Otherwise,
                 this will be a list of dictionaries containing the token, logprob, and bytes for each token in the message.
-            'choices' (list[dict], optional): A list of chat completion choices if n > 1 where each dict contains the above fields.
-            'completion_tokens' (int): The number of tokens used by the model to generate the completion.
+            choices (list[dict], optional): A list of chat completion choices if n > 1 where each dict contains the above fields.
+            completion_tokens (int): The number of tokens used by the model to generate the completion.
                 NOTE: If n > 1 this is the sum of all completions.
-            'prompt_tokens' (int): The number of tokens in the messages sent to the model.
-            'system_fingerprint' (str, optional): If seed is set, a unique identifier for the model used to generate the response.
-            'response_duration' (float): The time, in seconds, taken to generate the response from the API.
+            prompt_tokens (int): The number of tokens in the messages sent to the model.
+            system_fingerprint (str, optional): If seed is set, a unique identifier for the model used to generate the response.
+            response_duration (float): The time, in seconds, taken to generate the response from the API.
     """
-    response_format = {"type": "json_object"} if json_mode else None
+
+    if json_mode and json_schema is not None:
+        raise ValueError("json_schema and json_mode cannot be used together.")
+
+    if json_mode:
+        response_format: dict[str, Any] = {"type": "json_object"}
+    elif json_schema is not None:
+        if isinstance(json_schema, dict):
+            response_format = {"type": "json_schema", "json_schema": json_schema}
+        elif issubclass(json_schema, BaseModel):
+            response_format = json_schema
+    else:
+        response_format = {"type": "text"}
 
     kwargs.update(
         {
@@ -126,7 +143,7 @@ def chat_completion(
             response_data_curr["tool_args_list"] = tool_args_list
         elif finish_reason == "stop" or finish_reason == "length":
             message = response_choice.message.content
-            if json_mode:
+            if json_mode or json_schema is not None:
                 with contextlib.suppress(json.JSONDecodeError):
                     message = json.loads(message)
             response_data_curr["message"] = message

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/src/not_again_ai/llm/openai_api/openai_client.py

@@ -1,4 +1,7 @@
-from openai import OpenAI
+from typing import Literal
+
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from openai import AzureOpenAI, OpenAI
 
 
 class InvalidOAIAPITypeError(Exception):
@@ -8,13 +11,18 @@ class InvalidOAIAPITypeError(Exception):
 
 
 def openai_client(
-    api_type: str = "openai",
+    api_type: Literal["openai", "azure_openai"] = "openai",
     api_key: str | None = None,
     organization: str | None = None,
+    aoai_api_version: str = "2024-06-01",
+    azure_endpoint: str | None = None,
     timeout: float | None = None,
     max_retries: int | None = None,
-) -> OpenAI:
-    """Create an OpenAI client instance based on the specified API type and other provided parameters.
+) -> OpenAI | AzureOpenAI:
+    """Create an OpenAI or Azure OpenAI client instance based on the specified API type and other provided parameters.
+
+    Azure OpenAI requires RBAC authentication. You must be signed in with the Azure CLI and have correct role assigned.
+    See https://techcommunity.microsoft.com/t5/microsoft-developer-community/using-keyless-authentication-with-azure-openai/ba-p/4111521
 
     Args:
         api_type (str, optional): Type of the API to be used. Accepted values are 'openai' or 'azure_openai'.
@@ -23,6 +31,9 @@ def openai_client(
             OpenAI automatically uses `OPENAI_API_KEY` from the environment.
         organization (str, optional): The ID of the organization. If not provided,
             OpenAI automotically uses `OPENAI_ORG_ID` from the environment.
+        aoai_api_version (str, optional): Only applicable if using Azure OpenAI https://learn.microsoft.com/azure/ai-services/openai/reference#rest-api-versioning
+        azure_endpoint (str, optional): The endpoint to use for Azure OpenAI.
+            If not provided, will be read from the `AZURE_OPENAI_ENDPOINT` environment variable.
         timeout (float, optional): By default requests time out after 10 minutes.
         max_retries (int, optional): Certain errors are automatically retried 2 times by default,
             with a short exponential backoff. Connection errors (for example, due to a network connectivity problem),
@@ -52,6 +63,16 @@ def openai_client(
         filtered_args = {k: v for k, v in args.items() if v is not None}
         return OpenAI(**filtered_args)  # type: ignore
     elif api_type == "azure_openai":
-        raise NotImplementedError("AzureOpenAI is not yet supported by not-again-ai.")
+        azure_credential = DefaultAzureCredential()
+        ad_token_provider = get_bearer_token_provider(azure_credential, "https://cognitiveservices.azure.com/.default")
+        args = {
+            "api_version": aoai_api_version,
+            "azure_endpoint": azure_endpoint,
+            "azure_ad_token_provider": ad_token_provider,  # type: ignore
+            "timeout": timeout,
+            "max_retries": max_retries,
+        }
+        filtered_args = {k: v for k, v in args.items() if v is not None}
+        return AzureOpenAI(**filtered_args)  # type: ignore
     else:
-        raise NotImplementedError("This should never happen.")
+        raise NotImplementedError(f"API type '{api_type}' is invalid.")

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/src/not_again_ai/llm/openai_api/prompts.py

@@ -87,7 +87,7 @@ def chat_prompt(messages_unformatted: list[dict[str, Any]], variables: dict[str,
         A list which represents messages in the format that OpenAI expects for its chat completions API.
         See here for details: https://platform.openai.com/docs/api-reference/chat/create
 
-    Examples:
+    Example:
         >>> # Assume cat_image and dog_image are Path objects to image files
         >>> messages = [
         ...     {"role": "system", "content": "You are a helpful assistant."},

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/src/not_again_ai/llm/openai_api/tokens.py

@@ -86,6 +86,7 @@ def num_tokens_from_messages(
         "gpt-4-turbo-2024-04-09",
         "gpt-4o",
         "gpt-4o-2024-05-13",
+        "gpt-4o-2024-08-06",
         "gpt-4o-mini",
         "gpt-4o-mini-2024-07-18",
     }:

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/src/not_again_ai/local_llm/chat_completion.py

@@ -1,8 +1,10 @@
 from typing import Any
 
+from azure.ai.inference import ChatCompletionsClient
 from ollama import Client
 from openai import OpenAI
 
+from not_again_ai.llm.gh_models import chat_completion as chat_completion_gh_models
 from not_again_ai.llm.openai_api import chat_completion as chat_completion_openai
 from not_again_ai.local_llm.ollama import chat_completion as chat_completion_ollama
 
@@ -10,7 +12,8 @@ from not_again_ai.local_llm.ollama import chat_completion as chat_completion_oll
 def chat_completion(
     messages: list[dict[str, Any]],
     model: str,
-    client: OpenAI | Client,
+    client: OpenAI | Client | ChatCompletionsClient,
+    tools: list[dict[str, Any]] | None = None,
     max_tokens: int | None = None,
     temperature: float = 0.7,
     json_mode: bool = False,
@@ -25,6 +28,8 @@ def chat_completion(
         messages (list[dict[str, Any]]): A list of messages to send to the model.
         model (str): The model name to use.
         client (OpenAI | Client): The client object to use for chat completion.
+        tools (list[dict[str, Any]], optional):A list of tools the model may call.
+            Use this to provide a list of functions the model may generate JSON inputs for. Defaults to None.
         max_tokens (int, optional): The maximum number of tokens to generate.
         temperature (float, optional): The temperature of the model. Increasing the temperature will make the model answer more creatively.
         json_mode (bool, optional): This will structure the response as a valid JSON object.
@@ -34,6 +39,9 @@ def chat_completion(
         dict[str, Any]: A dictionary with the following keys
             message (str | dict): The content of the generated assistant message.
                 If json_mode is True, this will be a dictionary.
+            tool_names (list[str], optional): The names of the tools called by the model.
+                If the model does not support tools, a ResponseError is raised.
+            tool_args_list (list[dict], optional): The arguments of the tools called by the model.
             prompt_tokens (int): The number of tokens in the messages sent to the model.
             completion_tokens (int): The number of tokens used by the model to generate the completion.
             response_duration (float): The time, in seconds, taken to generate the response by using the model.
@@ -45,6 +53,7 @@ def chat_completion(
             messages=messages,
             model=model,
             client=client,
+            tools=tools,
             max_tokens=max_tokens,
             temperature=temperature,
             json_mode=json_mode,
@@ -56,6 +65,19 @@ def chat_completion(
             messages=messages,
             model=model,
             client=client,
+            tools=tools,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            json_mode=json_mode,
+            seed=seed,
+            **kwargs,
+        )
+    elif isinstance(client, ChatCompletionsClient):
+        response = chat_completion_gh_models.chat_completion(
+            messages=messages,  # type: ignore
+            model=model,
+            client=client,
+            tools=tools,  # type: ignore
             max_tokens=max_tokens,
             temperature=temperature,
             json_mode=json_mode,
@@ -68,6 +90,9 @@ def chat_completion(
     # Parse the responses to be consistent
     response_data = {}
     response_data["message"] = response.get("message")
+    if response.get("tool_names") and response.get("tool_args_list"):
+        response_data["tool_names"] = response.get("tool_names")
+        response_data["tool_args_list"] = response.get("tool_args_list")
     response_data["completion_tokens"] = response.get("completion_tokens")
     response_data["prompt_tokens"] = response.get("prompt_tokens")
     response_data["response_duration"] = response.get("response_duration")

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/src/not_again_ai/local_llm/ollama/chat_completion.py

@@ -13,6 +13,7 @@ def chat_completion(
     messages: list[dict[str, Any]],
     model: str,
     client: Client,
+    tools: list[dict[str, Any]] | None = None,
     max_tokens: int | None = None,
     context_window: int | None = None,
     temperature: float = 0.8,
@@ -27,6 +28,8 @@ def chat_completion(
         messages (list[dict[str, Any]]): A list of messages to send to the model.
         model (str): The model to use.
         client (Client): The Ollama client.
+        tools (list[dict[str, Any]], optional):A list of tools the model may call.
+            Use this to provide a list of functions the model may generate JSON inputs for. Defaults to None.
         max_tokens (int, optional): The maximum number of tokens to generate. Ollama calls this `num_predict`.
         context_window (int, optional): The number of tokens to consider as context. Ollama calls this `num_ctx`.
         temperature (float, optional): The temperature of the model. Increasing the temperature will make the model answer more creatively.
@@ -38,6 +41,9 @@ def chat_completion(
         dict[str, Any]: A dictionary with the following keys
             message (str | dict): The content of the generated assistant message.
                 If json_mode is True, this will be a dictionary.
+            tool_names (list[str], optional): The names of the tools called by the model.
+                If the model does not support tools, a ResponseError is raised.
+            tool_args_list (list[dict], optional): The arguments of the tools called by the model.
             prompt_tokens (int): The number of tokens in the messages sent to the model.
             completion_tokens (int): The number of tokens used by the model to generate the completion.
             response_duration (float): The time, in seconds, taken to generate the response by using the model.
@@ -59,6 +65,8 @@ def chat_completion(
     }
     if json_mode:
         all_args["format"] = "json"
+    if tools:
+        all_args["tools"] = tools
 
     try:
         start_time = time.time()
@@ -77,24 +85,27 @@ def chat_completion(
 
     response_data: dict[str, Any] = {}
 
-    # Handle getting the message returned by the model
-    message = response["message"].get("content", None)  # type: ignore
+    message = response["message"].get("content", "")
     if message and json_mode:
         with contextlib.suppress(json.JSONDecodeError):
             message = json.loads(message)
-    if message:
-        response_data["message"] = message
+    response_data["message"] = message
+
+    if response["message"].get("tool_calls"):
+        tool_calls = response["message"]["tool_calls"]
+        tool_names = [tool_call["function"]["name"] for tool_call in tool_calls]
+        tool_args_list = [tool_call["function"]["arguments"] for tool_call in tool_calls]
+        response_data["tool_names"] = tool_names
+        response_data["tool_args_list"] = tool_args_list
 
     tokenizer = load_tokenizer(model)
     prompt_tokens = num_tokens_from_messages(messages, tokenizer)
     response_data["prompt_tokens"] = prompt_tokens
 
-    # Get the number of tokens generated
-    response_data["completion_tokens"] = response.get("eval_count", None)  # type: ignore
+    response_data["completion_tokens"] = response.get("eval_count", None)
     if response_data["completion_tokens"] is None:
         response_data["completion_tokens"] = num_tokens_in_string(str(response_data["message"]), tokenizer)
 
-    # Get the latency of the response
     response_data["response_duration"] = round(response_duration, 4)
 
     return response_data

{not_again_ai-0.10.3 → not_again_ai-0.12.0}/src/not_again_ai/local_llm/ollama/model_mapping.py

@@ -7,6 +7,7 @@ i.e. all phi3 models will start with "phi3".
 OLLAMA_MODEL_MAPPING = {
     "phi3": "microsoft/Phi-3-mini-4k-instruct",
     "llama3:": "nvidia/Llama3-ChatQA-1.5-8B",  # Using this version to get around needed to accept an agreement to get access to the tokenizer
+    "llama3.1": "unsloth/Meta-Llama-3.1-8B-Instruct",
     "gemma": "google/gemma-1.1-7b-it",  # Requires HF_TOKEN set and accepting the agreement on the HF model page
     "qwen2": "Qwen/Qwen2-7B-Instruct",
     "granite-code": "ibm-granite/granite-34b-code-instruct",