krrishg/indie
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

live update while generating site

Browse Source
This commit is contained in:
Krrish Ghimire
2026-07-22 19:10:01 +05:45
parent ec4d030c0a
commit d279503dd7
17 changed files with 932 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
backend/src/main/java/com/krrishg/config/AsyncConfig.java Normal file
View File

@@ -0,0 +1,24 @@
package com.krrishg.config;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import java.util.concurrent.Executor;
@Configuration
@EnableAsync
public class AsyncConfig {
@Bean(name = "generationExecutor")
public Executor generationExecutor() {
ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
executor.setCorePoolSize(2);
executor.setMaxPoolSize(4);
executor.setQueueCapacity(10);
executor.setThreadNamePrefix("gen-");
executor.initialize();
return executor;
}
}

6
backend/src/main/java/com/krrishg/config/GlobalExceptionHandler.java
View File

@@ -10,6 +10,7 @@ import org.springframework.validation.FieldError;
import org.springframework.web.bind.MethodArgumentNotValidException;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.server.ResponseStatusException;
import java.util.stream.Collectors;
@@ -36,6 +37,11 @@ public class GlobalExceptionHandler {
return ProblemDetail.forStatusAndDetail(HttpStatus.BAD_REQUEST, message);
}
@ExceptionHandler(ResponseStatusException.class)
public ProblemDetail handleResponseStatus(ResponseStatusException ex) {
return ProblemDetail.forStatusAndDetail(ex.getStatusCode(), ex.getReason());
}
@ExceptionHandler(ClientAbortException.class)
public void handleClientAbort(ClientAbortException ex) {
if (log.isDebugEnabled()) {

19
backend/src/main/java/com/krrishg/config/SecurityConfig.java
View File

@@ -45,6 +45,8 @@ public class SecurityConfig {
.requestMatchers("/api/auth/**").permitAll()
.requestMatchers("/swagger-ui/**", "/api-docs/**").permitAll()
.requestMatchers(HttpMethod.GET, "/api/sites/**").authenticated()
.requestMatchers("/api/llm/generate/*/events").permitAll()
.requestMatchers("/api/llm/refine/*/events").permitAll()
.anyRequest().authenticated()
)
.exceptionHandling(ex -> ex
@@ -84,10 +86,9 @@ public class SecurityConfig {
HttpServletResponse response,
FilterChain filterChain)
throws ServletException, IOException {
String header = request.getHeader("Authorization");
String token = extractToken(request);
if (header != null && header.startsWith("Bearer ")) {
String token = header.substring(7);
if (token != null) {
try {
Claims claims = jwtConfig.validateToken(token);
UUID userId = UUID.fromString(claims.getSubject());
@@ -107,6 +108,18 @@ public class SecurityConfig {
filterChain.doFilter(request, response);
}
private String extractToken(HttpServletRequest request) {
String header = request.getHeader("Authorization");
if (header != null && header.startsWith("Bearer ")) {
return header.substring(7);
}
String queryToken = request.getParameter("token");
if (queryToken != null && !queryToken.isBlank()) {
return queryToken;
}
return null;
}
};
}
}

58
backend/src/main/java/com/krrishg/controller/LLMController.java
View File

@@ -1,5 +1,6 @@
package com.krrishg.controller;
import com.krrishg.config.JwtConfig;
import com.krrishg.config.RateLimitConfig;
import com.krrishg.config.SuggestRateLimitConfig;
import com.krrishg.config.UserPrincipal;
@@ -8,12 +9,16 @@ import com.krrishg.dto.SiteStructure;
import com.krrishg.dto.SuggestSectionsResponse;
import com.krrishg.model.GenerationVersion;
import com.krrishg.service.LLMService;
import com.krrishg.service.ProgressService;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.security.core.annotation.AuthenticationPrincipal;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.server.ResponseStatusException;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.util.List;
import java.util.Map;
@@ -24,18 +29,37 @@ import java.util.UUID;
public class LLMController {
private final LLMService llmService;
private final ProgressService progressService;
private final JwtConfig jwtConfig;
private final RateLimitConfig rateLimitConfig;
private final SuggestRateLimitConfig suggestRateLimitConfig;
private final ObjectMapper objectMapper;
public LLMController(LLMService llmService, RateLimitConfig rateLimitConfig,
public LLMController(LLMService llmService, ProgressService progressService,
JwtConfig jwtConfig,
RateLimitConfig rateLimitConfig,
SuggestRateLimitConfig suggestRateLimitConfig, ObjectMapper objectMapper) {
this.llmService = llmService;
this.progressService = progressService;
this.jwtConfig = jwtConfig;
this.rateLimitConfig = rateLimitConfig;
this.suggestRateLimitConfig = suggestRateLimitConfig;
this.objectMapper = objectMapper;
}
private UserPrincipal resolveToken(String token) {
if (token == null || token.isBlank()) return null;
try {
var claims = jwtConfig.validateToken(token);
UUID userId = UUID.fromString(claims.getSubject());
String email = claims.get("email", String.class);
String name = claims.get("name", String.class);
return new UserPrincipal(userId, email, name);
} catch (Exception e) {
return null;
}
}
@PostMapping("/generate")
public ResponseEntity<?> generate(
@AuthenticationPrincipal UserPrincipal principal,
@@ -73,13 +97,17 @@ public class LLMController {
request.get("references"), new TypeReference<List<Reference>>() {});
}
try {
SiteStructure structure = llmService.generateSite(
principal.id(), siteId, prompt, references, provider, apiKey, baseUrl, model);
return ResponseEntity.ok(structure);
} catch (IllegalStateException | IllegalArgumentException e) {
return ResponseEntity.badRequest().body(Map.of("error", e.getMessage()));
llmService.generateSiteAsync(principal.id(), siteId, prompt, references, provider, apiKey, baseUrl, model);
return ResponseEntity.accepted().body(Map.of("siteId", siteIdStr));
}
@GetMapping(value = "/generate/{siteId}/events", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public SseEmitter generateEvents(@PathVariable UUID siteId, @RequestParam(required = false) String token) {
UserPrincipal principal = resolveToken(token);
if (principal == null) {
throw new ResponseStatusException(HttpStatus.UNAUTHORIZED, "Invalid or missing token");
}
return progressService.register(siteId);
}
@PostMapping("/suggest-sections")
@@ -165,13 +193,17 @@ public class LLMController {
request.get("references"), new TypeReference<List<Reference>>() {});
}
try {
SiteStructure structure = llmService.refineSite(principal.id(), siteId,
currentStructure, prompt, references, provider, apiKey, baseUrl, model);
return ResponseEntity.ok(structure);
} catch (IllegalStateException | IllegalArgumentException e) {
return ResponseEntity.badRequest().body(Map.of("error", e.getMessage()));
llmService.refineSiteAsync(principal.id(), siteId, currentStructure, prompt, references, provider, apiKey, baseUrl, model);
return ResponseEntity.accepted().body(Map.of("siteId", siteIdStr));
}
@GetMapping(value = "/refine/{siteId}/events", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public SseEmitter refineEvents(@PathVariable UUID siteId, @RequestParam(required = false) String token) {
UserPrincipal principal = resolveToken(token);
if (principal == null) {
throw new ResponseStatusException(HttpStatus.UNAUTHORIZED, "Invalid or missing token");
}
return progressService.register(siteId);
}
@GetMapping("/versions/{siteId}")

2
backend/src/main/java/com/krrishg/service/LLMResponseParser.java
View File

@@ -1,5 +1,6 @@
package com.krrishg.service;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
@@ -25,6 +26,7 @@ public class LLMResponseParser {
public LLMResponseParser(ObjectMapper objectMapper) {
this.objectMapper = objectMapper.copy()
.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
this.objectMapper.configure(JsonParser.Feature.ALLOW_BACKSLASH_ESCAPING_ANY_CHARACTER, true);
}
public SiteStructure parse(String llmOutput) {

159
backend/src/main/java/com/krrishg/service/LLMService.java
View File

@@ -25,6 +25,7 @@ import org.springframework.data.domain.Sort;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import java.io.BufferedReader;
@@ -39,6 +40,9 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.function.Consumer;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
@Service
@@ -49,6 +53,7 @@ public class LLMService {
private final ProviderFactory providerFactory;
private final LLMResponseParser responseParser;
private final ReferenceService referenceService;
private final ProgressService progressService;
private final MongoTemplate mongoTemplate;
private final UserRepository userRepository;
private final EncryptionService encryptionService;
@@ -62,6 +67,7 @@ public class LLMService {
public LLMService(ProviderFactory providerFactory,
LLMResponseParser responseParser,
ReferenceService referenceService,
ProgressService progressService,
MongoTemplate mongoTemplate,
UserRepository userRepository,
EncryptionService encryptionService,
@@ -74,6 +80,7 @@ public class LLMService {
this.providerFactory = providerFactory;
this.responseParser = responseParser;
this.referenceService = referenceService;
this.progressService = progressService;
this.mongoTemplate = mongoTemplate;
this.userRepository = userRepository;
this.encryptionService = encryptionService;
@@ -115,6 +122,62 @@ public class LLMService {
return structure;
}
@Async("generationExecutor")
public void generateSiteAsync(UUID userId, UUID siteId, String prompt,
List<Reference> references, String provider, String apiKey, String baseUrl, String model) {
try {
progressService.sendProgress(siteId, "resolve_config", "Resolving API configuration...");
String resolvedKey = resolveApiKey(userId, provider, apiKey);
String resolvedBaseUrl = resolveBaseUrl(userId, provider, baseUrl);
String resolvedModel = resolveModel(userId, provider, model);
progressService.sendProgress(siteId, "prepare_prompt", "Preparing prompt context...");
String referenceContext = referenceService.buildReferenceContext(references);
String fullPrompt = prompt + referenceContext;
progressService.sendProgress(siteId, "calling_ai", "Contacting " + provider + " AI...");
progressService.sendProgress(siteId, "generating_section", "Generating pages...");
LLMOptions options = LLMOptions.builder().build();
StringBuilder jsonBuffer = new StringBuilder();
Set<String> seenTitles = new HashSet<>();
LLMResult result = callLLM(provider, resolvedKey, resolvedBaseUrl, resolvedModel,
systemPrompt, fullPrompt, references, options,
chunk -> {
progressService.sendStreamingChunk(siteId, chunk);
jsonBuffer.append(chunk);
detectAndReportNewPages(siteId, jsonBuffer.toString(), seenTitles);
});
log.info("LLM generation completed: model={}, latency={}ms, tokens={}+{}",
result.getModel(), result.getLatencyMs(),
result.getPromptTokens(), result.getCompletionTokens());
progressService.sendProgress(siteId, "generating_section", "");
progressService.sendProgress(siteId, "parsing", "Processing AI response...");
SiteStructure structure = responseParser.parse(result.getContent());
for (SiteStructure.PageStructure page : structure.getPages()) {
if (seenTitles.add(page.getTitle())) {
log.info("Detected page from final JSON: title={}, site={}", page.getTitle(), siteId);
progressService.sendProgress(siteId, "generating_section", "Generating page: " + page.getTitle());
}
}
progressService.sendProgress(siteId, "saving", "Saving generation version...");
saveVersion(userId, siteId, prompt, references, result.getContent());
progressService.sendComplete(siteId, structure);
} catch (Exception e) {
log.error("Async generation failed for site {}: {}", siteId, e.getMessage(), e);
progressService.sendError(siteId, e.getMessage());
}
}
public SuggestSectionsResponse suggestSections(UUID userId, String briefDescription,
String provider, String apiKey) {
return suggestSections(userId, briefDescription, provider, apiKey, null, null);
@@ -133,7 +196,7 @@ public class LLMService {
LLMOptions options = LLMOptions.builder()
.temperature(0.2)
.maxTokens(2000)
.maxTokens(8192)
.build();
LLMResult result = callLLM(provider, resolvedKey, resolvedBaseUrl, resolvedModel,
suggestSectionsPrompt, briefDescription, null, options);
@@ -187,6 +250,65 @@ public class LLMService {
return merged;
}
@Async("generationExecutor")
public void refineSiteAsync(UUID userId, UUID siteId, SiteStructure currentStructure,
String refinementPrompt, List<Reference> references,
String provider, String apiKey, String baseUrl, String model) {
try {
progressService.sendProgress(siteId, "resolve_config", "Resolving API configuration...");
String resolvedKey = resolveApiKey(userId, provider, apiKey);
String resolvedBaseUrl = resolveBaseUrl(userId, provider, baseUrl);
String resolvedModel = resolveModel(userId, provider, model);
progressService.sendProgress(siteId, "prepare_prompt", "Preparing refinement context...");
String currentJson = serializeStructure(currentStructure);
String referenceContext = referenceService.buildReferenceContext(references);
String combinedPrompt = "Current site structure:\n" + currentJson
+ "\n\nRefinement request: " + refinementPrompt
+ referenceContext
+ "\n\nOutput only the fields that changed. Fields you omit will remain unchanged."
+ " Include the \"id\" field for any page you modify or add."
+ " To delete a page, include it with \"id\" and \"deleted\": true."
+ " Pages not included in your output will stay as-is.";
progressService.sendProgress(siteId, "calling_ai", "Contacting " + provider + " AI...");
progressService.sendProgress(siteId, "generating_section", "Generating pages...");
LLMOptions options = LLMOptions.builder()
.temperature(refineTemperature)
.build();
StringBuilder jsonBuffer = new StringBuilder();
Set<String> seenTitles = new HashSet<>();
LLMResult result = callLLM(provider, resolvedKey, resolvedBaseUrl, resolvedModel, refinerPrompt, combinedPrompt, references, options,
chunk -> {
progressService.sendStreamingChunk(siteId, chunk);
jsonBuffer.append(chunk);
detectAndReportNewPages(siteId, jsonBuffer.toString(), seenTitles);
});
log.info("LLM refinement completed: model={}, latency={}ms, tokens={}+{}",
result.getModel(), result.getLatencyMs(),
result.getPromptTokens(), result.getCompletionTokens());
progressService.sendProgress(siteId, "generating_section", "");
progressService.sendProgress(siteId, "parsing", "Processing AI response...");
SiteStructure patch = responseParser.parseLenient(result.getContent());
SiteStructure merged = mergePatches(currentStructure, patch);
progressService.sendProgress(siteId, "saving", "Saving refinement version...");
saveVersion(userId, siteId, refinementPrompt, references, serializeStructure(merged));
progressService.sendRefineComplete(siteId, merged);
} catch (Exception e) {
log.error("Async refinement failed for site {}: {}", siteId, e.getMessage(), e);
progressService.sendError(siteId, e.getMessage());
}
}
public List<GenerationVersion> getVersions(UUID siteId) {
Query query = new Query(Criteria.where("siteId").is(siteId))
.with(Sort.by(Sort.Direction.DESC, "versionNumber"))
@@ -397,6 +519,41 @@ public class LLMService {
}
}
private LLMResult callLLM(String provider, String apiKey, String baseUrl, String model, String systemPrompt,
String userPrompt, List<Reference> references, LLMOptions options, Consumer<String> onChunk) {
LLMProvider p = providerFactory.createProvider(provider, apiKey, baseUrl, model);
try {
return p.generate(systemPrompt, userPrompt, references, options, onChunk);
} catch (ResourceAccessException e) {
throw new IllegalStateException(
provider + " timed out. Please try again or use a simpler prompt.");
} catch (HttpClientErrorException e) {
String detail = e.getResponseBodyAsString();
String msg = switch (e.getStatusCode().value()) {
case 401 -> provider + " rejected the API key (401 Unauthorized). Verify the key is valid in Settings.";
case 429 -> provider + " rate limit exceeded (429). Try again later.";
default -> {
String body = (detail != null && !detail.isBlank()) ? ": " + detail : "";
yield provider + " API error: " + e.getStatusCode() + " " + e.getStatusText() + body;
}
};
throw new IllegalStateException(msg);
}
}
private static final Pattern TITLE_PATTERN = Pattern.compile("\"title\"\\s*:\\s*\"([^\"]+)\"");
private void detectAndReportNewPages(UUID siteId, String json, Set<String> seenTitles) {
Matcher matcher = TITLE_PATTERN.matcher(json);
while (matcher.find()) {
String title = matcher.group(1);
if (seenTitles.add(title)) {
log.info("Detected page from stream: title={}, site={}", title, siteId);
progressService.sendProgress(siteId, "generating_section", "Generating page: " + title);
}
}
}
private void saveVersion(UUID userId, UUID siteId, String prompt,
List<Reference> references, String generatedJson) {
int nextVersion = getNextVersionNumber(siteId);

149
backend/src/main/java/com/krrishg/service/ProgressService.java Normal file
View File

@@ -0,0 +1,149 @@
package com.krrishg.service;
import com.krrishg.dto.SiteStructure;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
@Service
public class ProgressService {
private static final Logger log = LoggerFactory.getLogger(ProgressService.class);
private static final long SSE_TIMEOUT = 300_000L;
private final Map<UUID, SseEmitter> emitters = new ConcurrentHashMap<>();
private final Map<UUID, List<Map<String, Object>>> pendingEvents = new ConcurrentHashMap<>();
private final Object bufferLock = new Object();
public SseEmitter register(UUID siteId) {
SseEmitter emitter = new SseEmitter(SSE_TIMEOUT);
List<Map<String, Object>> buffer;
synchronized (bufferLock) {
emitters.put(siteId, emitter);
buffer = pendingEvents.remove(siteId);
}
if (buffer != null) {
for (Map<String, Object> event : buffer) {
String type = (String) event.get("type");
Object data = event.get("data");
try {
emitter.send(SseEmitter.event().name(type).data(data));
} catch (IOException e) {
emitters.remove(siteId);
log.debug("Failed to replay event for site {}: {}", siteId, e.getMessage());
return emitter;
}
}
}
emitter.onCompletion(() -> {
emitters.remove(siteId);
log.debug("SSE completed for site {}", siteId);
});
emitter.onTimeout(() -> {
emitters.remove(siteId);
log.debug("SSE timed out for site {}", siteId);
});
emitter.onError(e -> {
emitters.remove(siteId);
log.debug("SSE error for site {}: {}", siteId, e.getMessage());
});
return emitter;
}
public void sendProgress(UUID siteId, String stage, String message) {
Map<String, Object> data = Map.of("stage", stage, "message", message);
SseEmitter emitter = emitters.get(siteId);
if (emitter != null) {
try {
emitter.send(SseEmitter.event().name("progress").data(data));
return;
} catch (IOException e) {
emitters.remove(siteId);
log.debug("Failed to send progress for site {}: {}", siteId, e.getMessage());
}
}
bufferEvent(siteId, "progress", data);
}
public void sendComplete(UUID siteId, SiteStructure result) {
SseEmitter emitter = emitters.get(siteId);
if (emitter != null) {
try {
emitter.send(SseEmitter.event().name("complete").data(result));
emitter.complete();
return;
} catch (IOException e) {
log.debug("Failed to send complete for site {}: {}", siteId, e.getMessage());
}
} else {
bufferEvent(siteId, "complete", result);
}
emitters.remove(siteId);
}
public void sendRefineComplete(UUID siteId, SiteStructure result) {
SseEmitter emitter = emitters.get(siteId);
if (emitter != null) {
try {
emitter.send(SseEmitter.event().name("complete").data(result));
emitter.complete();
return;
} catch (IOException e) {
log.debug("Failed to send refine complete for site {}: {}", siteId, e.getMessage());
}
} else {
bufferEvent(siteId, "complete", result);
}
emitters.remove(siteId);
}
public void sendStreamingChunk(UUID siteId, String chunk) {
Map<String, String> data = Map.of("chunk", chunk);
SseEmitter emitter = emitters.get(siteId);
if (emitter != null) {
try {
emitter.send(SseEmitter.event().name("streaming").data(data));
} catch (IOException e) {
emitters.remove(siteId);
log.debug("Failed to send streaming chunk for site {}: {}", siteId, e.getMessage());
}
}
}
public void sendError(UUID siteId, String error) {
Map<String, String> data = Map.of("error", error);
SseEmitter emitter = emitters.get(siteId);
if (emitter != null) {
try {
emitter.send(SseEmitter.event().name("error").data(data));
emitter.complete();
return;
} catch (IOException e) {
log.debug("Failed to send error for site {}: {}", siteId, e.getMessage());
}
} else {
bufferEvent(siteId, "error", data);
}
emitters.remove(siteId);
}
private void bufferEvent(UUID siteId, String type, Object data) {
synchronized (bufferLock) {
if (!emitters.containsKey(siteId)) {
pendingEvents.computeIfAbsent(siteId, k -> new ArrayList<>())
.add(Map.of("type", type, "data", data));
}
}
}
}

180
backend/src/main/java/com/krrishg/service/llm/GeminiProvider.java
View File

@@ -1,14 +1,24 @@
package com.krrishg.service.llm;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.krrishg.dto.LLMResult;
import com.krrishg.dto.LLMOptions;
import com.krrishg.dto.Reference;
import org.springframework.http.*;
import org.springframework.web.client.RestTemplate;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
public class GeminiProvider implements LLMProvider {
@@ -16,64 +26,29 @@ public class GeminiProvider implements LLMProvider {
private final String apiKey;
private final String model;
private final String baseUrl;
private final ObjectMapper objectMapper;
public GeminiProvider(String apiKey, String model, String baseUrl, RestTemplate restTemplate) {
this.apiKey = apiKey;
this.model = model;
this.baseUrl = baseUrl;
this.restTemplate = restTemplate;
this.objectMapper = new ObjectMapper();
}
@Override
public LLMResult generate(String systemPrompt, String userPrompt, List<Reference> references, LLMOptions options) {
long start = System.currentTimeMillis();
List<Map<String, Object>> parts = new ArrayList<>();
parts.add(Map.of("text", userPrompt));
if (references != null) {
for (Reference ref : references) {
if (ref.getType() == Reference.ReferenceType.IMAGE && ref.getUrl() != null) {
parts.add(Map.of(
"inlineData", Map.of(
"mimeType", "image/jpeg",
"data", ref.getUrl()
)
));
}
}
}
Map<String, Object> systemInstruction = Map.of(
"parts", List.of(Map.of("text", systemPrompt))
);
Map<String, Object> contents = Map.of(
"parts", parts
);
Map<String, Object> generationConfig = Map.of(
"temperature", options.getTemperature(),
"maxOutputTokens", options.getMaxTokens()
);
Map<String, Object> requestBody = Map.of(
"system_instruction", systemInstruction,
"contents", List.of(contents),
"generationConfig", generationConfig
);
Map<String, Object> requestBody = buildRequestBody(systemPrompt, userPrompt, references, options);
HttpHeaders headers = new HttpHeaders();
headers.setContentType(MediaType.APPLICATION_JSON);
HttpEntity<Map<String, Object>> request = new HttpEntity<>(requestBody, headers);
String url = baseUrl + "/" + model + ":generateContent?key=" + apiKey;
ResponseEntity<Map> response = restTemplate.exchange(url, HttpMethod.POST, request, Map.class);
ResponseEntity<Map> response = restTemplate.exchange(buildUrl(), HttpMethod.POST, request, Map.class);
long latencyMs = System.currentTimeMillis() - start;
Map<String, Object> responseBody = response.getBody();
String content = extractContent(responseBody);
@@ -94,6 +69,133 @@ public class GeminiProvider implements LLMProvider {
.build();
}
@Override
public LLMResult generate(String systemPrompt, String userPrompt,
List<Reference> references, LLMOptions options, Consumer<String> onChunk) {
long start = System.currentTimeMillis();
try {
HttpURLConnection conn = openConnection(buildStreamingUrl(), options);
writeRequestBody(conn, buildRequestBody(systemPrompt, userPrompt, references, options));
verifyResponse(conn);
return readStream(conn, onChunk, start);
} catch (Exception e) {
throw new RuntimeException("Gemini streaming failed: " + e.getMessage(), e);
}
}
private Map<String, Object> buildRequestBody(String systemPrompt, String userPrompt,
List<Reference> references, LLMOptions options) {
List<Map<String, Object>> parts = new ArrayList<>();
parts.add(Map.of("text", userPrompt));
if (references != null) {
for (Reference ref : references) {
if (ref.getType() == Reference.ReferenceType.IMAGE && ref.getUrl() != null) {
parts.add(Map.of("inlineData", Map.of("mimeType", "image/jpeg", "data", ref.getUrl())));
}
}
}
return Map.of(
"system_instruction", Map.of("parts", List.of(Map.of("text", systemPrompt))),
"contents", List.of(Map.of("parts", parts)),
"generationConfig", Map.of("temperature", options.getTemperature(), "maxOutputTokens", options.getMaxTokens())
);
}
private String buildUrl() {
return baseUrl + "/" + model + ":generateContent?key=" + apiKey;
}
private String buildStreamingUrl() {
return baseUrl + "/" + model + ":streamGenerateContent?alt=sse&key=" + apiKey;
}
private HttpURLConnection openConnection(String url, LLMOptions options) throws IOException {
HttpURLConnection conn = (HttpURLConnection) URI.create(url).toURL().openConnection();
conn.setRequestMethod("POST");
conn.setRequestProperty("Content-Type", "application/json");
conn.setRequestProperty("Accept", "text/event-stream");
conn.setDoOutput(true);
conn.setConnectTimeout(30000);
conn.setReadTimeout(options.getTimeoutSeconds() * 1000);
return conn;
}
private void writeRequestBody(HttpURLConnection conn, Map<String, Object> requestBody) throws IOException {
byte[] bodyBytes = objectMapper.writeValueAsBytes(requestBody);
try (OutputStream os = conn.getOutputStream()) {
os.write(bodyBytes);
os.flush();
}
}
private void verifyResponse(HttpURLConnection conn) throws IOException {
int status = conn.getResponseCode();
if (status != 200) {
String errorBody = new String(conn.getErrorStream().readAllBytes(), StandardCharsets.UTF_8);
throw new RuntimeException("Gemini API error " + status + ": " + errorBody);
}
}
private static class StreamState {
StringBuilder content = new StringBuilder();
int promptTokens;
int completionTokens;
}
private LLMResult readStream(HttpURLConnection conn, Consumer<String> onChunk, long start) throws IOException {
StreamState state = new StreamState();
try (BufferedReader reader = new BufferedReader(new InputStreamReader(conn.getInputStream(), StandardCharsets.UTF_8))) {
String line;
while ((line = reader.readLine()) != null) {
processSseLine(line, state, onChunk);
}
}
return LLMResult.builder()
.content(state.content.toString())
.promptTokens(state.promptTokens)
.completionTokens(state.completionTokens)
.model(model)
.latencyMs(System.currentTimeMillis() - start)
.build();
}
private void processSseLine(String line, StreamState state, Consumer<String> onChunk) {
if (!line.startsWith("data: ")) return;
String data = line.substring(6).trim();
if (data.isEmpty()) return;
try {
JsonNode node = objectMapper.readTree(data);
if (node.has("usageMetadata")) {
JsonNode usage = node.get("usageMetadata");
state.promptTokens = usage.has("promptTokenCount") ? usage.get("promptTokenCount").asInt() : 0;
state.completionTokens = usage.has("candidatesTokenCount") ? usage.get("candidatesTokenCount").asInt() : 0;
}
JsonNode candidates = node.get("candidates");
if (candidates != null && candidates.isArray() && candidates.size() > 0) {
JsonNode content = candidates.get(0).get("content");
if (content != null && content.has("parts") && content.get("parts").isArray()) {
JsonNode partsNode = content.get("parts").get(0);
if (partsNode != null) {
JsonNode textNode = partsNode.get("text");
if (textNode != null && !textNode.isNull()) {
String text = textNode.asText();
if (!text.isEmpty()) {
state.content.append(text);
onChunk.accept(text);
}
}
}
}
}
} catch (Exception ignored) {
}
}
@SuppressWarnings("unchecked")
private String extractContent(Map<String, Object> responseBody) {
try {

6
backend/src/main/java/com/krrishg/service/llm/LLMProvider.java
View File

@@ -5,10 +5,16 @@ import com.krrishg.dto.LLMOptions;
import com.krrishg.dto.Reference;
import java.util.List;
import java.util.function.Consumer;
public interface LLMProvider {
LLMResult generate(String systemPrompt, String userPrompt, List<Reference> references, LLMOptions options);
default LLMResult generate(String systemPrompt, String userPrompt,
List<Reference> references, LLMOptions options, Consumer<String> onChunk) {
return generate(systemPrompt, userPrompt, references, options);
}
String getName();
}

139
backend/src/main/java/com/krrishg/service/llm/OpenAIProvider.java
View File

@@ -1,13 +1,23 @@
package com.krrishg.service.llm;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.krrishg.dto.LLMResult;
import com.krrishg.dto.LLMOptions;
import com.krrishg.dto.Reference;
import org.springframework.http.*;
import org.springframework.web.client.RestTemplate;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
public class OpenAIProvider implements LLMProvider {
@@ -15,27 +25,21 @@ public class OpenAIProvider implements LLMProvider {
private final String apiKey;
private final String model;
private final String url;
private final ObjectMapper objectMapper;
public OpenAIProvider(String apiKey, String model, String url, RestTemplate restTemplate) {
this.apiKey = apiKey;
this.model = model;
this.url = url;
this.restTemplate = restTemplate;
this.objectMapper = new ObjectMapper();
}
@Override
public LLMResult generate(String systemPrompt, String userPrompt, List<Reference> references, LLMOptions options) {
long start = System.currentTimeMillis();
Map<String, Object> requestBody = Map.of(
"model", model,
"messages", List.of(
Map.of("role", "system", "content", systemPrompt),
Map.of("role", "user", "content", userPrompt)
),
"temperature", options.getTemperature(),
"max_tokens", options.getMaxTokens()
);
Map<String, Object> requestBody = buildRequestBody(systemPrompt, userPrompt, options);
HttpHeaders headers = new HttpHeaders();
headers.setContentType(MediaType.APPLICATION_JSON);
@@ -67,6 +71,123 @@ public class OpenAIProvider implements LLMProvider {
.build();
}
@Override
public LLMResult generate(String systemPrompt, String userPrompt,
List<Reference> references, LLMOptions options, Consumer<String> onChunk) {
long start = System.currentTimeMillis();
try {
HttpURLConnection conn = openStreamingConnection(options);
writeRequestBody(conn, buildStreamingRequest(systemPrompt, userPrompt, options));
verifyResponse(conn);
String fullContent = readStream(conn, onChunk);
return buildResult(fullContent, start);
} catch (Exception e) {
throw new RuntimeException("OpenAI streaming failed: " + e.getMessage(), e);
}
}
private Map<String, Object> buildRequestBody(String systemPrompt, String userPrompt, LLMOptions options) {
return Map.of(
"model", model,
"messages", List.of(
Map.of("role", "system", "content", systemPrompt),
Map.of("role", "user", "content", userPrompt)
),
"temperature", options.getTemperature(),
"max_tokens", options.getMaxTokens()
);
}
private Map<String, Object> buildStreamingRequest(String systemPrompt, String userPrompt, LLMOptions options) {
return Map.of(
"model", model,
"messages", List.of(
Map.of("role", "system", "content", systemPrompt),
Map.of("role", "user", "content", userPrompt)
),
"temperature", options.getTemperature(),
"max_tokens", options.getMaxTokens(),
"stream", true
);
}
private HttpURLConnection openStreamingConnection(LLMOptions options) throws IOException {
HttpURLConnection conn = (HttpURLConnection) URI.create(url).toURL().openConnection();
conn.setRequestMethod("POST");
conn.setRequestProperty("Content-Type", "application/json");
conn.setRequestProperty("Authorization", "Bearer " + apiKey);
conn.setRequestProperty("Accept", "text/event-stream");
conn.setDoOutput(true);
conn.setConnectTimeout(30000);
conn.setReadTimeout(options.getTimeoutSeconds() * 1000);
return conn;
}
private void writeRequestBody(HttpURLConnection conn, Map<String, Object> requestBody) throws IOException {
byte[] bodyBytes = objectMapper.writeValueAsBytes(requestBody);
try (OutputStream os = conn.getOutputStream()) {
os.write(bodyBytes);
os.flush();
}
}
private void verifyResponse(HttpURLConnection conn) throws IOException {
int status = conn.getResponseCode();
if (status != 200) {
String errorBody = new String(conn.getErrorStream().readAllBytes(), StandardCharsets.UTF_8);
throw new RuntimeException("OpenAI API error " + status + ": " + errorBody);
}
}
private String readStream(HttpURLConnection conn, Consumer<String> onChunk) throws IOException {
StringBuilder fullContent = new StringBuilder();
try (BufferedReader reader = new BufferedReader(new InputStreamReader(conn.getInputStream(), StandardCharsets.UTF_8))) {
String line;
while ((line = reader.readLine()) != null) {
if (!processSseLine(line, fullContent, onChunk)) {
break;
}
}
}
return fullContent.toString();
}
private boolean processSseLine(String line, StringBuilder fullContent, Consumer<String> onChunk) {
if (!line.startsWith("data: ")) return true;
String data = line.substring(6).trim();
if ("[DONE]".equals(data)) return false;
try {
JsonNode node = objectMapper.readTree(data);
JsonNode choices = node.get("choices");
if (choices != null && choices.isArray() && choices.size() > 0) {
JsonNode delta = choices.get(0).get("delta");
if (delta != null) {
JsonNode content = delta.get("content");
if (content != null && !content.isNull()) {
String text = content.asText();
if (!text.isEmpty()) {
fullContent.append(text);
onChunk.accept(text);
}
}
}
}
} catch (Exception ignored) {
}
return true;
}
private LLMResult buildResult(String fullContent, long start) {
return LLMResult.builder()
.content(fullContent)
.promptTokens(0)
.completionTokens(0)
.model(model)
.latencyMs(System.currentTimeMillis() - start)
.build();
}
@SuppressWarnings("unchecked")
private String extractContent(Map<String, Object> responseBody) {
try {

101
backend/src/test/java/com/krrishg/controller/LLMControllerTest.java
View File

@@ -2,6 +2,7 @@ package com.krrishg.controller;
import com.krrishg.config.GlobalExceptionHandler;
import com.krrishg.config.JwtAuthenticationToken;
import com.krrishg.config.JwtConfig;
import com.krrishg.config.RateLimitConfig;
import com.krrishg.config.SuggestRateLimitConfig;
import com.krrishg.config.UserPrincipal;
@@ -10,6 +11,7 @@ import com.krrishg.dto.SiteStructure;
import com.krrishg.dto.SuggestSectionsResponse;
import com.krrishg.model.GenerationVersion;
import com.krrishg.service.LLMService;
import com.krrishg.service.ProgressService;
import com.krrishg.support.TestSecurityConfig;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
@@ -44,6 +46,12 @@ class LLMControllerTest {
@MockBean
private LLMService llmService;
@MockBean
private ProgressService progressService;
@MockBean
private JwtConfig jwtConfig;
@MockBean
private RateLimitConfig rateLimitConfig;
@@ -70,22 +78,13 @@ class LLMControllerTest {
class Generate {
@Test
void returns200WithSiteStructure() throws Exception {
void returns202Accepted() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
SiteStructure structure = SiteStructure.builder()
.pages(List.of(
SiteStructure.PageStructure.builder()
.title("Home").slug("home")
.rawHtml("<h1>Home</h1>")
.build()
))
.build();
when(llmService.generateSite(any(), any(), anyString(), any(), anyString(), any(), any(), any())).thenReturn(structure);
String siteId = UUID.randomUUID().toString();
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString(),
"siteId", siteId,
"provider", "gemini"
);
@@ -93,27 +92,20 @@ class LLMControllerTest {
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isOk())
.andExpect(jsonPath("$.pages[0].title").value("Home"));
.andExpect(status().isAccepted())
.andExpect(jsonPath("$.siteId").value(siteId));
verify(llmService).generateSiteAsync(any(), any(), anyString(), any(), anyString(), any(), any(), any());
}
@Test
void returns200WithReferences() throws Exception {
void returns202WithReferences() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
SiteStructure structure = SiteStructure.builder()
.pages(List.of(
SiteStructure.PageStructure.builder()
.title("Home").slug("home")
.rawHtml("<h1>Home</h1>")
.build()
))
.build();
when(llmService.generateSite(any(), any(), anyString(), any(), anyString(), any(), any(), any())).thenReturn(structure);
String siteId = UUID.randomUUID().toString();
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString(),
"siteId", siteId,
"provider", "gemini",
"references", List.of(
Map.of("type", "IMAGE", "url", "data:img", "altText", "Photo")
@@ -124,26 +116,20 @@ class LLMControllerTest {
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isOk());
.andExpect(status().isAccepted())
.andExpect(jsonPath("$.siteId").value(siteId));
verify(llmService).generateSiteAsync(any(), any(), anyString(), any(), anyString(), any(), any(), any());
}
@Test
void returns200WithPerRequestApiKey() throws Exception {
void returns202WithPerRequestApiKey() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
SiteStructure structure = SiteStructure.builder()
.pages(List.of(
SiteStructure.PageStructure.builder()
.title("Home").slug("home")
.rawHtml("<h1>Home</h1>")
.build()
))
.build();
when(llmService.generateSite(any(), any(), anyString(), any(), eq("gemini"), eq("custom-key"), any(), any())).thenReturn(structure);
String siteId = UUID.randomUUID().toString();
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString(),
"siteId", siteId,
"provider", "gemini",
"apiKey", "custom-key"
);
@@ -152,7 +138,10 @@ class LLMControllerTest {
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isOk());
.andExpect(status().isAccepted())
.andExpect(jsonPath("$.siteId").value(siteId));
verify(llmService).generateSiteAsync(any(), any(), anyString(), any(), eq("gemini"), eq("custom-key"), any(), any());
}
@Test
@@ -206,14 +195,13 @@ class LLMControllerTest {
}
@Test
void returns400WhenNoKeyConfigured() throws Exception {
void returns202WhenNoKeyConfigured() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
when(llmService.generateSite(any(), any(), anyString(), any(), anyString(), any(), any(), any()))
.thenThrow(new IllegalStateException("No API key configured for gemini"));
String siteId = UUID.randomUUID().toString();
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString(),
"siteId", siteId,
"provider", "gemini"
);
@@ -221,8 +209,8 @@ class LLMControllerTest {
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isBadRequest())
.andExpect(jsonPath("$.error").value("No API key configured for gemini"));
.andExpect(status().isAccepted())
.andExpect(jsonPath("$.siteId").value(siteId));
}
@Test
@@ -246,9 +234,10 @@ class LLMControllerTest {
class Refine {
@Test
void returns200WithRefinedStructure() throws Exception {
void returns202Accepted() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
String siteId = UUID.randomUUID().toString();
SiteStructure current = SiteStructure.builder()
.pages(List.of(
SiteStructure.PageStructure.builder()
@@ -257,19 +246,9 @@ class LLMControllerTest {
.build()
))
.build();
SiteStructure refined = SiteStructure.builder()
.pages(List.of(
SiteStructure.PageStructure.builder()
.title("Refined").slug("refined")
.rawHtml("<h1>Refined</h1>")
.build()
))
.build();
when(llmService.refineSite(any(), any(), any(), anyString(), any(), anyString(), any(), any(), any())).thenReturn(refined);
Map<String, Object> request = Map.of(
"prompt", "Make it better",
"siteId", UUID.randomUUID().toString(),
"siteId", siteId,
"provider", "gemini",
"currentStructure", current
);
@@ -278,8 +257,10 @@ class LLMControllerTest {
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isOk())
.andExpect(jsonPath("$.pages[0].title").value("Refined"));
.andExpect(status().isAccepted())
.andExpect(jsonPath("$.siteId").value(siteId));
verify(llmService).refineSiteAsync(any(), any(), any(), anyString(), any(), anyString(), any(), any(), any());
}
@Test

15
backend/src/test/java/com/krrishg/service/LLMResponseParserTest.java
View File

@@ -239,5 +239,20 @@ class LLMResponseParserTest {
SiteStructure structure = parser.parse(json);
assertNull(structure.getPages().get(0).getRawHtml());
}
@Test
void parsesRawHtmlWithInvalidBackslashEscape() {
String rawHtml = "before\\>after";
String json = """
{
"pages": [
{"title": "Home", "slug": "home", "rawHtml": "%s"}
]
}
""".formatted(rawHtml);
SiteStructure structure = parser.parse(json);
assertEquals("before>after", structure.getPages().get(0).getRawHtml());
}
}
}

5
backend/src/test/java/com/krrishg/service/LLMServiceTest.java
View File

@@ -41,6 +41,9 @@ class LLMServiceTest {
@Mock
private ResourceLoader resourceLoader;
@Mock
private ProgressService progressService;
private FakeLLMProvider geminiProvider;
private FakeProviderFactory providerFactory;
private LLMResponseParser responseParser;
@@ -80,7 +83,7 @@ class LLMServiceTest {
new ByteArrayInputStream("You are a site generator.".getBytes(StandardCharsets.UTF_8)));
when(resourceLoader.getResource(anyString())).thenReturn(mockResource);
llmService = new LLMService(providerFactory, responseParser, referenceService,
llmService = new LLMService(providerFactory, responseParser, referenceService, progressService,
mongoTemplate, userRepository, encryptionService,
resourceLoader, "classpath:prompts/site-generator.txt",
"classpath:prompts/site-refiner.txt",