Merge pull request #27 from bananaml/fix/endpoint-collision

raise Exception on endpoint collision + handle empty init

example.py

@@ -17,8 +17,7 @@ def init():
 
     return context
 
-
-@app.handler()
+@app.handler(route = "/")
 def handler(context: dict, request: Request) -> Response:
     prompt = request.json.get("prompt")
     model = context.get("model")
@@ -29,6 +28,5 @@ def handler(context: dict, request: Request) -> Response:
         status=200
     )
 
-
 if __name__ == "__main__":
-    app.serve()
+    app.serve()

potassium/potassium.py

@@ -31,11 +31,12 @@ class Potassium():
     def __init__(self, name):
         self.name = name
 
-        def default_func():
-            return
+        # default init function, if the user doesn't specify one
+        def empty_init():
+            return {}
 
         # semi-private vars, not intended for users to modify
-        self._init_func = default_func
+        self._init_func = empty_init
         self._endpoints = {}  # dictionary to store unlimited Endpoints, by unique route
         self._context = {}
         self._lock = Lock()
@@ -53,9 +54,25 @@ def init(self, func):
         """
 
         def wrapper():
+            print(colored("Running init()", 'yellow'))
             self._context = func()
+            if not isinstance(self._context, dict):
+                raise Exception("Potassium init() must return a dictionary")
+
         self._init_func = wrapper
         return wrapper
+
+    @staticmethod
+    def _standardize_route(route):
+        # handle empty or None case
+        if len(route) == 0 or route == None:
+            route = "/"
+
+        # prepend with "/" if not already, as router expects
+        if route[0] != "/":
+            route = "/" + route
+
+        return route
 
     # handler is a blocking http POST handler
     def handler(self, route: str = "/"):
@@ -76,19 +93,29 @@ def wrapper(request):
 
                 return out
 
+            nonlocal route # we need to modify the route variable in the outer scope
+            route = self._standardize_route(route)
+            if route in self._endpoints:
+                raise Exception("Route already in use")
+
             self._endpoints[route] = Endpoint(type="handler", func=wrapper)
             return wrapper
         return actual_decorator
 
     # background is a non-blocking http POST handler
-    def background(self, route: str = "/"):
+    def background(self, route: str = "/"):    
         "background is a non-blocking http POST handler"
         def actual_decorator(func):
             @functools.wraps(func)
             def wrapper(request):
                 # send in app's stateful context if GPU, and the request
                 return func(self._context, request)
-
+
+            nonlocal route # we need to modify the route variable in the outer scope
+            route = self._standardize_route(route)
+            if route in self._endpoints:
+                raise Exception("Route already in use")
+
             self._endpoints[route] = Endpoint(
                 type="background", func=wrapper)
             return wrapper
@@ -179,7 +206,6 @@ def handle(path):
     # serve runs the http server
     def serve(self, host="0.0.0.0", port=8000):
         print(colored("------\nStarting Potassium Server 🍌", 'yellow'))
-        print(colored("Running init()", 'yellow'))
         self._init_func()
         flask_app = self._create_flask_app()
         server = make_server(host, port, flask_app)

setup.py

@@ -8,7 +8,7 @@
 setup(
     name='potassium',
     packages=['potassium'],
-    version='0.1.1',
+    version='0.1.2',
     license='Apache License 2.0',
     # Give a short description about your library
     description='The potassium package is a flask-like HTTP server for serving large AI models',