AskMe / KshitijAIProxyServer.java
kvashist's picture
Create KshitijAIProxyServer.java
72e8ca8 verified
import com.sun.net.httpserver.Headers;
import com.sun.net.httpserver.HttpServer;
import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;
public class KshitijAIProxyServer {
private static final int PORT = 7860;
// Secrets (set in HF Space settings)
private static final String BACKEND_URL = getenvOrThrow("BACKEND_URL"); // e.g. https://user-kshitijai-backend.hf.space
private static final String HF_TOKEN = getenvOrThrow("HF_TOKEN"); // token that can access private space
public static void main(String[] args) throws Exception {
final Path baseDir = Paths.get(System.getProperty("user.dir")).toRealPath();
HttpServer server = HttpServer.create(new InetSocketAddress(PORT), 0);
server.createContext("/health", exchange -> {
addCors(exchange.getResponseHeaders());
if (!exchange.getRequestMethod().equalsIgnoreCase("GET")) {
exchange.sendResponseHeaders(405, -1);
return;
}
byte[] ok = "ok".getBytes(StandardCharsets.UTF_8);
exchange.getResponseHeaders().set("Content-Type", "text/plain; charset=utf-8");
exchange.sendResponseHeaders(200, ok.length);
try (OutputStream os = exchange.getResponseBody()) { os.write(ok); }
});
// Proxy: browser -> public space -> private backend space
server.createContext("/api/chat", exchange -> {
Headers h = exchange.getResponseHeaders();
addCors(h);
// Handle preflight
if (exchange.getRequestMethod().equalsIgnoreCase("OPTIONS")) {
exchange.sendResponseHeaders(204, -1);
return;
}
if (!exchange.getRequestMethod().equalsIgnoreCase("POST")) {
exchange.sendResponseHeaders(405, -1);
return;
}
byte[] requestBytes = exchange.getRequestBody().readAllBytes();
HttpClient client = HttpClient.newBuilder()
.connectTimeout(Duration.ofSeconds(10))
.build();
String target = BACKEND_URL.endsWith("/")
? BACKEND_URL + "api/chat"
: BACKEND_URL + "/api/chat";
HttpRequest req = HttpRequest.newBuilder()
.uri(URI.create(target))
.timeout(Duration.ofSeconds(300))
.header("Content-Type", "application/json")
.header("Authorization", "Bearer " + HF_TOKEN)
.POST(HttpRequest.BodyPublishers.ofByteArray(requestBytes))
.build();
HttpResponse<byte[]> resp;
try {
resp = client.send(req, HttpResponse.BodyHandlers.ofByteArray());
} catch (Exception e) {
byte[] err = ("Failed to reach backend at " + target + "\n" + e)
.getBytes(StandardCharsets.UTF_8);
h.set("Content-Type", "text/plain; charset=utf-8");
exchange.sendResponseHeaders(502, err.length);
try (OutputStream os = exchange.getResponseBody()) { os.write(err); }
return;
}
h.set("Content-Type", "application/json; charset=utf-8");
exchange.sendResponseHeaders(resp.statusCode(), resp.body().length);
try (OutputStream os = exchange.getResponseBody()) { os.write(resp.body()); }
});
// Static files: / -> index.html, /style.css, /app.js
server.createContext("/", exchange -> {
addCors(exchange.getResponseHeaders());
if (!exchange.getRequestMethod().equalsIgnoreCase("GET")) {
exchange.sendResponseHeaders(405, -1);
return;
}
String urlPath = exchange.getRequestURI().getPath();
if (urlPath == null || urlPath.isBlank() || urlPath.equals("/")) urlPath = "/index.html";
String rel = urlPath.startsWith("/") ? urlPath.substring(1) : urlPath;
Path requested = baseDir.resolve(rel).normalize();
if (!requested.startsWith(baseDir)) { exchange.sendResponseHeaders(403, -1); return; }
if (!Files.exists(requested) || Files.isDirectory(requested)) { exchange.sendResponseHeaders(404, -1); return; }
String ct = guessContentType(requested);
byte[] bytes = Files.readAllBytes(requested);
exchange.getResponseHeaders().set("Content-Type", ct);
exchange.sendResponseHeaders(200, bytes.length);
try (OutputStream os = exchange.getResponseBody()) { os.write(bytes); }
});
server.start();
System.out.println("Public UI + Proxy running on port " + PORT);
}
private static void addCors(Headers h) {
h.set("Access-Control-Allow-Origin", "*");
h.set("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
h.set("Access-Control-Allow-Headers", "Content-Type, Authorization");
}
private static String guessContentType(Path file) throws IOException {
String name = file.getFileName().toString().toLowerCase();
if (name.endsWith(".html")) return "text/html; charset=utf-8";
if (name.endsWith(".css")) return "text/css; charset=utf-8";
if (name.endsWith(".js")) return "application/javascript; charset=utf-8";
String probe = Files.probeContentType(file);
return probe != null ? probe : "application/octet-stream";
}
private static String getenvOrThrow(String key) {
String v = System.getenv(key);
if (v == null || v.isBlank()) throw new IllegalStateException("Missing env var: " + key);
return v;
}
}