From fb043a74d524a7fa7ab337128ab74e10d59b78d7 Mon Sep 17 00:00:00 2001
From: artitw <artitw@gmail.com>
Date: Sun, 22 Sep 2024 22:27:01 +0000
Subject: [PATCH] Assistant startup fix

---
 setup.py               |  2 +-
 text2text/assistant.py | 14 ++++++++------
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/setup.py b/setup.py
index 3e28258..93e75a5 100755
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
   name="text2text",
-  version="1.4.9",
+  version="1.5.0",
   author="artitw",
   author_email="artitw@gmail.com",
   description="Text2Text: Crosslingual NLP/G toolkit",
diff --git a/text2text/assistant.py b/text2text/assistant.py
index 5baa3af..ef9f362 100644
--- a/text2text/assistant.py
+++ b/text2text/assistant.py
@@ -1,7 +1,8 @@
 import os
 import ollama
 import psutil
-import time
+import subprocess
+
 from llama_index.llms.ollama import Ollama
 from llama_index.core.llms import ChatMessage
 
@@ -21,15 +22,16 @@ def __init__(self, **kwargs):
     return_code = os.system("curl -fsSL https://ollama.com/install.sh | sh")
     if return_code != 0:
       print("Cannot install ollama.")
+    return_code = os.system("sudo systemctl enable ollama")
     self.load_model()
     self.client = ollama.Client(host=self.model_url)
 
   def load_model(self):
-    return_code = os.system("sudo service ollama stop")
-    return_code = os.system(f"ollama serve & ollama pull {self.model_name}")
-    time.sleep(5.0)
-    if return_code != 0:
-      print(f"{self.model_name} is not loading up. Restarting and trying again might help. Maybe needs more memory.")
+    sub = subprocess.Popen(
+      f"ollama serve & ollama pull {self.model_name}", 
+      shell=True, 
+      stdout=subprocess.PIPE
+    )
 
   def chat_completion(self, messages=[{"role": "user", "content": "hello"}], stream=False, schema=None, **kwargs):
     if is_port_in_use(self.port):