diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 7de7373b20b..593a8167cc4 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -702,15 +702,9 @@ } }, "src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx": { - "local/no-complex-jsx-arrow": { - "count": 1 - }, "max-lines": { "count": 1 }, - "no-nested-ternary": { - "count": 6 - }, "react-hooks/set-state-in-effect": { "count": 4 } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index f07b66efdf8..e82eab63f7e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -490,6 +490,49 @@ describe("ChatUI", () => { expect(customProxyInput).toHaveValue(testProxyUrl); }); + it("revokes the active OCR preview URL on unmount", async () => { + const createObjectURL = vi.fn(() => "blob:http://localhost/ocr-preview"); + const revokeObjectURL = vi.fn(); + URL.createObjectURL = createObjectURL; + URL.revokeObjectURL = revokeObjectURL; + + const { getByText, unmount } = render( + , + ); + + await waitFor(() => { + expect(getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/ocr"); + + await waitFor(() => { + expect(screen.getByText("Click or drag a document or image to upload")).toBeInTheDocument(); + }); + + const uploadInput = screen.getByLabelText(/Click or drag a document or image to upload/); + + const imageFile = new File(["ocr"], "ocr.png", { type: "image/png" }); + await act(async () => { + fireEvent.change(uploadInput, { target: { files: [imageFile] } }); + }); + + await waitFor(() => { + expect(createObjectURL).toHaveBeenCalledWith(imageFile); + expect(screen.getByAltText("Upload preview")).toBeInTheDocument(); + }); + + unmount(); + + expect(revokeObjectURL).toHaveBeenCalledWith("blob:http://localhost/ocr-preview"); + }); + it("should enable search functionality for MCP server selector", async () => { const user = userEvent.setup(); render( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index e6257907918..8e77961faae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -45,6 +45,7 @@ import { Agent, fetchAvailableAgents } from "../../llm_calls/fetch_agents"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import { makeOpenAIImageEditsRequest } from "../../llm_calls/image_edits"; import { makeOpenAIImageGenerationRequest } from "../../llm_calls/image_generation"; +import { makeOpenAIOcrRequest } from "../../llm_calls/OcrApi"; import { makeOpenAIResponsesRequest } from "@/components/llm_calls/responses_api"; import { makeInteractionsRequest } from "../../llm_calls/interactions_api"; import AdditionalModelSettings from "./AdditionalModelSettings"; @@ -115,6 +116,63 @@ const MCP_SUPPORTED_ENDPOINTS = new Set([ const CUSTOM_MODEL_DEBOUNCE_WAIT_MS = 500; +interface SendDisabledOptions { + endpointType: EndpointType; + inputMessage: string; + isLoading: boolean; + selectedMCPDirectTool: string | undefined; + selectedMCPServers: readonly string[]; + uploadedAudio: File | null; + uploadedOcrFile: File | null; +} + +const getInputPlaceholder = (endpointType: EndpointType): string => { + switch (endpointType) { + case EndpointType.CHAT: + case EndpointType.EMBEDDINGS: + case EndpointType.RESPONSES: + case EndpointType.ANTHROPIC_MESSAGES: + case EndpointType.INTERACTIONS: + return "Type your message... (Shift+Enter for new line)"; + case EndpointType.A2A_AGENTS: + return "Send a message to the A2A agent..."; + case EndpointType.IMAGE_EDITS: + return "Describe how you want to edit the image..."; + case EndpointType.SPEECH: + return "Enter text to convert to speech..."; + case EndpointType.TRANSCRIPTION: + return "Optional: Add context or prompt for transcription..."; + case EndpointType.OCR: + return "Upload a document or image to run OCR"; + default: + return "Describe the image you want to generate..."; + } +}; + +const getSendDisabled = ({ + endpointType, + inputMessage, + isLoading, + selectedMCPDirectTool, + selectedMCPServers, + uploadedAudio, + uploadedOcrFile, +}: SendDisabledOptions): boolean => { + if (isLoading) { + return true; + } + if (endpointType === EndpointType.MCP) { + return !(selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" && selectedMCPDirectTool); + } + if (endpointType === EndpointType.TRANSCRIPTION) { + return !uploadedAudio; + } + if (endpointType === EndpointType.OCR) { + return !uploadedOcrFile; + } + return !inputMessage.trim(); +}; + const ChatUI: React.FC = ({ accessToken, token, @@ -263,6 +321,8 @@ const ChatUI: React.FC = ({ const [chatUploadedImage, setChatUploadedImage] = useState(null); const [chatImagePreviewUrl, setChatImagePreviewUrl] = useState(null); const [uploadedAudio, setUploadedAudio] = useState(null); + const [uploadedOcrFile, setUploadedOcrFile] = useState(null); + const [ocrFilePreviewUrl, setOcrFilePreviewUrl] = useState(null); const [isGetCodeModalVisible, setIsGetCodeModalVisible] = useState(false); const [generatedCode, setGeneratedCode] = useState(""); const [selectedSdk, setSelectedSdk] = useState<"openai" | "azure">("openai"); @@ -280,6 +340,24 @@ const ChatUI: React.FC = ({ const codeInterpreter = useCodeInterpreter(); const chatEndRef = useRef(null); + const ocrFilePreviewUrlRef = useRef(null); + + const revokeOcrFilePreviewUrl = () => { + if (ocrFilePreviewUrlRef.current) { + URL.revokeObjectURL(ocrFilePreviewUrlRef.current); + } + ocrFilePreviewUrlRef.current = null; + setOcrFilePreviewUrl(null); + }; + + useEffect(() => { + return () => { + if (ocrFilePreviewUrlRef.current) { + URL.revokeObjectURL(ocrFilePreviewUrlRef.current); + ocrFilePreviewUrlRef.current = null; + } + }; + }, []); // Fetch MCP servers and toolsets const loadMCPServers = async () => { @@ -646,6 +724,14 @@ const ChatUI: React.FC = ({ event.target.value = ""; }; + const handleAudioFileDrop = (event: React.DragEvent) => { + event.preventDefault(); + const file = event.dataTransfer.files[0]; + if (file) { + handleAudioUpload(file); + } + }; + const mcpServerOptions = useMemo((): MultiSelectOption[] => { const options: MultiSelectOption[] = []; if (endpointType !== EndpointType.MCP) { @@ -708,8 +794,43 @@ const ChatUI: React.FC = ({ setUploadedAudio(null); }; + const handleOcrFileUpload = (file: File): false => { + revokeOcrFilePreviewUrl(); + setUploadedOcrFile(file); + const previewUrl = file.type.startsWith("image/") ? URL.createObjectURL(file) : null; + ocrFilePreviewUrlRef.current = previewUrl; + setOcrFilePreviewUrl(previewUrl); + return false; + }; + + const handleOcrFileInputChange = (event: React.ChangeEvent) => { + const file = event.target.files?.[0]; + if (file) { + handleOcrFileUpload(file); + } + event.target.value = ""; + }; + + const handleOcrFileDrop = (event: React.DragEvent) => { + event.preventDefault(); + const file = event.dataTransfer.files[0]; + if (file) { + handleOcrFileUpload(file); + } + }; + + const handleRemoveOcrFile = () => { + setUploadedOcrFile(null); + revokeOcrFilePreviewUrl(); + }; + const handleSendMessage = async () => { - if (inputMessage.trim() === "" && endpointType !== EndpointType.TRANSCRIPTION && endpointType !== EndpointType.MCP) + if ( + inputMessage.trim() === "" && + endpointType !== EndpointType.TRANSCRIPTION && + endpointType !== EndpointType.MCP && + endpointType !== EndpointType.OCR + ) return; // For image edits, require both image and prompt @@ -724,6 +845,11 @@ const ChatUI: React.FC = ({ return; } + if (endpointType === EndpointType.OCR && !uploadedOcrFile) { + toast.fromError("Please upload a document or image for OCR"); + return; + } + // For A2A agents, require agent selection if (endpointType === EndpointType.A2A_AGENTS && !selectedAgent) { toast.fromError("Please select an agent to send a message"); @@ -779,6 +905,7 @@ const ChatUI: React.FC = ({ EndpointType.ANTHROPIC_MESSAGES, EndpointType.EMBEDDINGS, EndpointType.TRANSCRIPTION, + EndpointType.OCR, EndpointType.INTERACTIONS, ]; @@ -791,7 +918,11 @@ const ChatUI: React.FC = ({ return; } - const effectiveApiKey = simplified ? accessToken : apiKeySource === "session" ? accessToken : apiKey; + const effectiveApiKey = (() => { + if (simplified) return accessToken; + if (apiKeySource === "session") return accessToken; + return apiKey; + })(); if (!effectiveApiKey) { toast.fromError("Please provide a Virtual Key or select Current UI Session"); @@ -854,6 +985,8 @@ const ChatUI: React.FC = ({ ? `🎵 Audio file: ${uploadedAudio.name}\nPrompt: ${inputMessage}` : `🎵 Audio file: ${uploadedAudio.name}`; displayMessage = createDisplayMessage(audioMessage, false); + } else if (endpointType === EndpointType.OCR && uploadedOcrFile) { + displayMessage = createDisplayMessage(`OCR file: ${uploadedOcrFile.name}`, false); } else if (endpointType === EndpointType.MCP && selectedMCPDirectTool) { // For MCP direct mode, show tool name and arguments from form const mcpMessage = `🔧 MCP Tool: ${selectedMCPDirectTool}\nArguments: ${JSON.stringify(mcpToolArguments, null, 2)}`; @@ -1051,6 +1184,18 @@ const ChatUI: React.FC = ({ customProxyBaseUrl || undefined, ); } + } else if (endpointType === EndpointType.OCR) { + if (uploadedOcrFile) { + await makeOpenAIOcrRequest({ + file: uploadedOcrFile, + updateUI: (text, model) => updateTextUI("assistant", text, model), + selectedModel, + accessToken: effectiveApiKey, + tags: selectedTags, + signal, + customBaseUrl: customProxyBaseUrl || undefined, + }); + } } else if (endpointType === EndpointType.INTERACTIONS) { await makeInteractionsRequest( inputMessage, @@ -1136,6 +1281,9 @@ const ChatUI: React.FC = ({ if (endpointType === EndpointType.TRANSCRIPTION && uploadedAudio) { handleRemoveAudio(); } + if (endpointType === EndpointType.OCR && uploadedOcrFile) { + handleRemoveOcrFile(); + } } setInputMessage(""); @@ -1147,6 +1295,7 @@ const ChatUI: React.FC = ({ handleRemoveResponsesImage(); handleRemoveChatImage(); handleRemoveAudio(); + handleRemoveOcrFile(); toast.success("Chat history cleared."); }; @@ -1187,30 +1336,16 @@ const ChatUI: React.FC = ({ modelEmptyText = "No models available for this endpoint"; } - const inputPlaceholder = - endpointType === EndpointType.CHAT || - endpointType === EndpointType.EMBEDDINGS || - endpointType === EndpointType.RESPONSES || - endpointType === EndpointType.ANTHROPIC_MESSAGES || - endpointType === EndpointType.INTERACTIONS - ? "Type your message... (Shift+Enter for new line)" - : endpointType === EndpointType.A2A_AGENTS - ? "Send a message to the A2A agent..." - : endpointType === EndpointType.IMAGE_EDITS - ? "Describe how you want to edit the image..." - : endpointType === EndpointType.SPEECH - ? "Enter text to convert to speech..." - : endpointType === EndpointType.TRANSCRIPTION - ? "Optional: Add context or prompt for transcription..." - : "Describe the image you want to generate..."; - - const sendDisabled = - isLoading || - (endpointType === EndpointType.MCP - ? !(selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" && selectedMCPDirectTool) - : endpointType === EndpointType.TRANSCRIPTION - ? !uploadedAudio - : !inputMessage.trim()); + const inputPlaceholder = getInputPlaceholder(endpointType as EndpointType); + const sendDisabled = getSendDisabled({ + endpointType: endpointType as EndpointType, + inputMessage, + isLoading, + selectedMCPDirectTool, + selectedMCPServers, + uploadedAudio, + uploadedOcrFile, + }); return (
@@ -1887,13 +2022,7 @@ const ChatUI: React.FC = ({
)} + {endpointType === EndpointType.OCR && !uploadedOcrFile && ( +
+ +
+ )} {endpointType === EndpointType.RESPONSES && responsesUploadedImage && ( = ({ /> )} + {endpointType === EndpointType.OCR && uploadedOcrFile && ( + + )} {endpointType === EndpointType.RESPONSES && codeInterpreter.enabled && (
@@ -1998,10 +2153,15 @@ const ChatUI: React.FC = ({ onSubmit={handleSendMessage} onCancel={handleCancelRequest} placeholder={inputPlaceholder} - disabled={isLoading} + disabled={isLoading || endpointType === EndpointType.OCR} isLoading={isLoading} submitDisabled={sendDisabled} - showSuggestions={chatHistory.length === 0 && !isLoading && endpointType !== EndpointType.MCP} + showSuggestions={ + chatHistory.length === 0 && + !isLoading && + endpointType !== EndpointType.MCP && + endpointType !== EndpointType.OCR + } suggestions={ endpointType === EndpointType.A2A_AGENTS ? ["What can you help me with?", "Tell me about yourself", "What tasks can you perform?"] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx index a1e88e51053..d09be961727 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx @@ -26,4 +26,18 @@ describe("EndpointSelector", () => { expect(await screen.findByText("/v1/audio/speech")).toBeInTheDocument(); expect(await screen.findByText("/v1/audio/transcriptions")).toBeInTheDocument(); }); + + it("should filter and show the OCR endpoint when user inputs 'ocr'", async () => { + const user = userEvent.setup(); + render( {}} />); + + const combobox = screen.getByRole("combobox"); + await user.click(combobox); + + const input = await screen.findByRole("combobox"); + await user.clear(input); + await user.type(input, "ocr"); + + expect(await screen.findByText("/v1/ocr")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/chatConstants.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/chatConstants.ts index 2b59fbad2ee..1779b92d36a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/chatConstants.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/chatConstants.ts @@ -42,6 +42,7 @@ export const ENDPOINT_OPTIONS = [ { value: EndpointType.EMBEDDINGS, label: "/v1/embeddings" }, { value: EndpointType.SPEECH, label: "/v1/audio/speech" }, { value: EndpointType.TRANSCRIPTION, label: "/v1/audio/transcriptions" }, + { value: EndpointType.OCR, label: "/v1/ocr" }, { value: EndpointType.A2A_AGENTS, label: "/v1/a2a/message/send" }, { value: EndpointType.MCP, label: "/mcp-rest/tools/call" }, { value: EndpointType.REALTIME, label: "/v1/realtime" }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/OcrApi.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/OcrApi.test.tsx new file mode 100644 index 00000000000..d19016fae72 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/OcrApi.test.tsx @@ -0,0 +1,62 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { makeOpenAIOcrRequest } from "./OcrApi"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => "https://example.com"), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), +})); + +describe("ocr_api", () => { + const mockUpdateUI = vi.fn(); + const mockFetch = vi.fn(); + + beforeEach(() => { + const responseBody = JSON.stringify({ + pages: [ + { + index: 0, + markdown: "# Extracted text", + }, + ], + model: "mistral-ocr-latest", + }); + mockFetch.mockResolvedValue({ + ok: true, + text: async () => responseBody, + } as Response); + + global.fetch = mockFetch; + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("posts OCR uploads to /v1/ocr with model and file multipart fields", async () => { + const file = new File(["document data"], "invoice.pdf", { type: "application/pdf" }); + + await makeOpenAIOcrRequest({ + file, + updateUI: mockUpdateUI, + selectedModel: "mistral-ocr-latest", + accessToken: "sk-1234567890", + tags: ["tag1", "tag2"], + }); + + expect(mockFetch).toHaveBeenCalledTimes(1); + expect(mockFetch).toHaveBeenCalledWith("https://example.com/v1/ocr", { + method: "POST", + headers: { + Authorization: "Bearer sk-1234567890", + "x-litellm-tags": "tag1,tag2", + }, + body: expect.any(FormData), + signal: undefined, + }); + + const formData = mockFetch.mock.calls[0][1].body as FormData; + expect(formData.get("model")).toBe("mistral-ocr-latest"); + expect(formData.get("file")).toBe(file); + expect(mockUpdateUI).toHaveBeenCalledWith("# Extracted text", "mistral-ocr-latest"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/OcrApi.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/OcrApi.tsx new file mode 100644 index 00000000000..bc13af33ead --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/OcrApi.tsx @@ -0,0 +1,66 @@ +import { toast } from "@/lib/toast"; +import { getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking"; +import { createApiClient } from "@/lib/http/client"; + +interface OcrRequestParams { + file: File; + updateUI: (text: string, model: string) => void; + selectedModel: string; + accessToken: string; + tags?: string[]; + signal?: AbortSignal; + customBaseUrl?: string; +} + +const formatOcrResponse = (response: unknown): string => { + if (typeof response !== "object" || response === null || !("pages" in response)) { + return JSON.stringify(response, null, 2); + } + + const pages = (response as { pages?: unknown }).pages; + if (!Array.isArray(pages)) { + return JSON.stringify(response, null, 2); + } + + const markdown = pages + .map((page) => + typeof page === "object" && page !== null && "markdown" in page + ? (page as { markdown?: unknown }).markdown + : undefined, + ) + .filter((pageMarkdown): pageMarkdown is string => typeof pageMarkdown === "string" && pageMarkdown.length > 0) + .join("\n\n"); + + return markdown || JSON.stringify(response, null, 2); +}; + +export async function makeOpenAIOcrRequest({ + file, + updateUI, + selectedModel, + accessToken, + tags, + signal, + customBaseUrl, +}: OcrRequestParams) { + const client = createApiClient({ + getBaseUrl: () => customBaseUrl || getProxyBaseUrl(), + getAuthHeaderName: getGlobalLitellmHeaderName, + onError: (message) => toast.fromError(`OCR failed: ${message}`), + }); + const formData = new FormData(); + formData.append("model", selectedModel); + formData.append("file", file); + + const responseJson = await client.post("/v1/ocr", { + accessToken, + rawBody: formData, + headers: { + ...(tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : {}), + }, + signal, + }); + + updateUI(formatOcrResponse(responseJson), selectedModel); + toast.success("OCR completed successfully"); +} diff --git a/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.test.ts b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.test.ts new file mode 100644 index 00000000000..d4a51d392b6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.test.ts @@ -0,0 +1,16 @@ +import { describe, expect, it } from "vitest"; +import { EndpointType, getEndpointType, litellmModeMapping, ModelMode } from "./mode_endpoint_mapping"; + +describe("mode_endpoint_mapping", () => { + it("maps OCR model mode to the OCR endpoint", () => { + expect(getEndpointType("ocr")).toBe(EndpointType.OCR); + expect(litellmModeMapping[ModelMode.OCR]).toBe(EndpointType.OCR); + }); + + it("preserves existing model mode mappings", () => { + expect(getEndpointType("chat")).toBe(EndpointType.CHAT); + expect(getEndpointType("responses")).toBe(EndpointType.RESPONSES); + expect(getEndpointType("image_generation")).toBe(EndpointType.IMAGE); + expect(getEndpointType("audio_transcription")).toBe(EndpointType.TRANSCRIPTION); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx index 930ded5d1a5..1125405e714 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx @@ -13,6 +13,7 @@ export enum ModelMode { ANTHROPIC_MESSAGES = "anthropic_messages", EMBEDDING = "embedding", REALTIME = "realtime", + OCR = "ocr", } // Define an enum for the endpoint types your UI calls @@ -26,6 +27,7 @@ export enum EndpointType { EMBEDDINGS = "embeddings", SPEECH = "speech", TRANSCRIPTION = "transcription", + OCR = "ocr", A2A_AGENTS = "a2a_agents", MCP = "mcp", REALTIME = "realtime", @@ -45,6 +47,7 @@ export const litellmModeMapping: Record = { [ModelMode.AUDIO_TRANSCRIPTION]: EndpointType.TRANSCRIPTION, [ModelMode.EMBEDDING]: EndpointType.EMBEDDINGS, [ModelMode.REALTIME]: EndpointType.REALTIME, + [ModelMode.OCR]: EndpointType.OCR, }; export const getEndpointType = (mode: string): EndpointType => {