diff --git a/AGENTS.md b/AGENTS.md index 5a48049ef45..d7989b4842c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -49,6 +49,14 @@ LiteLLM is a unified interface for 100+ LLMs that: - Test provider-specific functionality thoroughly - Consider adding load tests for performance-critical changes +4. **Code Style**: Follow [Google's Python Style Guide](https://google.github.io/styleguide/pyguide.html) for Python; the project also uses Black, Ruff, and MyPy. + +### Performance and Database + +- **Hot paths:** Do not add duplicate DB queries in proxy/MCP auth or other high-frequency request paths. Reuse already-loaded auth/context (key, team, end-user objects) instead of re-querying. +- **List endpoints:** Avoid N+1 queries when implementing list endpoints (e.g. agents, MCP servers). Use batch loads, `include`/joins, or a single query with needed relations so DB round-trips stay constant. +- **Pre-merge:** For proxy and MCP code, verify that new DB usage in request-handling paths does not introduce duplicate or N+1 queries. + ### MAKING CODE CHANGES FOR THE UI (IGNORE FOR BACKEND) 1. **Tremor is DEPRECATED, do not use Tremor components in new features/changes** @@ -174,6 +182,8 @@ When opening issues or pull requests, follow these templates: 3. **Rate Limits**: Respect provider rate limits in tests 4. **Memory Usage**: Be mindful of memory usage in streaming scenarios 5. **Dependencies**: Keep dependencies minimal and well-justified +6. **Duplicate DB queries in hot paths**: Adding redundant DB calls in MCP auth, proxy auth, or other high-frequency paths causes performance degradation at scale. +7. **N+1 queries in list endpoints**: Looping over a list and performing a separate DB query per item (e.g. loading agents or related entities) — use batch/joined queries instead. ## HELPFUL RESOURCES @@ -187,4 +197,5 @@ When opening issues or pull requests, follow these templates: - Check similar provider implementations - Ensure comprehensive test coverage - Update documentation appropriately -- Consider backward compatibility impact \ No newline at end of file +- Consider backward compatibility impact +- For proxy/MCP code, confirm no new duplicate or N+1 DB queries in request paths \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md index 3cb67908076..7565560cb72 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -91,6 +91,13 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components: - Async/await patterns throughout - Type hints required for all public APIs - **Avoid imports within methods** — place all imports at the top of the file (module-level). Inline imports inside functions/methods make dependencies harder to trace and hurt readability. The only exception is avoiding circular imports where absolutely necessary. +- Follow [Google's Python Style Guide](https://google.github.io/styleguide/pyguide.html) for Python style and structure; the repo uses Black/Ruff/MyPy on top of that. + +### Performance and Database + +- **Hot paths:** Avoid adding duplicate or redundant DB queries in hot request paths (e.g. proxy auth, MCP request path). Each extra query per request multiplies load at scale. Prefer reusing already-fetched auth/context (e.g. key/team/end-user objects) instead of re-querying. +- **List endpoints and N+1:** When implementing list endpoints (e.g. agents, MCP servers), avoid N+1 patterns: do not loop over a list and perform a separate DB query per item. Use batch loads, `include`/joins, or a single query with the needed relations so the number of DB round-trips is constant. +- **Review before merging:** For proxy and MCP code, check that new DB access in request-handling paths does not introduce duplicate or N+1 queries. ### Testing Strategy - Unit tests in `tests/test_litellm/` diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225095818_baseline_diff/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225095818_baseline_diff/migration.sql new file mode 100644 index 00000000000..f999b933d01 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225095818_baseline_diff/migration.sql @@ -0,0 +1,15 @@ +-- DropForeignKey +ALTER TABLE "LiteLLM_AgentsTable" DROP CONSTRAINT "LiteLLM_AgentsTable_object_permission_id_fkey"; + +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" DROP COLUMN "object_permission_id"; + +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "spec_path" TEXT; + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" DROP COLUMN "agent_id"; + +-- DropTable +DROP TABLE "LiteLLM_ToolTable"; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225101119_baseline_diff/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225101119_baseline_diff/migration.sql new file mode 100644 index 00000000000..2f725d83806 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225101119_baseline_diff/migration.sql @@ -0,0 +1,2 @@ +-- This is an empty migration. + diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py index 3af8a435a39..3a373b5677a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -13,8 +13,7 @@ from typing import (TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Literal, from fastapi import HTTPException -from litellm.integrations.custom_guardrail import (CustomGuardrail, - log_guardrail_information) +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.base import \ GuardrailConfigModel @@ -282,7 +281,6 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): detection_info={"language": language}, ) - @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, @@ -339,6 +337,11 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): } if max_confidence is not None: tracing_kw["confidence_score"] = max_confidence + event_type = ( + GuardrailEventHooks.pre_call + if input_type == "request" + else GuardrailEventHooks.post_call + ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider="block_code_execution", guardrail_json_response=guardrail_response, @@ -347,6 +350,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, tracing_detail=GuardrailTracingDetail(**tracing_kw), # type: ignore[typeddict-item] ) @@ -356,28 +360,24 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): response: Any, request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: - """Accumulate streamed content and block if a complete fenced code block is detected.""" + """Accumulate streamed content and block as soon as a complete fenced code block is detected (before yielding that chunk).""" accumulated = "" async for item in response: if isinstance(item, ModelResponseStream) and item.choices: delta_content = "" - is_final = False for choice in item.choices: if hasattr(choice, "delta") and choice.delta: content = getattr(choice.delta, "content", None) if content and isinstance(content, str): delta_content += content - if getattr(choice, "finish_reason", None): - is_final = True accumulated += delta_content - if is_final: - # Run detection on full accumulated text (streaming: block only, no mask) - blocks = self._find_blocks(accumulated) - for _tag, _body, confidence, action_taken in blocks: - if ( - action_taken == "block" - and confidence >= self.confidence_threshold - ): - lang = _tag or "unknown" - self._raise_block_error(lang, True, request_data) + # Check after every chunk so we block before yielding the chunk that completes a blocked block + blocks = self._find_blocks(accumulated) + for _tag, _body, confidence, action_taken in blocks: + if ( + action_taken == "block" + and confidence >= self.confidence_threshold + ): + lang = _tag or "unknown" + self._raise_block_error(lang, True, request_data) yield item diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py index 0a54f12b535..c12943910d6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py @@ -5,13 +5,9 @@ from fastapi import HTTPException from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy.guardrails.guardrail_hooks.block_code_execution import ( - DEFAULT_EVENT_HOOKS, - BlockCodeExecutionGuardrail, - initialize_guardrail, -) -from litellm.proxy.guardrails.guardrail_hooks.block_code_execution.block_code_execution import ( - _normalize_escaped_newlines, -) + DEFAULT_EVENT_HOOKS, BlockCodeExecutionGuardrail, initialize_guardrail) +from litellm.proxy.guardrails.guardrail_hooks.block_code_execution.block_code_execution import \ + _normalize_escaped_newlines from litellm.types.guardrails import GuardrailEventHooks @@ -328,3 +324,46 @@ print(factorial(5)) # Output: 120 assert "python" in str(exc_info.value).lower() or "code" in str( exc_info.value ).lower() + + @pytest.mark.asyncio + async def test_streaming_hook_blocks_before_yielding_chunk_that_completes_block( + self, + ): + """Streaming hook runs block check after every chunk and raises before yielding the chunk that completes a blocked fenced block.""" + from litellm.types.utils import (Delta, ModelResponseStream, + StreamingChoices) + + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.5, + ) + + # Chunks that form "```python\nprint(1)\n```" when concatenated + async def mock_stream(): + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="```python\n"))], + ) + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="print(1)\n"))], + ) + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="```"))], + ) + + request_data = {"model": "gpt-4", "metadata": {}} + yielded_chunks = [] + + with pytest.raises(HTTPException) as exc_info: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, + response=mock_stream(), + request_data=request_data, + ): + yielded_chunks.append(chunk) + + assert exc_info.value.status_code == 400 + assert "code block" in (exc_info.value.detail or {}).get("error", "") + # The chunk that completes the block (third chunk) must not have been yielded + assert len(yielded_chunks) == 2