sanitize llm keys and handle updates properly (#4270)

* sanitize llm keys and handle updates properly

* fix llm provider testing

* fix test

* mypy

* fix default model editing

---------

Co-authored-by: Richard Kuo (Danswer) <rkuo@onyx.app>
Co-authored-by: Richard Kuo <rkuo@rkuo.com>

backend/onyx/server/manage/llm/api.py

@@ -9,9 +9,9 @@ from sqlalchemy.orm import Session
 from onyx.auth.users import current_admin_user
 from onyx.auth.users import current_chat_accessible_user
 from onyx.db.engine import get_session
+from onyx.db.llm import fetch_existing_llm_provider
 from onyx.db.llm import fetch_existing_llm_providers
 from onyx.db.llm import fetch_existing_llm_providers_for_user
-from onyx.db.llm import fetch_provider
 from onyx.db.llm import remove_llm_provider
 from onyx.db.llm import update_default_provider
 from onyx.db.llm import update_default_vision_provider
@@ -24,9 +24,9 @@ from onyx.llm.llm_provider_options import WellKnownLLMProviderDescriptor
 from onyx.llm.utils import litellm_exception_to_error_msg
 from onyx.llm.utils import model_supports_image_input
 from onyx.llm.utils import test_llm
-from onyx.server.manage.llm.models import FullLLMProvider
 from onyx.server.manage.llm.models import LLMProviderDescriptor
 from onyx.server.manage.llm.models import LLMProviderUpsertRequest
+from onyx.server.manage.llm.models import LLMProviderView
 from onyx.server.manage.llm.models import TestLLMRequest
 from onyx.server.manage.llm.models import VisionProviderResponse
 from onyx.utils.logger import setup_logger
@@ -49,11 +49,27 @@ def fetch_llm_options(
 def test_llm_configuration(
     test_llm_request: TestLLMRequest,
     _: User | None = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
 ) -> None:
+    """Test regular llm and fast llm settings"""
+
+    # the api key is sanitized if we are testing a provider already in the system
+
+    test_api_key = test_llm_request.api_key
+    if test_llm_request.name:
+        # NOTE: we are querying by name. we probably should be querying by an invariant id, but
+        # as it turns out the name is not editable in the UI and other code also keys off name,
+        # so we won't rock the boat just yet.
+        existing_provider = fetch_existing_llm_provider(
+            test_llm_request.name, db_session
+        )
+        if existing_provider:
+            test_api_key = existing_provider.api_key
+
     llm = get_llm(
         provider=test_llm_request.provider,
         model=test_llm_request.default_model_name,
-        api_key=test_llm_request.api_key,
+        api_key=test_api_key,
         api_base=test_llm_request.api_base,
         api_version=test_llm_request.api_version,
         custom_config=test_llm_request.custom_config,
@@ -69,7 +85,7 @@ def test_llm_configuration(
         fast_llm = get_llm(
             provider=test_llm_request.provider,
             model=test_llm_request.fast_default_model_name,
-            api_key=test_llm_request.api_key,
+            api_key=test_api_key,
             api_base=test_llm_request.api_base,
             api_version=test_llm_request.api_version,
             custom_config=test_llm_request.custom_config,
@@ -119,11 +135,17 @@ def test_default_provider(
 def list_llm_providers(
     _: User | None = Depends(current_admin_user),
     db_session: Session = Depends(get_session),
-) -> list[FullLLMProvider]:
-    return [
-        FullLLMProvider.from_model(llm_provider_model)
-        for llm_provider_model in fetch_existing_llm_providers(db_session)
-    ]
+) -> list[LLMProviderView]:
+    llm_provider_list: list[LLMProviderView] = []
+    for llm_provider_model in fetch_existing_llm_providers(db_session):
+        full_llm_provider = LLMProviderView.from_model(llm_provider_model)
+        if full_llm_provider.api_key:
+            full_llm_provider.api_key = (
+                full_llm_provider.api_key[:4] + "****" + full_llm_provider.api_key[-4:]
+            )
+        llm_provider_list.append(full_llm_provider)
+
+    return llm_provider_list
 
 
 @admin_router.put("/provider")
@@ -135,11 +157,11 @@ def put_llm_provider(
     ),
     _: User | None = Depends(current_admin_user),
     db_session: Session = Depends(get_session),
-) -> FullLLMProvider:
+) -> LLMProviderView:
     # validate request (e.g. if we're intending to create but the name already exists we should throw an error)
     # NOTE: may involve duplicate fetching to Postgres, but we're assuming SQLAlchemy is smart enough to cache
     # the result
-    existing_provider = fetch_provider(db_session, llm_provider.name)
+    existing_provider = fetch_existing_llm_provider(llm_provider.name, db_session)
     if existing_provider and is_creation:
         raise HTTPException(
             status_code=400,
@@ -161,6 +183,11 @@ def put_llm_provider(
                 llm_provider.fast_default_model_name
             )
 
+    # the llm api key is sanitized when returned to clients, so the only time we
+    # should get a real key is when it is explicitly changed
+    if existing_provider and not llm_provider.api_key_changed:
+        llm_provider.api_key = existing_provider.api_key
+
     try:
         return upsert_llm_provider(
             llm_provider=llm_provider,
@@ -234,7 +261,7 @@ def get_vision_capable_providers(
 
         # Only include providers with at least one vision-capable model
         if vision_models:
-            provider_dict = FullLLMProvider.from_model(provider).model_dump()
+            provider_dict = LLMProviderView.from_model(provider).model_dump()
             provider_dict["vision_models"] = vision_models
             logger.info(
                 f"Vision provider: {provider.provider} with models: {vision_models}"

backend/onyx/server/manage/llm/models.py

@@ -12,6 +12,7 @@ if TYPE_CHECKING:
 
 class TestLLMRequest(BaseModel):
     # provider level
+    name: str | None = None
     provider: str
     api_key: str | None = None
     api_base: str | None = None
@@ -76,16 +77,19 @@ class LLMProviderUpsertRequest(LLMProvider):
     # should only be used for a "custom" provider
     # for default providers, the built-in model names are used
     model_names: list[str] | None = None
+    api_key_changed: bool = False
 
 
-class FullLLMProvider(LLMProvider):
+class LLMProviderView(LLMProvider):
+    """Stripped down representation of LLMProvider for display / limited access info only"""
+
     id: int
     is_default_provider: bool | None = None
     is_default_vision_provider: bool | None = None
     model_names: list[str]
 
     @classmethod
-    def from_model(cls, llm_provider_model: "LLMProviderModel") -> "FullLLMProvider":
+    def from_model(cls, llm_provider_model: "LLMProviderModel") -> "LLMProviderView":
         return cls(
             id=llm_provider_model.id,
             name=llm_provider_model.name,
@@ -111,7 +115,7 @@ class FullLLMProvider(LLMProvider):
         )
 
 
-class VisionProviderResponse(FullLLMProvider):
+class VisionProviderResponse(LLMProviderView):
     """Response model for vision providers endpoint, including vision-specific fields."""
 
     vision_models: list[str]