├── .DS_Store ├── .gitignore ├── .swiftpm └── xcode │ ├── package.xcworkspace │ └── contents.xcworkspacedata │ └── xcshareddata │ └── xcschemes │ └── SwiftOpenAIProxy.xcscheme ├── Dockerfile ├── Images └── EditScheme.png ├── LICENSE ├── Package.resolved ├── Package.swift ├── README.md ├── Sources └── App │ ├── .DS_Store │ ├── App.swift │ ├── Application+configure.swift │ ├── Controllers │ └── AppStoreController.swift │ ├── Extensions │ └── Status.swift │ ├── Middleware │ ├── APIKeyMiddleware.swift │ ├── AppStoreAuthenticator.swift │ ├── JWTAuthenticator.swift │ ├── MessageRouterMiddleware.swift │ ├── ProxyServerMiddleware.swift │ ├── RateLimiterMiddleware.swift │ └── StreamingResponseDelegate.swift │ ├── Models │ ├── Chat.swift │ ├── Config.swift │ ├── LLMModel.swift │ ├── LLMProvider.swift │ ├── User.swift │ ├── UserMigration.swift │ ├── UserRequest.swift │ └── UserRequestMigration.swift │ ├── Providers.swift │ └── Resources │ ├── .DS_Store │ ├── AppleComputerRootCertificate.cer │ ├── AppleIncRootCertificate.cer │ ├── AppleRootCA-G2.cer │ ├── AppleRootCA-G3.cer │ └── Providers.plist └── Tests └── AppTests └── AppTests.swift /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PicoMLX/PicoAIProxy/5ecaaf24675bf0b9a51aa14323b4b31138cd8f86/.DS_Store -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Xcode 2 | # 3 | # gitignore contributors: remember to update Global/Xcode.gitignore, Objective-C.gitignore & Swift.gitignore 4 | 5 | ## User settings 6 | xcuserdata/ 7 | 8 | ## compatibility with Xcode 8 and earlier (ignoring not required starting Xcode 9) 9 | *.xcscmblueprint 10 | *.xccheckout 11 | 12 | ## compatibility with Xcode 3 and earlier (ignoring not required starting Xcode 4) 13 | build/ 14 | DerivedData/ 15 | *.moved-aside 16 | *.pbxuser 17 | !default.pbxuser 18 | *.mode1v3 19 | !default.mode1v3 20 | *.mode2v3 21 | !default.mode2v3 22 | *.perspectivev3 23 | !default.perspectivev3 24 | 25 | ## Obj-C/Swift specific 26 | *.hmap 27 | 28 | ## App packaging 29 | *.ipa 30 | *.dSYM.zip 31 | *.dSYM 32 | 33 | ## Playgrounds 34 | timeline.xctimeline 35 | playground.xcworkspace 36 | 37 | # Swift Package Manager 38 | # 39 | # Add this line if you want to avoid checking in source code from Swift Package Manager dependencies. 40 | Packages/ 41 | Package.pins 42 | Package.resolved 43 | *.xcodeproj 44 | 45 | # Xcode automatically generates this directory with a .xcworkspacedata file and xcuserdata 46 | # hence it is not needed unless you have added a package configuration file to your project 47 | .swiftpm 48 | 49 | .build/ 50 | 51 | # CocoaPods 52 | # 53 | # We recommend against adding the Pods directory to your .gitignore. However 54 | # you should judge for yourself, the pros and cons are mentioned at: 55 | # https://guides.cocoapods.org/using/using-cocoapods.html#should-i-check-the-pods-directory-into-source-control 56 | 57 | Pods/ 58 | 59 | # Add this line if you want to avoid checking in source code from the Xcode workspace 60 | *.xcworkspace 61 | 62 | # Carthage 63 | # 64 | # Add this line if you want to avoid checking in source code from Carthage dependencies. 65 | Carthage/Checkouts 66 | 67 | Carthage/Build/ 68 | 69 | # Accio dependency management 70 | Dependencies/ 71 | .accio/ 72 | 73 | fastlane 74 | 75 | # It is recommended to not store the screenshots in the git repo. 76 | # Instead, use fastlane to re-generate the screenshots whenever they are needed. 77 | # For more information about the recommended setup visit: 78 | # https://docs.fastlane.tools/best-practices/source-control/#source-control 79 | 80 | fastlane/report.xml 81 | fastlane/Preview.html 82 | fastlane/screenshots/**/*.png 83 | fastlane/test_output 84 | 85 | # Code Injection 86 | # 87 | # After new code Injection tools there's a generated folder /iOSInjectionProject 88 | # https://github.com/johnno1962/injectionforxcode 89 | 90 | iOSInjectionProject/ 91 | 92 | .DS_Store 93 | -------------------------------------------------------------------------------- /.swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata: -------------------------------------------------------------------------------- 1 | 2 | 4 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /.swiftpm/xcode/xcshareddata/xcschemes/SwiftOpenAIProxy.xcscheme: -------------------------------------------------------------------------------- 1 | 2 | 5 | 8 | 9 | 15 | 21 | 22 | 23 | 29 | 35 | 36 | 37 | 38 | 39 | 45 | 46 | 48 | 54 | 55 | 56 | 57 | 58 | 68 | 70 | 76 | 77 | 78 | 79 | 82 | 83 | 84 | 85 | 89 | 90 | 94 | 95 | 99 | 100 | 104 | 105 | 109 | 110 | 114 | 115 | 119 | 120 | 124 | 125 | 129 | 130 | 134 | 135 | 139 | 140 | 144 | 145 | 149 | 150 | 154 | 155 | 159 | 160 | 164 | 165 | 169 | 170 | 174 | 175 | 179 | 180 | 181 | 182 | 188 | 190 | 196 | 197 | 198 | 199 | 201 | 202 | 205 | 206 | 207 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # ================================ 2 | # Build image 3 | # ================================ 4 | FROM swift:5.9 as build 5 | 6 | # Install OS updates 7 | RUN export DEBIAN_FRONTEND=noninteractive DEBCONF_NONINTERACTIVE_SEEN=true \ 8 | && apt-get -q update \ 9 | && apt-get -q dist-upgrade -y \ 10 | && rm -rf /var/lib/apt/lists/* 11 | 12 | # Set up a build area 13 | WORKDIR /build 14 | 15 | # First just resolve dependencies. 16 | # This creates a cached layer that can be reused 17 | # as long as your Package.swift/Package.resolved 18 | # files do not change. 19 | COPY ./Package.* ./ 20 | RUN swift package resolve 21 | 22 | # Copy entire repo into container 23 | COPY . . 24 | 25 | # Build everything, with optimizations and test discovery 26 | RUN swift build --enable-test-discovery -c release 27 | 28 | # Switch to the staging area 29 | WORKDIR /staging 30 | 31 | # Copy main executable to staging area 32 | RUN cp "$(swift build --package-path /build -c release --show-bin-path)/App" ./ 33 | 34 | # Copy any resources from the public directory and views directory if the directories exist 35 | # Ensure that by default, neither the directory nor any of its contents are writable. 36 | RUN [ -d /build/Public ] && { mv /build/Public ./Public && chmod -R a-w ./Public; } || true 37 | RUN [ -d /build/Sources/App/Resources ] && { mv /build/Sources/App/Resources ./Resources && chmod -R a-w ./Resources; } || true 38 | 39 | # ================================ 40 | # Run image 41 | # ================================ 42 | FROM swift:5.9-slim 43 | 44 | # Make sure all system packages are up to date. 45 | RUN export DEBIAN_FRONTEND=noninteractive DEBCONF_NONINTERACTIVE_SEEN=true && \ 46 | apt-get -q update && apt-get -q dist-upgrade -y && rm -r /var/lib/apt/lists/* 47 | 48 | # Create a hummingbird user and group with /app as its home directory 49 | RUN useradd --user-group --create-home --system --skel /dev/null --home-dir /app hummingbird 50 | 51 | # Switch to the new home directory 52 | WORKDIR /app 53 | 54 | # Copy built executable and any staged resources from builder 55 | COPY --from=build --chown=hummingbird:hummingbird /staging /app 56 | 57 | # Ensure all further commands run as the vapor user 58 | USER hummingbird:hummingbird 59 | 60 | # Let Docker bind to port 8080 61 | EXPOSE 8080 62 | 63 | # Start the Vapor service when the image is run, default to listening on 8080 in production environment 64 | ENTRYPOINT ["./App"] 65 | CMD ["--hostname", "0.0.0.0", "--port", "8080"] 66 | -------------------------------------------------------------------------------- /Images/EditScheme.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PicoMLX/PicoAIProxy/5ecaaf24675bf0b9a51aa14323b4b31138cd8f86/Images/EditScheme.png -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Ronald Mannak 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Package.resolved: -------------------------------------------------------------------------------- 1 | { 2 | "pins" : [ 3 | { 4 | "identity" : "app-store-server-library-swift", 5 | "kind" : "remoteSourceControl", 6 | "location" : "https://github.com/apple/app-store-server-library-swift.git", 7 | "state" : { 8 | "revision" : "ec82178a1759b06a8f09becd1cd626318e00a7a2", 9 | "version" : "1.1.0" 10 | } 11 | }, 12 | { 13 | "identity" : "async-http-client", 14 | "kind" : "remoteSourceControl", 15 | "location" : "https://github.com/swift-server/async-http-client.git", 16 | "state" : { 17 | "revision" : "0ae99db85b2b9d1e79b362bd31fd1ffe492f7c47", 18 | "version" : "1.21.2" 19 | } 20 | }, 21 | { 22 | "identity" : "async-kit", 23 | "kind" : "remoteSourceControl", 24 | "location" : "https://github.com/vapor/async-kit.git", 25 | "state" : { 26 | "revision" : "7ece208cd401687641c88367a00e3ea2b04311f1", 27 | "version" : "1.19.0" 28 | } 29 | }, 30 | { 31 | "identity" : "fluent-kit", 32 | "kind" : "remoteSourceControl", 33 | "location" : "https://github.com/vapor/fluent-kit.git", 34 | "state" : { 35 | "revision" : "d69efce21242ad4dba6935cc1b8d5637281604d5", 36 | "version" : "1.48.5" 37 | } 38 | }, 39 | { 40 | "identity" : "fluent-sqlite-driver", 41 | "kind" : "remoteSourceControl", 42 | "location" : "https://github.com/vapor/fluent-sqlite-driver.git", 43 | "state" : { 44 | "revision" : "40303a20bc39c270c8e50339ada30f9750e2a681", 45 | "version" : "4.7.3" 46 | } 47 | }, 48 | { 49 | "identity" : "hummingbird", 50 | "kind" : "remoteSourceControl", 51 | "location" : "https://github.com/hummingbird-project/hummingbird.git", 52 | "state" : { 53 | "revision" : "c33220bf0d229eb127cb4a998c99c0d7554fd641", 54 | "version" : "1.12.2" 55 | } 56 | }, 57 | { 58 | "identity" : "hummingbird-auth", 59 | "kind" : "remoteSourceControl", 60 | "location" : "https://github.com/hummingbird-project/hummingbird-auth", 61 | "state" : { 62 | "revision" : "754cae807dcdab057e63e722c9572068edc7597f", 63 | "version" : "1.3.0" 64 | } 65 | }, 66 | { 67 | "identity" : "hummingbird-core", 68 | "kind" : "remoteSourceControl", 69 | "location" : "https://github.com/hummingbird-project/hummingbird-core.git", 70 | "state" : { 71 | "revision" : "2a16f079f78b2a746ad57b3931c366cd14961927", 72 | "version" : "1.6.1" 73 | } 74 | }, 75 | { 76 | "identity" : "hummingbird-fluent", 77 | "kind" : "remoteSourceControl", 78 | "location" : "https://github.com/hummingbird-project/hummingbird-fluent.git", 79 | "state" : { 80 | "revision" : "f99dc930c8c1b1e49d2ba2728efa76424293e0a8", 81 | "version" : "1.1.0" 82 | } 83 | }, 84 | { 85 | "identity" : "jwt-kit", 86 | "kind" : "remoteSourceControl", 87 | "location" : "https://github.com/vapor/jwt-kit.git", 88 | "state" : { 89 | "revision" : "c2595b9ad7f512d7f334830b4df1fed6e917946a", 90 | "version" : "4.13.4" 91 | } 92 | }, 93 | { 94 | "identity" : "sql-kit", 95 | "kind" : "remoteSourceControl", 96 | "location" : "https://github.com/vapor/sql-kit.git", 97 | "state" : { 98 | "revision" : "f697d3289c628acd241e3b2c7d3ff068adcc52d1", 99 | "version" : "3.31.1" 100 | } 101 | }, 102 | { 103 | "identity" : "sqlite-kit", 104 | "kind" : "remoteSourceControl", 105 | "location" : "https://github.com/vapor/sqlite-kit.git", 106 | "state" : { 107 | "revision" : "f35a863ecc2da5d563b836a9a696b148b0f4169f", 108 | "version" : "4.5.2" 109 | } 110 | }, 111 | { 112 | "identity" : "sqlite-nio", 113 | "kind" : "remoteSourceControl", 114 | "location" : "https://github.com/vapor/sqlite-nio.git", 115 | "state" : { 116 | "revision" : "1b03dafcd8b86047650925325a2bd4d20f6205fd", 117 | "version" : "1.10.1" 118 | } 119 | }, 120 | { 121 | "identity" : "swift-algorithms", 122 | "kind" : "remoteSourceControl", 123 | "location" : "https://github.com/apple/swift-algorithms", 124 | "state" : { 125 | "revision" : "f6919dfc309e7f1b56224378b11e28bab5bccc42", 126 | "version" : "1.2.0" 127 | } 128 | }, 129 | { 130 | "identity" : "swift-argument-parser", 131 | "kind" : "remoteSourceControl", 132 | "location" : "https://github.com/apple/swift-argument-parser.git", 133 | "state" : { 134 | "revision" : "0fbc8848e389af3bb55c182bc19ca9d5dc2f255b", 135 | "version" : "1.4.0" 136 | } 137 | }, 138 | { 139 | "identity" : "swift-asn1", 140 | "kind" : "remoteSourceControl", 141 | "location" : "https://github.com/apple/swift-asn1.git", 142 | "state" : { 143 | "revision" : "c7e239b5c1492ffc3ebd7fbcc7a92548ce4e78f0", 144 | "version" : "1.1.0" 145 | } 146 | }, 147 | { 148 | "identity" : "swift-atomics", 149 | "kind" : "remoteSourceControl", 150 | "location" : "https://github.com/apple/swift-atomics.git", 151 | "state" : { 152 | "revision" : "cd142fd2f64be2100422d658e7411e39489da985", 153 | "version" : "1.2.0" 154 | } 155 | }, 156 | { 157 | "identity" : "swift-backtrace", 158 | "kind" : "remoteSourceControl", 159 | "location" : "https://github.com/swift-server/swift-backtrace.git", 160 | "state" : { 161 | "revision" : "7277ee0e0378d27465b304c668d8e914192d986f", 162 | "version" : "1.3.5" 163 | } 164 | }, 165 | { 166 | "identity" : "swift-certificates", 167 | "kind" : "remoteSourceControl", 168 | "location" : "https://github.com/apple/swift-certificates.git", 169 | "state" : { 170 | "revision" : "4688f242811d21a9c7a8ad669b3bc5b336759929", 171 | "version" : "1.4.0" 172 | } 173 | }, 174 | { 175 | "identity" : "swift-collections", 176 | "kind" : "remoteSourceControl", 177 | "location" : "https://github.com/apple/swift-collections.git", 178 | "state" : { 179 | "revision" : "ee97538f5b81ae89698fd95938896dec5217b148", 180 | "version" : "1.1.1" 181 | } 182 | }, 183 | { 184 | "identity" : "swift-crypto", 185 | "kind" : "remoteSourceControl", 186 | "location" : "https://github.com/apple/swift-crypto.git", 187 | "state" : { 188 | "revision" : "bc1c29221f6dfeb0ebbfbc98eb95cd3d4967868e", 189 | "version" : "3.4.0" 190 | } 191 | }, 192 | { 193 | "identity" : "swift-distributed-tracing", 194 | "kind" : "remoteSourceControl", 195 | "location" : "https://github.com/apple/swift-distributed-tracing.git", 196 | "state" : { 197 | "revision" : "11c756c5c4d7de0eeed8595695cadd7fa107aa19", 198 | "version" : "1.1.1" 199 | } 200 | }, 201 | { 202 | "identity" : "swift-extras-base64", 203 | "kind" : "remoteSourceControl", 204 | "location" : "https://github.com/swift-extras/swift-extras-base64.git", 205 | "state" : { 206 | "revision" : "97237cf1bc1feebaeb0febec91c1e1b9e4d839b3", 207 | "version" : "0.7.0" 208 | } 209 | }, 210 | { 211 | "identity" : "swift-http-types", 212 | "kind" : "remoteSourceControl", 213 | "location" : "https://github.com/apple/swift-http-types", 214 | "state" : { 215 | "revision" : "1ddbea1ee34354a6a2532c60f98501c35ae8edfa", 216 | "version" : "1.2.0" 217 | } 218 | }, 219 | { 220 | "identity" : "swift-log", 221 | "kind" : "remoteSourceControl", 222 | "location" : "https://github.com/apple/swift-log.git", 223 | "state" : { 224 | "revision" : "9cb486020ebf03bfa5b5df985387a14a98744537", 225 | "version" : "1.6.1" 226 | } 227 | }, 228 | { 229 | "identity" : "swift-metrics", 230 | "kind" : "remoteSourceControl", 231 | "location" : "https://github.com/apple/swift-metrics.git", 232 | "state" : { 233 | "revision" : "e0165b53d49b413dd987526b641e05e246782685", 234 | "version" : "2.5.0" 235 | } 236 | }, 237 | { 238 | "identity" : "swift-nio", 239 | "kind" : "remoteSourceControl", 240 | "location" : "https://github.com/apple/swift-nio.git", 241 | "state" : { 242 | "revision" : "e5a216ba89deba84356bad9d4c2eab99071c745b", 243 | "version" : "2.67.0" 244 | } 245 | }, 246 | { 247 | "identity" : "swift-nio-extras", 248 | "kind" : "remoteSourceControl", 249 | "location" : "https://github.com/apple/swift-nio-extras.git", 250 | "state" : { 251 | "revision" : "a3b640d7dc567225db7c94386a6e71aded1bfa63", 252 | "version" : "1.22.0" 253 | } 254 | }, 255 | { 256 | "identity" : "swift-nio-http2", 257 | "kind" : "remoteSourceControl", 258 | "location" : "https://github.com/apple/swift-nio-http2.git", 259 | "state" : { 260 | "revision" : "8d8eb609929aee75336a0a3d2417280786265868", 261 | "version" : "1.32.0" 262 | } 263 | }, 264 | { 265 | "identity" : "swift-nio-ssl", 266 | "kind" : "remoteSourceControl", 267 | "location" : "https://github.com/apple/swift-nio-ssl.git", 268 | "state" : { 269 | "revision" : "2b09805797f21c380f7dc9bedaab3157c5508efb", 270 | "version" : "2.27.0" 271 | } 272 | }, 273 | { 274 | "identity" : "swift-nio-transport-services", 275 | "kind" : "remoteSourceControl", 276 | "location" : "https://github.com/apple/swift-nio-transport-services.git", 277 | "state" : { 278 | "revision" : "38ac8221dd20674682148d6451367f89c2652980", 279 | "version" : "1.21.0" 280 | } 281 | }, 282 | { 283 | "identity" : "swift-numerics", 284 | "kind" : "remoteSourceControl", 285 | "location" : "https://github.com/apple/swift-numerics.git", 286 | "state" : { 287 | "revision" : "0a5bc04095a675662cf24757cc0640aa2204253b", 288 | "version" : "1.0.2" 289 | } 290 | }, 291 | { 292 | "identity" : "swift-service-context", 293 | "kind" : "remoteSourceControl", 294 | "location" : "https://github.com/apple/swift-service-context.git", 295 | "state" : { 296 | "revision" : "ce0141c8f123132dbd02fd45fea448018762df1b", 297 | "version" : "1.0.0" 298 | } 299 | }, 300 | { 301 | "identity" : "swift-service-lifecycle", 302 | "kind" : "remoteSourceControl", 303 | "location" : "https://github.com/swift-server/swift-service-lifecycle.git", 304 | "state" : { 305 | "revision" : "22363fed316cd9942b56bcd1a1df8875df79b794", 306 | "version" : "1.0.0-alpha.11" 307 | } 308 | }, 309 | { 310 | "identity" : "swift-system", 311 | "kind" : "remoteSourceControl", 312 | "location" : "https://github.com/apple/swift-system.git", 313 | "state" : { 314 | "revision" : "f9266c85189c2751589a50ea5aec72799797e471", 315 | "version" : "1.3.0" 316 | } 317 | } 318 | ], 319 | "version" : 2 320 | } 321 | -------------------------------------------------------------------------------- /Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version:5.9 2 | // The swift-tools-version declares the minimum version of Swift required to build this package. 3 | 4 | import PackageDescription 5 | 6 | let package = Package( 7 | name: "SwiftOpenAIProxy", 8 | platforms: [.macOS(.v12)], 9 | products: [ 10 | .executable(name: "App", targets: ["App"]), 11 | ], 12 | dependencies: [ 13 | .package(url: "https://github.com/hummingbird-project/hummingbird.git", from: "1.12.1"), 14 | .package(url: "https://github.com/hummingbird-project/hummingbird-auth", from: "1.3.0"), 15 | .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.3.0"), 16 | .package(url: "https://github.com/swift-server/async-http-client.git", from: "1.19.0"), 17 | .package(url: "https://github.com/apple/app-store-server-library-swift.git", from: "1.0.1"), 18 | .package(url: "https://github.com/vapor/jwt-kit.git", from: "4.0.0"), 19 | .package(url: "https://github.com/hummingbird-project/hummingbird-fluent.git", from: "1.0.0"), 20 | .package(url: "https://github.com/vapor/fluent-kit.git", from: "1.16.0"), 21 | .package(url: "https://github.com/vapor/fluent-sqlite-driver.git", from: "4.0.0"), 22 | ], 23 | targets: [ 24 | .executableTarget(name: "App", 25 | dependencies: [ 26 | .product(name: "Hummingbird", package: "hummingbird"), 27 | .product(name: "HummingbirdAuth", package: "hummingbird-auth"), 28 | .product(name: "ArgumentParser", package: "swift-argument-parser"), 29 | .product(name: "AsyncHTTPClient", package: "async-http-client"), 30 | .product(name: "AppStoreServerLibrary", package: "app-store-server-library-swift"), 31 | .product(name: "JWTKit", package: "jwt-kit"), 32 | .product(name: "HummingbirdFluent", package: "hummingbird-fluent"), 33 | .product(name: "FluentKit", package: "fluent-kit"), 34 | .product(name: "FluentSQLiteDriver", package: "fluent-sqlite-driver"), 35 | ], 36 | resources: [ 37 | // Copy Apple root certificates 38 | .process("Resources/") 39 | ], 40 | swiftSettings: [ 41 | // Enable better optimizations when building in Release configuration. Despite the use of 42 | // the `.unsafeFlags` construct required by SwiftPM, this flag is recommended for Release 43 | // builds. See for details. 44 | .unsafeFlags(["-cross-module-optimization"], .when(configuration: .release)) 45 | ] 46 | ), 47 | .testTarget(name: "AppTests", 48 | dependencies: [ 49 | .byName(name: "App"), 50 | .product(name: "HummingbirdXCT", package: "hummingbird") 51 | ] 52 | ) 53 | ] 54 | ) 55 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pico AI Proxy 2 | 3 | ### Introduction 4 | 5 | Pico AI Proxy is a reverse proxy created specifically for iOS, macOS, iPadOS, and VisionOS developers. 6 | 7 | Pico AI Proxy was previously called SwiftOpenAIProxy. 8 | 9 | ### Highlights 10 | 11 | - Prevents hackers from stealing your API keys 12 | - Makes sure every user has a in-app subscription by validating App Store store receipts using Apple's [App Store Server Library](https://github.com/apple/app-store-server-library-swift) 13 | - Accepts the standard OpenAI Chat API your chat app already uses, so no changes need to be made to your client app 14 | - Supports multiple LLM providers such as OpenAI and Anthropic (more to come, stay tuned). PicoProxy automatically converts OpenAI Chat API calls to the different providers 15 | - One click install on [Railway](https://railway.app) and can be installed manually on many other hosting providers. [See How to Deploy](#How-to-deploy) 16 | 17 | PicoProxy is written in server-side Swift and uses [HummingBird](https://github.com/hummingbird-project/hummingbird) for its HTTP-server. 18 | 19 | ### Background 20 | 21 | In December 2023, I faced a [significant hack](https://youtu.be/_ueiYhLwwBc?si=8UC_7VZOrhgcXoKV) when my OpenAI key was compromised. They quickly used up my entire $2,500 monthly limit, resulting in an unexpected bill from OpenAI. This incident also forced me to take my app, Pico, offline, causing me to miss the lucrative Christmas sales period. 22 | 23 | As a response to this incident, I developed Pico AI Server, the first OpenAI proxy created in Server-side Swift. This tool is especially convenient for Swift developers, as it allows easy customization to meet their specific requirements. 24 | 25 | Pico AI Proxy is designed to be compatible with any existing OpenAI library. It works seamlessly with libraries such as [CleverBird](https://github.com/btfranklin/CleverBird), [OpenAISwift](https://github.com/adamrushy/OpenAISwift), [OpenAI-Kit](https://github.com/dylanshine/openai-kit), [MacPaw OpenAI](https://github.com/MacPaw/OpenAI), and can also integrate with your custom code. 26 | 27 | ### Key features 28 | 29 | - Pico AI Proxy uses Apple's [App Store Server Library](https://github.com/apple/app-store-server-library-swift) for receipt validation 30 | - Once the receipt is validated, Pico AI Proxy issues a JWT token the client can use for subsequent calls 31 | - Pico AI Proxy is API agnostic and forwards any request to https://api.openai.com. 32 | - The forwarding endpoint is customizable, allowing redirection to various non-OpenAI API endpoints 33 | - Optionally forward calls with a valid OpenAI key and org without validation 34 | - Pico AI Server can optionally track individual users through App Account IDs. This requires the client app to send a unique UUID to the [purchase](https://developer.apple.com/documentation/storekit/product/3791971-purchase) method. 35 | 36 | ### Supported APIs 37 | 38 | | API | Chat async | Chat streaming | Embeddings | Audio | Images | 39 | | --- | --- | --- | --- | --- | --- | 40 | | [OpenAI](https://platform.openai.com/docs/models) | ✅ | ✅ | ✅ | ✅ | ✅ | 41 | | [Anthropic](https://docs.anthropic.com/claude/docs/) | ❌ | ✅ | ❌ | ❌ | ❌ | 42 | 43 | ### Supported Models and endpoints 44 | 45 | OpenAI: 46 | 47 | - Pico AI Proxy supports all OpenAI models and endpoints: `chat`, `audio`, `embeddings`, `fine-tune`, `image` 48 | 49 | Anthropic: 50 | 51 | - API version `2023-06-01` 52 | - `claude-3-opus-20240229` 53 | - `claude-3-sonnet-20240229` 54 | - `claude-3-haiku-20240307` 55 | 56 | 57 | ### What's implemented 58 | - [x] Reverse proxy server forwarding calls to OpenAI (or any other endpoint) 59 | - [x] Authenticates using App Store receipt validation 60 | - [x] Rate limiter and black list 61 | - [x] Automatically translate and forward traffic to other AI APIs based on model setting in API call 62 | - [ ] [App Attestation](https://developer.apple.com/documentation/devicecheck/preparing_to_use_the_app_attest_service) is on hold, as macOS doesn't support app attestation 63 | - [ ] Account management 64 | 65 | ## Requirements 66 | 67 | - [Apple Developer Account](https://developer.apple.com) 68 | - [OpenAI API key](https://openai.com/blog/openai-api) 69 | - [Anthropic API key](https://docs.anthropic.com/claude/reference/getting-started-with-the-api) 70 | 71 | ## How to Set Up Pico AI Proxy 72 | 73 | To set up Pico AI Proxy, you need: 74 | - Your OpenAI API key and organization 75 | - A JWT private key, which can be generated in the terminal 76 | - Your app bundle Id, Apple app Id and team Id 77 | - App Store Server API key, Issuer Id, and Key Id 78 | - Apple root certificates, which are included in the repository but should be updated if Apple updates their certificates 79 | 80 | #### OpenAI API key and organization 81 | Generate an OpenAI API key at [OpenAI](https://platform.openai.com) 82 | Optionally generate an Anthropic Claude API key at [Anthropic](https://docs.anthropic.com/claude/reference/getting-started-with-the-api) 83 | 84 | #### JWT Private key 85 | Create a new JWT private key in macOS Terminal using `openssl rand -base64 32` 86 | 87 | Note: This JWT token is used to authenticate your client app. It is a different JWT token the App Store Server Library uses to communicate with the Apple App Store API. 88 | 89 | #### App Ids 90 | Find your App bundle Id, Apple app Id, and team Id on https://appstoreconnect.apple.com/apps Under **App Information** in the **General** section, you will find these details. 91 | 92 | Team Id is a 10-character string and can be found in https://developer.apple.com/account under **Membership Details**. 93 | 94 | #### App Store Server API key 95 | Generate the key under the **Users and Access** tab in App Store Connect, specifically under **In-app Purchase** [here](https://appstoreconnect.apple.com/access/api/subs). You will also find the Issuer Id and Key Id on the same page. 96 | 97 | See for more details [Creating API Keys for App Store Connect API](https://developer.apple.com/documentation/appstoreconnectapi/creating_api_keys_for_app_store_connect_api) 98 | 99 | ## Run Pico AI Proxy from Xcode 100 | 101 | To run Pico AI Proxy from Xcode, set the environment variables and arguments listed below to the information listed in How to Set Up Pico AI Proxy 102 | 103 | Both environment variables and arguments can be edited in Xcode using the Target -> Edit scheme. 104 | 105 | ![Xcode screenshot of edit scheme menu](Images/EditScheme.png) 106 | 107 | ### Arguments passed on launch 108 | 109 | | Argument | Default value | Default in scheme | 110 | | --- | --- | --- | 111 | | --hostname | 0.0.0.0 | | 112 | | --port | 8080 | 8080 | 113 | | --target | https://api.openai.com | | 114 | 115 | When launched from Xcode, Pico AI Proxy is accessible at http://localhost:8080. When deployed on Railway, Pico AI Proxy will default to port 443 (https). 116 | 117 | All traffic will be forwarded to `target`. The 'target' can be modified to direct traffic to any API, regardless of whether it conforms to the OpenAI API. , as long as your client application is compatible. 118 | The target is the site where all traffic is forwarded to. You can change the target to any API, even if the API doesn't conform OpenAI (so long as your client app does). 119 | 120 | ## Environment variables 121 | 122 | ### LLM providers environment variables. 123 | | Variable | Description | reference | 124 | | --- | --- | --- | 125 | | OpenAI-APIKey | OpenAI API key (sk-...) | https://platform.openai.com | 126 | | OpenAI-Organization | OpenAI org identifier (org-...) | https://platform.openai.com | 127 | | Anthropic-APIKey | Anthropic API key (sk-ant-api3-...) | https://docs.anthropic.com/claude/docs/ | 128 | | ~~allowKeyPassthrough~~ | if 1, requests with a valid OpenAI key and org in the header will be forwarded to OpenAI without modifications (deprecated) | 129 | 130 | ### App Store Connect environment variables 131 | | Variable | Description | reference | 132 | | --- | --- | --- | 133 | | appTeamId | Apple Team ID | https://appstoreconnect.apple.com/ | 134 | | appBundleId | E.g. com.example.myapp | https://appstoreconnect.apple.com/ | 135 | | appAppleId | Apple Id under App Information -> General Information | https://appstoreconnect.apple.com/ | 136 | 137 | ### App Store Server API environment variables 138 | | Variable | Description | reference | 139 | | --- | --- | --- | 140 | | IAPPrivateKey | IAP private key | https://appstoreconnect.apple.com/access/api/subs | 141 | | IAPIssuerId | IAP Issuer Id | https://appstoreconnect.apple.com/access/api/subs | 142 | | IAPKeyId | IAP Key Id | https://appstoreconnect.apple.com/access/api/subs | 143 | 144 | The `IAPPrivateKey` in Pico AI Proxy is formatted in PKCS #8, which is a multi-line format. The format begins with `-----BEGIN PRIVATE KEY-----` and ends with `-----END PRIVATE KEY-----`. Between these markers, the key comprises four lines of base64-encoded data. However, while Xcode supports environment variables with newlines, many hosting services, such as [Railway](https://railway.app), do not. 145 | 146 | To ensure compatibility across different environments, Pico AI Proxy requires the private key to be condensed into a single line. This is achieved by replacing all newline characters with `\\n` (double backslash followed by `n`). 147 | 148 | A correctly formatted `IAPPrivateKey` for Pico AI Proxy should appear as a single line: `-----BEGIN PRIVATE KEY-----\\n\\n\\n\\n\\n-----END PRIVATE KEY-----`, where ``, ``, ``, and `` represent the base64-encoded data of the key. 149 | 150 | ### JWT environment variables 151 | | Variable | Description | reference | 152 | | --- | --- | --- | 153 | | JWTPrivateKey | | https://jwt.io/introduction | 154 | 155 | ### Rate limiter environment variables 156 | | Variable | Default value | Description | 157 | | --- | --- | --- | 158 | | enableRateLimiter | 0 | Set to 1 to activate the rate limiter | 159 | | userMinuteRateLimit | 15 | Max queries per minute per registered user 160 | | userHourlyRateLimit | 50 | Max queries per hour per registered user 161 | | userPermanentBlock | 50 | Blocked request threshold for permanent user ban 162 | | anonMinuteRateLimit | 60 | Combined max queries per minute for all anonymous users 163 | | anonHourlyRateLimit | 200 | Combined max queries per hour for all anonymous users 164 | | anonPermanentBlock | 50 | Blocked request threshold for banning all anonymous users 165 | 166 | #### Guidelines and behavior 167 | 168 | By default, the rate limiter is off. To activate, set `enableRateLimiter` to 1. 169 | 170 | The rate limiter counts requests and doesn't distinguish between different models or LLM providers. It's primarily a safeguard against abusive traffic. 171 | 172 | Users are identified by their app account tokens from the StoreKit 2 Transaction.purchase() call. Unidentified users are considered anonymous. For apps where all users are identified, consider removing the anonymous user limits (`anonHourlyRateLimit`, `anonMinuteRateLimit`, and `anonPermanentBlock`). 173 | 174 | #### Rate limits 175 | 176 | There are three rate levels that can be individually set or disabled: 177 | 178 | - A maximum number queries per hour (`userMinuteRateLimit` and `anonHourlyRateLimit`) 179 | - A maximum number of queries per minute (`userHourlyRateLimit` and `anonMinuteRateLimit`) 180 | - A maximum number of blocked messages (`userPermanentBlock` and `anonPermanentBlock`) 181 | 182 | If the 1 minute limit is reached, the user will be blocked for 5 minutes. If the hourly limit is reached, the user will be blocked for 60 minutes. If a user has exceeded the value set in `userPermanentBlock` or `anonPermanentBlock` they will be banned permanently. These values are hardcoded in Pico AI Proxy. 183 | 184 | Note 185 | Pico AI Proxy currently does not persist data. Upon server restart, any permanently blocked users will be unblocked. 186 | 187 | ## How to call Pico AI Proxy from your iOS or macOS App 188 | 189 | **Note:** It is highly recommended to set the [`appAccountToken`](https://developer.apple.com/documentation/storekit/product/purchaseoption/3749440-appaccounttoken) to a user-identifying UUID when the user subscribes, like so: 190 | ```swift 191 | let result = try await subscription.product.purchase(options: [.appAccountToken(userIdentifyingUUID)]) 192 | ``` 193 | Setting the `appAccountToken` enables Pico Proxy to use the value for user-specific rate limiting. 194 | 195 | ### Client Apps Using StoreKit 2 196 | 197 | Here is an overview of the steps to have Pico Proxy validate a purchase and grant a user access: 198 | 1. The client app receives a `Transaction` from the App Store Server after a purchase or via a push notification from the App Store to StoreKit 2. 199 | 2. The client app fetches the signed JWS transaction stored in StoreKit 2's `VerificationResult.jwsRepresentation`. 200 | 3. The client app sends the raw signed JWS transaction in the body of an HTTP POST request to `https://.up.railway.app/appstore`. 201 | 4. The proxy server validates the authenticity and validity of the JWS transaction using Apple's [App Store Server Library](https://github.com/apple/app-store-server-library-swift). 202 | 5. The proxy server creates and returns the session token to the client app. 203 | 6. The client app includes the session token in every call until the session token expires. When it expires, the server will return a 401 Unauthorized error. 204 | 205 | **Note:** The client app should always fetch a new session token when it receives a 401 Unauthorized error. 206 | 207 | The code below is based on the WWDC StoreKit 2 [Backyard Birds](https://developer.apple.com/documentation/SwiftUI/Backyard-birds-sample) example code. 208 | 209 | ```swift 210 | import StoreKit 211 | 212 | @MainActor 213 | public final class StoreSubscriptionController: ObservableObject { 214 | @Published public private(set) var jwsTransaction: String? 215 | 216 | public func purchase(option subscription: Subscription) async -> PurchaseFinishedAction { 217 | let action: PurchaseFinishedAction 218 | do { 219 | // Add user identifier to transaction 220 | let idUUID = UUID() 221 | 222 | let result = try await subscription.product.purchase(options: [.appAccountToken(idUUID)]) 223 | switch result { 224 | case .success(let verificationResult): 225 | // Set the JWS token after purchase 226 | jwsTransaction = verificationResult.jwsRepresentation 227 | ... 228 | } 229 | } 230 | } 231 | 232 | // Handle push notification from App Store 233 | internal func handle(update status: Product.SubscriptionInfo.Status) { 234 | guard case .verified(let transaction) = status.transaction, 235 | case .verified(let renewalInfo) = status.renewalInfo else { 236 | return 237 | } 238 | if status.state == .subscribed || status.state == .inGracePeriod { 239 | jwsTransaction = status.transaction.jwsRepresentation 240 | } 241 | ... 242 | } 243 | 244 | // Handle updated entitlement 245 | func updateEntitlement(groupID: String) async { 246 | guard let statuses = try? await Product.SubscriptionInfo.status(for: groupID) else { 247 | return 248 | } 249 | for status in statuses { 250 | guard case .verified(let transaction) = status.transaction, 251 | case .verified(let renewalInfo) = status.renewalInfo else { 252 | continue 253 | } 254 | if status.state == .subscribed || status.state == .inGracePeriod { 255 | jwsTransaction = status.transaction.jwsRepresentation 256 | } 257 | ... 258 | } 259 | } 260 | ``` 261 | 262 | To authenticate a user, call the `appstore` endpoint of Pico Proxy: 263 | ```swift 264 | 265 | class PicoClient { 266 | 267 | var authToken: String? 268 | 269 | func authenticate() async throws { 270 | // Set body to `jwsTransaction` property of `StoreSubscriptionController` 271 | guard let body = await StoreActor.shared.subscriptionController.jwsTransaction else { 272 | // User has no subscription 273 | throw YourClientError.noSubscription 274 | } 275 | 276 | let tokenRequest = Request( 277 | path: "appstore", 278 | method: .post, 279 | body: body, 280 | headers: nil) 281 | let clientConfiguration = APIClient.Configuration(baseURL: "") 282 | let client = APIClient(configuration: clientConfiguration) 283 | let tokenResponse = try await client.send(tokenRequest) 284 | self.authToken = tokenResponse.value.token 285 | } 286 | 287 | func chatConnection() -> OpenAIAPIConnection { 288 | return OpenAIAPIConnection(apiKey: authToken ?? "NO_KEY", 289 | organization: organization, 290 | scheme: scheme.rawValue, 291 | host: host, 292 | chatCompletionPath: chatCompletionPath, 293 | port: port) 294 | } 295 | ... 296 | } 297 | ``` 298 | 299 | ### Client Apps Using Deprecated App Store Receipts 300 | 301 | **Note:** This method is not recommended as App Store receipts are deprecated, and the process is slower and more error-prone. Instead, use StoreKit 2 as described above. While it is technically possible to use StoreKit 2 *in combination with* App Store receipts, it is not recommended because there is a delay between a purchase and the transaction being included in the App Store receipt, which can lead to incorrect Unauthorized errors from Pico Proxy. 302 | 303 | This flow is slightly different: 304 | 1. The client loads the App Store receipt from disk. 305 | 2. The client sends the base64-encoded app receipt in the body of an HTTP POST request to `https://.up.railway.app/appstore`. 306 | 3. Pico Proxy extracts the transaction ID from the App Store receipt and verifies the transaction ID with Apple's App Store Server API. 307 | 4. If the transaction is found, Pico Proxy will create and return a session token to the client app. 308 | 5. The client app includes the session token in every call until the session token expires. When it expires, the server will return a 401 Unauthorized error. 309 | 310 | Using [CleverBird](http://github.com/btfranklin/CleverBird/issues) 311 | ```Swift 312 | import Get 313 | import CleverBird 314 | 315 | var token: Token? = nil 316 | 317 | func completion(prompt: String) async await { 318 | let openAIConnection = OpenAIAPIConnection(apiKey: token.token, organization: "", scheme: "http", host: "localhost", port: 8080) 319 | let chatThread = ChatThread() 320 | .addSystemMessage(content: "You are a helpful assistant.") 321 | .addUserMessage(content: "Who won the world series in 2020?") 322 | do { 323 | let completion = try await chatThread.complete(using: openAIAPIConnection) 324 | } catch CleverBird.proxyAuthenticationRequired { 325 | // Client needs to re-authenticate 326 | token = try await fetchToken() 327 | try await completion(prompt: String) 328 | } catch CleverBirdError.unauthorized { 329 | // Prompt user to buy a subscription 330 | } 331 | } 332 | 333 | func fetchToken() async throws -> Token { 334 | let body: String? 335 | 336 | /* 337 | // Fetch app store receipt 338 | if let appStoreReceiptURL = Bundle.main.appStoreReceiptURL, 339 | FileManager.default.fileExists(atPath: appStoreReceiptURL.path), 340 | let receiptData = try? Data(contentsOf: appStoreReceiptURL, options: .alwaysMapped) { 341 | body = receiptData.base64EncodedString(options: []) 342 | } else { 343 | // when running the app in Xcode Sandbox, there will be no receipt. In sandbox, Pico AI Proxy will accept 344 | // the receipt Id. 345 | body = "transaction Id here" 346 | } 347 | */ 348 | 349 | // Validating receipts is temporary disabled 350 | body = "transaction Id here" 351 | 352 | let tokenRequest = Request( 353 | path: "appstore", 354 | method: .post, 355 | body: body, 356 | headers: nil) 357 | let tokenResponse = try await AIClient.openAIAPIConnection.client.send(tokenRequest) 358 | return tokenResponse.value 359 | } 360 | 361 | struct Token: Codable { 362 | let token: String 363 | } 364 | ``` 365 | 366 | Optionally: Track users using [app account token](https://developer.apple.com/documentation/storekit/product/purchaseoption/3749440-appaccounttoken) 367 | ```Swift 368 | // Create new UUID 369 | let id = UUID() 370 | // Add id to user account 371 | 372 | // Purchase subscription 373 | let result = try await product.purchase(options: [.appAccountToken(idUUID)]) 374 | ``` 375 | 376 | Pico AI Proxy will automatically extract the app account token from the receipts. 377 | 378 | Pico AI Proxy may generate two distinct error codes related to authorization issues: unauthorized (401) and proxyAuthenticationRequired (407). The unauthorized error indicates a lack of a valid App Store subscription on the user's part. On the other hand, the proxyAuthenticationRequired error signifies that the client's authentication token is no longer valid, a situation that may arise following a server reboot. In the latter case, reauthorization can be achieved through a straightforward re-authentication process that does not require user intervention. 379 | 380 | ## How to deploy 381 | 382 | Use link below to deploy Pico AI Proxy on Railway. The link includes a referral code. 383 | 384 | [![Deploy on Railway](https://railway.app/button.svg)](https://railway.app/template/ocPcV2?referralCode=WKPLp3) 385 | 386 | Alternatively, Pico AI Proxy can be installed manually on any other hosting provider. 387 | 388 | 389 | ## Support 390 | 391 | - [Pico Discord](https://discord.gg/Nrf5y8Uaxw) 392 | - [Twitter](https://twitter.com/picoGPT) 393 | 394 | ## Apps using Pico AI Proxy 395 | 396 | - [Pico](https://apps.apple.com/us/app/pico-professional-ai-assistant/id1668205047) 397 | 398 | ## Contributors 399 | 400 | 401 | 402 | -------------------------------------------------------------------------------- /Sources/App/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PicoMLX/PicoAIProxy/5ecaaf24675bf0b9a51aa14323b4b31138cd8f86/Sources/App/.DS_Store -------------------------------------------------------------------------------- /Sources/App/App.swift: -------------------------------------------------------------------------------- 1 | // 2 | // App.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 12/19/23. 6 | // 7 | // Based on https://opticalaberration.com/2021/12/proxy-server.html 8 | 9 | import ArgumentParser 10 | import Hummingbird 11 | 12 | @main 13 | struct HummingbirdArguments: AsyncParsableCommand, AppArguments { 14 | @Option(name: .shortAndLong) 15 | var hostname: String = "0.0.0.0" 16 | 17 | @Option(name: .shortAndLong) 18 | var port: Int = 8080 19 | 20 | @Option(name: .shortAndLong) 21 | var location: String = "" // Note: this is ignored 22 | 23 | @Option(name: .shortAndLong) 24 | var target: String = "" // Note: this is ignored 25 | 26 | func run() async throws { 27 | 28 | // Load models and providers 29 | LLMModel.load() 30 | 31 | // Use Railway.app's port 32 | let port = Int(HBEnvironment().get("PORT") ?? "8080") ?? 8080 33 | let app = HBApplication( 34 | configuration: .init( 35 | address: .hostname(self.hostname, port: port), 36 | serverName: "PicoAIProxy" 37 | ) 38 | ) 39 | try await app.configure(self) 40 | try app.start() 41 | await app.asyncWait() 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /Sources/App/Application+configure.swift: -------------------------------------------------------------------------------- 1 | import AsyncHTTPClient 2 | import Hummingbird 3 | import Foundation 4 | import JWTKit 5 | import HummingbirdAuth 6 | import FluentKit 7 | import FluentSQLiteDriver 8 | import HummingbirdFluent 9 | 10 | public protocol AppArguments { 11 | var location: String { get } 12 | var target: String { get } 13 | } 14 | 15 | extension HBApplication { 16 | /// configure your application 17 | /// add middleware 18 | /// setup the encoder/decoder 19 | /// add your routes 20 | func configure(_ args: AppArguments) async throws { 21 | 22 | // 1. Set up JSON encoder and decoder 23 | self.encoder = JSONEncoder() 24 | self.decoder = JSONDecoder() 25 | (self.encoder as! JSONEncoder).dateEncodingStrategy = .iso8601 26 | self.httpClient = HTTPClient(eventLoopGroupProvider: .shared(self.eventLoopGroup)) 27 | 28 | // 2. Add logging middleware 29 | let logLevel = HBEnvironment().get("logLevel") ?? "info" 30 | self.logger.logLevel = .init(rawValue: logLevel) ?? .info 31 | self.logger.info("Logger level is set to \(self.logger.logLevel)") 32 | self.middleware.add(HBLogRequestsMiddleware(.info)) 33 | self.middleware.add(HBLogRequestsMiddleware(.debug)) 34 | self.middleware.add(HBLogRequestsMiddleware(.error)) 35 | 36 | // 3. Set up database 37 | self.addFluent() 38 | // if let inMemory = HBEnvironment().get("inMemoryDatabase"), inMemory == "1" { 39 | self.fluent.databases.use(.sqlite(.memory), as: .sqlite) 40 | // } else { 41 | // self.fluent.databases.use(.sqlite(.file("db.sqlite")), as: .sqlite) 42 | // } 43 | 44 | // 4. Add migrations 45 | self.fluent.migrations.add(UserMigration()) 46 | self.fluent.migrations.add(UserRequestMigration()) 47 | try await self.fluent.migrate() 48 | 49 | // 5. Fetch JWT private key from environment and set up JWT Signers 50 | guard let jwtKey = HBEnvironment().get("JWTPrivateKey"), 51 | !jwtKey.isEmpty else { 52 | self.logger.error("JWTPrivateKey environment variable must be set") 53 | throw HBHTTPError(.internalServerError) 54 | } 55 | let jwtAuthenticator = JWTAuthenticator() 56 | let jwtLocalSignerKid = JWKIdentifier("_aiproxy_local_") 57 | jwtAuthenticator.useSigner(.hs256(key: jwtKey), kid: jwtLocalSignerKid) 58 | 59 | // 6. Add AppStoreController routes to verify client's purchase and send JWT token to client 60 | let appStoreController = AppStoreController(jwtSigners: jwtAuthenticator.jwtSigners, kid: jwtLocalSignerKid) 61 | appStoreController.addRoutes(to: self.router.group("appstore")) 62 | 63 | // 7. Add JWT authenticator. Will return unauthorized error if no or invalid JWT token was received 64 | self.middleware.add(jwtAuthenticator) 65 | 66 | // 8. Add rate limiter 67 | self.middleware.add(RateLimiterMiddleware()) 68 | 69 | // 9. Route message to right provider 70 | self.middleware.add(MessageRouterMiddleware()) 71 | 72 | // 10. Add OpenAI API key middleware. This middleware will add the OpenAI org and API key in the header of the request 73 | self.middleware.add(APIKeyMiddleware()) 74 | 75 | // 11. Add Proxy middleware. If you don't need any authentication, you can remove steps 3 through 6 above 76 | self.middleware.add( 77 | HBProxyServerMiddleware( 78 | httpClient: httpClient //, 79 | // proxy: .init(location: args.location, target: args.target) // Note: This is ignored 80 | ) 81 | ) 82 | } 83 | } 84 | 85 | extension HBApplication { 86 | var httpClient: HTTPClient { 87 | get { self.extensions.get(\.httpClient) } 88 | set { self.extensions.set(\.httpClient, value: newValue) { httpClient in 89 | try httpClient.syncShutdown() 90 | }} 91 | } 92 | } 93 | -------------------------------------------------------------------------------- /Sources/App/Controllers/AppStoreController.swift: -------------------------------------------------------------------------------- 1 | // 2 | // AppStoreController.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/6/24. 6 | // 7 | 8 | import Foundation 9 | import Hummingbird 10 | import JWTKit 11 | import FluentKit 12 | 13 | struct AppStoreController { 14 | 15 | let jwtSigners: JWTSigners 16 | let kid: JWKIdentifier 17 | 18 | /// Add routes for /appstore 19 | func addRoutes(to group: HBRouterGroup) { 20 | 21 | // Set app store authenticator that returns valid User instance 22 | // If the setup fails because the environment variables aren't set, 23 | // calls will be forwarded to an error message 24 | let appStoreAuthenticator: AppStoreAuthenticator 25 | do { 26 | appStoreAuthenticator = try AppStoreAuthenticator() 27 | } catch { 28 | group 29 | .post("/", use: incorrectSetup) 30 | return 31 | } 32 | 33 | group 34 | .add(middleware: appStoreAuthenticator) 35 | .post("/", use: login) 36 | } 37 | 38 | private func incorrectSetup(_ request: HBRequest) async throws -> String { 39 | request.logger.error("Missing environment variable(s): IAPPrivateKey, IAPIssuerId, IAPKeyId, appBundleId and/or appAppleId") 40 | throw HBHTTPError(.internalServerError, message: "IAPPrivateKey, IAPIssuerId, IAPKeyId and/or appBundleId environment variables not set") 41 | } 42 | 43 | /// Note: appAccountToken can be nil. All users with an empty appAccountToken 44 | /// will be treated as a single user 45 | private func login(_ request: HBRequest) async throws -> [String: String] { 46 | 47 | // 1. Fetch user from AppStoreAuthenticator middleware 48 | let user = try request.authRequire(User.self) 49 | 50 | request.logger.info("Starting login for user \(user.appAccountToken?.uuidString ?? "anon")") 51 | 52 | // 2. If user is a new user, add user to database 53 | if try await User.query(on: request.db) 54 | .filter(\.$appAccountToken == user.appAccountToken) 55 | .first() 56 | == nil { 57 | try await user.save(on: request.db) 58 | request.logger.info("Saved user with app account token \(user.appAccountToken?.uuidString ?? "anon")") 59 | } 60 | 61 | let payload = JWTPayloadData( 62 | subject: .init(value: user.appAccountToken?.uuidString ?? "NO_ACCOUNT"), 63 | expiration: .init(value: Date(timeIntervalSinceNow: 12 * 60 * 60)) // 12 hours 64 | ) 65 | 66 | let token = try self.jwtSigners.sign(payload, kid: self.kid) 67 | 68 | user.jwtToken = token 69 | try await user.save(on: request.db) 70 | 71 | return [ 72 | "token": token, 73 | ] 74 | } 75 | } 76 | 77 | 78 | 79 | -------------------------------------------------------------------------------- /Sources/App/Extensions/Status.swift: -------------------------------------------------------------------------------- 1 | // 2 | // File.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 3/23/24. 6 | // 7 | 8 | import Foundation 9 | import AppStoreServerLibrary 10 | 11 | extension AppStoreServerLibrary.Status: CustomStringConvertible { 12 | 13 | public var description: String { 14 | switch self { 15 | case .active: 16 | return "active" 17 | case .expired: 18 | return "expired" 19 | case .billingRetry: 20 | return "billing retry" 21 | case .billingGracePeriod: 22 | return "billing grace period" 23 | case .revoked: 24 | return "revoked" 25 | } 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /Sources/App/Middleware/APIKeyMiddleware.swift: -------------------------------------------------------------------------------- 1 | // 2 | // APIKeyMiddleware.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 12/29/23. 6 | // 7 | 8 | import Hummingbird 9 | import NIOHTTP1 10 | 11 | struct APIKeyMiddleware: HBMiddleware { 12 | func apply(to request: Hummingbird.HBRequest, next: Hummingbird.HBResponder) -> NIOCore.EventLoopFuture { 13 | 14 | var headers = request.headers 15 | 16 | // Make sure not to send user's Pico Proxy auth token to provider 17 | headers.remove(name: "Authorization") 18 | 19 | do { 20 | if let modelName = headers.first(name: "model"), let model = LLMModel.fetch(model: modelName) { 21 | try model.provider.setHeaders(headers: &headers) 22 | } else { 23 | // Default to OpenAI 24 | try LLMProvider.openAI.setHeaders(headers: &headers) 25 | } 26 | } catch { 27 | request.logger.error("Error OpenAIKeyMiddleware: \(error.localizedDescription)") 28 | return request.failure(error) 29 | } 30 | 31 | let head = HTTPRequestHead(version: request.version, method: request.method, uri: request.uri.string, headers: headers) 32 | let request = HBRequest(head: head, body: request.body, application: request.application, context: request.context) 33 | 34 | return next.respond(to: request) 35 | } 36 | } 37 | -------------------------------------------------------------------------------- /Sources/App/Middleware/AppStoreAuthenticator.swift: -------------------------------------------------------------------------------- 1 | // 2 | // AppStoreAuthenticator.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/21/24. 6 | // 7 | 8 | import Foundation 9 | import FluentKit 10 | import Hummingbird 11 | import HummingbirdAuth 12 | import AppStoreServerLibrary 13 | #if os(Linux) 14 | import FoundationNetworking 15 | #endif 16 | 17 | /// Defines a custom authenticator for App Store transactions 18 | struct AppStoreAuthenticator: HBAsyncAuthenticator { 19 | 20 | // Properties to hold App Store credentials and app-specific information 21 | let iapKey: String 22 | let iapIssuerId: String 23 | let iapKeyId: String 24 | let bundleId: String 25 | let appAppleId: Int64 26 | let environment: Environment 27 | 28 | /// Initializer to load necessary credentials and configuration from environment variables 29 | init() throws { 30 | 31 | // Fetch IAP private key, issuer ID, and Key ID from environment variables 32 | // Information about creating a private key is available in Apple's documentation 33 | // Failing to find required environment variables results in an error 34 | // To create a private key, see: 35 | // https://developer.apple.com/documentation/appstoreserverapi/creating_api_keys_to_use_with_the_app_store_server_api 36 | // and https://developer.apple.com/wwdc23/10143 37 | guard let iapKey = HBEnvironment().get("IAPPrivateKey")?.replacingOccurrences(of: "\\\\n", with: "\n"), 38 | !iapKey.isEmpty, 39 | let iapIssuerId = HBEnvironment().get("IAPIssuerId"), 40 | !iapIssuerId.isEmpty, 41 | let iapKeyId = HBEnvironment().get("IAPKeyId"), 42 | !iapKeyId.isEmpty, 43 | let bundleId = HBEnvironment().get("appBundleId"), 44 | !bundleId.isEmpty, 45 | let appAppleIdString = HBEnvironment().get("appAppleId"), 46 | let appAppleId = Int64(appAppleIdString) 47 | else { 48 | // If the environment variables are not set, SwiftProxyAIServer will throw an internal server error. 49 | // Check your server's logs for the message below 50 | throw HBHTTPError(.internalServerError, message: "IAPPrivateKey, IAPIssuerId, IAPKeyId and/or appBundleId, appAppleId, environment variable(s) not set") 51 | } 52 | 53 | self.iapKey = iapKey 54 | self.iapIssuerId = iapIssuerId 55 | self.iapKeyId = iapKeyId 56 | self.bundleId = bundleId 57 | self.appAppleId = appAppleId 58 | 59 | // making the instance environment specific 60 | let environmentString = HBEnvironment().get("environment") ?? "Production" 61 | self.environment = Environment(rawValue: environmentString) ?? .production 62 | 63 | } 64 | 65 | /// Authenticates incoming requests based on App Store receipt or transaction ID 66 | /// - Parameter request: HBRequest 67 | /// - Returns: User model if app store receipt is valid 68 | func authenticate(request: HBRequest) async throws -> User? { 69 | 70 | // 1. The server expects an app store receipt 71 | // (in the iOS and macOS client app: Bundle.main.appStoreReceiptURL) 72 | // However, the receipt is not available when testing in Xcode Sandbox, 73 | // so the server accepts a transaction Id in sandbox mode as well 74 | let request = try await request.collateBody().get() 75 | guard let buffer = request.body.buffer, let body = buffer.getString(at: buffer.readerIndex, length: buffer.readableBytes) else { 76 | request.logger.error("AppStoreAuthenticator: /appstore invoked with empty body") 77 | throw HBHTTPError(.badRequest) 78 | } 79 | 80 | // same proxy instance can be used as both production and sandbox environment 81 | request.logger.debug("Attempting to validate StoreKit2 JWS using \(self.environment) environment") 82 | if environment == .production { 83 | if let payload = try await validateJWS(jws: body, environment: .production, request: request) { 84 | // Validated production transaction 85 | return try await addUser(request: request, payload: payload, environment: .production) 86 | 87 | } else if let payload = try await validateJWS(jws: body, environment: .sandbox, request: request) { 88 | // Validated sandbox transaction 89 | return try await addUser(request: request, payload: payload, environment: .sandbox) 90 | } 91 | } else { 92 | if let payload = try await validateJWS(jws: body, environment: self.environment, request: request) { 93 | // Validated sandbox transaction 94 | return try await addUser(request: request, payload: payload, environment: self.environment) 95 | } 96 | } 97 | 98 | // If StoreKit 2 verification didn't work - 99 | // Use Store Kit 1 app receipts 100 | request.logger.debug("Attempting to validate StoreKit 1 app receipt") 101 | 102 | // Extract transaction ID 103 | guard let id = ReceiptUtility.extractTransactionId(appReceipt: body) else { 104 | throw HBHTTPError(.unauthorized) 105 | } 106 | let transactionId = id 107 | 108 | // Try validating the transaction in production environment first 109 | // If not found (404 error), retries in the sandbox environment for TestFlight users 110 | do { 111 | return try await validate(request, transactionId: transactionId, environment: .production) 112 | } catch let error as HBHTTPError where error.status == .notFound { 113 | request.logger.error("AppStoreAuthenticator: Caught HBHTTPError.notFound. Validating transaction in sandbox.") 114 | return try await validate(request, transactionId: transactionId, environment: .sandbox) 115 | } 116 | } 117 | 118 | // MARK: - Store Kit 2 JWS authentication 119 | 120 | /// Validates JWS string and checks expiry date 121 | /// - Parameters: 122 | /// - jws: the jws string from the client app 123 | /// - environment: Either .production, .sandbox, or .xcode 124 | /// - request: HBRequest 125 | /// - Returns: Payload if successful or nil if jws has a different environment than provided 126 | private func validateJWS(jws: String, environment: Environment, request: HBRequest) async throws -> JWSTransactionDecodedPayload? { 127 | 128 | // 1. Set up JWT verifier 129 | let rootCertificates = try loadAppleRootCertificates(request: request) 130 | let verifier = try SignedDataVerifier(rootCertificates: rootCertificates, bundleId: bundleId, appAppleId: appAppleId, environment: environment, enableOnlineChecks: true) 131 | request.logger.debug("environment: \(environment)") 132 | 133 | // 2. Parse JWS transaction 134 | let verifyResponse = await verifier.verifyAndDecodeTransaction(signedTransaction: jws) 135 | request.logger.debug("verifyResponse: \(verifyResponse)") 136 | 137 | switch verifyResponse { 138 | case .valid(let payload): 139 | 140 | // Check expiry date 141 | if let date = payload.expiresDate, date < Date() { 142 | request.logger.error("Subscription of user \(payload.appAccountToken?.uuidString ?? "anon") expired on \(date)") 143 | throw HBHTTPError(.unauthorized) 144 | } 145 | 146 | // Check revocation date (for refunds and family removal) 147 | if let date = payload.revocationDate, date < Date() { 148 | request.logger.error("Subscription of user \(payload.appAccountToken?.uuidString ?? "anon") was revoked on \(date)") 149 | throw HBHTTPError(.unauthorized) 150 | } 151 | 152 | request.logger.info("AppStoreAuthenticator: validated tx for user \(payload.appAccountToken?.uuidString ?? "anon") in \(environment)") 153 | return payload 154 | 155 | case .invalid(let error): 156 | 157 | switch error { 158 | case .INVALID_JWT_FORMAT: 159 | request.logger.error("validateJWS failed: invalid JWT format") 160 | case .INVALID_CERTIFICATE: 161 | request.logger.error("validateJWS failed: invalid certificate") 162 | case .VERIFICATION_FAILURE: 163 | request.logger.error("validateJWS failed: verification failed") 164 | case .INVALID_APP_IDENTIFIER: 165 | request.logger.error("validateJWS failed: invalid app identifier") 166 | case .INVALID_ENVIRONMENT: 167 | // Return nil so caller can try a different environment 168 | return nil 169 | } 170 | // We don't want to throw an error here as it may be caused by an environment mismatch 171 | // throw HBHTTPError(.unauthorized) 172 | return nil 173 | } 174 | } 175 | 176 | /// Finds existing user or creates new user in database 177 | /// - Parameters: 178 | /// - request: HBRequest 179 | /// - payload: JWS payload 180 | /// - environment: Either .production, .sandbox, or .xcode 181 | /// - Returns: user 182 | private func addUser(request: HBRequest, payload: JWSTransactionDecodedPayload, environment: Environment) async throws -> User { 183 | 184 | // 1. create user 185 | let user = User(appAccountToken: nil, environment: environment.rawValue, productId: "", status: .expired) 186 | user.appAccountToken = payload.appAccountToken 187 | if let productId = payload.productId { 188 | user.productId = productId 189 | } 190 | 191 | // 2. If user doesn't exist, add to database 192 | // Otherwise, return existing user 193 | if let existingUser = try await User.query(on: request.db) 194 | .filter(\.$appAccountToken == user.appAccountToken) 195 | .first() { 196 | request.logger.info("addUser found existing user \(payload.appAccountToken?.uuidString ?? "anon") in database") 197 | return existingUser 198 | } 199 | try await user.save(on: request.db) 200 | request.logger.info("addUser added user \(payload.appAccountToken?.uuidString ?? "anon") to database") 201 | return user 202 | } 203 | 204 | // MARK: - Pre-Store Kit 2 App Store Receipt authentication 205 | 206 | /// Validates the transaction ID with the App Store and returns a User if successful 207 | /// - Parameters: 208 | /// - request: http request 209 | /// - transactionId: transaction Id. Can be original transaction Id 210 | /// - environment: e.g. .sandbox 211 | /// - Returns: new or existing user 212 | private func validate(_ request: HBRequest, transactionId: String, environment: Environment) async throws -> User? { 213 | 214 | // 1. Create App Store API client 215 | let appStoreClient = try AppStoreServerAPIClient(signingKey: iapKey, keyId: iapKeyId, issuerId: iapIssuerId, bundleId: bundleId, environment: environment) 216 | 217 | request.logger.info("AppStoreAuthenticator: Created API Client for keyId: \(iapKeyId), issuer: \(iapIssuerId), bundleId: \(bundleId), env: \(environment.rawValue)") 218 | 219 | // 3. create user 220 | let user = User(appAccountToken: nil, environment: environment.rawValue, productId: "", status: .expired) 221 | 222 | // 4. Use transactionId to fetch active subscriptions from App Store 223 | let allSubs = await appStoreClient.getAllSubscriptionStatuses(transactionId: transactionId, status: [.active, .billingGracePeriod]) 224 | switch allSubs { 225 | case .success(let response): 226 | 227 | request.logger.info("AppStoreAuthenticator: TxId: \(transactionId) \(environment.rawValue): number of data: \(response.data?.count ?? 0)") 228 | 229 | // SwiftProxyServer assumes app has a single subscription group 230 | guard let subscriptionGroup = response.data?.first, 231 | let lastTransactions = subscriptionGroup.lastTransactions else { 232 | request.logger.error("AppStoreAuthenticator: TxId: \(transactionId) \(environment.rawValue): Get all subscriptions succeeded but returned no transactions. No subscription group or no last transactions in \(environment.rawValue) for \(transactionId)") 233 | throw HBHTTPError(.unauthorized, message: "No active or grace period subscription status found") 234 | } 235 | 236 | request.logger.info("AppStoreAuthenticator: TxId: \(transactionId) \(environment.rawValue): Found \(lastTransactions.count) transactions") 237 | 238 | // Loop through the transactions in the subscription group 239 | for transaction in lastTransactions { 240 | 241 | request.logger.info("AppStoreAuthenticator: Parsing transaction \(transaction.originalTransactionId ?? "(No original tx id)"), status: \(transaction.status?.description ?? "(no known status)")") 242 | 243 | guard let signedTransactionInfo = transaction.signedTransactionInfo else { continue } 244 | 245 | let appAccountToken: UUID? 246 | let productId: String? 247 | do { 248 | let (token, product) = try await fetchUserAppAccountToken(signedTransactionInfo: signedTransactionInfo, environment: environment, request: request) 249 | appAccountToken = token 250 | productId = product 251 | } catch { 252 | // Skip to next product in case of an error 253 | continue 254 | } 255 | 256 | if let productId { 257 | user.productId = productId 258 | } 259 | 260 | if let appAccountToken { 261 | user.appAccountToken = appAccountToken 262 | } 263 | 264 | // We're only saving one subscription, an active one if available. 265 | // Update user.status. Make sure we don't overwrite it if the status is already .active 266 | if user.subscriptionStatus != Status.active.rawValue, let status = transaction.status { 267 | user.subscriptionStatus = status.rawValue 268 | } 269 | 270 | request.logger.info("AppStoreAuthenticator: TxId: \(transactionId) \(environment.rawValue): Found token \(appAccountToken?.uuidString ?? "(no token)") for \(productId ?? "(no product ID)") in \(environment.rawValue)") 271 | 272 | if let _ = user.appAccountToken, !user.productId.isEmpty { 273 | // We have all information 274 | break 275 | } 276 | } 277 | 278 | case .failure(let statusCode, let rawApiError, let apiError, let errorMessage, let causedBy): 279 | 280 | request.logger.error("AppStoreAuthenticator: TxId: \(transactionId) \(environment.rawValue): Get all subscriptions returned an error: code \(statusCode ?? 0), api error: \(rawApiError ?? 0), error msg: \(errorMessage ?? ""), caused by: \(causedBy?.localizedDescription ?? "")") 281 | 282 | if statusCode == 404 { 283 | // No transaction was found. 284 | request.logger.error("AppStoreAuthenticator: TxId: \(transactionId) \(environment.rawValue) not found. Error: \(statusCode ?? -1): \(errorMessage ?? "Unknown error"), \(String(describing: rawApiError)) \(String(describing: apiError)), \(String(describing: causedBy))") 285 | throw HBHTTPError(.notFound) 286 | } else { 287 | // Other error occurred 288 | request.logger.error("AppStoreAuthenticator: TxId: \(transactionId) \(environment.rawValue): Get all subscriptions failed in \(environment.rawValue) for \(transactionId). Error: \(statusCode ?? -1): \(errorMessage ?? "Unknown error"), \(String(describing: rawApiError)) \(String(describing: apiError)), \(String(describing: causedBy))") 289 | throw HBHTTPError(HTTPResponseStatus(statusCode: statusCode ?? 500, reasonPhrase: errorMessage ?? "Unknown error")) 290 | } 291 | } 292 | 293 | // 5. If user doesn't exist, add to database 294 | // Otherwise, return existing user 295 | if let existingUser = try await User.query(on: request.db) 296 | .filter(\.$appAccountToken == user.appAccountToken) 297 | .first() { 298 | return existingUser 299 | } 300 | try await user.save(on: request.db) 301 | return user 302 | } 303 | 304 | /// Fetches app account token set by client app and productID 305 | /// Note: Token is nil when client app doesn't set appAccountToken during the purchase 306 | /// See https://developer.apple.com/documentation/storekit/product/3791971-purchase 307 | /// - Parameters: 308 | /// - signedTransactionInfo: SignedTransactionInfo fetched by getAllSubscriptionStatuses 309 | /// - environment: E.g. .sandbox 310 | /// - request: The request (to access logger) 311 | /// - Returns: Tuple of optional appAccountToken and optional productID 312 | private func fetchUserAppAccountToken(signedTransactionInfo: String, environment: Environment, request: HBRequest) async throws -> (UUID?, String?) { 313 | 314 | // 1. Set up JWT verifier 315 | let rootCertificates = try loadAppleRootCertificates(request: request) 316 | let verifier = try SignedDataVerifier(rootCertificates: rootCertificates, bundleId: bundleId, appAppleId: appAppleId, environment: environment, enableOnlineChecks: true) 317 | 318 | // 5. Parse signed transaction 319 | let verifyResponse = await verifier.verifyAndDecodeTransaction(signedTransaction: signedTransactionInfo) 320 | 321 | switch verifyResponse { 322 | case .valid(let payload): 323 | 324 | return (payload.appAccountToken, payload.productId) 325 | 326 | case .invalid(let error): 327 | 328 | request.logger.error("Verifying transaction failed. Error: \(error)") 329 | throw HBHTTPError(.unauthorized) 330 | } 331 | } 332 | 333 | private func loadAppleRootCertificates(request: HBRequest) throws -> [Foundation.Data] { 334 | #if os(Linux) 335 | // Linux doesn't have app bundles, so we're copying the certificates in the Dockerfile to /app/Resources and load them manually 336 | return [ 337 | try loadData(url: URL(string: "/app/Resources/AppleComputerRootCertificate.cer"), request: request), 338 | try loadData(url: URL(string: "/app/Resources/AppleIncRootCertificate.cer"), request: request), 339 | try loadData(url: URL(string: "/app/Resources/AppleRootCA-G2.cer"), request: request), 340 | try loadData(url: URL(string: "/app/Resources/AppleRootCA-G3.cer"), request: request), 341 | ].compactMap { $0 } 342 | #else 343 | return [ 344 | try loadData(url: Bundle.module.url(forResource: "AppleComputerRootCertificate", withExtension: "cer"), request: request), 345 | try loadData(url: Bundle.module.url(forResource: "AppleIncRootCertificate", withExtension: "cer"), request: request), 346 | try loadData(url: Bundle.module.url(forResource: "AppleRootCA-G2", withExtension: "cer"), request: request), 347 | try loadData(url: Bundle.module.url(forResource: "AppleRootCA-G3", withExtension: "cer"), request: request), 348 | ].compactMap { $0 } 349 | #endif 350 | } 351 | 352 | private func loadData(url: URL?, request: HBRequest) throws -> Foundation.Data? { 353 | let fs = FileManager() 354 | guard let url = url, fs.fileExists(atPath: url.path) else { 355 | request.logger.error("File missing: \(url?.absoluteString ?? "invalid url")") 356 | throw HBHTTPError(.internalServerError) 357 | } 358 | 359 | guard let data = fs.contents(atPath: url.path) else { 360 | request.logger.error("Can't read data from \(url.absoluteString)") 361 | return nil 362 | } 363 | return data 364 | } 365 | } 366 | -------------------------------------------------------------------------------- /Sources/App/Middleware/JWTAuthenticator.swift: -------------------------------------------------------------------------------- 1 | // 2 | // File.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/8/24. 6 | // 7 | 8 | //import FluentKit 9 | import Foundation 10 | import Hummingbird 11 | import HummingbirdAuth 12 | import JWTKit 13 | import NIOFoundationCompat 14 | import FluentKit 15 | 16 | struct JWTPayloadData: JWTPayload, Equatable, HBAuthenticatable { 17 | enum CodingKeys: String, CodingKey { 18 | case subject = "sub" 19 | case expiration = "exp" 20 | } 21 | 22 | var subject: SubjectClaim 23 | var expiration: ExpirationClaim 24 | // Define additional JWT Attributes here 25 | 26 | func verify(using signer: JWTSigner) throws { 27 | try self.expiration.verifyNotExpired() 28 | } 29 | } 30 | 31 | struct JWTAuthenticator: HBAsyncAuthenticator { 32 | 33 | let jwtSigners: JWTSigners 34 | 35 | init() { 36 | self.jwtSigners = JWTSigners() 37 | } 38 | 39 | init(_ signer: JWTSigner, kid: JWKIdentifier) { 40 | self.jwtSigners = JWTSigners() 41 | self.jwtSigners.use(signer, kid: kid) 42 | } 43 | 44 | init(jwksData: ByteBuffer) throws { 45 | let jwks = try JSONDecoder().decode(JWKS.self, from: jwksData) 46 | self.jwtSigners = JWTSigners() 47 | try self.jwtSigners.use(jwks: jwks) 48 | } 49 | 50 | func useSigner(_ signer: JWTSigner, kid: JWKIdentifier) { 51 | self.jwtSigners.use(signer, kid: kid) 52 | } 53 | 54 | func authenticate(request: HBRequest) async throws -> User? { 55 | 56 | // 1. Get JWT token from bearer authorization header 57 | // If no token is present, return unauthorized error 58 | guard let jwtToken = request.authBearer?.token else { 59 | request.logger.error("No jwtToken found") 60 | throw HBHTTPError(.unauthorized) 61 | } 62 | 63 | // 2. If passthrough is enabled, and OpenAI key and org is found in headers 64 | // then forward request 65 | if let passthrough = HBEnvironment().get("allowKeyPassthrough"), 66 | passthrough == "1", 67 | let org = request.headers["OpenAI-Organization"].first, 68 | org.hasPrefix("org-") == true, 69 | jwtToken.hasPrefix("sk-") == true { 70 | request.logger.info("OpenAI API key Passthrough") 71 | return nil 72 | } 73 | 74 | // 3. Verify token is a valid token created by Pico AI Proxy 75 | let payload: String 76 | let appAccountToken: UUID? 77 | do { 78 | payload = try self.jwtSigners.verify(jwtToken, as: JWTPayloadData.self).subject.value 79 | appAccountToken = UUID(uuidString: payload) 80 | } catch { 81 | request.logger.error("Invalid jwtToken received: \(jwtToken)") 82 | throw HBHTTPError(.unauthorized) 83 | } 84 | 85 | // 4. See if user is in database 86 | guard let user = try await User.query(on: request.db) 87 | .filter(\.$appAccountToken == appAccountToken) 88 | .first() else { 89 | 90 | // The user has a valid jwtToken but isn't in the database 91 | // (This can happen after the server restarted) 92 | // Ask user to re-authenticate 93 | request.logger.error("User with app account token \(appAccountToken?.uuidString ?? "anon") not found in database") 94 | throw HBHTTPError(.proxyAuthenticationRequired) 95 | } 96 | 97 | // 5. Return user 98 | request.logger.info("Verified user with app account token \(user.appAccountToken?.uuidString ?? "anon")") 99 | return user 100 | } 101 | } 102 | -------------------------------------------------------------------------------- /Sources/App/Middleware/MessageRouterMiddleware.swift: -------------------------------------------------------------------------------- 1 | // 2 | // MessageParserMiddleware.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 4/6/24. 6 | // 7 | 8 | import Foundation 9 | import Hummingbird 10 | import NIOHTTP1 11 | 12 | /// Note: The Anthropic API is very picky. `max_tokens` is required (an option in OpenAI) and the 13 | /// roles must alternate between `user` and `assistant`. It won't accept multiple `user` roles in a row 14 | /// and extra inputs are not permitted (e.g. `user`) 15 | struct MessageRouterMiddleware: HBAsyncMiddleware { 16 | 17 | func apply(to request: Hummingbird.HBRequest, next: any Hummingbird.HBResponder) async throws -> Hummingbird.HBResponse { 18 | 19 | // 1. Fetch body 20 | let request = try await request.collateBody().get() 21 | guard let buffer = request.body.buffer, let body = buffer.getString(at: buffer.readerIndex, length: buffer.readableBytes), let data = body.data(using: .utf8) else { 22 | request.logger.error("Unable to decode body in MessageRouterMiddleware") 23 | throw HBHTTPError(.badRequest) 24 | } 25 | 26 | // 2. Find model in body. Default to OpenAI if body isn't a chat (but e.g. an embedding) 27 | var headers = request.headers 28 | var uri = request.uri.string 29 | if let model = LLMModel.fetchModel(from: data) { 30 | request.logger.info("Rerouting \(model.name) to \(model.provider.name)") 31 | headers.replaceOrAdd(name: "model", value: model.name) 32 | if !model.proxy().location.isEmpty { 33 | uri = model.proxy().location 34 | } 35 | } 36 | 37 | // 3. Update header 38 | let head = HTTPRequestHead(version: request.version, method: request.method, uri: uri, headers: headers) 39 | let convertedRequest = HBRequest(head: head, body: request.body, application: request.application, context: request.context) 40 | 41 | let response = try await next.respond(to: convertedRequest) 42 | // print("---") 43 | // switch response.body { 44 | // case .stream(let streamer): 45 | // let output = streamer.read(on: request.eventLoop) 46 | // output.whenSuccess { output in 47 | // print(output) 48 | // } 49 | // default: 50 | // break 51 | // } 52 | // print("---") 53 | return response 54 | // return try await next.respond(to: convertedRequest) 55 | } 56 | } 57 | 58 | -------------------------------------------------------------------------------- /Sources/App/Middleware/ProxyServerMiddleware.swift: -------------------------------------------------------------------------------- 1 | //===----------------------------------------------------------------------===// 2 | // 3 | // This source file is part of the Hummingbird server framework project 4 | // 5 | // Copyright (c) 2021-2021 the Hummingbird authors 6 | // Licensed under Apache License v2.0 7 | // 8 | // See LICENSE.txt for license information 9 | // See hummingbird/CONTRIBUTORS.txt for the list of Hummingbird authors 10 | // 11 | // SPDX-License-Identifier: Apache-2.0 12 | // 13 | //===----------------------------------------------------------------------===// 14 | 15 | import AsyncHTTPClient 16 | import Hummingbird 17 | import HummingbirdCore 18 | import Logging 19 | 20 | /// Middleware forwarding requests onto another server 21 | public struct HBProxyServerMiddleware: HBMiddleware { 22 | public struct Proxy { 23 | let location: String 24 | let target: String 25 | 26 | init(location: String, target: String) { 27 | self.location = location.dropSuffix("/") 28 | self.target = target.dropSuffix("/") 29 | } 30 | } 31 | 32 | let httpClient: HTTPClient 33 | 34 | public init(httpClient: HTTPClient) { 35 | self.httpClient = httpClient 36 | } 37 | 38 | public func apply(to request: HBRequest, next: HBResponder) -> EventLoopFuture { 39 | 40 | let proxy: Proxy 41 | if let modelName = request.headers.first(name: "model"), let model = LLMModel.fetch(model: modelName) { 42 | proxy = model.proxy() 43 | } else { 44 | proxy = Proxy(location: "", target: "https://api.openai.com") 45 | } 46 | 47 | guard let responseFuture = forward(request: request, to: proxy) else { 48 | return next.respond(to: request) 49 | } 50 | return responseFuture 51 | } 52 | 53 | func forward(request: HBRequest, to proxy: Proxy) -> EventLoopFuture? { 54 | guard request.uri.description.hasPrefix(proxy.location) else { return nil } 55 | let newURI = request.uri.description //.dropFirst(proxy.location.count) // TODO: not sure if this boilerplate code did anything we want it to do. Refactor after HummingBird 2.0 56 | guard newURI.first == nil || newURI.first == "/" else { return nil } 57 | 58 | do { 59 | // create request 60 | let ahcRequest = try request.ahcRequest(uri: String(newURI), host: proxy.target, eventLoop: request.eventLoop) 61 | 62 | request.logger.info("\(request.uri) -> \(ahcRequest.url)") 63 | 64 | // create response body streamer 65 | let streamer = HBByteBufferStreamer(eventLoop: request.eventLoop, maxSize: 2048 * 1024, maxStreamingBufferSize: 128 * 1024) 66 | // delegate for streaming bytebuffers from AsyncHTTPClient 67 | let delegate = StreamingResponseDelegate(on: request.eventLoop, streamer: streamer) 68 | // execute request 69 | _ = self.httpClient.execute( 70 | request: ahcRequest, 71 | delegate: delegate, 72 | eventLoop: .delegateAndChannel(on: request.eventLoop), 73 | logger: request.logger 74 | ) 75 | // when delegate receives header then signal completion 76 | return delegate.responsePromise.futureResult 77 | } catch { 78 | return request.failure(.badRequest) 79 | } 80 | } 81 | } 82 | 83 | extension HBRequest { 84 | /// create AsyncHTTPClient request from Hummingbird Request 85 | func ahcRequest(uri: String, host: String, eventLoop: EventLoop) throws -> HTTPClient.Request { 86 | 87 | var headers = self.headers 88 | headers.remove(name: "host") 89 | 90 | switch self.body { 91 | case .byteBuffer(let buffer): 92 | return try HTTPClient.Request( 93 | url: host + uri, 94 | method: self.method, 95 | headers: headers, 96 | body: buffer.map { .byteBuffer($0) } 97 | ) 98 | 99 | case .stream(let stream): 100 | let contentLength = self.headers["content-length"].first.map { Int($0) } ?? nil 101 | return try HTTPClient.Request( 102 | url: host + uri, 103 | method: self.method, 104 | headers: headers, 105 | body: .stream(length: contentLength) { writer in 106 | return stream.consumeAll(on: eventLoop) { byteBuffer in 107 | writer.write(.byteBuffer(byteBuffer)) 108 | } 109 | } 110 | ) 111 | } 112 | } 113 | } 114 | 115 | extension String { 116 | private func addSuffix(_ suffix: String) -> String { 117 | if hasSuffix(suffix) { 118 | return self 119 | } else { 120 | return self + suffix 121 | } 122 | } 123 | 124 | fileprivate func dropSuffix(_ suffix: String) -> String { 125 | if hasSuffix(suffix) { 126 | return String(self.dropLast(suffix.count)) 127 | } else { 128 | return self 129 | } 130 | } 131 | } 132 | -------------------------------------------------------------------------------- /Sources/App/Middleware/RateLimiterMiddleware.swift: -------------------------------------------------------------------------------- 1 | // 2 | // RateLimiterMiddleware.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/23/24. 6 | // 7 | 8 | import Foundation 9 | import FluentKit 10 | import Hummingbird 11 | 12 | /// The rate limiter is designed to mitigate the risk of excessive usage for a given key. 13 | /// 14 | /// The rate limiter is off by default. To enable the rate limiter, set the value of environment variable `enableRateLimiter` to 1 (default is 0) 15 | /// 16 | /// It allows developers to configure usage thresholds on two levels: 17 | /// 18 | /// 1. Maximum messages per minute: The default limit is 15 messages. 19 | /// 2. Maximum messages per hour: The default limit is 50 messages. 20 | /// 21 | /// User identification is based on their app account token. 22 | /// See for more info: 23 | /// https://developer.apple.com/documentation/storekit/product/purchaseoption/3749440-appaccounttoken 24 | /// 25 | /// The system also provides separate rate limits for users who don't have an app account token, 26 | /// catering to scenarios where the client application may not assign these tokens. 27 | /// The default rate limits for these users are set higher, at 60 messages per minute and 200 messages per hour. 28 | /// It's recommended to tailor these default values based on the proportion of users operating without an app account token. 29 | /// 30 | /// Users will be blocked permanently if the number of blocked messages exceeds environment variable `userPermanentBlock` 31 | /// for users with an app account token or `anonPermanentBlock` for all users without a token combined. 32 | /// The default value is 50 for `userPermanentBlock` and `anonPermanentBlock` 33 | /// 34 | /// To disable any or multiple rules, remove the environment variable or its value (i.e. `userPermanentBlock`, `userHourlyRateLimit`, `userMinuteRateLimit`, 35 | /// `anonPermanentBlock`, `anonHourlyRateLimit`, and `anonMinuteRateLimit`) 36 | /// 37 | /// Note that the current version of Pico AI Proxy doesn't store data persistently. Blocks and message request history will be reset 38 | /// whenever Pico AI Proxy is deployed. 39 | /// 40 | /// For more precise control, such as downgrading the model from GPT-4 to GPT-3.5 when a user exceeds a specific limit, 41 | /// you need to handle this logic on the client side 42 | struct RateLimiterMiddleware: HBAsyncMiddleware { 43 | 44 | func apply(to request: Hummingbird.HBRequest, next: Hummingbird.HBResponder) async throws -> Hummingbird.HBResponse { 45 | 46 | // 1. Fetch user 47 | let user = try request.authRequire(User.self) 48 | 49 | // 2. If rate limiter isn't enabled, we're done 50 | guard let enableUserRateLimiter = HBEnvironment().get("enableRateLimiter"), enableUserRateLimiter == "1" else { 51 | 52 | let userRequest = try? UserRequest(endpoint: request.uri.path, wasBlocked: false, userId: user.requireID()) 53 | try? await userRequest?.save(on: request.db) 54 | request.logger.info("User \(user.appAccountToken?.uuidString ?? "anon") request \(request.uri.path). Rate limiter is disabled") 55 | return try await next.respond(to: request) 56 | 57 | } 58 | 59 | // 3. Fetch user request history 60 | let lastMinute = try await user.numberOfRecentRequests(hours: 0, minutes: 1, db: request.db) 61 | let lastHour = try await user.numberOfRecentRequests(hours: 1, minutes: 0, db: request.db) 62 | let blockedRequests = try await user.numberOfBlockedRequests(db: request.db) 63 | let now = Date() 64 | var blockUntil: Date? = nil 65 | 66 | if let blockedDate = user.blockedUntil, blockedDate == Date.distantFuture { 67 | 68 | // 4. If user is blocked permanently, set blockUntil so the request gets logged 69 | blockUntil = blockedDate 70 | 71 | } else if let _ = user.appAccountToken { 72 | 73 | // 5. Check rate limit for users with an app account token 74 | 75 | if let limit = rateLimit(for: "userPermanentBlock"), blockedRequests >= limit { 76 | 77 | // 5.a User will be blocked permanently if the number of blocked requests exceeds userPermanentBlock 78 | blockUntil = Date.distantFuture 79 | 80 | } else if let limit = rateLimit(for: "userHourlyRateLimit"), lastHour >= limit { 81 | 82 | // 5.b User will be blocked for one hour if they exceeded hourly limit 83 | blockUntil = Calendar.current.date(byAdding: .hour, value: 1, to: now) 84 | 85 | } else if let limit = rateLimit(for: "userMinuteRateLimit"), lastHour >= limit { 86 | 87 | // 5.c User will be blocked for one 5 minutes if they exceeded minute limit 88 | blockUntil = Calendar.current.date(byAdding: .minute, value: 5, to: now) 89 | 90 | } 91 | 92 | } else { 93 | 94 | // 6. Check rate limit for users with an app account token (anonymous users) 95 | 96 | if let limit = rateLimit(for: "anonPermanentBlock"), blockedRequests >= limit { 97 | 98 | // 6.a All anonymous users will be blocked 99 | blockUntil = Date.distantFuture 100 | 101 | } else if let limit = rateLimit(for: "anonHourlyRateLimit"), lastHour >= limit { 102 | 103 | // 6.c User will be blocked for one hour if they exceeded hourly limit 104 | blockUntil = Calendar.current.date(byAdding: .hour, value: 1, to: now) 105 | 106 | } else if let limit = rateLimit(for: "anonMinuteRateLimit"), lastHour >= limit { 107 | 108 | // 6.d User will be blocked for one 5 minutes if they exceeded minute limit 109 | blockUntil = Calendar.current.date(byAdding: .minute, value: 5, to: now) 110 | 111 | } 112 | 113 | } 114 | 115 | request.logger.info("User \(user.appAccountToken?.uuidString ?? "anon"). Requests last minute: \(lastMinute). Requests last hour: \(lastHour)") 116 | 117 | // 7. If request is blocked, register request attempt 118 | 119 | if let blockUntil { 120 | 121 | // 8.a Update blockedUntil if new block is further into the future 122 | if let alreadyBlockedUntil = user.blockedUntil, blockUntil > alreadyBlockedUntil { 123 | user.blockedUntil = blockUntil 124 | try await user.save(on: request.db) 125 | } 126 | 127 | // 8.b Log request 128 | let userRequest = try UserRequest(endpoint: request.uri.path, wasBlocked: true, userId: user.requireID()) 129 | try await userRequest.save(on: request.db) 130 | request.logger.info("User \(user.appAccountToken?.uuidString ?? "anon") requested \(request.uri.path) and is blocked until \(user.blockedUntil?.description ?? "(no date)")") 131 | 132 | // 8.c Reject request by throwing too many requests error 133 | if let blockedUntil = user.blockedUntil, 134 | blockedUntil != Date.distantFuture, 135 | let minutes = Calendar.current.dateComponents([.minute], from: now, to: blockedUntil).minute { 136 | throw HBHTTPError(.tooManyRequests, message: "Rate limit reached. Try again in \(minutes) minutes") 137 | } else { 138 | throw HBHTTPError(.tooManyRequests, message: "Rate limit reached") 139 | } 140 | 141 | } else { 142 | 143 | // 9. Request was allowed. Log request and forward request to the next step 144 | let userRequest = try UserRequest(endpoint: request.uri.path, wasBlocked: false, userId: user.requireID()) 145 | try await userRequest.save(on: request.db) 146 | request.logger.info("User \(user.appAccountToken?.uuidString ?? "anon") requested \(request.uri.path)") 147 | 148 | return try await next.respond(to: request) 149 | } 150 | } 151 | 152 | private func rateLimit(for key: String) -> Int? { 153 | guard let userHourlyRateLimiter = HBEnvironment().get(key), 154 | let limit = Int(userHourlyRateLimiter), 155 | limit > 0 else { 156 | // This limit is not set, we're done 157 | return nil 158 | } 159 | 160 | return limit 161 | } 162 | } 163 | -------------------------------------------------------------------------------- /Sources/App/Middleware/StreamingResponseDelegate.swift: -------------------------------------------------------------------------------- 1 | // 2 | // File.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 12/19/23. 6 | // 7 | 8 | import AsyncHTTPClient 9 | import Hummingbird 10 | import HummingbirdCore 11 | import NIOCore 12 | import NIOHTTP1 13 | 14 | final class StreamingResponseDelegate: HTTPClientResponseDelegate { 15 | typealias Response = HBResponse 16 | 17 | enum State { 18 | case idle 19 | case head(HTTPResponseHead) 20 | case error(Error) 21 | } 22 | 23 | let streamer: HBByteBufferStreamer 24 | let responsePromise: EventLoopPromise 25 | let eventLoop: EventLoop 26 | var state: State 27 | 28 | init(on eventLoop: EventLoop, streamer: HBByteBufferStreamer) { 29 | self.eventLoop = eventLoop 30 | self.streamer = streamer 31 | self.responsePromise = eventLoop.makePromise() 32 | self.state = .idle 33 | } 34 | 35 | func didReceiveHead(task: HTTPClient.Task, _ head: HTTPResponseHead) -> EventLoopFuture { 36 | switch self.state { 37 | case .idle: 38 | let response = Response(status: head.status, headers: head.headers, body: .stream(self.streamer)) 39 | self.responsePromise.succeed(response) 40 | self.state = .head(head) 41 | case .error: 42 | break 43 | default: 44 | preconditionFailure("Unexpected state \(self.state)") 45 | } 46 | return self.eventLoop.makeSucceededVoidFuture() 47 | } 48 | 49 | func didReceiveBodyPart(task: HTTPClient.Task, _ buffer: ByteBuffer) -> EventLoopFuture { 50 | switch self.state { 51 | case .head: 52 | return self.streamer.feed(buffer: buffer) 53 | case .error: 54 | break 55 | default: 56 | preconditionFailure("Unexpected state \(self.state)") 57 | } 58 | return self.eventLoop.makeSucceededVoidFuture() 59 | } 60 | 61 | func didFinishRequest(task: HTTPClient.Task) throws -> HBResponse { 62 | switch self.state { 63 | case .head(let head): 64 | self.state = .idle 65 | self.streamer.feed(.end) 66 | return .init(status: head.status, headers: head.headers, body: .stream(self.streamer)) 67 | case .error(let error): 68 | print(error.localizedDescription) 69 | throw error 70 | default: 71 | preconditionFailure("Unexpected state \(self.state)") 72 | } 73 | } 74 | 75 | func didReceiveError(task: HTTPClient.Task, _ error: Error) { 76 | self.streamer.feed(.error(error)) 77 | self.responsePromise.fail(error) 78 | self.state = .error(error) 79 | } 80 | } 81 | -------------------------------------------------------------------------------- /Sources/App/Models/Chat.swift: -------------------------------------------------------------------------------- 1 | // 2 | // Chat.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 4/6/24. 6 | // 7 | 8 | import Foundation 9 | 10 | public struct Chat: Codable { 11 | 12 | /// Name of the model to use. 13 | /// E.g. `gpt-3.5-turbo` 14 | public let model: String 15 | 16 | // TODO: add rest of model 17 | } 18 | -------------------------------------------------------------------------------- /Sources/App/Models/Config.swift: -------------------------------------------------------------------------------- 1 | // 2 | // Config.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/25/24. 6 | // 7 | 8 | import Foundation 9 | /* 10 | struct Config: Codable { 11 | 12 | let proxies: [ProxyConfig] 13 | 14 | let spendingLimits: [SpendingLimit]? 15 | } 16 | 17 | struct ProxyConfig: Codable { 18 | 19 | /// E.g. ["gpt-4-0125-preview", "gpt-3.5-turbo", "text-embedding-3-small"] 20 | let models: [LLMModel] 21 | 22 | /// The label of the environment variable containing the API key value. 23 | /// Warning: do not store the API key itself here. The API key should be stored in 24 | /// an environment variable named the keyLabel value 25 | /// E.g. OpenAILabel 26 | let keyLabel: String 27 | 28 | /// The label of the environment variable containing the organization key value. 29 | /// orgLabel should be nil for APIs that don't require a organization 30 | /// Warning: do not store the organization key itself here. The organization key should be stored in 31 | /// an environment variable named the orgLabel value 32 | /// E.g. OpenAIKey 33 | let orgLabel: String? 34 | 35 | /// E.g. https://api.openai.com/v1/chat/completions 36 | let endpoint: URL 37 | 38 | } 39 | 40 | struct LLMModel: Codable { 41 | 42 | /// E.g. "gpt-4-0125-preview" 43 | let model: String 44 | 45 | /// Price in USD per 1K tokens. For example, for gpt-3.5-turbo the value is 0.0015 46 | let price: Float 47 | } 48 | 49 | /// Type of APIs supported 50 | enum LLMAPI: Int, Codable { 51 | case openAI 52 | } 53 | 54 | struct SpendingLimit: Codable { 55 | 56 | // E.g. "pico.subscription.lite" 57 | let productIdLabel: String 58 | 59 | // E.g. 4 60 | let monthlyLimitLabel: String 61 | } 62 | */ 63 | -------------------------------------------------------------------------------- /Sources/App/Models/LLMModel.swift: -------------------------------------------------------------------------------- 1 | // 2 | // LLMModel.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 4/9/24. 6 | // 7 | 8 | import Foundation 9 | 10 | struct LLMModel: Codable { 11 | 12 | /// E.g. `gpt-3.5` or `gpt-4` 13 | let name: String 14 | 15 | /// E.g. `https://api.openai.com/v1/chat/completions` or `https://api.openai.com/v1/embeddings` 16 | let endpoint: String 17 | 18 | let provider: LLMProvider 19 | } 20 | 21 | extension LLMModel { 22 | 23 | func proxy() -> HBProxyServerMiddleware.Proxy { 24 | HBProxyServerMiddleware.Proxy(location: endpoint, target: provider.host) 25 | } 26 | 27 | static func fetch(model: String) -> LLMModel? { 28 | return models.filter({ $0.name == model }).first 29 | } 30 | 31 | static func fetchModel(from chat: Data) -> LLMModel? { 32 | let decoder = JSONDecoder() 33 | if let chat = try? decoder.decode(Chat.self, from: chat) { 34 | return fetch(model: chat.model) 35 | } else { 36 | return nil 37 | } 38 | } 39 | } 40 | -------------------------------------------------------------------------------- /Sources/App/Models/LLMProvider.swift: -------------------------------------------------------------------------------- 1 | // 2 | // LLMProvider.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 4/8/24. 6 | // 7 | 8 | import Foundation 9 | import Hummingbird 10 | 11 | struct LLMProvider: Codable { 12 | 13 | /// The name of the provider, e.g. `OpenAI` 14 | let name: String 15 | 16 | /// E.g. `api.openai.com` 17 | let host: String 18 | 19 | /// The name of the environment key that stores the API key 20 | /// E.g. `OpenAI-APIKey` 21 | let apiEnvKey: String 22 | 23 | /// The HTTP header key for the API key 24 | /// E.g. `Authorization` 25 | let apiHeaderKey: String 26 | 27 | /// Adds `Bearer` to API key header value if true 28 | let bearer: Bool 29 | 30 | /// The name of the environment key that stores the API key 31 | /// E.g. `OpenAI-Organization` 32 | let orgEnvKey: String? 33 | 34 | /// The HTTP header key for the organization key 35 | /// E.g. `OpenAI-Organization` 36 | let orgHeaderKey: String? 37 | 38 | let additionalHeaders: [String: String]? 39 | } 40 | 41 | extension LLMProvider { 42 | 43 | 44 | /// Updates headers with API key, org and/or other required headers for this provider 45 | /// - Parameter headers: headers 46 | func setHeaders(headers: inout HTTPHeaders) throws { 47 | 48 | // 1. Set API key 49 | guard let apiKey = HBEnvironment().get(apiEnvKey) else { 50 | throw HBHTTPError(.internalServerError, message: "Environment API Key \(apiEnvKey) not set") 51 | } 52 | headers.replaceOrAdd(name: apiHeaderKey, value: (bearer ? "Bearer " + apiKey : apiKey)) 53 | 54 | if let orgEnvKey, let org = HBEnvironment().get(orgEnvKey) { 55 | headers.replaceOrAdd(name: orgEnvKey, value: org) 56 | } 57 | 58 | if let additionalHeaders { 59 | for header in additionalHeaders { 60 | headers.replaceOrAdd(name: header.key, value: header.value) 61 | } 62 | } 63 | } 64 | } 65 | 66 | -------------------------------------------------------------------------------- /Sources/App/Models/User.swift: -------------------------------------------------------------------------------- 1 | // 2 | // User.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/21/24. 6 | // 7 | 8 | import FluentKit 9 | import Foundation 10 | import Hummingbird 11 | import HummingbirdAuth 12 | import AppStoreServerLibrary 13 | 14 | final class User: Model, HBAuthenticatable { 15 | 16 | static let schema = "user" 17 | 18 | @ID(key: .id) 19 | var id: UUID? 20 | 21 | /// app account Id to identify individual users 22 | /// See https://developer.apple.com/documentation/appstoreserverapi/appaccounttoken 23 | @Field(key: "appAccountToken") 24 | var appAccountToken: UUID? 25 | 26 | /// Can be Sandbox, Production, Xcode, LocalTesting 27 | @Field(key: "environment") 28 | var environment: String 29 | 30 | /// The user's active subscription 31 | @Field(key: "productId") 32 | var productId: String 33 | 34 | /// Subscription status 35 | /// See https://developer.apple.com/documentation/appstoreserverapi/status 36 | @Field(key: "subscriptionStatus") 37 | var subscriptionStatus: Int32 38 | 39 | /// If not nil contains the date until user is rate limited 40 | @Field(key: "blockedUntil") 41 | var blockedUntil: Date? 42 | 43 | /// JWT token 44 | @Field(key: "jwtToken") 45 | var jwtToken: String? 46 | 47 | /// Overview of requests this user had made 48 | @Children(for: \.$user) 49 | var messages: [UserRequest] 50 | 51 | internal init() {} 52 | 53 | internal init(id: UUID? = nil, appAccountToken: UUID?, environment: String, productId: String, status: AppStoreServerLibrary.Status, token: String? = nil) { 54 | self.id = id 55 | self.appAccountToken = appAccountToken 56 | self.environment = environment 57 | self.productId = productId 58 | self.subscriptionStatus = status.rawValue 59 | self.jwtToken = token 60 | } 61 | } 62 | 63 | 64 | extension User { 65 | 66 | /// Fetches all messages belonging to user that were made in the last hours and/or minutes 67 | /// - Parameters: 68 | /// - hours: Filters messages in the last x hours 69 | /// - minutes: Filters messages in the last x minutes 70 | /// - userId: id of the user 71 | /// - db: database in the request 72 | /// - Returns: array of messages in the timeframe provided 73 | func numberOfRecentRequests(hours: Int, minutes: Int, db: Database) async throws -> Int { 74 | 75 | // 1. Get the current date and time 76 | let now = Date() 77 | 78 | // 2. Calculate the start time 79 | let hours = Calendar.current.date(byAdding: .hour, value: -hours, to: now) ?? now 80 | let startTime = Calendar.current.date(byAdding: .minute, value: -minutes, to: hours) ?? hours 81 | 82 | // 3. Query messages where the date is greater than one hour ago and the user ID matches 83 | return try await UserRequest.query(on: db) 84 | .filter(\.$user.$id == self.requireID()) 85 | .filter(\.$date > startTime) 86 | .all() 87 | .count 88 | } 89 | 90 | func numberOfBlockedRequests(db: Database) async throws -> Int { 91 | return try await UserRequest.query(on: db) 92 | .filter(\.$user.$id == self.requireID()) 93 | .filter(\.$wasBlocked == true) 94 | .all() 95 | .count 96 | } 97 | } 98 | -------------------------------------------------------------------------------- /Sources/App/Models/UserMigration.swift: -------------------------------------------------------------------------------- 1 | // 2 | // UserMigration.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/21/24. 6 | // 7 | 8 | import Foundation 9 | import FluentKit 10 | 11 | struct UserMigration: AsyncMigration { 12 | 13 | func prepare(on database: FluentKit.Database) async throws { 14 | try await database.schema("user") 15 | .id() 16 | .field("appAccountToken", .uuid) 17 | .field("environment", .string, .required) 18 | .field("productId", .string, .required) 19 | .field("subscriptionStatus", .int32, .required) 20 | .field("blockedUntil", .date) 21 | .field("jwtToken", .string) 22 | .unique(on: "appAccountToken") 23 | .create() 24 | } 25 | 26 | func revert(on database: FluentKit.Database) async throws { 27 | try await database.schema("user").delete() 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /Sources/App/Models/UserRequest.swift: -------------------------------------------------------------------------------- 1 | // 2 | // UserRequest.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/23/24. 6 | // 7 | 8 | import Foundation 9 | import FluentKit 10 | 11 | /// Every request will be logged for the rate limiter middleware 12 | final class UserRequest: Model { 13 | 14 | static let schema = "userrequest" 15 | 16 | @ID(key: .id) 17 | var id: UUID? 18 | 19 | /// Date and time of the request 20 | @Field(key: "date") 21 | var date: Date 22 | 23 | /// Endpoint called in this request 24 | @Field(key: "endpoint") 25 | var endpoint: String 26 | 27 | /// If true, this request was blocked by the rate limiter 28 | @Field(key: "wasBlocked") 29 | var wasBlocked: Bool 30 | 31 | /// Model called. Currently not used 32 | @Field(key: "model") 33 | var model: String? 34 | 35 | /// Length of the request in characters (not tokens). Currently not used 36 | @Field(key: "requestLen") 37 | var requestLen: Int? 38 | 39 | /// Length of the response in characters (not tokens). Currently not used 40 | @Field(key: "responseLen") 41 | var responseLen: Int? 42 | 43 | /// User who made the request 44 | @Parent(key: "user") 45 | var user: User 46 | 47 | internal init() {} 48 | 49 | internal init(endpoint: String, wasBlocked: Bool, model: String? = nil, requestLen: Int? = nil, responseLen: Int? = nil, userId: UUID) { 50 | self.id = id 51 | self.date = Date() 52 | self.endpoint = endpoint 53 | self.wasBlocked = wasBlocked 54 | self.model = model 55 | self.requestLen = requestLen 56 | self.responseLen = responseLen 57 | self.$user.id = userId 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /Sources/App/Models/UserRequestMigration.swift: -------------------------------------------------------------------------------- 1 | // 2 | // File.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 1/24/24. 6 | // 7 | 8 | import Foundation 9 | import FluentKit 10 | 11 | struct UserRequestMigration: AsyncMigration { 12 | 13 | func prepare(on database: FluentKit.Database) async throws { 14 | try await database.schema("userrequest") 15 | .id() 16 | .field("date", .datetime, .required) 17 | .field("endpoint", .string) 18 | .field("model", .string) 19 | .field("wasBlocked", .bool) 20 | .field("requestLen", .int) 21 | .field("responseLen", .int) 22 | .field("token", .string) 23 | .field("user", .uuid, .required, .references("user", "id")) 24 | .create() 25 | } 26 | 27 | func revert(on database: FluentKit.Database) async throws { 28 | try await database.schema("user").delete() 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /Sources/App/Providers.swift: -------------------------------------------------------------------------------- 1 | // 2 | // File.swift 3 | // 4 | // 5 | // Created by Ronald Mannak on 4/8/24. 6 | // 7 | 8 | import Foundation 9 | 10 | extension LLMProvider { 11 | 12 | static var providers = [openAI, anthropic] 13 | 14 | static var openAI = LLMProvider( 15 | name: "OpenAI", 16 | host: "https://api.openai.com", 17 | apiEnvKey: "OpenAI-APIKey", 18 | apiHeaderKey: "Authorization", 19 | bearer: true, 20 | orgEnvKey: "OpenAI-Organization", 21 | orgHeaderKey: "OpenAI-Organization", 22 | additionalHeaders: nil 23 | ) 24 | 25 | static var anthropic = LLMProvider( 26 | name: "Anthropic", 27 | host: "https://api.anthropic.com", 28 | apiEnvKey: "Anthropic-APIKey", 29 | apiHeaderKey: "x-api-key", 30 | bearer: false, 31 | orgEnvKey: nil, 32 | orgHeaderKey: nil, 33 | additionalHeaders: ["anthropic-version": "2023-06-01"] 34 | ) 35 | 36 | static var groq = LLMProvider( 37 | name: "Groq", 38 | host: "https://api.groq.com", 39 | apiEnvKey: "Groq-APIKey", 40 | apiHeaderKey: "Authorization", 41 | bearer: true, 42 | orgEnvKey: nil, 43 | orgHeaderKey: nil, 44 | additionalHeaders: nil) 45 | } 46 | 47 | extension LLMModel { 48 | static var models = [LLMModel]() 49 | 50 | private static let gptModels = ["gpt-4o", "gpt-4o-2024-05-13", "gpt-4-turbo", "gpt-4-turbo-2024-04-09", "gpt-4-0125-preview", "gpt-4-turbo-preview", "gpt-4-1106-preview", "gpt-4-vision-preview", "gpt-4-1106-vision-preview", "gpt-4", "gpt-4-0613", "gpt-4-32k", "gpt-4-32k-0613", "gpt-3.5-turbo-0125", "gpt-3.5-turbo", "gpt-3.5-turbo-1106", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-0613", "gpt-3.5-turbo-16k-0613"] 51 | private static let claudeModels = ["claude-3-opus-20240229", "claude-3-sonnet-20240229", "claude-3-haiku-20240307"] 52 | private static let groqModels = ["llama3-8b-8192", "llama3-70b-8192", "mixtral-8x7b-32768", "gemma-7b-it"] 53 | 54 | static func load() { 55 | for model in gptModels { 56 | models.append(LLMModel(name: model, endpoint: "", provider: LLMProvider.openAI)) // leaving endpoint empty so all calls will be forwarded to original path /v1/chat/completions 57 | } 58 | for model in claudeModels { 59 | models.append(LLMModel(name: model, endpoint: "/v1/messages", provider: LLMProvider.anthropic)) 60 | } 61 | for model in groqModels { 62 | models.append(LLMModel(name: model, endpoint: "/openai/v1/chat/completions", provider: LLMProvider.groq)) 63 | } 64 | } 65 | } 66 | -------------------------------------------------------------------------------- /Sources/App/Resources/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PicoMLX/PicoAIProxy/5ecaaf24675bf0b9a51aa14323b4b31138cd8f86/Sources/App/Resources/.DS_Store -------------------------------------------------------------------------------- /Sources/App/Resources/AppleComputerRootCertificate.cer: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PicoMLX/PicoAIProxy/5ecaaf24675bf0b9a51aa14323b4b31138cd8f86/Sources/App/Resources/AppleComputerRootCertificate.cer -------------------------------------------------------------------------------- /Sources/App/Resources/AppleIncRootCertificate.cer: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PicoMLX/PicoAIProxy/5ecaaf24675bf0b9a51aa14323b4b31138cd8f86/Sources/App/Resources/AppleIncRootCertificate.cer -------------------------------------------------------------------------------- /Sources/App/Resources/AppleRootCA-G2.cer: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PicoMLX/PicoAIProxy/5ecaaf24675bf0b9a51aa14323b4b31138cd8f86/Sources/App/Resources/AppleRootCA-G2.cer -------------------------------------------------------------------------------- /Sources/App/Resources/AppleRootCA-G3.cer: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PicoMLX/PicoAIProxy/5ecaaf24675bf0b9a51aa14323b4b31138cd8f86/Sources/App/Resources/AppleRootCA-G3.cer -------------------------------------------------------------------------------- /Sources/App/Resources/Providers.plist: -------------------------------------------------------------------------------- 1 | ( 2 | {}, 3 | ) -------------------------------------------------------------------------------- /Tests/AppTests/AppTests.swift: -------------------------------------------------------------------------------- 1 | @testable import App 2 | import Hummingbird 3 | import HummingbirdXCT 4 | import XCTest 5 | 6 | final class AppTests: XCTestCase { 7 | func createProxy(port: Int) async throws -> HBApplication { 8 | struct Arguments: AppArguments { 9 | var location: String 10 | var target: String 11 | } 12 | return try await createProxy(args: Arguments(location: "", target: "http://localhost:\(port)")) 13 | } 14 | 15 | func createProxy(args: AppArguments) async throws -> HBApplication { 16 | let app = HBApplication(testing: .live) 17 | try await app.configure(args) 18 | return app 19 | } 20 | 21 | func randomBuffer(size: Int) -> ByteBuffer { 22 | var data = [UInt8](repeating: 0, count: size) 23 | data = data.map { _ in UInt8.random(in: 0...255) } 24 | return ByteBufferAllocator().buffer(bytes: data) 25 | } 26 | 27 | func testSimple() async throws { 28 | let app = HBApplication(configuration: .init(address: .hostname(port: 0))) 29 | app.router.get("hello") { _ in 30 | return "Hello" 31 | } 32 | try app.start() 33 | defer { app.stop() } 34 | 35 | let proxy = try await createProxy(port: app.server.port!) 36 | try proxy.XCTStart() 37 | defer { proxy.XCTStop() } 38 | 39 | try proxy.XCTExecute(uri: "/hello", method: .GET) { response in 40 | let body = try XCTUnwrap(response.body) 41 | XCTAssertEqual(String(buffer: body), "Hello") 42 | } 43 | } 44 | 45 | func testEchoBody() async throws { 46 | let app = HBApplication(configuration: .init(address: .hostname(port: 0))) 47 | app.router.post("echo") { request in 48 | return request.body.buffer 49 | } 50 | try app.start() 51 | defer { app.stop() } 52 | 53 | let proxy = try await createProxy(port: app.server.port!) 54 | try proxy.XCTStart() 55 | defer { proxy.XCTStop() } 56 | 57 | let bodyString = "This is a test body" 58 | try proxy.XCTExecute(uri: "/echo", method: .POST, body: ByteBuffer(string: bodyString)) { response in 59 | let body = try XCTUnwrap(response.body) 60 | XCTAssertEqual(String(buffer: body), bodyString) 61 | } 62 | } 63 | 64 | func testLargeBody() async throws { 65 | let app = HBApplication(configuration: .init(address: .hostname(port: 0))) 66 | app.router.post("echo") { request in 67 | return request.body.buffer 68 | } 69 | try app.start() 70 | defer { app.stop() } 71 | 72 | let proxy = try await createProxy(port: app.server.port!) 73 | try proxy.XCTStart() 74 | defer { proxy.XCTStop() } 75 | 76 | let buffer = randomBuffer(size: 1024 * 1500) 77 | try proxy.XCTExecute(uri: "/echo", method: .POST, body: buffer) { response in 78 | XCTAssertEqual(response.body, buffer) 79 | } 80 | } 81 | } 82 | --------------------------------------------------------------------------------