added the ability to call Ollama client seamlessly

This commit is contained in:
Ayman Hamed 2025-01-06 17:11:19 +02:00
parent 05b78e7ce1
commit 2f655da810
2 changed files with 36 additions and 16 deletions

View file

@ -42,6 +42,7 @@ dependencies = [
"pathvalidate", "pathvalidate",
"charset-normalizer", "charset-normalizer",
"openai", "openai",
"ollama"
] ]
[project.urls] [project.urls]

View file

@ -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,24 +1086,42 @@ 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]
messages = [ }
{ ]
"role": "user",
"content": [ response = client.chat(
{"type": "text", "text": prompt}, model = model,
{ messages = messages,
"type": "image_url",
"image_url": { )
"url": data_uri,
return response.message.content
else:# use openai
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{
"type": "image_url",
"image_url": {
"url": data_uri,
},
}, },
}, ],
], }
} ]
] response = client.chat.completions.create(model=model, messages=messages)
return response.choices[0].message.content
response = client.chat.completions.create(model=model, messages=messages)
return response.choices[0].message.content
class OutlookMsgConverter(DocumentConverter): class OutlookMsgConverter(DocumentConverter):