sentient-jimmy/jimmy/utils.py

import asyncio
import typing
from functools import partial
from fuzzywuzzy.fuzz import ratio
from ollama import Message

if typing.TYPE_CHECKING:
    from .config import ServerConfig


__all__ = (
    'async_ratio',
    'create_ollama_message',
    'find_suitable_server',
    'decorate_server_name'
)


def decorate_server_name(_s: "ServerConfig") -> str:
    if _s.gpu:
        return f"{_s.name} (\u26A1)"
    return _s.name


async def async_ratio(a: str, b: str) -> int:
    """
    Wraps fuzzywuzzy ratio in an async function

    :param a: str - first string
    :param b: str - second string
    :return: int - ratio of similarity
    """
    return await asyncio.to_thread(partial(ratio, a, b))


async def create_ollama_message(
        content: str,
        role: typing.Literal["system", "assistant", "user"] = "user",
        images: typing.List[str | bytes] = None
) -> Message:
    """
    Create a message for ollama.

    :param content: str - the content of the message
    :param role: str - the role of the message
    :param images: list - the images to attach to the message
    :return: dict - the message
    """
    def factory(**kwargs):
        return Message(**kwargs)
    return await asyncio.to_thread(
        partial(
            factory,
            role=role,
            content=content,
            images=images
        )
    )


async def find_suitable_server(cpu_fallback: bool = True) -> "ServerConfig":
    """
    Finds a suitable server to use for Ollama.

    :param cpu_fallback: bool - whether to fall back to CPU servers if GPU servers are unavailable.
    :return: ServerConfig - the server to use
    """
    from .config import get_servers
    servers = get_servers()
    if not servers:
        raise ValueError("No servers configured.")
    for server in servers:
        if cpu_fallback is False and server.gpu is False:
            continue
        if not await server.is_online():
            continue
        return server
    raise ValueError("No servers available.")