remove flask dependency and serve extension through gateway

Author: thrau

openai/Makefile

@@ -0,0 +1,39 @@
+VENV_BIN = python3 -m venv
+VENV_DIR ?= .venv
+VENV_ACTIVATE = $(VENV_DIR)/bin/activate
+VENV_RUN = . $(VENV_ACTIVATE)
+
+venv: $(VENV_ACTIVATE)
+
+$(VENV_ACTIVATE): setup.py setup.cfg
+	test -d .venv || $(VENV_BIN) .venv
+	$(VENV_RUN); pip install --upgrade pip setuptools plux wheel
+	$(VENV_RUN); pip install --upgrade black isort pyproject-flake8 flake8-black flake8-isort
+	$(VENV_RUN); pip install -e .
+	touch $(VENV_DIR)/bin/activate
+
+clean:
+	rm -rf .venv/
+	rm -rf build/
+	rm -rf .eggs/
+	rm -rf *.egg-info/
+
+lint:              		  ## Run code linter to check code style
+	($(VENV_RUN); python -m pflake8 --show-source)
+
+format:            		  ## Run black and isort code formatter
+	$(VENV_RUN); python -m isort .; python -m black .
+
+install: venv
+	$(VENV_RUN); python setup.py develop
+
+dist: venv
+	$(VENV_RUN); python setup.py sdist bdist_wheel
+
+publish: clean-dist venv dist
+	$(VENV_RUN); pip install --upgrade twine; twine upload dist/*
+
+clean-dist: clean
+	rm -rf dist/
+
+.PHONY: clean clean-dist dist install publish

openai/localstack_openai/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.0"
+__version__ = "0.1.0"

openai/localstack_openai/extension.py

@@ -1,33 +1,38 @@
-import atexit
 import logging
 
-from localstack.extensions.api import Extension, http, services
+from localstack import config
+from localstack.extensions.api import Extension, http
+from rolo.router import RuleAdapter, WithHost
+from werkzeug.routing import Submount
 
 LOG = logging.getLogger(__name__)
 
 
 class LocalstackOpenAIExtension(Extension):
     name = "openai"
 
-    backend_url: str
+    submount = "/_extension/openai"
+    subdomain = "openai"
 
-    def on_platform_start(self):
-        # start localstripe when localstack starts
-        from . import mock_openai
-
-        port = services.external_service_ports.reserve_port()
-        self.backend_url = f"http://localhost:{port}"
-
-        print(f"Starting mock OpenAI service on {self.backend_url}")
-        mock_openai.run(port)
-        atexit.register(mock_openai.stop)
+    def on_extension_load(self):
+        logging.getLogger("localstack_openai").setLevel(
+            logging.DEBUG if config.DEBUG else logging.INFO
+        )
 
     def update_gateway_routes(self, router: http.Router[http.RouteHandler]):
-        # a ProxyHandler forwards all incoming requests to the backend URL
-        endpoint = http.ProxyHandler(self.backend_url)
+        from localstack_openai.mock_openai import Api
+
+        api = RuleAdapter(Api())
 
         # add path routes for localhost:4566/v1/chat/completion
         router.add(
-            "/v1/chat/completion",
-            endpoint=endpoint,
-        )
+            [
+                Submount(self.submount, [api]),
+                WithHost(f"{self.subdomain}.{config.LOCALSTACK_HOST.host}<__host__>", [api]),
+            ]
+        )
+
+        LOG.info(
+            "OpenAI mock available at %s%s", str(config.LOCALSTACK_HOST).rstrip("/"), self.submount
+        )
+        LOG.info("OpenAI mock available at %s", f"{self.subdomain}.{config.LOCALSTACK_HOST}")

openai/localstack_openai/mock_openai.py

@@ -1,14 +1,14 @@
 import json
 import time
-from faker import Faker
-from flask import Flask, request, Response
 
+from faker import Faker
+from rolo import Request, Response, route
 
 faker = Faker()
-app = Flask(__name__)
 
 res_len = 20
 
+
 class ChunkReader:
     def __init__(self, chunks, delay):
         self.ID = ""
@@ -18,12 +18,15 @@ def __init__(self, chunks, delay):
         self.SentDone = False
         self.Delay = delay
 
+
 def new_chunk_reader(cs, d):
     return ChunkReader(cs, d)
 
+
 def done(r):
     return r.SentFinished and r.SentDone
 
+
 def next_chunk(r):
     if r.SentDone:
         return None, None
@@ -74,6 +77,7 @@ def next_chunk(r):
     b = b"data: " + b + b"\n\n"
     return b, None
 
+
 def read(r, p):
     if done(r):
         return 0, None
@@ -130,106 +134,106 @@ def read(r, p):
     time.sleep(r.Delay)
     return n, None
 
-@app.route('/v1/chat/completions', methods=['POST'])
-def chat_completions():
-    data = request.get_data()
-    req = json.loads(data)
 
-    ws = [faker.word() for _ in range(res_len)]
-    ws = [" " + w if i > 0 else w for i, w in enumerate(ws)]
+class Api:
+
+    @route("/v1/chat/completions", methods=["POST"])
+    def chat_completions(self, request: Request):
+        data = request.get_data()
+        req = json.loads(data)
+
+        ws = [faker.word() for _ in range(res_len)]
+        ws = [" " + w if i > 0 else w for i, w in enumerate(ws)]
+
+        if not req.get("stream"):
+            m = "".join(ws)
+            return {
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": m,
+                        },
+                    }
+                ]
+            }
 
-    if not req.get("stream"):
-        m = "".join(ws)
-        return {
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": m,
-                    },
-                }
-            ]
-        }
+        id = faker.uuid4()
+        ct = int(time.time())
+        sd = 0.5
 
-    id = faker.uuid4()
-    ct = int(time.time())
-    sd = 0.5
+        def generate():
+            for w in ws:
+                b, _ = next_chunk(chunk_reader)
+                if b is not None:
+                    yield b
+                time.sleep(sd)
 
-    def generate():
-        for w in ws:
             b, _ = next_chunk(chunk_reader)
             if b is not None:
                 yield b
-            time.sleep(sd)
-
-        b, _ = next_chunk(chunk_reader)
-        if b is not None:
-            yield b
 
-        yield b"[done]\n"
+            yield b"[done]\n"
 
-    chunk_reader = new_chunk_reader(ws, sd)
-    return Response(generate(), content_type='text/event-stream')
+        chunk_reader = new_chunk_reader(ws, sd)
+        return Response(generate(), content_type="text/event-stream")
 
+    @route("/v1/audio/transcriptions", methods=["POST"])
+    def transcribe(self, request: Request):
+        return {
+            "text": faker.sentence(),
+        }
 
-@app.route('/v1/audio/transcriptions', methods=['POST'])
-def transcribe():
-    return {
-        "text": faker.sentence(),
-    }
-
-
-@app.route('/v1/audio/translations', methods=['POST'])
-def translate():
-    return {
-        "text": faker.sentence(),
-    }
-
+    @route("/v1/audio/translations", methods=["POST"])
+    def translate(self, request: Request):
+        return {
+            "text": faker.sentence(),
+        }
 
-@app.route('/v1/images/generations', methods=['POST'])
-def generate_image():
-    return {
-        "created": int(time.time()),
-        "data": [
-            {"url": faker.image_url()}
-        ]
-    }
+    @route("/v1/images/generations", methods=["POST"])
+    def generate_image(self, request: Request):
+        return {"created": int(time.time()), "data": [{"url": faker.image_url()}]}
 
-@app.route('/v1/engines', methods=['GET'])
-def list_engines():
-    return {
-        "object": "list",
-        "data": [
-            {
-                "id": "model-id-0",
-                "object": "model",
-                "created": 1686935002,
-                "owned_by": "organization-owner"
-            },
-            {
-                "id": "model-id-1",
-                "object": "model",
-                "created": 1686935002,
-                "owned_by": "organization-owner",
-            },
-            {
-                "id": "model-id-2",
-                "object": "model",
-                "created": 1686935002,
-                "owned_by": "openai"
-            },
-        ],
-        "object": "list"
-    }
+    @route("/v1/engines", methods=["GET"])
+    def list_engines(self, request: Request):
+        return {
+            "object": "list",
+            "data": [
+                {
+                    "id": "model-id-0",
+                    "object": "model",
+                    "created": 1686935002,
+                    "owned_by": "organization-owner",
+                },
+                {
+                    "id": "model-id-1",
+                    "object": "model",
+                    "created": 1686935002,
+                    "owned_by": "organization-owner",
+                },
+                {
+                    "id": "model-id-2",
+                    "object": "model",
+                    "created": 1686935002,
+                    "owned_by": "openai",
+                },
+            ],
+            "object": "list",
+        }
 
 
 def run(port=1323):
-    app.run(host="0.0.0.0",port=port, debug=True)
+    from rolo import Router
+    from rolo.dispatcher import handler_dispatcher
+    from werkzeug import Request, run_simple
+
+    r = Router(dispatcher=handler_dispatcher())
+    r.add(Api())
 
+    app = Request.application(r.dispatch)
 
-def stop():
-    app.stop()
+    run_simple("0.0.0.0", port, app)
 
 
 if __name__ == "__main__":

openai/setup.cfg

@@ -23,16 +23,17 @@ classifiers =
 zip_safe = False
 packages = find:
 install_requires =
-    flask>=2.0.1
     faker>=8.12.1
-    localstack>=1.0
+    localstack>=3.1
     plux>=1.3
+    rolo>=0.3
 test_requires =
-    openai>=0.10.2
+    openai>=0.10.2,<1.0
     pytest>=6.2.4
 
 [options.extras_require]
 dev =
+    openai>=0.10.2,<1.0
     pytest>=6.2.4
     black==22.3.0
     isort==5.10.1

openai/tests/tests.py renamed to openai/tests/test_api.py

@@ -1,8 +1,9 @@
 import openai
+
 openai.organization = "org-test"
 openai.api_key = "test"
 # openai.api_base = "http://localhost:1323/v1"
-openai.api_base = "http://localhost:4510/v1"
+openai.api_base = "http://localhost:4566/_extension/openai/v1"
 
 
 def test_list_models():
@@ -12,14 +13,15 @@ def test_list_models():
 
 def test_chat_completion():
     completion = openai.ChatCompletion.create(
-      model="gpt-3.5-turbo",
-      messages=[
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": "Hello!"}
-      ]
+        model="gpt-3.5-turbo",
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hello!"},
+        ],
     )
     assert len(completion.choices) > 0
 
+
 def test_transcribe():
     transcript = openai.Audio.transcribe("whisper-1", open("sample.wav", "rb"))
     assert len(transcript.text) > 0
@@ -31,9 +33,5 @@ def test_translate():
 
 
 def test_generate_image():
-    response = openai.Image.create(
-        prompt="a white siamese cat",
-        n=1,
-        size="1024x1024"
-    )
-    assert response['data'][0]['url']
+    response = openai.Image.create(prompt="a white siamese cat", n=1, size="1024x1024")
+    assert response["data"][0]["url"]