Initial commit of openai-proxy

Dockerfile

@@ -0,0 +1,21 @@
+# Use an official Python runtime as a parent image
+FROM python:3.11-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the requirements file into the container at /app
+COPY requirements.txt .
+
+# Install any needed packages specified in requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application's code into the container at /app
+COPY . .
+
+# Back4app will set the PORT environment variable, but we expose 8000 as a default
+EXPOSE 8000
+
+# Run uvicorn when the container launches
+# Use 0.0.0.0 to be accessible from outside the container
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 fangwentong
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,8 @@
 ---
-title: Openai Proxy
-emoji: 🌖
+title: OpenAI Proxy
+emoji: 🚀
 colorFrom: green
-colorTo: green
+colorTo: blue
 sdk: docker
-pinned: false
-license: mit
-short_description: openai-proxy
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+app_port: 8000
+---

log.py

@@ -0,0 +1,52 @@
+import databases
+
+from sqlalchemy import create_engine, Column, Integer, String, BigInteger
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+# database config
+DATABASE_URL = 'sqlite:///./openai_log.db'
+database = databases.Database(DATABASE_URL)
+Base = declarative_base()
+
+
+# database model
+class OpenAILog(Base):
+    """
+    OpenAI API call log
+    """
+    __tablename__ = 'openai_logs'
+
+    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
+    request_url = Column(String)
+    request_method = Column(String)
+    request_time = Column(BigInteger)
+    response_time = Column(BigInteger)
+    status_code = Column(Integer)
+    request_content = Column(String)
+    response_header = Column(String)
+    response_content = Column(String)
+
+    def to_dict(self):
+        return {
+            'id': self.id,
+            'request_url': self.request_url,
+            'request_method': self.request_method,
+            'request_time': self.request_time,
+            'response_time': self.response_time,
+            'status_code': self.status_code,
+            'request_content': self.request_content,
+            'response_header': self.response_header,
+            'response_content': self.response_content,
+        }
+
+
+engine = create_engine(DATABASE_URL)
+Base.metadata.create_all(bind=engine)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+
+async def save_log(log: OpenAILog):
+    async with database.transaction():
+        query = OpenAILog.__table__.insert().values(**log.to_dict())
+        await database.execute(query)

main.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+import json
+import time
+from datetime import datetime
+
+import httpx
+from fastapi import FastAPI, Request, HTTPException
+from starlette.background import BackgroundTask
+
+from log import OpenAILog, save_log
+from utils import PathMatchingTree, OverrideStreamResponse
+
+proxied_hosts = PathMatchingTree({
+    "/": "https://api.openai.com",
+    "/backend-api/conversation": "https://chat.openai.com",
+})
+
+# FastAPI app
+app = FastAPI()
+
+
+async def proxy_openai_api(request: Request):
+    # proxy request to OpenAI API
+    headers = {k: v for k, v in request.headers.items() if
+               k not in {'host', 'content-length', 'x-forwarded-for', 'x-real-ip', 'connection'}}
+    url = f'{proxied_hosts.get_matching(request.url.path)}{request.url.path}'
+
+    start_time = datetime.now().microsecond
+    # create httpx async client
+    client = httpx.AsyncClient()
+
+    request_body = await request.json() if request.method in {'POST', 'PUT'} else None
+
+    log = OpenAILog()
+
+    async def stream_api_response():
+        nonlocal log
+        try:
+            st = client.stream(request.method, url, headers=headers, params=request.query_params, json=request_body)
+            async with st as res:
+                response.status_code = res.status_code
+                response.init_headers({k: v for k, v in res.headers.items() if
+                                       k not in {'content-length', 'content-encoding', 'alt-svc'}})
+
+                content = bytearray()
+                async for chunk in res.aiter_bytes():
+                    yield chunk
+                    content.extend(chunk)
+
+                # gather log data
+                log.request_url = url
+                log.request_method = request.method
+                log.request_time = start_time
+                log.response_time = time.time() - start_time
+                log.status_code = res.status_code
+                log.request_content = (await request.body()).decode('utf-8') if request.method == 'POST' else None
+                log.response_content = content.decode('utf-8')
+                log.response_header = json.dumps([[k, v] for k, v in res.headers.items()])
+
+        except httpx.RequestError as exc:
+            raise HTTPException(status_code=500, detail=f'An error occurred while requesting: {exc}')
+
+    async def update_log():
+        nonlocal log
+        log.response_time = datetime.now().microsecond - start_time
+        await save_log(log)
+
+    response = OverrideStreamResponse(stream_api_response(), background=BackgroundTask(update_log))
+    return response
+
+
+@app.route('/{path:path}', methods=['GET', 'POST', 'PUT', 'DELETE'])
+async def request_handler(request: Request):
+    return await proxy_openai_api(request)
+
+
+if __name__ == '__main__':
+    import uvicorn
+
+    uvicorn.run("main:app", host="127.0.0.1", port=8000, log_level="info", reload=True)

pyproject.toml

@@ -0,0 +1,16 @@
+[tool.poetry]
+name = "openai-gateway"
+version = "0.1.0"
+description = ""
+authors = ["fangwentong <fangwentong2012@gmail.com>"]
+readme = "README.md"
+packages = [{ include = "openai_gateway" }]
+
+[tool.poetry.dependencies]
+python = "^3.10"
+databases = "^0.7.0"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

requirements.txt

@@ -0,0 +1,6 @@
+fastapi==0.74.0
+httpx==0.23.0
+sqlalchemy~=1.4.47
+databases~=0.7.0
+uvicorn==0.15.0
+aiosqlite==0.18.0

utils.py

@@ -0,0 +1,102 @@
+import typing
+from functools import partial
+
+import anyio
+from fastapi.responses import StreamingResponse
+from starlette.types import Send, Scope, Receive
+
+
+class PathMatchingTree:
+    """
+    PathMatchingTree is a data structure that can be used to match a path with a value.
+    It supports exact match, partial match, and wildcard match.
+    For example, if the tree is built with the following config:
+      {
+          "/foo/bar": "value1",
+          "/baz/qux": "value2",
+          "/foo/*": "value3",
+          "/foo/*/bar": "value4"
+      }
+    Then the following path will match the corresponding value:
+      /foo/bar -> value1
+      /baz/qux -> value2
+      /foo/baz -> value3
+      /foo/baz/bar -> value4
+      /foo/baz/bar2 -> value3
+    """
+    child = dict
+    value = None
+
+    def __init__(self, config):
+        self.child = {}
+        self._build_tree(config)
+
+    def _build_tree(self, config):
+        for k, v in config.items():
+            parts = k.split('/')
+            self._add(parts, v)
+
+    def _add(self, parts, value):
+        node = self
+        for part in parts:
+            if part == '':
+                continue
+            if part not in node.child:
+                node.child[part] = PathMatchingTree(dict())
+            node = node.child[part]
+        node.value = value
+
+    def get_matching(self, path):
+        parts = path.split('/')
+        matched = self
+        for part in parts:
+            if part == '':
+                continue
+            if part in matched.child:
+                matched = matched.child[part]
+            elif '*' in matched.child:
+                matched = matched.child['*']
+            else:
+                break
+        return matched.value
+
+
+class OverrideStreamResponse(StreamingResponse):
+    """
+    Override StreamingResponse to support lazy send response status_code and response headers
+    """
+
+    async def stream_response(self, send: Send) -> None:
+        first_chunk = True
+        async for chunk in self.body_iterator:
+            if first_chunk:
+                await self.send_request_header(send)
+                first_chunk = False
+            if not isinstance(chunk, bytes):
+                chunk = chunk.encode(self.charset)
+            await send({'type': 'http.response.body', 'body': chunk, 'more_body': True})
+
+        if first_chunk:
+            await self.send_request_header(send)
+        await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
+
+    async def send_request_header(self, send: Send) -> None:
+        await send(
+            {
+                'type': 'http.response.start',
+                'status': self.status_code,
+                'headers': self.raw_headers,
+            }
+        )
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        async with anyio.create_task_group() as task_group:
+            async def wrap(func: typing.Callable[[], typing.Coroutine]) -> None:
+                await func()
+                await task_group.cancel_scope.cancel()
+
+            task_group.start_soon(wrap, partial(self.stream_response, send))
+            await wrap(partial(self.listen_for_disconnect, receive))
+
+        if self.background is not None:
+            await self.background()

utils_test.py

@@ -0,0 +1,62 @@
+import unittest
+from utils import PathMatchingTree
+
+
+class TestPathMatchingTree(unittest.TestCase):
+
+    def test_get_matching_exact_match(self):
+        config = {
+            "foo/bar": "value1",
+            "baz/qux": "value2"
+        }
+        pmt = PathMatchingTree(config)
+        result = pmt.get_matching("foo/bar")
+        self.assertEqual(result, "value1")
+
+    def test_get_matching_partial_match(self):
+        config = {
+            "foo/bar": "value1",
+            "baz/qux": "value2"
+        }
+        pmt = PathMatchingTree(config)
+        self.assertIsNone(pmt.get_matching("foo"))
+
+    def test_get_matching_wildcard_match(self):
+        config = {
+            "/foo/*": "value1",
+            "/baz/qux": "value2"
+        }
+        pmt = PathMatchingTree(config)
+        self.assertEqual(pmt.get_matching("foo/bar"), "value1")
+
+    def test_get_matching_multiple_wildcard_match(self):
+        config = {
+            "/foo/*": "value1",
+            "/foo/*/bar": "value2"
+        }
+        pmt = PathMatchingTree(config)
+        self.assertIsNone(pmt.get_matching("/foo"))
+        self.assertEqual(pmt.get_matching("/foo/baz"), "value1")
+        self.assertEqual(pmt.get_matching("/foo/baz/bar2"), "value1")
+        self.assertEqual(pmt.get_matching("/foo/baz/bar"), "value2")
+
+    def test_get_matching_no_match(self):
+        config = {
+            "/foo/bar": "value1",
+            "/baz/qux": "value2"
+        }
+        pmt = PathMatchingTree(config)
+        self.assertIsNone(pmt.get_matching("/foo"))
+        self.assertIsNone(pmt.get_matching("/baz"))
+
+    def test_get_matching_empty_string_match(self):
+        config = {
+            "/": "value1"
+        }
+        pmt = PathMatchingTree(config)
+        self.assertEqual(pmt.get_matching("/"), "value1")
+        self.assertEqual(pmt.get_matching("/test"), "value1")
+
+
+if __name__ == "__main__":
+    unittest.main()