Temporary agents and chats

backend/src/inference.py

@@ -1,4 +1,5 @@
 import threading
+import asyncio
 import ollama
 from stream import Stream
 
@@ -23,22 +24,33 @@ class Models:
 
 # llama cpp models
 model_cache = {}
+last_model_file = None
+last_model_instance = None
+
+def get_llama_cpp_model(app, model_file):
+    global last_model_file, last_model_instance
+    if last_model_file == model_file:
+        return last_model_instance
+    
+    if last_model_instance is not None:
+        #last_model_instance = None
+        del last_model_instance
+
+    last_model_file = model_file
+    last_model_instance = Llama(    
+        model_path=f"{app.data_directory}/models/{model_file}",
+        n_gpu_layers=-1, # Uncomment to use GPU acceleration
+        # seed=1337, # Uncomment to set a specific seed
+        n_ctx=CONTEXT_WINDOW,
+    )
+    return last_model_instance
 
 # llama.cpp 
 async def local_chat_completion(app, messages, model, opts={}):
     content_stream = Stream(app)
 
     def worker():
-        if model in model_cache:
-            llm = model_cache[model]
-        else:
-            llm = Llama(    
-                model_path=f"{app.data_directory}/models/{model}",
-                n_gpu_layers=-1, # Uncomment to use GPU acceleration
-                # seed=1337, # Uncomment to set a specific seed
-                n_ctx=CONTEXT_WINDOW,
-            )
-            model_cache[model] = llm
+        llm = get_llama_cpp_model(app, model)
 
         temp = opts.get("temperature")
         if temp is None:
@@ -77,23 +89,23 @@ async def ollama_list(app, messages, model, opts={}):
 async def ollama_chat_completion(app, messages, model, opts={}):
     content_stream = Stream(app)
 
-    print("START OLLAMA INFERENCE")
-    stream_response = await ollama.AsyncClient().chat(
-        model=model,
-        messages=[{'role': msg.role, 'content': msg.content} for msg in messages],
-        stream=True,
-    )
+    async def worker():
+        print("START OLLAMA INFERENCE")
+        stream_response = await ollama.AsyncClient().chat(
+            model=model,
+            messages=[{'role': msg.role, 'content': msg.content} for msg in messages],
+            stream=True,
+        )
 
-    async for part in stream_response:
-        content_stream.feed(part.message.content)
-        print(part['message']['content'], end='', flush=True)
+        async for part in stream_response:
+            content_stream.feed(part.message.content)
 
-    content_stream.close()
+        content_stream.close()
+        print("END OLLAMA INFERENCE")
+    
+    asyncio.create_task(worker())
 
-    print("END OLLAMA INFERENCE")
-    return await content_stream.content_wait()
-
-    #return content_stream
+    return content_stream
 
 
 async def replicate_completion(app, messages, model=None, opts={}):
@@ -104,7 +116,7 @@ async def replicate_completion(app, messages, model=None, opts={}):
 
     arguments = {}
     arguments[parameter_names.get('prompt', 'prompt')] = format_instruct(messages)
-    arguments[parameter_names.get('temperature', 'temperature')] = opts.get('temperature', 0.75)
+    arguments[parameter_names.get('temperature', 'temperature')] = opts.get('temperature') or 0.75
     arguments[parameter_names.get('max_tokens', 'max_tokens')] = opts.get('max_tokens', 2048)
 
     arguments[parameter_names.get('system_prompt', 'system_prompt')] = ""