added the ability to call Ollama client seamlessly
This commit is contained in:
parent
05b78e7ce1
commit
2f655da810
2 changed files with 36 additions and 16 deletions
|
|
@ -42,6 +42,7 @@ dependencies = [
|
||||||
"pathvalidate",
|
"pathvalidate",
|
||||||
"charset-normalizer",
|
"charset-normalizer",
|
||||||
"openai",
|
"openai",
|
||||||
|
"ollama"
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.urls]
|
[project.urls]
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,7 @@ import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
import traceback
|
import traceback
|
||||||
import zipfile
|
import zipfile
|
||||||
|
from http.client import responses
|
||||||
from xml.dom import minidom
|
from xml.dom import minidom
|
||||||
from typing import Any, Dict, List, Optional, Union
|
from typing import Any, Dict, List, Optional, Union
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
@ -1085,7 +1086,26 @@ class ImageConverter(MediaConverter):
|
||||||
content_type = "image/jpeg"
|
content_type = "image/jpeg"
|
||||||
image_base64 = base64.b64encode(image_file.read()).decode("utf-8")
|
image_base64 = base64.b64encode(image_file.read()).decode("utf-8")
|
||||||
data_uri = f"data:{content_type};base64,{image_base64}"
|
data_uri = f"data:{content_type};base64,{image_base64}"
|
||||||
|
# check if Ollama client
|
||||||
|
if str(type(client)) == "<class 'ollama._client.Client'>":
|
||||||
|
messages = [
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": prompt,
|
||||||
|
'images': [local_path]
|
||||||
|
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
||||||
|
response = client.chat(
|
||||||
|
model = model,
|
||||||
|
messages = messages,
|
||||||
|
|
||||||
|
)
|
||||||
|
|
||||||
|
return response.message.content
|
||||||
|
|
||||||
|
else:# use openai
|
||||||
messages = [
|
messages = [
|
||||||
{
|
{
|
||||||
"role": "user",
|
"role": "user",
|
||||||
|
|
@ -1100,7 +1120,6 @@ class ImageConverter(MediaConverter):
|
||||||
],
|
],
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|
||||||
response = client.chat.completions.create(model=model, messages=messages)
|
response = client.chat.completions.create(model=model, messages=messages)
|
||||||
return response.choices[0].message.content
|
return response.choices[0].message.content
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue