mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(rag): keep forwarded provider keys on /rag/query and mixed JWT team claims working
A forwarded x-api-key now counts as a caller credential, so the /rag/query search uses it over the team deployment as at the merge base. Sorting the JWT team claim compares entries as text, so a stray number or null no longer turns into a 401. A Gemini store with no proxy key is a 500 again, and a non-string vector_store_id keeps the merge base's 400 message.
This commit is contained in:
mateo-berri
parent
7155c37c16
commit
f0b056949f
7 changed files with 806 additions and 80 deletions
|
|
@ -1705,7 +1705,7 @@ class JWTAuthManager:
|
|||
return None, None
|
||||
|
||||
any_claim_team_resolved = False
|
||||
for team_id in sorted(team_ids):
|
||||
for team_id in sorted(team_ids, key=str):
|
||||
try:
|
||||
team_object = await get_team_object(
|
||||
team_id=team_id,
|
||||
|
|
|
|||
|
|
@ -837,7 +837,7 @@ async def rag_query(
|
|||
if not isinstance(vector_store_id, str):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "retrieval_config must contain a string 'vector_store_id'"},
|
||||
detail={"error": "vector_store_id must be a non-empty string"},
|
||||
)
|
||||
reject_caller_embedding_selection_params(payload=retrieval_config, source="retrieval_config")
|
||||
resolved_stores: Final = await _authorize_nested_vector_store_ids(
|
||||
|
|
@ -867,25 +867,25 @@ async def rag_query(
|
|||
{key: value for key, value in retrieval_config.items() if key not in _EXPLICIT_CREDENTIAL_KEYS}
|
||||
)
|
||||
)
|
||||
team_credentials: Final = _team_provider_credentials(
|
||||
effective_vector_store_config,
|
||||
llm_router,
|
||||
user_api_key_dict,
|
||||
)
|
||||
vector_store_params: Final = (
|
||||
{**team_credentials, **managed_store_params} if managed_store is not None else team_credentials
|
||||
)
|
||||
|
||||
# Add litellm data
|
||||
request_data: dict[str, object] = {}
|
||||
request_data = await add_litellm_data_to_request(
|
||||
data=request_data,
|
||||
request_data: Final[dict[str, object]] = await add_litellm_data_to_request(
|
||||
data={},
|
||||
request=request,
|
||||
general_settings=general_settings,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
version=version,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
forwarded_credentials: Final = _without_blank_credentials(
|
||||
{key: value for key, value in request_data.items() if key in _EXPLICIT_CREDENTIAL_KEYS}
|
||||
)
|
||||
team_credentials: Final = _team_provider_credentials(
|
||||
{**effective_vector_store_config, **forwarded_credentials},
|
||||
llm_router,
|
||||
user_api_key_dict,
|
||||
)
|
||||
vector_store_params: Final = (
|
||||
{**team_credentials, **managed_store_params} if managed_store is not None else team_credentials
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"RAG Query - model: %s, vector_store_id: %s, custom_llm_provider: %s",
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.llms.gemini.common_utils import GeminiModelInfo
|
||||
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError
|
||||
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -101,10 +101,10 @@ class GeminiRAGIngestion(BaseRAGIngestion):
|
|||
api_base: Final = cast(str | None, vector_store_config.get("api_base")) or GeminiModelInfo.get_api_base()
|
||||
|
||||
if not api_key:
|
||||
raise RAGIngestBadRequestError("GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search")
|
||||
raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search")
|
||||
|
||||
if not api_base:
|
||||
raise RAGIngestBadRequestError("GEMINI_API_BASE is required")
|
||||
raise ValueError("GEMINI_API_BASE is required")
|
||||
|
||||
api_version: Final = "v1beta"
|
||||
base_url: Final = f"{api_base}/{api_version}"
|
||||
|
|
|
|||
|
|
@ -16,8 +16,10 @@ from types import MappingProxyType
|
|||
from typing import Final, Literal
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import jwt
|
||||
import openai
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
|
|
@ -60,12 +62,31 @@ BATCH_INPUT: Final = ("batch.jsonl", (json.dumps(dict(BATCH_INPUT_ROW)) + "\n").
|
|||
FAILING_STORE_NAMES: Final = MappingProxyType({"rag-upstream-401": 401, "rag-upstream-429": 429})
|
||||
OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
|
||||
STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
PREFIXED: Final = re.compile(r"^/(team|wildcard|shared|environment|credential|caller|gemini|azure)(/[^?]*)")
|
||||
PREFIXED: Final = re.compile(r"^/(team|wildcard|shared|environment|credential|caller|gemini|azure|bedrock)(/[^?]*)")
|
||||
STORE_FILES: Final = re.compile(r"^/v1/vector_stores/([^/]+)/files$")
|
||||
FILE_CONTENT: Final = re.compile(r"^/v1/files/([^/]+)/content$")
|
||||
FILE_ITEM: Final = re.compile(r"^/v1/files/([^/]+)$")
|
||||
BATCH_ITEM: Final = re.compile(r"^/v1/batches/([^/]+)$")
|
||||
BATCH_CANCEL: Final = re.compile(r"^/v1/batches/([^/]+)/cancel$")
|
||||
BEDROCK_DATA_SOURCES: Final = re.compile(r"^/knowledgebases/([^/]+)/datasources/$")
|
||||
BEDROCK_DATA_SOURCE: Final = re.compile(r"^/knowledgebases/([^/]+)/datasources/([^/]+)$")
|
||||
EMPTY_KNOWLEDGE_BASE: Final = "EMPTYKB000"
|
||||
BUCKETLESS_KNOWLEDGE_BASE: Final = "NOBUCKET00"
|
||||
WEB_DATA_SOURCE: Final = "WEBSOURCE0"
|
||||
STORE_SEARCH: Final = re.compile(r"^(?:/v1|/openai)/vector_stores/([^/]+)/search$")
|
||||
GEMINI_UPLOAD: Final = re.compile(r"^/upload/v1beta/fileSearchStores/([^:]+):uploadToFileSearchStore$")
|
||||
|
||||
JWT_TEAM_COUNT: Final = 10
|
||||
AWS_AMBIENT_CREDENTIALS: Final = (
|
||||
"AWS_PROFILE",
|
||||
"AWS_SESSION_TOKEN",
|
||||
"AWS_BEARER_TOKEN_BEDROCK",
|
||||
"AWS_ROLE_ARN",
|
||||
"AWS_ROLE_NAME",
|
||||
"AWS_SESSION_NAME",
|
||||
"AWS_WEB_IDENTITY_TOKEN_FILE",
|
||||
)
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(OWNED_PROXY_CELL_SECONDS)
|
||||
|
||||
Signature = tuple[str, str, str, str | None]
|
||||
|
|
@ -100,8 +121,13 @@ class Tenants:
|
|||
wildcard_bare_key: str
|
||||
wildcard_prefixed_key: str
|
||||
wildcard_global_key: str
|
||||
team_s: str
|
||||
team_s_key: str
|
||||
teamless_key: str
|
||||
jwt_teams: tuple[str, ...]
|
||||
jwt_team_model: str
|
||||
jwt_multi_member: str
|
||||
azure_restricted_key: str
|
||||
credential_name: str
|
||||
unreachable_credential_name: str
|
||||
|
||||
|
|
@ -115,10 +141,10 @@ class RagRig:
|
|||
signing_key: rsa.RSAPrivateKey
|
||||
scenario: Scenario
|
||||
|
||||
def jwt(self, subject: str, team_id: str) -> str:
|
||||
def jwt(self, subject: str, *team_ids: str) -> str:
|
||||
issued_at: Final = int(time.time())
|
||||
return jwt.encode(
|
||||
{"sub": subject, "groups": [team_id], "iat": issued_at, "exp": issued_at + 300},
|
||||
{"sub": subject, "groups": list(team_ids), "iat": issued_at, "exp": issued_at + 300},
|
||||
self.signing_key,
|
||||
algorithm="RS256",
|
||||
headers={"kid": JWT_KEY_ID},
|
||||
|
|
@ -181,44 +207,127 @@ def _vector_store_reply(request: Request, outage: threading.Event) -> Reply:
|
|||
)
|
||||
|
||||
|
||||
def _file_reply() -> Reply:
|
||||
return _json_reply(
|
||||
{
|
||||
"id": f"file-{uuid.uuid4().hex}",
|
||||
"object": "file",
|
||||
"bytes": len(DOCUMENT[1]),
|
||||
"created_at": int(time.time()),
|
||||
"filename": DOCUMENT[0],
|
||||
"purpose": "assistants",
|
||||
"status": "processed",
|
||||
}
|
||||
)
|
||||
def _owned_by(prefix: str, identity: str) -> bool:
|
||||
return identity.startswith((f"file-{prefix}-", f"batch_{prefix}_"))
|
||||
|
||||
|
||||
def _file_object(identity: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": identity,
|
||||
"object": "file",
|
||||
"bytes": len(DOCUMENT[1]),
|
||||
"created_at": int(time.time()),
|
||||
"filename": DOCUMENT[0],
|
||||
"purpose": "assistants",
|
||||
"status": "processed",
|
||||
}
|
||||
|
||||
|
||||
def _file_reply(prefix: str) -> Reply:
|
||||
return _json_reply(_file_object(f"file-{prefix}-{uuid.uuid4().hex}"))
|
||||
|
||||
|
||||
def _batch_input_reply() -> Reply:
|
||||
return Reply(status=200, body=BATCH_INPUT[1])
|
||||
|
||||
|
||||
def _batch_reply(request: Request) -> Reply:
|
||||
def _batch_object(identity: str, input_file_id: str, status: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": identity,
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"input_file_id": input_file_id,
|
||||
"completion_window": "24h",
|
||||
"status": status,
|
||||
"created_at": int(time.time()),
|
||||
"errors": None,
|
||||
"output_file_id": None,
|
||||
"error_file_id": None,
|
||||
"request_counts": {"total": 0, "completed": 0, "failed": 0},
|
||||
"metadata": None,
|
||||
}
|
||||
|
||||
|
||||
def _batch_reply(request: Request, prefix: str) -> Reply:
|
||||
body: Final = _json_body(request)
|
||||
return _json_reply(
|
||||
_batch_object(f"batch_{prefix}_{uuid.uuid4().hex}", str(body.get("input_file_id")), "validating")
|
||||
)
|
||||
|
||||
|
||||
def _list_reply() -> Reply:
|
||||
return _json_reply({"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False})
|
||||
|
||||
|
||||
def _owned_reply(prefix: str, identity: str, payload: Mapping[str, JsonValue]) -> Reply:
|
||||
return _json_reply(payload) if _owned_by(prefix, identity) else _error_reply(404)
|
||||
|
||||
|
||||
def _openai_get(route: str, prefix: str) -> Reply:
|
||||
if route in {"/v1/files", "/v1/batches"}:
|
||||
return _list_reply()
|
||||
content: Final = FILE_CONTENT.match(route)
|
||||
if content is not None:
|
||||
return _batch_input_reply() if _owned_by(prefix, content[1]) else _error_reply(404)
|
||||
file_item: Final = FILE_ITEM.match(route)
|
||||
if file_item is not None:
|
||||
return _owned_reply(prefix, file_item[1], _file_object(file_item[1]))
|
||||
batch_item: Final = BATCH_ITEM.match(route)
|
||||
if batch_item is not None:
|
||||
return _owned_reply(prefix, batch_item[1], _batch_object(batch_item[1], f"file-{prefix}-input", "validating"))
|
||||
return _error_reply(404)
|
||||
|
||||
|
||||
def _openai_delete(route: str, prefix: str) -> Reply:
|
||||
file_item: Final = FILE_ITEM.match(route)
|
||||
if file_item is None:
|
||||
return _error_reply(404)
|
||||
return _owned_reply(prefix, file_item[1], {"id": file_item[1], "object": "file", "deleted": True})
|
||||
|
||||
|
||||
def _embedding_reply(request: Request) -> Reply:
|
||||
body: Final = _json_body(request)
|
||||
inputs: Final = body.get("input")
|
||||
count: Final = len(inputs) if isinstance(inputs, list) else 1
|
||||
return _json_reply(
|
||||
{
|
||||
"id": f"batch_{uuid.uuid4().hex}",
|
||||
"object": "batch",
|
||||
"endpoint": str(body.get("endpoint")),
|
||||
"input_file_id": str(body.get("input_file_id")),
|
||||
"completion_window": str(body.get("completion_window")),
|
||||
"status": "validating",
|
||||
"created_at": int(time.time()),
|
||||
"errors": None,
|
||||
"output_file_id": None,
|
||||
"error_file_id": None,
|
||||
"request_counts": {"total": 0, "completed": 0, "failed": 0},
|
||||
"metadata": None,
|
||||
"object": "list",
|
||||
"data": [{"object": "embedding", "index": index, "embedding": [0.1, 0.2, 0.3]} for index in range(count)],
|
||||
"model": str(body.get("model")),
|
||||
"usage": {"prompt_tokens": count, "total_tokens": count},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _bedrock_reply(request: Request, route: str) -> Reply:
|
||||
listed: Final = BEDROCK_DATA_SOURCES.match(route)
|
||||
if request.method == "POST" and listed is not None:
|
||||
summary: Final[dict[str, JsonValue]] = {
|
||||
"knowledgeBaseId": listed[1],
|
||||
"dataSourceId": WEB_DATA_SOURCE,
|
||||
"name": "integration-web-source",
|
||||
"status": "AVAILABLE",
|
||||
"updatedAt": "2026-01-01T00:00:00Z",
|
||||
}
|
||||
return _json_reply({"dataSourceSummaries": [] if listed[1] == EMPTY_KNOWLEDGE_BASE else [summary]})
|
||||
detail: Final = BEDROCK_DATA_SOURCE.match(route)
|
||||
if request.method == "GET" and detail is not None:
|
||||
return _json_reply(
|
||||
{
|
||||
"dataSource": {
|
||||
"knowledgeBaseId": detail[1],
|
||||
"dataSourceId": detail[2],
|
||||
"name": "integration-web-source",
|
||||
"status": "AVAILABLE",
|
||||
"dataSourceConfiguration": {"type": "WEB"},
|
||||
"createdAt": "2026-01-01T00:00:00Z",
|
||||
"updatedAt": "2026-01-01T00:00:00Z",
|
||||
}
|
||||
}
|
||||
)
|
||||
return _error_reply(404)
|
||||
|
||||
|
||||
def _attach_reply(request: Request, store_id: str) -> Reply:
|
||||
return _json_reply(
|
||||
{
|
||||
|
|
@ -294,28 +403,59 @@ def _chat_reply(request: Request) -> Reply:
|
|||
|
||||
def _responses_reply(request: Request) -> Reply:
|
||||
identity: Final = uuid.uuid4().hex
|
||||
return _json_reply(
|
||||
{
|
||||
"id": f"resp_{identity}",
|
||||
"object": "response",
|
||||
"created_at": int(time.time()),
|
||||
"status": "completed",
|
||||
"model": str(_json_body(request).get("model")),
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": f"msg_{identity}",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "integration response", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"parallel_tool_calls": False,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
completed: Final = _responses_object(identity, str(_json_body(request).get("model")))
|
||||
if _json_body(request).get("stream") is not True:
|
||||
return _json_reply(completed)
|
||||
events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = (
|
||||
(
|
||||
"response.created",
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {**completed, "status": "in_progress", "output": []},
|
||||
},
|
||||
),
|
||||
(
|
||||
"response.output_text.delta",
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": f"msg_{identity}",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "integration response",
|
||||
"logprobs": [],
|
||||
},
|
||||
),
|
||||
("response.completed", {"type": "response.completed", "sequence_number": 2, "response": completed}),
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events),
|
||||
)
|
||||
|
||||
|
||||
def _responses_object(identity: str, model: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"resp_{identity}",
|
||||
"object": "response",
|
||||
"created_at": int(time.time()),
|
||||
"status": "completed",
|
||||
"model": model,
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": f"msg_{identity}",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "integration response", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"parallel_tool_calls": False,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
def _gemini_reply(request: Request, route: str) -> Reply:
|
||||
|
|
@ -330,21 +470,28 @@ def _gemini_reply(request: Request, route: str) -> Reply:
|
|||
return _error_reply(404)
|
||||
|
||||
|
||||
def _openai_reply(request: Request, route: str, outage: threading.Event) -> Reply:
|
||||
def _openai_reply(request: Request, prefix: str, route: str, outage: threading.Event) -> Reply:
|
||||
if request.method == "GET":
|
||||
return _batch_input_reply() if FILE_CONTENT.match(route) is not None else _error_reply(404)
|
||||
return _openai_get(route, prefix)
|
||||
if request.method == "DELETE":
|
||||
return _openai_delete(route, prefix)
|
||||
if request.method != "POST":
|
||||
return _error_reply(404)
|
||||
if route == "/v1/vector_stores":
|
||||
return _vector_store_reply(request, outage)
|
||||
if route == "/v1/files":
|
||||
return _file_reply()
|
||||
if route in {"/v1/files", "/openai/files"}:
|
||||
return _file_reply(prefix)
|
||||
if route == "/v1/batches":
|
||||
return _batch_reply(request)
|
||||
return _batch_reply(request, prefix)
|
||||
cancelled: Final = BATCH_CANCEL.match(route)
|
||||
if cancelled is not None:
|
||||
return _owned_reply(prefix, cancelled[1], _batch_object(cancelled[1], f"file-{prefix}-input", "cancelling"))
|
||||
if route == "/v1/chat/completions":
|
||||
return _chat_reply(request)
|
||||
if route == "/v1/responses":
|
||||
return _responses_reply(request)
|
||||
if route == "/v1/embeddings":
|
||||
return _embedding_reply(request)
|
||||
attach: Final = STORE_FILES.match(route)
|
||||
if attach is not None:
|
||||
return _attach_reply(request, attach[1])
|
||||
|
|
@ -360,7 +507,9 @@ def _responder(outage: threading.Event) -> Callable[[Request], Reply]:
|
|||
return _error_reply(404)
|
||||
if matched[1] == "gemini":
|
||||
return _gemini_reply(request, matched[2])
|
||||
return _openai_reply(request, matched[2], outage)
|
||||
if matched[1] == "bedrock":
|
||||
return _bedrock_reply(request, matched[2])
|
||||
return _openai_reply(request, matched[1], matched[2], outage)
|
||||
|
||||
return respond
|
||||
|
||||
|
|
@ -375,7 +524,7 @@ def _proxy_config(directory: Path) -> Path:
|
|||
"enable_jwt_auth": True,
|
||||
"litellm_jwtauth": {"user_id_jwt_field": "sub", "team_ids_jwt_field": "groups"},
|
||||
},
|
||||
"files_settings": [{"custom_llm_provider": "openai"}],
|
||||
"files_settings": [{"custom_llm_provider": "openai"}, {"custom_llm_provider": "azure"}],
|
||||
}
|
||||
path: Final = directory / "rag_team_provider_credentials.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
|
|
@ -397,6 +546,38 @@ def _provider_environment(upstream: str) -> Mapping[str, str]:
|
|||
)
|
||||
|
||||
|
||||
def _aws_environment(upstream: str) -> Mapping[str, str]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
"AWS_ENDPOINT_URL_BEDROCK_AGENT": f"{upstream}/bedrock",
|
||||
"AWS_ACCESS_KEY_ID": "integration-rag-aws-access-key",
|
||||
"AWS_SECRET_ACCESS_KEY": "integration-rag-aws-secret-key",
|
||||
"AWS_REGION_NAME": "us-east-1",
|
||||
"AWS_DEFAULT_REGION": "us-east-1",
|
||||
"AWS_EC2_METADATA_DISABLED": "true",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _jwt_team_key(team_id: str) -> str:
|
||||
return f"integration-rag-jwt-{team_id}"
|
||||
|
||||
|
||||
def _jwt_team(gateway: Gateway, scenario: Scenario, upstream: str, model: str, member: str) -> str:
|
||||
team_id: Final = scenario.team(models=[model])
|
||||
_deployment(
|
||||
gateway,
|
||||
scenario,
|
||||
name=model,
|
||||
model="openai/gpt-4o-mini",
|
||||
api_key=_jwt_team_key(team_id),
|
||||
api_base=f"{upstream}/team/v1",
|
||||
team_id=team_id,
|
||||
)
|
||||
gateway.post("/team/member_add", {"team_id": team_id, "member": {"user_id": member, "role": "user"}})
|
||||
return team_id
|
||||
|
||||
|
||||
def _deployment(
|
||||
gateway: Gateway,
|
||||
scenario: Scenario,
|
||||
|
|
@ -486,6 +667,18 @@ def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants:
|
|||
team_id=team_w,
|
||||
)
|
||||
team_s: Final = scenario.team(models=[shared_model])
|
||||
team_z: Final = scenario.team(models=[shared_model])
|
||||
_deployment(
|
||||
gateway,
|
||||
scenario,
|
||||
name="azure/*",
|
||||
model="azure/*",
|
||||
api_key=WILDCARD_API_KEY,
|
||||
api_base=f"{upstream}/azure",
|
||||
team_id=team_z,
|
||||
)
|
||||
jwt_team_model: Final = f"integration-rag-jwt-{uuid.uuid4().hex}"
|
||||
jwt_multi_member: Final = scenario.user(user_role="internal_user")
|
||||
return Tenants(
|
||||
shared_model=shared_model,
|
||||
global_model=global_model,
|
||||
|
|
@ -500,8 +693,15 @@ def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants:
|
|||
wildcard_bare_key=scenario.key(team_id=team_w, models=[BARE_WILDCARD_MODEL]),
|
||||
wildcard_prefixed_key=scenario.key(team_id=team_w, models=[WILDCARD_MODEL]),
|
||||
wildcard_global_key=scenario.key(team_id=team_w, models=[global_model]),
|
||||
team_s=team_s,
|
||||
team_s_key=scenario.key(team_id=team_s),
|
||||
teamless_key=scenario.key(),
|
||||
jwt_teams=tuple(
|
||||
_jwt_team(gateway, scenario, upstream, jwt_team_model, jwt_multi_member) for _ in range(JWT_TEAM_COUNT)
|
||||
),
|
||||
jwt_team_model=jwt_team_model,
|
||||
jwt_multi_member=jwt_multi_member,
|
||||
azure_restricted_key=scenario.key(team_id=team_z, models=["azure/gpt-4o-mini"]),
|
||||
credential_name=_credential(gateway, scenario, f"{upstream}/credential/v1"),
|
||||
unreachable_credential_name=_credential(gateway, scenario, _closed_url()),
|
||||
)
|
||||
|
|
@ -540,9 +740,13 @@ def rag_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[RagRig]:
|
|||
owned_proxy(
|
||||
admin_gateway,
|
||||
directory,
|
||||
{**_provider_environment(upstream.url), "JWT_PUBLIC_KEY_URL": jwks.url},
|
||||
{
|
||||
**_provider_environment(upstream.url),
|
||||
**_aws_environment(upstream.url),
|
||||
"JWT_PUBLIC_KEY_URL": jwks.url,
|
||||
},
|
||||
config=_proxy_config(directory),
|
||||
remove_environment=("GOOGLE_API_KEY",),
|
||||
remove_environment=("GOOGLE_API_KEY", *AWS_AMBIENT_CREDENTIALS),
|
||||
workers=2,
|
||||
) as gateway,
|
||||
gateway.scenario() as owned_scenario,
|
||||
|
|
@ -550,6 +754,12 @@ def rag_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[RagRig]:
|
|||
yield RagRig(gateway, upstream, outage, tenants, signing_key, owned_scenario)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _drain_calls_an_earlier_failed_cell_left(request: pytest.FixtureRequest) -> None:
|
||||
if "rag_rig" in request.fixturenames:
|
||||
request.getfixturevalue("rag_rig").upstream.drain()
|
||||
|
||||
|
||||
def _ingest_request(name: str, vector_store: Mapping[str, JsonValue]) -> Mapping[str, str]:
|
||||
return MappingProxyType(
|
||||
{"request": json.dumps({"ingest_options": {"name": name, "vector_store": dict(vector_store)}})}
|
||||
|
|
@ -1056,7 +1266,7 @@ def test_non_string_vector_store_id_is_rejected_before_search(rag_rig: RagRig, v
|
|||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "retrieval_config must contain a string 'vector_store_id'" in response.text, response.text
|
||||
assert "vector_store_id must be a non-empty string" in response.text, response.text
|
||||
assert rag_rig.calls() == ()
|
||||
|
||||
|
||||
|
|
@ -1234,7 +1444,7 @@ async def test_worker_sigkill_mid_query_burst_leaves_the_sibling_on_team_credent
|
|||
held_searches.put(request.target)
|
||||
assert release.wait(timeout=60), "The burst was never released"
|
||||
return _search_reply(request)
|
||||
return _openai_reply(request, matched[2], threading.Event())
|
||||
return _openai_reply(request, matched[1], matched[2], threading.Event())
|
||||
|
||||
with (
|
||||
gateway_from_environment() as admin_gateway,
|
||||
|
|
@ -1364,3 +1574,456 @@ def test_a_team_with_two_openai_deployments_keeps_ingesting_through_the_older_on
|
|||
assert [calls for _, calls in ingests] == [
|
||||
_ingest_signatures(picked_prefix, picked_key, store_id) for store_id, _ in ingests
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("stored", "prefix", "api_key"),
|
||||
(("own_key", "caller", CALLER_API_KEY), ("null_key", "team", TEAM_API_KEY)),
|
||||
ids=("own_key", "null_key"),
|
||||
)
|
||||
def test_ingest_into_a_managed_store_uses_its_own_key_else_the_team_key(
|
||||
rag_rig: RagRig, stored: Literal["own_key", "null_key"], prefix: str, api_key: str
|
||||
) -> None:
|
||||
litellm_params: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{
|
||||
"own_key": {"api_key": CALLER_API_KEY, "api_base": f"{rag_rig.upstream.url}/caller/v1"},
|
||||
"null_key": {"api_key": None},
|
||||
}[stored]
|
||||
)
|
||||
with rag_rig.gateway.scenario() as scenario:
|
||||
store_id: Final = _managed_store(rag_rig, scenario, litellm_params)
|
||||
|
||||
response: Final = _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, {**OPENAI, "vector_store_id": store_id})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert object_value(response.json())["vector_store_id"] == store_id, response.text
|
||||
bearer: Final = f"Bearer {api_key}"
|
||||
assert rag_rig.signatures() == (
|
||||
(prefix, "POST", "/v1/files", bearer),
|
||||
(prefix, "POST", f"/v1/vector_stores/{store_id}/files", bearer),
|
||||
)
|
||||
|
||||
|
||||
def _file_id_ingest(gateway: Gateway, key: str, vector_store: Mapping[str, JsonValue]) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/rag/ingest",
|
||||
{
|
||||
"file_id": "file-team-existing",
|
||||
"ingest_options": {"name": "litellm-rag-ingest", "vector_store": dict(vector_store)},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("vector_store", "ingestion"),
|
||||
(
|
||||
({"custom_llm_provider": "gemini"}, "GeminiRAGIngestion"),
|
||||
({"custom_llm_provider": "bedrock", "vector_store_id": BUCKETLESS_KNOWLEDGE_BASE}, "BedrockRAGIngestion"),
|
||||
({"custom_llm_provider": "s3_vectors", "vector_bucket_name": "integration-bucket"}, "S3VectorsRAGIngestion"),
|
||||
({"custom_llm_provider": "vertex_ai"}, "VertexAIRAGIngestion"),
|
||||
),
|
||||
ids=("gemini", "bedrock", "s3_vectors", "vertex_ai"),
|
||||
)
|
||||
def test_existing_file_id_on_a_store_that_cannot_take_one_is_rejected_before_any_upstream_call(
|
||||
rag_rig: RagRig, vector_store: Mapping[str, JsonValue], ingestion: str
|
||||
) -> None:
|
||||
before: Final = _store_rows("")
|
||||
|
||||
response: Final = _file_id_ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, vector_store)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert f"{ingestion} does not support ingesting an existing file_id" in response.text, response.text
|
||||
assert rag_rig.calls() == ()
|
||||
assert _store_rows("") == before
|
||||
|
||||
|
||||
def test_openai_file_id_without_a_vector_store_id_is_rejected_before_any_upstream_call(rag_rig: RagRig) -> None:
|
||||
before: Final = _store_rows("")
|
||||
|
||||
response: Final = _file_id_ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, OPENAI)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "vector_store_id is required when ingesting an existing file_id" in response.text, response.text
|
||||
assert rag_rig.calls() == ()
|
||||
assert _store_rows("") == before
|
||||
|
||||
|
||||
def test_openai_file_id_attaches_to_an_existing_store_with_the_team_key(rag_rig: RagRig) -> None:
|
||||
store_id: Final = f"vs_existing_{uuid.uuid4().hex}"
|
||||
|
||||
response: Final = _file_id_ingest(
|
||||
rag_rig.gateway, rag_rig.tenants.team_a_key, {**OPENAI, "vector_store_id": store_id}
|
||||
)
|
||||
|
||||
assert _completed_store(rag_rig, response) == store_id
|
||||
assert rag_rig.signatures() == (("team", "POST", f"/v1/vector_stores/{store_id}/files", f"Bearer {TEAM_API_KEY}"),)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("knowledge_base", "error", "routes"),
|
||||
(
|
||||
(
|
||||
EMPTY_KNOWLEDGE_BASE,
|
||||
f"No data sources found for Knowledge Base {EMPTY_KNOWLEDGE_BASE}",
|
||||
(("POST", f"/knowledgebases/{EMPTY_KNOWLEDGE_BASE}/datasources/"),),
|
||||
),
|
||||
(
|
||||
BUCKETLESS_KNOWLEDGE_BASE,
|
||||
f"Could not auto-detect S3 bucket for data source {WEB_DATA_SOURCE}",
|
||||
(
|
||||
("POST", f"/knowledgebases/{BUCKETLESS_KNOWLEDGE_BASE}/datasources/"),
|
||||
("GET", f"/knowledgebases/{BUCKETLESS_KNOWLEDGE_BASE}/datasources/{WEB_DATA_SOURCE}"),
|
||||
),
|
||||
),
|
||||
),
|
||||
ids=("no_data_source", "no_s3_bucket"),
|
||||
)
|
||||
def test_bedrock_knowledge_base_the_caller_named_without_a_usable_data_source_fails_ingest_with_400(
|
||||
rag_rig: RagRig, knowledge_base: str, error: str, routes: tuple[tuple[str, str], ...]
|
||||
) -> None:
|
||||
before: Final = _store_rows("")
|
||||
|
||||
response: Final = _ingest(
|
||||
rag_rig.gateway,
|
||||
rag_rig.tenants.team_a_key,
|
||||
{"custom_llm_provider": "bedrock", "vector_store_id": knowledge_base},
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert error in response.text, response.text
|
||||
assert tuple((call.prefix, call.method, call.route) for call in rag_rig.calls()) == tuple(
|
||||
("bedrock", method, route) for method, route in routes
|
||||
)
|
||||
assert _store_rows("") == before
|
||||
|
||||
|
||||
def test_s3_vectors_bucket_name_the_caller_sent_is_rejected_with_400_after_embedding(rag_rig: RagRig) -> None:
|
||||
tenants: Final = rag_rig.tenants
|
||||
before: Final = _store_rows("")
|
||||
request: Final = {
|
||||
"request": json.dumps(
|
||||
{
|
||||
"ingest_options": {
|
||||
"name": "litellm-rag-ingest",
|
||||
"embedding": {"model": tenants.shared_model},
|
||||
"vector_store": {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "ab", "dimension": 3},
|
||||
}
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
response: Final = rag_rig.gateway.request_multipart(
|
||||
"/v1/rag/ingest", request, {"file": DOCUMENT}, key=tenants.team_a_key
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "Invalid vector_bucket_name 'ab'" in response.text, response.text
|
||||
assert rag_rig.signatures() == (("shared", "POST", "/v1/embeddings", f"Bearer {SHARED_API_KEY}"),)
|
||||
assert _store_rows("") == before
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class KeylessRig:
|
||||
gateway: Gateway
|
||||
upstream: Wire
|
||||
key: str
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def keyless_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[KeylessRig]:
|
||||
directory: Final = tmp_path_factory.mktemp("rag_keyless_providers")
|
||||
with (
|
||||
gateway_from_environment() as admin_gateway,
|
||||
admin_gateway.scenario() as scenario,
|
||||
wire_server(_responder(threading.Event())) as upstream,
|
||||
):
|
||||
key: Final = scenario.key()
|
||||
with owned_proxy(
|
||||
admin_gateway,
|
||||
directory,
|
||||
{"OPENAI_API_KEY": ENVIRONMENT_API_KEY, "OPENAI_API_BASE": f"{upstream.url}/environment/v1"},
|
||||
remove_environment=("GEMINI_API_KEY", "GOOGLE_API_KEY", "VERTEXAI_PROJECT"),
|
||||
workers=2,
|
||||
) as gateway:
|
||||
yield KeylessRig(gateway, upstream, key)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("provider", "status", "error"),
|
||||
(
|
||||
("gemini", 500, "GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search"),
|
||||
("vertex_ai", 400, "vertex_project is required for Vertex AI RAG ingestion"),
|
||||
),
|
||||
ids=("gemini_without_a_proxy_key", "vertex_ai_without_a_project"),
|
||||
)
|
||||
def test_ingest_status_says_whether_the_caller_or_the_proxy_must_fix_the_missing_setting(
|
||||
keyless_rig: KeylessRig, provider: str, status: int, error: str
|
||||
) -> None:
|
||||
before: Final = _store_rows("")
|
||||
|
||||
response: Final = _ingest(keyless_rig.gateway, keyless_rig.key, {"custom_llm_provider": provider})
|
||||
|
||||
assert response.status_code == status, response.text
|
||||
assert error in response.text, response.text
|
||||
assert tuple(keyless_rig.upstream.drain()) == ()
|
||||
assert _store_rows("") == before
|
||||
|
||||
|
||||
async def _jwt_chat_burst(rig: RagRig, token: str, model: str, count: int) -> tuple[httpx.Response, ...]:
|
||||
async with httpx.AsyncClient(base_url=str(rig.gateway.client.base_url), timeout=60, trust_env=False) as client:
|
||||
return tuple(
|
||||
await asyncio.gather(
|
||||
*(
|
||||
client.post(
|
||||
"/v1/chat/completions",
|
||||
json=_llm_body("/v1/chat/completions", model),
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
for _ in range(count)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def test_jwt_in_ten_teams_that_hold_the_model_always_resolves_the_lowest_team_id(rag_rig: RagRig) -> None:
|
||||
tenants: Final = rag_rig.tenants
|
||||
token: Final = rag_rig.jwt(tenants.jwt_multi_member, *reversed(tenants.jwt_teams))
|
||||
|
||||
responses: Final = await _jwt_chat_burst(rag_rig, token, tenants.jwt_team_model, 16)
|
||||
|
||||
assert [response.status_code for response in responses] == [200] * 16, [response.text for response in responses]
|
||||
assert (
|
||||
rag_rig.signatures()
|
||||
== (("team", "POST", "/v1/chat/completions", f"Bearer {_jwt_team_key(min(tenants.jwt_teams))}"),) * 16
|
||||
)
|
||||
|
||||
|
||||
def _prompt() -> str:
|
||||
return f"hello {uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def _sync_sdk_text(route: str, base_url: str, token: str, model: str, stream: bool) -> str:
|
||||
if route == "/v1/messages":
|
||||
messages_client: Final = anthropic.Anthropic(base_url=base_url, api_key=token, max_retries=0)
|
||||
request: Final = {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": _prompt()}]}
|
||||
if stream:
|
||||
return "".join(
|
||||
event.delta.text
|
||||
for event in messages_client.messages.create(**request, stream=True)
|
||||
if event.type == "content_block_delta" and event.delta.type == "text_delta"
|
||||
)
|
||||
message: Final = messages_client.messages.create(**request)
|
||||
return "".join(block.text for block in message.content if block.type == "text")
|
||||
client: Final = openai.OpenAI(base_url=f"{base_url}/v1", api_key=token, max_retries=0)
|
||||
if route == "/v1/responses":
|
||||
if stream:
|
||||
return "".join(
|
||||
event.delta
|
||||
for event in client.responses.create(model=model, input=_prompt(), stream=True)
|
||||
if event.type == "response.output_text.delta"
|
||||
)
|
||||
return client.responses.create(model=model, input=_prompt()).output_text
|
||||
chat_messages: Final = [{"role": "user", "content": _prompt()}]
|
||||
if stream:
|
||||
return "".join(
|
||||
chunk.choices[0].delta.content or ""
|
||||
for chunk in client.chat.completions.create(model=model, messages=chat_messages, stream=True)
|
||||
if chunk.choices
|
||||
)
|
||||
return client.chat.completions.create(model=model, messages=chat_messages).choices[0].message.content or ""
|
||||
|
||||
|
||||
async def _async_sdk_text(route: str, base_url: str, token: str, model: str, stream: bool) -> str:
|
||||
if route == "/v1/messages":
|
||||
messages_client: Final = anthropic.AsyncAnthropic(base_url=base_url, api_key=token, max_retries=0)
|
||||
request: Final = {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": _prompt()}]}
|
||||
if stream:
|
||||
return "".join(
|
||||
[
|
||||
event.delta.text
|
||||
async for event in await messages_client.messages.create(**request, stream=True)
|
||||
if event.type == "content_block_delta" and event.delta.type == "text_delta"
|
||||
]
|
||||
)
|
||||
message: Final = await messages_client.messages.create(**request)
|
||||
return "".join(block.text for block in message.content if block.type == "text")
|
||||
client: Final = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=token, max_retries=0)
|
||||
if route == "/v1/responses":
|
||||
if stream:
|
||||
return "".join(
|
||||
[
|
||||
event.delta
|
||||
async for event in await client.responses.create(model=model, input=_prompt(), stream=True)
|
||||
if event.type == "response.output_text.delta"
|
||||
]
|
||||
)
|
||||
return (await client.responses.create(model=model, input=_prompt())).output_text
|
||||
chat_messages: Final = [{"role": "user", "content": _prompt()}]
|
||||
if stream:
|
||||
return "".join(
|
||||
[
|
||||
chunk.choices[0].delta.content or ""
|
||||
async for chunk in await client.chat.completions.create(
|
||||
model=model, messages=chat_messages, stream=True
|
||||
)
|
||||
if chunk.choices
|
||||
]
|
||||
)
|
||||
completion: Final = await client.chat.completions.create(model=model, messages=chat_messages)
|
||||
return completion.choices[0].message.content or ""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
|
||||
@pytest.mark.parametrize("client", ("sync", "async"))
|
||||
@pytest.mark.parametrize("route", ("/v1/chat/completions", "/v1/messages", "/v1/responses"))
|
||||
async def test_jwt_team_member_reaches_its_llm_routes_through_the_sdks(
|
||||
rag_rig: RagRig, route: str, client: Literal["sync", "async"], stream: bool
|
||||
) -> None:
|
||||
tenants: Final = rag_rig.tenants
|
||||
base_url: Final = str(rag_rig.gateway.client.base_url).rstrip("/")
|
||||
token: Final = rag_rig.jwt(tenants.jwt_member, tenants.team_a)
|
||||
|
||||
text: Final = (
|
||||
await asyncio.to_thread(_sync_sdk_text, route, base_url, token, tenants.team_a_model, stream)
|
||||
if client == "sync"
|
||||
else await _async_sdk_text(route, base_url, token, tenants.team_a_model, stream)
|
||||
)
|
||||
|
||||
assert text == "integration response", text
|
||||
calls: Final = rag_rig.calls()
|
||||
assert [(call.prefix, call.authorization) for call in calls] == [("team", f"Bearer {TEAM_API_KEY}")], calls
|
||||
|
||||
|
||||
def _batch_for(rig: RagRig, key: str) -> tuple[str, str]:
|
||||
uploaded: Final = rig.gateway.request_multipart("/v1/files", {"purpose": "batch"}, {"file": BATCH_INPUT}, key=key)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
file_id: Final = string_value(object_value(uploaded.json())["id"])
|
||||
created: Final = rig.gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
|
||||
key=key,
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
rig.upstream.drain()
|
||||
return file_id, string_value(object_value(created.json())["id"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "path", "route"),
|
||||
(
|
||||
("GET", "/v1/files/{file_id}", "/v1/files/{file_id}"),
|
||||
("GET", "/v1/files/{file_id}/content", "/v1/files/{file_id}/content"),
|
||||
("DELETE", "/v1/files/{file_id}", "/v1/files/{file_id}"),
|
||||
("GET", "/v1/files?purpose=batch&custom_llm_provider=openai", "/v1/files"),
|
||||
("GET", "/v1/batches/{batch_id}", "/v1/batches/{batch_id}"),
|
||||
("POST", "/v1/batches/{batch_id}/cancel", "/v1/batches/{batch_id}/cancel"),
|
||||
),
|
||||
ids=("file_retrieve", "file_content", "file_delete", "file_list", "batch_retrieve", "batch_cancel"),
|
||||
)
|
||||
def test_files_and_batches_calls_by_a_key_restricted_under_the_team_wildcard_use_the_team_deployment(
|
||||
rag_rig: RagRig, method: str, path: str, route: str
|
||||
) -> None:
|
||||
key: Final = rag_rig.tenants.wildcard_bare_key
|
||||
file_id, batch_id = _batch_for(rag_rig, key)
|
||||
|
||||
response: Final = rag_rig.gateway.request(method, path.format(file_id=file_id, batch_id=batch_id), key=key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert rag_rig.signatures() == (
|
||||
("wildcard", method, route.format(file_id=file_id, batch_id=batch_id), f"Bearer {WILDCARD_API_KEY}"),
|
||||
)
|
||||
|
||||
|
||||
def test_batch_list_by_a_key_restricted_under_the_team_wildcard_is_answered_by_the_proxy_without_an_upstream_call(
|
||||
rag_rig: RagRig,
|
||||
) -> None:
|
||||
key: Final = rag_rig.tenants.wildcard_bare_key
|
||||
_batch_for(rag_rig, key)
|
||||
|
||||
response: Final = rag_rig.gateway.request("GET", "/v1/batches?custom_llm_provider=openai", key=key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert object_value(response.json())["object"] == "list", response.text
|
||||
assert rag_rig.signatures() == ()
|
||||
|
||||
|
||||
def _upload_batch_input(rig: RagRig, key: str) -> str:
|
||||
uploaded: Final = rig.gateway.request_multipart("/v1/files", {"purpose": "batch"}, {"file": BATCH_INPUT}, key=key)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
return string_value(object_value(uploaded.json())["id"])
|
||||
|
||||
|
||||
def _create_batch(rig: RagRig, key: str, file_id: str) -> httpx.Response:
|
||||
return rig.gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("team", "prefix", "api_key"),
|
||||
(("team_a", "team", TEAM_API_KEY), ("team_s", "shared", SHARED_API_KEY)),
|
||||
ids=("team_deployment", "accessible_deployment"),
|
||||
)
|
||||
def test_batch_over_the_key_tpm_limit_is_rejected_after_reading_its_input_where_it_was_uploaded(
|
||||
rag_rig: RagRig, team: Literal["team_a", "team_s"], prefix: str, api_key: str
|
||||
) -> None:
|
||||
key: Final = rag_rig.scenario.key(team_id=getattr(rag_rig.tenants, team), tpm_limit=1)
|
||||
file_id: Final = _upload_batch_input(rag_rig, key)
|
||||
|
||||
response: Final = _create_batch(rag_rig, key, file_id)
|
||||
|
||||
assert response.status_code == 429, response.text
|
||||
bearer: Final = f"Bearer {api_key}"
|
||||
assert rag_rig.signatures() == (
|
||||
(prefix, "POST", "/v1/files", bearer),
|
||||
(prefix, "GET", f"/v1/files/{file_id}/content", bearer),
|
||||
)
|
||||
|
||||
|
||||
def test_batch_whose_input_names_a_model_outside_the_key_allowlist_is_rejected_before_creation(
|
||||
rag_rig: RagRig,
|
||||
) -> None:
|
||||
key: Final = rag_rig.tenants.restricted_key
|
||||
file_id: Final = _upload_batch_input(rag_rig, key)
|
||||
|
||||
response: Final = _create_batch(rag_rig, key, file_id)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert BARE_WILDCARD_MODEL in response.text, response.text
|
||||
bearer: Final = f"Bearer {SHARED_API_KEY}"
|
||||
assert rag_rig.signatures() == (
|
||||
("shared", "POST", "/v1/files", bearer),
|
||||
("shared", "GET", f"/v1/files/{file_id}/content", bearer),
|
||||
)
|
||||
|
||||
|
||||
def test_azure_files_upload_by_a_key_restricted_under_the_team_azure_wildcard_uses_the_team_deployment(
|
||||
rag_rig: RagRig,
|
||||
) -> None:
|
||||
response: Final = rag_rig.gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "assistants", "custom_llm_provider": "azure"},
|
||||
{"file": DOCUMENT},
|
||||
key=rag_rig.tenants.azure_restricted_key,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert tuple((call.prefix, call.method, call.route, call.azure_key) for call in rag_rig.calls()) == (
|
||||
("azure", "POST", "/openai/files", WILDCARD_API_KEY),
|
||||
)
|
||||
|
||||
|
||||
def test_query_without_a_vector_store_id_is_rejected_before_search(rag_rig: RagRig) -> None:
|
||||
tenants: Final = rag_rig.tenants
|
||||
|
||||
response: Final = _query(rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, OPENAI)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "vector_store_id" in response.text, response.text
|
||||
assert rag_rig.calls() == ()
|
||||
|
|
|
|||
|
|
@ -338,6 +338,41 @@ async def test_a_jwt_in_two_teams_resolves_the_same_team_in_any_claim_order(clai
|
|||
assert team_id == "team-a"
|
||||
assert team_obj == teams["team-a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stray_claim", [1, None])
|
||||
async def test_a_jwt_team_claim_with_a_non_string_entry_still_resolves_the_string_team(stray_claim: object):
|
||||
jwt_handler: Final = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups")
|
||||
team: Final = LiteLLM_TeamTable(team_id="aa-team", models=["gpt-4"], metadata={})
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=lambda team_id, **_: {"aa-team": team}[team_id],
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.handle_jwt.can_team_access_model",
|
||||
new_callable=AsyncMock,
|
||||
return_value=True,
|
||||
),
|
||||
):
|
||||
team_id, team_obj = await JWTAuthManager.find_team_with_model_access(
|
||||
team_ids=set(jwt_handler.get_team_ids_from_jwt({"groups": ["aa-team", stray_claim]})),
|
||||
requested_model="gpt-4",
|
||||
route="/chat/completions",
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert team_id == "aa-team"
|
||||
assert team_obj == team
|
||||
|
||||
|
||||
_AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = {
|
||||
"test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": {
|
||||
"endpoint_id": "test-uuid-1",
|
||||
|
|
|
|||
|
|
@ -385,6 +385,7 @@ def _post_rag_query(
|
|||
retrieval_config: Mapping[str, object],
|
||||
managed_stores: Sequence[LiteLLM_ManagedVectorStore],
|
||||
router: Router,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
) -> httpx.Response:
|
||||
with (
|
||||
patch.object( # test-quality-ok: seeds the real managed-store registry the query path reads
|
||||
|
|
@ -402,6 +403,7 @@ def _post_rag_query(
|
|||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"retrieval_config": dict(retrieval_config),
|
||||
},
|
||||
headers=dict(headers or {}),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -688,6 +690,32 @@ def test_rag_query_without_a_team_deployment_keeps_the_environment_openai_creden
|
|||
assert search.calls.last.request.headers["authorization"] == f"Bearer {_ENVIRONMENT_OPENAI_KEY}"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("httpx_transport")
|
||||
def test_rag_query_search_uses_the_forwarded_provider_key_over_team_provider_credentials(
|
||||
client_team_a: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
forwarded_provider_key: Final = "sk-forwarded-provider-key"
|
||||
monkeypatch.setenv("OPENAI_API_KEY", _ENVIRONMENT_OPENAI_KEY)
|
||||
for unset in ("OPENAI_BASE_URL", "OPENAI_API_BASE"):
|
||||
monkeypatch.delenv(unset, raising=False)
|
||||
search: Final = respx_mock.post("https://api.openai.com/v1/vector_stores/provider-native-store/search").respond(
|
||||
json=_OPENAI_SEARCH_PAGE
|
||||
)
|
||||
|
||||
with patch.dict("litellm.proxy.proxy_server.general_settings", {"forward_llm_provider_auth_headers": True}):
|
||||
response: Final = _post_rag_query(
|
||||
client_team_a,
|
||||
retrieval_config={"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"},
|
||||
managed_stores=(),
|
||||
router=_team_provider_router(),
|
||||
headers={"Authorization": "Bearer sk-litellm-virtual-key", "x-api-key": forwarded_provider_key},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.json()
|
||||
assert search.call_count == 1
|
||||
assert search.calls.last.request.headers["authorization"] == f"Bearer {forwarded_provider_key}"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("httpx_transport")
|
||||
def test_rag_query_bedrock_search_keeps_the_callers_region_and_credentials_over_team_deployment(
|
||||
client_team_a: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
|
|
|
|||
|
|
@ -229,7 +229,7 @@ async def test_a_gemini_upstream_error_keeps_its_status_on_the_failed_ingest(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_gemini_store_with_no_api_key_anywhere_fails_as_a_bad_request(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_a_gemini_store_with_no_api_key_anywhere_fails_without_a_caller_error_status(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("GEMINI_API_KEY", raising=False)
|
||||
monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
|
||||
ingestion: Final = GeminiRAGIngestion(ingest_options={"vector_store": {"custom_llm_provider": "gemini"}})
|
||||
|
|
@ -237,7 +237,7 @@ async def test_a_gemini_store_with_no_api_key_anywhere_fails_as_a_bad_request(mo
|
|||
response: Final = await ingestion.ingest(file_data=_TEXT_FILE)
|
||||
|
||||
assert response["status"] == "failed"
|
||||
assert response["error_status_code"] == 400
|
||||
assert response.get("error_status_code") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue