From 2f655da81062fb0cbdfc69a7fc1fe97985d82a0b Mon Sep 17 00:00:00 2001
From: Ayman Hamed <ayman3000@gmail.com>
Date: Mon, 6 Jan 2025 17:11:19 +0200
Subject: [PATCH] added the ability to call Ollama client seamlessly

---
 pyproject.toml                |  1 +
 src/markitdown/_markitdown.py | 51 ++++++++++++++++++++++++-----------
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 9c113ad..43322ae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,6 +42,7 @@ dependencies = [
   "pathvalidate",
   "charset-normalizer",
   "openai",
+  "ollama"
 ]
 
 [project.urls]
diff --git a/src/markitdown/_markitdown.py b/src/markitdown/_markitdown.py
index b6acfe8..c7fb3e4 100644
--- a/src/markitdown/_markitdown.py
+++ b/src/markitdown/_markitdown.py
@@ -13,6 +13,7 @@
 import tempfile
 import traceback
 import zipfile
+from http.client import responses
 from xml.dom import minidom
 from typing import Any, Dict, List, Optional, Union
 from pathlib import Path
@@ -1085,24 +1086,42 @@ def _get_llm_description(self, local_path, extension, client, model, prompt=None
                 content_type = "image/jpeg"
             image_base64 = base64.b64encode(image_file.read()).decode("utf-8")
             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": [
-                    {"type": "text", "text": prompt},
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": data_uri,
-                        },
-                    },
-                ],
-            }
-        ]
+                }
+            ]
+
+            response = client.chat(
+            model = model,
+            messages = messages,
 
-        response = client.chat.completions.create(model=model, messages=messages)
-        return response.choices[0].message.content
+        )
+
+            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
 
 
 class OutlookMsgConverter(DocumentConverter):