From 30bc90e611be49b7c9922f68ea47bc13173e7716 Mon Sep 17 00:00:00 2001
From: artitw <artitw@gmail.com>
Date: Mon, 23 Sep 2024 01:26:00 +0000
Subject: [PATCH] verbosity and error handling

---
 setup.py               |  2 +-
 text2text/assistant.py | 26 ++++++++++++++------------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/setup.py b/setup.py
index c3e135d..2118c5b 100755
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
   name="text2text",
-  version="1.5.4",
+  version="1.5.6",
   author="artitw",
   author_email="artitw@gmail.com",
   description="Text2Text: Crosslingual NLP/G toolkit",
diff --git a/text2text/assistant.py b/text2text/assistant.py
index 08e52cc..7ce8ab4 100644
--- a/text2text/assistant.py
+++ b/text2text/assistant.py
@@ -15,7 +15,7 @@ def __init__(self, **kwargs):
     self.model_name = kwargs.get("model_name", "llama3.1")
     self.load_model()
     self.client = ollama.Client(host=self.model_url)
-    self.llama_index_client = Ollama(model=self.model_name, request_timeout=120.0)
+    self.structured_client = Ollama(model=self.model_name, request_timeout=120.0)
 
   def __del__(self):
     ollama.delete(self.model_name)
@@ -40,20 +40,22 @@ def load_model(self):
       
     result = ollama.pull(self.model_name)
     if result["status"] != "success":
-      print(f"Cannot pull {self.model_name}.")
+      print(f"Did not pull {self.model_name}.")
     
     time.sleep(10)
-
+    
   def chat_completion(self, messages=[{"role": "user", "content": "hello"}], stream=False, schema=None, **kwargs):
-    try:
-      if schema:
-        msgs = [ChatMessage(**m) for m in messages]
-        return self.llama_index_client.as_structured_llm(schema).chat(messages=msgs).raw
-      return self.client.chat(model=self.model_name, messages=messages, stream=stream)
-    except Exception as e:
-      print(e)
-      self.load_model()
-      return self.chat_completion(messages=messages, stream=stream, **kwargs)
+    if not ollama.ps():
+      result = ollama.pull(self.model_name)
+      if result["status"] == "success":
+        time.sleep(10)
+        return self.chat_completion(messages=messages, stream=stream, **kwargs)
+      raise Exception(f"Cannot pull {self.model_name}. Try restarting.")
+
+    if schema:
+      msgs = [ChatMessage(**m) for m in messages]
+      return self.structured_client.as_structured_llm(schema).chat(messages=msgs).raw
+    return self.client.chat(model=self.model_name, messages=messages, stream=stream)
 
   def embed(self, texts):
     return ollama.embed(model=self.model_name, input=texts)