2024-06-10 03:14:52 +01:00
|
|
|
import asyncio
|
|
|
|
import typing
|
|
|
|
from functools import partial
|
|
|
|
from fuzzywuzzy.fuzz import ratio
|
|
|
|
from ollama import Message
|
|
|
|
|
2024-06-10 16:56:12 +01:00
|
|
|
if typing.TYPE_CHECKING:
|
|
|
|
from .config import ServerConfig
|
|
|
|
|
2024-06-10 03:14:52 +01:00
|
|
|
|
|
|
|
__all__ = (
|
|
|
|
'async_ratio',
|
|
|
|
'create_ollama_message',
|
2024-06-10 16:56:12 +01:00
|
|
|
'find_suitable_server',
|
|
|
|
'decorate_server_name'
|
2024-06-10 03:14:52 +01:00
|
|
|
)
|
|
|
|
|
|
|
|
|
2024-06-10 16:56:12 +01:00
|
|
|
def decorate_server_name(_s: "ServerConfig") -> str:
|
|
|
|
if _s.gpu:
|
|
|
|
return f"{_s.name} (\u26A1)"
|
|
|
|
return _s.name
|
|
|
|
|
|
|
|
|
2024-06-10 03:14:52 +01:00
|
|
|
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
|
|
|
|
)
|
|
|
|
)
|
2024-06-10 16:56:12 +01:00
|
|
|
|
|
|
|
|
|
|
|
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.")
|