not-again-ai 0.7.0__tar.gz → 0.8.0__tar.gz

{not_again_ai-0.7.0 → not_again_ai-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: not-again-ai
-Version: 0.7.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

{not_again_ai-0.7.0 → not_again_ai-0.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "not-again-ai"
-version = "0.7.0"
+version = "0.8.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"

not_again_ai-0.8.0/src/not_again_ai/llm/chat_completion.py

@@ -0,0 +1,76 @@
+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, Any]],
+    model: str,
+    client: OpenAI | Client,
+    max_tokens: int | None = None,
+    temperature: float = 0.7,
+    json_mode: bool = False,
+    seed: int | None = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """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[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
+            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.
+            extras (dict): This will contain any additional fields returned by corresponding provider.
+    """
+    # 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-0.7.0/src/not_again_ai/llm/openai → not_again_ai-0.8.0/src/not_again_ai/llm/openai_api}/context_management.py

@@ -1,6 +1,6 @@
 import copy
 
-from not_again_ai.llm.openai.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(