From 0fb7c9f6a992385acfb04604b2c21ebad5bc136a Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Thu, 21 Dec 2023 16:37:07 -0800
Subject: [PATCH] examples

---
 examples/simple-chat-stream/main.py     | 6 ++----
 examples/simple-chat/main.py            | 2 +-
 examples/simple-generate-stream/main.py | 5 +++++
 examples/simple-generate/main.py        | 5 +++++
 examples/simple-multimodal/main.py      | 2 +-
 5 files changed, 14 insertions(+), 6 deletions(-)
 create mode 100644 examples/simple-generate-stream/main.py
 create mode 100644 examples/simple-generate/main.py

diff --git a/examples/simple-chat-stream/main.py b/examples/simple-chat-stream/main.py
index 33d1740..2a57346 100644
--- a/examples/simple-chat-stream/main.py
+++ b/examples/simple-chat-stream/main.py
@@ -8,10 +8,8 @@
   },
 ]
 
-for message in chat('mistral', messages=messages, stream=True):
-  if message := message.get('message'):
-    if message.get('role') == 'assistant':
-      print(message.get('content', ''), end='', flush=True)
+for part in chat('mistral', messages=messages, stream=True):
+  print(part['message']['content'], end='', flush=True)
 
 # end with a newline
 print()
diff --git a/examples/simple-chat/main.py b/examples/simple-chat/main.py
index 5019307..90c5f90 100644
--- a/examples/simple-chat/main.py
+++ b/examples/simple-chat/main.py
@@ -9,4 +9,4 @@
 ]
 
 response = chat('mistral', messages=messages)
-print(response['message'])
+print(response['message']['content'])
diff --git a/examples/simple-generate-stream/main.py b/examples/simple-generate-stream/main.py
new file mode 100644
index 0000000..a24b410
--- /dev/null
+++ b/examples/simple-generate-stream/main.py
@@ -0,0 +1,5 @@
+from ollama import generate
+
+
+for part in generate('mistral', 'Why is the sky blue?', stream=True):
+  print(part['response'], end='', flush=True)
diff --git a/examples/simple-generate/main.py b/examples/simple-generate/main.py
new file mode 100644
index 0000000..e39e295
--- /dev/null
+++ b/examples/simple-generate/main.py
@@ -0,0 +1,5 @@
+from ollama import generate
+
+
+response = generate('mistral', 'Why is the sky blue?')
+print(response['response'])
diff --git a/examples/simple-multimodal/main.py b/examples/simple-multimodal/main.py
index 97eba59..44b3716 100644
--- a/examples/simple-multimodal/main.py
+++ b/examples/simple-multimodal/main.py
@@ -24,6 +24,6 @@
 raw.raise_for_status()
 
 for response in generate('llava', 'explain this comic:', images=[raw.content], stream=True):
-  print(response.get('response'), end='', flush=True)
+  print(response['response'], end='', flush=True)
 
 print()