diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index ac82b3dbada..cdea83554d6 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -559,20 +559,28 @@ def get_team_provider_credentials( return {key: value for key, value in credentials.items() if key != "model"} return None - # 1. Prefer the team's own BYOK deployment, matched by model_info.team_id. + def _usable_team_deployment_id(deployment: dict) -> str | None: + model_info: Final = deployment.get("model_info") or {} + if model_info.get("team_id") != team_id: + return None + if not _key_may_use(model_info.get("team_public_model_name") or deployment.get("model_name")): + return None + return model_info.get("id") + + # 1. Prefer the team's own BYOK deployment, matched by model_info.team_id. The lowest + # deployment id wins so every worker picks the same one in any load or edit order. if team_id is not None: - for deployment in llm_router.model_list or []: - model_info = deployment.get("model_info") or {} - if model_info.get("team_id") != team_id: - continue - deployment_id = model_info.get("id") - if deployment_id is None: - continue - if not _key_may_use(model_info.get("team_public_model_name") or deployment.get("model_name")): - continue - credentials = _provider_credentials(deployment_id) - if credentials is not None: - return credentials + team_deployment_ids: Final = sorted( + deployment_id + for deployment in llm_router.model_list or [] + if (deployment_id := _usable_team_deployment_id(deployment)) is not None + ) + team_credentials: Final = next( + (credentials for credentials in map(_provider_credentials, team_deployment_ids) if credentials is not None), + None, + ) + if team_credentials is not None: + return team_credentials if not include_accessible_deployments: return None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3eb11738af6..bdaa4af4e26 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8307,7 +8307,7 @@ class ProxyConfig: try: new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository( WriterPinnedClient(prisma_client.db) - ).table.find_many(order=[{"created_at": "asc"}, {"model_id": "asc"}]) + ).table.find_many() self.auto_router_db_catalog = build_auto_router_catalog(new_models) return new_models except Exception as e: diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 92ddb423086..b607cc3fa7f 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -690,7 +690,7 @@ async def rag_ingest( if response.get("status") == "failed": raise HTTPException( status_code=response.get("error_status_code") or 500, - detail={"error": response.get("error")}, + detail={"id": response.get("id"), "status": "failed", "error": response.get("error")}, ) # Save vector store to database if it was newly created and prisma_client is available diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index 0d9ba2f3ddb..75eb6eab887 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -34,6 +34,10 @@ if TYPE_CHECKING: from litellm import Router +class RAGIngestBadRequestError(ValueError): + status_code: Final = 400 + + class BaseRAGIngestion(ABC): """ Base class for RAG ingestion. diff --git a/litellm/rag/ingestion/openai_ingestion.py b/litellm/rag/ingestion/openai_ingestion.py index 6c7c619237d..e6a3f0bb74f 100644 --- a/litellm/rag/ingestion/openai_ingestion.py +++ b/litellm/rag/ingestion/openai_ingestion.py @@ -10,7 +10,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Final import litellm -from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError from litellm.vector_store_files.main import acreate as vector_store_file_acreate from litellm.vector_stores.main import acreate as vector_store_acreate @@ -87,7 +87,7 @@ class OpenAIRAGIngestion(BaseRAGIngestion): api_base: Final = self.vector_store_config.get("api_base") if existing_file_id and not vector_store_id: - raise ValueError("vector_store_id is required when ingesting an existing file_id") + raise RAGIngestBadRequestError("vector_store_id is required when ingesting an existing file_id") # Create vector store if not provided if not vector_store_id: diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index eff8ad1b8cb..cceddf31117 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -22,7 +22,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError if TYPE_CHECKING: from litellm import Router @@ -134,7 +134,9 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): Tuple of (rag_corpus_id, file_id) """ if not self.project_id: - raise ValueError("vertex_project is required for Vertex AI RAG ingestion. Set it in vector_store config.") + raise RAGIngestBadRequestError( + "vertex_project is required for Vertex AI RAG ingestion. Set it in vector_store config." + ) # Get or create RAG corpus rag_corpus_id = self.vector_store_config.get("vector_store_id") diff --git a/tests/integration/authorization/test_rag_team_provider_credentials.py b/tests/integration/authorization/test_rag_team_provider_credentials.py index 8ee38fe6060..92a6805f5e9 100644 --- a/tests/integration/authorization/test_rag_team_provider_credentials.py +++ b/tests/integration/authorization/test_rag_team_provider_credentials.py @@ -1275,35 +1275,64 @@ async def test_worker_sigkill_mid_query_burst_leaves_the_sibling_on_team_credent assert sum(1 for call in calls if call.route == "/v1/chat/completions") == held_by[survivor_pid] + 1, calls -def test_a_team_with_two_openai_deployments_ingests_through_the_older_one_after_it_is_edited(tmp_path: Path) -> None: +def _ingest_and_drain( + gateway: Gateway, scenario: Scenario, key: str, upstream: Wire +) -> tuple[str, tuple[Signature, ...]]: + upstream.drain() + response: Final = _ingest(gateway, key) + assert response.status_code == 200, response.text + store_id: Final = string_value(object_value(response.json())["vector_store_id"]) + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store_id}) + calls: Final = tuple(_call(request) for request in upstream.drain()) + return store_id, tuple(call.signature for call in calls if call.route != "/v1/models") + + +def _ingests_after_editing( + gateway: Gateway, scenario: Scenario, key: str, upstream: Wire, model_id: str +) -> tuple[tuple[str, tuple[Signature, ...]], ...]: + edited: Final = gateway.request("PATCH", f"/model/{model_id}/update", {"litellm_params": {"timeout": 30}}) + assert edited.status_code == 200, edited.text + return tuple(_ingest_and_drain(gateway, scenario, key, upstream) for _ in range(4)) + + +def test_a_team_with_two_openai_deployments_ingests_through_the_same_one_after_either_is_edited( + tmp_path: Path, +) -> None: with ( gateway_from_environment() as admin_gateway, admin_gateway.scenario() as scenario, wire_server(_responder(threading.Event())) as upstream, ): - older_model: Final = f"integration-rag-older-{uuid.uuid4().hex}" - newer_model: Final = f"integration-rag-newer-{uuid.uuid4().hex}" - team_id: Final = scenario.team(models=[older_model, newer_model]) - older_id: Final = _deployment( - admin_gateway, - scenario, - name=older_model, - model="openai/gpt-4o-mini", - api_key=TEAM_API_KEY, - api_base=f"{upstream.url}/team/v1", - team_id=team_id, + team_model: Final = f"integration-rag-team-{uuid.uuid4().hex}" + wildcard_model: Final = f"integration-rag-wildcard-{uuid.uuid4().hex}" + team_id: Final = scenario.team(models=[team_model, wildcard_model]) + team_deployment: Final = ( + _deployment( + admin_gateway, + scenario, + name=team_model, + model="openai/gpt-4o-mini", + api_key=TEAM_API_KEY, + api_base=f"{upstream.url}/team/v1", + team_id=team_id, + ), + "team", + TEAM_API_KEY, ) - _deployment( - admin_gateway, - scenario, - name=newer_model, - model="openai/gpt-4o-mini", - api_key=WILDCARD_API_KEY, - api_base=f"{upstream.url}/wildcard/v1", - team_id=team_id, + wildcard_deployment: Final = ( + _deployment( + admin_gateway, + scenario, + name=wildcard_model, + model="openai/gpt-4o-mini", + api_key=WILDCARD_API_KEY, + api_base=f"{upstream.url}/wildcard/v1", + team_id=team_id, + ), + "wildcard", + WILDCARD_API_KEY, ) - edited: Final = admin_gateway.request("PATCH", f"/model/{older_id}/update", {"litellm_params": {"timeout": 30}}) - assert edited.status_code == 200, edited.text + (picked_id, picked_prefix, picked_key), (other_id, _, _) = sorted((team_deployment, wildcard_deployment)) key: Final = scenario.key(team_id=team_id) with ( owned_proxy( @@ -1315,12 +1344,12 @@ def test_a_team_with_two_openai_deployments_ingests_through_the_older_one_after_ ) as gateway, gateway.scenario() as owned_scenario, ): - upstream.drain() - response: Final = _ingest(gateway, key) - assert response.status_code == 200, response.text - store_id: Final = string_value(object_value(response.json())["vector_store_id"]) - owned_scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store_id}) - - assert tuple(_call(request).signature for request in upstream.drain()) == _ingest_signatures( - "team", TEAM_API_KEY, store_id + ingests: Final = ( + _ingest_and_drain(gateway, owned_scenario, key, upstream), + *_ingests_after_editing(gateway, owned_scenario, key, upstream, picked_id), + *_ingests_after_editing(gateway, owned_scenario, key, upstream, other_id), ) + + assert [calls for _, calls in ingests] == [ + _ingest_signatures(picked_prefix, picked_key, store_id) for store_id, _ in ingests + ] diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index 5c1c371fd32..9e2fb6f1eaf 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -464,9 +464,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( ) -def _completed_batch_for_retire( - output_file_id: str | None, counts: BatchRequestCounts | None = None -) -> LiteLLMBatch: +def _completed_batch_for_retire(output_file_id: str | None, counts: BatchRequestCounts | None = None) -> LiteLLMBatch: kwargs = dict( id="batch-1", completion_window="24h", @@ -714,3 +712,30 @@ def test_a_restricted_key_resolves_the_team_deployment_its_wildcard_or_model_cov @pytest.mark.parametrize("key_models", [["claude-opus-4-6"], ["anthropic/*"]]) def test_a_restricted_key_outside_the_team_wildcard_never_resolves_its_credentials(key_models: list[str]): assert _team_only_openai_credentials("openai/*", "openai/*", key_models) is None + + +def _team_openai_deployment(deployment_id: str, api_key: str) -> dict: + return { + "model_name": f"team-gpt-{deployment_id}", + "litellm_params": {"model": "openai/gpt-4o", "api_key": api_key}, + "model_info": {"id": deployment_id, "team_id": "team-a"}, + } + + +@pytest.mark.parametrize("deployment_ids", [("deployment-a", "deployment-b"), ("deployment-b", "deployment-a")]) +def test_a_team_with_two_openai_deployments_resolves_the_same_one_in_any_router_order( + deployment_ids: tuple[str, str], +): + router: Final = Router( + model_list=[_team_openai_deployment(deployment_id, f"{deployment_id}-key") for deployment_id in deployment_ids] + ) + + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="team-a"), + custom_llm_provider="openai", + include_accessible_deployments=False, + ) + + assert credentials is not None + assert credentials["api_key"] == "deployment-a-key" diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 6f894dab878..6a4525560f9 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -556,6 +556,7 @@ def test_rag_ingest_failed_upstream_status_is_returned_without_persisting_store( "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", new=AsyncMock( return_value={ + "id": "ingest_abc", "status": "failed", "error": "unauthorized", "error_status_code": 401, @@ -571,7 +572,7 @@ def test_rag_ingest_failed_upstream_status_is_returned_without_persisting_store( ) assert response.status_code == 401, response.json() - assert response.json() == {"detail": {"error": "unauthorized"}} + assert response.json() == {"detail": {"id": "ingest_abc", "status": "failed", "error": "unauthorized"}} managed_vector_store_table.create.assert_not_called() diff --git a/tests/unit/rag/ingestion/test_base_ingestion.py b/tests/unit/rag/ingestion/test_base_ingestion.py index 50b1823767f..8bf54085f6f 100644 --- a/tests/unit/rag/ingestion/test_base_ingestion.py +++ b/tests/unit/rag/ingestion/test_base_ingestion.py @@ -8,6 +8,8 @@ import respx import litellm from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion +from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion +from litellm.rag.ingestion.vertex_ai_ingestion import VertexAIRAGIngestion from litellm.types.utils import CredentialItem _FILE_URL: Final = "https://files.example/docs/report.pdf" @@ -148,3 +150,24 @@ async def test_an_existing_file_id_on_a_store_that_cannot_take_one_fails_as_a_ba assert response["status"] == "failed" assert response["error_status_code"] == 400 assert "does not support ingesting an existing file_id" in str(response["error"]) + + +@pytest.mark.parametrize( + ("ingestion_class", "ingest_kwargs"), + [ + (OpenAIRAGIngestion, {"file_id": "file-abc"}), + (VertexAIRAGIngestion, {"file_data": ("report.txt", b"plain text", "text/plain")}), + ], +) +async def test_a_vector_store_config_missing_a_caller_field_fails_as_a_bad_request( + ingestion_class: type[BaseRAGIngestion], + ingest_kwargs: dict[str, object], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + ingestion: Final = ingestion_class(ingest_options={"vector_store": {"api_key": "caller-key"}}) + + response: Final = await ingestion.ingest(**ingest_kwargs) + + assert response["status"] == "failed" + assert response["error_status_code"] == 400 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx index ff0706b24cb..9ee41af6b84 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx @@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import CreateVectorStore from "./CreateVectorStore"; import * as networking from "@/components/networking"; +import { toast } from "@/lib/toast"; // Mock the networking module vi.mock("@/components/networking", () => ({ @@ -215,6 +216,43 @@ describe("CreateVectorStore", () => { }); }); + it("should report the created store and the failed file when a later file fails to ingest", async () => { + vi.mocked(networking.ragIngestCall) + .mockResolvedValueOnce({ id: "ingest-1", status: "completed", vector_store_id: "vs_123", file_id: "file_1" }) + .mockRejectedValueOnce(new Error("upstream rejected the file")); + const onSuccess = vi.fn(); + render(); + + const uploadInput = document.querySelector('input[type="file"]') as HTMLInputElement; + await act(async () => { + fireEvent.change(uploadInput, { + target: { + files: [ + new File(["first"], "first.txt", { type: "text/plain" }), + new File(["second"], "second.txt", { type: "text/plain" }), + ], + }, + }); + }); + await waitFor(() => { + expect(screen.getByText("Uploaded Documents (2)")).toBeInTheDocument(); + }); + + await act(async () => { + fireEvent.click(screen.getByRole("button", { name: /Create Vector Store/i })); + }); + + await waitFor(() => { + expect(toast.warning).toHaveBeenCalledWith( + "Vector store vs_123 was created with 1 document(s), but second.txt failed to ingest", + { description: "upstream rejected the file" }, + ); + }); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(onSuccess).toHaveBeenCalledWith("vs_123"); + expect(screen.getByText("Vector Store Created Successfully")).toBeInTheDocument(); + }); + it("should exclude valkey from the provider dropdown since it has no RAG ingestion", async () => { render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx index 8f24e47340c..2422e83355d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx @@ -175,6 +175,7 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu setIsCreating(true); const results: RAGIngestResponse[] = []; let vectorStoreId: string | undefined; + let failure: { name: string; error: unknown } | undefined; try { // Ingest each document @@ -208,19 +209,35 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu console.error(`Error ingesting ${doc.name}:`, error); // Update document status to error setDocuments((prev) => prev.map((d) => (d.uid === doc.uid ? { ...d, status: "error" as const } : d))); - throw error; // Stop processing on first error + failure = { name: doc.name, error }; + break; } } + if (failure && !vectorStoreId) { + throw failure.error; + } + setIngestResults(results); - toast.success( - `Successfully created vector store with ${results.length} document(s). Vector Store ID: ${vectorStoreId}`, - ); + if (failure) { + toast.warning( + `Vector store ${vectorStoreId} was created with ${results.length} document(s), but ${failure.name} failed to ingest`, + { description: failure.error instanceof Error ? failure.error.message : String(failure.error) }, + ); + } else { + toast.success( + `Successfully created vector store with ${results.length} document(s). Vector Store ID: ${vectorStoreId}`, + ); + } if (onSuccess && vectorStoreId) { onSuccess(vectorStoreId); } + if (failure) { + return; + } + // Clear documents after successful creation setTimeout(() => { setDocuments([]);