From 64127ee821f11b8812e2b9a2b8ddd31bd5518caa Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Wed, 10 Jan 2024 17:02:49 -0800
Subject: [PATCH] parse json error

---
 README.md        | 15 +++++++++++++++
 ollama/_types.py |  8 ++++++++
 2 files changed, 23 insertions(+)

diff --git a/README.md b/README.md
index ec7801c..770cc6a 100644
--- a/README.md
+++ b/README.md
@@ -68,3 +68,18 @@ async def chat():
 
 asyncio.run(chat())
 ```
+
+## Handling Errors
+
+Errors are raised if requests return an error status or if an error is detected while streaming.
+
+```python
+model = 'does-not-yet-exist'
+
+try:
+  ollama.chat(model)
+except ollama.ResponseError as e:
+  print('Error:', e.content)
+  if e.status_code == 404:
+    ollama.pull(model)
+```
diff --git a/ollama/_types.py b/ollama/_types.py
index b8fa06c..fca614a 100644
--- a/ollama/_types.py
+++ b/ollama/_types.py
@@ -1,3 +1,4 @@
+import json
 from typing import Any, TypedDict, Sequence, Literal
 
 import sys
@@ -145,6 +146,13 @@ class ResponseError(Exception):
   """
 
   def __init__(self, content: str, status_code: int = -1):
+    try:
+      # try to parse content as JSON and extract 'error'
+      # fallback to raw content if JSON parsing fails
+      content = json.loads(content).get('error', content)
+    except json.JSONDecodeError:
+      ...
+
     super().__init__(content)
     self.content = content
     "Reason for the error."