From 27a38cbd0d991caecf430685b68fd3804ed2ba42 Mon Sep 17 00:00:00 2001 From: rkritika1508 Date: Thu, 21 May 2026 15:19:38 +0530 Subject: [PATCH 1/2] added output parameter --- backend/app/services/llm/guardrails.py | 10 ++- backend/app/services/llm/jobs.py | 7 +- .../app/tests/services/llm/test_guardrails.py | 48 ++++++++++ backend/app/tests/services/llm/test_jobs.py | 87 +++++++++++++++++++ 4 files changed, 148 insertions(+), 4 deletions(-) diff --git a/backend/app/services/llm/guardrails.py b/backend/app/services/llm/guardrails.py index 916c8bd94..c1b16cbdc 100644 --- a/backend/app/services/llm/guardrails.py +++ b/backend/app/services/llm/guardrails.py @@ -17,17 +17,20 @@ def run_guardrails_validation( project_id: int | None, organization_id: int | None, suppress_pass_logs: bool = True, + output_text: str | None = None, ) -> dict[str, Any]: """ Call the Kaapi guardrails service to validate and process input text. Args: - input_text: Text to validate and process. + input_text: User query text, maps to payload["input"]. guardrail_config: List of validator configurations to apply. job_id: Unique identifier for the request. project_id: Project identifier expected by guardrails API. organization_id: Organization identifier expected by guardrails API. suppress_pass_logs: Whether to suppress successful validation logs in guardrails service. + output_text: LLM response text, maps to payload["output"]. Required for validators + that evaluate input/output pairs. Returns: JSON response from the guardrails service with validation results. @@ -39,7 +42,7 @@ def run_guardrails_validation( for validator in guardrail_config ] - payload = { + payload: dict[str, Any] = { "request_id": str(job_id), "project_id": project_id, "organization_id": organization_id, @@ -47,6 +50,9 @@ def run_guardrails_validation( "validators": validators, } + if output_text is not None: + payload["output"] = output_text + headers = { "accept": "application/json", "Authorization": f"Bearer {settings.KAAPI_GUARDRAILS_AUTH}", diff --git a/backend/app/services/llm/jobs.py b/backend/app/services/llm/jobs.py index f818ea489..410f95792 100644 --- a/backend/app/services/llm/jobs.py +++ b/backend/app/services/llm/jobs.py @@ -425,6 +425,7 @@ def apply_output_guardrails( job_id: UUID, project_id: int, organization_id: int, + input_text: str | None = None, ) -> tuple[BlockResult, str | None]: """Apply output guardrails from a config_blob. Shared by /llm/call and /llm/chain. @@ -451,14 +452,15 @@ def apply_output_guardrails( if not output_guardrails: return result, None - output_text = result.response.response.output.content.value + llm_output = result.response.response.output.content.value safe = run_guardrails_validation( - output_text, + input_text or "", output_guardrails, job_id, project_id, organization_id, suppress_pass_logs=True, + output_text=llm_output, ) logger.info( @@ -956,6 +958,7 @@ def execute_llm_call( job_id=job_id, project_id=project_id, organization_id=organization_id, + input_text=original_input_value, ) if output_error: out_guard_span.set_status( diff --git a/backend/app/tests/services/llm/test_guardrails.py b/backend/app/tests/services/llm/test_guardrails.py index 990b7364f..22004d179 100644 --- a/backend/app/tests/services/llm/test_guardrails.py +++ b/backend/app/tests/services/llm/test_guardrails.py @@ -136,6 +136,54 @@ def test_run_guardrails_validation_serializes_validator_models(mock_client_cls) assert kwargs["json"]["validators"] == [{"validator_config_id": str(vid)}] +@patch("app.services.llm.guardrails.httpx.Client") +def test_run_guardrails_validation_includes_output_in_payload(mock_client_cls) -> None: + mock_response = MagicMock() + mock_response.raise_for_status.return_value = None + mock_response.json.return_value = {"success": True} + + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_client_cls.return_value.__enter__.return_value = mock_client + + run_guardrails_validation( + TEST_TEXT, + TEST_CONFIG, + TEST_JOB_ID, + TEST_PROJECT_ID, + TEST_ORGANIZATION_ID, + output_text="some llm response", + ) + + _, kwargs = mock_client.post.call_args + assert kwargs["json"]["input"] == TEST_TEXT + assert kwargs["json"]["output"] == "some llm response" + + +@patch("app.services.llm.guardrails.httpx.Client") +def test_run_guardrails_validation_omits_output_from_payload_by_default( + mock_client_cls, +) -> None: + mock_response = MagicMock() + mock_response.raise_for_status.return_value = None + mock_response.json.return_value = {"success": True} + + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_client_cls.return_value.__enter__.return_value = mock_client + + run_guardrails_validation( + TEST_TEXT, + TEST_CONFIG, + TEST_JOB_ID, + TEST_PROJECT_ID, + TEST_ORGANIZATION_ID, + ) + + _, kwargs = mock_client.post.call_args + assert "output" not in kwargs["json"] + + @patch("app.services.llm.guardrails.httpx.Client") def test_run_guardrails_validation_allows_disable_suppress_pass_logs( mock_client_cls, diff --git a/backend/app/tests/services/llm/test_jobs.py b/backend/app/tests/services/llm/test_jobs.py index 4f134e782..fe1b4e3d1 100644 --- a/backend/app/tests/services/llm/test_jobs.py +++ b/backend/app/tests/services/llm/test_jobs.py @@ -995,6 +995,93 @@ def test_guardrails_sanitize_output_after_provider( assert "REDACTED" in result["data"]["response"]["output"]["content"]["value"] + def test_guardrails_output_validation_sends_input_output_pair( + self, db, job_env, job_for_execution + ): + env = job_env + user_query = "What is my Aadhar number?" + llm_output = "Your Aadhar number is 1234-5678-9012" + + env["mock_llm_response"].response.output.content.value = llm_output + env["provider"].execute.return_value = (env["mock_llm_response"], None) + + with ( + patch("app.services.llm.jobs.run_guardrails_validation") as mock_guardrails, + patch("app.services.llm.jobs.list_validators_config") as mock_fetch_configs, + ): + mock_guardrails.return_value = { + "success": True, + "bypassed": False, + "data": {"safe_text": "Your Aadhar number is [REDACTED]", "rephrase_needed": False}, + } + mock_fetch_configs.return_value = ( + [], + [{"type": "pii_remover", "stage": "output"}], + ) + + request_data = { + "query": {"input": user_query}, + "config": { + "blob": { + "completion": { + "provider": "openai-native", + "type": "text", + "params": {"model": "gpt-4o"}, + }, + "input_guardrails": [], + "output_guardrails": [{"validator_config_id": VALIDATOR_CONFIG_ID_2}], + } + }, + } + self._execute_job(job_for_execution, db, request_data) + + mock_guardrails.assert_called_once() + _, kwargs = mock_guardrails.call_args + assert kwargs.get("output_text") == llm_output + assert mock_guardrails.call_args[0][0] == user_query + + def test_guardrails_bypass_does_not_modify_output( + self, db, job_env, job_for_execution + ): + env = job_env + original_llm_output = "Your Aadhar number is 1234-5678-9012" + + env["mock_llm_response"].response.output.content.value = original_llm_output + env["provider"].execute.return_value = (env["mock_llm_response"], None) + + with ( + patch("app.services.llm.jobs.run_guardrails_validation") as mock_guardrails, + patch("app.services.llm.jobs.list_validators_config") as mock_fetch_configs, + ): + mock_guardrails.return_value = { + "success": False, + "bypassed": True, + "data": {"safe_text": original_llm_output, "rephrase_needed": False}, + } + mock_fetch_configs.return_value = ( + [], + [{"type": "pii_remover", "stage": "output"}], + ) + + request_data = { + "query": {"input": "some question"}, + "config": { + "blob": { + "completion": { + "provider": "openai-native", + "type": "text", + "params": {"model": "gpt-4o"}, + }, + "input_guardrails": [], + "output_guardrails": [{"validator_config_id": VALIDATOR_CONFIG_ID_2}], + } + }, + } + result = self._execute_job(job_for_execution, db, request_data) + + assert result["success"] is True + assert result["data"]["response"]["output"]["content"]["value"] == original_llm_output + def test_guardrails_skip_output_validation_for_audio_output( self, db, job_env, job_for_execution ): From 14f3297c17a12fcb6b1682f7f1e438134a3eccb3 Mon Sep 17 00:00:00 2001 From: rkritika1508 Date: Thu, 21 May 2026 15:22:37 +0530 Subject: [PATCH 2/2] pre-commit --- backend/app/tests/services/llm/test_jobs.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/backend/app/tests/services/llm/test_jobs.py b/backend/app/tests/services/llm/test_jobs.py index fe1b4e3d1..f9cb9ab62 100644 --- a/backend/app/tests/services/llm/test_jobs.py +++ b/backend/app/tests/services/llm/test_jobs.py @@ -1012,7 +1012,10 @@ def test_guardrails_output_validation_sends_input_output_pair( mock_guardrails.return_value = { "success": True, "bypassed": False, - "data": {"safe_text": "Your Aadhar number is [REDACTED]", "rephrase_needed": False}, + "data": { + "safe_text": "Your Aadhar number is [REDACTED]", + "rephrase_needed": False, + }, } mock_fetch_configs.return_value = ( [], @@ -1029,7 +1032,9 @@ def test_guardrails_output_validation_sends_input_output_pair( "params": {"model": "gpt-4o"}, }, "input_guardrails": [], - "output_guardrails": [{"validator_config_id": VALIDATOR_CONFIG_ID_2}], + "output_guardrails": [ + {"validator_config_id": VALIDATOR_CONFIG_ID_2} + ], } }, } @@ -1073,14 +1078,19 @@ def test_guardrails_bypass_does_not_modify_output( "params": {"model": "gpt-4o"}, }, "input_guardrails": [], - "output_guardrails": [{"validator_config_id": VALIDATOR_CONFIG_ID_2}], + "output_guardrails": [ + {"validator_config_id": VALIDATOR_CONFIG_ID_2} + ], } }, } result = self._execute_job(job_for_execution, db, request_data) assert result["success"] is True - assert result["data"]["response"]["output"]["content"]["value"] == original_llm_output + assert ( + result["data"]["response"]["output"]["content"]["value"] + == original_llm_output + ) def test_guardrails_skip_output_validation_for_audio_output( self, db, job_env, job_for_execution