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

not_again_ai/llm/chat_completion/types.py

@@ -0,0 +1,145 @@
+from enum import Enum
+from typing import Any, Generic, Literal, TypeVar
+
+from pydantic import BaseModel, Field
+
+
+class Role(str, Enum):
+    ASSISTANT = "assistant"
+    DEVELOPER = "developer"
+    SYSTEM = "system"
+    TOOL = "tool"
+    USER = "user"
+
+
+class ContentPartType(str, Enum):
+    TEXT = "text"
+    IMAGE = "image_url"
+
+
+class TextContent(BaseModel):
+    type: Literal[ContentPartType.TEXT] = ContentPartType.TEXT
+    text: str
+
+
+class ImageDetail(str, Enum):
+    AUTO = "auto"
+    LOW = "low"
+    HIGH = "high"
+
+
+class ImageUrl(BaseModel):
+    url: str
+    detail: ImageDetail = ImageDetail.AUTO
+
+
+class ImageContent(BaseModel):
+    type: Literal[ContentPartType.IMAGE] = ContentPartType.IMAGE
+    image_url: ImageUrl
+
+
+ContentT = TypeVar("ContentT", bound=str | list[TextContent | ImageContent])
+
+
+class BaseMessage(BaseModel, Generic[ContentT]):
+    content: ContentT
+    role: Role
+    name: str | None = None
+
+
+class Function(BaseModel):
+    name: str
+    arguments: dict[str, Any]
+
+
+class ToolCall(BaseModel):
+    id: str
+    function: Function
+    type: Literal["function"] = "function"
+
+
+class DeveloperMessage(BaseMessage[str]):
+    role: Literal[Role.DEVELOPER] = Role.DEVELOPER
+
+
+class SystemMessage(BaseMessage[str]):
+    role: Literal[Role.SYSTEM] = Role.SYSTEM
+
+
+class UserMessage(BaseMessage[str | list[TextContent | ImageContent]]):
+    role: Literal[Role.USER] = Role.USER
+
+
+class AssistantMessage(BaseMessage[str]):
+    role: Literal[Role.ASSISTANT] = Role.ASSISTANT
+    refusal: str | None = None
+    tool_calls: list[ToolCall] | None = None
+
+
+class ToolMessage(BaseMessage[str]):
+    # A tool message's name field will be interpreted as "tool_call_id"
+    role: Literal[Role.TOOL] = Role.TOOL
+
+
+MessageT = AssistantMessage | DeveloperMessage | SystemMessage | ToolMessage | UserMessage
+
+
+class ChatCompletionRequest(BaseModel):
+    messages: list[MessageT]
+    model: str
+
+    max_completion_tokens: int | None = Field(default=None)
+    context_window: int | None = Field(default=None)
+    logprobs: bool | None = Field(default=None)
+    n: int | None = Field(default=None)
+
+    tools: list[dict[str, Any]] | None = Field(default=None)
+    tool_choice: str | None = Field(default=None)
+    parallel_tool_calls: bool | None = Field(default=None)
+    json_mode: bool | None = Field(default=None)
+    structured_outputs: dict[str, Any] | None = Field(default=None)
+
+    temperature: float | None = Field(default=None)
+    reasoning_effort: Literal["low", "medium", "high"] | None = Field(default=None)
+    top_p: float | None = Field(default=None)
+    logit_bias: dict[str, float] | None = Field(default=None)
+    top_logprobs: int | None = Field(default=None)
+    frequency_penalty: float | None = Field(default=None)
+    presence_penalty: float | None = Field(default=None)
+    stop: str | list[str] | None = Field(default=None)
+
+    seed: int | None = Field(default=None)
+
+    mirostat: int | None = Field(default=None)
+    mirostat_eta: float | None = Field(default=None)
+    mirostat_tau: float | None = Field(default=None)
+    repeat_last_n: int | None = Field(default=None)
+    tfs_z: float | None = Field(default=None)
+    top_k: int | None = Field(default=None)
+    min_p: float | None = Field(default=None)
+
+
+class ChatCompletionChoice(BaseModel):
+    message: AssistantMessage
+    finish_reason: Literal["stop", "length", "tool_calls", "content_filter"]
+
+    json_message: dict[str, Any] | None = Field(default=None)
+    logprobs: list[dict[str, Any] | list[dict[str, Any]]] | None = Field(default=None)
+
+    extras: Any | None = Field(default=None)
+
+
+class ChatCompletionResponse(BaseModel):
+    choices: list[ChatCompletionChoice]
+
+    errors: str = Field(default="")
+
+    completion_tokens: int
+    prompt_tokens: int
+    completion_detailed_tokens: dict[str, int] | None = Field(default=None)
+    prompt_detailed_tokens: dict[str, int] | None = Field(default=None)
+    response_duration: float
+
+    system_fingerprint: str | None = Field(default=None)
+
+    extras: Any | None = Field(default=None)

not_again_ai/llm/embedding/__init__.py

@@ -0,0 +1,4 @@
+from not_again_ai.llm.embedding.interface import create_embeddings
+from not_again_ai.llm.embedding.types import EmbeddingRequest
+
+__all__ = ["EmbeddingRequest", "create_embeddings"]

not_again_ai/llm/embedding/interface.py

@@ -0,0 +1,28 @@
+from collections.abc import Callable
+from typing import Any
+
+from not_again_ai.llm.embedding.providers.ollama_api import ollama_create_embeddings
+from not_again_ai.llm.embedding.providers.openai_api import openai_create_embeddings
+from not_again_ai.llm.embedding.types import EmbeddingRequest, EmbeddingResponse
+
+
+def create_embeddings(request: EmbeddingRequest, provider: str, client: Callable[..., Any]) -> EmbeddingResponse:
+    """Get a embedding 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:
+        EmbeddingResponse: The embedding response.
+    """
+    if provider == "openai" or provider == "azure_openai":
+        return openai_create_embeddings(request, client)
+    elif provider == "ollama":
+        return ollama_create_embeddings(request, client)
+    else:
+        raise ValueError(f"Provider {provider} not supported")

not_again_ai/llm/embedding/providers/ollama_api.py

@@ -0,0 +1,87 @@
+from collections.abc import Callable
+import os
+import re
+import time
+from typing import Any
+
+from loguru import logger
+from ollama import Client, EmbedResponse, ResponseError
+
+from not_again_ai.llm.embedding.types import EmbeddingObject, EmbeddingRequest, EmbeddingResponse
+
+OLLAMA_PARAMETER_MAP = {
+    "dimensions": None,
+}
+
+
+def validate(request: EmbeddingRequest) -> None:
+    # 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.")
+
+
+def ollama_create_embeddings(request: EmbeddingRequest, client: Callable[..., Any]) -> EmbeddingResponse:
+    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]
+
+    # Explicitly set truncate to True (it is the default)
+    kwargs["truncate"] = True
+
+    try:
+        start_time = time.time()
+        response: EmbedResponse = 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
+
+    embeddings: list[EmbeddingObject] = []
+    for index, embedding in enumerate(response.embeddings):
+        embeddings.append(EmbeddingObject(embedding=list(embedding), index=index))
+
+    return EmbeddingResponse(
+        embeddings=embeddings,
+        response_duration=response_duration,
+        total_tokens=response.prompt_eval_count,
+    )
+
+
+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.embed(**kwargs)
+
+    return client_callable

not_again_ai/llm/embedding/providers/openai_api.py

@@ -0,0 +1,126 @@
+from collections.abc import Callable
+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.embedding.types import EmbeddingObject, EmbeddingRequest, EmbeddingResponse
+
+
+def openai_create_embeddings(request: EmbeddingRequest, client: Callable[..., Any]) -> EmbeddingResponse:
+    kwargs = request.model_dump(mode="json", exclude_none=True)
+
+    start_time = time.time()
+    response = client(**kwargs)
+    end_time = time.time()
+    response_duration = round(end_time - start_time, 4)
+
+    embeddings: list[EmbeddingObject] = []
+    for data in response["data"]:
+        embeddings.append(EmbeddingObject(embedding=data["embedding"], index=data["index"]))
+
+    return EmbeddingResponse(
+        embeddings=embeddings,
+        response_duration=response_duration,
+        total_tokens=response["usage"]["total_tokens"],
+    )
+
+
+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.embeddings.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.")

not_again_ai/llm/embedding/types.py

@@ -0,0 +1,23 @@
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class EmbeddingRequest(BaseModel):
+    input: str | list[str]
+    model: str
+    dimensions: int | None = Field(default=None)
+
+
+class EmbeddingObject(BaseModel):
+    embedding: list[float]
+    index: int
+
+
+class EmbeddingResponse(BaseModel):
+    embeddings: list[EmbeddingObject]
+    total_tokens: int | None = Field(default=None)
+    response_duration: float
+
+    errors: str = Field(default="")
+    extras: Any | None = Field(default=None)

not_again_ai/llm/prompting/__init__.py

@@ -0,0 +1,3 @@
+from not_again_ai.llm.prompting.interface import Tokenizer
+
+__all__ = ["Tokenizer"]

not_again_ai/llm/prompting/compile_prompt.py

@@ -0,0 +1,125 @@
+import base64
+from collections.abc import Sequence
+from copy import deepcopy
+import mimetypes
+from pathlib import Path
+from typing import Any
+
+from liquid import Template
+from openai.lib._pydantic import to_strict_json_schema
+from pydantic import BaseModel
+
+from not_again_ai.llm.chat_completion.types import MessageT
+
+
+def _apply_templates(value: Any, variables: dict[str, str]) -> Any:
+    """Recursively applies Liquid templating to all string fields within the given value."""
+    if isinstance(value, str):
+        return Template(value).render(**variables)
+    elif isinstance(value, list):
+        return [_apply_templates(item, variables) for item in value]
+    elif isinstance(value, dict):
+        return {key: _apply_templates(val, variables) for key, val in value.items()}
+    elif isinstance(value, BaseModel):
+        # Process each field in the BaseModel by converting it to a dict,
+        # applying templating to its values, and then re-instantiating the model.
+        processed_data = {key: _apply_templates(val, variables) for key, val in value.model_dump().items()}
+        return value.__class__(**processed_data)
+    else:
+        return value
+
+
+def compile_messages(messages: Sequence[MessageT], variables: dict[str, str]) -> Sequence[MessageT]:
+    """Compiles messages using Liquid templating and the provided variables.
+    Calls Template(content_part).render(**variables) on each text content part.
+
+    Args:
+        messages: List of MessageT where content can contain Liquid templates.
+        variables: The variables to inject into the templates.
+
+    Returns:
+        The same list of messages with the content parts injected with the variables.
+    """
+    messages_formatted = deepcopy(messages)
+    messages_formatted = [_apply_templates(message, variables) for message in messages_formatted]
+    return messages_formatted
+
+
+def compile_tools(tools: Sequence[dict[str, Any]], variables: dict[str, str]) -> Sequence[dict[str, Any]]:
+    """Compiles a list of tool argument dictionaries using Liquid templating and provided variables.
+
+    Each dictionary in the list is deep copied and processed recursively to substitute any Liquid
+    templates present in its data structure.
+
+    Args:
+        tools: A list of dictionaries representing tool arguments, where values can include Liquid templates.
+        variables: A dictionary of variables to substitute into the Liquid templates.
+
+    Returns:
+        A new list of dictionaries with the Liquid templates replaced by their corresponding variable values.
+    """
+    tools_formatted = deepcopy(tools)
+    tools_formatted = [_apply_templates(tool, variables) for tool in tools_formatted]
+    return tools_formatted
+
+
+def encode_image(image_path: Path) -> str:
+    """Encodes an image file at the given Path to base64.
+
+    Args:
+        image_path: The path to the image file to encode.
+
+    Returns:
+        The base64 encoded image as a string.
+    """
+    with Path.open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+
+def create_image_url(image_path: Path) -> str:
+    """Creates a data URL for an image file at the given Path.
+
+    Args:
+        image_path: The path to the image file to encode.
+
+    Returns:
+        The data URL for the image.
+    """
+    image_data = encode_image(image_path)
+
+    valid_mime_types = ["image/jpeg", "image/png", "image/webp", "image/gif"]
+
+    # Get the MIME type from the image file extension
+    mime_type = mimetypes.guess_type(image_path)[0]
+
+    # Check if the MIME type is valid
+    # List of valid types is here: https://platform.openai.com/docs/guides/vision/what-type-of-files-can-i-upload
+    if mime_type not in valid_mime_types:
+        raise ValueError(f"Invalid MIME type for image: {mime_type}")
+
+    return f"data:{mime_type};base64,{image_data}"
+
+
+def pydantic_to_json_schema(
+    pydantic_model: type[BaseModel], schema_name: str, description: str | None = None
+) -> dict[str, Any]:
+    """Converts a Pydantic model to a JSON schema expected by Structured Outputs.
+    Must adhere to the supported schemas: https://platform.openai.com/docs/guides/structured-outputs/supported-schemas
+
+    Args:
+        pydantic_model: The Pydantic model to convert.
+        schema_name: The name of the schema.
+        description: An optional description of the schema.
+
+    Returns:
+        A JSON schema dictionary representing the Pydantic model.
+    """
+    converted_pydantic = to_strict_json_schema(pydantic_model)
+    schema = {
+        "name": schema_name,
+        "strict": True,
+        "schema": converted_pydantic,
+    }
+    if description:
+        schema["description"] = description
+    return schema

not_again_ai/llm/prompting/interface.py

@@ -0,0 +1,46 @@
+from collections.abc import Collection, Set
+from typing import Literal
+
+from loguru import logger
+
+from not_again_ai.llm.chat_completion.types import MessageT
+from not_again_ai.llm.prompting.providers.openai_tiktoken import TokenizerOpenAI
+from not_again_ai.llm.prompting.types import BaseTokenizer
+
+
+class Tokenizer(BaseTokenizer):
+    def __init__(
+        self,
+        model: str,
+        provider: str,
+        allowed_special: Literal["all"] | Set[str] | None = None,
+        disallowed_special: Literal["all"] | Collection[str] | None = None,
+    ):
+        self.model = model
+        self.provider = provider
+        self.allowed_special = allowed_special
+        self.disallowed_special = disallowed_special
+
+        self.init_tokenizer(model, provider, allowed_special, disallowed_special)
+
+    def init_tokenizer(
+        self,
+        model: str,
+        provider: str,
+        allowed_special: Literal["all"] | Set[str] | None = None,
+        disallowed_special: Literal["all"] | Collection[str] | None = None,
+    ) -> None:
+        if provider == "openai" or provider == "azure_openai":
+            self.tokenizer = TokenizerOpenAI(model, provider, allowed_special, disallowed_special)
+        else:
+            logger.warning(f"Provider {provider} not supported. Initializing using tiktoken and gpt-4o.")
+            self.tokenizer = TokenizerOpenAI("gpt-4o", "openai", allowed_special, disallowed_special)
+
+    def truncate_str(self, text: str, max_len: int) -> str:
+        return self.tokenizer.truncate_str(text, max_len)
+
+    def num_tokens_in_str(self, text: str) -> int:
+        return self.tokenizer.num_tokens_in_str(text)
+
+    def num_tokens_in_messages(self, messages: list[MessageT]) -> int:
+        return self.tokenizer.num_tokens_in_messages(messages)