diff --git a/backend/src/main/java/com/krrishg/config/AsyncConfig.java b/backend/src/main/java/com/krrishg/config/AsyncConfig.java new file mode 100644 index 0000000..9d8d198 --- /dev/null +++ b/backend/src/main/java/com/krrishg/config/AsyncConfig.java @@ -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; + } +} diff --git a/backend/src/main/java/com/krrishg/config/GlobalExceptionHandler.java b/backend/src/main/java/com/krrishg/config/GlobalExceptionHandler.java index 6490709..35edc8f 100644 --- a/backend/src/main/java/com/krrishg/config/GlobalExceptionHandler.java +++ b/backend/src/main/java/com/krrishg/config/GlobalExceptionHandler.java @@ -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()) { diff --git a/backend/src/main/java/com/krrishg/config/SecurityConfig.java b/backend/src/main/java/com/krrishg/config/SecurityConfig.java index bb3f8eb..4228238 100644 --- a/backend/src/main/java/com/krrishg/config/SecurityConfig.java +++ b/backend/src/main/java/com/krrishg/config/SecurityConfig.java @@ -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; + } }; } } diff --git a/backend/src/main/java/com/krrishg/controller/LLMController.java b/backend/src/main/java/com/krrishg/controller/LLMController.java index d2867a1..7664a29 100644 --- a/backend/src/main/java/com/krrishg/controller/LLMController.java +++ b/backend/src/main/java/com/krrishg/controller/LLMController.java @@ -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>() {}); } - 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>() {}); } - 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}") diff --git a/backend/src/main/java/com/krrishg/service/LLMResponseParser.java b/backend/src/main/java/com/krrishg/service/LLMResponseParser.java index 24d757b..9557b33 100644 --- a/backend/src/main/java/com/krrishg/service/LLMResponseParser.java +++ b/backend/src/main/java/com/krrishg/service/LLMResponseParser.java @@ -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) { diff --git a/backend/src/main/java/com/krrishg/service/LLMService.java b/backend/src/main/java/com/krrishg/service/LLMService.java index f38a594..54fe399 100644 --- a/backend/src/main/java/com/krrishg/service/LLMService.java +++ b/backend/src/main/java/com/krrishg/service/LLMService.java @@ -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 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 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 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 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 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 references, LLMOptions options, Consumer 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 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 references, String generatedJson) { int nextVersion = getNextVersionNumber(siteId); diff --git a/backend/src/main/java/com/krrishg/service/ProgressService.java b/backend/src/main/java/com/krrishg/service/ProgressService.java new file mode 100644 index 0000000..c7e7bb1 --- /dev/null +++ b/backend/src/main/java/com/krrishg/service/ProgressService.java @@ -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 emitters = new ConcurrentHashMap<>(); + private final Map>> pendingEvents = new ConcurrentHashMap<>(); + private final Object bufferLock = new Object(); + + public SseEmitter register(UUID siteId) { + SseEmitter emitter = new SseEmitter(SSE_TIMEOUT); + + List> buffer; + synchronized (bufferLock) { + emitters.put(siteId, emitter); + buffer = pendingEvents.remove(siteId); + } + + if (buffer != null) { + for (Map 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 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 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 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)); + } + } + } +} diff --git a/backend/src/main/java/com/krrishg/service/llm/GeminiProvider.java b/backend/src/main/java/com/krrishg/service/llm/GeminiProvider.java index bc2e9f2..42fb3a0 100644 --- a/backend/src/main/java/com/krrishg/service/llm/GeminiProvider.java +++ b/backend/src/main/java/com/krrishg/service/llm/GeminiProvider.java @@ -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 references, LLMOptions options) { long start = System.currentTimeMillis(); - List> 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 systemInstruction = Map.of( - "parts", List.of(Map.of("text", systemPrompt)) - ); - - Map contents = Map.of( - "parts", parts - ); - - Map generationConfig = Map.of( - "temperature", options.getTemperature(), - "maxOutputTokens", options.getMaxTokens() - ); - - Map requestBody = Map.of( - "system_instruction", systemInstruction, - "contents", List.of(contents), - "generationConfig", generationConfig - ); + Map requestBody = buildRequestBody(systemPrompt, userPrompt, references, options); HttpHeaders headers = new HttpHeaders(); headers.setContentType(MediaType.APPLICATION_JSON); - HttpEntity> request = new HttpEntity<>(requestBody, headers); - String url = baseUrl + "/" + model + ":generateContent?key=" + apiKey; - - ResponseEntity response = restTemplate.exchange(url, HttpMethod.POST, request, Map.class); + ResponseEntity response = restTemplate.exchange(buildUrl(), HttpMethod.POST, request, Map.class); long latencyMs = System.currentTimeMillis() - start; - Map 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 references, LLMOptions options, Consumer 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 buildRequestBody(String systemPrompt, String userPrompt, + List references, LLMOptions options) { + List> 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 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 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 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 responseBody) { try { diff --git a/backend/src/main/java/com/krrishg/service/llm/LLMProvider.java b/backend/src/main/java/com/krrishg/service/llm/LLMProvider.java index 5d8bc44..7b0fb60 100644 --- a/backend/src/main/java/com/krrishg/service/llm/LLMProvider.java +++ b/backend/src/main/java/com/krrishg/service/llm/LLMProvider.java @@ -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 references, LLMOptions options); + default LLMResult generate(String systemPrompt, String userPrompt, + List references, LLMOptions options, Consumer onChunk) { + return generate(systemPrompt, userPrompt, references, options); + } + String getName(); } diff --git a/backend/src/main/java/com/krrishg/service/llm/OpenAIProvider.java b/backend/src/main/java/com/krrishg/service/llm/OpenAIProvider.java index 282ed7e..269a0c3 100644 --- a/backend/src/main/java/com/krrishg/service/llm/OpenAIProvider.java +++ b/backend/src/main/java/com/krrishg/service/llm/OpenAIProvider.java @@ -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 references, LLMOptions options) { long start = System.currentTimeMillis(); - Map 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 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 references, LLMOptions options, Consumer 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 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 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 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 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 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 responseBody) { try { diff --git a/backend/src/test/java/com/krrishg/controller/LLMControllerTest.java b/backend/src/test/java/com/krrishg/controller/LLMControllerTest.java index 5f8fe92..53997c2 100644 --- a/backend/src/test/java/com/krrishg/controller/LLMControllerTest.java +++ b/backend/src/test/java/com/krrishg/controller/LLMControllerTest.java @@ -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("

Home

") - .build() - )) - .build(); - when(llmService.generateSite(any(), any(), anyString(), any(), anyString(), any(), any(), any())).thenReturn(structure); - + String siteId = UUID.randomUUID().toString(); Map 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("

Home

") - .build() - )) - .build(); - when(llmService.generateSite(any(), any(), anyString(), any(), anyString(), any(), any(), any())).thenReturn(structure); - + String siteId = UUID.randomUUID().toString(); Map 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("

Home

") - .build() - )) - .build(); - when(llmService.generateSite(any(), any(), anyString(), any(), eq("gemini"), eq("custom-key"), any(), any())).thenReturn(structure); - + String siteId = UUID.randomUUID().toString(); Map 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 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("

Refined

") - .build() - )) - .build(); - when(llmService.refineSite(any(), any(), any(), anyString(), any(), anyString(), any(), any(), any())).thenReturn(refined); - Map 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 diff --git a/backend/src/test/java/com/krrishg/service/LLMResponseParserTest.java b/backend/src/test/java/com/krrishg/service/LLMResponseParserTest.java index ab03e1d..327a5f4 100644 --- a/backend/src/test/java/com/krrishg/service/LLMResponseParserTest.java +++ b/backend/src/test/java/com/krrishg/service/LLMResponseParserTest.java @@ -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()); + } } } diff --git a/backend/src/test/java/com/krrishg/service/LLMServiceTest.java b/backend/src/test/java/com/krrishg/service/LLMServiceTest.java index 451ca5e..691e2e8 100644 --- a/backend/src/test/java/com/krrishg/service/LLMServiceTest.java +++ b/backend/src/test/java/com/krrishg/service/LLMServiceTest.java @@ -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", diff --git a/frontend/src/app/core/services/llm-api.service.ts b/frontend/src/app/core/services/llm-api.service.ts index 54da045..3b378fe 100644 --- a/frontend/src/app/core/services/llm-api.service.ts +++ b/frontend/src/app/core/services/llm-api.service.ts @@ -1,24 +1,40 @@ -import { Injectable } from '@angular/core'; +import { Injectable, inject } from '@angular/core'; import { HttpClient } from '@angular/common/http'; import { Observable } from 'rxjs'; +import { AuthService } from './auth.service'; import { LLMGenerateRequest, LLMRefineRequest, SiteStructure, GenerationVersion, SuggestSectionsRequest, SuggestSectionsResponse } from '../models/llm'; @Injectable({ providedIn: 'root' }) export class LlmApiService { + private authService = inject(AuthService); + constructor(private http: HttpClient) {} - generate(request: LLMGenerateRequest): Observable { - return this.http.post('/api/llm/generate', request); + private withToken(url: string): string { + const token = this.authService.getAccessToken(); + return token ? `${url}?token=${encodeURIComponent(token)}` : url; } - refine(request: LLMRefineRequest): Observable { - return this.http.post('/api/llm/refine', request); + generate(request: LLMGenerateRequest): Observable<{siteId: string}> { + return this.http.post<{siteId: string}>('/api/llm/generate', request); + } + + refine(request: LLMRefineRequest): Observable<{siteId: string}> { + return this.http.post<{siteId: string}>('/api/llm/refine', request); } suggestSections(request: SuggestSectionsRequest): Observable { return this.http.post('/api/llm/suggest-sections', request); } + getGenerationEvents(siteId: string): EventSource { + return new EventSource(this.withToken(`/api/llm/generate/${siteId}/events`)); + } + + getRefineEvents(siteId: string): EventSource { + return new EventSource(this.withToken(`/api/llm/refine/${siteId}/events`)); + } + getVersions(siteId: string): Observable { return this.http.get(`/api/llm/versions/${siteId}`); } diff --git a/frontend/src/app/llm-workspace/generation.service.ts b/frontend/src/app/llm-workspace/generation.service.ts index d99232f..aa79791 100644 --- a/frontend/src/app/llm-workspace/generation.service.ts +++ b/frontend/src/app/llm-workspace/generation.service.ts @@ -1,4 +1,4 @@ -import { Injectable, inject } from '@angular/core'; +import { Injectable, NgZone, inject } from '@angular/core'; import { MatSnackBar } from '@angular/material/snack-bar'; import { BehaviorSubject, Subject, Subscription } from 'rxjs'; import { LlmApiService } from '../core/services/llm-api.service'; @@ -16,8 +16,10 @@ export class GenerationService { private api = inject(LlmApiService); private session = inject(LlmSessionService); private snackBar = inject(MatSnackBar); + private zone = inject(NgZone); private generationSubscription: Subscription | null = null; + private eventSource: EventSource | null = null; private stateSubject = new BehaviorSubject({ status: 'idle', result: null, error: null, }); @@ -31,17 +33,17 @@ export class GenerationService { this.stateSubject.next({ status: 'generating', result: null, error: null }); this.session.setGenerating(true); this.session.clearError(); + this.session.clearProgressMessages(); + this.session.clearStreamingContent(); + + this.openGenerationEventSource(siteId); this.generationSubscription = this.api.generate(request).subscribe({ - next: (data) => { - this.stateSubject.next({ status: 'completed', result: data, error: null }); - this.session.setResult(data); - this.session.setCurrentVersionId(null); - this.refreshVersionsSubject.next(siteId); - this.snackBar.open('Site generated successfully!', 'Close', { duration: 5000 }); + next: () => { this.generationSubscription = null; }, error: (err) => { + this.closeEventSource(); const msg = err.error?.error || err.statusText || 'Failed to connect to server'; this.stateSubject.next({ status: 'error', result: null, error: msg }); this.session.setError(msg); @@ -56,17 +58,17 @@ export class GenerationService { this.stateSubject.next({ status: 'generating', result: null, error: null }); this.session.setGenerating(true); this.session.clearError(); + this.session.clearProgressMessages(); + this.session.clearStreamingContent(); + + this.openRefineEventSource(siteId); this.generationSubscription = this.api.refine(request).subscribe({ - next: (data) => { - this.stateSubject.next({ status: 'completed', result: data, error: null }); - this.session.setResult(data); - this.session.setCurrentVersionId(null); - this.refreshVersionsSubject.next(siteId); - this.snackBar.open('Site refined successfully!', 'Close', { duration: 5000 }); + next: () => { this.generationSubscription = null; }, error: (err) => { + this.closeEventSource(); const msg = err.error?.error || err.statusText || 'Failed to connect to server'; this.stateSubject.next({ status: 'error', result: null, error: msg }); this.session.setError(msg); @@ -76,14 +78,123 @@ export class GenerationService { }); } + private openGenerationEventSource(siteId: string): void { + this.closeEventSource(); + this.eventSource = this.api.getGenerationEvents(siteId); + + this.eventSource.addEventListener('progress', (event: MessageEvent) => { + const data = JSON.parse(event.data); + this.zone.run(() => { + if (data.stage === 'generating_section') { + this.session.setCurrentSection(data.message); + } else { + this.session.addProgressMessage({ stage: data.stage, message: data.message }); + } + }); + }); + + this.eventSource.addEventListener('streaming', (event: MessageEvent) => { + const data = JSON.parse(event.data); + if (data.chunk) { + this.zone.run(() => this.session.appendStreamingContent(data.chunk)); + } + }); + + this.eventSource.addEventListener('complete', (event: MessageEvent) => { + const data = JSON.parse(event.data); + this.zone.run(() => { + this.stateSubject.next({ status: 'completed', result: data, error: null }); + this.session.setResult(data); + this.session.setCurrentVersionId(null); + this.refreshVersionsSubject.next(siteId); + this.snackBar.open('Site generated successfully!', 'Close', { duration: 5000 }); + }); + this.closeEventSource(); + }); + + this.eventSource.addEventListener('error', (event: MessageEvent) => { + let msg = 'Generation failed'; + try { + const data = JSON.parse((event as any).data); + msg = data.error || msg; + } catch {} + this.zone.run(() => { + this.stateSubject.next({ status: 'error', result: null, error: msg }); + this.session.setError(msg); + this.snackBar.open('Generation failed: ' + msg, 'Dismiss', { duration: 8000 }); + }); + this.closeEventSource(); + }); + } + + private openRefineEventSource(siteId: string): void { + this.closeEventSource(); + this.eventSource = this.api.getRefineEvents(siteId); + + this.eventSource.addEventListener('progress', (event: MessageEvent) => { + const data = JSON.parse(event.data); + this.zone.run(() => { + if (data.stage === 'generating_section') { + this.session.setCurrentSection(data.message); + } else { + this.session.addProgressMessage({ stage: data.stage, message: data.message }); + } + }); + }); + + this.eventSource.addEventListener('streaming', (event: MessageEvent) => { + const data = JSON.parse(event.data); + if (data.chunk) { + this.zone.run(() => this.session.appendStreamingContent(data.chunk)); + } + }); + + this.eventSource.addEventListener('complete', (event: MessageEvent) => { + const data = JSON.parse(event.data); + this.zone.run(() => { + this.stateSubject.next({ status: 'completed', result: data, error: null }); + this.session.setResult(data); + this.session.setCurrentVersionId(null); + this.refreshVersionsSubject.next(siteId); + this.snackBar.open('Site refined successfully!', 'Close', { duration: 5000 }); + }); + this.closeEventSource(); + }); + + this.eventSource.addEventListener('error', (event: MessageEvent) => { + let msg = 'Refinement failed'; + try { + const data = JSON.parse((event as any).data); + msg = data.error || msg; + } catch {} + this.zone.run(() => { + this.stateSubject.next({ status: 'error', result: null, error: msg }); + this.session.setError(msg); + this.snackBar.open('Refinement failed: ' + msg, 'Dismiss', { duration: 8000 }); + }); + this.closeEventSource(); + }); + } + + private closeEventSource(): void { + if (this.eventSource) { + this.eventSource.close(); + this.eventSource = null; + } + } + cancel(): void { if (this.generationSubscription) { this.generationSubscription.unsubscribe(); this.generationSubscription = null; } + this.closeEventSource(); this.stateSubject.next({ status: 'idle', result: null, error: null }); this.session.setGenerating(false); this.session.clearError(); + this.session.clearProgressMessages(); + this.session.clearStreamingContent(); + this.session.setCurrentSection(''); } resetState(): void { diff --git a/frontend/src/app/llm-workspace/llm-session.service.ts b/frontend/src/app/llm-workspace/llm-session.service.ts index 1416c52..8810ff2 100644 --- a/frontend/src/app/llm-workspace/llm-session.service.ts +++ b/frontend/src/app/llm-workspace/llm-session.service.ts @@ -2,6 +2,11 @@ import { Injectable } from '@angular/core'; import { BehaviorSubject, Observable } from 'rxjs'; import { SiteStructure, GenerationVersion } from '../core/models/llm'; +export interface ProgressMessage { + stage: string; + message: string; +} + export interface PromptEntry { originalPrompt: string; refinedPrompt: string | null; @@ -19,6 +24,9 @@ export class LlmSessionService { private versionsSubject = new BehaviorSubject([]); private currentVersionIdSubject = new BehaviorSubject(null); private loadingVersionsSubject = new BehaviorSubject(false); + private progressMessagesSubject = new BehaviorSubject([]); + private streamingContentSubject = new BehaviorSubject(''); + private currentSectionSubject = new BehaviorSubject(''); currentPrompt$ = this.currentPromptSubject.asObservable(); currentResult$ = this.currentResultSubject.asObservable(); @@ -28,6 +36,9 @@ export class LlmSessionService { versions$ = this.versionsSubject.asObservable(); currentVersionId$ = this.currentVersionIdSubject.asObservable(); loadingVersions$ = this.loadingVersionsSubject.asObservable(); + progressMessages$ = this.progressMessagesSubject.asObservable(); + streamingContent$ = this.streamingContentSubject.asObservable(); + currentSection$ = this.currentSectionSubject.asObservable(); get currentPrompt(): string { return this.currentPromptSubject.value; } get currentResult(): SiteStructure | null { return this.currentResultSubject.value; } @@ -35,6 +46,9 @@ export class LlmSessionService { get error(): string | null { return this.errorSubject.value; } get versions(): GenerationVersion[] { return this.versionsSubject.value; } get currentVersionId(): string | null { return this.currentVersionIdSubject.value; } + get progressMessages(): ProgressMessage[] { return this.progressMessagesSubject.value; } + get streamingContent(): string { return this.streamingContentSubject.value; } + get currentSection(): string { return this.currentSectionSubject.value; } setPrompt(prompt: string): void { this.currentPromptSubject.next(prompt); @@ -72,12 +86,37 @@ export class LlmSessionService { this.errorSubject.next(null); } + addProgressMessage(msg: ProgressMessage): void { + const msgs = this.progressMessagesSubject.value; + this.progressMessagesSubject.next([...msgs, msg]); + } + + clearProgressMessages(): void { + this.progressMessagesSubject.next([]); + this.currentSectionSubject.next(''); + } + + setCurrentSection(section: string): void { + this.currentSectionSubject.next(section); + } + + appendStreamingContent(chunk: string): void { + this.streamingContentSubject.next(this.streamingContentSubject.value + chunk); + } + + clearStreamingContent(): void { + this.streamingContentSubject.next(''); + } + clearHistory(): void { this.historySubject.next([]); this.currentResultSubject.next(null); this.currentPromptSubject.next(''); this.currentVersionIdSubject.next(null); this.errorSubject.next(null); + this.progressMessagesSubject.next([]); + this.streamingContentSubject.next(''); + this.currentSectionSubject.next(''); } setVersions(versions: GenerationVersion[]): void { diff --git a/frontend/src/app/llm-workspace/llm-workspace.component.ts b/frontend/src/app/llm-workspace/llm-workspace.component.ts index 65f43f5..8e71bcc 100644 --- a/frontend/src/app/llm-workspace/llm-workspace.component.ts +++ b/frontend/src/app/llm-workspace/llm-workspace.component.ts @@ -33,7 +33,7 @@ type Phase = 'idle' | 'suggesting' | 'configuring' | 'generating'; imports: [ NgIf, NgFor, AsyncPipe, KeyValuePipe, DatePipe, FormsModule, RouterModule, MatToolbarModule, MatButtonModule, MatIconModule, - MatInputModule, MatFormFieldModule, MatProgressSpinnerModule, MatSelectModule, + MatInputModule, MatFormFieldModule, MatProgressSpinnerModule, MatSelectModule, MatSnackBarModule, MatCardModule, MatDividerModule, MatExpansionModule, MatDialogModule, SectionConfiguratorComponent, @@ -105,7 +105,6 @@ type Phase = 'idle' | 'suggesting' | 'configuring' | 'generating'; -
@@ -129,7 +128,6 @@ type Phase = 'idle' | 'suggesting' | 'configuring' | 'generating'; >
-

Configure AI Provider

@@ -172,11 +170,17 @@ type Phase = 'idle' | 'suggesting' | 'configuring' | 'generating';
-
+
-
-

Generating your site...

-

This usually takes 15-30 seconds

+
+
+ check_circle + {{ msg.message }} +
+
+ arrow_forward + {{ section }} +