Update auth for litellm proxy (#2316)

* update for auth * validated embedding model names * remove embedding provider * remove logs * add ability to delete search setting * add abiility to delete models + more streamlined API endpoints * remove upsert * minor typing fix * add connector utils
2025-06-20 13:01:34 +02:00 · 2024-09-04 13:59:07 -07:00 · 2024-09-04 13:59:07 -07:00 · 34ba3181ff
commit 34ba3181ff
parent 630e2248bd
13 changed files with 282 additions and 91 deletions
--- a/backend/danswer/db/search_settings.py
+++ b/backend/danswer/db/search_settings.py
@ -1,3 +1,5 @@
+from sqlalchemy import and_
+from sqlalchemy import delete
 from sqlalchemy import select
 from sqlalchemy.orm import Session

@ -13,6 +15,7 @@ from danswer.configs.model_configs import OLD_DEFAULT_MODEL_NORMALIZE_EMBEDDINGS
 from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.llm import fetch_embedding_provider
 from danswer.db.models import CloudEmbeddingProvider
+from danswer.db.models import IndexAttempt
 from danswer.db.models import IndexModelStatus
 from danswer.db.models import SearchSettings
 from danswer.indexing.models import IndexingSetting
@ -89,6 +92,30 @@ def get_current_db_embedding_provider(
    return current_embedding_provider


+def delete_search_settings(db_session: Session, search_settings_id: int) -> None:
+    current_settings = get_current_search_settings(db_session)
+
+    if current_settings.id == search_settings_id:
+        raise ValueError("Cannot delete currently active search settings")
+
+    # First, delete associated index attempts
+    index_attempts_query = delete(IndexAttempt).where(
+        IndexAttempt.search_settings_id == search_settings_id
+    )
+    db_session.execute(index_attempts_query)
+
+    # Then, delete the search settings
+    search_settings_query = delete(SearchSettings).where(
+        and_(
+            SearchSettings.id == search_settings_id,
+            SearchSettings.status != IndexModelStatus.PRESENT,
+        )
+    )
+
+    db_session.execute(search_settings_query)
+    db_session.commit()
+
+
 def get_current_search_settings(db_session: Session) -> SearchSettings:
    query = (
        select(SearchSettings)
--- a/backend/danswer/indexing/models.py
+++ b/backend/danswer/indexing/models.py
@ -95,6 +95,7 @@ class DocMetadataAwareIndexChunk(IndexChunk):


 class EmbeddingModelDetail(BaseModel):
+    id: int | None = None
    model_name: str
    normalize: bool
    query_prefix: str | None
@ -112,6 +113,7 @@ class EmbeddingModelDetail(BaseModel):
        search_settings: "SearchSettings",
    ) -> "EmbeddingModelDetail":
        return cls(
+            id=search_settings.id,
            model_name=search_settings.model_name,
            normalize=search_settings.normalize,
            query_prefix=search_settings.query_prefix,
--- a/backend/danswer/server/manage/embedding/api.py
+++ b/backend/danswer/server/manage/embedding/api.py
@ -42,10 +42,10 @@ def test_embedding_configuration(
            api_key=test_llm_request.api_key,
            api_url=test_llm_request.api_url,
            provider_type=test_llm_request.provider_type,
+            model_name=test_llm_request.model_name,
            normalize=False,
            query_prefix=None,
            passage_prefix=None,
-            model_name=None,
        )
        test_model.encode(["Testing Embedding"], text_type=EmbedTextType.QUERY)

--- a/backend/danswer/server/manage/embedding/models.py
+++ b/backend/danswer/server/manage/embedding/models.py
@ -8,10 +8,15 @@ if TYPE_CHECKING:
    from danswer.db.models import CloudEmbeddingProvider as CloudEmbeddingProviderModel


+class SearchSettingsDeleteRequest(BaseModel):
+    search_settings_id: int
+
+
 class TestEmbeddingRequest(BaseModel):
    provider_type: EmbeddingProvider
    api_key: str | None = None
    api_url: str | None = None
+    model_name: str | None = None


 class CloudEmbeddingProvider(BaseModel):
--- a/backend/danswer/server/manage/search_settings.py
+++ b/backend/danswer/server/manage/search_settings.py
@ -14,6 +14,7 @@ from danswer.db.index_attempt import expire_index_attempts
 from danswer.db.models import IndexModelStatus
 from danswer.db.models import User
 from danswer.db.search_settings import create_search_settings
+from danswer.db.search_settings import delete_search_settings
 from danswer.db.search_settings import get_current_search_settings
 from danswer.db.search_settings import get_embedding_provider_from_provider_type
 from danswer.db.search_settings import get_secondary_search_settings
@ -23,6 +24,7 @@ from danswer.document_index.factory import get_default_document_index
 from danswer.natural_language_processing.search_nlp_models import clean_model_name
 from danswer.search.models import SavedSearchSettings
 from danswer.search.models import SearchSettingsCreationRequest
+from danswer.server.manage.embedding.models import SearchSettingsDeleteRequest
 from danswer.server.manage.models import FullModelVersionResponse
 from danswer.server.models import IdReturn
 from danswer.utils.logger import setup_logger
@ -97,6 +99,7 @@ def set_new_search_settings(
        primary_index_name=search_settings.index_name,
        secondary_index_name=new_search_settings.index_name,
    )
+
    document_index.ensure_indices_exist(
        index_embedding_dim=search_settings.model_dim,
        secondary_index_embedding_dim=new_search_settings.model_dim,
@ -132,6 +135,21 @@ def cancel_new_embedding(
        )


+@router.delete("/delete-search-settings")
+def delete_search_settings_endpoint(
+    deletion_request: SearchSettingsDeleteRequest,
+    _: User | None = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+) -> None:
+    try:
+        delete_search_settings(
+            db_session=db_session,
+            search_settings_id=deletion_request.search_settings_id,
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
+
+
@router.get("/get-current-search-settings")
 def get_current_search_settings_endpoint(
    _: User | None = Depends(current_user),
--- a/backend/model_server/encoders.py
+++ b/backend/model_server/encoders.py
@ -237,15 +237,18 @@ def get_local_reranking_model(


 def embed_with_litellm_proxy(
-    texts: list[str], api_url: str, model: str
+    texts: list[str], api_url: str, model_name: str, api_key: str | None
 ) -> list[Embedding]:
+    headers = {} if not api_key else {"Authorization": f"Bearer {api_key}"}
+
    with httpx.Client() as client:
        response = client.post(
            api_url,
            json={
-                "model": model,
+                "model": model_name,
                "input": texts,
            },
+            headers=headers,
        )
        response.raise_for_status()
        result = response.json()
@ -280,7 +283,12 @@ def embed_text(
            logger.error("API URL not provided for LiteLLM proxy")
            raise ValueError("API URL is required for LiteLLM proxy embedding.")
        try:
-            return embed_with_litellm_proxy(texts, api_url, model_name or "")
+            return embed_with_litellm_proxy(
+                texts=texts,
+                api_url=api_url,
+                model_name=model_name or "",
+                api_key=api_key,
+            )
        except Exception as e:
            logger.exception(f"Error during LiteLLM proxy embedding: {str(e)}")
            raise
--- a/backend/shared_configs/configs.py
+++ b/backend/shared_configs/configs.py
@ -58,6 +58,7 @@ LOG_LEVEL = os.environ.get("LOG_LEVEL", "notice")

 # Fields which should only be set on new search setting
 PRESERVED_SEARCH_FIELDS = [
+    "id",
    "provider_type",
    "api_key",
    "model_name",
--- a/web/src/app/admin/connectors/[connector]/AddConnectorPage.tsx
+++ b/web/src/app/admin/connectors/[connector]/AddConnectorPage.tsx
@ -91,7 +91,7 @@ export default function AddConnector({
  >({
    name: "",
    groups: [],
-    is_public: false,
+    is_public: true,
    ...configuration.values.reduce(
      (acc, field) => {
        if (field.type === "list") {
--- a/web/src/app/admin/embeddings/modals/ChangeCredentialsModal.tsx
+++ b/web/src/app/admin/embeddings/modals/ChangeCredentialsModal.tsx
@ -153,34 +153,22 @@ export function ChangeCredentialsModal({
      title={`Modify your ${provider.provider_type} ${isProxy ? "URL" : "key"}`}
      onOutsideClick={onCancel}
    >
-      <div className="mb-4">
-        <Subtitle className="font-bold text-lg">
-          Want to swap out your {isProxy ? "URL" : "key"}?
-        </Subtitle>
-        <a
-          href={provider.apiLink}
-          target="_blank"
-          rel="noopener noreferrer"
-          className="underline cursor-pointer mt-2 mb-4"
-        >
-          Visit API
-        </a>
+      <>
+        {isProxy && (
+          <div className="mb-4">
+            <Subtitle className="font-bold text-lg">
+              Want to swap out your URL?
+            </Subtitle>
+            <a
+              href={provider.apiLink}
+              target="_blank"
+              rel="noopener noreferrer"
+              className="underline cursor-pointer mt-2 mb-4"
+            >
+              Visit API
+            </a>

-        <div className="flex flex-col mt-4 gap-y-2">
-          {useFileUpload ? (
-            <>
-              <Label>Upload JSON File</Label>
-              <input
-                ref={fileInputRef}
-                type="file"
-                accept=".json"
-                onChange={handleFileUpload}
-                className="text-lg w-full p-1"
-              />
-              {fileName && <p>Uploaded file: {fileName}</p>}
-            </>
-          ) : (
-            <>
+            <div className="flex flex-col mt-4 gap-y-2">
              <input
                className={`
                    border 
@ -193,46 +181,116 @@ export function ChangeCredentialsModal({
                `}
                value={apiKeyOrUrl}
                onChange={(e: any) => setApiKeyOrUrl(e.target.value)}
-                placeholder={`Paste your ${isProxy ? "API URL" : "API key"} here`}
+                placeholder="Paste your API URL here"
              />
-            </>
+            </div>
+
+            {testError && (
+              <Callout title="Error" color="red" className="mt-4">
+                {testError}
+              </Callout>
+            )}
+
+            <div className="flex mt-4 justify-between">
+              <Button
+                color="blue"
+                onClick={() => handleSubmit()}
+                disabled={!apiKeyOrUrl}
+              >
+                Swap URL
+              </Button>
+            </div>
+
+            {deletionError && (
+              <Callout title="Error" color="red" className="mt-4">
+                {deletionError}
+              </Callout>
+            )}
+            <Divider />
+          </div>
+        )}
+
+        <div className="mb-4">
+          <Subtitle className="font-bold text-lg">
+            Want to swap out your key?
+          </Subtitle>
+          <a
+            href={provider.apiLink}
+            target="_blank"
+            rel="noopener noreferrer"
+            className="underline cursor-pointer mt-2 mb-4"
+          >
+            Visit API
+          </a>
+
+          <div className="flex flex-col mt-4 gap-y-2">
+            {useFileUpload ? (
+              <>
+                <Label>Upload JSON File</Label>
+                <input
+                  ref={fileInputRef}
+                  type="file"
+                  accept=".json"
+                  onChange={handleFileUpload}
+                  className="text-lg w-full p-1"
+                />
+                {fileName && <p>Uploaded file: {fileName}</p>}
+              </>
+            ) : (
+              <>
+                <input
+                  className={`
+                    border 
+                    border-border 
+                    rounded 
+                    w-full 
+                    py-2 
+                    px-3 
+                    bg-background-emphasis
+                `}
+                  value={apiKeyOrUrl}
+                  onChange={(e: any) => setApiKeyOrUrl(e.target.value)}
+                  placeholder="Paste your API key here"
+                />
+              </>
+            )}
+          </div>
+
+          {testError && (
+            <Callout title="Error" color="red" className="mt-4">
+              {testError}
+            </Callout>
+          )}
+
+          <div className="flex mt-4 justify-between">
+            <Button
+              color="blue"
+              onClick={() => handleSubmit()}
+              disabled={!apiKeyOrUrl}
+            >
+              Swap Key
+            </Button>
+          </div>
+          <Divider />
+
+          <Subtitle className="mt-4 font-bold text-lg mb-2">
+            You can also delete your configuration.
+          </Subtitle>
+          <Text className="mb-2">
+            This is only possible if you have already switched to a different
+            embedding type!
+          </Text>
+
+          <Button onClick={handleDelete} color="red">
+            Delete Configuration
+          </Button>
+          {deletionError && (
+            <Callout title="Error" color="red" className="mt-4">
+              {deletionError}
+            </Callout>
          )}
        </div>
-
-        {testError && (
-          <Callout title="Error" color="red" className="mt-4">
-            {testError}
-          </Callout>
-        )}
-
-        <div className="flex mt-4 justify-between">
-          <Button
-            color="blue"
-            onClick={() => handleSubmit()}
-            disabled={!apiKeyOrUrl}
-          >
-            Swap {isProxy ? "URL" : "Key"}
-          </Button>
-        </div>
-        <Divider />
-
-        <Subtitle className="mt-4 font-bold text-lg mb-2">
-          You can also delete your {isProxy ? "URL" : "key"}.
-        </Subtitle>
-        <Text className="mb-2">
-          This is only possible if you have already switched to a different
-          embedding type!
-        </Text>
-
-        <Button onClick={handleDelete} color="red">
-          Delete {isProxy ? "URL" : "key"}
-        </Button>
-        {deletionError && (
-          <Callout title="Error" color="red" className="mt-4">
-            {deletionError}
-          </Callout>
-        )}
-      </div>
+      </>
    </Modal>
  );
 }
--- a/web/src/app/admin/embeddings/modals/ProviderCreationModal.tsx
+++ b/web/src/app/admin/embeddings/modals/ProviderCreationModal.tsx
@ -36,6 +36,7 @@ export function ProviderCreationModal({
      ? Object.entries(existingProvider.custom_config)
      : [],
    model_id: 0,
+    model_name: null,
  };

  const validationSchema = Yup.object({
@ -45,6 +46,9 @@ export function ProviderCreationModal({
      : useFileUpload
        ? Yup.string()
        : Yup.string().required("API Key is required"),
+    model_name: isProxy
+      ? Yup.string().required("Model name is required")
+      : Yup.string().nullable(),
    api_url: isProxy
      ? Yup.string().required("API URL is required")
      : Yup.string(),
@ -96,6 +100,7 @@ export function ProviderCreationModal({
            provider_type: values.provider_type.toLowerCase().split(" ")[0],
            api_key: values.api_key,
            api_url: values.api_url,
+            model_name: values.model_name,
          }),
        }
      );
@ -182,15 +187,25 @@ export function ProviderCreationModal({
                </a>
              </Text>

-              <div className="flex w-full flex-col gap-y-2">
-                {isProxy ? (
-                  <TextFormField
-                    name="api_url"
-                    label="API URL"
-                    placeholder="API URL"
-                    type="text"
-                  />
-                ) : useFileUpload ? (
+              <div className="flex w-full flex-col gap-y-6">
+                {isProxy && (
+                  <>
+                    <TextFormField
+                      name="api_url"
+                      label="API URL"
+                      placeholder="API URL"
+                      type="text"
+                    />
+                    <TextFormField
+                      name="model_name"
+                      label="Model Name (for testing)"
+                      placeholder="Model Name"
+                      type="text"
+                    />
+                  </>
+                )}
+
+                {useFileUpload ? (
                  <>
                    <Label>Upload JSON File</Label>
                    <input
@ -205,7 +220,7 @@ export function ProviderCreationModal({
                ) : (
                  <TextFormField
                    name="api_key"
-                    label="API Key"
+                    label={`API Key ${isProxy && "(for non-local deployments)"}`}
                    placeholder="API Key"
                    type="password"
                  />
--- a/web/src/app/admin/embeddings/pages/CloudEmbeddingPage.tsx
+++ b/web/src/app/admin/embeddings/pages/CloudEmbeddingPage.tsx
@ -12,10 +12,13 @@ import {
  LITELLM_CLOUD_PROVIDER,
 } from "../../../../components/embedding/interfaces";
 import { EmbeddingDetails } from "../EmbeddingModelSelectionForm";
-import { FiExternalLink, FiInfo } from "react-icons/fi";
+import { FiExternalLink, FiInfo, FiTrash } from "react-icons/fi";
 import { HoverPopup } from "@/components/HoverPopup";
 import { Dispatch, SetStateAction, useEffect, useState } from "react";
 import { LiteLLMModelForm } from "@/components/embedding/LiteLLMModelForm";
+import { deleteSearchSettings } from "./utils";
+import { usePopup } from "@/components/admin/connectors/Popup";
+import { DeleteEntityModal } from "@/components/modals/DeleteEntityModal";

 export default function CloudEmbeddingPage({
  currentModel,
@ -181,7 +184,7 @@ export default function CloudEmbeddingPage({
                onClick={() => setShowTentativeProvider(LITELLM_CLOUD_PROVIDER)}
                className="mb-2 px-4 py-2 bg-blue-500 text-white rounded hover:bg-blue-600 text-sm cursor-pointer"
              >
-                Provide API URL
+                Set API Configuration
              </button>
            ) : (
              <button
@ -190,7 +193,7 @@ export default function CloudEmbeddingPage({
                }
                className="mb-2 hover:underline text-sm cursor-pointer"
              >
-                Modify API URL
+                Modify API Configuration
              </button>
            )}

@ -283,11 +286,33 @@ export function CloudModelCard({
    React.SetStateAction<CloudEmbeddingProvider | null>
  >;
 }) {
+  const { popup, setPopup } = usePopup();
+  const [showDeleteModel, setShowDeleteModel] = useState(false);
  const enabled =
    model.model_name === currentModel.model_name &&
    model.provider_type?.toLowerCase() ==
      currentModel.provider_type?.toLowerCase();

+  const deleteModel = async () => {
+    if (!model.id) {
+      setPopup({ message: "Model cannot be deleted", type: "error" });
+      return;
+    }
+
+    const response = await deleteSearchSettings(model.id);
+
+    if (response.ok) {
+      setPopup({ message: "Model deleted successfully", type: "success" });
+      setShowDeleteModel(false);
+    } else {
+      setPopup({
+        message:
+          "Failed to delete model. Ensure you are not attempting to delete a curently active model.",
+        type: "error",
+      });
+    }
+  };
+
  return (
    <div
      className={`p-4 w-96 border rounded-lg transition-all duration-200 ${
@ -296,17 +321,38 @@ export function CloudModelCard({
          : "border-gray-300 hover:border-blue-300 hover:shadow-sm"
      } ${!provider.configured && "opacity-80 hover:opacity-100"}`}
    >
+      {popup}
+      {showDeleteModel && (
+        <DeleteEntityModal
+          entityName={model.model_name}
+          entityType="embedding model configuration"
+          onSubmit={() => deleteModel()}
+          onClose={() => setShowDeleteModel(false)}
+        />
+      )}
+
      <div className="flex items-center justify-between mb-3">
        <h3 className="font-bold text-lg">{model.model_name}</h3>
-        <a
-          href={provider.website}
-          target="_blank"
-          rel="noopener noreferrer"
-          onClick={(e) => e.stopPropagation()}
-          className="text-blue-500 hover:text-blue-700 transition-colors duration-200"
-        >
-          <FiExternalLink size={18} />
-        </a>
+        <div className="flex gap-x-2">
+          {model.provider_type == EmbeddingProvider.LITELLM.toLowerCase() && (
+            <button
+              onClickCapture={() => setShowDeleteModel(true)}
+              onClick={(e) => e.stopPropagation()}
+              className="text-blue-500 hover:text-blue-700 transition-colors duration-200"
+            >
+              <FiTrash size={18} />
+            </button>
+          )}
+          <a
+            href={provider.website}
+            target="_blank"
+            rel="noopener noreferrer"
+            onClick={(e) => e.stopPropagation()}
+            className="text-blue-500 hover:text-blue-700 transition-colors duration-200"
+          >
+            <FiExternalLink size={18} />
+          </a>
+        </div>
      </div>
      <p className="text-sm text-gray-600 mb-2">{model.description}</p>
      {model?.provider_type?.toLowerCase() !=
--- a/web/src/app/admin/embeddings/pages/utils.ts
+++ b/web/src/app/admin/embeddings/pages/utils.ts
@ -0,0 +1,10 @@
+export const deleteSearchSettings = async (search_settings_id: number) => {
+  const response = await fetch(`/api/search-settings/delete-search-settings`, {
+    method: "DELETE",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify({ search_settings_id }),
+  });
+  return response;
+};
--- a/web/src/components/embedding/interfaces.tsx
+++ b/web/src/components/embedding/interfaces.tsx
@ -39,6 +39,7 @@ export interface CloudEmbeddingProvider {

 // Embedding Models
 export interface EmbeddingModelDescriptor {
+  id?: number;
  model_name: string;
  model_dim: number;
  normalize: boolean;