Files
bifrost/ui/components/prompts/context.tsx
Beyhan Oğur 880f412e2c first commit
2026-04-26 21:52:23 +03:00

623 lines
22 KiB
TypeScript

import { extractVariablesFromMessages, mergeVariables, Message, MessageRole, MessageType, type VariableMap } from "@/lib/message";
import { getErrorMessage } from "@/lib/store";
import {
useDeleteFolderMutation,
useDeletePromptMutation,
useGetFoldersQuery,
useGetPromptsQuery,
useGetPromptVersionQuery,
useGetSessionsQuery,
useUpdatePromptMutation,
} from "@/lib/store/apis/promptsApi";
import { useGetModelParametersQuery } from "@/lib/store/apis/providersApi";
import { Folder, ModelParams, Prompt, PromptSession, PromptVersion } from "@/lib/types/prompts";
import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib";
import { parseAsInteger, parseAsString, useQueryStates } from "nuqs";
import { createContext, useCallback, useContext, useEffect, useMemo, useRef, useState, type ReactNode } from "react";
import { toast } from "sonner";
import { executePrompt } from "./utils/executor";
interface PromptContextValue {
// Data
folders: Folder[];
prompts: Prompt[];
selectedPrompt?: Prompt;
sessions: PromptSession[];
selectedSession?: PromptSession;
selectedVersion?: PromptVersion;
// Loading states
foldersLoading: boolean;
promptsLoading: boolean;
foldersError: unknown;
promptsError: unknown;
isLoadingPlayground: boolean;
isStreaming: boolean;
// URL state
selectedPromptId: string | null;
selectedSessionId: number | null;
selectedVersionId: number | null;
setUrlState: (state: { promptId?: string | null; sessionId?: number | null; versionId?: number | null }) => void;
// Playground state
messages: Message[];
setMessages: React.Dispatch<React.SetStateAction<Message[]>>;
provider: string;
setProvider: React.Dispatch<React.SetStateAction<string>>;
model: string;
setModel: React.Dispatch<React.SetStateAction<string>>;
modelParams: ModelParams;
setModelParams: React.Dispatch<React.SetStateAction<ModelParams>>;
apiKeyId: string;
setApiKeyId: React.Dispatch<React.SetStateAction<string>>;
// Jinja2 variables
variables: VariableMap;
setVariables: React.Dispatch<React.SetStateAction<VariableMap>>;
// Sheet states
folderSheet: { open: boolean; folder?: Folder };
setFolderSheet: React.Dispatch<React.SetStateAction<{ open: boolean; folder?: Folder }>>;
promptSheet: { open: boolean; prompt?: Prompt; folderId?: string };
setPromptSheet: React.Dispatch<React.SetStateAction<{ open: boolean; prompt?: Prompt; folderId?: string }>>;
commitSheet: { open: boolean; session?: PromptSession };
setCommitSheet: React.Dispatch<React.SetStateAction<{ open: boolean; session?: PromptSession }>>;
// Delete dialog states
deleteFolderDialog: { open: boolean; folder?: Folder };
setDeleteFolderDialog: React.Dispatch<React.SetStateAction<{ open: boolean; folder?: Folder }>>;
deletePromptDialog: { open: boolean; prompt?: Prompt };
setDeletePromptDialog: React.Dispatch<React.SetStateAction<{ open: boolean; prompt?: Prompt }>>;
// Mutation loading states
isDeletingFolder: boolean;
isDeletingPrompt: boolean;
// Model capabilities
supportsVision: boolean;
// Diff detection
hasChanges: boolean;
hasVersionChanges: boolean;
hasSessionChanges: boolean;
// Handlers
handleSelectPrompt: (id: string) => void;
handleMovePrompt: (promptId: string, folderId: string | null) => Promise<void>;
handleDeleteFolder: () => Promise<void>;
handleDeletePrompt: () => Promise<void>;
handleSendMessage: (pendingMessage?: Message) => Promise<void>;
handleSubmitToolResult: (afterIndex: number, toolCallId: string, content: string) => Promise<void>;
// RBAC permissions
canCreate: boolean;
canUpdate: boolean;
canDelete: boolean;
}
const PromptContext = createContext<PromptContextValue | null>(null);
export function usePromptContext() {
const context = useContext(PromptContext);
if (!context) {
throw new Error("usePromptContext must be used within a PromptProvider");
}
return context;
}
export function PromptProvider({ children }: { children: ReactNode }) {
// RBAC permissions
const canCreate = useRbac(RbacResource.PromptRepository, RbacOperation.Create);
const canUpdate = useRbac(RbacResource.PromptRepository, RbacOperation.Update);
const canDelete = useRbac(RbacResource.PromptRepository, RbacOperation.Delete);
// API queries
const { data: foldersData, isLoading: foldersLoading, error: foldersError } = useGetFoldersQuery();
const { data: promptsData, isLoading: promptsLoading, error: promptsError } = useGetPromptsQuery();
// Mutations
const [deleteFolder, { isLoading: isDeletingFolder }] = useDeleteFolderMutation();
const [deletePrompt, { isLoading: isDeletingPrompt }] = useDeletePromptMutation();
const [updatePrompt] = useUpdatePromptMutation();
// UI state — persisted in URL query params
const [{ promptId: selectedPromptId, sessionId: selectedSessionId, versionId: selectedVersionId }, setUrlState] = useQueryStates(
{
promptId: parseAsString,
sessionId: parseAsInteger,
versionId: parseAsInteger,
},
{ history: "replace" },
);
// Sheet states
const [folderSheet, setFolderSheet] = useState<{ open: boolean; folder?: Folder }>({ open: false });
const [promptSheet, setPromptSheet] = useState<{ open: boolean; prompt?: Prompt; folderId?: string }>({ open: false });
const [commitSheet, setCommitSheet] = useState<{ open: boolean; session?: PromptSession }>({ open: false });
// Delete dialog states
const [deleteFolderDialog, setDeleteFolderDialog] = useState<{ open: boolean; folder?: Folder }>({ open: false });
const [deletePromptDialog, setDeletePromptDialog] = useState<{ open: boolean; prompt?: Prompt }>({ open: false });
// Playground state
const [messages, setMessagesRaw] = useState<Message[]>([Message.system("")]);
const setMessages = useCallback<React.Dispatch<React.SetStateAction<Message[]>>>((action) => {
setMessagesRaw((prev) => {
const next = typeof action === "function" ? action(prev) : action;
return next.map((msg, i) => msg.withIndex(i));
});
}, []);
const [provider, setProvider] = useState("");
const [model, setModel] = useState("");
const [modelParams, setModelParams] = useState<ModelParams>({ stream: true });
const [apiKeyId, setApiKeyId] = useState("__auto__");
const [isStreaming, setIsStreaming] = useState(false);
const activeRunRef = useRef<symbol | null>(null);
const [variables, setVariables] = useState<VariableMap>({});
// Fetch model datasheet for capabilities
const { data: datasheetData } = useGetModelParametersQuery(model, { skip: !model });
const supportsVision = datasheetData?.supports_vision ?? false;
// Derived data
const folders = useMemo(() => foldersData?.folders ?? [], [foldersData]);
const prompts = useMemo(() => promptsData?.prompts ?? [], [promptsData]);
const selectedPrompt = useMemo(() => prompts.find((p) => p.id === selectedPromptId), [prompts, selectedPromptId]);
// Fetch versions and sessions for selected prompt
const { data: sessionsData, isLoading: isSessionsLoading } = useGetSessionsQuery(selectedPromptId ?? "", { skip: !selectedPromptId });
// Filter sessions to current prompt — RTK Query may briefly return stale cached data from the previous prompt
const sessions = useMemo(() => {
const all = sessionsData?.sessions ?? [];
if (!selectedPromptId) return [];
return all.filter((s) => s.prompt_id === selectedPromptId);
}, [sessionsData, selectedPromptId]);
const selectedSession = useMemo(() => sessions.find((s) => s.id === selectedSessionId), [sessions, selectedSessionId]);
// Fetch full version data (with messages) when a version is selected
const {
currentData: selectedVersionData,
isLoading: isVersionLoading,
isFetching: isVersionFetching,
} = useGetPromptVersionQuery(selectedVersionId ?? 0, {
skip: !selectedVersionId,
});
const selectedVersion = selectedVersionData?.version;
// Show loader only on initial fetch, not on cache refetches (avoids flicker on save)
const isLoadingPlayground = !!(
selectedPromptId &&
(isSessionsLoading ||
(selectedVersionId && isVersionLoading) ||
// Sessions loaded but auto-select hasn't happened yet
(sessions.length > 0 && !selectedSessionId && !selectedVersionId))
);
// Load session or version data when selection changes
useEffect(() => {
// Don't reset state while waiting for data that hasn't arrived yet
if (selectedSessionId && !selectedSession) return;
if (selectedVersionId && !selectedVersion) return;
const loadFromParams = (params: ModelParams, prov: string, mod: string) => {
const { api_key_id, ...rest } = params || ({} as ModelParams);
setModelParams({ stream: true, ...rest });
setApiKeyId(api_key_id || "__auto__");
setProvider(prov || "");
setModel(mod || "");
};
const loadMessages = (msgs: Message[]) => {
setMessages(msgs);
const varNames = extractVariablesFromMessages(msgs);
setVariables((prev) => mergeVariables(prev, varNames));
};
if (selectedSession) {
const raw = (selectedSession.messages ?? []).map((m) => m.message);
const loaded = Message.fromLegacyAll(raw);
loadMessages(loaded.length > 0 ? loaded : [Message.system("")]);
loadFromParams(selectedSession.model_params, selectedSession.provider, selectedSession.model);
// Restore variables (key:value) from session
if (selectedSession.variables && Object.keys(selectedSession.variables).length > 0) {
setVariables(selectedSession.variables);
}
} else if (selectedVersion) {
// If sessions are still loading and no session is explicitly selected,
// wait — a session may auto-select and take priority
if (isSessionsLoading && !selectedSessionId) return;
const raw = (selectedVersion.messages ?? []).map((m) => m.message);
const loaded = Message.fromLegacyAll(raw);
loadMessages(loaded.length > 0 ? loaded : [Message.system("")]);
loadFromParams(selectedVersion.model_params, selectedVersion.provider, selectedVersion.model);
// Initialize variables from version (keys with empty values)
if (selectedVersion.variables && Object.keys(selectedVersion.variables).length > 0) {
setVariables((prev) => mergeVariables(prev, Object.keys(selectedVersion.variables!)));
}
} else if (selectedPrompt?.latest_version) {
// Only fall back to latest_version after sessions have settled
// to avoid racing with the session auto-select effect
if (isSessionsLoading) return;
const version = selectedPrompt.latest_version;
const raw = (version.messages ?? []).map((m) => m.message);
const loaded = Message.fromLegacyAll(raw);
loadMessages(loaded.length > 0 ? loaded : [Message.system("")]);
loadFromParams(version.model_params, version.provider, version.model);
// Initialize variables from version (keys with empty values)
if (version.variables && Object.keys(version.variables).length > 0) {
setVariables((prev) => mergeVariables(prev, Object.keys(version.variables!)));
}
if (sessions.length === 0) {
setUrlState({ versionId: version.id });
}
} else {
setMessages([Message.system("")]);
setProvider("");
setModel("");
setModelParams({ stream: true });
setApiKeyId("__auto__");
}
}, [
selectedSession,
selectedVersion,
selectedPrompt,
selectedSessionId,
selectedVersionId,
setUrlState,
isSessionsLoading,
sessions.length,
]);
// Auto-select the most recent session when sessions load and none is selected
// Sessions take priority over versions for initial loading
useEffect(() => {
if (sessions.length > 0 && !selectedSessionId && !selectedVersionId) {
setUrlState({ sessionId: sessions[0].id });
}
}, [selectedPromptId, sessions, selectedSessionId, selectedVersionId, setUrlState]);
// Diff detection helper — compares current playground state against a reference config
const diffAgainst = useCallback(
(ref: { messages?: any[]; model_params?: ModelParams; provider?: string; model?: string } | undefined) => {
if (!ref) return true; // No reference — treat as changed
const refMessages = ref.messages ?? [];
const refProvider = ref.provider;
const refModel = ref.model;
const refParams = ref.model_params;
if (provider !== refProvider) return true;
if (model !== refModel) return true;
const { api_key_id: refApiKeyId, ...refParamsRest } = refParams || ({} as ModelParams);
const currentApiKeyId = apiKeyId !== "__auto__" ? apiKeyId : undefined;
if (currentApiKeyId !== (refApiKeyId || undefined)) return true;
// Normalize: treat missing stream as stream: true so legacy params without stream don't appear changed
const normalizeParams = (p: ModelParams): ModelParams => {
const { stream = true, ...rest } = p;
return { stream, ...rest };
};
const normalizedCurrent = normalizeParams(modelParams);
const normalizedRef = normalizeParams(refParamsRest);
if (
JSON.stringify(normalizedCurrent, Object.keys(normalizedCurrent).sort()) !==
JSON.stringify(normalizedRef, Object.keys(normalizedRef).sort())
)
return true;
const currentSerialized = Message.serializeAll(messages);
if (JSON.stringify(currentSerialized) !== JSON.stringify(refMessages)) return true;
return false;
},
[provider, model, modelParams, apiKeyId, messages],
);
// Diff detection — compare current playground state against the loaded session/version
const hasChanges = useMemo(() => {
// Suppress diff while version data is in flight to avoid flicker
if (selectedVersionId && (isVersionFetching || selectedVersion?.id !== selectedVersionId)) return false;
if (selectedSession) {
return diffAgainst({
messages: selectedSession.messages?.map((m) => m.message) ?? [],
model_params: selectedSession.model_params,
provider: selectedSession.provider,
model: selectedSession.model,
});
}
if (selectedVersion) {
return diffAgainst({
messages: selectedVersion.messages?.map((m) => m.message) ?? [],
model_params: selectedVersion.model_params,
provider: selectedVersion.provider,
model: selectedVersion.model,
});
}
return true;
}, [selectedSession, selectedVersion, diffAgainst, selectedVersionId, isVersionFetching]);
// Diff against the active version — drives "unpublished changes" badge & commit button
// Uses the explicitly selected version if available, otherwise falls back to latest_version
const activeVersionRef = selectedVersion ?? selectedPrompt?.latest_version;
const hasVersionChanges = useMemo(() => {
// Suppress diff while version data is in flight or mismatched to avoid flash
if (selectedVersionId && (isVersionFetching || selectedVersion?.id !== selectedVersionId)) return false;
if (!activeVersionRef) return true; // No versions yet — always allow commit
return diffAgainst({
messages: activeVersionRef.messages?.map((m) => m.message) ?? [],
model_params: activeVersionRef.model_params,
provider: activeVersionRef.provider,
model: activeVersionRef.model,
});
}, [activeVersionRef, diffAgainst, selectedVersionId, isVersionFetching, selectedVersion?.id]);
// Diff against the selected session — drives red asterisk indicator
const hasSessionChanges = useMemo(() => {
if (!selectedSession) return false;
return diffAgainst({
messages: selectedSession.messages?.map((m) => m.message) ?? [],
model_params: selectedSession.model_params,
provider: selectedSession.provider,
model: selectedSession.model,
});
}, [selectedSession, diffAgainst]);
// Handlers
const handleSelectPrompt = useCallback(
(id: string) => {
setMessages([Message.system("")]);
setProvider("");
setModel("");
setModelParams({ stream: true });
setApiKeyId("__auto__");
setUrlState({ promptId: id, sessionId: null, versionId: null });
},
[setUrlState],
);
const handleMovePrompt = useCallback(
async (promptId: string, folderId: string | null) => {
try {
await updatePrompt({ id: promptId, data: { folder_id: folderId } }).unwrap();
toast.success("Prompt moved successfully");
} catch (err) {
toast.error(getErrorMessage(err) || "Failed to move prompt");
}
},
[updatePrompt],
);
const handleDeleteFolder = useCallback(async () => {
if (!deleteFolderDialog.folder) return;
try {
await deleteFolder(deleteFolderDialog.folder.id).unwrap();
toast.success("Folder deleted");
setDeleteFolderDialog({ open: false });
if (selectedPrompt?.folder_id === deleteFolderDialog.folder.id) {
setUrlState({ promptId: null, sessionId: null, versionId: null });
}
} catch (err) {
toast.error("Failed to delete folder", { description: getErrorMessage(err) });
}
}, [deleteFolderDialog.folder, deleteFolder, selectedPrompt, setUrlState]);
const handleDeletePrompt = useCallback(async () => {
if (!deletePromptDialog.prompt) return;
try {
await deletePrompt(deletePromptDialog.prompt.id).unwrap();
toast.success("Prompt deleted");
setDeletePromptDialog({ open: false });
if (selectedPromptId === deletePromptDialog.prompt.id) {
setUrlState({ promptId: null, sessionId: null, versionId: null });
}
} catch (err) {
toast.error("Failed to delete prompt", { description: getErrorMessage(err) });
}
}, [deletePromptDialog.prompt, deletePrompt, selectedPromptId, setUrlState]);
const handleSendMessage = useCallback(
async (pendingMessage?: Message) => {
const runToken = Symbol();
activeRunRef.current = runToken;
const isActive = () => activeRunRef.current === runToken;
setIsStreaming(true);
await executePrompt(
messages,
pendingMessage,
{ provider, model, modelParams, apiKeyId, variables },
{
onStreamingStart: (allMessages, placeholder) => {
if (!isActive()) return;
setMessages([...allMessages, placeholder]);
},
onStreamChunk: (content) => {
if (!isActive()) return;
setMessages((prev) => {
const updated = [...prev];
const last = updated[updated.length - 1];
const clone = last.clone();
clone.content = content;
updated[updated.length - 1] = clone;
return updated;
});
},
onComplete: (content, usage) => {
if (!isActive()) return;
setMessages((prev) => {
const updated = [...prev];
updated[updated.length - 1] = Message.response(content, 0, usage);
return updated;
});
},
onToolCallComplete: (content, toolCalls, usage) => {
if (!isActive()) return;
setMessages((prev) => {
const updated = [...prev];
updated[updated.length - 1] = Message.toolCallResponse(content, toolCalls, 0, usage);
return updated;
});
},
onEmptyResponse: () => {
if (!isActive()) return;
setMessages((prev) => prev.slice(0, -1));
},
onError: (error) => {
if (!isActive()) return;
setMessages((prev) => {
const withoutPlaceholder = prev.slice(0, -1);
return [...withoutPlaceholder, Message.error(error)];
});
},
onFinally: () => {
if (!isActive()) return;
setIsStreaming(false);
},
},
);
},
[messages, provider, model, modelParams, apiKeyId, variables],
);
const handleSubmitToolResult = useCallback(
async (afterIndex: number, toolCallId: string, content: string) => {
const runToken = Symbol();
activeRunRef.current = runToken;
const isActive = () => activeRunRef.current === runToken;
const toolResultMsg = new Message(crypto.randomUUID(), 0, MessageType.ToolResult, {
role: MessageRole.TOOL,
content,
tool_call_id: toolCallId,
});
const newMessages = [...messages];
// Insert after any existing tool results that follow the assistant message
let insertAt = afterIndex + 1;
while (insertAt < newMessages.length && newMessages[insertAt].type === MessageType.ToolResult) {
insertAt++;
}
newMessages.splice(insertAt, 0, toolResultMsg);
setMessages(newMessages);
// Execute with the updated messages
setIsStreaming(true);
await executePrompt(
newMessages,
undefined,
{ provider, model, modelParams, apiKeyId, variables },
{
onStreamingStart: (allMessages, placeholder) => {
if (!isActive()) return;
setMessages([...allMessages, placeholder]);
},
onStreamChunk: (content) => {
if (!isActive()) return;
setMessages((prev) => {
const updated = [...prev];
const last = updated[updated.length - 1];
const clone = last.clone();
clone.content = content;
updated[updated.length - 1] = clone;
return updated;
});
},
onComplete: (content, usage) => {
if (!isActive()) return;
setMessages((prev) => {
const updated = [...prev];
updated[updated.length - 1] = Message.response(content, 0, usage);
return updated;
});
},
onToolCallComplete: (content, toolCalls, usage) => {
if (!isActive()) return;
setMessages((prev) => {
const updated = [...prev];
updated[updated.length - 1] = Message.toolCallResponse(content, toolCalls, 0, usage);
return updated;
});
},
onEmptyResponse: () => {
if (!isActive()) return;
setMessages((prev) => prev.slice(0, -1));
},
onError: (error) => {
if (!isActive()) return;
setMessages((prev) => {
const withoutPlaceholder = prev.slice(0, -1);
return [...withoutPlaceholder, Message.error(error)];
});
},
onFinally: () => {
if (!isActive()) return;
setIsStreaming(false);
},
},
);
},
[messages, provider, model, modelParams, apiKeyId, variables],
);
const value: PromptContextValue = {
folders,
prompts,
selectedPrompt,
sessions,
selectedSession,
selectedVersion,
foldersLoading,
promptsLoading,
foldersError,
promptsError,
isLoadingPlayground,
isStreaming,
selectedPromptId,
selectedSessionId,
selectedVersionId,
setUrlState,
messages,
setMessages,
provider,
setProvider,
model,
setModel,
modelParams,
setModelParams,
apiKeyId,
setApiKeyId,
variables,
setVariables,
folderSheet,
setFolderSheet,
promptSheet,
setPromptSheet,
commitSheet,
setCommitSheet,
deleteFolderDialog,
setDeleteFolderDialog,
deletePromptDialog,
setDeletePromptDialog,
isDeletingFolder,
isDeletingPrompt,
supportsVision,
hasChanges,
hasVersionChanges,
hasSessionChanges,
handleSelectPrompt,
handleMovePrompt,
handleDeleteFolder,
handleDeletePrompt,
handleSendMessage,
handleSubmitToolResult,
canCreate,
canUpdate,
canDelete,
};
return <PromptContext.Provider value={value}>{children}</PromptContext.Provider>;
}