not-again-ai 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

not_again_ai/base/file_system.py

@@ -1,12 +1,34 @@
-import pathlib
+from pathlib import Path
 
 
-def create_file_dir(filepath: str) -> None:
+def create_file_dir(filepath: str | Path) -> None:
     """Creates the parent directories for the specified filepath.
     Does not throw any errors if the directories already exist.
 
     Args:
-        filepath (str): path to a file
+        filepath (str | Path): path to a file
     """
-    root_path = pathlib.Path(filepath).parent
+    root_path = Path(filepath).parent
     root_path.mkdir(parents=True, exist_ok=True)
+
+
+def readable_size(size: float) -> str:
+    """Convert a file size given in bytes to a human-readable format.
+
+    Args:
+        size (int): file size in bytes
+
+    Returns:
+        str: human-readable file size
+    """
+    # Define the suffixes for each size unit
+    suffixes = ["B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"]
+
+    # Start with bytes
+    count = 0
+    while size >= 1024 and count < len(suffixes) - 1:
+        count += 1
+        size /= 1024
+
+    # Format the size to two decimal places and append the appropriate suffix
+    return f"{size:.2f} {suffixes[count]}"

not_again_ai/llm/chat_completion.py

@@ -1,167 +1,76 @@
-import contextlib
-import json
 from typing import Any
 
+from ollama import Client
 from openai import OpenAI
 
+from not_again_ai.llm.ollama import chat_completion as chat_completion_ollama
+from not_again_ai.llm.openai_api import chat_completion as chat_completion_openai
+
 
 def chat_completion(
-    messages: list[dict[str, str]],
+    messages: list[dict[str, Any]],
     model: str,
-    client: OpenAI,
-    tools: list[dict[str, Any]] | None = None,
-    tool_choice: str = "auto",
+    client: OpenAI | Client,
     max_tokens: int | None = None,
     temperature: float = 0.7,
     json_mode: bool = False,
     seed: int | None = None,
-    logprobs: tuple[bool, int | None] | None = None,
-    n: int = 1,
     **kwargs: Any,
 ) -> dict[str, Any]:
-    """Get an OpenAI chat completion response: https://platform.openai.com/docs/api-reference/chat/create
-
-    NOTE: Depending on the model, certain parameters may not be supported,
-    particularly for older vision-enabled models like gpt-4-1106-vision-preview.
-    Be sure to check the documentation: https://platform.openai.com/docs/models/gpt-4-turbo-and-gpt-4
+    """Creates a common wrapper around chat completion models from different providers.
+    Currently supports the OpenAI API and Ollama local models.
+    All input parameters are supported by all providers in similar ways and the output is standardized.
 
     Args:
-        messages (list): A list of messages comprising the conversation so far.
-        model (str): ID of the model to use. See the model endpoint compatibility table:
-            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.
-        tool_choice (str, optional): The tool choice to use. Can be "auto", "none", or a specific function name.
-            Defaults to "auto".
-        max_tokens (int, optional): The maximum number of tokens to generate in the chat completion.
-            Defaults to None, which automatically limits to the model's maximum context length.
-        temperature (float, optional): What sampling temperature to use, between 0 and 2.
-            Higher values like 0.8 will make the output more random,
-            while lower values like 0.2 will make it more focused and deterministic. Defaults to 0.7.
-        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
-        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
-            parameter to monitor changes in the backend.
-        logprobs (tuple[bool, int], optional): Whether to return log probabilities of the output tokens or not.
-            If `logprobs[0]` is true, returns the log probabilities of each output token returned in the content of message.
-            `logprobs[1]` is an integer between 0 and 5 specifying the number of most likely tokens to return at each token position,
-            each with an associated log probability. `logprobs[0]` must be set to true if this parameter is used.
-        n (int, optional): How many chat completion choices to generate for each input message.
-            Defaults to 1.
-        **kwargs: Additional keyword arguments to pass to the OpenAI client 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.
+        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.
+        seed (int, optional): The seed to use for the model for reproducible outputs.
 
     Returns:
-        dict[str, Any]: A dictionary with the following keys:
-            '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.
+        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.
-            '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.
-                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.
+            completion_tokens (int): The number of tokens used by the model to generate the completion.
+            extras (dict): This will contain any additional fields returned by corresponding provider.
     """
-    response_format = {"type": "json_object"} if json_mode else None
-
-    kwargs.update(
-        {
-            "messages": messages,
-            "model": model,
-            "max_tokens": max_tokens,
-            "temperature": temperature,
-            "response_format": response_format,
-            "n": n,
-        }
-    )
-
-    if tools is not None:
-        kwargs["tools"] = tools
-        if tool_choice not in ["none", "auto"]:
-            kwargs["tool_choice"] = {"type": "function", "function": {"name": tool_choice}}
-        else:
-            kwargs["tool_choice"] = tool_choice
-
-    if seed is not None:
-        kwargs["seed"] = seed
-
-    if logprobs is not None:
-        kwargs["logprobs"] = logprobs[0]
-        if logprobs[0] and logprobs[1] is not None:
-            kwargs["top_logprobs"] = logprobs[1]
-
-    response = client.chat.completions.create(**kwargs)
-
-    response_data: dict[str, Any] = {"choices": []}
-    for response_choice in response.choices:
-        response_data_curr = {}
-        finish_reason = response_choice.finish_reason
-        response_data_curr["finish_reason"] = finish_reason
-
-        # We first check for tool calls because even if the finish_reason is stop, the model may have called a tool
-        tool_calls = response_choice.message.tool_calls
-        if tool_calls:
-            tool_names = []
-            tool_args_list = []
-            for tool_call in tool_calls:
-                tool_names.append(tool_call.function.name)
-                tool_args_list.append(json.loads(tool_call.function.arguments))
-            response_data_curr["message"] = response_choice.message.content
-            response_data_curr["tool_names"] = tool_names
-            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:
-                with contextlib.suppress(json.JSONDecodeError):
-                    message = json.loads(message)
-            response_data_curr["message"] = message
-
-        if response_choice.logprobs and response_choice.logprobs.content is not None:
-            logprobs_list: list[dict[str, Any] | list[dict[str, Any]]] = []
-            for logprob in response_choice.logprobs.content:
-                if logprob.top_logprobs:
-                    curr_logprob_infos = []
-                    for top_logprob in logprob.top_logprobs:
-                        curr_logprob_infos.append(
-                            {
-                                "token": top_logprob.token,
-                                "logprob": top_logprob.logprob,
-                                "bytes": top_logprob.bytes,
-                            }
-                        )
-                    logprobs_list.append(curr_logprob_infos)
-                else:
-                    logprobs_list.append(
-                        {
-                            "token": logprob.token,
-                            "logprob": logprob.logprob,
-                            "bytes": logprob.bytes,
-                        }
-                    )
-
-            response_data_curr["logprobs"] = logprobs_list
-        response_data["choices"].append(response_data_curr)
-
-    usage = response.usage
-    if usage is not None:
-        response_data["completion_tokens"] = usage.completion_tokens
-        response_data["prompt_tokens"] = usage.prompt_tokens
-
-    if seed is not None and response.system_fingerprint is not None:
-        response_data["system_fingerprint"] = response.system_fingerprint
-
-    if len(response_data["choices"]) == 1:
-        response_data.update(response_data["choices"][0])
-        del response_data["choices"]
+    # Determine which chat_completion function to call based on the client type
+    if isinstance(client, OpenAI):
+        response = chat_completion_openai.chat_completion(
+            messages=messages,
+            model=model,
+            client=client,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            json_mode=json_mode,
+            seed=seed,
+            **kwargs,
+        )
+    elif isinstance(client, Client):
+        response = chat_completion_ollama.chat_completion(
+            messages=messages,
+            model=model,
+            client=client,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            json_mode=json_mode,
+            seed=seed,
+            **kwargs,
+        )
+    else:
+        raise ValueError("Invalid client type")
+
+    # Parse the responses to be consistent
+    response_data = {}
+    response_data["message"] = response.get("message", None)
+    response_data["completion_tokens"] = response.get("completion_tokens", None)
+
+    # Return any additional fields from the response in an "extras" dictionary
+    extras = {k: v for k, v in response.items() if k not in response_data}
+    if extras:
+        response_data["extras"] = extras
 
     return response_data

not_again_ai/llm/ollama/chat_completion.py

@@ -0,0 +1,95 @@
+import contextlib
+import json
+import re
+from typing import Any
+
+from ollama import Client, ResponseError
+
+
+def _convert_duration(nanoseconds: int) -> float:
+    seconds = nanoseconds / 1_000_000_000
+    return round(seconds, 5)
+
+
+def chat_completion(
+    messages: list[dict[str, Any]],
+    model: str,
+    client: Client,
+    max_tokens: int | None = None,
+    context_window: int | None = None,
+    temperature: float = 0.8,
+    json_mode: bool = False,
+    seed: int | None = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Gets a Ollama chat completion response, see https://github.com/ollama/ollama/blob/main/docs/api.md#generate-a-chat-completion
+    For a full list of valid parameters: https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values
+
+    Args:
+        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.
+        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.
+        json_mode (bool, optional): This will structure the response as a valid JSON object.
+            It is important to instruct the model to use JSON in the prompt. Otherwise, the model may generate large amounts whitespace.
+        seed (int, optional): The seed to use for the model for reproducible outputs. Defaults to None.
+
+    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.
+            completion_tokens (int): The number of tokens used by the model to generate the completion.
+            response_duration (float): The time taken to generate the response in seconds.
+    """
+
+    options = {
+        "num_predict": max_tokens,
+        "num_ctx": context_window,
+        "temperature": temperature,
+    }
+    if seed is not None:
+        options["seed"] = seed
+    options.update(kwargs)
+
+    all_args = {
+        "model": model,
+        "messages": messages,
+        "options": options,
+    }
+    if json_mode:
+        all_args["format"] = "json"
+
+    try:
+        response = client.chat(**all_args)
+    except ResponseError as e:
+        # If the error says "model 'model' not found" use regex then raise a more specific error
+        expected_pattern = f"model '{model}' not found"
+        if re.search(expected_pattern, e.error):
+            raise ResponseError(
+                f"Model '{model}' not found. Please use not_again_ai.llm.ollama.service.pull() first."
+            ) from e
+        else:
+            raise ResponseError(e.message) from e
+
+    response_data: dict[str, Any] = {}
+
+    # Handle getting the message returned by the model
+    message = response["message"].get("content", None)
+    if message and json_mode:
+        with contextlib.suppress(json.JSONDecodeError):
+            message = json.loads(message)
+    if message:
+        response_data["message"] = message
+
+    # Get the number of tokens generated
+    response_data["completion_tokens"] = response.get("eval_count", None)
+
+    # Get the latency of the response
+    if response.get("total_duration", None):
+        response_data["response_duration"] = _convert_duration(response["total_duration"])
+    else:
+        response_data["response_duration"] = None
+
+    return response_data

not_again_ai/llm/ollama/ollama_client.py

@@ -0,0 +1,24 @@
+import os
+
+from ollama import Client
+
+
+def ollama_client(host: str | None = None, timeout: float | None = None) -> Client:
+    """Create an Ollama client instance based on the specified host or will read from the OLLAMA_HOST environment variable.
+
+    Args:
+        host (str, optional): The host URL of the Ollama server.
+        timeout (float, optional): The timeout for requests
+
+    Returns:
+        Client: An instance of the Ollama client.
+
+    Examples:
+        >>> client = client(host="http://localhost:11434")
+    """
+    if host is None:
+        host = os.getenv("OLLAMA_HOST")
+        if host is None:
+            raise ValueError("Host must be provided or OLLAMA_HOST environment variable must be set.")
+
+    return Client(host=host, timeout=timeout)

not_again_ai/llm/ollama/service.py

@@ -0,0 +1,81 @@
+from typing import Any
+
+from ollama import Client
+
+from not_again_ai.base.file_system import readable_size
+
+
+def list_models(client: Client) -> list[dict[str, Any]]:
+    """List models that are available locally.
+
+    Args:
+        client (Client): The Ollama client.
+
+    Returns:
+        list[dict[str, Any]]: A list of dictionaries (each corresponding to an available model) with the following keys:
+            name (str): Name of the model
+            model (str): Name of the model. This should be the same as the name.
+            modified_at (str): The date and time the model was last modified.
+            size (int): The size of the model in bytes.
+            size_readable (str): The size of the model in a human-readable format.
+            details (dict[str, Any]): Additional details about the model.
+    """
+    response = client.list().get("models", [])
+
+    response_data = []
+    for model_data in response:
+        curr_model_data = {}
+        curr_model_data["name"] = model_data["name"]
+        curr_model_data["model"] = model_data["model"]
+        curr_model_data["modified_at"] = model_data["modified_at"]
+        curr_model_data["size"] = model_data["size"]
+        curr_model_data["size_readable"] = readable_size(model_data["size"])
+        curr_model_data["details"] = model_data["details"]
+
+        response_data.append(curr_model_data)
+
+    return response_data
+
+
+def is_model_available(model_name: str, client: Client) -> bool:
+    """Check if a model is available locally.
+
+    Args:
+        model_name (str): The name of the model.
+        client (Client): The Ollama client.
+
+    Returns:
+        bool: True if the model is available locally, False otherwise.
+    """
+    # If model_name does not have a ":", append ":latest"
+    if ":" not in model_name:
+        model_name = f"{model_name}:latest"
+    models = list_models(client)
+    return any(model["name"] == model_name for model in models)
+
+
+def show(model_name: str, client: Client) -> dict[str, Any]:
+    """Show information about a model including the modelfile, available parameters, template, and additional details.
+
+    Args:
+        model_name (str): The name of the model.
+        client (Client): The Ollama client.
+    """
+    response = client.show(model_name)
+
+    response_data = {}
+    response_data["modelfile"] = response["modelfile"]
+    response_data["parameters"] = response["parameters"]
+    response_data["template"] = response["template"]
+    response_data["details"] = response["details"]
+    return response_data
+
+
+def pull(model_name: str, client: Client) -> Any:
+    """Pull a model from the Ollama server and returns the status of the pull operation."""
+    return client.pull(model_name)
+
+
+def delete(model_name: str, client: Client) -> Any:
+    """Delete a model from the local filesystem and returns the status of the delete operation."""
+    return client.delete(model_name)

not_again_ai/llm/openai_api/chat_completion.py

@@ -0,0 +1,167 @@
+import contextlib
+import json
+from typing import Any
+
+from openai import OpenAI
+
+
+def chat_completion(
+    messages: list[dict[str, Any]],
+    model: str,
+    client: OpenAI,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: str = "auto",
+    max_tokens: int | None = None,
+    temperature: float = 0.7,
+    json_mode: bool = False,
+    seed: int | None = None,
+    logprobs: tuple[bool, int | None] | None = None,
+    n: int = 1,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Get an OpenAI chat completion response: https://platform.openai.com/docs/api-reference/chat/create
+
+    NOTE: Depending on the model, certain parameters may not be supported,
+    particularly for older vision-enabled models like gpt-4-1106-vision-preview.
+    Be sure to check the documentation: https://platform.openai.com/docs/models/gpt-4-turbo-and-gpt-4
+
+    Args:
+        messages (list): A list of messages comprising the conversation so far.
+        model (str): ID of the model to use. See the model endpoint compatibility table:
+            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.
+        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.
+            Defaults to None, which automatically limits to the model's maximum context length.
+        temperature (float, optional): What sampling temperature to use, between 0 and 2.
+            Higher values like 0.8 will make the output more random,
+            while lower values like 0.2 will make it more focused and deterministic. Defaults to 0.7.
+        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
+        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
+            parameter to monitor changes in the backend.
+        logprobs (tuple[bool, int], optional): Whether to return log probabilities of the output tokens or not.
+            If `logprobs[0]` is true, returns the log probabilities of each output token returned in the content of message.
+            `logprobs[1]` is an integer between 0 and 5 specifying the number of most likely tokens to return at each token position,
+            each with an associated log probability. `logprobs[0]` must be set to true if this parameter is used.
+        n (int, optional): How many chat completion choices to generate for each input message.
+            Defaults to 1.
+        **kwargs: Additional keyword arguments to pass to the OpenAI client chat completion.
+
+    Returns:
+        dict[str, Any]: A dictionary with the following keys:
+            '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.
+                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
+                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.
+                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_format = {"type": "json_object"} if json_mode else None
+
+    kwargs.update(
+        {
+            "messages": messages,
+            "model": model,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "response_format": response_format,
+            "n": n,
+        }
+    )
+
+    if tools is not None:
+        kwargs["tools"] = tools
+        if tool_choice not in ["none", "auto", "required"]:
+            kwargs["tool_choice"] = {"type": "function", "function": {"name": tool_choice}}
+        else:
+            kwargs["tool_choice"] = tool_choice
+
+    if seed is not None:
+        kwargs["seed"] = seed
+
+    if logprobs is not None:
+        kwargs["logprobs"] = logprobs[0]
+        if logprobs[0] and logprobs[1] is not None:
+            kwargs["top_logprobs"] = logprobs[1]
+
+    response = client.chat.completions.create(**kwargs)
+
+    response_data: dict[str, Any] = {"choices": []}
+    for response_choice in response.choices:
+        response_data_curr = {}
+        finish_reason = response_choice.finish_reason
+        response_data_curr["finish_reason"] = finish_reason
+
+        # We first check for tool calls because even if the finish_reason is stop, the model may have called a tool
+        tool_calls = response_choice.message.tool_calls
+        if tool_calls:
+            tool_names = []
+            tool_args_list = []
+            for tool_call in tool_calls:
+                tool_names.append(tool_call.function.name)
+                tool_args_list.append(json.loads(tool_call.function.arguments))
+            response_data_curr["message"] = response_choice.message.content
+            response_data_curr["tool_names"] = tool_names
+            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:
+                with contextlib.suppress(json.JSONDecodeError):
+                    message = json.loads(message)
+            response_data_curr["message"] = message
+
+        if response_choice.logprobs and response_choice.logprobs.content is not None:
+            logprobs_list: list[dict[str, Any] | list[dict[str, Any]]] = []
+            for logprob in response_choice.logprobs.content:
+                if logprob.top_logprobs:
+                    curr_logprob_infos = []
+                    for top_logprob in logprob.top_logprobs:
+                        curr_logprob_infos.append(
+                            {
+                                "token": top_logprob.token,
+                                "logprob": top_logprob.logprob,
+                                "bytes": top_logprob.bytes,
+                            }
+                        )
+                    logprobs_list.append(curr_logprob_infos)
+                else:
+                    logprobs_list.append(
+                        {
+                            "token": logprob.token,
+                            "logprob": logprob.logprob,
+                            "bytes": logprob.bytes,
+                        }
+                    )
+
+            response_data_curr["logprobs"] = logprobs_list
+        response_data["choices"].append(response_data_curr)
+
+    usage = response.usage
+    if usage is not None:
+        response_data["completion_tokens"] = usage.completion_tokens
+        response_data["prompt_tokens"] = usage.prompt_tokens
+
+    if seed is not None and response.system_fingerprint is not None:
+        response_data["system_fingerprint"] = response.system_fingerprint
+
+    if len(response_data["choices"]) == 1:
+        response_data.update(response_data["choices"][0])
+        del response_data["choices"]
+
+    return response_data

not_again_ai/llm/{context_management.py → openai_api/context_management.py}

@@ -1,6 +1,6 @@
 import copy
 
-from not_again_ai.llm.tokens import num_tokens_from_messages, truncate_str
+from not_again_ai.llm.openai_api.tokens import num_tokens_from_messages, truncate_str
 
 
 def _inject_variable(

{not_again_ai-0.6.0.dist-info → not_again_ai-0.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: not-again-ai
-Version: 0.6.0
+Version: 0.8.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,7 +21,8 @@ Provides-Extra: llm
 Provides-Extra: statistics
 Provides-Extra: viz
 Requires-Dist: numpy (>=1.26.4,<2.0.0) ; extra == "statistics" or extra == "viz"
-Requires-Dist: openai (>=1.23.2,<2.0.0) ; extra == "llm"
+Requires-Dist: ollama (>=0.1.9,<0.2.0) ; extra == "llm"
+Requires-Dist: openai (>=1.25.1,<2.0.0) ; extra == "llm"
 Requires-Dist: pandas (>=2.2.2,<3.0.0) ; extra == "viz"
 Requires-Dist: python-liquid (>=1.12.1,<2.0.0) ; extra == "llm"
 Requires-Dist: scikit-learn (>=1.4.2,<2.0.0) ; extra == "statistics"
@@ -47,9 +48,9 @@ Description-Content-Type: text/markdown
 [ruff-badge]: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json
 [mypy-badge]: https://www.mypy-lang.org/static/mypy_badge.svg
 
-**not-again-ai** is a collection of various functionalities that come up over and over again when developing AI projects as a Data/Applied/Research Scientist. It is designed to be simple and minimize dependencies first and foremost. This is not meant to reinvent the wheel, but instead be a home for functions that don’t belong well elsewhere. Additionally, feel free to **a)** use this as a template for your own Python package. **b)** instead of installing the package, copy and paste functions into your own projects (this is made easier with the limited amount of dependencies and the MIT license).
+**not-again-ai** is a collection of various building blocks that come up over and over again when developing AI products. The key goals of this package are to have simple, but flexible interfaces and to minimize dependencies. Feel free to **a)** use this as a template for your own Python package. **b)** instead of installing the package, copy and paste functions into your own projects (this is made possible with the limited amount of dependencies and the MIT license).
 
-**Documentation** available within the [readmes](readmes) or auto-generated at [DaveCoDev.github.io/not-again-ai/](https://DaveCoDev.github.io/not-again-ai/).
+**Documentation** available within individual **[notebooks](notebooks)**, docstrings within the source, or auto-generated at [DaveCoDev.github.io/not-again-ai/](https://DaveCoDev.github.io/not-again-ai/).
 
 # Installation
 
@@ -61,82 +62,26 @@ Install the entire package from [PyPI](https://pypi.org/project/not-again-ai/) w
 $ pip install not_again_ai[llm,statistics,viz]
 ```
 
-The package is split into subpackages, so you can install only the parts you need.
+The package is split into subpackages, so you can install only the parts you need. See the **[notebooks](notebooks)** for examples.
 * **Base only**: `pip install not_again_ai`
-* **LLM only**: `pip install not_again_ai[llm]`
+* **LLM**: `pip install not_again_ai[llm]`
+    1. If you wish to use OpenAI
+        1. Go to https://platform.openai.com/settings/profile?tab=api-keys to get your API key.
+        1. (Optionally) Set the `OPENAI_API_KEY` and the `OPENAI_ORG_ID` environment variables.
+    1. If you wish to use Ollama:
+        1. follow the instructions to install ollama for your system: https://github.com/ollama/ollama
+        1. [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. If you'd like to make the ollama service accessible on your local network and it is hosted on Linux, add the following to the `/etc/systemd/system/ollama.service` file:
+            ```bash
+            [Service]
+            ...
+            Environment="OLLAMA_HOST=0.0.0.0"
+            ```
+        Now ollama will be available at `http://<local_address>:11434`
 * **Statistics**: `pip install not_again_ai[statistics]`
 * **Visualization**: `pip install not_again_ai[viz]`
 
 
-# Quick Tour
-
-## Base
-[README](https://github.com/DaveCoDev/not-again-ai/blob/main/readmes/base.md)
-
-The base package includes only functions that have minimal external dependencies and are useful in a variety of situations such as parallelization and filesystem operations.
-
-## LLM (Large Language Model)
-[README](https://github.com/DaveCoDev/not-again-ai/blob/main/readmes/llm.md), [Example Notebooks](https://github.com/DaveCoDev/not-again-ai/blob/main/notebooks/llm/)
-
-Supports OpenAI chat completions and text embeddings. Includes functions for creating chat completion prompts, token management, and context management.
-
-One example: 
-```python
-client = openai_client()
-messages = [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello!"}]
-response = chat_completion(messages=messages, model="gpt-3.5-turbo", max_tokens=100, client=client)["message"]
->>> "Hello! How can I help you today?"
-```
-
-## Statistics
-[README](https://github.com/DaveCoDev/not-again-ai/blob/main/readmes/statistics.md)
-
-We provide a few helpers for data analysis such as:
-
-```python
-from not_again_ai.statistics.dependence import pearson_correlation
-# quadratic dependence
->>> x = (rs.rand(500) * 4) - 2
->>> y = x**2 + (rs.randn(500) * 0.2)
->>> pearson_correlation(x, y)
-0.05
-```
-
-## Visualization
-[README](https://github.com/DaveCoDev/not-again-ai/blob/main/readmes/viz.md)
-
-We offer opinionated wrappers around seaborn to make common visualizations easier to create and customize.
-
-```python
->>> import numpy as np
->>> import pandas as pd
->>> from not_again_ai.viz.time_series import ts_lineplot
->>> from not_again_ai.viz.distributions import univariate_distplot
-
-# get some time series data
->>> rs = np.random.RandomState(365)
->>> values = rs.randn(365, 4).cumsum(axis=0)
->>> dates = pd.date_range('1 1 2021', periods=365, freq='D')
-# plot the time series and save it to a file
->>> ts_lineplot(ts_data=values, save_pathname='myplot.png', ts_x=dates, ts_names=['A', 'B', 'C', 'D'])
-
-# get a random distribution
->>> distrib = np.random.beta(a=0.5, b=0.5, size=1000)
-# plot the distribution and save it to a file
->>> univariate_distplot(
-...     data=distrib, 
-...     save_pathname='mydistribution.svg', 
-...     print_summary=False, bins=100, 
-...     title=r'Beta Distribution $\alpha=0.5, \beta=0.5$'
-... )
-```
-
-<p float="center">
-  <img src="https://raw.githubusercontent.com/DaveCoDev/not-again-ai/44c53fb7fb07234aaceea40c90d8cb74e5fa6c15/assets/distributions_test4.svg" width="404" />
-  <img src="https://raw.githubusercontent.com/DaveCoDev/not-again-ai/44c53fb7fb07234aaceea40c90d8cb74e5fa6c15/assets/ts_lineplot5.svg" width="404" /> 
-</p>
-
-
 # Development Information
 
 The following information is relevant if you would like to contribute or use this package as a template for yourself. 

not_again_ai-0.8.0.dist-info/RECORD

@@ -0,0 +1,31 @@
+not_again_ai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+not_again_ai/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+not_again_ai/base/file_system.py,sha256=KNQmacO4Q__CQuq2oPzWrg3rQO48n3evglc9bNiP7KM,949
+not_again_ai/base/parallel.py,sha256=fcYhKBYBWvob84iKp3O93wvFFdXeidljZsShgBLTNGA,3448
+not_again_ai/llm/__init__.py,sha256=_wNUL6FDaT369Z8W48FsaC_NkcOZ-ib2MMUvnaLOS-0,451
+not_again_ai/llm/chat_completion.py,sha256=a5jqce5qQzwdDFK8W7XRqGjTdFHAqKcWDmrTDtKls9U,3070
+not_again_ai/llm/ollama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+not_again_ai/llm/ollama/chat_completion.py,sha256=EKvqpHztsEISH9skYaLgKABEeoDhUlUyrakz_v6zvTw,3682
+not_again_ai/llm/ollama/ollama_client.py,sha256=dktyw7aKFq4EA3dU7Le5UpfsSq3Oh_POmYSrAI4qLi8,765
+not_again_ai/llm/ollama/service.py,sha256=XczbxISTAp4KHnIkqRZaMsfBohH-TAHrjZ8T9x3cRAY,2900
+not_again_ai/llm/openai_api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+not_again_ai/llm/openai_api/chat_completion.py,sha256=UJljrAV1lS2UvA5Sjt6N7S-9DaAgKRBpswLNDDP9RBI,8623
+not_again_ai/llm/openai_api/context_management.py,sha256=SKksCHs-C-UXtSksvXSvWWR4eeUyuyVc6XIhGgqJUF4,3435
+not_again_ai/llm/openai_api/embeddings.py,sha256=4OBnxZicrY6q4dQhuPqMdAnifyjwrsKMTDj-kVre0yc,2500
+not_again_ai/llm/openai_api/openai_client.py,sha256=6pZw2xw9X-ceV22rhApwFJ2tAKCxi-SxkjxBsTBZ2Nw,2470
+not_again_ai/llm/openai_api/prompts.py,sha256=7cDfvIKCTYM0t5lK34FLLqYf-SR_cynDXIXw3zWDizA,7094
+not_again_ai/llm/openai_api/tokens.py,sha256=2WGHzZJ0mIFAGpkuS_otjFzwhqjaSMgLoP2FVMnJTiE,4301
+not_again_ai/py.typed,sha256=UaCuPFa3H8UAakbt-5G8SPacldTOGvJv18pPjUJ5gDY,93
+not_again_ai/statistics/__init__.py,sha256=gA8r9JQFbFSN0ykrHy4G1IQgcky4f2eM5Oo24oVI5Ik,466
+not_again_ai/statistics/dependence.py,sha256=yZDk_e3ng96mp4hu8dDtQ0-uIn6KdSuGRS9uyM0O3x0,4429
+not_again_ai/viz/__init__.py,sha256=MeaWae_QRbDEHJ4MWYoY1-Ad6S0FhSDaRhQncS2cpSc,447
+not_again_ai/viz/barplots.py,sha256=xhxEXf4mfaymxDQtpphg9auy03zQ9e29R6L9yHixW1Y,3382
+not_again_ai/viz/distributions.py,sha256=OyWwJaNI6lMRm_iSrhq-CORLNvXfeuLSgDtVo3umnzU,4354
+not_again_ai/viz/scatterplot.py,sha256=eBtIf0Tf_1EcN-akRNJgvwLU0zpRx1zOl0VF9QTnbZA,2290
+not_again_ai/viz/time_series.py,sha256=pOGZqXp_2nd6nKo-PUQNCtmMh__69jxQ6bQibTGLwZA,5212
+not_again_ai/viz/utils.py,sha256=hN7gwxtBt3U6jQni2K8j5m5pCXpaJDoNzGhBBikEU28,238
+not_again_ai-0.8.0.dist-info/LICENSE,sha256=btjOgNGpp-ux5xOo1Gx1MddxeWtT9sof3s3Nui29QfA,1071
+not_again_ai-0.8.0.dist-info/METADATA,sha256=xM1pNeu0diA32AzTmWV_u8KNXpqkdjRBzrzsfJ9VGr0,14203
+not_again_ai-0.8.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+not_again_ai-0.8.0.dist-info/entry_points.txt,sha256=EMJegugnmJUd-jMUA_qIRMIPAasbei8gP6O4-ER0BxQ,61
+not_again_ai-0.8.0.dist-info/RECORD,,

{not_again_ai-0.6.0.dist-info → not_again_ai-0.8.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.8.1
+Generator: poetry-core 1.9.0
 Root-Is-Purelib: true
 Tag: py3-none-any

not_again_ai-0.6.0.dist-info/RECORD

@@ -1,25 +0,0 @@
-not_again_ai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-not_again_ai/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-not_again_ai/base/file_system.py,sha256=SX1ab2igdcFGjvdh4rDJIYVXbYC-jsYnCiM4oZtutAU,344
-not_again_ai/base/parallel.py,sha256=fcYhKBYBWvob84iKp3O93wvFFdXeidljZsShgBLTNGA,3448
-not_again_ai/llm/__init__.py,sha256=_wNUL6FDaT369Z8W48FsaC_NkcOZ-ib2MMUvnaLOS-0,451
-not_again_ai/llm/chat_completion.py,sha256=rjgaswSX4lf5ufPsv5SKWsXvCglLMFbX4y6bUvueUNQ,8528
-not_again_ai/llm/context_management.py,sha256=LXNd_L-jxurIFR1bO7gWakEH8I-cZmvLW83dFtvfJa4,3424
-not_again_ai/llm/embeddings.py,sha256=4OBnxZicrY6q4dQhuPqMdAnifyjwrsKMTDj-kVre0yc,2500
-not_again_ai/llm/openai_client.py,sha256=6pZw2xw9X-ceV22rhApwFJ2tAKCxi-SxkjxBsTBZ2Nw,2470
-not_again_ai/llm/prompts.py,sha256=7cDfvIKCTYM0t5lK34FLLqYf-SR_cynDXIXw3zWDizA,7094
-not_again_ai/llm/tokens.py,sha256=2WGHzZJ0mIFAGpkuS_otjFzwhqjaSMgLoP2FVMnJTiE,4301
-not_again_ai/py.typed,sha256=UaCuPFa3H8UAakbt-5G8SPacldTOGvJv18pPjUJ5gDY,93
-not_again_ai/statistics/__init__.py,sha256=gA8r9JQFbFSN0ykrHy4G1IQgcky4f2eM5Oo24oVI5Ik,466
-not_again_ai/statistics/dependence.py,sha256=yZDk_e3ng96mp4hu8dDtQ0-uIn6KdSuGRS9uyM0O3x0,4429
-not_again_ai/viz/__init__.py,sha256=MeaWae_QRbDEHJ4MWYoY1-Ad6S0FhSDaRhQncS2cpSc,447
-not_again_ai/viz/barplots.py,sha256=xhxEXf4mfaymxDQtpphg9auy03zQ9e29R6L9yHixW1Y,3382
-not_again_ai/viz/distributions.py,sha256=OyWwJaNI6lMRm_iSrhq-CORLNvXfeuLSgDtVo3umnzU,4354
-not_again_ai/viz/scatterplot.py,sha256=eBtIf0Tf_1EcN-akRNJgvwLU0zpRx1zOl0VF9QTnbZA,2290
-not_again_ai/viz/time_series.py,sha256=pOGZqXp_2nd6nKo-PUQNCtmMh__69jxQ6bQibTGLwZA,5212
-not_again_ai/viz/utils.py,sha256=hN7gwxtBt3U6jQni2K8j5m5pCXpaJDoNzGhBBikEU28,238
-not_again_ai-0.6.0.dist-info/LICENSE,sha256=btjOgNGpp-ux5xOo1Gx1MddxeWtT9sof3s3Nui29QfA,1071
-not_again_ai-0.6.0.dist-info/METADATA,sha256=3OVUTI8dpkQ9qCxLt2r7nFHk8tXPWDh1jnGAhlHIOUQ,15966
-not_again_ai-0.6.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-not_again_ai-0.6.0.dist-info/entry_points.txt,sha256=EMJegugnmJUd-jMUA_qIRMIPAasbei8gP6O4-ER0BxQ,61
-not_again_ai-0.6.0.dist-info/RECORD,,