2023-12-16 19:30:20 -07:00
/** This whole module kinda sucks */
2023-04-14 19:21:04 -06:00
import fs from "fs" ;
2024-02-04 13:04:46 -07:00
import express , { Router , Request , Response } from "express" ;
2023-04-08 07:50:03 -06:00
import showdown from "showdown" ;
2023-12-16 19:30:20 -07:00
import { config } from "./config" ;
import { buildInfo , ServiceInfo } from "./service-info" ;
2023-11-13 22:41:19 -07:00
import { getLastNImages } from "./shared/file-storage/image-history" ;
2023-12-16 19:30:20 -07:00
import { keyPool } from "./shared/key-management" ;
import { MODEL_FAMILY_SERVICE , ModelFamily } from "./shared/models" ;
2024-02-04 13:04:46 -07:00
import { withSession } from "./shared/with-session" ;
import { checkCsrfToken , injectCsrfToken } from "./shared/inject-csrf" ;
2023-05-09 17:11:57 -06:00
2023-08-30 18:20:29 -06:00
const INFO_PAGE_TTL = 2000 ;
2023-12-16 19:30:20 -07:00
const MODEL_FAMILY_FRIENDLY_NAME : { [ f in ModelFamily ] : string } = {
2024-03-09 12:03:50 -07:00
turbo : "GPT-3.5 Turbo" ,
gpt4 : "GPT-4" ,
2023-12-16 19:30:20 -07:00
"gpt4-32k" : "GPT-4 32k" ,
"gpt4-turbo" : "GPT-4 Turbo" ,
"dall-e" : "DALL-E" ,
2024-03-09 12:03:50 -07:00
claude : "Claude (Sonnet)" ,
2024-03-04 13:08:59 -07:00
"claude-opus" : "Claude (Opus)" ,
2023-12-16 19:30:20 -07:00
"gemini-pro" : "Gemini Pro" ,
2023-12-25 11:33:16 -07:00
"mistral-tiny" : "Mistral 7B" ,
2024-02-26 17:12:08 -07:00
"mistral-small" : "Mixtral Small" , // Originally 8x7B, but that now refers to the older open-weight version. Mixtral Small is a newer closed-weight update to the 8x7B model.
"mistral-medium" : "Mistral Medium" ,
"mistral-large" : "Mistral Large" ,
2024-03-04 15:25:06 -07:00
"aws-claude" : "AWS Claude (Sonnet)" ,
2023-12-16 19:30:20 -07:00
"azure-turbo" : "Azure GPT-3.5 Turbo" ,
"azure-gpt4" : "Azure GPT-4" ,
"azure-gpt4-32k" : "Azure GPT-4 32k" ,
"azure-gpt4-turbo" : "Azure GPT-4 Turbo" ,
2024-03-09 12:03:50 -07:00
"azure-dall-e" : "Azure DALL-E" ,
2023-08-30 18:20:29 -06:00
} ;
2023-12-16 19:30:20 -07:00
const converter = new showdown . Converter ( ) ;
const customGreeting = fs . existsSync ( "greeting.md" )
? ` \ n## Server Greeting \ n ${ fs . readFileSync ( "greeting.md" , "utf8" ) } `
: "" ;
let infoPageHtml : string | undefined ;
let infoPageLastUpdated = 0 ;
2023-08-30 18:20:29 -06:00
2023-04-08 04:43:28 -06:00
export const handleInfoPage = ( req : Request , res : Response ) = > {
2023-05-09 17:11:57 -06:00
if ( infoPageLastUpdated + INFO_PAGE_TTL > Date . now ( ) ) {
2023-12-03 21:21:18 -07:00
return res . send ( infoPageHtml ) ;
2023-05-09 17:11:57 -06:00
}
2023-06-08 12:50:04 -06:00
const baseUrl =
process . env . SPACE_ID && ! req . get ( "host" ) ? . includes ( "hf.space" )
? getExternalUrlForHuggingfaceSpaceId ( process . env . SPACE_ID )
: req . protocol + "://" + req . get ( "host" ) ;
2023-06-06 19:47:45 -06:00
2024-02-26 17:20:34 -07:00
const info = buildInfo ( baseUrl + config . proxyEndpointRoute ) ;
2023-12-16 19:30:20 -07:00
infoPageHtml = renderPage ( info ) ;
2023-11-15 16:12:07 -07:00
infoPageLastUpdated = Date . now ( ) ;
res . send ( infoPageHtml ) ;
2023-04-08 04:43:28 -06:00
} ;
2023-12-16 19:30:20 -07:00
export function renderPage ( info : ServiceInfo ) {
2023-05-15 15:47:30 -06:00
const title = getServerTitle ( ) ;
2023-12-16 19:30:20 -07:00
const headerHtml = buildInfoPageHeader ( info ) ;
2023-04-08 04:43:28 -06:00
2023-11-15 16:12:07 -07:00
return ` <!DOCTYPE html>
2023-04-08 04:43:28 -06:00
< html lang = "en" >
< head >
< meta charset = "utf-8" / >
2023-05-15 16:05:04 -06:00
< meta name = "robots" content = "noindex" / >
2023-04-17 15:19:33 -06:00
< title > $ { title } < / title >
2024-01-25 15:24:11 -07:00
< style >
body {
font - family : sans - serif ;
background - color : # f0f0f0 ;
padding : 1em ;
}
@media ( prefers - color - scheme : dark ) {
body {
background - color : # 222 ;
color : # eee ;
}
a :link , a :visited {
color : # bbe ;
}
}
< / style >
2023-04-08 04:43:28 -06:00
< / head >
2024-01-25 15:24:11 -07:00
< body >
2023-05-09 17:11:57 -06:00
$ { headerHtml }
2023-04-08 07:50:03 -06:00
< hr / >
< h2 > Service Info < / h2 >
2023-04-08 04:43:28 -06:00
< pre > $ { JSON . stringify ( info , null , 2 ) } < / pre >
2023-09-02 14:09:11 -06:00
$ { getSelfServiceLinks ( ) }
2023-04-08 04:43:28 -06:00
< / body >
< / html > ` ;
}
2023-04-14 19:21:04 -06:00
/ * *
* If the server operator provides a ` greeting.md ` file , it will be included in
* the rendered info page .
* * /
2023-12-16 19:30:20 -07:00
function buildInfoPageHeader ( info : ServiceInfo ) {
const title = getServerTitle ( ) ;
2023-05-16 00:19:13 -06:00
// TODO: use some templating engine instead of this mess
2023-12-16 19:30:20 -07:00
let infoBody = ` # ${ title } ` ;
2023-04-14 19:21:04 -06:00
if ( config . promptLogging ) {
2023-12-03 21:21:18 -07:00
infoBody += ` \ n## Prompt Logging Enabled
This proxy keeps full logs of all prompts and AI responses . Prompt logs are anonymous and do not contain IP addresses or timestamps .
2023-04-14 19:21:04 -06:00
2023-12-03 21:21:18 -07:00
[ You can see the type of data logged here , along with the rest of the code . ] ( https : //gitgud.io/khanon/oai-reverse-proxy/-/blob/main/src/shared/prompt-logging/index.ts).
2023-04-14 19:21:04 -06:00
2023-04-15 02:11:00 -06:00
* * If you are uncomfortable with this , don ' t send prompts to this proxy ! * * ` ;
2023-04-14 19:21:04 -06:00
}
2023-05-09 17:11:57 -06:00
2023-11-15 16:12:07 -07:00
if ( config . staticServiceInfo ) {
return converter . makeHtml ( infoBody + customGreeting ) ;
}
2023-08-09 17:20:19 -06:00
const waits : string [ ] = [ ] ;
2023-08-29 16:56:54 -06:00
2023-12-16 19:30:20 -07:00
for ( const modelFamily of config . allowedModelFamilies ) {
const service = MODEL_FAMILY_SERVICE [ modelFamily ] ;
2023-08-29 16:56:54 -06:00
2023-12-16 19:30:20 -07:00
const hasKeys = keyPool . list ( ) . some ( ( k ) = > {
return k . service === service && k . modelFamilies . includes ( modelFamily ) ;
} ) ;
2023-11-14 00:28:59 -07:00
2023-12-16 19:30:20 -07:00
const wait = info [ modelFamily ] ? . estimatedQueueTime ;
if ( hasKeys && wait ) {
2024-03-09 12:03:50 -07:00
waits . push (
` ** ${ MODEL_FAMILY_FRIENDLY_NAME [ modelFamily ] || modelFamily } **: ${ wait } `
) ;
2023-11-14 00:28:59 -07:00
}
2023-05-09 17:11:57 -06:00
}
2023-08-09 17:20:19 -06:00
infoBody += "\n\n" + waits . join ( " / " ) ;
2023-11-15 16:12:07 -07:00
infoBody += customGreeting ;
2023-11-13 22:41:19 -07:00
infoBody += buildRecentImageSection ( ) ;
2023-04-14 19:21:04 -06:00
return converter . makeHtml ( infoBody ) ;
}
2023-05-09 17:11:57 -06:00
2023-09-02 14:09:11 -06:00
function getSelfServiceLinks() {
if ( config . gatekeeper !== "user_token" ) return "" ;
2023-09-02 14:23:30 -06:00
return ` <footer style="font-size: 0.8em;"><hr /><a target="_blank" href="/user/lookup">Check your user token info</a></footer> ` ;
2023-09-02 14:09:11 -06:00
}
2023-05-15 15:47:30 -06:00
function getServerTitle() {
// Use manually set title if available
if ( process . env . SERVER_TITLE ) {
return process . env . SERVER_TITLE ;
}
// Huggingface
if ( process . env . SPACE_ID ) {
return ` ${ process . env . SPACE_AUTHOR_NAME } / ${ process . env . SPACE_TITLE } ` ;
}
// Render
if ( process . env . RENDER ) {
return ` Render / ${ process . env . RENDER_SERVICE_NAME } ` ;
}
return "OAI Reverse Proxy" ;
}
2023-06-06 19:47:45 -06:00
2023-11-13 22:41:19 -07:00
function buildRecentImageSection() {
2024-03-09 12:03:50 -07:00
const dalleModels : ModelFamily [ ] = [ "azure-dall-e" , "dall-e" ] ;
2023-11-13 22:41:19 -07:00
if (
2024-03-09 12:03:50 -07:00
! config . showRecentImages ||
dalleModels . every ( ( f ) = > ! config . allowedModelFamilies . includes ( f ) )
2023-11-13 22:41:19 -07:00
) {
return "" ;
}
let html = ` <h2>Recent DALL-E Generations</h2> ` ;
const recentImages = getLastNImages ( 12 ) . reverse ( ) ;
if ( recentImages . length === 0 ) {
html += ` <p>No images yet.</p> ` ;
return html ;
}
html += ` <div style="display: flex; flex-wrap: wrap;" id="recent-images"> ` ;
for ( const { url , prompt } of recentImages ) {
const thumbUrl = url . replace ( /\.png$/ , "_t.jpg" ) ;
const escapedPrompt = escapeHtml ( prompt ) ;
html += ` <div style="margin: 0.5em;" class="recent-image">
< a href = "${url}" target = "_blank" > < img src = "${thumbUrl}" title = "${escapedPrompt}" alt = "${escapedPrompt}" style = "max-width: 150px; max-height: 150px;" / > < / a >
< / div > ` ;
}
html += ` </div> ` ;
2024-03-10 13:53:11 -06:00
html += ` <p style="clear: both; text-align: center;"><a href="/user/image-history">View all recent images</a></p> `
2023-11-13 22:41:19 -07:00
return html ;
}
function escapeHtml ( unsafe : string ) {
return unsafe
2023-11-15 16:12:07 -07:00
. replace ( /&/g , "&" )
. replace ( /</g , "<" )
. replace ( />/g , ">" )
. replace ( /"/g , """ )
. replace ( /'/g , "'" ) ;
2023-11-13 22:41:19 -07:00
}
2023-06-06 19:47:45 -06:00
function getExternalUrlForHuggingfaceSpaceId ( spaceId : string ) {
try {
const [ username , spacename ] = spaceId . split ( "/" ) ;
return ` https:// ${ username } - ${ spacename . replace ( /_/g , "-" ) } .hf.space ` ;
} catch ( e ) {
return "" ;
}
}
2024-02-04 13:04:46 -07:00
2024-03-09 12:03:50 -07:00
function checkIfUnlocked (
req : Request ,
res : Response ,
next : express.NextFunction
) {
2024-02-04 13:04:46 -07:00
if ( config . serviceInfoPassword ? . length && ! req . session ? . unlocked ) {
return res . redirect ( "/unlock-info" ) ;
}
next ( ) ;
}
const infoPageRouter = Router ( ) ;
if ( config . serviceInfoPassword ? . length ) {
infoPageRouter . use (
express . json ( { limit : "1mb" } ) ,
express . urlencoded ( { extended : true , limit : "1mb" } )
) ;
infoPageRouter . use ( withSession ) ;
infoPageRouter . use ( injectCsrfToken , checkCsrfToken ) ;
2024-03-09 12:03:50 -07:00
infoPageRouter . post ( "/unlock-info" , ( req , res ) = > {
if ( req . body . password !== config . serviceInfoPassword ) {
return res . status ( 403 ) . send ( "Incorrect password" ) ;
}
req . session ! . unlocked = true ;
res . redirect ( "/" ) ;
} ) ;
2024-02-04 13:04:46 -07:00
infoPageRouter . get ( "/unlock-info" , ( _req , res ) = > {
if ( _req . session ? . unlocked ) return res . redirect ( "/" ) ;
res . send ( `
< form method = "post" action = "/unlock-info" >
< h1 > Unlock Service Info < / h1 >
< input type = "hidden" name = "_csrf" value = "${res.locals.csrfToken}" / >
< input type = "password" name = "password" placeholder = "Password" / >
< button type = "submit" > Unlock < / button >
< / form >
` );
} ) ;
infoPageRouter . use ( checkIfUnlocked ) ;
}
infoPageRouter . get ( "/" , handleInfoPage ) ;
infoPageRouter . get ( "/status" , ( req , res ) = > {
res . json ( buildInfo ( req . protocol + "://" + req . get ( "host" ) , false ) ) ;
} ) ;
export { infoPageRouter } ;