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

add bring your own key support

Browse Source
This commit is contained in:
Krrish Ghimire
2026-07-06 22:45:24 +05:45
parent a6a9ac6e7e
commit 25db5b2970
23 changed files with 950 additions and 353 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -76,11 +76,12 @@ class LLMControllerTest {
.build()
))
.build();
when(llmService.generateSite(any(), any(), anyString(), any())).thenReturn(structure);
when(llmService.generateSite(any(), any(), anyString(), any(), anyString(), any(), any(), any())).thenReturn(structure);
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString()
"siteId", UUID.randomUUID().toString(),
"provider", "gemini"
);
mockMvc.perform(post("/api/llm/generate")
@@ -103,11 +104,12 @@ class LLMControllerTest {
.build()
))
.build();
when(llmService.generateSite(any(), any(), anyString(), any())).thenReturn(structure);
when(llmService.generateSite(any(), any(), anyString(), any(), anyString(), any(), any(), any())).thenReturn(structure);
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString(),
"provider", "gemini",
"references", List.of(
Map.of("type", "IMAGE", "url", "data:img", "altText", "Photo")
)
@@ -120,6 +122,34 @@ class LLMControllerTest {
.andExpect(status().isOk());
}
@Test
void returns200WithPerRequestApiKey() 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);
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString(),
"provider", "gemini",
"apiKey", "custom-key"
);
mockMvc.perform(post("/api/llm/generate")
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isOk());
}
@Test
void returns429WhenRateLimited() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(false);
@@ -127,7 +157,8 @@ class LLMControllerTest {
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString()
"siteId", UUID.randomUUID().toString(),
"provider", "gemini"
);
mockMvc.perform(post("/api/llm/generate")
@@ -142,7 +173,8 @@ class LLMControllerTest {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
Map<String, Object> request = Map.of(
"siteId", UUID.randomUUID().toString()
"siteId", UUID.randomUUID().toString(),
"provider", "gemini"
);
mockMvc.perform(post("/api/llm/generate")
@@ -153,43 +185,12 @@ class LLMControllerTest {
}
@Test
void returns400WhenPromptBlank() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
Map<String, Object> request = Map.of(
"prompt", " ",
"siteId", UUID.randomUUID().toString()
);
mockMvc.perform(post("/api/llm/generate")
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isBadRequest());
}
@Test
void returns400WhenSiteIdMissing() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
Map<String, Object> request = Map.of(
"prompt", "Build a site"
);
mockMvc.perform(post("/api/llm/generate")
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isBadRequest());
}
@Test
void returns400WhenSiteIdBlank() throws Exception {
void returns400WhenProviderMissing() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", ""
"siteId", UUID.randomUUID().toString()
);
mockMvc.perform(post("/api/llm/generate")
@@ -198,6 +199,26 @@ class LLMControllerTest {
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isBadRequest());
}
@Test
void returns400WhenNoKeyConfigured() 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"));
Map<String, Object> request = Map.of(
"prompt", "Build a site",
"siteId", UUID.randomUUID().toString(),
"provider", "gemini"
);
mockMvc.perform(post("/api/llm/generate")
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isBadRequest())
.andExpect(jsonPath("$.error").value("No API key configured for gemini"));
}
}
@Nested
@@ -223,11 +244,12 @@ class LLMControllerTest {
.build()
))
.build();
when(llmService.refineSite(any(), any(), any(), anyString(), any())).thenReturn(refined);
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(),
"provider", "gemini",
"currentStructure", current
);
@@ -247,6 +269,7 @@ class LLMControllerTest {
Map<String, Object> request = Map.of(
"prompt", "Refine",
"siteId", UUID.randomUUID().toString(),
"provider", "gemini",
"currentStructure", Map.of()
);
@@ -258,10 +281,11 @@ class LLMControllerTest {
}
@Test
void returns400WhenPromptMissing() throws Exception {
void returns400WhenProviderMissing() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
Map<String, Object> request = Map.of(
"prompt", "Refine",
"siteId", UUID.randomUUID().toString(),
"currentStructure", SiteStructure.builder().build()
);
@@ -272,38 +296,6 @@ class LLMControllerTest {
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isBadRequest());
}
@Test
void returns400WhenSiteIdMissing() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
Map<String, Object> request = Map.of(
"prompt", "Refine",
"currentStructure", SiteStructure.builder().build()
);
mockMvc.perform(post("/api/llm/refine")
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isBadRequest());
}
@Test
void returns400WhenCurrentStructureMissing() throws Exception {
when(rateLimitConfig.tryConsume(userId)).thenReturn(true);
Map<String, Object> request = Map.of(
"prompt", "Refine",
"siteId", UUID.randomUUID().toString()
);
mockMvc.perform(post("/api/llm/refine")
.with(authentication(auth))
.contentType(MediaType.APPLICATION_JSON)
.content(objectMapper.writeValueAsString(request)))
.andExpect(status().isBadRequest());
}
}
@Nested

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

@@ -3,8 +3,11 @@ package com.krrishg.service;
import com.krrishg.dto.Reference;
import com.krrishg.dto.SiteStructure;
import com.krrishg.model.GenerationVersion;
import com.krrishg.service.llm.LLMFactory;
import com.krrishg.model.User;
import com.krrishg.support.FakeEncryptionService;
import com.krrishg.support.FakeLLMProvider;
import com.krrishg.support.FakeProviderFactory;
import com.krrishg.support.FakeUserRepository;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -19,9 +22,9 @@ import org.springframework.data.mongodb.core.query.Query;
import java.io.ByteArrayInputStream;
import java.nio.charset.StandardCharsets;
import java.lang.reflect.Method;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import static org.junit.jupiter.api.Assertions.*;
@@ -37,11 +40,12 @@ class LLMServiceTest {
@Mock
private ResourceLoader resourceLoader;
private FakeLLMProvider activeProvider;
private FakeLLMProvider fallbackProvider;
private LLMFactory llmFactory;
private FakeLLMProvider geminiProvider;
private FakeProviderFactory providerFactory;
private LLMResponseParser responseParser;
private FakeReferenceService referenceService;
private FakeUserRepository userRepository;
private FakeEncryptionService encryptionService;
private LLMService llmService;
private UUID userId;
@@ -52,15 +56,22 @@ class LLMServiceTest {
userId = UUID.randomUUID();
siteId = UUID.randomUUID();
activeProvider = new FakeLLMProvider("gemini");
fallbackProvider = new FakeLLMProvider("openai");
llmFactory = new LLMFactory(List.of(activeProvider, fallbackProvider), "gemini", "openai");
Method init = LLMFactory.class.getDeclaredMethod("init");
init.setAccessible(true);
init.invoke(llmFactory);
geminiProvider = new FakeLLMProvider("gemini");
providerFactory = new FakeProviderFactory()
.withProvider("gemini", geminiProvider);
responseParser = new LLMResponseParser(new ObjectMapper());
referenceService = new FakeReferenceService();
userRepository = new FakeUserRepository();
encryptionService = new FakeEncryptionService();
User user = User.builder()
.id(userId)
.email("user@test.com")
.name("Test User")
.llmApiKeys(Map.of("gemini", encryptionService.encrypt("test-key")))
.build();
userRepository.save(user);
Resource mockResource = mock(Resource.class);
when(mockResource.exists()).thenReturn(true);
@@ -68,14 +79,15 @@ class LLMServiceTest {
new ByteArrayInputStream("You are a site generator.".getBytes(StandardCharsets.UTF_8)));
when(resourceLoader.getResource(anyString())).thenReturn(mockResource);
llmService = new LLMService(llmFactory, responseParser, referenceService,
mongoTemplate, resourceLoader, "classpath:prompts/site-generator.txt",
llmService = new LLMService(providerFactory, responseParser, referenceService,
mongoTemplate, userRepository, encryptionService,
resourceLoader, "classpath:prompts/site-generator.txt",
"classpath:prompts/site-refiner.txt");
}
@Test
void generateSiteReturnsParsedStructure() {
activeProvider.withContent("""
geminiProvider.withContent("""
{
"pages": [
{"title": "Home", "slug": "home", "rawHtml": "<h1>Welcome</h1>"}
@@ -85,7 +97,7 @@ class LLMServiceTest {
when(mongoTemplate.findOne(any(Query.class), eq(GenerationVersion.class))).thenReturn(null);
SiteStructure result = llmService.generateSite(userId, siteId, "Build a site", null);
SiteStructure result = llmService.generateSite(userId, siteId, "Build a site", null, "gemini", null);
assertEquals(1, result.getPages().size());
assertEquals("Home", result.getPages().get(0).getTitle());
@@ -94,7 +106,7 @@ class LLMServiceTest {
@Test
void generateSiteSavesVersionWithCorrectData() {
activeProvider.withContent("""
geminiProvider.withContent("""
{
"pages": [
{"title": "About", "slug": "about", "rawHtml": "<p>About</p>"}
@@ -104,7 +116,7 @@ class LLMServiceTest {
when(mongoTemplate.findOne(any(Query.class), eq(GenerationVersion.class))).thenReturn(null);
llmService.generateSite(userId, siteId, "Build about page", null);
llmService.generateSite(userId, siteId, "Build about page", null, "gemini", null);
ArgumentCaptor<GenerationVersion> captor = ArgumentCaptor.forClass(GenerationVersion.class);
verify(mongoTemplate).save(captor.capture());
@@ -119,7 +131,7 @@ class LLMServiceTest {
@Test
void generateSiteIncrementsVersionNumber() {
activeProvider.withContent("""
geminiProvider.withContent("""
{
"pages": [
{"title": "Home", "slug": "home", "rawHtml": "<p>Hero</p>"}
@@ -134,7 +146,7 @@ class LLMServiceTest {
.build();
when(mongoTemplate.findOne(any(Query.class), eq(GenerationVersion.class))).thenReturn(existing);
llmService.generateSite(userId, siteId, "Build a site", null);
llmService.generateSite(userId, siteId, "Build a site", null, "gemini", null);
ArgumentCaptor<GenerationVersion> captor = ArgumentCaptor.forClass(GenerationVersion.class);
verify(mongoTemplate).save(captor.capture());
@@ -144,7 +156,7 @@ class LLMServiceTest {
@Test
void generateSiteWithReferences() {
activeProvider.withContent("""
geminiProvider.withContent("""
{
"pages": [
{"title": "Home", "slug": "home", "rawHtml": "<p>Hero</p>"}
@@ -158,7 +170,7 @@ class LLMServiceTest {
when(mongoTemplate.findOne(any(Query.class), eq(GenerationVersion.class))).thenReturn(null);
llmService.generateSite(userId, siteId, "Build a site", references);
llmService.generateSite(userId, siteId, "Build a site", references, "gemini", null);
ArgumentCaptor<GenerationVersion> captor = ArgumentCaptor.forClass(GenerationVersion.class);
verify(mongoTemplate).save(captor.capture());
@@ -169,31 +181,30 @@ class LLMServiceTest {
}
@Test
void generateSiteFallsBackOnProviderFailure() {
activeProvider.setFailure(new RuntimeException("API error"));
fallbackProvider.withContent("""
void generateSiteUsesProvidedApiKeyOverStored() {
geminiProvider.withContent("""
{
"pages": [
{"title": "Fallback", "slug": "fallback", "rawHtml": "<p>Hero</p>"}
{"title": "Home", "slug": "home", "rawHtml": "<p>Hero</p>"}
]
}
""");
when(mongoTemplate.findOne(any(Query.class), eq(GenerationVersion.class))).thenReturn(null);
SiteStructure result = llmService.generateSite(userId, siteId, "Build", null);
llmService.generateSite(userId, siteId, "Build", null, "gemini", "override-key");
assertEquals("Fallback", result.getPages().get(0).getTitle());
verify(mongoTemplate).save(any(GenerationVersion.class));
}
@Test
void generateSiteThrowsWhenAllProvidersFail() {
activeProvider.setFailure(new RuntimeException("API error"));
fallbackProvider.setFailure(new RuntimeException("Fallback also failed"));
void generateSiteThrowsWhenNoKeyConfigured() {
User user = userRepository.findById(userId).orElseThrow();
user.setLlmApiKeys(null);
userRepository.save(user);
assertThrows(RuntimeException.class,
() -> llmService.generateSite(userId, siteId, "Build", null));
assertThrows(IllegalStateException.class,
() -> llmService.generateSite(userId, siteId, "Build", null, "gemini", null));
verify(mongoTemplate, never()).save(any());
}
@@ -209,7 +220,7 @@ class LLMServiceTest {
))
.build();
activeProvider.withContent("""
geminiProvider.withContent("""
{
"pages": [
{"title": "Refined", "slug": "refined", "rawHtml": "<p>Hero</p>"}
@@ -219,7 +230,7 @@ class LLMServiceTest {
when(mongoTemplate.findOne(any(Query.class), eq(GenerationVersion.class))).thenReturn(null);
SiteStructure result = llmService.refineSite(userId, siteId, current, "Make it better", null);
SiteStructure result = llmService.refineSite(userId, siteId, current, "Make it better", null, "gemini", null);
assertEquals("Refined", result.getPages().get(0).getTitle());
verify(mongoTemplate).save(any(GenerationVersion.class));

20
backend/src/test/java/com/krrishg/service/llm/GeminiProviderTest.java
View File

@@ -8,12 +8,10 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.boot.web.client.RestTemplateBuilder;
import org.springframework.http.HttpEntity;
import org.springframework.http.ResponseEntity;
import org.springframework.web.client.RestTemplate;
import java.time.Duration;
import java.util.List;
import java.util.Map;
@@ -34,25 +32,9 @@ class GeminiProviderTest {
@BeforeEach
void setUp() {
RestTemplateBuilder builder = new RestTemplateBuilder()
.setConnectTimeout(Duration.ofSeconds(1));
try {
var field = RestTemplateBuilder.class.getDeclaredField("interceptors");
field.setAccessible(true);
} catch (Exception e) {
}
provider = new GeminiProvider("test-key", "gemini-2.0-flash",
"https://generativelanguage.googleapis.com/v1beta/models",
builder);
try {
var restTemplateField = GeminiProvider.class.getDeclaredField("restTemplate");
restTemplateField.setAccessible(true);
restTemplateField.set(provider, restTemplate);
} catch (Exception e) {
throw new RuntimeException(e);
}
restTemplate);
defaultOptions = LLMOptions.builder().build();
}

89
backend/src/test/java/com/krrishg/service/llm/LLMFactoryTest.java
View File

@@ -1,89 +0,0 @@
package com.krrishg.service.llm;
import com.krrishg.support.FakeLLMProvider;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import java.util.List;
import static org.junit.jupiter.api.Assertions.*;
class LLMFactoryTest {
private FakeLLMProvider gemini;
private FakeLLMProvider openai;
private List<LLMProvider> providers;
@BeforeEach
void setUp() {
gemini = new FakeLLMProvider("gemini");
openai = new FakeLLMProvider("openai");
providers = List.of(gemini, openai);
}
@Test
void getActiveProviderReturnsConfiguredProvider() {
LLMFactory factory = new LLMFactory(providers, "gemini", "openai");
factory.init();
LLMProvider active = factory.getActiveProvider();
assertEquals("gemini", active.getName());
}
@Test
void getActiveProviderThrowsWhenNotFound() {
LLMFactory factory = new LLMFactory(providers, "nonexistent", "openai");
factory.init();
assertThrows(IllegalStateException.class, factory::getActiveProvider);
}
@Test
void getActiveProviderThrowsWhenNoProviders() {
LLMFactory factory = new LLMFactory(List.of(), "gemini", "openai");
factory.init();
assertThrows(IllegalStateException.class, factory::getActiveProvider);
}
@Test
void getFallbackProviderReturnsConfiguredFallback() {
LLMFactory factory = new LLMFactory(providers, "gemini", "openai");
factory.init();
assertTrue(factory.getFallbackProvider().isPresent());
assertEquals("openai", factory.getFallbackProvider().get().getName());
}
@Test
void getFallbackProviderReturnsEmptyWhenNotFound() {
LLMFactory factory = new LLMFactory(providers, "gemini", "nonexistent");
factory.init();
assertTrue(factory.getFallbackProvider().isEmpty());
}
@Test
void getFallbackProviderReturnsEmptyWhenNotConfigured() {
LLMFactory factory = new LLMFactory(providers, "gemini", "");
factory.init();
assertTrue(factory.getFallbackProvider().isEmpty());
}
@Test
void activeAndFallbackCanBeSameProvider() {
LLMFactory factory = new LLMFactory(providers, "gemini", "gemini");
factory.init();
assertEquals("gemini", factory.getActiveProvider().getName());
assertTrue(factory.getFallbackProvider().isPresent());
assertEquals("gemini", factory.getFallbackProvider().get().getName());
}
@Test
void initLogsAvailableProviders() {
LLMFactory factory = new LLMFactory(providers, "gemini", "openai");
assertDoesNotThrow(factory::init);
}
}

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

@@ -7,12 +7,10 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.boot.web.client.RestTemplateBuilder;
import org.springframework.http.HttpEntity;
import org.springframework.http.ResponseEntity;
import org.springframework.web.client.RestTemplate;
import java.time.Duration;
import java.util.List;
import java.util.Map;
@@ -32,19 +30,8 @@ class OpenAIProviderTest {
@BeforeEach
void setUp() {
RestTemplateBuilder builder = new RestTemplateBuilder()
.setConnectTimeout(Duration.ofSeconds(1));
provider = new OpenAIProvider("sk-test", "gpt-4o-mini",
"https://api.openai.com/v1/chat/completions", builder);
try {
var field = OpenAIProvider.class.getDeclaredField("restTemplate");
field.setAccessible(true);
field.set(provider, restTemplate);
} catch (Exception e) {
throw new RuntimeException(e);
}
"https://api.openai.com/v1/chat/completions", restTemplate);
defaultOptions = LLMOptions.builder().build();
}

25
backend/src/test/java/com/krrishg/support/FakeEncryptionService.java Normal file
View File

@@ -0,0 +1,25 @@
package com.krrishg.support;
import com.krrishg.config.EncryptionService;
import java.util.Base64;
public class FakeEncryptionService extends EncryptionService {
public FakeEncryptionService() {
super("config/test-encryption.key");
}
public void init() {
}
@Override
public String encrypt(String plaintext) {
return Base64.getEncoder().encodeToString(plaintext.getBytes());
}
@Override
public String decrypt(String encrypted) {
return new String(Base64.getDecoder().decode(encrypted));
}
}

49
backend/src/test/java/com/krrishg/support/FakeProviderFactory.java Normal file
View File

@@ -0,0 +1,49 @@
package com.krrishg.support;
import com.krrishg.service.llm.LLMProvider;
import com.krrishg.service.llm.ProviderFactory;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class FakeProviderFactory extends ProviderFactory {
private final Map<String, FakeLLMProvider> providers = new HashMap<>();
public FakeProviderFactory() {
super("gemini-2.0-flash", "https://generativelanguage.googleapis.com/v1beta/models",
"gpt-4o-mini", "https://api.openai.com/v1/chat/completions",
null);
}
public FakeProviderFactory withProvider(String name, FakeLLMProvider provider) {
providers.put(name, provider);
return this;
}
@Override
public LLMProvider createProvider(String name, String apiKey) {
return createProvider(name, apiKey, null, null);
}
@Override
public LLMProvider createProvider(String name, String apiKey, String baseUrl) {
return createProvider(name, apiKey, baseUrl, null);
}
@Override
public LLMProvider createProvider(String name, String apiKey, String baseUrl, String model) {
FakeLLMProvider p = providers.get(name);
if (p == null) {
throw new IllegalArgumentException("Unknown LLM provider: " + name
+ ". Available: " + providers.keySet());
}
return p;
}
@Override
public List<String> getAvailableProviders() {
return List.copyOf(providers.keySet());
}
}