Skip to content

Commit b55499d

Browse files
pk910pk910
committed
allow client side AI key
1 parent 1d3245b commit b55499d

File tree

8 files changed

+695
-73
lines changed

8 files changed

+695
-73
lines changed

‎pkg/web/api/ai_handler.go‎

Lines changed: 59 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -69,9 +69,10 @@ func (h *AIHandler) checkAuth(r *http.Request) bool {
6969

7070
// AIConfigResponse is the response for GetConfig endpoint.
7171
type AIConfigResponse struct {
72-
Enabled bool `json:"enabled"`
73-
DefaultModel string `json:"defaultModel"`
74-
AllowedModels []string `json:"allowedModels"`
72+
Enabled bool `json:"enabled"`
73+
DefaultModel string `json:"defaultModel"`
74+
AllowedModels []string `json:"allowedModels"`
75+
ServerKeyConfigured bool `json:"serverKeyConfigured"`
7576
}
7677

7778
// GetConfig returns AI configuration (enabled, available models).
@@ -80,20 +81,72 @@ func (h *AIHandler) GetConfig(w http.ResponseWriter, _ *http.Request) {
8081
w.Header().Set("Content-Type", contentTypeJSON)
8182

8283
response := AIConfigResponse{
83-
Enabled: false,
84-
DefaultModel: "",
85-
AllowedModels: []string{},
84+
Enabled: false,
85+
DefaultModel: "",
86+
AllowedModels: []string{},
87+
ServerKeyConfigured: false,
8688
}
8789

8890
if h.config != nil && h.config.Enabled {
8991
response.Enabled = true
9092
response.DefaultModel = h.config.DefaultModel
9193
response.AllowedModels = h.config.AllowedModels
94+
response.ServerKeyConfigured = h.client != nil
9295
}
9396

9497
h.sendOKResponse(w, "/api/v1/ai/config", response)
9598
}
9699

100+
// SystemPromptResponse is the response for GetSystemPrompt endpoint.
101+
type SystemPromptResponse struct {
102+
Prompt string `json:"prompt"`
103+
}
104+
105+
// GetSystemPrompt returns the AI system prompt text.
106+
// GET /api/v1/ai/system_prompt
107+
func (h *AIHandler) GetSystemPrompt(w http.ResponseWriter, _ *http.Request) {
108+
w.Header().Set("Content-Type", contentTypeJSON)
109+
110+
prompt := ai.BuildSystemPrompt()
111+
h.sendOKResponse(w, "/api/v1/ai/system_prompt", SystemPromptResponse{
112+
Prompt: prompt,
113+
})
114+
}
115+
116+
// ValidateYamlRequest is the request body for ValidateYaml endpoint.
117+
type ValidateYamlRequest struct {
118+
Yaml string `json:"yaml"`
119+
}
120+
121+
// ValidateYaml validates AI-generated YAML and returns validation results.
122+
// POST /api/v1/ai/validate
123+
func (h *AIHandler) ValidateYaml(w http.ResponseWriter, r *http.Request) {
124+
w.Header().Set("Content-Type", contentTypeJSON)
125+
126+
body, err := io.ReadAll(r.Body)
127+
if err != nil {
128+
h.sendErrorResponse(
129+
w, "/api/v1/ai/validate", "failed to read request body", http.StatusBadRequest,
130+
)
131+
132+
return
133+
}
134+
135+
var req ValidateYamlRequest
136+
137+
err = json.Unmarshal(body, &req)
138+
if err != nil {
139+
h.sendErrorResponse(
140+
w, "/api/v1/ai/validate", "invalid request body", http.StatusBadRequest,
141+
)
142+
143+
return
144+
}
145+
146+
result := ai.ValidateGeneratedYaml(req.Yaml)
147+
h.sendOKResponse(w, "/api/v1/ai/validate", result)
148+
}
149+
97150
// AIUsageResponse is the response for GetUsage endpoint.
98151
type AIUsageResponse struct {
99152
LastDay *db.AIUsageStats `json:"lastDay"`

‎pkg/web/server.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,8 @@ func (ws *Server) ConfigureRoutes(frontendConfig *types.FrontendConfig, apiConfi
160160
disableAuth,
161161
)
162162
ws.router.HandleFunc("/api/v1/ai/config", aiHandler.GetConfig).Methods("GET")
163+
ws.router.HandleFunc("/api/v1/ai/system_prompt", aiHandler.GetSystemPrompt).Methods("GET")
164+
ws.router.HandleFunc("/api/v1/ai/validate", aiHandler.ValidateYaml).Methods("POST")
163165
ws.router.HandleFunc("/api/v1/ai/usage", aiHandler.GetUsage).Methods("GET")
164166
ws.router.HandleFunc("/api/v1/ai/chat", aiHandler.Chat).Methods("POST")
165167
ws.router.HandleFunc("/api/v1/ai/chat/{sessionId}", aiHandler.GetChatSession).Methods("GET")

‎web-ui/src/api/ai.ts‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ export interface AIConfig {
77
enabled: boolean;
88
defaultModel: string;
99
allowedModels: string[];
10+
serverKeyConfigured: boolean;
1011
}
1112

1213
export interface AIUsageStats {
@@ -208,3 +209,160 @@ export async function pollAISession(
208209

209210
throw new Error('Session polling timed out');
210211
}
212+
213+
// Get AI system prompt (for client-side mode)
214+
export async function fetchSystemPrompt(): Promise<string> {
215+
const data = await fetchAIApi<{ prompt: string }>('/ai/system_prompt');
216+
return data.prompt;
217+
}
218+
219+
// Validate YAML via backend (for client-side mode)
220+
export async function validateYaml(yaml: string): Promise<ValidationResult> {
221+
return fetchAIApi<ValidationResult>('/ai/validate', {
222+
method: 'POST',
223+
body: JSON.stringify({ yaml }),
224+
});
225+
}
226+
227+
// Extract YAML code blocks from AI response
228+
export function extractYamlFromResponse(response: string): string {
229+
const re = /```ya?ml\s*\n([\s\S]*?)```/;
230+
const match = re.exec(response);
231+
return match ? match[1] : '';
232+
}
233+
234+
// Build a fix prompt for validation issues (mirrors Go logic)
235+
export function buildFixPrompt(
236+
brokenYaml: string,
237+
validation: ValidationResult,
238+
warningsOnly: boolean
239+
): string {
240+
let prompt: string;
241+
242+
if (warningsOnly) {
243+
prompt =
244+
'The YAML you generated has validation warnings. Please review and fix the following issues if appropriate:\n\n';
245+
for (const issue of validation.issues ?? []) {
246+
if (issue.type === 'warning') {
247+
prompt += `- ${issue.path}: ${issue.message}\n`;
248+
}
249+
}
250+
} else {
251+
prompt = 'The YAML you generated has validation errors. Please fix the following issues:\n\n';
252+
for (const issue of validation.issues ?? []) {
253+
if (issue.type === 'error') {
254+
prompt += `- ${issue.path}: ${issue.message}\n`;
255+
}
256+
}
257+
}
258+
259+
prompt += '\nHere is the YAML:\n```yaml\n' + brokenYaml + '\n```\n\n';
260+
261+
if (warningsOnly) {
262+
prompt +=
263+
'Please provide an improved version of the YAML that addresses the warnings where appropriate. ';
264+
} else {
265+
prompt += 'Please provide a corrected version of the YAML that fixes all the errors. ';
266+
}
267+
268+
prompt += 'Only output the corrected YAML in a code block, no explanations needed.';
269+
return prompt;
270+
}
271+
272+
// Stream chat directly to OpenRouter from the browser (client-side mode)
273+
export async function clientSideStreamChat(
274+
apiKey: string,
275+
model: string,
276+
messages: ChatMessage[],
277+
maxTokens: number,
278+
onChunk: (chunk: string) => void
279+
): Promise<string> {
280+
const response = await fetch('https://openrouter.ai/api/v1/chat/completions', {
281+
method: 'POST',
282+
headers: {
283+
Authorization: `Bearer ${apiKey}`,
284+
'Content-Type': 'application/json',
285+
'HTTP-Referer': window.location.origin,
286+
},
287+
body: JSON.stringify({
288+
model,
289+
messages,
290+
max_tokens: maxTokens,
291+
stream: true,
292+
}),
293+
});
294+
295+
if (!response.ok) {
296+
const errorBody = await response.text();
297+
throw new Error(`OpenRouter API error: ${response.status} - ${errorBody}`);
298+
}
299+
300+
const reader = response.body?.getReader();
301+
if (!reader) {
302+
throw new Error('No response body');
303+
}
304+
305+
const decoder = new TextDecoder();
306+
let fullResponse = '';
307+
let buffer = '';
308+
309+
for (;;) {
310+
const { done, value } = await reader.read();
311+
if (done) break;
312+
313+
buffer += decoder.decode(value, { stream: true });
314+
const lines = buffer.split('\n');
315+
buffer = lines.pop() ?? '';
316+
317+
for (const line of lines) {
318+
const trimmed = line.trim();
319+
if (!trimmed.startsWith('data: ')) continue;
320+
321+
const data = trimmed.slice(6);
322+
if (data === '[DONE]') continue;
323+
324+
try {
325+
const parsed = JSON.parse(data);
326+
const content = parsed.choices?.[0]?.delta?.content;
327+
if (content) {
328+
fullResponse += content;
329+
onChunk(content);
330+
}
331+
} catch {
332+
// skip malformed SSE chunks
333+
}
334+
}
335+
}
336+
337+
return fullResponse;
338+
}
339+
340+
// Non-streaming chat call to OpenRouter (for fix attempts)
341+
export async function clientSideChat(
342+
apiKey: string,
343+
model: string,
344+
messages: ChatMessage[],
345+
maxTokens: number
346+
): Promise<string> {
347+
const response = await fetch('https://openrouter.ai/api/v1/chat/completions', {
348+
method: 'POST',
349+
headers: {
350+
Authorization: `Bearer ${apiKey}`,
351+
'Content-Type': 'application/json',
352+
'HTTP-Referer': window.location.origin,
353+
},
354+
body: JSON.stringify({
355+
model,
356+
messages,
357+
max_tokens: maxTokens,
358+
}),
359+
});
360+
361+
if (!response.ok) {
362+
const errorBody = await response.text();
363+
throw new Error(`OpenRouter API error: ${response.status} - ${errorBody}`);
364+
}
365+
366+
const result = await response.json();
367+
return result.choices?.[0]?.message?.content ?? '';
368+
}

‎web-ui/src/components/builder/ai/AIPanel.tsx‎

Lines changed: 48 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
1-
import React, { useEffect } from 'react';
1+
import React, { useEffect, useState } from 'react';
22
import { useAIStore } from '../../../stores/aiStore';
33
import { ModelSelector } from './ModelSelector';
44
import { TokenUsageDisplay } from './TokenUsageDisplay';
55
import { ChatHistory } from './ChatHistory';
66
import { MessageInput } from './MessageInput';
77
import { YamlPreview } from './YamlPreview';
8+
import { ApiKeyInput } from './ApiKeyInput';
89

910
interface AIPanelProps {
1011
testName: string;
@@ -30,8 +31,13 @@ export const AIPanel: React.FC<AIPanelProps> = ({
3031
error,
3132
clearError,
3233
clearMessages,
34+
userApiKey,
35+
isAvailable,
36+
isClientSide,
3337
} = useAIStore();
3438

39+
const [showApiKeyInput, setShowApiKeyInput] = useState(false);
40+
3541
useEffect(() => {
3642
if (!config && !configLoading) {
3743
loadConfig();
@@ -49,13 +55,16 @@ export const AIPanel: React.FC<AIPanelProps> = ({
4955
setPendingYaml(null);
5056
};
5157

58+
const aiAvailable = isAvailable();
59+
const clientSide = isClientSide();
60+
5261
return (
5362
<div className="flex flex-col h-full bg-[var(--color-bg-secondary)] overflow-hidden">
5463
{/* Header */}
5564
<div className="flex items-center justify-between px-4 py-3 border-b border-[var(--color-border)] bg-[var(--color-bg-tertiary)]">
5665
<div className="flex items-center gap-2">
5766
<svg
58-
className="w-5 h-5 text-blue-500"
67+
className="size-5 text-blue-500"
5968
fill="none"
6069
viewBox="0 0 24 24"
6170
stroke="currentColor"
@@ -75,7 +84,7 @@ export const AIPanel: React.FC<AIPanelProps> = ({
7584
className="p-1 text-[var(--color-text-tertiary)] hover:text-[var(--color-text-secondary)]"
7685
title="Clear chat"
7786
>
78-
<svg className="w-5 h-5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
87+
<svg className="size-5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
7988
<path
8089
strokeLinecap="round"
8190
strokeLinejoin="round"
@@ -88,7 +97,7 @@ export const AIPanel: React.FC<AIPanelProps> = ({
8897
onClick={onClose}
8998
className="p-1 text-[var(--color-text-tertiary)] hover:text-[var(--color-text-secondary)]"
9099
>
91-
<svg className="w-5 h-5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
100+
<svg className="size-5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
92101
<path
93102
strokeLinecap="round"
94103
strokeLinejoin="round"
@@ -132,13 +141,44 @@ export const AIPanel: React.FC<AIPanelProps> = ({
132141
</div>
133142
)}
134143

144+
{/* Enabled but no key available - prompt user to enter key */}
145+
{config && config.enabled && !aiAvailable && (
146+
<div className="flex-1 flex items-center justify-center p-4">
147+
<div className="text-center max-w-sm">
148+
<p className="text-[var(--color-text-secondary)] mb-3">
149+
No server API key is configured. Enter your own OpenRouter API key to use AI
150+
features.
151+
</p>
152+
<p className="text-xs text-[var(--color-text-tertiary)] mb-4">
153+
Your key is stored locally in your browser and never sent to this server.
154+
</p>
155+
<ApiKeyInput />
156+
</div>
157+
</div>
158+
)}
159+
135160
{/* Main content */}
136-
{config && config.enabled && (
161+
{config && config.enabled && aiAvailable && (
137162
<>
138163
{/* Config bar */}
139-
<div className="px-3 py-2 border-b border-[var(--color-border)] space-y-2">
164+
<div className="px-3 py-2 border-b border-[var(--color-border)] flex flex-col gap-2">
140165
<ModelSelector />
141-
<TokenUsageDisplay />
166+
{/* Always show ApiKeyInput when no server key (required) or user already set a key */}
167+
{(!config.serverKeyConfigured || userApiKey) && <ApiKeyInput />}
168+
{/* When server key exists and no user key, show optional toggle */}
169+
{config.serverKeyConfigured && !userApiKey && (
170+
showApiKeyInput ? (
171+
<ApiKeyInput />
172+
) : (
173+
<button
174+
onClick={() => setShowApiKeyInput(true)}
175+
className="text-xs text-[var(--color-text-tertiary)] hover:text-[var(--color-text-secondary)] text-left"
176+
>
177+
Use your own OpenRouter API key instead...
178+
</button>
179+
)
180+
)}
181+
{!clientSide && <TokenUsageDisplay />}
142182
</div>
143183

144184
{/* Error banner */}
@@ -150,7 +190,7 @@ export const AIPanel: React.FC<AIPanelProps> = ({
150190
onClick={clearError}
151191
className="ml-2 text-red-400 hover:text-red-600"
152192
>
153-
<svg className="w-4 h-4" fill="none" viewBox="0 0 24 24" stroke="currentColor">
193+
<svg className="size-4" fill="none" viewBox="0 0 24 24" stroke="currentColor">
154194
<path
155195
strokeLinecap="round"
156196
strokeLinejoin="round"

0 commit comments

Comments
 (0)