not-again-ai 0.14.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

not_again_ai/llm/chat_completion/__init__.py

@@ -0,0 +1,4 @@
+from not_again_ai.llm.chat_completion.interface import chat_completion
+from not_again_ai.llm.chat_completion.types import ChatCompletionRequest
+
+__all__ = ["ChatCompletionRequest", "chat_completion"]

not_again_ai/llm/chat_completion/interface.py

@@ -0,0 +1,32 @@
+from collections.abc import Callable
+from typing import Any
+
+from not_again_ai.llm.chat_completion.providers.ollama_api import ollama_chat_completion
+from not_again_ai.llm.chat_completion.providers.openai_api import openai_chat_completion
+from not_again_ai.llm.chat_completion.types import ChatCompletionRequest, ChatCompletionResponse
+
+
+def chat_completion(
+    request: ChatCompletionRequest,
+    provider: str,
+    client: Callable[..., Any],
+) -> ChatCompletionResponse:
+    """Get a chat completion response from the given provider. Currently supported providers:
+    - `openai` - OpenAI
+    - `azure_openai` - Azure OpenAI
+    - `ollama` - Ollama
+
+    Args:
+        request: Request parameter object
+        provider: The supported provider name
+        client: Client information, see the provider's implementation for what can be provided
+
+    Returns:
+        ChatCompletionResponse: The chat completion response.
+    """
+    if provider == "openai" or provider == "azure_openai":
+        return openai_chat_completion(request, client)
+    elif provider == "ollama":
+        return ollama_chat_completion(request, client)
+    else:
+        raise ValueError(f"Provider {provider} not supported")

not_again_ai/llm/chat_completion/providers/ollama_api.py

@@ -0,0 +1,227 @@
+from collections.abc import Callable
+import json
+import os
+import re
+import time
+from typing import Any, Literal, cast
+
+from loguru import logger
+from ollama import ChatResponse, Client, ResponseError
+
+from not_again_ai.llm.chat_completion.types import (
+    AssistantMessage,
+    ChatCompletionChoice,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    Function,
+    ToolCall,
+)
+
+OLLAMA_PARAMETER_MAP = {
+    "frequency_penalty": "repeat_penalty",
+    "max_completion_tokens": "num_predict",
+    "context_window": "num_ctx",
+    "n": None,
+    "tool_choice": None,
+    "reasoning_effort": None,
+    "parallel_tool_calls": None,
+    "logit_bias": None,
+    "top_logprobs": None,
+    "presence_penalty": None,
+}
+
+
+def validate(request: ChatCompletionRequest) -> None:
+    if request.json_mode and request.structured_outputs is not None:
+        raise ValueError("json_schema and json_mode cannot be used together.")
+
+    # Check if any of the parameters set to OLLAMA_PARAMETER_MAP are not None
+    for key, value in OLLAMA_PARAMETER_MAP.items():
+        if value is None and getattr(request, key) is not None:
+            logger.warning(f"Parameter {key} is not supported by Ollama and will be ignored.")
+
+    # If "stop" is not None, check if it is just a string
+    if isinstance(request.stop, list):
+        logger.warning("Parameter 'stop' needs to be a string and not a list. It will be ignored.")
+        request.stop = None
+
+
+def ollama_chat_completion(
+    request: ChatCompletionRequest,
+    client: Callable[..., Any],
+) -> ChatCompletionResponse:
+    validate(request)
+
+    kwargs = request.model_dump(mode="json", exclude_none=True)
+
+    # For each key in OLLAMA_PARAMETER_MAP
+    # If it is not None, set the key in kwargs to the value of the corresponding value in OLLAMA_PARAMETER_MAP
+    # If it is None, remove that key from kwargs
+    for key, value in OLLAMA_PARAMETER_MAP.items():
+        if value is not None and key in kwargs:
+            kwargs[value] = kwargs.pop(key)
+        elif value is None and key in kwargs:
+            del kwargs[key]
+
+    # If json_mode is True, set the format to json
+    json_mode = kwargs.get("json_mode", None)
+    if json_mode:
+        kwargs["format"] = "json"
+        kwargs.pop("json_mode")
+    elif json_mode is not None and not json_mode:
+        kwargs.pop("json_mode")
+
+    # If structured_outputs is not None, set the format to structured_outputs
+    if kwargs.get("structured_outputs", None):
+        # Check if the schema is in the OpenAI and pull out the schema
+        if "schema" in kwargs["structured_outputs"]:
+            kwargs["format"] = kwargs["structured_outputs"]["schema"]
+            kwargs.pop("structured_outputs")
+        else:
+            kwargs["format"] = kwargs.pop("structured_outputs")
+
+    option_fields = [
+        "mirostat",
+        "mirostat_eta",
+        "mirostat_tau",
+        "num_ctx",
+        "repeat_last_n",
+        "repeat_penalty",
+        "temperature",
+        "seed",
+        "stop",
+        "tfs_z",
+        "num_predict",
+        "top_k",
+        "top_p",
+        "min_p",
+    ]
+    # For each field in option_fields, if it is in kwargs, make it under an options dictionary
+    options = {}
+    for field in option_fields:
+        if field in kwargs:
+            options[field] = kwargs.pop(field)
+    kwargs["options"] = options
+
+    for message in kwargs["messages"]:
+        role = message.get("role", None)
+        # For each ToolMessage, remove the name field
+        if role is not None and role == "tool":
+            message.pop("name")
+
+        # For each AssistantMessage with tool calls, remove the id field
+        if role is not None and role == "assistant" and message.get("tool_calls", None):
+            for tool_call in message["tool_calls"]:
+                tool_call.pop("id")
+
+        # Content and images need to be separated
+        images = []
+        content = ""
+        if isinstance(message["content"], list):
+            for item in message["content"]:
+                if item["type"] == "image_url":
+                    image_url = item["image_url"]["url"]
+                    # Remove the data URL prefix if present
+                    if image_url.startswith("data:"):
+                        image_url = image_url.split("base64,", 1)[1]
+                    images.append(image_url)
+                else:
+                    content += item["text"]
+        else:
+            content = message["content"]
+
+        message["content"] = content
+        if len(images) > 1:
+            images = images[:1]
+            logger.warning("Ollama model only supports a single image per message. Using only the first images.")
+        message["images"] = images
+
+    try:
+        start_time = time.time()
+        response: ChatResponse = client(**kwargs)
+        end_time = time.time()
+        response_duration = round(end_time - start_time, 4)
+    except ResponseError as e:
+        # If the error says "model 'model' not found" use regex then raise a more specific error
+        expected_pattern = f"model '{request.model}' not found"
+        if re.search(expected_pattern, e.error):
+            raise ResponseError(f"Model '{request.model}' not found.") from e
+        else:
+            raise ResponseError(e.error) from e
+
+    errors = ""
+
+    # Handle tool calls
+    tool_calls: list[ToolCall] | None = None
+    if response.message.tool_calls:
+        parsed_tool_calls: list[ToolCall] = []
+        for tool_call in response.message.tool_calls:
+            tool_name = tool_call.function.name
+            if request.tools and tool_name not in [tool["function"]["name"] for tool in request.tools]:
+                errors += f"Tool call {tool_call} has an invalid tool name: {tool_name}\n"
+            tool_args = tool_call.function.arguments
+            parsed_tool_calls.append(
+                ToolCall(
+                    id="",
+                    function=Function(
+                        name=tool_name,
+                        arguments=tool_args,
+                    ),
+                )
+            )
+        tool_calls = parsed_tool_calls
+
+    json_message = None
+    if (request.json_mode or (request.structured_outputs is not None)) and response.message.content:
+        try:
+            json_message = json.loads(response.message.content)
+        except json.JSONDecodeError:
+            errors += "Message failed to parse into JSON\n"
+
+    finish_reason = cast(
+        Literal["stop", "length", "tool_calls", "content_filter"],
+        "stop" if response.done_reason is None else response.done_reason or "stop",
+    )
+
+    choice = ChatCompletionChoice(
+        message=AssistantMessage(
+            content=response.message.content or "",
+            tool_calls=tool_calls,
+        ),
+        finish_reason=finish_reason,
+        json_message=json_message,
+    )
+
+    return ChatCompletionResponse(
+        choices=[choice],
+        errors=errors.strip(),
+        completion_tokens=response.get("eval_count", -1),
+        prompt_tokens=response.get("prompt_eval_count", -1),
+        response_duration=response_duration,
+    )
+
+
+def ollama_client(host: str | None = None, timeout: float | None = None) -> Callable[..., Any]:
+    """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:
+            logger.warning("OLLAMA_HOST environment variable not set, using default host: http://localhost:11434")
+            host = "http://localhost:11434"
+
+    def client_callable(**kwargs: Any) -> Any:
+        client = Client(host=host, timeout=timeout)
+        return client.chat(**kwargs)
+
+    return client_callable

not_again_ai/llm/chat_completion/providers/openai_api.py

@@ -0,0 +1,290 @@
+from collections.abc import Callable
+import json
+import time
+from typing import Any, Literal
+
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from openai import AzureOpenAI, OpenAI
+
+from not_again_ai.llm.chat_completion.types import (
+    AssistantMessage,
+    ChatCompletionChoice,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    Function,
+    ToolCall,
+)
+
+OPENAI_PARAMETER_MAP = {
+    "context_window": None,
+    "mirostat": None,
+    "mirostat_eta": None,
+    "mirostat_tau": None,
+    "repeat_last_n": None,
+    "tfs_z": None,
+    "top_k": None,
+    "min_p": None,
+}
+
+
+def validate(request: ChatCompletionRequest) -> None:
+    if request.json_mode and request.structured_outputs is not None:
+        raise ValueError("json_schema and json_mode cannot be used together.")
+
+
+def openai_chat_completion(
+    request: ChatCompletionRequest,
+    client: Callable[..., Any],
+) -> ChatCompletionResponse:
+    validate(request)
+
+    # Format the response format parameters to be compatible with OpenAI API
+    if request.json_mode:
+        response_format: dict[str, Any] = {"type": "json_object"}
+    elif request.structured_outputs is not None:
+        response_format = {"type": "json_schema", "json_schema": request.structured_outputs}
+    else:
+        response_format = {"type": "text"}
+
+    kwargs = request.model_dump(mode="json", exclude_none=True)
+
+    # For each key in OPENAI_PARAMETER_MAP
+    # If it is not None, set the key in kwargs to the value of the corresponding value in OPENAI_PARAMETER_MAP
+    # If it is None, remove that key from kwargs
+    for key, value in OPENAI_PARAMETER_MAP.items():
+        if value is not None and key in kwargs:
+            kwargs[value] = kwargs.pop(key)
+        elif value is None and key in kwargs:
+            del kwargs[key]
+
+    # Iterate over each message and
+    for message in kwargs["messages"]:
+        role = message.get("role", None)
+        # For each ToolMessage, change the "name" field to be named "tool_call_id" instead
+        if role is not None and role == "tool":
+            message["tool_call_id"] = message.pop("name")
+
+        # For each AssistantMessage with tool calls, make the function arguments a string
+        if role is not None and role == "assistant" and message.get("tool_calls", None):
+            for tool_call in message["tool_calls"]:
+                tool_call["function"]["arguments"] = str(tool_call["function"]["arguments"])
+
+    # Delete the json_mode and structured_outputs from kwargs
+    kwargs.pop("json_mode", None)
+    kwargs.pop("structured_outputs", None)
+
+    # Add the response_format to kwargs
+    kwargs["response_format"] = response_format
+
+    # Handle tool_choice when the provided tool_choice the name of the required tool.
+    if request.tool_choice is not None and request.tool_choice not in ["none", "auto", "required"]:
+        kwargs["tool_choice"] = {"type": "function", "function": {"name": request.tool_choice}}
+
+    start_time = time.time()
+    response = client(**kwargs)
+    end_time = time.time()
+    response_duration = round(end_time - start_time, 4)
+
+    errors = ""
+    extras: dict[str, Any] = {}
+    choices: list[ChatCompletionChoice] = []
+    for index, choice in enumerate(response["choices"]):
+        choice_extras: dict[str, Any] = {}
+        finish_reason = choice["finish_reason"]
+
+        message = choice["message"]
+        tool_calls: list[ToolCall] | None = None
+        if message.get("tool_calls", None):
+            parsed_tool_calls: list[ToolCall] = []
+            for tool_call in message["tool_calls"]:
+                tool_name = tool_call.get("function", {}).get("name", None)
+                # Check if the tool name is valid (one of the tool names in the request)
+                if request.tools and tool_name not in [tool["function"]["name"] for tool in request.tools]:
+                    errors += f"Choice {index}: Tool call {tool_call} has an invalid tool name: {tool_name}\n"
+
+                tool_args = tool_call.get("function", {}).get("arguments", None)
+                try:
+                    tool_args = json.loads(tool_args)
+                except json.JSONDecodeError:
+                    errors += f"Choice {index}: Tool call {tool_call} failed to parse arguments into JSON\n"
+
+                parsed_tool_calls.append(
+                    ToolCall(
+                        id=tool_call["id"],
+                        function=Function(
+                            name=tool_name,
+                            arguments=tool_args,
+                        ),
+                    )
+                )
+            tool_calls = parsed_tool_calls
+
+        json_message = None
+        if request.json_mode or (request.structured_outputs is not None):
+            try:
+                json_message = json.loads(message.get("content", "{}"))
+            except json.JSONDecodeError:
+                errors += f"Choice {index}: Message failed to parse into JSON\n"
+
+        # Handle logprobs
+        logprobs: list[dict[str, Any] | list[dict[str, Any]]] | None = None
+        if choice.get("logprobs", None) and choice["logprobs"].get("content", None) is not None:
+            logprobs_list: list[dict[str, Any] | list[dict[str, Any]]] = []
+            for logprob in choice["logprobs"]["content"]:
+                if logprob.get("top_logprobs", None):
+                    curr_logprob_infos: list[dict[str, Any]] = []
+                    for top_logprob in logprob.get("top_logprobs", []):
+                        curr_logprob_infos.append(
+                            {
+                                "token": top_logprob.get("token", ""),
+                                "logprob": top_logprob.get("logprob", 0),
+                                "bytes": top_logprob.get("bytes", 0),
+                            }
+                        )
+                    logprobs_list.append(curr_logprob_infos)
+                else:
+                    logprobs_list.append(
+                        {
+                            "token": logprob.get("token", ""),
+                            "logprob": logprob.get("logprob", 0),
+                            "bytes": logprob.get("bytes", 0),
+                        }
+                    )
+            logprobs = logprobs_list
+
+        # Handle extras that OpenAI or Azure OpenAI return
+        if choice.get("content_filter_results", None):
+            choice_extras["content_filter_results"] = choice["content_filter_results"]
+
+        choices.append(
+            ChatCompletionChoice(
+                message=AssistantMessage(
+                    content=message.get("content") or "",
+                    refusal=message.get("refusal", None),
+                    tool_calls=tool_calls,
+                ),
+                finish_reason=finish_reason,
+                json_message=json_message,
+                logprobs=logprobs,
+                extras=choice_extras,
+            )
+        )
+
+    completion_tokens = response["usage"].get("completion_tokens", -1)
+    prompt_tokens = response["usage"].get("prompt_tokens", -1)
+    completion_detailed_tokens = response["usage"].get("completion_detailed_tokens", None)
+    prompt_detailed_tokens = response["usage"].get("prompt_detailed_tokens", None)
+    system_fingerprint = response.get("system_fingerprint", None)
+
+    extras["prompt_filter_results"] = response.get("prompt_filter_results", None)
+
+    return ChatCompletionResponse(
+        choices=choices,
+        errors=errors.strip(),
+        extras=extras,
+        completion_tokens=completion_tokens,
+        prompt_tokens=prompt_tokens,
+        completion_detailed_tokens=completion_detailed_tokens,
+        prompt_detailed_tokens=prompt_detailed_tokens,
+        system_fingerprint=system_fingerprint,
+        response_duration=response_duration,
+    )
+
+
+def create_client_callable(client_class: type[OpenAI | AzureOpenAI], **client_args: Any) -> Callable[..., Any]:
+    """Creates a callable that instantiates and uses an OpenAI client.
+
+    Args:
+        client_class: The OpenAI client class to instantiate (OpenAI or AzureOpenAI)
+        **client_args: Arguments to pass to the client constructor
+
+    Returns:
+        A callable that creates a client and returns completion results
+    """
+    filtered_args = {k: v for k, v in client_args.items() if v is not None}
+
+    def client_callable(**kwargs: Any) -> Any:
+        client = client_class(**filtered_args)
+        completion = client.chat.completions.create(**kwargs)
+        return completion.to_dict()
+
+    return client_callable
+
+
+class InvalidOAIAPITypeError(Exception):
+    """Raised when an invalid OAIAPIType string is provided."""
+
+
+def openai_client(
+    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,
+) -> Callable[..., Any]:
+    """Create an OpenAI or Azure OpenAI client instance based on the specified API type and other provided parameters.
+
+    It is preferred to use RBAC authentication for Azure OpenAI. 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'.
+            Defaults to 'openai'.
+        api_key (str, optional): The API key to authenticate the client. If not provided,
+            OpenAI automatically uses `OPENAI_API_KEY` from the environment.
+            If provided for Azure OpenAI, it will be used for authentication instead of the Azure AD token provider.
+        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.
+        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),
+            408 Request Timeout, 409 Conflict, 429 Rate Limit, and >=500 Internal errors are all retried by default.
+
+    Returns:
+        Callable[..., Any]: A callable that creates a client and returns completion results
+
+
+    Raises:
+        InvalidOAIAPITypeError: If an invalid API type string is provided.
+        NotImplementedError: If the specified API type is recognized but not yet supported (e.g., 'azure_openai').
+    """
+    if api_type not in ["openai", "azure_openai"]:
+        raise InvalidOAIAPITypeError(f"Invalid OAIAPIType: {api_type}. Must be 'openai' or 'azure_openai'.")
+
+    if api_type == "openai":
+        return create_client_callable(
+            OpenAI,
+            api_key=api_key,
+            organization=organization,
+            timeout=timeout,
+            max_retries=max_retries,
+        )
+    elif api_type == "azure_openai":
+        if api_key:
+            return create_client_callable(
+                AzureOpenAI,
+                api_version=aoai_api_version,
+                azure_endpoint=azure_endpoint,
+                api_key=api_key,
+                timeout=timeout,
+                max_retries=max_retries,
+            )
+        else:
+            azure_credential = DefaultAzureCredential()
+            ad_token_provider = get_bearer_token_provider(
+                azure_credential, "https://cognitiveservices.azure.com/.default"
+            )
+            return create_client_callable(
+                AzureOpenAI,
+                api_version=aoai_api_version,
+                azure_endpoint=azure_endpoint,
+                azure_ad_token_provider=ad_token_provider,
+                timeout=timeout,
+                max_retries=max_retries,
+            )
+    else:
+        raise NotImplementedError(f"API type '{api_type}' is invalid.")