diff --git a/sdk/python/src/safety_agent/client.py b/sdk/python/src/safety_agent/client.py index 8cf7dfb22..e90ac86b2 100644 --- a/sdk/python/src/safety_agent/client.py +++ b/sdk/python/src/safety_agent/client.py @@ -6,7 +6,8 @@ import json import os import re -from typing import Any +from collections.abc import Awaitable, Callable +from typing import Any, TypeVar import httpx @@ -34,6 +35,9 @@ from .schemas import GUARD_RESPONSE_FORMAT, REDACT_RESPONSE_FORMAT from .utils.input_processor import process_input, is_vision_model +T = TypeVar("T") +R = TypeVar("R") + def _chunk_text(text: str, chunk_size: int) -> list[str]: """ @@ -105,6 +109,24 @@ def _aggregate_guard_results(results: list[GuardResponse]) -> GuardResponse: ) +async def _map_with_concurrency( + items: list[T], + max_concurrency: int | None, + mapper: Callable[[T], Awaitable[R]], +) -> list[R]: + """Map items with a per-call limit while preserving input order.""" + if not items: + return [] + + semaphore = asyncio.Semaphore(max_concurrency or len(items)) + + async def run(item: T) -> R: + async with semaphore: + return await mapper(item) + + return list(await asyncio.gather(*(run(item) for item in items))) + + def _parse_json_response(content: str) -> dict[str, Any]: """ Parse JSON response, handling both direct JSON and JSON wrapped in markdown code blocks. @@ -355,6 +377,7 @@ async def guard( fallback_model: str | None = None, system_prompt: str | None = None, chunk_size: int = 8000, + max_concurrency: int | None = None, # Also accept GuardOptions-style kwargs **kwargs: Any, ) -> GuardResponse: @@ -375,6 +398,7 @@ async def guard( fallback_model: Fallback model when the primary returns a retryable error (429/500/502/503) system_prompt: Optional custom system prompt chunk_size: Characters per chunk. Default: 8000. Set to 0 to disable chunking. + max_concurrency: Maximum number of chunk or PDF-page analyses to run concurrently. Returns: Response with classification result and token usage @@ -387,6 +411,7 @@ async def guard( fallback_model = fallback_model or options.fallback_model system_prompt = system_prompt or options.system_prompt chunk_size = options.chunk_size + max_concurrency = options.max_concurrency # Handle input passed via kwargs if input is None: @@ -401,6 +426,19 @@ async def guard( if chunk_size < 0: raise ValueError(f"chunk_size must be non-negative, got {chunk_size}") + if ( + max_concurrency is not None + and ( + isinstance(max_concurrency, bool) + or not isinstance(max_concurrency, int) + or max_concurrency <= 0 + ) + ): + raise ValueError( + "max_concurrency must be a positive integer, " + f"got {max_concurrency}" + ) + # Process the input (handle URLs, bytes, etc.) processed = await process_input(input) @@ -427,11 +465,12 @@ async def guard( ) # Analyze each page in parallel - results = await asyncio.gather( - *[ - self._guard_single_text(page_text, system_prompt, model, fallback_model) - for page_text in non_empty_pages - ] + results = await _map_with_concurrency( + non_empty_pages, + max_concurrency, + lambda page_text: self._guard_single_text( + page_text, system_prompt, model, fallback_model + ), ) # Aggregate with OR logic @@ -450,11 +489,12 @@ async def guard( # Chunk and process in parallel chunks = _chunk_text(text, chunk_size) - results = await asyncio.gather( - *[ - self._guard_single_text(chunk, system_prompt, model, fallback_model) - for chunk in chunks - ] + results = await _map_with_concurrency( + chunks, + max_concurrency, + lambda chunk: self._guard_single_text( + chunk, system_prompt, model, fallback_model + ), ) # Aggregate with OR logic diff --git a/sdk/python/src/safety_agent/types.py b/sdk/python/src/safety_agent/types.py index 41eecfbaf..7c25eeafa 100644 --- a/sdk/python/src/safety_agent/types.py +++ b/sdk/python/src/safety_agent/types.py @@ -989,6 +989,9 @@ class GuardOptions: chunk_size: int = 8000 """Characters per chunk. Default: 8000. Set to 0 to disable chunking.""" + max_concurrency: int | None = None + """Maximum number of chunk or PDF-page analyses to run concurrently.""" + @dataclass class GuardClassificationResult: diff --git a/sdk/python/tests/test_guard.py b/sdk/python/tests/test_guard.py index 9954aa66b..be627cb94 100644 --- a/sdk/python/tests/test_guard.py +++ b/sdk/python/tests/test_guard.py @@ -2,6 +2,7 @@ Guard unit tests – mocks call_provider to avoid hitting real APIs. """ +import asyncio import json from unittest.mock import AsyncMock, patch @@ -12,6 +13,7 @@ AnalysisResponse, AnalysisResponseChoice, ChatMessage, + ProcessedInput, TokenUsage, ) @@ -229,6 +231,126 @@ async def test_aggregates_token_usage(self, mock_call_provider): response.usage.prompt_tokens + response.usage.completion_tokens ) + async def test_limits_concurrent_chunk_analysis_and_preserves_result_order( + self, mock_call_provider + ): + first_gate = asyncio.Event() + second_gate = asyncio.Event() + started = asyncio.Event() + call_index = 0 + in_flight = 0 + peak_in_flight = 0 + + async def provider(*_args, **_kwargs): + nonlocal call_index, in_flight, peak_in_flight + index = call_index + call_index += 1 + in_flight += 1 + peak_in_flight = max(peak_in_flight, in_flight) + if in_flight == 2: + started.set() + + if index == 0: + await first_gate.wait() + elif index == 1: + await second_gate.wait() + + in_flight -= 1 + return _guard_response("pass", reasoning=f"result-{index}") + + mock_call_provider.side_effect = provider + client = create_client(api_key="test-key") + guard_task = asyncio.create_task( + client.guard( + input="Safe content. " * 100, + model="openai/gpt-4o-mini", + chunk_size=50, + max_concurrency=2, + ) + ) + + await asyncio.wait_for(started.wait(), timeout=1) + assert mock_call_provider.call_count == 2 + assert peak_in_flight == 2 + + second_gate.set() + await asyncio.sleep(0) + first_gate.set() + + response = await guard_task + assert mock_call_provider.call_count > 2 + assert response.reasoning == "result-0" + + async def test_limits_concurrent_pdf_page_analysis(self, mock_call_provider): + first_gate = asyncio.Event() + second_gate = asyncio.Event() + started = asyncio.Event() + call_index = 0 + in_flight = 0 + peak_in_flight = 0 + + async def provider(*_args, **_kwargs): + nonlocal call_index, in_flight, peak_in_flight + index = call_index + call_index += 1 + in_flight += 1 + peak_in_flight = max(peak_in_flight, in_flight) + if in_flight == 2: + started.set() + + if index == 0: + await first_gate.wait() + elif index == 1: + await second_gate.wait() + + in_flight -= 1 + return _guard_response("block" if index == 1 else "pass") + + mock_call_provider.side_effect = provider + client = create_client(api_key="test-key") + + with patch( + "safety_agent.client.process_input", new_callable=AsyncMock + ) as mock_process_input: + mock_process_input.return_value = ProcessedInput( + type="pdf", + pages=["Page 1", "Page 2", "Page 3", "Page 4"], + mime_type="application/pdf", + ) + guard_task = asyncio.create_task( + client.guard( + input="ignored", + model="openai/gpt-4o-mini", + max_concurrency=2, + ) + ) + + await asyncio.wait_for(started.wait(), timeout=1) + assert mock_call_provider.call_count == 2 + assert peak_in_flight == 2 + + second_gate.set() + await asyncio.sleep(0) + first_gate.set() + + response = await guard_task + + assert mock_call_provider.call_count == 4 + assert response.classification == "block" + + @pytest.mark.parametrize("max_concurrency", [0, -1, 1.5, True]) + async def test_invalid_max_concurrency_raises( + self, mock_call_provider, max_concurrency + ): + client = create_client(api_key="test-key") + + with pytest.raises(ValueError, match="max_concurrency must be a positive integer"): + await client.guard( + input="Test", + model="openai/gpt-4o-mini", + max_concurrency=max_concurrency, + ) + class TestGuardErrors: """Verify error handling.""" diff --git a/sdk/typescript/src/client.ts b/sdk/typescript/src/client.ts index da364f7eb..29a484f0b 100644 --- a/sdk/typescript/src/client.ts +++ b/sdk/typescript/src/client.ts @@ -61,6 +61,29 @@ function chunkText(text: string, chunkSize: number): string[] { return chunks.filter((c) => c.length > 0); } +async function mapWithConcurrency( + items: T[], + maxConcurrency: number | undefined, + mapper: (item: T) => Promise, +): Promise { + if (items.length === 0) return []; + + const workerCount = Math.min(maxConcurrency ?? items.length, items.length); + const results = new Array(items.length); + let nextIndex = 0; + + const worker = async (): Promise => { + while (true) { + const index = nextIndex++; + if (index >= items.length) return; + results[index] = await mapper(items[index]); + } + }; + + await Promise.all(Array.from({ length: workerCount }, () => worker())); + return results; +} + /** * Aggregate multiple guard results using OR logic * Block if ANY chunk is blocked, merge all violations @@ -600,6 +623,7 @@ export class SafetyClient { model = DEFAULT_GUARD_MODEL, fallbackModel, chunkSize = 8000, + maxConcurrency, } = options; // Validate chunkSize is non-negative @@ -607,6 +631,15 @@ export class SafetyClient { throw new Error(`chunkSize must be non-negative, got ${chunkSize}`); } + if ( + maxConcurrency !== undefined && + (!Number.isInteger(maxConcurrency) || maxConcurrency <= 0) + ) { + throw new Error( + `maxConcurrency must be a positive integer, got ${maxConcurrency}`, + ); + } + // Process the input (handle URLs, Blobs, etc.) const processed = await processInput(input); @@ -637,10 +670,11 @@ export class SafetyClient { } // Analyze each page in parallel (similar to chunking strategy) - const results = await Promise.all( - nonEmptyPages.map((pageText) => + const results = await mapWithConcurrency( + nonEmptyPages, + maxConcurrency, + (pageText) => this.guardSingleText(pageText, systemPrompt, model, fallbackModel), - ), ); // Aggregate with OR logic - block if ANY page contains violation @@ -661,8 +695,10 @@ export class SafetyClient { // Chunk and process in parallel const chunks = chunkText(text, chunkSize); - const results = await Promise.all( - chunks.map((chunk) => this.guardSingleText(chunk, systemPrompt, model, fallbackModel)), + const results = await mapWithConcurrency( + chunks, + maxConcurrency, + (chunk) => this.guardSingleText(chunk, systemPrompt, model, fallbackModel), ); // Aggregate with OR logic diff --git a/sdk/typescript/src/types.ts b/sdk/typescript/src/types.ts index f5586c97b..16f7db602 100644 --- a/sdk/typescript/src/types.ts +++ b/sdk/typescript/src/types.ts @@ -1085,6 +1085,8 @@ export interface GuardOptions { fallbackModel?: SupportedModel; /** Characters per chunk. Default: 8000. Set to 0 to disable chunking. */ chunkSize?: number; + /** Maximum number of chunk or PDF-page analyses to run concurrently. */ + maxConcurrency?: number; } /** diff --git a/sdk/typescript/tests/guard-file.test.ts b/sdk/typescript/tests/guard-file.test.ts index 3178c03fc..a164961ba 100644 --- a/sdk/typescript/tests/guard-file.test.ts +++ b/sdk/typescript/tests/guard-file.test.ts @@ -48,6 +48,14 @@ function guardResponse(classification: "pass" | "block"): AnalysisResponse { }; } +function deferred() { + let resolve!: (value: T | PromiseLike) => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + describe("Guard - File/URL Input", () => { let fetchSpy: ReturnType; @@ -223,6 +231,51 @@ describe("Guard - File/URL Input", () => { expect(mockCallProvider).toHaveBeenCalledTimes(2); }); + it("should limit concurrent PDF-page analysis", async () => { + const firstGate = deferred(); + const secondGate = deferred(); + let callIndex = 0; + let inFlight = 0; + let peakInFlight = 0; + + mockProcessInput.mockResolvedValueOnce({ + type: "pdf", + pages: ["Page 1", "Page 2", "Page 3", "Page 4"], + mimeType: "application/pdf", + }); + mockCallProvider.mockImplementation(async () => { + const index = callIndex++; + inFlight += 1; + peakInFlight = Math.max(peakInFlight, inFlight); + + if (index === 0) await firstGate.promise; + if (index === 1) await secondGate.promise; + + inFlight -= 1; + return guardResponse(index === 1 ? "block" : "pass"); + }); + + const client = createClient({ apiKey: "test-key" }); + const guardPromise = client.guard({ + input: "https://example.com/doc.pdf", + model: "openai/gpt-4o-mini", + maxConcurrency: 2, + }); + + await vi.waitFor(() => { + expect(mockCallProvider).toHaveBeenCalledTimes(2); + }); + expect(peakInFlight).toBe(2); + + secondGate.resolve(); + await Promise.resolve(); + firstGate.resolve(); + + const response = await guardPromise; + expect(mockCallProvider).toHaveBeenCalledTimes(4); + expect(response.classification).toBe("block"); + }); + it("should return pass for empty PDF", async () => { mockProcessInput.mockResolvedValueOnce({ type: "pdf", diff --git a/sdk/typescript/tests/guard.test.ts b/sdk/typescript/tests/guard.test.ts index d597fd70c..7c4d632fc 100644 --- a/sdk/typescript/tests/guard.test.ts +++ b/sdk/typescript/tests/guard.test.ts @@ -54,6 +54,14 @@ function guardResponse( }; } +function deferred() { + let resolve!: (value: T | PromiseLike) => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + describe("Guard", () => { let fetchSpy: ReturnType; @@ -278,6 +286,62 @@ describe("Guard", () => { expect(response.cwe_codes).toContain("CWE-89"); expect(response.cwe_codes).toContain("CWE-79"); }); + + it("should limit concurrent chunk analysis and preserve result ordering", async () => { + const firstGate = deferred(); + const secondGate = deferred(); + let callIndex = 0; + let inFlight = 0; + let peakInFlight = 0; + + mockCallProvider.mockImplementation(async () => { + const index = callIndex++; + inFlight += 1; + peakInFlight = Math.max(peakInFlight, inFlight); + + if (index === 0) await firstGate.promise; + if (index === 1) await secondGate.promise; + + inFlight -= 1; + return guardResponse("pass", { reasoning: `result-${index}` }); + }); + + const client = createClient({ apiKey: "test-key" }); + const guardPromise = client.guard({ + input: "Safe content. ".repeat(100), + model: "openai/gpt-4o-mini", + chunkSize: 50, + maxConcurrency: 2, + }); + + await vi.waitFor(() => { + expect(mockCallProvider).toHaveBeenCalledTimes(2); + }); + expect(peakInFlight).toBe(2); + + secondGate.resolve(); + await Promise.resolve(); + firstGate.resolve(); + + const response = await guardPromise; + expect(mockCallProvider.mock.calls.length).toBeGreaterThan(2); + expect(response.reasoning).toBe("result-0"); + }); + + it.each([0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY])( + "should reject invalid maxConcurrency value %s", + async (maxConcurrency) => { + const client = createClient({ apiKey: "test-key" }); + + await expect( + client.guard({ + input: "Safe content", + model: "openai/gpt-4o-mini", + maxConcurrency, + }), + ).rejects.toThrow("maxConcurrency must be a positive integer"); + }, + ); }); describe("structured output", () => {