From 4663ffb1fc2be61911bc78e9a4efefcaadd708a5 Mon Sep 17 00:00:00 2001 From: nishika26 Date: Mon, 11 May 2026 18:53:50 +0530 Subject: [PATCH 1/7] adding llm call content to db --- backend/app/services/llm/jobs.py | 54 ++++++++++++++++++++ backend/app/tests/services/llm/test_jobs.py | 56 +++++++++++++++++++++ 2 files changed, 110 insertions(+) diff --git a/backend/app/services/llm/jobs.py b/backend/app/services/llm/jobs.py index 27fe4e28b..2bec5944e 100644 --- a/backend/app/services/llm/jobs.py +++ b/backend/app/services/llm/jobs.py @@ -451,6 +451,12 @@ def execute_llm_call( else: config_blob = config.blob + original_input_value = ( + query.input.content.value + if isinstance(query.input, TextInput) + else None + ) + if config_blob.prompt_template and isinstance(query.input, TextInput): template = config_blob.prompt_template.template interpolated = template.replace("{{input}}", query.input.content.value) @@ -482,10 +488,58 @@ def execute_llm_call( ), usage=guardrail_usage, ) + if original_input_value is not None: + query.input.content.value = original_input_value + try: + rephrase_call_request = LLMCallRequest( + query=query, + config=config, + request_metadata=request_metadata, + ) + rephrase_resolved_config = ConfigBlob( + completion=config_blob.completion, + prompt_template=config_blob.prompt_template, + input_guardrails=config_blob.input_guardrails, + output_guardrails=config_blob.output_guardrails, + ) + rephrase_llm_call = create_llm_call( + session, + request=rephrase_call_request, + job_id=job_id, + project_id=project_id, + organization_id=organization_id, + resolved_config=rephrase_resolved_config, + original_provider=str(config_blob.completion.provider), + chain_id=chain_id, + ) + llm_call_id = rephrase_llm_call.id + update_llm_call_response( + session, + llm_call_id=llm_call_id, + provider_response_id=str(job_id), + content={ + "type": "text", + "content": { + "format": "text", + "value": guardrail_direct_response, + }, + }, + usage={ + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + ) + except Exception as e: + logger.error( + f"[execute_llm_call] Failed to record rephrase guardrail call: {e} | job_id={job_id}", + exc_info=True, + ) return BlockResult( response=llm_response, usage=guardrail_usage, metadata=request_metadata, + llm_call_id=llm_call_id, ) if input_error: guard_span.set_status( diff --git a/backend/app/tests/services/llm/test_jobs.py b/backend/app/tests/services/llm/test_jobs.py index c15f28fcb..4f134e782 100644 --- a/backend/app/tests/services/llm/test_jobs.py +++ b/backend/app/tests/services/llm/test_jobs.py @@ -1174,6 +1174,62 @@ def test_guardrails_rephrase_needed_allows_job_with_sanitized_input( result["data"]["response"]["output"]["content"]["value"] == "Rephrased text" ) + def test_guardrails_rephrase_saves_original_input_and_safe_text_in_db( + self, db, job_env, job_for_execution + ): + from app.models.llm.request import LlmCall + + env = job_env + + 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": "Please rephrase the query without unsafe content. Input is outside the allowed topic scope.", + "rephrase_needed": True, + }, + } + mock_fetch_configs.return_value = ( + [{"type": "policy", "stage": "input"}], + [], + ) + + request_data = { + "query": {"input": "unsafe user query"}, + "config": { + "blob": { + "completion": { + "provider": "openai-native", + "type": "text", + "params": {"model": "gpt-4o"}, + }, + "input_guardrails": [ + {"validator_config_id": VALIDATOR_CONFIG_ID_1} + ], + "output_guardrails": [], + } + }, + } + self._execute_job(job_for_execution, db, request_data) + + llm_call = db.exec( + select(LlmCall).where(LlmCall.job_id == job_for_execution.id) + ).first() + + assert llm_call is not None + assert llm_call.input == "unsafe user query" + assert llm_call.content == { + "type": "text", + "content": { + "format": "text", + "value": "Please rephrase the query without unsafe content. Input is outside the allowed topic scope.", + }, + } + def test_execute_job_fetches_validator_configs_from_blob_refs( self, db, job_env, job_for_execution ): From d22f61ae962c5a452819629d8c0aaebf889b54f4 Mon Sep 17 00:00:00 2001 From: nishika26 Date: Tue, 12 May 2026 10:51:23 +0530 Subject: [PATCH 2/7] test cases fix --- backend/app/crud/llm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/app/crud/llm.py b/backend/app/crud/llm.py index c7f5b1aee..a2c5b293f 100644 --- a/backend/app/crud/llm.py +++ b/backend/app/crud/llm.py @@ -116,7 +116,7 @@ def create_llm_call( } else: config_dict = { - "config_blob": resolved_config.model_dump(), + "config_blob": resolved_config.model_dump(mode="json"), } # Extract conversation info if present From b84617a8fa7f0c15bd4f5db5895ec4202ae9911e Mon Sep 17 00:00:00 2001 From: nishika26 Date: Tue, 12 May 2026 14:21:56 +0530 Subject: [PATCH 3/7] passing test cases --- backend/app/services/llm/guardrails.py | 75 +++++++++++++++++++++++++- backend/app/services/llm/jobs.py | 58 +++++--------------- 2 files changed, 87 insertions(+), 46 deletions(-) diff --git a/backend/app/services/llm/guardrails.py b/backend/app/services/llm/guardrails.py index 7ba8d72fe..38dd389ff 100644 --- a/backend/app/services/llm/guardrails.py +++ b/backend/app/services/llm/guardrails.py @@ -3,9 +3,17 @@ import logging import httpx +from sqlmodel import Session from app.core.config import settings -from app.models.llm.request import Validator +from app.crud.llm import create_llm_call, update_llm_call_response +from app.models.llm.request import ( + ConfigBlob, + LLMCallConfig, + LLMCallRequest, + QueryParams, + Validator, +) logger = logging.getLogger(__name__) @@ -167,3 +175,68 @@ def _fetch_by_ids(validator_ids: list[UUID]) -> list[dict[str, Any]]: f"project_id={project_id}, endpoint={endpoint}, error={e}" ) return [], [] + + +def save_rephrase_guardrail_call( + *, + session: Session, + query: QueryParams, + config: LLMCallConfig, + request_metadata: dict | None, + config_blob: ConfigBlob, + guardrail_direct_response: str, + job_id: UUID, + project_id: int, + organization_id: int, + chain_id: UUID | None, +) -> UUID | None: + """Persist the LLM call record for a guardrail rephrase response. + + Returns the llm_call_id on success, None if the DB write fails (non-fatal). + """ + try: + rephrase_call_request = LLMCallRequest( + query=query, + config=config, + request_metadata=request_metadata, + ) + rephrase_resolved_config = ConfigBlob( + completion=config_blob.completion, + prompt_template=config_blob.prompt_template, + input_guardrails=config_blob.input_guardrails, + output_guardrails=config_blob.output_guardrails, + ) + rephrase_llm_call = create_llm_call( + session, + request=rephrase_call_request, + job_id=job_id, + project_id=project_id, + organization_id=organization_id, + resolved_config=rephrase_resolved_config, + original_provider=str(config_blob.completion.provider), + chain_id=chain_id, + ) + update_llm_call_response( + session, + llm_call_id=rephrase_llm_call.id, + provider_response_id=str(job_id), + content={ + "type": "text", + "content": { + "format": "text", + "value": guardrail_direct_response, + }, + }, + usage={ + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + ) + return rephrase_llm_call.id + except Exception as e: + logger.error( + f"[save_rephrase_guardrail_call] Failed to record rephrase guardrail call: {e} | job_id={job_id}", + exc_info=True, + ) + return None diff --git a/backend/app/services/llm/jobs.py b/backend/app/services/llm/jobs.py index 2bec5944e..64eb9e368 100644 --- a/backend/app/services/llm/jobs.py +++ b/backend/app/services/llm/jobs.py @@ -46,6 +46,7 @@ from app.services.llm.guardrails import ( list_validators_config, run_guardrails_validation, + save_rephrase_guardrail_call, ) from app.services.llm.mappers import transform_kaapi_config_to_native from app.services.llm.providers.registry import get_llm_provider @@ -490,51 +491,18 @@ def execute_llm_call( ) if original_input_value is not None: query.input.content.value = original_input_value - try: - rephrase_call_request = LLMCallRequest( - query=query, - config=config, - request_metadata=request_metadata, - ) - rephrase_resolved_config = ConfigBlob( - completion=config_blob.completion, - prompt_template=config_blob.prompt_template, - input_guardrails=config_blob.input_guardrails, - output_guardrails=config_blob.output_guardrails, - ) - rephrase_llm_call = create_llm_call( - session, - request=rephrase_call_request, - job_id=job_id, - project_id=project_id, - organization_id=organization_id, - resolved_config=rephrase_resolved_config, - original_provider=str(config_blob.completion.provider), - chain_id=chain_id, - ) - llm_call_id = rephrase_llm_call.id - update_llm_call_response( - session, - llm_call_id=llm_call_id, - provider_response_id=str(job_id), - content={ - "type": "text", - "content": { - "format": "text", - "value": guardrail_direct_response, - }, - }, - usage={ - "input_tokens": 0, - "output_tokens": 0, - "total_tokens": 0, - }, - ) - except Exception as e: - logger.error( - f"[execute_llm_call] Failed to record rephrase guardrail call: {e} | job_id={job_id}", - exc_info=True, - ) + llm_call_id = save_rephrase_guardrail_call( + session=session, + query=query, + config=config, + request_metadata=request_metadata, + config_blob=config_blob, + guardrail_direct_response=guardrail_direct_response, + job_id=job_id, + project_id=project_id, + organization_id=organization_id, + chain_id=chain_id, + ) return BlockResult( response=llm_response, usage=guardrail_usage, From 2b66b6e50ca7c89e9b12f239b05b4fae652994c0 Mon Sep 17 00:00:00 2001 From: nishika26 Date: Tue, 12 May 2026 15:04:34 +0530 Subject: [PATCH 4/7] adding test coverage --- .../app/tests/services/llm/test_guardrails.py | 121 +++++++++++++++++- 1 file changed, 120 insertions(+), 1 deletion(-) diff --git a/backend/app/tests/services/llm/test_guardrails.py b/backend/app/tests/services/llm/test_guardrails.py index 161056980..ed73d8a9e 100644 --- a/backend/app/tests/services/llm/test_guardrails.py +++ b/backend/app/tests/services/llm/test_guardrails.py @@ -1,14 +1,28 @@ import uuid from unittest.mock import MagicMock, patch +from uuid import UUID, uuid4 import httpx +import pytest +from sqlmodel import Session, select from app.core.config import settings -from app.models.llm.request import Validator +from app.crud.jobs import JobCrud +from app.models import JobType +from app.models.llm.request import ( + ConfigBlob, + LLMCallConfig, + LlmCall, + NativeCompletionConfig, + QueryParams, + Validator, +) from app.services.llm.guardrails import ( list_validators_config, run_guardrails_validation, + save_rephrase_guardrail_call, ) +from app.tests.utils.utils import get_project TEST_JOB_ID = uuid.uuid4() @@ -253,3 +267,108 @@ def test_list_validators_config_network_error_fails_open(mock_client_cls) -> Non assert input_guardrails == [] assert output_guardrails == [] + + +_SAFE_TEXT = "Please rephrase: content not allowed." +_CONFIG_BLOB = ConfigBlob( + completion=NativeCompletionConfig( + provider="openai-native", + params={"model": "gpt-4o"}, + type="text", + ), + input_guardrails=[Validator(validator_config_id=uuid4())], +) +_CONFIG = LLMCallConfig(blob=_CONFIG_BLOB) + + +class TestSaveRephraseGuardrailCall: + @pytest.fixture + def job(self, db: Session): + j = JobCrud(session=db).create( + job_type=JobType.LLM_API, trace_id="rephrase-test" + ) + db.commit() + return j + + def _call(self, db: Session, job, **overrides): + project = get_project(db) + kwargs = dict( + session=db, + query=QueryParams(input="original unsafe input"), + config=_CONFIG, + request_metadata=None, + config_blob=_CONFIG_BLOB, + guardrail_direct_response=_SAFE_TEXT, + job_id=job.id, + project_id=project.id, + organization_id=project.organization_id, + chain_id=None, + ) + kwargs.update(overrides) + return save_rephrase_guardrail_call(**kwargs) + + def test_success_returns_uuid(self, db: Session, job) -> None: + result = self._call(db, job) + assert isinstance(result, UUID) + + def test_success_saves_original_input_and_job_id(self, db: Session, job) -> None: + llm_call_id = self._call(db, job) + llm_call = db.exec(select(LlmCall).where(LlmCall.id == llm_call_id)).first() + assert llm_call is not None + assert llm_call.input == "original unsafe input" + assert llm_call.job_id == job.id + + def test_success_saves_safe_text_as_content(self, db: Session, job) -> None: + llm_call_id = self._call(db, job) + llm_call = db.exec(select(LlmCall).where(LlmCall.id == llm_call_id)).first() + assert llm_call.content == { + "type": "text", + "content": {"format": "text", "value": _SAFE_TEXT}, + } + + def test_success_saves_zero_usage(self, db: Session, job) -> None: + llm_call_id = self._call(db, job) + llm_call = db.exec(select(LlmCall).where(LlmCall.id == llm_call_id)).first() + assert llm_call.usage == { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + } + + def test_create_llm_call_error_returns_none(self, db: Session, job) -> None: + with patch( + "app.services.llm.guardrails.create_llm_call", + side_effect=Exception("DB insert failed"), + ): + result = self._call(db, job) + assert result is None + + def test_update_llm_call_response_error_returns_none( + self, db: Session, job + ) -> None: + with patch( + "app.services.llm.guardrails.update_llm_call_response", + side_effect=Exception("DB update failed"), + ): + result = self._call(db, job) + assert result is None + + def test_chain_id_forwarded_to_create_llm_call(self, db: Session, job) -> None: + chain_id = uuid4() + with patch("app.services.llm.guardrails.create_llm_call") as mock_create: + mock_create.return_value = MagicMock(id=uuid4()) + with patch("app.services.llm.guardrails.update_llm_call_response"): + self._call(db, job, chain_id=chain_id) + _, kwargs = mock_create.call_args + assert kwargs["chain_id"] == chain_id + + def test_request_metadata_forwarded_to_llm_call_request( + self, db: Session, job + ) -> None: + metadata = {"request_id": "abc", "user": "test"} + with patch("app.services.llm.guardrails.create_llm_call") as mock_create: + mock_create.return_value = MagicMock(id=uuid4()) + with patch("app.services.llm.guardrails.update_llm_call_response"): + self._call(db, job, request_metadata=metadata) + _, kwargs = mock_create.call_args + assert kwargs["request"].request_metadata == metadata From c80af4b4cba36485ec583ac15cc6763348af1f01 Mon Sep 17 00:00:00 2001 From: nishika26 Date: Thu, 14 May 2026 17:13:23 +0530 Subject: [PATCH 5/7] pr review and test cases --- backend/app/crud/llm.py | 62 +++++++++++++++ backend/app/services/llm/guardrails.py | 76 +------------------ backend/app/services/llm/jobs.py | 2 +- .../app/tests/services/llm/test_guardrails.py | 14 ++-- 4 files changed, 73 insertions(+), 81 deletions(-) diff --git a/backend/app/crud/llm.py b/backend/app/crud/llm.py index a3d8143c4..18a6bbb46 100644 --- a/backend/app/crud/llm.py +++ b/backend/app/crud/llm.py @@ -14,6 +14,8 @@ QueryInput, ImageInput, PDFInput, + LLMCallConfig, + QueryParams, ) logger = logging.getLogger(__name__) @@ -249,6 +251,66 @@ def update_llm_call_input( ) +def save_rephrase_guardrail_call( + *, + session: Session, + query: QueryParams, + config: LLMCallConfig, + request_metadata: dict | None, + config_blob: ConfigBlob, + guardrail_direct_response: str, + job_id: UUID, + project_id: int, + organization_id: int, + chain_id: UUID | None, +) -> UUID | None: + """Persist the LLM call record for a guardrail rephrase response. + + Returns the llm_call_id on success, None if the DB write fails (non-fatal). + """ + try: + rephrase_call_request = LLMCallRequest( + query=query, + config=config, + request_metadata=request_metadata, + ) + rephrase_llm_call = create_llm_call( + session, + request=rephrase_call_request, + job_id=job_id, + project_id=project_id, + organization_id=organization_id, + resolved_config=config_blob, + original_provider=str(config_blob.completion.provider), + chain_id=chain_id, + ) + update_llm_call_response( + session, + llm_call_id=rephrase_llm_call.id, + provider_response_id=str(job_id), + content={ + "type": "text", + "content": { + "format": "text", + "value": guardrail_direct_response, + }, + }, + # No LLM was invoked, so token counts are genuinely zero. + usage={ + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + ) + return rephrase_llm_call.id + except Exception as e: + logger.error( + f"[save_rephrase_guardrail_call] Failed to record rephrase guardrail call: {e} | job_id={job_id}", + exc_info=True, + ) + return None + + def get_llm_call_by_id( session: Session, llm_call_id: UUID, diff --git a/backend/app/services/llm/guardrails.py b/backend/app/services/llm/guardrails.py index 38dd389ff..1ad79bc3b 100644 --- a/backend/app/services/llm/guardrails.py +++ b/backend/app/services/llm/guardrails.py @@ -3,17 +3,12 @@ import logging import httpx -from sqlmodel import Session from app.core.config import settings from app.crud.llm import create_llm_call, update_llm_call_response -from app.models.llm.request import ( - ConfigBlob, - LLMCallConfig, - LLMCallRequest, - QueryParams, - Validator, -) +from app.models.llm import LLMCallRequest +from app.models.llm.request import ConfigBlob, LLMCallConfig, QueryParams, Validator +from sqlmodel import Session logger = logging.getLogger(__name__) @@ -175,68 +170,3 @@ def _fetch_by_ids(validator_ids: list[UUID]) -> list[dict[str, Any]]: f"project_id={project_id}, endpoint={endpoint}, error={e}" ) return [], [] - - -def save_rephrase_guardrail_call( - *, - session: Session, - query: QueryParams, - config: LLMCallConfig, - request_metadata: dict | None, - config_blob: ConfigBlob, - guardrail_direct_response: str, - job_id: UUID, - project_id: int, - organization_id: int, - chain_id: UUID | None, -) -> UUID | None: - """Persist the LLM call record for a guardrail rephrase response. - - Returns the llm_call_id on success, None if the DB write fails (non-fatal). - """ - try: - rephrase_call_request = LLMCallRequest( - query=query, - config=config, - request_metadata=request_metadata, - ) - rephrase_resolved_config = ConfigBlob( - completion=config_blob.completion, - prompt_template=config_blob.prompt_template, - input_guardrails=config_blob.input_guardrails, - output_guardrails=config_blob.output_guardrails, - ) - rephrase_llm_call = create_llm_call( - session, - request=rephrase_call_request, - job_id=job_id, - project_id=project_id, - organization_id=organization_id, - resolved_config=rephrase_resolved_config, - original_provider=str(config_blob.completion.provider), - chain_id=chain_id, - ) - update_llm_call_response( - session, - llm_call_id=rephrase_llm_call.id, - provider_response_id=str(job_id), - content={ - "type": "text", - "content": { - "format": "text", - "value": guardrail_direct_response, - }, - }, - usage={ - "input_tokens": 0, - "output_tokens": 0, - "total_tokens": 0, - }, - ) - return rephrase_llm_call.id - except Exception as e: - logger.error( - f"[save_rephrase_guardrail_call] Failed to record rephrase guardrail call: {e} | job_id={job_id}", - exc_info=True, - ) - return None diff --git a/backend/app/services/llm/jobs.py b/backend/app/services/llm/jobs.py index 0838c6407..f818ea489 100644 --- a/backend/app/services/llm/jobs.py +++ b/backend/app/services/llm/jobs.py @@ -33,6 +33,7 @@ serialize_input, update_llm_call_input, update_llm_call_response, + save_rephrase_guardrail_call, ) from app.crud.llm_chain import create_llm_chain, update_llm_chain_status from app.models import JobStatus, JobType, JobUpdate, LLMCallRequest, LLMChainRequest @@ -61,7 +62,6 @@ from app.services.llm.guardrails import ( list_validators_config, run_guardrails_validation, - save_rephrase_guardrail_call, ) from app.services.llm.mappers import transform_kaapi_config_to_native from app.services.llm.providers.registry import get_llm_provider diff --git a/backend/app/tests/services/llm/test_guardrails.py b/backend/app/tests/services/llm/test_guardrails.py index ed73d8a9e..335e8fa37 100644 --- a/backend/app/tests/services/llm/test_guardrails.py +++ b/backend/app/tests/services/llm/test_guardrails.py @@ -8,6 +8,7 @@ from app.core.config import settings from app.crud.jobs import JobCrud +from app.crud.llm import save_rephrase_guardrail_call from app.models import JobType from app.models.llm.request import ( ConfigBlob, @@ -20,7 +21,6 @@ from app.services.llm.guardrails import ( list_validators_config, run_guardrails_validation, - save_rephrase_guardrail_call, ) from app.tests.utils.utils import get_project @@ -337,7 +337,7 @@ def test_success_saves_zero_usage(self, db: Session, job) -> None: def test_create_llm_call_error_returns_none(self, db: Session, job) -> None: with patch( - "app.services.llm.guardrails.create_llm_call", + "app.crud.llm.create_llm_call", side_effect=Exception("DB insert failed"), ): result = self._call(db, job) @@ -347,7 +347,7 @@ def test_update_llm_call_response_error_returns_none( self, db: Session, job ) -> None: with patch( - "app.services.llm.guardrails.update_llm_call_response", + "app.crud.llm.update_llm_call_response", side_effect=Exception("DB update failed"), ): result = self._call(db, job) @@ -355,9 +355,9 @@ def test_update_llm_call_response_error_returns_none( def test_chain_id_forwarded_to_create_llm_call(self, db: Session, job) -> None: chain_id = uuid4() - with patch("app.services.llm.guardrails.create_llm_call") as mock_create: + with patch("app.crud.llm.create_llm_call") as mock_create: mock_create.return_value = MagicMock(id=uuid4()) - with patch("app.services.llm.guardrails.update_llm_call_response"): + with patch("app.crud.llm.update_llm_call_response"): self._call(db, job, chain_id=chain_id) _, kwargs = mock_create.call_args assert kwargs["chain_id"] == chain_id @@ -366,9 +366,9 @@ def test_request_metadata_forwarded_to_llm_call_request( self, db: Session, job ) -> None: metadata = {"request_id": "abc", "user": "test"} - with patch("app.services.llm.guardrails.create_llm_call") as mock_create: + with patch("app.crud.llm.create_llm_call") as mock_create: mock_create.return_value = MagicMock(id=uuid4()) - with patch("app.services.llm.guardrails.update_llm_call_response"): + with patch("app.crud.llm.update_llm_call_response"): self._call(db, job, request_metadata=metadata) _, kwargs = mock_create.call_args assert kwargs["request"].request_metadata == metadata From 198e74a843870afff8b53deec6c784daf38d50f4 Mon Sep 17 00:00:00 2001 From: nishika26 Date: Thu, 14 May 2026 17:43:24 +0530 Subject: [PATCH 6/7] coderabbit reviews --- backend/app/crud/llm.py | 42 +++++++++++-------- .../app/tests/services/llm/test_guardrails.py | 25 ++++++----- 2 files changed, 39 insertions(+), 28 deletions(-) diff --git a/backend/app/crud/llm.py b/backend/app/crud/llm.py index 18a6bbb46..44e605fb7 100644 --- a/backend/app/crud/llm.py +++ b/backend/app/crud/llm.py @@ -284,24 +284,32 @@ def save_rephrase_guardrail_call( original_provider=str(config_blob.completion.provider), chain_id=chain_id, ) - update_llm_call_response( - session, - llm_call_id=rephrase_llm_call.id, - provider_response_id=str(job_id), - content={ - "type": "text", - "content": { - "format": "text", - "value": guardrail_direct_response, + try: + update_llm_call_response( + session, + llm_call_id=rephrase_llm_call.id, + provider_response_id=None, + content={ + "type": "text", + "content": { + "format": "text", + "value": guardrail_direct_response, + }, }, - }, - # No LLM was invoked, so token counts are genuinely zero. - usage={ - "input_tokens": 0, - "output_tokens": 0, - "total_tokens": 0, - }, - ) + # No LLM was invoked, so token counts are genuinely zero. + usage={ + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + ) + except Exception: + try: + session.delete(rephrase_llm_call) + session.commit() + except Exception: + pass + raise return rephrase_llm_call.id except Exception as e: logger.error( diff --git a/backend/app/tests/services/llm/test_guardrails.py b/backend/app/tests/services/llm/test_guardrails.py index 335e8fa37..990b7364f 100644 --- a/backend/app/tests/services/llm/test_guardrails.py +++ b/backend/app/tests/services/llm/test_guardrails.py @@ -1,4 +1,5 @@ import uuid +from typing import Any from unittest.mock import MagicMock, patch from uuid import UUID, uuid4 @@ -9,7 +10,7 @@ from app.core.config import settings from app.crud.jobs import JobCrud from app.crud.llm import save_rephrase_guardrail_call -from app.models import JobType +from app.models import Job, JobType from app.models.llm.request import ( ConfigBlob, LLMCallConfig, @@ -283,14 +284,14 @@ def test_list_validators_config_network_error_fails_open(mock_client_cls) -> Non class TestSaveRephraseGuardrailCall: @pytest.fixture - def job(self, db: Session): + def job(self, db: Session) -> Job: j = JobCrud(session=db).create( job_type=JobType.LLM_API, trace_id="rephrase-test" ) db.commit() return j - def _call(self, db: Session, job, **overrides): + def _call(self, db: Session, job: Job, **overrides: Any) -> UUID | None: project = get_project(db) kwargs = dict( session=db, @@ -307,18 +308,20 @@ def _call(self, db: Session, job, **overrides): kwargs.update(overrides) return save_rephrase_guardrail_call(**kwargs) - def test_success_returns_uuid(self, db: Session, job) -> None: + def test_success_returns_uuid(self, db: Session, job: Job) -> None: result = self._call(db, job) assert isinstance(result, UUID) - def test_success_saves_original_input_and_job_id(self, db: Session, job) -> None: + def test_success_saves_original_input_and_job_id( + self, db: Session, job: Job + ) -> None: llm_call_id = self._call(db, job) llm_call = db.exec(select(LlmCall).where(LlmCall.id == llm_call_id)).first() assert llm_call is not None assert llm_call.input == "original unsafe input" assert llm_call.job_id == job.id - def test_success_saves_safe_text_as_content(self, db: Session, job) -> None: + def test_success_saves_safe_text_as_content(self, db: Session, job: Job) -> None: llm_call_id = self._call(db, job) llm_call = db.exec(select(LlmCall).where(LlmCall.id == llm_call_id)).first() assert llm_call.content == { @@ -326,7 +329,7 @@ def test_success_saves_safe_text_as_content(self, db: Session, job) -> None: "content": {"format": "text", "value": _SAFE_TEXT}, } - def test_success_saves_zero_usage(self, db: Session, job) -> None: + def test_success_saves_zero_usage(self, db: Session, job: Job) -> None: llm_call_id = self._call(db, job) llm_call = db.exec(select(LlmCall).where(LlmCall.id == llm_call_id)).first() assert llm_call.usage == { @@ -335,7 +338,7 @@ def test_success_saves_zero_usage(self, db: Session, job) -> None: "total_tokens": 0, } - def test_create_llm_call_error_returns_none(self, db: Session, job) -> None: + def test_create_llm_call_error_returns_none(self, db: Session, job: Job) -> None: with patch( "app.crud.llm.create_llm_call", side_effect=Exception("DB insert failed"), @@ -344,7 +347,7 @@ def test_create_llm_call_error_returns_none(self, db: Session, job) -> None: assert result is None def test_update_llm_call_response_error_returns_none( - self, db: Session, job + self, db: Session, job: Job ) -> None: with patch( "app.crud.llm.update_llm_call_response", @@ -353,7 +356,7 @@ def test_update_llm_call_response_error_returns_none( result = self._call(db, job) assert result is None - def test_chain_id_forwarded_to_create_llm_call(self, db: Session, job) -> None: + def test_chain_id_forwarded_to_create_llm_call(self, db: Session, job: Job) -> None: chain_id = uuid4() with patch("app.crud.llm.create_llm_call") as mock_create: mock_create.return_value = MagicMock(id=uuid4()) @@ -363,7 +366,7 @@ def test_chain_id_forwarded_to_create_llm_call(self, db: Session, job) -> None: assert kwargs["chain_id"] == chain_id def test_request_metadata_forwarded_to_llm_call_request( - self, db: Session, job + self, db: Session, job: Job ) -> None: metadata = {"request_id": "abc", "user": "test"} with patch("app.crud.llm.create_llm_call") as mock_create: From f1a997c301263dbb11a21f784f9fde1ea99b0a53 Mon Sep 17 00:00:00 2001 From: nishika26 Date: Thu, 14 May 2026 18:48:35 +0530 Subject: [PATCH 7/7] increasing timeout for bigger validators --- backend/app/services/llm/guardrails.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/backend/app/services/llm/guardrails.py b/backend/app/services/llm/guardrails.py index 1ad79bc3b..916c8bd94 100644 --- a/backend/app/services/llm/guardrails.py +++ b/backend/app/services/llm/guardrails.py @@ -5,10 +5,7 @@ import httpx from app.core.config import settings -from app.crud.llm import create_llm_call, update_llm_call_response -from app.models.llm import LLMCallRequest -from app.models.llm.request import ConfigBlob, LLMCallConfig, QueryParams, Validator -from sqlmodel import Session +from app.models.llm.request import Validator logger = logging.getLogger(__name__) @@ -57,7 +54,7 @@ def run_guardrails_validation( } try: - with httpx.Client(timeout=10.0) as client: + with httpx.Client(timeout=45.0) as client: response = client.post( f"{settings.KAAPI_GUARDRAILS_URL}/", json=payload,