48 lines
1.1 KiB
Python
48 lines
1.1 KiB
Python
|
import asyncio
|
||
|
|
import typing
|
||
|
|
from functools import partial
|
||
|
|
from fuzzywuzzy.fuzz import ratio
|
||
|
|
from ollama import Message
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = (
|
||
|
|
'async_ratio',
|
||
|
|
'create_ollama_message',
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
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
|
||
|
|
)
|
||
|
|
)
|