From b2b1dd853c8bce3becd96d0b22ee74b2eaf1eaa8 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Sun, 26 May 2024 16:17:51 -0700
Subject: [PATCH] feat: ollama manifold pipeline example

---
 .../examples/ollama_manifold_pipeline.py      | 52 +++++++++++++------
 1 file changed, 35 insertions(+), 17 deletions(-)

diff --git a/pipelines/examples/ollama_manifold_pipeline.py b/pipelines/examples/ollama_manifold_pipeline.py
index b5d141c..752fb6e 100644
--- a/pipelines/examples/ollama_manifold_pipeline.py
+++ b/pipelines/examples/ollama_manifold_pipeline.py
@@ -1,27 +1,30 @@
 from typing import List, Union, Generator, Iterator
 from schemas import OpenAIChatMessage
+import requests
 
 
 class Pipeline:
     def __init__(self):
-        self.id = "manifold_pipeline"
-        self.name = "Manifold Pipeline"
         # You can also set the pipelines that are available in this pipeline.
         # Set manifold to True if you want to use this pipeline as a manifold.
         # Manifold pipelines can have multiple pipelines.
         self.manifold = True
-        self.pipelines = [
-            {
-                "id": "pipeline-1",  # This will turn into `manifold_pipeline.pipeline-1`
-                "name": "Manifold: Pipeline 1",
-            },
-            {
-                "id": "pipeline-2",
-                "name": "Manifold: Pipeline 2",
-            },
-        ]
+        self.id = "ollama_manifold"
+        # Optionally, you can set the name of the manifold pipeline.
+        self.name = "Ollama: "
+
+        self.OLLAMA_BASE_URL = "http://localhost:11434"
+        self.pipelines = self.get_ollama_models()
         pass
 
+    def get_ollama_models(self):
+        r = requests.get(f"{self.OLLAMA_BASE_URL}/api/tags")
+        models = r.json()
+
+        return [
+            {"id": model["model"], "name": model["name"]} for model in models["models"]
+        ]
+
     async def on_startup(self):
         # This function is called when the server is started.
         print(f"on_startup:{__name__}")
@@ -36,10 +39,25 @@ class Pipeline:
         self, user_message: str, model_id: str, messages: List[dict], body: dict
     ) -> Union[str, Generator, Iterator]:
         # This is where you can add your custom pipelines like RAG.'
-        print(f"get_response:{__name__}")
 
-        print(messages)
-        print(user_message)
-        print(body)
+        if "user" in body:
+            print("######################################")
+            print(f'# User: {body["user"]["name"]} ({body["user"]["id"]})')
+            print(f"# Message: {user_message}")
+            print("######################################")
 
-        return f"{model_id} response to: {user_message}"
+        try:
+            r = requests.post(
+                url=f"{self.OLLAMA_BASE_URL}/v1/chat/completions",
+                json={**body, "model": model_id},
+                stream=True,
+            )
+
+            r.raise_for_status()
+
+            if body["stream"]:
+                return r.iter_lines()
+            else:
+                return r.json()
+        except Exception as e:
+            return f"Error: {e}"