OpenWebUI/pipelines
1
0
Fork 0
mirror of https://github.com/open-webui/pipelines synced 2025-05-15 09:55:45 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: ollama manifold pipeline example

Browse Source
This commit is contained in:
Timothy J. Baek 2024-05-26 16:17:51 -07:00
parent 918b709f30
commit b2b1dd853c
1 changed files with 35 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
pipelines/examples/ollama_manifold_pipeline.py
View File

@ -1,27 +1,30 @@
from typing import List, Union, Generator, Iterator from typing import List, Union, Generator, Iterator
from schemas import OpenAIChatMessage from schemas import OpenAIChatMessage
import requests
class Pipeline: class Pipeline:
def __init__(self): def __init__(self):
self.id = "manifold_pipeline"
self.name = "Manifold Pipeline"
# You can also set the pipelines that are available in this 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. # Set manifold to True if you want to use this pipeline as a manifold.
# Manifold pipelines can have multiple pipelines. # Manifold pipelines can have multiple pipelines.
self.manifold = True self.manifold = True
self.pipelines = [ self.id = "ollama_manifold"
{ # Optionally, you can set the name of the manifold pipeline.
"id": "pipeline-1", # This will turn into `manifold_pipeline.pipeline-1` self.name = "Ollama: "
"name": "Manifold: Pipeline 1",
}, self.OLLAMA_BASE_URL = "http://localhost:11434"
{ self.pipelines = self.get_ollama_models()
"id": "pipeline-2",
"name": "Manifold: Pipeline 2",
},
]
pass 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): async def on_startup(self):
# This function is called when the server is started. # This function is called when the server is started.
print(f"on_startup:{__name__}") print(f"on_startup:{__name__}")
@ -36,10 +39,25 @@ class Pipeline:
self, user_message: str, model_id: str, messages: List[dict], body: dict self, user_message: str, model_id: str, messages: List[dict], body: dict
) -> Union[str, Generator, Iterator]: ) -> Union[str, Generator, Iterator]:
# This is where you can add your custom pipelines like RAG.' # This is where you can add your custom pipelines like RAG.'
print(f"get_response:{__name__}")
print(messages) if "user" in body:
print(user_message) print("######################################")
print(body) 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}"