From 6ba7c27171c7f3980aee9f87d36dd5453cf9a0ac Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 20 May 2026 14:54:20 +0530 Subject: [PATCH 1/8] feat: Add evaluation batch submission functionality - Introduced `run_evaluation_batch_submission` task to handle evaluation batch submissions. - Created `start_evaluation_batch_submission` utility to enqueue evaluation tasks. - Updated `start_evaluation` to utilize the new Celery task for batch submissions. - Enhanced error handling in evaluation run creation and processing. - Added support for Google Gemini in evaluation batch processing. - Refactored model configuration to support multiple providers and completion types. - Improved logging for better traceability during evaluation processes. - Updated tests to reflect changes in evaluation batch submission logic. --- backend/app/celery/tasks/job_execution.py | 20 ++ backend/app/celery/utils.py | 18 ++ backend/app/core/batch/gemini.py | 2 +- backend/app/core/batch/operations.py | 3 +- backend/app/crud/evaluations/batch.py | 268 +++++++++++------- backend/app/crud/evaluations/langfuse.py | 16 +- backend/app/crud/evaluations/processing.py | 265 +++++++++++------ backend/app/crud/model_config.py | 75 ++++- backend/app/models/model_config.py | 6 +- backend/app/services/evaluations/batch_job.py | 92 ++++++ .../app/services/evaluations/evaluation.py | 69 ++--- .../app/tests/api/routes/test_evaluation.py | 2 +- 12 files changed, 592 insertions(+), 244 deletions(-) create mode 100644 backend/app/services/evaluations/batch_job.py diff --git a/backend/app/celery/tasks/job_execution.py b/backend/app/celery/tasks/job_execution.py index adadf1c9c..a106b88da 100644 --- a/backend/app/celery/tasks/job_execution.py +++ b/backend/app/celery/tasks/job_execution.py @@ -172,6 +172,26 @@ def run_delete_collection_job( ) +@celery_app.task(bind=True, queue="low_priority", priority=1) +@gevent_timeout(settings.CELERY_TASK_SOFT_TIME_LIMIT, "run_evaluation_batch_submission") +def run_evaluation_batch_submission( + self, project_id: int, job_id: str, trace_id: str, **kwargs +): + from app.services.evaluations.batch_job import execute_evaluation_batch_submission + + _set_trace(trace_id) + return _run_with_otel_parent( + self, + lambda: execute_evaluation_batch_submission( + project_id=project_id, + job_id=job_id, + task_id=current_task.request.id, + task_instance=self, + **kwargs, + ), + ) + + @celery_app.task(bind=True, queue="low_priority", priority=1) @gevent_timeout(settings.CELERY_TASK_SOFT_TIME_LIMIT, "run_stt_batch_submission") def run_stt_batch_submission( diff --git a/backend/app/celery/utils.py b/backend/app/celery/utils.py index 288cba7c4..0337c696d 100644 --- a/backend/app/celery/utils.py +++ b/backend/app/celery/utils.py @@ -126,6 +126,24 @@ def start_delete_collection_job( return task_id +def start_evaluation_batch_submission( + project_id: int, job_id: str, trace_id: str = "N/A", **kwargs +) -> str: + from app.celery.tasks.job_execution import run_evaluation_batch_submission + + task_id = _enqueue_with_trace_context( + run_evaluation_batch_submission, + project_id=project_id, + job_id=job_id, + trace_id=trace_id, + **kwargs, + ) + logger.info( + f"[start_evaluation_batch_submission] Started job {job_id} with Celery task {task_id}" + ) + return task_id + + def start_stt_batch_submission( project_id: int, job_id: str, trace_id: str = "N/A", **kwargs ) -> str: diff --git a/backend/app/core/batch/gemini.py b/backend/app/core/batch/gemini.py index f121bb9ee..4222afbb6 100644 --- a/backend/app/core/batch/gemini.py +++ b/backend/app/core/batch/gemini.py @@ -349,7 +349,7 @@ def upload_file(self, content: str, purpose: str = "batch") -> str: uploaded_file = self._upload_to_gemini( content=content, suffix=".jsonl", - mime_type="jsonl", + mime_type="application/jsonl", display_name=f"batch-input-{int(time.time())}", ) diff --git a/backend/app/core/batch/operations.py b/backend/app/core/batch/operations.py index d02884fdb..bcd51866b 100644 --- a/backend/app/core/batch/operations.py +++ b/backend/app/core/batch/operations.py @@ -74,7 +74,8 @@ def start_batch_job( logger.error(f"[start_batch_job] Failed | {e}", exc_info=True) batch_job_update = BatchJobUpdate( - error_message=f"Batch creation failed: {str(e)}" + provider_status="failed", + error_message=f"Batch creation failed: {str(e)}", ) update_batch_job( session=session, batch_job=batch_job, batch_job_update=batch_job_update diff --git a/backend/app/crud/evaluations/batch.py b/backend/app/crud/evaluations/batch.py index 4181910b8..604cc146a 100644 --- a/backend/app/crud/evaluations/batch.py +++ b/backend/app/crud/evaluations/batch.py @@ -11,15 +11,23 @@ from typing import Any from langfuse import Langfuse -from openai import OpenAI from sqlmodel import Session -from app.core.batch.base import BATCH_KEY - -from app.core.batch import OpenAIBatchProvider, start_batch_job +from app.core.batch import ( + BATCH_KEY, + GeminiBatchProvider, + OpenAIBatchProvider, + start_batch_job, +) +from app.core.batch.client import GeminiClient from app.models import EvaluationRun from app.models.batch_job import BatchJobType -from app.models.llm.request import KaapiLLMParams +from app.services.llm.mappers import ( + map_kaapi_to_google_params, + map_kaapi_to_openai_params, +) +from app.services.llm.providers.registry import LLMProvider +from app.utils import get_openai_client logger = logging.getLogger(__name__) @@ -62,149 +70,202 @@ def fetch_dataset_items(langfuse: Langfuse, dataset_name: str) -> list[dict[str, return items -def build_evaluation_jsonl( - dataset_items: list[dict[str, Any]], config: KaapiLLMParams +def build_openai_evaluation_jsonl( + dataset_items: list[dict[str, Any]], openai_params: dict[str, Any] ) -> list[dict[str, Any]]: + """Build OpenAI Responses API batch JSONL from Langfuse dataset items. + + Each line: + { + BATCH_KEY: , + "method": "POST", + "url": "/v1/responses", + "body": { ...openai_params, "input": } + } """ - Build JSONL data for evaluation batch using OpenAI Responses API. + jsonl_data: list[dict[str, Any]] = [] + for item in dataset_items: + question = item["input"].get("question", "") + if not question: + logger.warning( + f"[build_openai_evaluation_jsonl] Skipping item - no question found | item_id={item['id']}" + ) + continue - Each line is a dict with: - - BATCH_KEY: Unique identifier for the request (dataset item ID) - - method: POST - - url: /v1/responses - - body: Response request using config as-is with input from dataset + body = dict(openai_params) + body["input"] = question - Args: - dataset_items: List of dataset items from Langfuse - config: Evaluation configuration dict with OpenAI Responses API parameters. - This config is used as-is in the body, with only "input" being added - from the dataset. Config can include any fields like: - - model (required) - - instructions - - tools - - reasoning - - text - - temperature - - include - etc. + jsonl_data.append( + { + BATCH_KEY: item["id"], + "method": "POST", + "url": "/v1/responses", + "body": body, + } + ) + return jsonl_data - Returns: - List of dictionaries (JSONL data) + +def build_google_evaluation_jsonl( + dataset_items: list[dict[str, Any]], google_params: dict[str, Any] +) -> list[dict[str, Any]]: + """Build Gemini batch JSONL from Langfuse dataset items. + + Each line: + { + "key": , + "request": { contents, systemInstruction?, generationConfig? } + } """ - jsonl_data = [] + jsonl_data: list[dict[str, Any]] = [] + system_instruction = google_params.get("instructions") + + generation_config: dict[str, Any] = {} + temperature = google_params.get("temperature") + if temperature is not None: + generation_config["temperature"] = temperature + reasoning = google_params.get("reasoning") + if reasoning: + generation_config["thinkingConfig"] = { + "includeThoughts": False, + "thinkingLevel": reasoning, + } + for item in dataset_items: - # Extract question from input question = item["input"].get("question", "") if not question: logger.warning( - f"[build_evaluation_jsonl] Skipping item - no question found | item_id={item['id']}" + f"[build_google_evaluation_jsonl] Skipping item - no question found | item_id={item['id']}" ) continue - # Build the batch request object for Responses API - # Use config as-is and only add the input field - body: dict[str, Any] = { - "model": config.model, - "instructions": config.instructions, - "input": question, # Add input from dataset + request: dict[str, Any] = { + "contents": [{"parts": [{"text": question}], "role": "user"}], } + if system_instruction: + request["systemInstruction"] = {"parts": [{"text": system_instruction}]} + if generation_config: + request["generationConfig"] = generation_config - if "temperature" in config.model_fields_set: - body["temperature"] = config.temperature - - # Add reasoning only if provided - if config.reasoning: - body["reasoning"] = {"effort": config.reasoning} - - # Add tools only if knowledge_base_ids are provided - if config.knowledge_base_ids: - body["tools"] = [ - { - "type": "file_search", - "vector_store_ids": config.knowledge_base_ids, - "max_num_results": config.max_num_results or 20, - } - ] - - batch_request = { - BATCH_KEY: item["id"], - "method": "POST", - "url": "/v1/responses", - "body": body, - } + jsonl_data.append({"key": item["id"], "request": request}) - jsonl_data.append(batch_request) return jsonl_data def start_evaluation_batch( langfuse: Langfuse, - openai_client: OpenAI, session: Session, eval_run: EvaluationRun, - config: KaapiLLMParams, + params: dict[str, Any], + provider: str, ) -> EvaluationRun: """ - Fetch data, build JSONL, and start evaluation batch. - - This function orchestrates the evaluation-specific logic and delegates - to the generic batch infrastructure for actual batch creation. + Fetch dataset, build JSONL, submit batch via the appropriate provider. Args: langfuse: Configured Langfuse client - openai_client: Configured OpenAI client session: Database session - eval_run: EvaluationRun database object (with run_name, dataset_name, config) - config: KaapiLLMParams with model, instructions, knowledge_base_ids, etc. + eval_run: EvaluationRun database object + params: Kaapi-standardized completion params (dict) + provider: Completion provider ("openai" or "google", with optional "-native" suffix) Returns: Updated EvaluationRun with batch_job_id populated - - Raises: - Exception: If any step fails """ try: - # Step 1: Fetch dataset items from Langfuse logger.info( - f"[start_evaluation_batch] Starting evaluation batch | run={eval_run.run_name}" + f"[start_evaluation_batch] Starting evaluation batch | run={eval_run.run_name} | provider={provider}" ) dataset_items = fetch_dataset_items( langfuse=langfuse, dataset_name=eval_run.dataset_name ) - # Step 2: Build evaluation-specific JSONL - jsonl_data = build_evaluation_jsonl(dataset_items=dataset_items, config=config) + base_provider = provider.replace("-native", "") - if not jsonl_data: - raise ValueError( - "Evaluation dataset did not produce any JSONL entries (missing questions?)." + if base_provider == LLMProvider.OPENAI: + mapped_params, warnings = map_kaapi_to_openai_params( + session=session, kaapi_params=params ) + if warnings: + logger.info("[start_evaluation_batch] Mapper warnings: %s", warnings) - # Step 3: Create batch provider - provider = OpenAIBatchProvider(client=openai_client) + jsonl_data = build_openai_evaluation_jsonl( + dataset_items=dataset_items, openai_params=mapped_params + ) + if not jsonl_data: + raise ValueError( + "Evaluation dataset did not produce any JSONL entries (missing questions?)." + ) + + openai_client = get_openai_client( + session=session, + org_id=eval_run.organization_id, + project_id=eval_run.project_id, + ) + batch_provider = OpenAIBatchProvider(client=openai_client) - # Step 4: Prepare batch configuration - batch_config = { - "endpoint": "/v1/responses", - "description": f"Evaluation: {eval_run.run_name}", - "completion_window": "24h", - # Store complete config for reference - "evaluation_config": config.model_dump(exclude_unset=True), - } + batch_config = { + "endpoint": "/v1/responses", + "description": f"Evaluation: {eval_run.run_name}", + "completion_window": "24h", + "evaluation_config": params, + } - # Step 5: Start batch job using generic infrastructure - batch_job = start_batch_job( - session=session, - provider=provider, - provider_name="openai", - job_type=BatchJobType.EVALUATION, - organization_id=eval_run.organization_id, - project_id=eval_run.project_id, - jsonl_data=jsonl_data, - config=batch_config, - ) + batch_job = start_batch_job( + session=session, + provider=batch_provider, + provider_name="openai", + job_type=BatchJobType.EVALUATION, + organization_id=eval_run.organization_id, + project_id=eval_run.project_id, + jsonl_data=jsonl_data, + config=batch_config, + ) + + elif base_provider == LLMProvider.GOOGLE: + mapped_params, warnings = map_kaapi_to_google_params( + kaapi_params=params, completion_type="text" + ) + if warnings: + logger.info("[start_evaluation_batch] Mapper warnings: %s", warnings) + + jsonl_data = build_google_evaluation_jsonl( + dataset_items=dataset_items, google_params=mapped_params + ) + if not jsonl_data: + raise ValueError( + "Evaluation dataset did not produce any JSONL entries (missing questions?)." + ) + + gemini_client = GeminiClient.from_credentials( + session=session, + org_id=eval_run.organization_id, + project_id=eval_run.project_id, + ) + model_name = mapped_params.get("model", "gemini-2.5-pro") + batch_provider = GeminiBatchProvider( + client=gemini_client.client, model=model_name + ) + + batch_config = { + "display_name": f"evaluation-{eval_run.run_name}", + "model": f"models/{model_name}", + } + + batch_job = start_batch_job( + session=session, + provider=batch_provider, + provider_name="google", + job_type=BatchJobType.EVALUATION, + organization_id=eval_run.organization_id, + project_id=eval_run.project_id, + jsonl_data=jsonl_data, + config=batch_config, + ) + + else: + raise ValueError(f"Unsupported provider for evaluation batches: {provider}") - # Step 6: Link batch_job to evaluation_run eval_run.batch_job_id = batch_job.id eval_run.status = "processing" eval_run.total_items = batch_job.total_items @@ -217,7 +278,8 @@ def start_evaluation_batch( f"[start_evaluation_batch] Successfully started evaluation batch | " f"batch_job_id={batch_job.id} | " f"provider_batch_id={batch_job.provider_batch_id} | " - f"run={eval_run.run_name} | items={batch_job.total_items}" + f"run={eval_run.run_name} | items={batch_job.total_items} | " + f"provider={base_provider}" ) return eval_run diff --git a/backend/app/crud/evaluations/langfuse.py b/backend/app/crud/evaluations/langfuse.py index be5da4eb3..c72326829 100644 --- a/backend/app/crud/evaluations/langfuse.py +++ b/backend/app/crud/evaluations/langfuse.py @@ -145,11 +145,17 @@ def create_langfuse_dataset_run( trace_id_mapping[item_id] = trace_id except Exception as e: - logger.error( - f"[create_langfuse_dataset_run] Failed to create trace | " - f"item_id={item_id} | {e}", - exc_info=True, - ) + if getattr(e, "status_code", None) == 429: + logger.error( + f"[create_langfuse_dataset_run] Langfuse rate limit (429) | " + f"item_id={item_id}" + ) + else: + logger.error( + f"[create_langfuse_dataset_run] Failed to create trace | " + f"item_id={item_id} | {e}", + exc_info=True, + ) continue langfuse.flush() diff --git a/backend/app/crud/evaluations/processing.py b/backend/app/crud/evaluations/processing.py index 607e29e49..87671943a 100644 --- a/backend/app/crud/evaluations/processing.py +++ b/backend/app/crud/evaluations/processing.py @@ -9,6 +9,7 @@ """ import ast +import asyncio import json import logging from collections import defaultdict @@ -20,13 +21,17 @@ from sqlmodel import Session, select from app.core.batch import ( + GeminiBatchProvider, OpenAIBatchProvider, download_batch_results, poll_batch_status, upload_batch_results_to_object_store, ) -from app.core.batch.base import BATCH_KEY +from app.core.batch.base import BATCH_KEY, BatchProvider +from app.core.batch.client import GeminiClient +from app.core.batch.gemini import BatchJobState, extract_text_from_response_dict from app.crud.evaluations.batch import fetch_dataset_items +from app.services.llm.providers.registry import LLMProvider from app.crud.evaluations.core import resolve_model_from_config, update_evaluation_run from app.crud.evaluations.cost import attach_cost from app.crud.evaluations.embeddings import ( @@ -47,6 +52,32 @@ logger = logging.getLogger(__name__) +def _get_batch_provider( + session: Session, + provider_name: str, + organization_id: int, + project_id: int, +) -> BatchProvider: + """Get appropriate batch provider instance for the main response batch.""" + if provider_name in (LLMProvider.OPENAI, LLMProvider.OPENAI_NATIVE): + openai_client = get_openai_client( + session=session, + org_id=organization_id, + project_id=project_id, + ) + return OpenAIBatchProvider(client=openai_client) + + if provider_name in (LLMProvider.GOOGLE, LLMProvider.GOOGLE_NATIVE): + gemini_client = GeminiClient.from_credentials( + session=session, + org_id=organization_id, + project_id=project_id, + ) + return GeminiBatchProvider(client=gemini_client.client) + + raise ValueError(f"Unsupported provider for evaluation polling: {provider_name}") + + def _extract_batch_error_message( provider: OpenAIBatchProvider, error_file_id: str, @@ -115,8 +146,35 @@ def _extract_batch_error_message( return error_msg +def _extract_gemini_usage(response: dict[str, Any]) -> dict[str, Any] | None: + """Map Gemini usageMetadata to OpenAI-style usage dict for cost tracking.""" + usage_metadata = response.get("usageMetadata") or response.get("usage_metadata") + if not isinstance(usage_metadata, dict): + return None + input_tokens = usage_metadata.get("promptTokenCount") or usage_metadata.get( + "prompt_token_count" + ) + output_tokens = usage_metadata.get("candidatesTokenCount") or usage_metadata.get( + "candidates_token_count" + ) + total_tokens = usage_metadata.get("totalTokenCount") or usage_metadata.get( + "total_token_count" + ) + reasoning_tokens = usage_metadata.get("thoughtsTokenCount") or usage_metadata.get( + "thoughts_token_count" + ) + return { + "input_tokens": input_tokens or 0, + "output_tokens": output_tokens or 0, + "total_tokens": total_tokens or 0, + "reasoning_tokens": reasoning_tokens or 0, + } + + def parse_evaluation_output( - raw_results: list[dict[str, Any]], dataset_items: list[dict[str, Any]] + raw_results: list[dict[str, Any]], + dataset_items: list[dict[str, Any]], + provider_name: str = LLMProvider.OPENAI, ) -> list[dict[str, Any]]: """ Parse batch output into evaluation results. @@ -148,20 +206,19 @@ def parse_evaluation_output( """ # Create lookup map for dataset items by ID dataset_map = {item["id"]: item for item in dataset_items} + is_google = provider_name in (LLMProvider.GOOGLE, LLMProvider.GOOGLE_NATIVE) results = [] for line_num, response in enumerate(raw_results, 1): try: - # Extract BATCH_KEY (which is our dataset item ID) - item_id = response.get(BATCH_KEY) + item_id = response.get(BATCH_KEY) or response.get("key") if not item_id: logger.warning( - f"[parse_evaluation_output] No {BATCH_KEY} found, skipping | line={line_num}" + f"[parse_evaluation_output] No item_id found, skipping | line={line_num}" ) continue - # Get original dataset item dataset_item = dataset_map.get(item_id) if not dataset_item: logger.warning( @@ -169,64 +226,70 @@ def parse_evaluation_output( ) continue - # Extract the response body - response_body = response.get("response", {}).get("body", {}) + generated_output = "" + response_id: str | None = None + usage: dict[str, Any] | None = None - # Extract response ID from response.body.id - response_id = response_body.get("id") + if is_google: + gemini_response = response.get("response") + error = response.get("error") + if error: + error_msg = str(error) + logger.warning( + f"[parse_evaluation_output] Item had error | item_id={item_id} | {error_msg}" + ) + generated_output = f"ERROR: {error_msg}" + elif isinstance(gemini_response, dict): + text = extract_text_from_response_dict(gemini_response) + generated_output = text or "" + response_id = gemini_response.get( + "responseId" + ) or gemini_response.get("response_id") + usage = _extract_gemini_usage(gemini_response) + else: + response_body = response.get("response", {}).get("body", {}) + response_id = response_body.get("id") + usage = response_body.get("usage") - # Extract usage information for cost tracking - usage = response_body.get("usage") + if response.get("error"): + error_msg = response["error"].get("message", "Unknown error") + logger.warning( + f"[parse_evaluation_output] Item had error | item_id={item_id} | {error_msg}" + ) + generated_output = f"ERROR: {error_msg}" + else: + output = response_body.get("output", "") - # Handle errors in batch processing - if response.get("error"): - error_msg = response["error"].get("message", "Unknown error") - logger.warning( - f"[parse_evaluation_output] Item had error | item_id={item_id} | {error_msg}" - ) - generated_output = f"ERROR: {error_msg}" - else: - # Extract text from output (can be string, list, or complex structure) - output = response_body.get("output", "") - - # If string, try to parse it (may be JSON or Python repr of list) - if isinstance(output, str): - try: - output = json.loads(output) - except (json.JSONDecodeError, ValueError): + if isinstance(output, str): try: - output = ast.literal_eval(output) - except (ValueError, SyntaxError): - # Keep as string if parsing fails - generated_output = output - output = None - - # If we have a list structure, extract text from message items - if isinstance(output, list): - generated_output = "" - for item in output: - if isinstance(item, dict) and item.get("type") == "message": - for content in item.get("content", []): - if ( - isinstance(content, dict) - and content.get("type") == "output_text" - ): - generated_output = content.get("text", "") + output = json.loads(output) + except (json.JSONDecodeError, ValueError): + try: + output = ast.literal_eval(output) + except (ValueError, SyntaxError): + generated_output = output + output = None + + if isinstance(output, list): + for item in output: + if isinstance(item, dict) and item.get("type") == "message": + for content in item.get("content", []): + if ( + isinstance(content, dict) + and content.get("type") == "output_text" + ): + generated_output = content.get("text", "") + break + if generated_output: break - if generated_output: - break - elif output is not None: - # output was not a string and not a list - generated_output = "" - logger.warning( - f"[parse_evaluation_output] Unexpected output type | item_id={item_id} | type={type(output)}" - ) + elif output is not None: + generated_output = "" + logger.warning( + f"[parse_evaluation_output] Unexpected output type | item_id={item_id} | type={type(output)}" + ) - # Extract question and ground truth from dataset item question = dataset_item["input"].get("question", "") ground_truth = dataset_item["expected_output"].get("answer", "") - - # Extract question_id from dataset item metadata question_id = dataset_item.get("metadata", {}).get("question_id") results.append( @@ -300,8 +363,15 @@ async def process_completed_evaluation( logger.info( f"[process_completed_evaluation] {log_prefix} Downloading batch results | batch_job_id={batch_job.id}" ) - provider = OpenAIBatchProvider(client=openai_client) - raw_results = download_batch_results(provider=provider, batch_job=batch_job) + provider = _get_batch_provider( + session=session, + provider_name=batch_job.provider, + organization_id=eval_run.organization_id, + project_id=eval_run.project_id, + ) + raw_results = await asyncio.to_thread( + download_batch_results, provider=provider, batch_job=batch_job + ) # Step 2a: Upload raw results to object store for evaluation_run object_store_url = None @@ -318,13 +388,17 @@ async def process_completed_evaluation( logger.info( f"[process_completed_evaluation] {log_prefix} Fetching dataset items | dataset={eval_run.dataset_name}" ) - dataset_items = fetch_dataset_items( - langfuse=langfuse, dataset_name=eval_run.dataset_name + dataset_items = await asyncio.to_thread( + fetch_dataset_items, + langfuse=langfuse, + dataset_name=eval_run.dataset_name, ) # Step 4: Parse evaluation results results = parse_evaluation_output( - raw_results=raw_results, dataset_items=dataset_items + raw_results=raw_results, + dataset_items=dataset_items, + provider_name=batch_job.provider, ) if not results: @@ -348,7 +422,8 @@ async def process_completed_evaluation( update=EvaluationRunUpdate(cost=eval_run.cost), ) - trace_id_mapping = create_langfuse_dataset_run( + trace_id_mapping = await asyncio.to_thread( + create_langfuse_dataset_run, langfuse=langfuse, dataset_name=eval_run.dataset_name, model=model, @@ -365,7 +440,8 @@ async def process_completed_evaluation( # Step 6: Start embedding batch for similarity scoring # Pass trace_id_mapping directly without storing in DB try: - eval_run = start_embedding_batch( + eval_run = await asyncio.to_thread( + start_embedding_batch, session=session, openai_client=openai_client, eval_run=eval_run, @@ -463,8 +539,10 @@ async def process_completed_embedding_batch( # Step 2: Create provider and download results provider = OpenAIBatchProvider(client=openai_client) - raw_results = download_batch_results( - provider=provider, batch_job=embedding_batch_job + raw_results = await asyncio.to_thread( + download_batch_results, + provider=provider, + batch_job=embedding_batch_job, ) # Step 3: Parse embedding results @@ -497,7 +575,8 @@ async def process_completed_embedding_batch( per_item_scores = similarity_stats.get("per_item_scores", []) if per_item_scores: try: - update_traces_with_cosine_scores( + await asyncio.to_thread( + update_traces_with_cosine_scores, langfuse=langfuse, per_item_scores=per_item_scores, ) @@ -594,8 +673,11 @@ async def check_and_process_evaluation( if embedding_batch_job: # Poll embedding batch status provider = OpenAIBatchProvider(client=openai_client) - poll_batch_status( - session=session, provider=provider, batch_job=embedding_batch_job + await asyncio.to_thread( + poll_batch_status, + session=session, + provider=provider, + batch_job=embedding_batch_job, ) session.refresh(embedding_batch_job) @@ -666,23 +748,41 @@ async def check_and_process_evaluation( f"BatchJob {eval_run.batch_job_id} not found for evaluation {eval_run.id}" ) - # IMPORTANT: Poll OpenAI to get the latest status before checking - provider = OpenAIBatchProvider(client=openai_client) - status_result = poll_batch_status( - session=session, provider=provider, batch_job=batch_job + # Build batch provider from batch_job.provider (OpenAI or Gemini) + provider = _get_batch_provider( + session=session, + provider_name=batch_job.provider, + organization_id=eval_run.organization_id, + project_id=eval_run.project_id, + ) + status_result = await asyncio.to_thread( + poll_batch_status, + session=session, + provider=provider, + batch_job=batch_job, ) # Refresh batch_job to get the updated provider_status session.refresh(batch_job) provider_status = batch_job.provider_status + is_openai = batch_job.provider in ( + LLMProvider.OPENAI, + LLMProvider.OPENAI_NATIVE, + ) # Handle different provider statuses - if provider_status == "completed": - # Check if batch completed but all requests failed - # (output_file_id is absent, error_file_id is present) - if not status_result.get( - "provider_output_file_id", batch_job.provider_output_file_id - ) and status_result.get("error_file_id"): + if ( + provider_status == "completed" + or provider_status == BatchJobState.SUCCEEDED.value + ): + # OpenAI-only: batch completed but all requests failed (error file present, output file absent) + if ( + is_openai + and not status_result.get( + "provider_output_file_id", batch_job.provider_output_file_id + ) + and status_result.get("error_file_id") + ): error_msg = _extract_batch_error_message( provider=provider, error_file_id=status_result["error_file_id"], @@ -730,7 +830,14 @@ async def check_and_process_evaluation( "action": "processed", } - elif provider_status in ["failed", "expired", "cancelled"]: + elif provider_status in ( + "failed", + "expired", + "cancelled", + BatchJobState.FAILED.value, + BatchJobState.CANCELLED.value, + BatchJobState.EXPIRED.value, + ): # Mark evaluation as failed based on provider status error_msg = batch_job.error_message or f"Provider batch {provider_status}" diff --git a/backend/app/crud/model_config.py b/backend/app/crud/model_config.py index 6d535240a..1ab7f923a 100644 --- a/backend/app/crud/model_config.py +++ b/backend/app/crud/model_config.py @@ -1,20 +1,28 @@ from typing import Any, Literal +from sqlalchemy.dialects.postgresql import ARRAY +from sqlalchemy.sql import sqltypes from sqlmodel import Session, select from app.models import ModelConfig +Provider = str +CompletionType = Literal["text", "stt", "tts"] + def list_active_model_configs( session: Session, - provider: Literal["openai", "google"] | None = None, + provider: Provider | None = None, skip: int = 0, limit: int = 100, ) -> tuple[list[ModelConfig], bool]: statement = select(ModelConfig).where(ModelConfig.is_active) if provider: - statement = statement.where(ModelConfig.provider == provider) + try: + statement = statement.where(ModelConfig.provider == provider) + except Exception: + return [], False statement = statement.order_by(ModelConfig.provider, ModelConfig.model_name) statement = statement.offset(skip).limit(limit + 1) @@ -30,19 +38,22 @@ def list_active_model_configs( def list_all_active_model_configs( session: Session, - provider: Literal["openai", "google"] | None = None, + provider: Provider | None = None, ) -> list[ModelConfig]: statement = select(ModelConfig).where(ModelConfig.is_active) if provider: - statement = statement.where(ModelConfig.provider == provider) + try: + statement = statement.where(ModelConfig.provider == provider) + except Exception: + return [] statement = statement.order_by(ModelConfig.provider, ModelConfig.model_name) return list(session.exec(statement).all()) def get_model_config( - session: Session, provider: Literal["openai", "google"], model_name: str + session: Session, provider: Provider, model_name: str ) -> ModelConfig | None: statement = select(ModelConfig).where( ModelConfig.provider == provider, @@ -52,9 +63,57 @@ def get_model_config( return session.exec(statement).first() -def is_reasoning_model( - session: Session, provider: Literal["openai", "google"], model_name: str +def _modality_filter(stmt, completion_type: CompletionType): + """Restrict query to models matching the completion type via modalities.""" + str_array = ARRAY(sqltypes.String) + input_col = ModelConfig.input_modalities + output_col = ModelConfig.output_modalities + + if completion_type == "stt": + return stmt.where( + input_col.cast(str_array).contains(["AUDIO"]), + output_col.cast(str_array).contains(["TEXT"]), + ) + if completion_type == "tts": + return stmt.where( + input_col.cast(str_array).contains(["TEXT"]), + output_col.cast(str_array).contains(["AUDIO"]), + ) + return stmt.where( + input_col.cast(str_array).contains(["TEXT"]), + output_col.cast(str_array).contains(["TEXT"]), + ) + + +def list_supported_models( + session: Session, provider: Provider, completion_type: CompletionType +) -> list[str]: + """Return active model names for a provider+completion type.""" + stmt = select(ModelConfig.model_name).where( + ModelConfig.provider == provider, + ModelConfig.is_active, + ) + stmt = _modality_filter(stmt, completion_type) + return list(session.exec(stmt).all()) + + +def is_model_supported( + session: Session, + provider: Provider, + completion_type: CompletionType, + model_name: str, ) -> bool: + """Check whether (provider, model_name) is active and matches the completion type.""" + stmt = select(ModelConfig.id).where( + ModelConfig.provider == provider, + ModelConfig.model_name == model_name, + ModelConfig.is_active, + ) + stmt = _modality_filter(stmt, completion_type) + return session.exec(stmt).first() is not None + + +def is_reasoning_model(session: Session, provider: Provider, model_name: str) -> bool: """Return True if the model is configured with a reasoning `effort` control. A model is considered reasoning-capable if its `config` JSON contains an @@ -69,7 +128,7 @@ def is_reasoning_model( def estimate_model_cost( session: Session, - provider: Literal["openai", "google"], + provider: Provider, model_name: str, input_tokens: int, output_tokens: int, diff --git a/backend/app/models/model_config.py b/backend/app/models/model_config.py index bc1d14115..c4a109cba 100644 --- a/backend/app/models/model_config.py +++ b/backend/app/models/model_config.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, Literal +from typing import Any import sqlalchemy as sa from sqlalchemy.dialects.postgresql import ARRAY, JSONB @@ -9,12 +9,12 @@ class ModelConfigBase(SQLModel): - provider: Literal["openai", "google"] = Field( + provider: str = Field( default="openai", sa_column=sa.Column( sa.String, nullable=False, - comment="provider name (e.g. openai, google)", + comment="provider name (e.g. openai, google, sarvamai, elevenlabs)", ), ) diff --git a/backend/app/services/evaluations/batch_job.py b/backend/app/services/evaluations/batch_job.py new file mode 100644 index 000000000..ec13045ee --- /dev/null +++ b/backend/app/services/evaluations/batch_job.py @@ -0,0 +1,92 @@ +import logging +from uuid import UUID + +from celery.exceptions import SoftTimeLimitExceeded +from gevent import Timeout +from sqlmodel import Session + +from app.core.db import engine +from app.crud.evaluations import ( + get_evaluation_run_by_id, + resolve_evaluation_config, + start_evaluation_batch, +) +from app.crud.evaluations.core import update_evaluation_run +from app.models.evaluation import EvaluationRunUpdate +from app.utils import get_langfuse_client + +logger = logging.getLogger(__name__) + + +def execute_evaluation_batch_submission( + project_id: int, + job_id: str, + task_id: str, + task_instance, + organization_id: int, + config_id: str, + config_version: int, + **kwargs, +) -> dict: + run_id = int(job_id) + logger.info( + f"[execute_evaluation_batch_submission] Starting | run_id={run_id} | task={task_id}" + ) + with Session(engine) as session: + run = get_evaluation_run_by_id( + session=session, + evaluation_id=run_id, + organization_id=organization_id, + project_id=project_id, + ) + if not run: + return {"success": False, "error": "Run not found"} + try: + config, error = resolve_evaluation_config( + session=session, + config_id=UUID(str(config_id)), + config_version=config_version, + project_id=project_id, + ) + if error: + update_evaluation_run( + session=session, + eval_run=run, + update=EvaluationRunUpdate(status="failed", error_message=error), + ) + return {"success": False, "error": error} + + langfuse = get_langfuse_client( + session=session, org_id=organization_id, project_id=project_id + ) + run = start_evaluation_batch( + langfuse=langfuse, + session=session, + eval_run=run, + params=config.completion.params, + provider=config.completion.provider, + ) + return {"success": True, "batch_job_id": run.batch_job_id} + except (Timeout, SoftTimeLimitExceeded): + logger.warning( + f"[execute_evaluation_batch_submission] Timeout | run_id={run_id}" + ) + update_evaluation_run( + session=session, + eval_run=run, + update=EvaluationRunUpdate( + status="failed", error_message="Task exceeded soft time limit" + ), + ) + raise + except Exception as e: + logger.error( + f"[execute_evaluation_batch_submission] Failed | run_id={run_id} | {e}", + exc_info=True, + ) + update_evaluation_run( + session=session, + eval_run=run, + update=EvaluationRunUpdate(status="failed", error_message=str(e)), + ) + raise diff --git a/backend/app/services/evaluations/evaluation.py b/backend/app/services/evaluations/evaluation.py index 9f78b9e37..1d23b9efc 100644 --- a/backend/app/services/evaluations/evaluation.py +++ b/backend/app/services/evaluations/evaluation.py @@ -3,9 +3,11 @@ import logging from uuid import UUID +from asgi_correlation_id import correlation_id from fastapi import HTTPException from sqlmodel import Session +from app.celery.utils import start_evaluation_batch_submission from app.crud.evaluations import ( create_evaluation_run, fetch_trace_scores_from_langfuse, @@ -13,12 +15,11 @@ get_evaluation_run_by_id, resolve_evaluation_config, save_score, - start_evaluation_batch, ) -from app.models.evaluation import EvaluationRun -from app.models.llm.request import TextLLMParams, STTLLMParams, TTSLLMParams +from app.models.evaluation import EvaluationRun, EvaluationRunUpdate +from app.crud.evaluations.core import update_evaluation_run from app.services.llm.providers import LLMProvider -from app.utils import get_langfuse_client, get_openai_client +from app.utils import get_langfuse_client from app.core.cloud.storage import get_cloud_storage from app.core.storage_utils import load_json_from_object_store @@ -106,29 +107,16 @@ def start_evaluation( status_code=400, detail=f"Failed to resolve config from stored config: {error}", ) - elif config.completion.provider != LLMProvider.OPENAI: + elif config.completion.provider not in (LLMProvider.OPENAI, LLMProvider.GOOGLE): raise HTTPException( status_code=422, - detail="Only 'openai' provider is supported for evaluation configs", + detail="Only 'openai' and 'google' providers are supported for evaluation configs", ) logger.info( "[start_evaluation] Successfully resolved config from config management" ) - # Get API clients - openai_client = get_openai_client( - session=session, - org_id=organization_id, - project_id=project_id, - ) - langfuse = get_langfuse_client( - session=session, - org_id=organization_id, - project_id=project_id, - ) - - # Step 3: Create EvaluationRun record with config references eval_run = create_evaluation_run( session=session, run_name=experiment_name, @@ -140,39 +128,34 @@ def start_evaluation( project_id=project_id, ) - # Step 4: Start the batch evaluation + trace_id = correlation_id.get() or "N/A" try: - # Convert params dict to appropriate model instance based on type - param_models = { - "text": TextLLMParams, - "stt": STTLLMParams, - "tts": TTSLLMParams, - } - model_class = param_models[config.completion.type] - validated_params = model_class.model_validate(config.completion.params) - - eval_run = start_evaluation_batch( - langfuse=langfuse, - openai_client=openai_client, - session=session, - eval_run=eval_run, - config=validated_params, + celery_task_id = start_evaluation_batch_submission( + project_id=project_id, + job_id=str(eval_run.id), + trace_id=trace_id, + organization_id=organization_id, + config_id=str(config_id), + config_version=config_version, ) - logger.info( - f"[start_evaluation] Evaluation started successfully | " - f"batch_job_id={eval_run.batch_job_id} | total_items={eval_run.total_items}" + f"[start_evaluation] Batch submission queued | " + f"run_id={eval_run.id} | celery_task_id={celery_task_id}" ) - return eval_run - except Exception as e: logger.error( - f"[start_evaluation] Failed to start evaluation | run_id={eval_run.id} | {e}", + f"[start_evaluation] Failed to queue batch submission | run_id={eval_run.id} | {e}", exc_info=True, ) - # Error is already handled in start_evaluation_batch - session.refresh(eval_run) + eval_run = update_evaluation_run( + session=session, + eval_run=eval_run, + update=EvaluationRunUpdate( + status="failed", + error_message=f"Failed to queue batch submission: {e}", + ), + ) return eval_run diff --git a/backend/app/tests/api/routes/test_evaluation.py b/backend/app/tests/api/routes/test_evaluation.py index 8b4e6679c..f324111c8 100644 --- a/backend/app/tests/api/routes/test_evaluation.py +++ b/backend/app/tests/api/routes/test_evaluation.py @@ -7,7 +7,7 @@ from fastapi.testclient import TestClient from sqlmodel import Session, select -from app.crud.evaluations.batch import build_evaluation_jsonl +from app.crud.evaluations.batch import build_openai_evaluation_jsonl from app.models import EvaluationDataset, EvaluationRun from app.models.llm.request import TextLLMParams from app.tests.utils.auth import TestAuthContext From ad191d3ebd67c03e98abfbc43daeba69e05409cd Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Thu, 21 May 2026 08:18:15 +0530 Subject: [PATCH 2/8] feat: Add Langfuse flush to prevent thread accumulation in evaluation polling --- backend/app/crud/evaluations/processing.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/backend/app/crud/evaluations/processing.py b/backend/app/crud/evaluations/processing.py index 87671943a..acf2db61d 100644 --- a/backend/app/crud/evaluations/processing.py +++ b/backend/app/crud/evaluations/processing.py @@ -1039,6 +1039,16 @@ async def poll_all_pending_evaluations(session: Session) -> dict[str, Any]: ) total_failed_count += 1 + # Flush background ingestion threads before discarding this client. + # Each Langfuse() instance spawns TaskManager threads; without flush, + # they accumulate across cron runs and pile up concurrent API calls. + try: + langfuse.flush() + except Exception as flush_err: + logger.warning( + f"[poll_all_pending_evaluations] Langfuse flush failed | project_id={project_id} | {flush_err}" + ) + except Exception as e: logger.error( f"[poll_all_pending_evaluations] Failed to process project | project_id={project_id} | {e}", From 783643a351543b12f6f2fc32415172f18605e975 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Mon, 15 Jun 2026 14:16:50 +0530 Subject: [PATCH 3/8] Refactor JSONL building tests to use openai_params dict instead of TextLLMParams class --- .../app/tests/api/routes/test_evaluation.py | 85 +++++++++++-------- 1 file changed, 49 insertions(+), 36 deletions(-) diff --git a/backend/app/tests/api/routes/test_evaluation.py b/backend/app/tests/api/routes/test_evaluation.py index f324111c8..10595af24 100644 --- a/backend/app/tests/api/routes/test_evaluation.py +++ b/backend/app/tests/api/routes/test_evaluation.py @@ -9,7 +9,6 @@ from app.crud.evaluations.batch import build_openai_evaluation_jsonl from app.models import EvaluationDataset, EvaluationRun -from app.models.llm.request import TextLLMParams from app.tests.utils.auth import TestAuthContext from app.tests.utils.test_data import create_test_config, create_test_evaluation_dataset @@ -597,10 +596,18 @@ def test_start_batch_evaluation_without_authentication( class TestBatchEvaluationJSONLBuilding: - """Test JSONL building logic for batch evaluation.""" + """Test JSONL building logic for batch evaluation. + + ``build_openai_evaluation_jsonl`` wraps an already-mapped ``openai_params`` + dict (the output of ``map_kaapi_to_openai_params``) into Responses API batch + lines: ``body`` is the params dict plus the per-item ``input``. Param shaping + (temperature inclusion, knowledge_base_ids -> tools) is the mapper's job and + is covered in tests/services/llm/test_mappers.py, so these tests pass plain + pre-mapped dicts and assert faithful pass-through. + """ def test_build_batch_jsonl_basic(self) -> None: - """Test basic JSONL building with minimal config.""" + """Test basic JSONL building with a fully populated params dict.""" dataset_items = [ { "id": "item1", @@ -610,13 +617,13 @@ def test_build_batch_jsonl_basic(self) -> None: } ] - config = TextLLMParams( - model="gpt-4o", - temperature=0.2, - instructions="You are a helpful assistant", - ) + openai_params = { + "model": "gpt-4o", + "temperature": 0.2, + "instructions": "You are a helpful assistant", + } - jsonl_data = build_evaluation_jsonl(dataset_items, config) + jsonl_data = build_openai_evaluation_jsonl(dataset_items, openai_params) assert len(jsonl_data) == 1 assert isinstance(jsonl_data[0], dict) @@ -631,7 +638,7 @@ def test_build_batch_jsonl_basic(self) -> None: assert request["body"]["input"] == "What is 2+2?" def test_build_batch_jsonl_with_tools(self) -> None: - """Test JSONL building with tools configuration.""" + """Tools present in the params dict are passed through to the body.""" dataset_items = [ { "id": "item1", @@ -641,13 +648,19 @@ def test_build_batch_jsonl_with_tools(self) -> None: } ] - config = TextLLMParams( - model="gpt-4o-mini", - instructions="Search documents", - knowledge_base_ids=["vs_abc123"], - ) + openai_params = { + "model": "gpt-4o-mini", + "instructions": "Search documents", + "tools": [ + { + "type": "file_search", + "vector_store_ids": ["vs_abc123"], + "max_num_results": 20, + } + ], + } - jsonl_data = build_evaluation_jsonl(dataset_items, config) + jsonl_data = build_openai_evaluation_jsonl(dataset_items, openai_params) assert len(jsonl_data) == 1 request = jsonl_data[0] @@ -655,7 +668,7 @@ def test_build_batch_jsonl_with_tools(self) -> None: assert "vs_abc123" in request["body"]["tools"][0]["vector_store_ids"] def test_build_batch_jsonl_minimal_config(self) -> None: - """Test JSONL building with minimal config (only model required).""" + """Test JSONL building with minimal params (only model).""" dataset_items = [ { "id": "item1", @@ -665,9 +678,9 @@ def test_build_batch_jsonl_minimal_config(self) -> None: } ] - config = TextLLMParams(model="gpt-4o") # Only model provided + openai_params = {"model": "gpt-4o"} - jsonl_data = build_evaluation_jsonl(dataset_items, config) + jsonl_data = build_openai_evaluation_jsonl(dataset_items, openai_params) assert len(jsonl_data) == 1 request = jsonl_data[0] @@ -697,9 +710,9 @@ def test_build_batch_jsonl_skips_empty_questions(self) -> None: }, ] - config = TextLLMParams(model="gpt-4o", instructions="Test") + openai_params = {"model": "gpt-4o", "instructions": "Test"} - jsonl_data = build_evaluation_jsonl(dataset_items, config) + jsonl_data = build_openai_evaluation_jsonl(dataset_items, openai_params) # Should only have 1 valid item assert len(jsonl_data) == 1 @@ -717,12 +730,12 @@ def test_build_batch_jsonl_multiple_items(self) -> None: for i in range(5) ] - config = TextLLMParams( - model="gpt-4o", - instructions="Answer questions", - ) + openai_params = { + "model": "gpt-4o", + "instructions": "Answer questions", + } - jsonl_data = build_evaluation_jsonl(dataset_items, config) + jsonl_data = build_openai_evaluation_jsonl(dataset_items, openai_params) assert len(jsonl_data) == 5 @@ -732,7 +745,7 @@ def test_build_batch_jsonl_multiple_items(self) -> None: assert request_dict["body"]["model"] == "gpt-4o" def test_build_batch_jsonl_temperature_included_when_explicitly_set(self) -> None: - """When temperature is explicitly set, it should appear in the JSONL body.""" + """A temperature present in the params dict appears in the JSONL body.""" dataset_items = [ { "id": "item1", @@ -742,16 +755,16 @@ def test_build_batch_jsonl_temperature_included_when_explicitly_set(self) -> Non } ] - config = TextLLMParams(model="gpt-4o", temperature=0.5) + openai_params = {"model": "gpt-4o", "temperature": 0.5} - jsonl_data = build_evaluation_jsonl(dataset_items, config) + jsonl_data = build_openai_evaluation_jsonl(dataset_items, openai_params) assert len(jsonl_data) == 1 assert "temperature" in jsonl_data[0]["body"] assert jsonl_data[0]["body"]["temperature"] == 0.5 def test_build_batch_jsonl_temperature_excluded_when_not_set(self) -> None: - """When temperature is not explicitly set, it should NOT appear in the JSONL body.""" + """A temperature absent from the params dict stays absent from the body.""" dataset_items = [ { "id": "item1", @@ -761,10 +774,10 @@ def test_build_batch_jsonl_temperature_excluded_when_not_set(self) -> None: } ] - # Only model provided — temperature not in model_fields_set - config = TextLLMParams(model="gpt-4o") + # Mapper omits temperature when it was never set, so it isn't in the dict. + openai_params = {"model": "gpt-4o"} - jsonl_data = build_evaluation_jsonl(dataset_items, config) + jsonl_data = build_openai_evaluation_jsonl(dataset_items, openai_params) assert len(jsonl_data) == 1 assert "temperature" not in jsonl_data[0]["body"] @@ -772,7 +785,7 @@ def test_build_batch_jsonl_temperature_excluded_when_not_set(self) -> None: def test_build_batch_jsonl_temperature_zero_included_when_explicitly_set( self, ) -> None: - """When temperature is explicitly set to 0.0, it should still appear in the body.""" + """A temperature of 0.0 in the params dict is preserved in the body.""" dataset_items = [ { "id": "item1", @@ -782,9 +795,9 @@ def test_build_batch_jsonl_temperature_zero_included_when_explicitly_set( } ] - config = TextLLMParams(model="gpt-4o", temperature=0.0) + openai_params = {"model": "gpt-4o", "temperature": 0.0} - jsonl_data = build_evaluation_jsonl(dataset_items, config) + jsonl_data = build_openai_evaluation_jsonl(dataset_items, openai_params) assert len(jsonl_data) == 1 assert "temperature" in jsonl_data[0]["body"] From 160386757ae7482ea9e2ee5422684d720464f669 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Mon, 15 Jun 2026 14:40:20 +0530 Subject: [PATCH 4/8] Update provider names to include "google-aistudio" and adjust related logic in evaluation batch processing --- backend/app/crud/evaluations/batch.py | 6 +++--- backend/app/crud/evaluations/processing.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/backend/app/crud/evaluations/batch.py b/backend/app/crud/evaluations/batch.py index 604cc146a..0ce20a4c2 100644 --- a/backend/app/crud/evaluations/batch.py +++ b/backend/app/crud/evaluations/batch.py @@ -167,7 +167,7 @@ def start_evaluation_batch( session: Database session eval_run: EvaluationRun database object params: Kaapi-standardized completion params (dict) - provider: Completion provider ("openai" or "google", with optional "-native" suffix) + provider: Completion provider ("openai" or "google-aistudio", with optional "-native" suffix) Returns: Updated EvaluationRun with batch_job_id populated @@ -222,7 +222,7 @@ def start_evaluation_batch( config=batch_config, ) - elif base_provider == LLMProvider.GOOGLE: + elif base_provider == LLMProvider.GOOGLE_AISTUDIO: mapped_params, warnings = map_kaapi_to_google_params( kaapi_params=params, completion_type="text" ) @@ -255,7 +255,7 @@ def start_evaluation_batch( batch_job = start_batch_job( session=session, provider=batch_provider, - provider_name="google", + provider_name="google-aistudio", job_type=BatchJobType.EVALUATION, organization_id=eval_run.organization_id, project_id=eval_run.project_id, diff --git a/backend/app/crud/evaluations/processing.py b/backend/app/crud/evaluations/processing.py index acf2db61d..e4b2db25c 100644 --- a/backend/app/crud/evaluations/processing.py +++ b/backend/app/crud/evaluations/processing.py @@ -67,7 +67,10 @@ def _get_batch_provider( ) return OpenAIBatchProvider(client=openai_client) - if provider_name in (LLMProvider.GOOGLE, LLMProvider.GOOGLE_NATIVE): + if provider_name in ( + LLMProvider.GOOGLE_AISTUDIO, + LLMProvider.GOOGLE_AISTUDIO_NATIVE, + ): gemini_client = GeminiClient.from_credentials( session=session, org_id=organization_id, @@ -206,7 +209,10 @@ def parse_evaluation_output( """ # Create lookup map for dataset items by ID dataset_map = {item["id"]: item for item in dataset_items} - is_google = provider_name in (LLMProvider.GOOGLE, LLMProvider.GOOGLE_NATIVE) + is_google = provider_name in ( + LLMProvider.GOOGLE_AISTUDIO, + LLMProvider.GOOGLE_AISTUDIO_NATIVE, + ) results = [] From 53d54bc114a53e54de385562981296e3afcce9c1 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 17 Jun 2026 09:02:52 +0530 Subject: [PATCH 5/8] Add tests for Google evaluation JSONL building and Langfuse dataset run handling --- .../app/tests/api/routes/test_evaluation.py | 339 +++++++++++++++++- .../tests/crud/evaluations/test_langfuse.py | 35 ++ .../tests/crud/evaluations/test_processing.py | 181 ++++++++++ .../evaluations/test_evaluation_service_s3.py | 119 +++++- 4 files changed, 669 insertions(+), 5 deletions(-) diff --git a/backend/app/tests/api/routes/test_evaluation.py b/backend/app/tests/api/routes/test_evaluation.py index 10595af24..bffd3c566 100644 --- a/backend/app/tests/api/routes/test_evaluation.py +++ b/backend/app/tests/api/routes/test_evaluation.py @@ -1,14 +1,27 @@ import io from typing import Any -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch from uuid import uuid4 import pytest from fastapi.testclient import TestClient from sqlmodel import Session, select -from app.crud.evaluations.batch import build_openai_evaluation_jsonl -from app.models import EvaluationDataset, EvaluationRun +from app.core.util import now +from app.crud.evaluations.batch import ( + build_google_evaluation_jsonl, + build_openai_evaluation_jsonl, + start_evaluation_batch, +) +from app.crud.evaluations.core import create_evaluation_run +from app.models import ( + BatchJob, + EvaluationDataset, + EvaluationRun, + Organization, + Project, +) +from app.models.batch_job import BatchJobType from app.tests.utils.auth import TestAuthContext from app.tests.utils.test_data import create_test_config, create_test_evaluation_dataset @@ -804,6 +817,326 @@ def test_build_batch_jsonl_temperature_zero_included_when_explicitly_set( assert jsonl_data[0]["body"]["temperature"] == 0.0 +class TestBuildGoogleEvaluationJSONL: + """Test JSONL building for Gemini batch evaluation. + + ``build_google_evaluation_jsonl`` wraps an already-mapped ``google_params`` + dict (output of ``map_kaapi_to_google_params``) into Gemini batch lines: + ``{"key": , "request": {...}}``. The request always carries the + per-item question in ``contents``; ``instructions``/``temperature``/ + ``reasoning`` are folded into ``systemInstruction``/``generationConfig`` only + when present, so these tests pass pre-mapped dicts and assert that shaping. + """ + + def _items(self, *questions: str) -> list[dict[str, Any]]: + return [ + { + "id": f"item{i}", + "input": {"question": q}, + "expected_output": {"answer": "a"}, + "metadata": {}, + } + for i, q in enumerate(questions) + ] + + def test_basic_question_only(self) -> None: + """A bare params dict produces contents but no system/generation config.""" + jsonl_data = build_google_evaluation_jsonl( + self._items("What is 2+2?"), {"model": "gemini-2.5-pro"} + ) + + assert len(jsonl_data) == 1 + line = jsonl_data[0] + assert line["key"] == "item0" + request = line["request"] + assert request["contents"] == [ + {"parts": [{"text": "What is 2+2?"}], "role": "user"} + ] + assert "systemInstruction" not in request + assert "generationConfig" not in request + + def test_instructions_become_system_instruction(self) -> None: + """``instructions`` are mapped into ``systemInstruction``.""" + jsonl_data = build_google_evaluation_jsonl( + self._items("Q"), + {"model": "gemini-2.5-pro", "instructions": "Be concise"}, + ) + + request = jsonl_data[0]["request"] + assert request["systemInstruction"] == {"parts": [{"text": "Be concise"}]} + + def test_temperature_in_generation_config(self) -> None: + """``temperature`` goes into ``generationConfig``.""" + jsonl_data = build_google_evaluation_jsonl( + self._items("Q"), {"model": "gemini-2.5-pro", "temperature": 0.3} + ) + + assert jsonl_data[0]["request"]["generationConfig"]["temperature"] == 0.3 + + def test_temperature_zero_included(self) -> None: + """A temperature of 0.0 is preserved (not treated as absent).""" + jsonl_data = build_google_evaluation_jsonl( + self._items("Q"), {"model": "gemini-2.5-pro", "temperature": 0.0} + ) + + assert jsonl_data[0]["request"]["generationConfig"]["temperature"] == 0.0 + + def test_reasoning_becomes_thinking_config(self) -> None: + """``reasoning`` is mapped into ``thinkingConfig``.""" + jsonl_data = build_google_evaluation_jsonl( + self._items("Q"), {"model": "gemini-2.5-pro", "reasoning": "high"} + ) + + thinking = jsonl_data[0]["request"]["generationConfig"]["thinkingConfig"] + assert thinking == {"includeThoughts": False, "thinkingLevel": "high"} + + def test_full_params(self) -> None: + """All optional fields together populate both config blocks.""" + jsonl_data = build_google_evaluation_jsonl( + self._items("Q"), + { + "model": "gemini-2.5-pro", + "instructions": "sys", + "temperature": 0.7, + "reasoning": "low", + }, + ) + + request = jsonl_data[0]["request"] + assert request["systemInstruction"] == {"parts": [{"text": "sys"}]} + assert request["generationConfig"]["temperature"] == 0.7 + assert request["generationConfig"]["thinkingConfig"]["thinkingLevel"] == "low" + + def test_skips_empty_and_missing_questions(self) -> None: + """Items with empty or missing questions are dropped.""" + dataset_items = [ + {"id": "ok", "input": {"question": "Real"}, "expected_output": {}, + "metadata": {}}, + {"id": "empty", "input": {"question": ""}, "expected_output": {}, + "metadata": {}}, + {"id": "missing", "input": {}, "expected_output": {}, "metadata": {}}, + ] + + jsonl_data = build_google_evaluation_jsonl( + dataset_items, {"model": "gemini-2.5-pro"} + ) + + assert len(jsonl_data) == 1 + assert jsonl_data[0]["key"] == "ok" + + def test_multiple_items(self) -> None: + """Each item yields one line keyed by its id.""" + jsonl_data = build_google_evaluation_jsonl( + self._items("Q0", "Q1", "Q2"), {"model": "gemini-2.5-pro"} + ) + + assert [line["key"] for line in jsonl_data] == ["item0", "item1", "item2"] + assert jsonl_data[1]["request"]["contents"][0]["parts"][0]["text"] == "Q1" + + +_BATCH_DATASET_ITEMS = [ + { + "id": "item1", + "input": {"question": "What is 2+2?"}, + "expected_output": {"answer": "4"}, + "metadata": {}, + } +] + + +class TestStartEvaluationBatch: + """Test start_evaluation_batch, the provider-dispatching orchestrator.""" + + @pytest.fixture + def eval_run(self, db: Session) -> EvaluationRun: + """A persisted EvaluationRun in the default (pending) state.""" + org = db.exec(select(Organization)).first() + project = db.exec( + select(Project).where(Project.organization_id == org.id) + ).first() + dataset = create_test_evaluation_dataset( + db=db, + organization_id=org.id, + project_id=project.id, + name="test_dataset_start_batch", + description="Test dataset", + original_items_count=1, + duplication_factor=1, + ) + config = create_test_config(db, project_id=project.id, use_kaapi_schema=True) + return create_evaluation_run( + session=db, + run_name="start_batch_run", + dataset_name=dataset.name, + dataset_id=dataset.id, + config_id=config.id, + config_version=1, + organization_id=org.id, + project_id=project.id, + ) + + def _make_batch_job( + self, db: Session, eval_run: EvaluationRun, provider: str + ) -> BatchJob: + """Persist a BatchJob to stand in for start_batch_job's return value.""" + batch_job = BatchJob( + provider=provider, + provider_batch_id="batch_xyz", + provider_status="submitted", + job_type=BatchJobType.EVALUATION, + total_items=1, + status="submitted", + organization_id=eval_run.organization_id, + project_id=eval_run.project_id, + inserted_at=now(), + updated_at=now(), + ) + db.add(batch_job) + db.commit() + db.refresh(batch_job) + return batch_job + + @patch("app.crud.evaluations.batch.start_batch_job") + @patch("app.crud.evaluations.batch.OpenAIBatchProvider") + @patch("app.crud.evaluations.batch.get_openai_client") + @patch("app.crud.evaluations.batch.map_kaapi_to_openai_params") + @patch("app.crud.evaluations.batch.fetch_dataset_items") + def test_openai_branch( + self, + mock_fetch, + mock_map, + mock_get_client, + mock_provider_cls, + mock_start_batch, + db: Session, + eval_run: EvaluationRun, + ) -> None: + mock_fetch.return_value = _BATCH_DATASET_ITEMS + mock_map.return_value = ({"model": "gpt-4o"}, []) + batch_job = self._make_batch_job(db, eval_run, "openai") + mock_start_batch.return_value = batch_job + + result = start_evaluation_batch( + langfuse=MagicMock(), + session=db, + eval_run=eval_run, + params={"model": "gpt-4o"}, + provider="openai", + ) + + assert result.status == "processing" + assert result.batch_job_id == batch_job.id + assert result.total_items == 1 + assert mock_start_batch.call_args.kwargs["provider_name"] == "openai" + mock_get_client.assert_called_once() + + @patch("app.crud.evaluations.batch.start_batch_job") + @patch("app.crud.evaluations.batch.OpenAIBatchProvider") + @patch("app.crud.evaluations.batch.get_openai_client") + @patch("app.crud.evaluations.batch.map_kaapi_to_openai_params") + @patch("app.crud.evaluations.batch.fetch_dataset_items") + def test_native_suffix_resolves_to_base_provider( + self, + mock_fetch, + mock_map, + mock_get_client, + mock_provider_cls, + mock_start_batch, + db: Session, + eval_run: EvaluationRun, + ) -> None: + """An 'openai-native' provider runs through the OpenAI branch.""" + mock_fetch.return_value = _BATCH_DATASET_ITEMS + mock_map.return_value = ({"model": "gpt-4o"}, ["warn"]) + mock_start_batch.return_value = self._make_batch_job(db, eval_run, "openai") + + result = start_evaluation_batch( + langfuse=MagicMock(), + session=db, + eval_run=eval_run, + params={"model": "gpt-4o"}, + provider="openai-native", + ) + + assert result.status == "processing" + assert mock_start_batch.call_args.kwargs["provider_name"] == "openai" + + @patch("app.crud.evaluations.batch.start_batch_job") + @patch("app.crud.evaluations.batch.GeminiBatchProvider") + @patch("app.crud.evaluations.batch.GeminiClient") + @patch("app.crud.evaluations.batch.map_kaapi_to_google_params") + @patch("app.crud.evaluations.batch.fetch_dataset_items") + def test_google_branch( + self, + mock_fetch, + mock_map, + mock_gemini_client, + mock_provider_cls, + mock_start_batch, + db: Session, + eval_run: EvaluationRun, + ) -> None: + mock_fetch.return_value = _BATCH_DATASET_ITEMS + mock_map.return_value = ({"model": "gemini-2.5-pro"}, []) + batch_job = self._make_batch_job(db, eval_run, "google-aistudio") + mock_start_batch.return_value = batch_job + + result = start_evaluation_batch( + langfuse=MagicMock(), + session=db, + eval_run=eval_run, + params={"model": "gemini-2.5-pro"}, + provider="google-aistudio", + ) + + assert result.status == "processing" + assert result.batch_job_id == batch_job.id + assert mock_start_batch.call_args.kwargs["provider_name"] == "google-aistudio" + mock_gemini_client.from_credentials.assert_called_once() + + @patch("app.crud.evaluations.batch.fetch_dataset_items") + def test_unsupported_provider_marks_failed( + self, mock_fetch, db: Session, eval_run: EvaluationRun + ) -> None: + mock_fetch.return_value = _BATCH_DATASET_ITEMS + + with pytest.raises(ValueError, match="Unsupported provider"): + start_evaluation_batch( + langfuse=MagicMock(), + session=db, + eval_run=eval_run, + params={"model": "claude"}, + provider="anthropic", + ) + + db.refresh(eval_run) + assert eval_run.status == "failed" + assert eval_run.error_message + + @patch("app.crud.evaluations.batch.map_kaapi_to_openai_params") + @patch("app.crud.evaluations.batch.fetch_dataset_items") + def test_empty_jsonl_marks_failed( + self, mock_fetch, mock_map, db: Session, eval_run: EvaluationRun + ) -> None: + """No questions in the dataset -> empty JSONL -> failure.""" + mock_fetch.return_value = [ + {"id": "x", "input": {}, "expected_output": {}, "metadata": {}} + ] + mock_map.return_value = ({"model": "gpt-4o"}, []) + + with pytest.raises(ValueError, match="did not produce any JSONL"): + start_evaluation_batch( + langfuse=MagicMock(), + session=db, + eval_run=eval_run, + params={"model": "gpt-4o"}, + provider="openai", + ) + + db.refresh(eval_run) + assert eval_run.status == "failed" + + class TestGetEvaluationRunStatus: """Test GET /evaluations/{evaluation_id} endpoint.""" diff --git a/backend/app/tests/crud/evaluations/test_langfuse.py b/backend/app/tests/crud/evaluations/test_langfuse.py index 31980b7fd..6344b8f98 100644 --- a/backend/app/tests/crud/evaluations/test_langfuse.py +++ b/backend/app/tests/crud/evaluations/test_langfuse.py @@ -176,6 +176,41 @@ def test_create_langfuse_dataset_run_handles_trace_error(self) -> None: assert "item_1" in trace_id_mapping assert "item_2" not in trace_id_mapping + def test_create_langfuse_dataset_run_handles_rate_limit(self) -> None: + """A 429 status_code error is handled (logged distinctly) and skipped.""" + mock_langfuse = MagicMock() + mock_dataset = MagicMock() + + rate_limit_error = Exception("Too Many Requests") + rate_limit_error.status_code = 429 + + mock_item = MagicMock() + mock_item.id = "item_1" + mock_item.observe.side_effect = rate_limit_error + + mock_dataset.items = [mock_item] + mock_langfuse.get_dataset.return_value = mock_dataset + + results = [ + { + "item_id": "item_1", + "question": "What is 2+2?", + "generated_output": "4", + "ground_truth": "4", + "response_id": "resp_123", + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ] + + trace_id_mapping = create_langfuse_dataset_run( + langfuse=mock_langfuse, + dataset_name="test_dataset", + run_name="test_run", + results=results, + ) + + assert trace_id_mapping == {} + def test_create_langfuse_dataset_run_empty_results(self) -> None: """Test with empty results list.""" mock_langfuse = MagicMock() diff --git a/backend/app/tests/crud/evaluations/test_processing.py b/backend/app/tests/crud/evaluations/test_processing.py index 527ff86f4..a25fe728f 100644 --- a/backend/app/tests/crud/evaluations/test_processing.py +++ b/backend/app/tests/crud/evaluations/test_processing.py @@ -1,4 +1,5 @@ import json +from typing import Any from unittest.mock import MagicMock, patch import pytest @@ -8,6 +9,8 @@ from app.crud.evaluations.core import create_evaluation_run from app.crud.evaluations.processing import ( _extract_batch_error_message, + _extract_gemini_usage, + _get_batch_provider, check_and_process_evaluation, parse_evaluation_output, poll_all_pending_evaluations, @@ -267,6 +270,184 @@ def test_parse_evaluation_output_multiple_items(self) -> None: assert result["ground_truth"] == f"A{i}" +class TestParseEvaluationOutputGoogle: + """Test parsing Gemini batch output (provider_name=google-aistudio). + + Gemini lines are keyed by ``key`` (not ``custom_id``) and carry a nested + ``response`` dict with ``candidates``/``usageMetadata`` rather than the + OpenAI ``response.body`` shape. + """ + + def _dataset_items(self) -> list[dict[str, Any]]: + return [ + { + "id": "item1", + "input": {"question": "What is 2+2?"}, + "expected_output": {"answer": "4"}, + "metadata": {"question_id": 7}, + } + ] + + def test_parse_gemini_success(self) -> None: + """Text is pulled from candidates and usage from usageMetadata.""" + raw_results = [ + { + "key": "item1", + "response": { + "responseId": "resp_g1", + "candidates": [ + {"content": {"parts": [{"text": "The answer is 4"}]}} + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + "thoughtsTokenCount": 2, + }, + }, + } + ] + + results = parse_evaluation_output( + raw_results, self._dataset_items(), provider_name="google-aistudio" + ) + + assert len(results) == 1 + result = results[0] + assert result["item_id"] == "item1" + assert result["generated_output"] == "The answer is 4" + assert result["response_id"] == "resp_g1" + assert result["question_id"] == 7 + assert result["usage"] == { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "reasoning_tokens": 2, + } + + def test_parse_gemini_error(self) -> None: + """An ``error`` on the line yields an ERROR: generated_output.""" + raw_results = [{"key": "item1", "error": {"message": "quota exceeded"}}] + + results = parse_evaluation_output( + raw_results, self._dataset_items(), provider_name="google-aistudio" + ) + + assert len(results) == 1 + assert results[0]["generated_output"].startswith("ERROR:") + assert "quota exceeded" in results[0]["generated_output"] + + def test_parse_gemini_native_provider(self) -> None: + """The -native provider variant is treated as Google too.""" + raw_results = [ + { + "key": "item1", + "response": { + "candidates": [{"content": {"parts": [{"text": "ok"}]}}], + }, + } + ] + + results = parse_evaluation_output( + raw_results, + self._dataset_items(), + provider_name="google-aistudio-native", + ) + + assert results[0]["generated_output"] == "ok" + assert results[0]["usage"] is None + + +class TestExtractGeminiUsage: + """Test ``_extract_gemini_usage`` mapping of Gemini usage to OpenAI shape.""" + + def test_camel_case(self) -> None: + usage = _extract_gemini_usage( + { + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + "thoughtsTokenCount": 3, + } + } + ) + assert usage == { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "reasoning_tokens": 3, + } + + def test_snake_case(self) -> None: + usage = _extract_gemini_usage( + { + "usage_metadata": { + "prompt_token_count": 1, + "candidates_token_count": 2, + "total_token_count": 3, + "thoughts_token_count": 4, + } + } + ) + assert usage == { + "input_tokens": 1, + "output_tokens": 2, + "total_tokens": 3, + "reasoning_tokens": 4, + } + + def test_missing_metadata_returns_none(self) -> None: + assert _extract_gemini_usage({}) is None + assert _extract_gemini_usage({"usageMetadata": None}) is None + + def test_partial_metadata_defaults_to_zero(self) -> None: + usage = _extract_gemini_usage({"usageMetadata": {"promptTokenCount": 9}}) + assert usage == { + "input_tokens": 9, + "output_tokens": 0, + "total_tokens": 0, + "reasoning_tokens": 0, + } + + +class TestGetBatchProvider: + """Test _get_batch_provider dispatch by provider name.""" + + @patch("app.crud.evaluations.processing.OpenAIBatchProvider") + @patch("app.crud.evaluations.processing.get_openai_client") + def test_openai(self, mock_get_client, mock_provider_cls) -> None: + provider = _get_batch_provider( + session=MagicMock(), + provider_name="openai", + organization_id=1, + project_id=2, + ) + mock_get_client.assert_called_once() + assert provider is mock_provider_cls.return_value + + @patch("app.crud.evaluations.processing.GeminiBatchProvider") + @patch("app.crud.evaluations.processing.GeminiClient") + def test_google(self, mock_gemini_client, mock_provider_cls) -> None: + provider = _get_batch_provider( + session=MagicMock(), + provider_name="google-aistudio", + organization_id=1, + project_id=2, + ) + mock_gemini_client.from_credentials.assert_called_once() + assert provider is mock_provider_cls.return_value + + def test_unsupported_raises(self) -> None: + with pytest.raises(ValueError, match="Unsupported provider"): + _get_batch_provider( + session=MagicMock(), + provider_name="anthropic", + organization_id=1, + project_id=2, + ) + + class TestProcessCompletedEvaluation: """Test processing completed evaluation batch.""" diff --git a/backend/app/tests/services/evaluations/test_evaluation_service_s3.py b/backend/app/tests/services/evaluations/test_evaluation_service_s3.py index 78089fbdc..722318ac2 100644 --- a/backend/app/tests/services/evaluations/test_evaluation_service_s3.py +++ b/backend/app/tests/services/evaluations/test_evaluation_service_s3.py @@ -1,13 +1,18 @@ -"""Tests for get_evaluation_with_scores() S3 retrieval.""" +"""Tests for app.services.evaluations.evaluation service functions.""" from collections.abc import Callable from typing import Optional from unittest.mock import MagicMock, patch +from uuid import uuid4 import pytest +from fastapi import HTTPException from app.models import EvaluationRun -from app.services.evaluations.evaluation import get_evaluation_with_scores +from app.services.evaluations.evaluation import ( + get_evaluation_with_scores, + validate_and_start_batch_evaluation, +) class TestGetEvaluationWithScoresS3: @@ -200,3 +205,113 @@ def test_resync_skipped_when_cache_unreadable( assert error is not None mock_fetch_langfuse.assert_not_called() # did not re-fetch mock_save_score.assert_not_called() # did not overwrite the cache + + +_MODULE = "app.services.evaluations.evaluation" + + +def _dataset() -> MagicMock: + dataset = MagicMock() + dataset.id = 1 + dataset.name = "ds" + dataset.langfuse_dataset_id = "lf_ds_1" + dataset.object_store_url = None + return dataset + + +def _config(provider: str) -> MagicMock: + config = MagicMock() + config.completion.provider = provider + config.completion.type = "text" + config.completion.params = {"model": "m"} + return config + + +class TestValidateAndStartBatchEvaluation: + """Test validate_and_start_batch_evaluation provider gating and queueing.""" + + @patch(f"{_MODULE}.resolve_evaluation_config") + @patch(f"{_MODULE}.get_dataset_by_id") + def test_unsupported_provider_raises_422( + self, mock_get_dataset, mock_resolve + ) -> None: + mock_get_dataset.return_value = _dataset() + mock_resolve.return_value = (_config("anthropic"), None) + + with pytest.raises(HTTPException) as exc: + validate_and_start_batch_evaluation( + session=MagicMock(), + dataset_id=1, + experiment_name="exp", + config_id=uuid4(), + config_version=1, + organization_id=1, + project_id=2, + ) + + assert exc.value.status_code == 422 + assert "not supported" in exc.value.detail + + @patch(f"{_MODULE}.update_evaluation_run") + @patch(f"{_MODULE}.start_evaluation_batch_submission") + @patch(f"{_MODULE}.create_evaluation_run_or_409") + @patch(f"{_MODULE}.resolve_evaluation_config") + @patch(f"{_MODULE}.get_dataset_by_id") + def test_queue_failure_marks_run_failed( + self, + mock_get_dataset, + mock_resolve, + mock_create_run, + mock_submit, + mock_update, + ) -> None: + """A failure enqueueing the celery task flips the run to failed (no raise).""" + mock_get_dataset.return_value = _dataset() + mock_resolve.return_value = (_config("openai"), None) + + eval_run = MagicMock() + eval_run.id = 99 + mock_create_run.return_value = eval_run + mock_submit.side_effect = Exception("broker down") + mock_update.return_value = eval_run + + result = validate_and_start_batch_evaluation( + session=MagicMock(), + dataset_id=1, + experiment_name="exp", + config_id=uuid4(), + config_version=1, + organization_id=1, + project_id=2, + ) + + assert result is eval_run + update_arg = mock_update.call_args.kwargs["update"] + assert update_arg.status == "failed" + assert "Failed to queue batch submission" in update_arg.error_message + + @patch(f"{_MODULE}.start_evaluation_batch_submission") + @patch(f"{_MODULE}.create_evaluation_run_or_409") + @patch(f"{_MODULE}.resolve_evaluation_config") + @patch(f"{_MODULE}.get_dataset_by_id") + def test_success_returns_run( + self, mock_get_dataset, mock_resolve, mock_create_run, mock_submit + ) -> None: + mock_get_dataset.return_value = _dataset() + mock_resolve.return_value = (_config("google-aistudio"), None) + eval_run = MagicMock() + mock_create_run.return_value = eval_run + mock_submit.return_value = "task-1" + + result = validate_and_start_batch_evaluation( + session=MagicMock(), + dataset_id=1, + experiment_name="exp", + config_id=uuid4(), + config_version=1, + organization_id=1, + project_id=2, + ) + + assert result is eval_run + mock_submit.assert_called_once() From 5c2435ba2dd7776d7c08e993058c3645a28f457e Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 17 Jun 2026 09:03:01 +0530 Subject: [PATCH 6/8] Refactor test data formatting in Google evaluation JSONL builder --- backend/app/tests/api/routes/test_evaluation.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/backend/app/tests/api/routes/test_evaluation.py b/backend/app/tests/api/routes/test_evaluation.py index bffd3c566..7653f6182 100644 --- a/backend/app/tests/api/routes/test_evaluation.py +++ b/backend/app/tests/api/routes/test_evaluation.py @@ -910,10 +910,18 @@ def test_full_params(self) -> None: def test_skips_empty_and_missing_questions(self) -> None: """Items with empty or missing questions are dropped.""" dataset_items = [ - {"id": "ok", "input": {"question": "Real"}, "expected_output": {}, - "metadata": {}}, - {"id": "empty", "input": {"question": ""}, "expected_output": {}, - "metadata": {}}, + { + "id": "ok", + "input": {"question": "Real"}, + "expected_output": {}, + "metadata": {}, + }, + { + "id": "empty", + "input": {"question": ""}, + "expected_output": {}, + "metadata": {}, + }, {"id": "missing", "input": {}, "expected_output": {}, "metadata": {}}, ] From 4c6a4cd56c331d41db75e36d3d916ec40426943c Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 1 Jul 2026 07:55:58 +0530 Subject: [PATCH 7/8] enhancement(evaluation): Add instructions field to TTSLLMParams and validation for unsupported instructions; Update evaluation process to enforce text completion type --- backend/app/models/llm/request.py | 7 +++++++ backend/app/services/evaluations/evaluation.py | 16 ++++++++++++++-- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/backend/app/models/llm/request.py b/backend/app/models/llm/request.py index 37b25224b..227b83cf6 100644 --- a/backend/app/models/llm/request.py +++ b/backend/app/models/llm/request.py @@ -101,6 +101,13 @@ class TTSLLMParams(SQLModel): voice: str = DEFAULT_TTS_VOICE language: str | None = None response_format: Literal["mp3", "wav", "ogg"] | None = "wav" + instructions: str | None = Field(default=None, exclude=True) + + @model_validator(mode="after") + def _reject_nonempty_instructions(self) -> Self: + if self.instructions: + raise ValueError("instructions is not supported for TTS completions") + return self class ProxyLLMParams(SQLModel): diff --git a/backend/app/services/evaluations/evaluation.py b/backend/app/services/evaluations/evaluation.py index ecbece716..c01bfbeb9 100644 --- a/backend/app/services/evaluations/evaluation.py +++ b/backend/app/services/evaluations/evaluation.py @@ -27,6 +27,7 @@ from app.crud.evaluations.merge import apply_cosine_breakdown from app.crud.evaluations.score import CategoryMetrics, TraceData from app.models.evaluation import EvaluationRun, EvaluationRunUpdate, RunModeEnum +from app.models.llm.constants import CompletionType from app.services.llm.providers import LLMProvider from app.utils import get_langfuse_client @@ -213,8 +214,10 @@ def validate_and_start_batch_evaluation( Steps: 1. Validate dataset exists and has Langfuse ID 2. Resolve config from stored config management - 3. Create evaluation run record - 4. Start batch processing + 3. Check the config provider is supported for batch evaluation + 4. Check the config type is 'text' + 5. Create evaluation run record + 6. Start batch processing Args: session: Database session @@ -289,6 +292,15 @@ def validate_and_start_batch_evaluation( ), ) + if config.completion.type != CompletionType.TEXT: + raise HTTPException( + status_code=422, + detail=( + f"Config type '{config.completion.type}' is not supported for " + f"evaluation configs. Only '{CompletionType.TEXT.value}' type is supported." + ), + ) + logger.info( "[validate_and_start_batch_evaluation] Successfully resolved config from config management" ) From 1487f444b157c41b4509a852fe23904d61f18dd5 Mon Sep 17 00:00:00 2001 From: Ayush <80516839+Ayush8923@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:19:37 +0530 Subject: [PATCH 8/8] feat(organization-management): Delete & Edit Functionality (#956) Co-authored-by: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> --- backend/app/api/docs/organization/delete.md | 13 +- backend/app/api/docs/organization/list.md | 10 +- backend/app/api/docs/projects/delete.md | 13 +- backend/app/api/docs/projects/list.md | 10 +- backend/app/api/docs/projects/list_by_org.md | 8 +- backend/app/api/docs/user_project/delete.md | 2 +- backend/app/api/routes/auth.py | 18 ++ backend/app/api/routes/login.py | 18 +- backend/app/api/routes/organization.py | 78 +++++++-- backend/app/api/routes/project.py | 119 ++++++++++--- backend/app/crud/organization.py | 122 +++++++++++++- backend/app/crud/project.py | 105 +++++++++++- backend/app/crud/user_project.py | 73 ++++++-- backend/app/models/__init__.py | 2 +- backend/app/models/message.py | 13 +- backend/app/tests/api/routes/test_login.py | 3 + backend/app/tests/api/routes/test_org.py | 37 +++- backend/app/tests/api/routes/test_project.py | 167 ++++++++++++++++++- backend/app/tests/api/routes/test_users.py | 9 +- backend/app/tests/api/test_auth.py | 22 +-- backend/app/tests/api/test_deps.py | 12 +- backend/app/tests/crud/test_org.py | 24 ++- backend/app/tests/crud/test_project.py | 30 +++- backend/app/tests/crud/test_user_project.py | 20 ++- backend/app/tests/utils/test_data.py | 70 +++++--- 25 files changed, 852 insertions(+), 146 deletions(-) diff --git a/backend/app/api/docs/organization/delete.md b/backend/app/api/docs/organization/delete.md index e0841c04a..bcaf2f66f 100644 --- a/backend/app/api/docs/organization/delete.md +++ b/backend/app/api/docs/organization/delete.md @@ -1,3 +1,12 @@ -Delete an organization. +Delete an organization. **Requires superuser access.** -Permanently deletes an organization and all associated data. +Supports two delete modes, selected by an optional request body: + +```json +{ "hard_delete": false } +``` + +- **`hard_delete: false` (default)** — *Soft delete.* The organization is marked **inactive** and all of its projects are deactivated as well. No data is removed, so the organization can be reactivated later. It simply stops appearing in listings and can no longer be used. Omitting the body entirely also performs a soft delete. +- **`hard_delete: true`** — *Permanent delete.* The organization and everything owned by it — projects, collections, documents, credentials, assistants, fine-tunings, conversations, and user-project mappings — are permanently removed. **This cannot be undone.** + +In both modes, user **accounts are never deleted** (a user may belong to other organizations). Any user left without an active project afterwards is marked **inactive** and can no longer log in until they are added to an active project again. diff --git a/backend/app/api/docs/organization/list.md b/backend/app/api/docs/organization/list.md index 0d128ac8e..44fbbda4a 100644 --- a/backend/app/api/docs/organization/list.md +++ b/backend/app/api/docs/organization/list.md @@ -1,3 +1,9 @@ -List all organizations. +List organizations. **Requires superuser access.** -Returns paginated list of all organizations in the system. The response includes a `has_more` field in `metadata` indicating whether additional pages are available. +Returns a paginated list of organizations. The response includes a `has_more` field in `metadata` indicating whether additional pages are available. + +**Query Parameters:** +- `search` (optional): Case-insensitive substring match on the organization **name**. For example, `?search=acme` returns every organization whose name contains "acme". +- `is_active` (optional, default `true`): Filter by active status. Pass `false` to list soft-deleted organizations. +- `skip` (optional, default `0`): Number of records to skip for pagination. +- `limit` (optional, default `100`, max `100`): Maximum number of records to return. diff --git a/backend/app/api/docs/projects/delete.md b/backend/app/api/docs/projects/delete.md index 8afee4da9..f107c8227 100644 --- a/backend/app/api/docs/projects/delete.md +++ b/backend/app/api/docs/projects/delete.md @@ -1,3 +1,12 @@ -Delete a project. +Delete a project. **Requires superuser access.** -Permanently deletes a project and all associated data including documents, collections, and configurations. +Supports two delete modes, selected by an optional request body: + +```json +{ "hard_delete": false } +``` + +- **`hard_delete: false` (default)** — *Soft delete.* The project is marked **inactive**. No data is removed, so the project can be reactivated later. It simply stops appearing in listings and can no longer be used. Omitting the body entirely also performs a soft delete. +- **`hard_delete: true`** — *Permanent delete.* The project and everything owned by it — collections, documents, credentials, assistants, fine-tunings, conversations, and user-project mappings — are permanently removed. **This cannot be undone.** + +In both modes, user **accounts are never deleted** (a user may belong to other projects). Any user left without an active project afterwards is marked **inactive** and can no longer log in until they are added to an active project again. diff --git a/backend/app/api/docs/projects/list.md b/backend/app/api/docs/projects/list.md index 911f7d1dc..c9b18af6b 100644 --- a/backend/app/api/docs/projects/list.md +++ b/backend/app/api/docs/projects/list.md @@ -1,3 +1,9 @@ -List all projects. +List projects across all organizations. **Requires superuser access.** -Returns paginated list of all projects across all organizations. +Returns a paginated list of projects. The response includes a `has_more` field in `metadata` indicating whether additional pages are available. + +**Query Parameters:** +- `search` (optional): Case-insensitive substring match on the project **name**. For example, `?search=onboarding` returns every project whose name contains "onboarding". +- `is_active` (optional, default `true`): Filter by active status. Pass `false` to list soft-deleted projects. +- `skip` (optional, default `0`): Number of records to skip for pagination. +- `limit` (optional, default `100`, max `100`): Maximum number of records to return. diff --git a/backend/app/api/docs/projects/list_by_org.md b/backend/app/api/docs/projects/list_by_org.md index b227312d0..c60d5bf59 100644 --- a/backend/app/api/docs/projects/list_by_org.md +++ b/backend/app/api/docs/projects/list_by_org.md @@ -1,3 +1,7 @@ -List all projects for a given organization. +List projects for a given organization. **Requires superuser access.** -Returns all projects belonging to the specified organization ID. The organization must exist and be active. +Returns the projects belonging to the specified organization ID. The organization must exist and be active. + +**Query Parameters:** +- `search` (optional): Case-insensitive substring match on the project **name**. For example, `?search=onboarding` returns only projects whose name contains "onboarding". +- `is_active` (optional, default `true`): Filter by active status. Pass `false` to list soft-deleted projects (e.g. to selectively reactivate them). diff --git a/backend/app/api/docs/user_project/delete.md b/backend/app/api/docs/user_project/delete.md index 4b765c47f..74a1c4372 100644 --- a/backend/app/api/docs/user_project/delete.md +++ b/backend/app/api/docs/user_project/delete.md @@ -6,4 +6,4 @@ Remove a user from a project. **Requires superuser access.** **Query Parameters:** - `project_id` (required): The ID of the project to remove the user from. -This only removes the user-project mapping — the user account itself is not deleted. You cannot remove yourself from a project. +This only removes the user-project mapping — the user account itself is never deleted. If this was the user's last project, the account is marked **inactive** and can no longer log in until they are added to a project again. You cannot remove yourself from a project. diff --git a/backend/app/api/routes/auth.py b/backend/app/api/routes/auth.py index cc557bfa8..f311c25da 100644 --- a/backend/app/api/routes/auth.py +++ b/backend/app/api/routes/auth.py @@ -114,6 +114,15 @@ def google_auth(session: SessionDep, body: GoogleAuthRequest) -> JSONResponse: available_projects = get_user_accessible_projects(session=session, user_id=user.id) + if not user.is_superuser and not available_projects: + logger.info( + f"[google_auth] User has no accessible projects | user_id: {user.id}" + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You are not assigned to any active project. Please contact your administrator.", + ) + if len(available_projects) == 1: proj = available_projects[0] logger.info( @@ -389,6 +398,15 @@ def verify_magic_link(session: SessionDep, token: str) -> JSONResponse: # Get user's projects to embed in token available_projects = get_user_accessible_projects(session=session, user_id=user.id) + if not user.is_superuser and not available_projects: + logger.info( + f"[verify_magic_link] User has no accessible projects | user_id: {user.id}" + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You are not assigned to any active project. Please contact your administrator.", + ) + organization_id = None project_id = None if len(available_projects) == 1: diff --git a/backend/app/api/routes/login.py b/backend/app/api/routes/login.py index 54c1dbeb1..665300ec3 100644 --- a/backend/app/api/routes/login.py +++ b/backend/app/api/routes/login.py @@ -1,7 +1,7 @@ from datetime import timedelta -from typing import Annotated, Any, Optional +from typing import Annotated, Any -from fastapi import APIRouter, Depends, HTTPException, Form +from fastapi import APIRouter, Depends, Form, HTTPException from fastapi.responses import HTMLResponse from fastapi.security import OAuth2PasswordRequestForm @@ -11,6 +11,7 @@ from app.core.config import settings from app.core.security import get_password_hash from app.crud import authenticate, get_user_by_email +from app.crud.auth import get_user_accessible_projects from app.models import Message, NewPassword, Token, UserPublic from app.utils import ( generate_password_reset_token, @@ -26,9 +27,8 @@ def login_access_token( session: SessionDep, form_data: Annotated[OAuth2PasswordRequestForm, Depends()], - token_expiry_minutes: Optional[int] = Form( - default=settings.ACCESS_TOKEN_EXPIRE_MINUTES, ge=1, le=60 * 24 * 360 - ), + token_expiry_minutes: int + | None = Form(default=settings.ACCESS_TOKEN_EXPIRE_MINUTES, ge=1, le=60 * 24 * 360), ) -> Token: """ OAuth2 compatible token login with customizable expiration time. @@ -42,6 +42,14 @@ def login_access_token( elif not user.is_active: raise HTTPException(status_code=400, detail="Inactive user") + if not user.is_superuser and not get_user_accessible_projects( + session=session, user_id=user.id + ): + raise HTTPException( + status_code=403, + detail="You are not assigned to any active project. Please contact your administrator.", + ) + access_token_expires = timedelta(minutes=token_expiry_minutes) return Token( diff --git a/backend/app/api/routes/organization.py b/backend/app/api/routes/organization.py index 40b2941ee..3bb3a74d5 100644 --- a/backend/app/api/routes/organization.py +++ b/backend/app/api/routes/organization.py @@ -1,19 +1,26 @@ import logging -from typing import List from fastapi import APIRouter, Depends, HTTPException, Query from sqlalchemy import func from sqlmodel import select +from app.api.deps import SessionDep +from app.api.permissions import Permission, require_permission +from app.crud.organization import ( + cascade_deactivate_organization, + create_organization, + get_organization_by_id, + get_organization_by_name, + hard_delete_organization, + soft_delete_organization, +) from app.models import ( + DeleteRequest, Organization, OrganizationCreate, - OrganizationUpdate, OrganizationPublic, + OrganizationUpdate, ) -from app.api.deps import SessionDep -from app.api.permissions import Permission, require_permission -from app.crud.organization import create_organization, get_organization_by_id from app.utils import APIResponse, load_description logger = logging.getLogger(__name__) @@ -24,18 +31,30 @@ @router.get( "", dependencies=[Depends(require_permission(Permission.SUPERUSER))], - response_model=APIResponse[List[OrganizationPublic]], + response_model=APIResponse[list[OrganizationPublic]], description=load_description("organization/list.md"), ) def read_organizations( session: SessionDep, skip: int = Query(0, ge=0), limit: int = Query(100, ge=1, le=100), -) -> APIResponse[List[OrganizationPublic]]: - count_statement = select(func.count()).select_from(Organization) + search: str + | None = Query( + None, description="Case-insensitive substring match on the organization name" + ), + is_active: bool = Query( + True, + description="Filter by active status. Pass false to list soft-deleted organizations.", + ), +) -> APIResponse[list[OrganizationPublic]]: + filters = [Organization.is_active.is_(is_active)] + if search and search.strip(): + filters.append(Organization.name.ilike(f"%{search.strip()}%")) + + count_statement = select(func.count()).select_from(Organization).where(*filters) count = session.exec(count_statement).one() - statement = select(Organization).offset(skip).limit(limit) + statement = select(Organization).where(*filters).offset(skip).limit(limit) organizations = session.exec(statement).all() has_more = (skip + limit) < count @@ -69,7 +88,7 @@ def read_organization( Retrieve an organization by ID. """ org = get_organization_by_id(session=session, org_id=org_id) - if org is None: + if org is None or not org.is_active: raise HTTPException(status_code=404, detail="Organization not found") return APIResponse.success_response(org) @@ -89,11 +108,28 @@ def update_organization( raise HTTPException(status_code=404, detail="Organization not found") org_data = org_in.model_dump(exclude_unset=True) - org = org.model_copy(update=org_data) + new_name = org_data.get("name") + if new_name and new_name != org.name: + existing = get_organization_by_name(session=session, name=new_name) + if existing and existing.id != org.id: + raise HTTPException( + status_code=409, + detail="An organization with this name already exists", + ) + + target_active = org_data.get("is_active") + deactivating = target_active is False and org.is_active + + org.sqlmodel_update(org_data) session.add(org) - session.commit() session.flush() + + if deactivating: + cascade_deactivate_organization(session=session, organization=org) + + session.commit() + session.refresh(org) logger.info( f"[update_organization] Organization Updated Successfully | 'org_id': {org.id}" ) @@ -105,17 +141,25 @@ def update_organization( "/{org_id}", dependencies=[Depends(require_permission(Permission.SUPERUSER))], response_model=APIResponse[None], - include_in_schema=False, description=load_description("organization/delete.md"), ) -def delete_organization(session: SessionDep, org_id: int) -> APIResponse[None]: +def delete_organization_endpoint( + session: SessionDep, + org_id: int, + body: DeleteRequest | None = None, +) -> APIResponse[None]: org = get_organization_by_id(session=session, org_id=org_id) if org is None: raise HTTPException(status_code=404, detail="Organization not found") - session.delete(org) - session.commit() + hard_delete = body.hard_delete if body else False + if hard_delete: + hard_delete_organization(session=session, organization=org) + else: + soft_delete_organization(session=session, organization=org) + logger.info( - f"[delete_organization] Organization Deleted Successfully | 'org_id': {org_id}" + f"[delete_organization_endpoint] Organization deleted | 'org_id': {org_id}, " + f"hard_delete: {hard_delete}" ) return APIResponse.success_response(None) diff --git a/backend/app/api/routes/project.py b/backend/app/api/routes/project.py index 62632fba2..9f01e9c3f 100644 --- a/backend/app/api/routes/project.py +++ b/backend/app/api/routes/project.py @@ -1,19 +1,32 @@ import logging -from typing import List from fastapi import APIRouter, Depends, HTTPException, Query from sqlalchemy import func from sqlmodel import select -from app.models import Project, ProjectCreate, ProjectUpdate, ProjectPublic from app.api.deps import SessionDep from app.api.permissions import Permission, require_permission +from app.crud.organization import get_organization_by_id, validate_organization from app.crud.project import ( create_project, get_project_by_id, + get_project_by_name, get_projects_by_organization, + hard_delete_project, + soft_delete_project, +) +from app.crud.user_project import ( + deactivate_users_without_projects, + get_user_ids_for_project, + reactivate_users_with_access, +) +from app.models import ( + DeleteRequest, + Project, + ProjectCreate, + ProjectPublic, + ProjectUpdate, ) -from app.crud.organization import validate_organization from app.utils import APIResponse, load_description logger = logging.getLogger(__name__) @@ -24,18 +37,30 @@ @router.get( "", dependencies=[Depends(require_permission(Permission.SUPERUSER))], - response_model=APIResponse[List[ProjectPublic]], + response_model=APIResponse[list[ProjectPublic]], description=load_description("projects/list.md"), ) def read_projects( session: SessionDep, skip: int = Query(0, ge=0), limit: int = Query(100, ge=1, le=100), + search: str + | None = Query( + None, description="Case-insensitive substring match on the project name" + ), + is_active: bool = Query( + True, + description="Filter by active status. Pass false to list soft-deleted projects.", + ), ): - count_statement = select(func.count()).select_from(Project) + filters = [Project.is_active.is_(is_active)] + if search and search.strip(): + filters.append(Project.name.ilike(f"%{search.strip()}%")) + + count_statement = select(func.count()).select_from(Project).where(*filters) count = session.exec(count_statement).one() - statement = select(Project).offset(skip).limit(limit) + statement = select(Project).where(*filters).offset(skip).limit(limit) projects = session.exec(statement).all() has_more = (skip + limit) < count @@ -65,7 +90,7 @@ def read_project(*, session: SessionDep, project_id: int): Retrieve a project by ID. """ project = get_project_by_id(session=session, project_id=project_id) - if project is None: + if project is None or not project.is_active: raise HTTPException(status_code=404, detail="Project not found") return APIResponse.success_response(project) @@ -83,9 +108,47 @@ def update_project(*, session: SessionDep, project_id: int, project_in: ProjectU raise HTTPException(status_code=404, detail="Project not found") project_data = project_in.model_dump(exclude_unset=True) - project.sqlmodel_update(project_data) + new_name = project_data.get("name") + if new_name and new_name != project.name: + existing = get_project_by_name( + session=session, + project_name=new_name, + organization_id=project.organization_id, + ) + if existing and existing.id != project.id: + raise HTTPException( + status_code=409, + detail="A project with this name already exists in this organization", + ) + + target_active = project_data.get("is_active") + activating = target_active is True and not project.is_active + deactivating = target_active is False and project.is_active + + if activating: + org = get_organization_by_id(session=session, org_id=project.organization_id) + if org is None or not org.is_active: + raise HTTPException( + status_code=400, + detail="Cannot activate a project whose organization is inactive. Reactivate the organization first.", + ) + + project.sqlmodel_update(project_data) session.add(project) + session.flush() + + if activating: + reactivate_users_with_access( + session=session, + user_ids=get_user_ids_for_project(session=session, project_id=project.id), + ) + elif deactivating: + deactivate_users_without_projects( + session=session, + user_ids=get_user_ids_for_project(session=session, project_id=project.id), + ) + session.commit() session.refresh(project) logger.info( @@ -94,36 +157,54 @@ def update_project(*, session: SessionDep, project_id: int, project_in: ProjectU return APIResponse.success_response(project) -# Delete a project @router.delete( "/{project_id}", dependencies=[Depends(require_permission(Permission.SUPERUSER))], - include_in_schema=False, + response_model=APIResponse[None], description=load_description("projects/delete.md"), ) -def delete_project(session: SessionDep, project_id: int): +def delete_project_endpoint( + session: SessionDep, + project_id: int, + body: DeleteRequest | None = None, +): project = get_project_by_id(session=session, project_id=project_id) if project is None: raise HTTPException(status_code=404, detail="Project not found") - session.delete(project) - session.commit() + hard_delete = body.hard_delete if body else False + if hard_delete: + hard_delete_project(session=session, project=project) + else: + soft_delete_project(session=session, project=project) + logger.info( - f"[delete_project] Project deleted successfully | project_id={project_id}" + f"[delete_project_endpoint] Project deleted | project_id={project_id}, " + f"hard_delete: {hard_delete}" ) return APIResponse.success_response(None) -# Get projects by organization @router.get( "/organization/{org_id}", dependencies=[Depends(require_permission(Permission.SUPERUSER))], - response_model=APIResponse[List[ProjectPublic]], + response_model=APIResponse[list[ProjectPublic]], description=load_description("projects/list_by_org.md"), ) def read_projects_by_organization( - session: SessionDep, org_id: int -) -> APIResponse[List[ProjectPublic]]: + session: SessionDep, + org_id: int, + is_active: bool = Query( + True, + description="Filter by active status. Pass false to list soft-deleted projects to reactivate.", + ), + search: str + | None = Query( + None, description="Case-insensitive substring match on the project name" + ), +) -> APIResponse[list[ProjectPublic]]: validate_organization(session=session, org_id=org_id) - projects = get_projects_by_organization(session=session, org_id=org_id) + projects = get_projects_by_organization( + session=session, org_id=org_id, is_active=is_active, search=search + ) return APIResponse.success_response(projects) diff --git a/backend/app/crud/organization.py b/backend/app/crud/organization.py index be7970c5a..60ef7691b 100644 --- a/backend/app/crud/organization.py +++ b/backend/app/crud/organization.py @@ -1,11 +1,11 @@ import logging -from typing import Any, Optional -from datetime import datetime, timezone -from sqlmodel import Session, select + from fastapi import HTTPException +from sqlmodel import Session, select -from app.models import Organization, OrganizationCreate from app.core.util import now +from app.crud.user_project import deactivate_users_without_projects +from app.models import Organization, OrganizationCreate, Project, UserProject logger = logging.getLogger(__name__) @@ -13,6 +13,22 @@ def create_organization( *, session: Session, org_create: OrganizationCreate ) -> Organization: + existing = get_organization_by_name(session=session, name=org_create.name) + if existing: + logger.warning( + f"[create_organization] Organization name already exists | " + f"'org_id': {existing.id}, 'name': {existing.name}, 'is_active': {existing.is_active}" + ) + if existing.is_active: + raise HTTPException( + status_code=409, + detail="An organization with this name already exists", + ) + raise HTTPException( + status_code=409, + detail="An organization with this name already exists but is inactive. Reactivate it instead of creating a new one.", + ) + db_org = Organization.model_validate(org_create) db_org.inserted_at = now() db_org.updated_at = now() @@ -26,16 +42,110 @@ def create_organization( # Get organization by ID -def get_organization_by_id(session: Session, org_id: int) -> Optional[Organization]: +def get_organization_by_id(session: Session, org_id: int) -> Organization | None: statement = select(Organization).where(Organization.id == org_id) return session.exec(statement).first() -def get_organization_by_name(*, session: Session, name: str) -> Optional[Organization]: +def get_organization_by_name(*, session: Session, name: str) -> Organization | None: statement = select(Organization).where(Organization.name == name) return session.exec(statement).first() +def cascade_deactivate_organization( + *, session: Session, organization: Organization +) -> tuple[int, list[int]]: + """ + Mark an organization inactive and cascade the deactivation to all of its + projects, then deactivate any user left without an accessible project. + + Does not commit — the caller is responsible for committing. Returns a tuple + of (number of projects deactivated, list of deactivated user IDs). + """ + org_id = organization.id + + affected_user_ids = session.exec( + select(UserProject.user_id).where(UserProject.organization_id == org_id) + ).all() + + organization.is_active = False + session.add(organization) + + projects = session.exec( + select(Project).where( + Project.organization_id == org_id, + Project.is_active.is_(True), + ) + ).all() + for project in projects: + project.is_active = False + session.add(project) + + session.flush() + + deactivated = deactivate_users_without_projects( + session=session, user_ids=affected_user_ids + ) + return len(projects), deactivated + + +def soft_delete_organization( + *, session: Session, organization: Organization +) -> list[int]: + """ + Soft-delete an organization by marking it inactive. All of its projects are + deactivated as well, so the org and everything under it stop appearing in + listings and can no longer be used. No rows are removed. + + User accounts are never deleted. Any user left without an accessible project + afterwards is marked inactive so they can no longer log in. Returns the list + of deactivated user IDs. + """ + org_id = organization.id + project_count, deactivated = cascade_deactivate_organization( + session=session, organization=organization + ) + session.commit() + logger.info( + f"[soft_delete_organization] Organization soft-deleted | 'org_id': {org_id}, " + f"deactivated_projects: {project_count}, deactivated_users: {len(deactivated)}" + ) + return deactivated + + +def hard_delete_organization( + *, session: Session, organization: Organization +) -> list[int]: + """ + Permanently delete an organization and everything owned by it — projects, + collections, documents, credentials, assistants, fine-tunings, conversations, + and user-project mappings all cascade away. This cannot be undone. + + User *accounts* are still never deleted (a user may belong to other orgs): + any user left without an accessible project after the cascade is marked + inactive instead. Returns the list of deactivated user IDs. + """ + org_id = organization.id + + # Capture affected users before the cascade removes their mappings. + affected_user_ids = session.exec( + select(UserProject.user_id).where(UserProject.organization_id == org_id) + ).all() + + session.delete(organization) + session.flush() + + deactivated = deactivate_users_without_projects( + session=session, user_ids=affected_user_ids + ) + session.commit() + logger.info( + f"[hard_delete_organization] Organization permanently deleted | " + f"'org_id': {org_id}, deactivated_users: {len(deactivated)}" + ) + return deactivated + + # Validate if organization exists and is active def validate_organization(session: Session, org_id: int) -> Organization: """ diff --git a/backend/app/crud/project.py b/backend/app/crud/project.py index 2f73982f7..be47e24cf 100644 --- a/backend/app/crud/project.py +++ b/backend/app/crud/project.py @@ -1,10 +1,11 @@ import logging -from typing import List, Optional -from sqlmodel import Session, select + from fastapi import HTTPException +from sqlmodel import Session, select from app.core.util import now -from app.models import Project, ProjectCreate, Organization +from app.crud.user_project import deactivate_users_without_projects +from app.models import Project, ProjectCreate, UserProject logger = logging.getLogger(__name__) @@ -17,9 +18,17 @@ def create_project(*, session: Session, project_create: ProjectCreate) -> Projec ) if project: logger.warning( - f"[create_project] Project already exists | 'project_id': {project.id}, 'name': {project.name}" + f"[create_project] Project name already exists | 'project_id': {project.id}, " + f"'name': {project.name}, 'is_active': {project.is_active}" + ) + if project.is_active: + raise HTTPException( + 409, "A project with this name already exists in this organization" + ) + raise HTTPException( + 409, + "A project with this name already exists in this organization but is inactive. Reactivate it instead of creating a new one.", ) - raise HTTPException(409, "Project already exists") db_project = Project.model_validate(project_create) db_project.inserted_at = now() @@ -33,24 +42,104 @@ def create_project(*, session: Session, project_create: ProjectCreate) -> Projec return db_project -def get_project_by_id(*, session: Session, project_id: int) -> Optional[Project]: +def get_project_by_id(*, session: Session, project_id: int) -> Project | None: return session.get(Project, project_id) def get_project_by_name( *, session: Session, project_name: str, organization_id: int -) -> Optional[Project]: +) -> Project | None: statement = select(Project).where( Project.name == project_name, Project.organization_id == organization_id ) return session.exec(statement).first() -def get_projects_by_organization(*, session: Session, org_id: int) -> List[Project]: +def get_projects_by_organization( + *, + session: Session, + org_id: int, + is_active: bool | None = True, + search: str | None = None, +) -> list[Project]: + """ + Return projects for an organization. + + By default only active projects are returned. Pass ``is_active=False`` to + list soft-deleted projects (e.g. to selectively reactivate them), or + ``is_active=None`` to return every project regardless of status. Pass + ``search`` for a case-insensitive substring match on the project name. + """ statement = select(Project).where(Project.organization_id == org_id) + if is_active is not None: + statement = statement.where(Project.is_active.is_(is_active)) + if search and search.strip(): + statement = statement.where(Project.name.ilike(f"%{search.strip()}%")) return session.exec(statement).all() +def soft_delete_project(*, session: Session, project: Project) -> list[int]: + """ + Soft-delete a project by marking it inactive, so it stops appearing in + listings and can no longer be used. No rows are removed. + + User accounts are never deleted. Any user left without an accessible project + afterwards is marked inactive so they can no longer log in. Returns the list + of deactivated user IDs. + """ + project_id = project.id + + affected_user_ids = session.exec( + select(UserProject.user_id).where(UserProject.project_id == project_id) + ).all() + + # Soft-delete the project. + project.is_active = False + session.add(project) + session.flush() + + deactivated = deactivate_users_without_projects( + session=session, user_ids=affected_user_ids + ) + session.commit() + logger.info( + f"[soft_delete_project] Project soft-deleted | 'project_id': {project_id}, " + f"deactivated_users: {len(deactivated)}" + ) + return deactivated + + +def hard_delete_project(*, session: Session, project: Project) -> list[int]: + """ + Permanently delete a project and everything owned by it — collections, + documents, credentials, assistants, fine-tunings, conversations, and + user-project mappings all cascade away. This cannot be undone. + + User *accounts* are still never deleted (a user may belong to other + projects): any user left without an accessible project after the cascade is + marked inactive instead. Returns the list of deactivated user IDs. + """ + project_id = project.id + + # Capture affected users before the cascade removes their mappings. + affected_user_ids = session.exec( + select(UserProject.user_id).where(UserProject.project_id == project_id) + ).all() + + session.delete(project) + session.flush() + + deactivated = deactivate_users_without_projects( + session=session, user_ids=affected_user_ids + ) + session.commit() + logger.info( + f"[hard_delete_project] Project permanently deleted | " + f"'project_id': {project_id}, deactivated_users: {len(deactivated)}" + ) + return deactivated + + def validate_project(session: Session, project_id: int) -> Project: """ Ensures that an project exists and is active. diff --git a/backend/app/crud/user_project.py b/backend/app/crud/user_project.py index fec57c01d..d9467626f 100644 --- a/backend/app/crud/user_project.py +++ b/backend/app/crud/user_project.py @@ -1,10 +1,11 @@ import logging import secrets -from typing import Sequence +from collections.abc import Iterable, Sequence from sqlmodel import Session, and_, select from app.core.security import get_password_hash +from app.crud.auth import get_user_accessible_projects from app.models import ( User, UserProject, @@ -94,12 +95,69 @@ def add_user_to_project( return user, "added" +def deactivate_users_without_projects( + *, session: Session, user_ids: Iterable[int] +) -> list[int]: + """ + Mark users inactive when they no longer belong to any *active* project. + """ + deactivated: list[int] = [] + for user_id in set(user_ids): + if get_user_accessible_projects(session=session, user_id=user_id): + continue + + user = session.get(User, user_id) + if user and not user.is_superuser and user.is_active: + user.is_active = False + session.add(user) + deactivated.append(user_id) + + if deactivated: + session.flush() + logger.info( + f"[deactivate_users_without_projects] Users deactivated | user_ids: {deactivated}" + ) + return deactivated + + +def reactivate_users_with_access( + *, session: Session, user_ids: Iterable[int] +) -> list[int]: + """ + Re-activate users who regained access to an active project. + """ + reactivated: list[int] = [] + for user_id in set(user_ids): + user = session.get(User, user_id) + if user and not user.is_active and not user.is_superuser: + if get_user_accessible_projects(session=session, user_id=user_id): + user.is_active = True + session.add(user) + reactivated.append(user_id) + + if reactivated: + session.flush() + logger.info( + f"[reactivate_users_with_access] Users reactivated | user_ids: {reactivated}" + ) + return reactivated + + +def get_user_ids_for_project(*, session: Session, project_id: int) -> list[int]: + """Return the IDs of all users mapped to a project.""" + return list( + session.exec( + select(UserProject.user_id).where(UserProject.project_id == project_id) + ).all() + ) + + def remove_user_from_project( *, session: Session, user_id: int, project_id: int ) -> bool: """ Remove a user from a project. If this was their last project, - deactivate the user account. + deactivate the user account (the user is never deleted). Returns True if removed, False if not found. """ @@ -118,16 +176,7 @@ def remove_user_from_project( session.delete(user_project) session.flush() - # Check if user has any remaining projects - remaining = session.exec( - select(UserProject.id).where(UserProject.user_id == user_id).limit(1) - ).first() - - if not remaining: - user = session.get(User, user_id) - if user and not user.is_superuser: - session.delete(user) - session.flush() + deactivate_users_without_projects(session=session, user_ids=[user_id]) return True diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 4f4bc5af0..f604d6175 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -136,7 +136,7 @@ LLMJobImmediatePublic, LLMJobPublic, ) -from .message import Message +from .message import DeleteRequest, Message from .model_config import ( ModelConfig, ModelConfigBase, diff --git a/backend/app/models/message.py b/backend/app/models/message.py index e0aadfeaa..1d7f4f5e1 100644 --- a/backend/app/models/message.py +++ b/backend/app/models/message.py @@ -1,6 +1,17 @@ -from sqlmodel import SQLModel +from sqlmodel import Field, SQLModel # Generic message class Message(SQLModel): message: str + + +class DeleteRequest(SQLModel): + hard_delete: bool = Field( + default=False, + description=( + "When true, permanently delete the record and all of its related " + "data. When false (the default), perform a soft delete by marking " + "the record inactive." + ), + ) diff --git a/backend/app/tests/api/routes/test_login.py b/backend/app/tests/api/routes/test_login.py index 67ce7feb1..95d5ca5e8 100644 --- a/backend/app/tests/api/routes/test_login.py +++ b/backend/app/tests/api/routes/test_login.py @@ -5,6 +5,7 @@ from app.core.security import verify_password from app.crud import create_user from app.models import UserCreate +from app.tests.utils.test_data import add_user_to_test_project from app.tests.utils.user import user_authentication_headers from app.tests.utils.utils import random_email, random_lower_string from app.utils import generate_password_reset_token @@ -67,6 +68,8 @@ def test_reset_password(client: TestClient, db: Session) -> None: is_superuser=False, ) user = create_user(session=db, user_create=user_create) + # Login requires an active project membership. + add_user_to_test_project(db, user) token = generate_password_reset_token(email=email) headers = user_authentication_headers(client=client, email=email, password=password) data = {"new_password": new_password, "token": token} diff --git a/backend/app/tests/api/routes/test_org.py b/backend/app/tests/api/routes/test_org.py index 7daae61be..1ae748b1d 100644 --- a/backend/app/tests/api/routes/test_org.py +++ b/backend/app/tests/api/routes/test_org.py @@ -3,9 +3,9 @@ from sqlmodel import Session from app.core.config import settings -from app.models import Organization -from app.main import app from app.crud.organization import get_organization_by_id +from app.main import app +from app.models import Organization from app.tests.utils.test_data import create_test_organization client = TestClient(app) @@ -74,23 +74,50 @@ def test_update_organization( assert updated_org["is_active"] == update_data["is_active"] -# Test deleting an organization +# Test soft deleting an organization (default) def test_delete_organization( db: Session, test_organization: Organization, superuser_token_headers: dict[str, str], ) -> None: + org_id = test_organization.id response = client.delete( - f"{settings.API_V1_STR}/organizations/{test_organization.id}", + f"{settings.API_V1_STR}/organizations/{org_id}", headers=superuser_token_headers, ) assert response.status_code == 200 + + # Soft-deleted orgs are hidden from the API... response = client.get( - f"{settings.API_V1_STR}/organizations/{test_organization.id}", + f"{settings.API_V1_STR}/organizations/{org_id}", headers=superuser_token_headers, ) assert response.status_code == 404 + db.expire_all() + org = get_organization_by_id(session=db, org_id=org_id) + assert org is not None + assert org.is_active is False + + +def test_hard_delete_organization( + db: Session, + test_organization: Organization, + superuser_token_headers: dict[str, str], +) -> None: + org_id = test_organization.id + response = client.request( + "DELETE", + f"{settings.API_V1_STR}/organizations/{org_id}", + json={"hard_delete": True}, + headers=superuser_token_headers, + ) + assert response.status_code == 200 + + # The row is permanently gone. + db.expire_all() + assert get_organization_by_id(session=db, org_id=org_id) is None + # Test pagination has_more metadata def test_read_organizations_has_more( diff --git a/backend/app/tests/api/routes/test_project.py b/backend/app/tests/api/routes/test_project.py index 487bec70f..4f0f9c8ad 100644 --- a/backend/app/tests/api/routes/test_project.py +++ b/backend/app/tests/api/routes/test_project.py @@ -2,11 +2,12 @@ from fastapi.testclient import TestClient from sqlmodel import Session -from app.main import app from app.core.config import settings -from app.models import Project, ProjectCreate +from app.crud.project import create_project, get_project_by_id +from app.main import app +from app.models import Organization, Project, ProjectCreate from app.tests.utils.test_data import create_test_organization, create_test_project - +from app.tests.utils.utils import random_lower_string client = TestClient(app) @@ -16,6 +17,21 @@ def test_project(db: Session) -> Project: return create_test_project(db) +def _make_project( + db: Session, org: Organization, name: str, is_active: bool +) -> Project: + """Create a project with a given name and active status under an org.""" + return create_project( + session=db, + project_create=ProjectCreate( + name=name, + description="Test project", + is_active=is_active, + organization_id=org.id, + ), + ) + + # Test creating a project def test_create_new_project( db: Session, superuser_token_headers: dict[str, str] @@ -86,6 +102,57 @@ def test_read_projects_has_more( assert response_data["metadata"]["has_more"] is False +# Test the search param filters projects by name (defaults to active only) +def test_read_projects_search( + db: Session, superuser_token_headers: dict[str, str] +) -> None: + org = create_test_organization(db) + prefix = random_lower_string() + _make_project(db, org, f"{prefix}-alpha", is_active=True) + _make_project(db, org, f"{prefix}-beta", is_active=True) + _make_project(db, org, "unrelated-name", is_active=True) + + response = client.get( + f"{settings.API_V1_STR}/projects?search={prefix}", + headers=superuser_token_headers, + ) + assert response.status_code == 200 + names = [p["name"] for p in response.json()["data"]] + assert f"{prefix}-alpha" in names + assert f"{prefix}-beta" in names + assert "unrelated-name" not in names + + +# Test the is_active param toggles active vs inactive projects, together with search +def test_read_projects_active_inactive_filter( + db: Session, superuser_token_headers: dict[str, str] +) -> None: + org = create_test_organization(db) + prefix = random_lower_string() + _make_project(db, org, f"{prefix}-active", is_active=True) + _make_project(db, org, f"{prefix}-inactive", is_active=False) + + # Default (is_active omitted) returns only active matches. + response = client.get( + f"{settings.API_V1_STR}/projects?search={prefix}", + headers=superuser_token_headers, + ) + assert response.status_code == 200 + names = [p["name"] for p in response.json()["data"]] + assert f"{prefix}-active" in names + assert f"{prefix}-inactive" not in names + + # is_active=false returns only inactive matches. + response = client.get( + f"{settings.API_V1_STR}/projects?search={prefix}&is_active=false", + headers=superuser_token_headers, + ) + assert response.status_code == 200 + names = [p["name"] for p in response.json()["data"]] + assert f"{prefix}-inactive" in names + assert f"{prefix}-active" not in names + + # Test updating a project def test_update_project( db: Session, test_project: Project, superuser_token_headers: dict[str, str] @@ -106,21 +173,45 @@ def test_update_project( assert updated_project["is_active"] == update_data["is_active"] -# Test deleting a project +# Test soft deleting a project (default) def test_delete_project( db: Session, test_project: Project, superuser_token_headers: dict[str, str] ) -> None: + project_id = test_project.id response = client.delete( - f"{settings.API_V1_STR}/projects/{test_project.id}", + f"{settings.API_V1_STR}/projects/{project_id}", headers=superuser_token_headers, ) assert response.status_code == 200 + response = client.get( - f"{settings.API_V1_STR}/projects/{test_project.id}", + f"{settings.API_V1_STR}/projects/{project_id}", headers=superuser_token_headers, ) assert response.status_code == 404 + db.expire_all() + project = get_project_by_id(session=db, project_id=project_id) + assert project is not None + assert project.is_active is False + + +def test_hard_delete_project( + db: Session, test_project: Project, superuser_token_headers: dict[str, str] +) -> None: + project_id = test_project.id + response = client.request( + "DELETE", + f"{settings.API_V1_STR}/projects/{project_id}", + json={"hard_delete": True}, + headers=superuser_token_headers, + ) + assert response.status_code == 200 + + # The row is permanently gone. + db.expire_all() + assert get_project_by_id(session=db, project_id=project_id) is None + # Test retrieving projects by organization def test_read_projects_by_organization( @@ -139,6 +230,70 @@ def test_read_projects_by_organization( assert any(p["id"] == project.id for p in response_data["data"]) +# Test the is_active param toggles active/inactive projects for an organization +def test_read_projects_by_organization_active_inactive( + db: Session, superuser_token_headers: dict[str, str] +) -> None: + org = create_test_organization(db) + active = _make_project(db, org, f"{random_lower_string()}-active", is_active=True) + inactive = _make_project( + db, org, f"{random_lower_string()}-inactive", is_active=False + ) + + # Default returns only active projects. + response = client.get( + f"{settings.API_V1_STR}/projects/organization/{org.id}", + headers=superuser_token_headers, + ) + assert response.status_code == 200 + ids = [p["id"] for p in response.json()["data"]] + assert active.id in ids + assert inactive.id not in ids + + # is_active=false returns only inactive projects. + response = client.get( + f"{settings.API_V1_STR}/projects/organization/{org.id}?is_active=false", + headers=superuser_token_headers, + ) + assert response.status_code == 200 + ids = [p["id"] for p in response.json()["data"]] + assert inactive.id in ids + assert active.id not in ids + + +# Test search combined with is_active for the org-scoped project list +def test_read_projects_by_organization_search( + db: Session, superuser_token_headers: dict[str, str] +) -> None: + org = create_test_organization(db) + prefix = random_lower_string() + active = _make_project(db, org, f"{prefix}-active", is_active=True) + inactive = _make_project(db, org, f"{prefix}-inactive", is_active=False) + _make_project(db, org, "other-project", is_active=True) + + # Search active projects by name. + response = client.get( + f"{settings.API_V1_STR}/projects/organization/{org.id}?search={prefix}", + headers=superuser_token_headers, + ) + assert response.status_code == 200 + names = [p["name"] for p in response.json()["data"]] + assert active.name in names + assert "other-project" not in names + assert inactive.name not in names + + # Search inactive projects by name. + response = client.get( + f"{settings.API_V1_STR}/projects/organization/{org.id}" + f"?search={prefix}&is_active=false", + headers=superuser_token_headers, + ) + assert response.status_code == 200 + names = [p["name"] for p in response.json()["data"]] + assert inactive.name in names + assert active.name not in names + + # Test retrieving projects by non-existent organization def test_read_projects_by_organization_not_found( db: Session, superuser_token_headers: dict[str, str] diff --git a/backend/app/tests/api/routes/test_users.py b/backend/app/tests/api/routes/test_users.py index 4b1e2113f..12712af17 100644 --- a/backend/app/tests/api/routes/test_users.py +++ b/backend/app/tests/api/routes/test_users.py @@ -1,5 +1,3 @@ -from unittest.mock import patch - from fastapi.testclient import TestClient from sqlmodel import Session, select @@ -7,7 +5,8 @@ from app.core.config import settings from app.core.security import verify_password from app.models import User, UserCreate -from app.tests.utils.utils import random_email, random_lower_string, get_non_existent_id +from app.tests.utils.test_data import add_user_to_test_project +from app.tests.utils.utils import get_non_existent_id, random_email, random_lower_string def test_get_users_superuser_me( @@ -57,6 +56,8 @@ def test_get_existing_user_current_user(client: TestClient, db: Session) -> None user_in = UserCreate(email=username, password=password) user = crud.create_user(session=db, user_create=user_in) user_id = user.id + # Login requires an active project membership. + add_user_to_test_project(db, user) login_data = { "username": username, @@ -372,6 +373,8 @@ def test_delete_user_me(client: TestClient, db: Session) -> None: user_in = UserCreate(email=username, password=password) user = crud.create_user(session=db, user_create=user_in) user_id = user.id + # Login requires an active project membership. + add_user_to_test_project(db, user) login_data = { "username": username, diff --git a/backend/app/tests/api/test_auth.py b/backend/app/tests/api/test_auth.py index b94782a98..04eaf015c 100644 --- a/backend/app/tests/api/test_auth.py +++ b/backend/app/tests/api/test_auth.py @@ -13,6 +13,7 @@ verify_magic_link_token, ) from app.tests.utils.auth import TestAuthContext +from app.tests.utils.test_data import add_user_to_test_project from app.tests.utils.user import create_random_user GOOGLE_AUTH_URL = f"{settings.API_V1_STR}/auth/google" @@ -100,6 +101,8 @@ def test_google_auth_activates_inactive_user( ): """Test that inactive user is activated on first Google login.""" user = create_random_user(db) + # Login requires an active project membership. + add_user_to_test_project(db, user) user.is_active = False db.add(user) db.commit() @@ -118,10 +121,10 @@ def test_google_auth_activates_inactive_user( @patch("app.api.routes.auth.id_token.verify_oauth2_token") @patch("app.api.routes.auth.settings") - def test_google_auth_success_no_projects( + def test_google_auth_no_projects_forbidden( self, mock_settings, mock_verify, db: Session, client: TestClient ): - """Test successful login for user with no projects.""" + """A non-superuser with no active project cannot log in via Google.""" user = create_random_user(db) mock_settings.GOOGLE_CLIENT_ID = "test-client-id" @@ -133,15 +136,8 @@ def test_google_auth_success_no_projects( mock_verify.return_value = _mock_idinfo(user.email) resp = client.post(GOOGLE_AUTH_URL, json={"token": "fake"}) - assert resp.status_code == 200 - - body = resp.json() - assert body["success"] is True - data = body["data"] - assert "access_token" in data - assert data["requires_project_selection"] is False - assert data["available_projects"] == [] - assert "access_token" in resp.cookies + assert resp.status_code == 403 + assert "not assigned to any active project" in resp.json()["error"] @patch("app.api.routes.auth.id_token.verify_oauth2_token") @patch("app.api.routes.auth.settings") @@ -454,6 +450,8 @@ def test_verify_user_not_found(self, client: TestClient): def test_verify_activates_inactive_user(self, db: Session, client: TestClient): """Test magic link verify activates inactive user.""" user = create_random_user(db) + # Login requires an active project membership. + add_user_to_test_project(db, user) user.is_active = False db.add(user) db.commit() @@ -469,6 +467,8 @@ def test_verify_activates_inactive_user(self, db: Session, client: TestClient): def test_verify_success(self, db: Session, client: TestClient): """Test successful magic link verification logs user in.""" user = create_random_user(db) + # Login requires an active project membership. + add_user_to_test_project(db, user) token = generate_magic_link_token(email=user.email) resp = client.get(f"{MAGIC_LINK_VERIFY_URL}?token={token}") diff --git a/backend/app/tests/api/test_deps.py b/backend/app/tests/api/test_deps.py index ddb09cb2e..881df45e8 100644 --- a/backend/app/tests/api/test_deps.py +++ b/backend/app/tests/api/test_deps.py @@ -1,20 +1,20 @@ from unittest.mock import MagicMock import pytest -from sqlmodel import Session from fastapi import HTTPException from fastapi.testclient import TestClient +from sqlmodel import Session from app.api.deps import get_auth_context +from app.core.config import settings +from app.core.security import create_access_token, create_refresh_token from app.models import ( - User, AuthContext, + User, ) from app.tests.utils.auth import TestAuthContext +from app.tests.utils.test_data import add_user_to_test_project, create_test_api_key from app.tests.utils.user import authentication_token_from_email, create_random_user -from app.core.config import settings -from app.core.security import create_access_token, create_refresh_token -from app.tests.utils.test_data import create_test_api_key def _mock_request(cookies: dict | None = None) -> MagicMock: @@ -127,6 +127,8 @@ def test_get_auth_context_with_inactive_user_via_token( ) -> None: """Test authentication fails when token belongs to inactive user""" user = create_random_user(db) + # Login requires an active project membership. + add_user_to_test_project(db, user) token_headers = authentication_token_from_email( client=client, email=user.email, db=db ) diff --git a/backend/app/tests/crud/test_org.py b/backend/app/tests/crud/test_org.py index db5bb2f2c..503fc76d0 100644 --- a/backend/app/tests/crud/test_org.py +++ b/backend/app/tests/crud/test_org.py @@ -8,8 +8,8 @@ validate_organization, ) from app.models import Organization, OrganizationCreate -from app.tests.utils.utils import random_lower_string, get_non_existent_id from app.tests.utils.test_data import create_test_organization +from app.tests.utils.utils import get_non_existent_id, random_lower_string def test_create_organization(db: Session) -> None: @@ -23,6 +23,28 @@ def test_create_organization(db: Session) -> None: assert organization.is_active is True +def test_create_organization_duplicate_active_name(db: Session) -> None: + """Creating an org whose name is already taken by an active org is rejected.""" + name = random_lower_string() + create_organization(session=db, org_create=OrganizationCreate(name=name)) + + with pytest.raises(HTTPException, match="already exists"): + create_organization(session=db, org_create=OrganizationCreate(name=name)) + + +def test_create_organization_duplicate_inactive_name(db: Session) -> None: + """Recreating a soft-deleted org name is rejected with a reactivate hint.""" + name = random_lower_string() + org = create_organization(session=db, org_create=OrganizationCreate(name=name)) + + org.is_active = False + db.add(org) + db.flush() + + with pytest.raises(HTTPException, match="inactive"): + create_organization(session=db, org_create=OrganizationCreate(name=name)) + + def test_get_organization_by_id(db: Session) -> None: """Test retrieving an organization by ID.""" organization = create_test_organization(db) diff --git a/backend/app/tests/crud/test_project.py b/backend/app/tests/crud/test_project.py index 9698b4662..d96b5fd77 100644 --- a/backend/app/tests/crud/test_project.py +++ b/backend/app/tests/crud/test_project.py @@ -1,8 +1,7 @@ import pytest -from sqlmodel import Session from fastapi import HTTPException +from sqlmodel import Session -from app.models import Project, ProjectCreate from app.crud.project import ( create_project, get_project_by_id, @@ -10,8 +9,9 @@ get_projects_by_organization, validate_project, ) -from app.tests.utils.utils import random_lower_string, get_non_existent_id -from app.tests.utils.test_data import create_test_project, create_test_organization +from app.models import Project, ProjectCreate +from app.tests.utils.test_data import create_test_organization, create_test_project +from app.tests.utils.utils import get_non_existent_id, random_lower_string def test_create_project(db: Session) -> None: @@ -45,8 +45,28 @@ def test_create_project_duplicate_name(db: Session) -> None: is_active=True, organization_id=organization.id, ) + create_project(session=db, project_create=project_data) + with pytest.raises(HTTPException, match="already exists in this organization"): + create_project(session=db, project_create=project_data) + + +def test_create_project_duplicate_inactive_name(db: Session) -> None: + """Recreating a soft-deleted project name is rejected with a reactivate hint.""" + organization = create_test_organization(db) + + project_data = ProjectCreate( + name=random_lower_string(), + description="Test description", + is_active=True, + organization_id=organization.id, + ) project = create_project(session=db, project_create=project_data) - with pytest.raises(HTTPException, match="Project already exists"): + + project.is_active = False + db.add(project) + db.flush() + + with pytest.raises(HTTPException, match="inactive"): create_project(session=db, project_create=project_data) diff --git a/backend/app/tests/crud/test_user_project.py b/backend/app/tests/crud/test_user_project.py index 917686e6b..65b592951 100644 --- a/backend/app/tests/crud/test_user_project.py +++ b/backend/app/tests/crud/test_user_project.py @@ -1,4 +1,3 @@ -import pytest from sqlmodel import Session from app.crud.user_project import ( @@ -162,8 +161,8 @@ def test_remove_nonexistent_mapping_returns_false(self, db: Session): removed = remove_user_from_project(session=db, user_id=99999, project_id=99999) assert removed is False - def test_remove_last_project_deletes_user(self, db: Session): - """Test removing user from their last project deletes the user.""" + def test_remove_last_project_deactivates_user(self, db: Session): + """Test removing user from their last project deactivates (not deletes) the user.""" project = create_test_project(db) email = random_email() @@ -174,16 +173,23 @@ def test_remove_last_project_deletes_user(self, db: Session): project_id=project.id, ) user_id = user.id + # Simulate an activated user before they lose their last project. + user.is_active = True + db.add(user) + db.flush() remove_user_from_project(session=db, user_id=user_id, project_id=project.id) - assert db.get(User, user_id) is None + deactivated_user = db.get(User, user_id) + assert deactivated_user is not None + assert deactivated_user.is_active is False def test_remove_last_project_preserves_superuser(self, db: Session): - """Test superuser is not deleted when removed from last project.""" + """Test superuser is not deactivated when removed from last project.""" project = create_test_project(db) user = create_random_user(db) user.is_superuser = True + user.is_active = True db.add(user) db.flush() @@ -197,7 +203,9 @@ def test_remove_last_project_preserves_superuser(self, db: Session): remove_user_from_project(session=db, user_id=user.id, project_id=project.id) - assert db.get(User, user.id) is not None + preserved_user = db.get(User, user.id) + assert preserved_user is not None + assert preserved_user.is_active is True class TestGetUserProjects: diff --git a/backend/app/tests/utils/test_data.py b/backend/app/tests/utils/test_data.py index b7b4320ae..2538b5380 100644 --- a/backend/app/tests/utils/test_data.py +++ b/backend/app/tests/utils/test_data.py @@ -1,43 +1,45 @@ from sqlmodel import Session +from app.core.providers import Provider +from app.crud import ( + APIKeyCrud, + create_fine_tuning_job, + create_model_evaluation, + create_organization, + create_project, + set_creds_for_org, +) +from app.crud.config import ConfigCrud, ConfigVersionCrud from app.models import ( - Organization, - Project, APIKeyCreateResponse, - Credential, - OrganizationCreate, - ProjectCreate, + Config, ConfigBlob, + ConfigCreate, + ConfigVersion, + ConfigVersionUpdate, + Credential, CredsCreate, - FineTuningJobCreate, + EvaluationDataset, FineTuning, + FineTuningJobCreate, ModelEvaluation, ModelEvaluationBase, ModelEvaluationStatus, - Config, - ConfigCreate, - ConfigVersion, - ConfigVersionUpdate, - EvaluationDataset, + Organization, + OrganizationCreate, + Project, + ProjectCreate, + User, + UserProject, ) from app.models.config.config import ConfigTag -from app.models.llm import KaapiLLMParams, KaapiCompletionConfig, NativeCompletionConfig -from app.crud import ( - create_organization, - create_project, - set_creds_for_org, - create_fine_tuning_job, - create_model_evaluation, - APIKeyCrud, -) -from app.crud.config import ConfigCrud, ConfigVersionCrud -from app.core.providers import Provider +from app.models.llm import KaapiCompletionConfig, NativeCompletionConfig from app.tests.utils.user import create_random_user from app.tests.utils.utils import ( - random_lower_string, generate_random_string, get_document, get_project, + random_lower_string, ) @@ -71,6 +73,25 @@ def create_test_project(db: Session) -> Project: return create_project(session=db, project_create=project_in) +def add_user_to_test_project(db: Session, user: User) -> Project: + """ + Map a user to a fresh active organization + project. + + Login requires at least one active project, so tests that create a bare + user and then log in must give the user a project first. + """ + project = create_test_project(db) + db.add( + UserProject( + user_id=user.id, + organization_id=project.organization_id, + project_id=project.id, + ) + ) + db.commit() + return project + + def test_credential_data(db: Session) -> CredsCreate: """ Returns credential data for a test project in the form of a CredsCreate schema. @@ -324,7 +345,8 @@ def create_test_version( """ if config_blob is None: # Fetch the latest version to maintain type consistency - from sqlmodel import select, and_ + from sqlmodel import and_, select + from app.models import ConfigVersion stmt = (