From 12e415615d0aa270ae4d8f0d3f82823d4827373a Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 28 May 2024 17:07:54 -0700
Subject: [PATCH] feat: pipelines function

---
 main.py                                               | 10 +++++++++-
 .../examples/litellm_subprocess_manifold_pipeline.py  | 11 ++++++-----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/main.py b/main.py
index f8fa864..d486d6a 100644
--- a/main.py
+++ b/main.py
@@ -49,7 +49,15 @@ def get_all_pipelines():
 
         if hasattr(pipeline, "type"):
             if pipeline.type == "manifold":
-                for p in pipeline.pipelines:
+                manifold_pipelines = []
+
+                # Check if pipelines is a function or a list
+                if callable(pipeline.pipelines):
+                    manifold_pipelines = pipeline.pipelines()
+                else:
+                    manifold_pipelines = pipeline.pipelines
+
+                for p in manifold_pipelines:
                     manifold_pipeline_id = f'{pipeline_id}.{p["id"]}'
 
                     manifold_pipeline_name = p["name"]
diff --git a/pipelines/examples/litellm_subprocess_manifold_pipeline.py b/pipelines/examples/litellm_subprocess_manifold_pipeline.py
index 88ff220..1f7fef7 100644
--- a/pipelines/examples/litellm_subprocess_manifold_pipeline.py
+++ b/pipelines/examples/litellm_subprocess_manifold_pipeline.py
@@ -32,10 +32,8 @@ class Pipeline:
             LITELLM_PROXY_HOST: str = "127.0.0.1"
             litellm_config: dict = {}
 
-        # Initialize rate limits
+        # Initialize Valves
         self.valves = Valves(**{"LITELLM_CONFIG_DIR": f"./litellm/config.yaml"})
-        self.pipelines = []
-
         self.background_process = None
         pass
 
@@ -102,8 +100,6 @@ class Pipeline:
             self.background_process = process
             print("Subprocess started successfully.")
 
-            self.pipelines = self.get_litellm_models()
-
             # Capture STDERR for debugging purposes
             stderr_output = await process.stderr.read()
             stderr_text = stderr_output.decode().strip()
@@ -172,6 +168,11 @@ class Pipeline:
         else:
             return []
 
+    # Pipelines are the models that are available in the manifold.
+    # It can be a list or a function that returns a list.
+    def pipelines(self) -> List[dict]:
+        return self.get_litellm_models()
+
     def pipe(
         self, user_message: str, model_id: str, messages: List[dict], body: dict
     ) -> Union[str, Generator, Iterator]: