diff --git a/main.py b/main.py index 61e50af..eaf291e 100644 --- a/main.py +++ b/main.py @@ -38,8 +38,10 @@ for loaded_module in load_modules_from_directory("./pipelines"): # Do something with the loaded module print("Loaded:", loaded_module.__name__) + pipeline = loaded_module.Pipeline() + PIPELINES[loaded_module.__name__] = { - "module": loaded_module.Pipeline(), + "module": pipeline, "id": loaded_module.__name__, "name": loaded_module.__name__, } diff --git a/pipelines/examples/pipeline_example.py b/pipelines/examples/pipeline_example.py index 2d987c1..4341c97 100644 --- a/pipelines/examples/pipeline_example.py +++ b/pipelines/examples/pipeline_example.py @@ -2,23 +2,27 @@ from typing import List, Union, Generator from schemas import OpenAIChatMessage -def get_response( - user_message: str, messages: List[OpenAIChatMessage] -) -> Union[str, Generator]: - # This is where you can add your custom pipelines like RAG. +class Pipeline: + def __init__(self): + pass - print(messages) - print(user_message) + async def on_startup(self): + # This function is called when the server is started. + print(f"on_startup:{__name__}") + pass - return f"pipeline response to: {user_message}" + async def on_shutdown(self): + # This function is called when the server is stopped. + print(f"on_shutdown:{__name__}") + pass + def get_response( + self, user_message: str, messages: List[OpenAIChatMessage] + ) -> Union[str, Generator]: + # This is where you can add your custom pipelines like RAG.' + print(f"get_response:{__name__}") -async def on_startup(): - # This function is called when the server is started. - print(f"on_startup:{__name__}") - pass + print(messages) + print(user_message) - -async def on_shutdown(): - # This function is called when the server is stopped. - pass + return f"{__name__} response to: {user_message}" diff --git a/pipelines/pipeline.py b/pipelines/pipeline.py index a70aabf..4341c97 100644 --- a/pipelines/pipeline.py +++ b/pipelines/pipeline.py @@ -8,21 +8,21 @@ class Pipeline: async def on_startup(self): # This function is called when the server is started. - print("onstartup") - print(__name__) - + print(f"on_startup:{__name__}") pass - async def on_shutdown(): + async def on_shutdown(self): # This function is called when the server is stopped. + print(f"on_shutdown:{__name__}") pass def get_response( self, user_message: str, messages: List[OpenAIChatMessage] ) -> Union[str, Generator]: - # 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) print(user_message) - return f"pipeline response to: {user_message}" + return f"{__name__} response to: {user_message}"