Test Key UI Embeddings (#16065)

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(
+      <ChatUI
+        accessToken="1234567890"
+        token="1234567890"
+        userRole="user"
+        userID="1234567890"
+        disabledPersonalKeyCreation={false}
+      />,
+    );
+    expect(getByText("Test Key")).toBeInTheDocument();
+  });
+});

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<ChatUIProps> = ({ 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<ChatUIProps> = ({ 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<ChatUIProps> = ({ 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)"

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)");
+  });
+});

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.";
   }

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(<EndpointSelector endpointType={endpointType.value} onEndpointChange={() => {}} />);
+      await waitFor(() => {
+        expect(getByText(endpointType.label)).toBeInTheDocument();
+      });
+    });
+  });
+});

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<EndpointSelectorProps> = ({ 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<EndpointSelectorProps> = ({ endpointType, onEndpointChange, className }) => {
   return (
     <div className={className}>
       <Text>Endpoint Type:</Text>

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!");
+    });
+  });
+});

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;
+}

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",
+    );
+  });
+});

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<string, string> = {};
+  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
+  }
+}

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
 }