diff --git a/ui/litellm-dashboard/src/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/ChatUI.test.tsx new file mode 100644 index 0000000000..22d8d1fd2a --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/ChatUI.test.tsx @@ -0,0 +1,18 @@ +import { render } from "@testing-library/react"; +import { describe, it, expect } from "vitest"; +import ChatUI from "./ChatUI"; + +describe("ChatUI", () => { + it("should render the chat UI", () => { + const { getByText } = render( + , + ); + expect(getByText("Test Key")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/chat_ui/ChatUI.tsx index 680f66702b..164295261a 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/ChatUI.tsx @@ -1,12 +1,6 @@ import React, { useState, useEffect, useRef } from "react"; import ReactMarkdown from "react-markdown"; -import { - Card, - Title, - Text, - TextInput, - Button as TremorButton, -} from "@tremor/react"; +import { Card, Title, Text, TextInput, Button as TremorButton } from "@tremor/react"; import { v4 as uuidv4 } from "uuid"; import { Select, Spin, Typography, Tooltip, Input, Upload, Modal, Button } from "antd"; @@ -57,6 +51,8 @@ import { ArrowUpOutlined, } from "@ant-design/icons"; import NotificationsManager from "../molecules/notifications_manager"; +import { makeOpenAIEmbeddingsRequest } from "./llm_calls/embeddings_api"; +import { truncateString } from "./chatUtils"; const { TextArea } = Input; const { Dragger } = Upload; @@ -497,6 +493,13 @@ const ChatUI: React.FC = ({ accessToken, token, userRole, userID, d setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]); }; + const updateEmbeddingsUI = (embeddings: string, model?: string) => { + setChatHistory((prevHistory) => [ + ...prevHistory, + { role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true }, + ]); + }; + const updateChatImageUI = (imageUrl: string, model?: string) => { setChatHistory((prev) => { const last = prev[prev.length - 1]; @@ -786,6 +789,14 @@ const ChatUI: React.FC = ({ accessToken, token, userRole, userID, d selectedGuardrails.length > 0 ? selectedGuardrails : undefined, selectedMCPTools, // Pass the selected tools array ); + } else if (endpointType === EndpointType.EMBEDDINGS) { + await makeOpenAIEmbeddingsRequest( + inputMessage, + (embeddings, model) => updateEmbeddingsUI(embeddings, model), + selectedModel, + effectiveApiKey, + selectedTags, + ); } } } catch (error) { @@ -1413,6 +1424,7 @@ const ChatUI: React.FC = ({ accessToken, token, userRole, userID, d onKeyDown={handleKeyDown} placeholder={ endpointType === EndpointType.CHAT || + endpointType === EndpointType.EMBEDDINGS || endpointType === EndpointType.RESPONSES || endpointType === EndpointType.ANTHROPIC_MESSAGES ? "Type your message... (Shift+Enter for new line)" diff --git a/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx new file mode 100644 index 0000000000..61f882025b --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx @@ -0,0 +1,26 @@ +import { describe, it, expect } from "vitest"; +import { generateCodeSnippet } from "./CodeSnippets"; +import { EndpointType } from "./mode_endpoint_mapping"; + +describe("CodeSnippets", () => { + it("should generate the correct code snippet for embeddings", () => { + const code = generateCodeSnippet({ + endpointType: EndpointType.EMBEDDINGS, + inputMessage: "Hello, world!", + selectedModel: "text-embedding-3-small", + apiKeySource: "session", + accessToken: "1234567890", + apiKey: "1234567890", + chatHistory: [], + selectedTags: [], + selectedVectorStores: [], + selectedGuardrails: [], + selectedMCPTools: [], + selectedSdk: "openai", + }); + expect(code).toContain("text-embedding-3-small"); + expect(code).toContain("Hello, world!"); + expect(code).toContain("client.embeddings.create"); + expect(code).toContain("print(response.data[0].embedding)"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx index d85f515bcc..44a0b198db 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx @@ -493,6 +493,17 @@ else: `; } break; + case EndpointType.EMBEDDINGS: + endpointSpecificCode = ` +response = client.embeddings.create( + input="${inputMessage || "Your string here"}", + model="${modelNameForCode}", + encoding_format="base64" # or "float" +) + +print(response.data[0].embedding) +`; + break; default: endpointSpecificCode = "\n# Code generation for this endpoint is not implemented yet."; } diff --git a/ui/litellm-dashboard/src/components/chat_ui/EndpointSelector.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/EndpointSelector.test.tsx new file mode 100644 index 0000000000..b32bcf8ea7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/EndpointSelector.test.tsx @@ -0,0 +1,14 @@ +import { render, waitFor } from "@testing-library/react"; +import { describe, it, expect } from "vitest"; +import EndpointSelector, { endpointOptions } from "./EndpointSelector"; + +describe("EndpointSelector", () => { + Object.values(endpointOptions).forEach((endpointType) => { + it(`should render the endpoint selector for ${endpointType.value}`, async () => { + const { getByText } = render( {}} />); + await waitFor(() => { + expect(getByText(endpointType.label)).toBeInTheDocument(); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat_ui/EndpointSelector.tsx b/ui/litellm-dashboard/src/components/chat_ui/EndpointSelector.tsx index 00dae23d6f..141016c81a 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/EndpointSelector.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/EndpointSelector.tsx @@ -12,16 +12,16 @@ interface EndpointSelectorProps { /** * A reusable component for selecting API endpoints */ -const EndpointSelector: React.FC = ({ endpointType, onEndpointChange, className }) => { - // Map endpoint types to their display labels - const endpointOptions = [ - { value: EndpointType.CHAT, label: "/v1/chat/completions" }, - { value: EndpointType.RESPONSES, label: "/v1/responses" }, - { value: EndpointType.ANTHROPIC_MESSAGES, label: "/v1/messages" }, - { value: EndpointType.IMAGE, label: "/v1/images/generations" }, - { value: EndpointType.IMAGE_EDITS, label: "/v1/images/edits" }, - ]; +export const endpointOptions = [ + { value: EndpointType.CHAT, label: "/v1/chat/completions" }, + { value: EndpointType.RESPONSES, label: "/v1/responses" }, + { value: EndpointType.ANTHROPIC_MESSAGES, label: "/v1/messages" }, + { value: EndpointType.IMAGE, label: "/v1/images/generations" }, + { value: EndpointType.IMAGE_EDITS, label: "/v1/images/edits" }, + { value: EndpointType.EMBEDDINGS, label: "/v1/embeddings" }, +]; +const EndpointSelector: React.FC = ({ endpointType, onEndpointChange, className }) => { return (
Endpoint Type: diff --git a/ui/litellm-dashboard/src/components/chat_ui/chatUtils.test.ts b/ui/litellm-dashboard/src/components/chat_ui/chatUtils.test.ts new file mode 100644 index 0000000000..547bc72e81 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/chatUtils.test.ts @@ -0,0 +1,14 @@ +import { describe, expect, it } from "vitest"; +import { truncateString } from "./chatUtils"; + +describe("chatUtils", () => { + describe("truncateString", () => { + it("should truncate a string", () => { + expect(truncateString("Hello, world!", 5)).toBe("Hello..."); + }); + + it("should return the original string if it is less than the max length", () => { + expect(truncateString("Hello, world!", 20)).toBe("Hello, world!"); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat_ui/chatUtils.ts b/ui/litellm-dashboard/src/components/chat_ui/chatUtils.ts new file mode 100644 index 0000000000..bd590a5ae1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/chatUtils.ts @@ -0,0 +1,3 @@ +export function truncateString(str: string, maxLength: number) { + return str.length > maxLength ? str.substring(0, maxLength) + "..." : str; +} diff --git a/ui/litellm-dashboard/src/components/chat_ui/llm_calls/embeddings_api.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/embeddings_api.test.tsx new file mode 100644 index 0000000000..c8d4d49cbe --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/embeddings_api.test.tsx @@ -0,0 +1,59 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { makeOpenAIEmbeddingsRequest } from "./embeddings_api"; +import OpenAI from "openai"; + +vi.mock("openai"); + +describe("embeddings_api", () => { + const mockCreate = vi.fn(); + const mockUpdateEmbeddingsUI = vi.fn(); + + beforeEach(() => { + // Mock the response structure from OpenAI embeddings API + mockCreate.mockResolvedValue({ + data: [ + { + embedding: [0.1, 0.2, 0.3, 0.4, 0.5], + index: 0, + object: "embedding", + }, + ], + model: "text-embedding-3-small", + object: "list", + usage: { + prompt_tokens: 5, + total_tokens: 5, + }, + }); + + // Mock the OpenAI constructor and its methods + (OpenAI as any).mockImplementation(() => ({ + embeddings: { + create: mockCreate, + }, + })); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("should make a request to the embeddings API", async () => { + await makeOpenAIEmbeddingsRequest( + "Hello, world!", + mockUpdateEmbeddingsUI, + "text-embedding-3-small", + "1234567890", + [], + ); + + expect(mockCreate).toHaveBeenCalledWith({ + model: "text-embedding-3-small", + input: "Hello, world!", + }); + expect(mockUpdateEmbeddingsUI).toHaveBeenCalledWith( + JSON.stringify([0.1, 0.2, 0.3, 0.4, 0.5]), + "text-embedding-3-small", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat_ui/llm_calls/embeddings_api.tsx b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/embeddings_api.tsx new file mode 100644 index 0000000000..7a5ab21d9f --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/embeddings_api.tsx @@ -0,0 +1,50 @@ +import NotificationManager from "@/components/molecules/notifications_manager"; +import { getProxyBaseUrl } from "@/components/networking"; +import openai from "openai"; + +export async function makeOpenAIEmbeddingsRequest( + input: string, + updateEmbeddingsUI: (embeddings: string, model?: string) => void, + selectedModel: string, + accessToken: string, + tags?: string[], +) { + if (!accessToken) { + throw new Error("API key is required"); + } + + // Base URL should be the current base_url + const isLocal = process.env.NODE_ENV === "development"; + if (isLocal !== true) { + console.log = function () {}; + } + + const proxyBaseUrl = getProxyBaseUrl(); + // Prepare headers with tags and trace ID + const headers: Record = {}; + if (tags && tags.length > 0) { + headers["x-litellm-tags"] = tags.join(","); + } + + const client = new openai.OpenAI({ + apiKey: accessToken, + baseURL: proxyBaseUrl, + dangerouslyAllowBrowser: true, + defaultHeaders: headers, + }); + + try { + const response = await client.embeddings.create({ + model: selectedModel, + input: input, + }); + + updateEmbeddingsUI(JSON.stringify(response.data[0].embedding), selectedModel); + } catch (error: unknown) { + NotificationManager.fromBackend( + `Error occurred while making embeddings request. Please try again. Error: ${error}`, + ); + + throw error; // Re-throw to allow the caller to handle the error + } +} 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 65b4c0b867..ef798c367c 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 @@ -17,6 +17,7 @@ export enum EndpointType { RESPONSES = "responses", IMAGE_EDITS = "image_edits", ANTHROPIC_MESSAGES = "anthropic_messages", + EMBEDDINGS = "embeddings", // add additional endpoint types if required }