pipelines/examples/scaffolds/example_pipeline_scaffold.py
Timothy J. Baek 8aa82f9eb9 chore
2024-06-01 11:45:29 -07:00

63 lines
2.0 KiB
Python

from typing import List, Union, Generator, Iterator
from schemas import OpenAIChatMessage
from pydantic import BaseModel
class Pipeline:
class Valves(BaseModel):
pass
def __init__(self):
# Optionally, you can set the id and name of the pipeline.
# Assign a unique identifier to the pipeline.
# The identifier must be unique across all pipelines.
# The identifier must be an alphanumeric string that can include underscores or hyphens. It cannot contain spaces, special characters, slashes, or backslashes.
self.id = "pipeline_example"
self.name = "Pipeline Example"
pass
async def on_startup(self):
# This function is called when the server is started.
print(f"on_startup:{__name__}")
pass
async def on_shutdown(self):
# This function is called when the server is stopped.
print(f"on_shutdown:{__name__}")
pass
async def on_valves_updated(self):
# This function is called when the valves are updated.
pass
async def inlet(self, body: dict, user: dict) -> dict:
# This function is called before the OpenAI API request is made. You can modify the form data before it is sent to the OpenAI API.
print(f"inlet:{__name__}")
print(body)
print(user)
return body
async def outlet(self, body: dict, user: dict) -> dict:
# This function is called after the OpenAI API response is completed. You can modify the messages after they are received from the OpenAI API.
print(f"outlet:{__name__}")
print(body)
print(user)
return body
def pipe(
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"pipe:{__name__}")
print(messages)
print(user_message)
print(body)
return f"{__name__} response to: {user_message}"