diff --git a/cli/test/test_config_generator.py b/cli/test/test_config_generator.py index dc4c3c44..17fa56cc 100644 --- a/cli/test/test_config_generator.py +++ b/cli/test/test_config_generator.py @@ -104,13 +104,13 @@ listeners: agents: - id: simple_tmobile_rag_agent description: t-mobile virtual assistant for device contracts. - filter_chain: + input_filters: - query_rewriter - context_builder - response_generator - id: research_agent description: agent to research and gather information from various sources. - filter_chain: + input_filters: - research_agent - response_generator port: 8000 diff --git a/demos/filter_chains/http_filter/src/rag_agent/context_builder.py b/demos/filter_chains/http_filter/src/rag_agent/context_builder.py index 75b73ef0..be67eb10 100644 --- a/demos/filter_chains/http_filter/src/rag_agent/context_builder.py +++ b/demos/filter_chains/http_filter/src/rag_agent/context_builder.py @@ -219,7 +219,12 @@ async def context_builder(path: str, request: Request) -> dict: messages, traceparent_header, request_id ) - return {**body, "messages": [{"role": msg.role, "content": msg.content} for msg in updated_messages]} + return { + **body, + "messages": [ + {"role": msg.role, "content": msg.content} for msg in updated_messages + ], + } # Register MCP tool only if mcp is available diff --git a/demos/filter_chains/mcp_filter/src/rag_agent/context_builder.py b/demos/filter_chains/mcp_filter/src/rag_agent/context_builder.py index df60f70c..f46d3f19 100644 --- a/demos/filter_chains/mcp_filter/src/rag_agent/context_builder.py +++ b/demos/filter_chains/mcp_filter/src/rag_agent/context_builder.py @@ -220,7 +220,12 @@ async def context_builder(body: dict, path: str) -> dict: messages, traceparent_header, request_id ) - return {**body, "messages": [{"role": msg.role, "content": msg.content} for msg in updated_messages]} + return { + **body, + "messages": [ + {"role": msg.role, "content": msg.content} for msg in updated_messages + ], + } # Register MCP tool only if mcp is available diff --git a/demos/filter_chains/mcp_filter/src/rag_agent/query_rewriter.py b/demos/filter_chains/mcp_filter/src/rag_agent/query_rewriter.py index d175d123..dedf5359 100644 --- a/demos/filter_chains/mcp_filter/src/rag_agent/query_rewriter.py +++ b/demos/filter_chains/mcp_filter/src/rag_agent/query_rewriter.py @@ -24,6 +24,7 @@ plano_client = AsyncOpenAI( api_key="EMPTY", # Plano doesn't require a real API key ) + async def rewrite_query_with_plano( messages: List[ChatMessage], traceparent_header: str, diff --git a/demos/filter_chains/pii_anonymizer/pii_anonymizer.py b/demos/filter_chains/pii_anonymizer/pii_anonymizer.py index 9adf16cd..f8e2e613 100644 --- a/demos/filter_chains/pii_anonymizer/pii_anonymizer.py +++ b/demos/filter_chains/pii_anonymizer/pii_anonymizer.py @@ -57,18 +57,32 @@ async def anonymize(path: str, request: Request) -> dict[str, Any]: body = {**body, "input": anonymized} elif isinstance(input_val, list): items = [ - {**item, "content": anonymize_message_content(item.get("content", ""), all_mappings)} - if isinstance(item, dict) and item.get("role") == "user" - else item + ( + { + **item, + "content": anonymize_message_content( + item.get("content", ""), all_mappings + ), + } + if isinstance(item, dict) and item.get("role") == "user" + else item + ) for item in input_val ] body = {**body, "input": items} else: # /v1/chat/completions and /v1/messages both use "messages" messages = [ - {**msg, "content": anonymize_message_content(msg.get("content", ""), all_mappings)} - if msg.get("role") == "user" - else msg + ( + { + **msg, + "content": anonymize_message_content( + msg.get("content", ""), all_mappings + ), + } + if msg.get("role") == "user" + else msg + ) for msg in body.get("messages", []) ] if messages: diff --git a/demos/filter_chains/pii_anonymizer/store.py b/demos/filter_chains/pii_anonymizer/store.py index 5be58f68..74cbd01c 100644 --- a/demos/filter_chains/pii_anonymizer/store.py +++ b/demos/filter_chains/pii_anonymizer/store.py @@ -50,7 +50,9 @@ def restore_streaming(request_id: str, content: str, mapping: Dict[str, str]) -> else: _buffers.pop(request_id, None) if restored != content: - logger.info("request_id=%s restored '%s' -> '%s'", request_id, content, restored) + logger.info( + "request_id=%s restored '%s' -> '%s'", request_id, content, restored + ) return restored @@ -70,13 +72,17 @@ def deanonymize_sse( if chunk.get("type") == "content_block_delta": delta = chunk.get("delta", {}) if delta.get("type") == "text_delta" and delta.get("text"): - delta["text"] = restore_streaming(request_id, delta["text"], mapping) + delta["text"] = restore_streaming( + request_id, delta["text"], mapping + ) else: # {"choices": [{"delta": {"content": "..."}}]} for choice in chunk.get("choices", []): delta = choice.get("delta", {}) if delta.get("content"): - delta["content"] = restore_streaming(request_id, delta["content"], mapping) + delta["content"] = restore_streaming( + request_id, delta["content"], mapping + ) result_lines.append("data: " + json.dumps(chunk)) except json.JSONDecodeError: result_lines.append(line) @@ -95,10 +101,19 @@ def deanonymize_json( if is_anthropic: # {"content": [{"type": "text", "text": "..."}]} for part in body.get("content", []): - if isinstance(part, dict) and part.get("type") == "text" and part.get("text"): + if ( + isinstance(part, dict) + and part.get("type") == "text" + and part.get("text") + ): restored, _ = deanonymize_text(part["text"], mapping) if restored != part["text"]: - logger.info("request_id=%s restored '%s' -> '%s'", request_id, part["text"], restored) + logger.info( + "request_id=%s restored '%s' -> '%s'", + request_id, + part["text"], + restored, + ) part["text"] = restored else: # {"choices": [{"message": {"content": "..."}}]} @@ -108,7 +123,12 @@ def deanonymize_json( if content and isinstance(content, str): restored, _ = deanonymize_text(content, mapping) if restored != content: - logger.info("request_id=%s restored '%s' -> '%s'", request_id, content, restored) + logger.info( + "request_id=%s restored '%s' -> '%s'", + request_id, + content, + restored, + ) message["content"] = restored return Response(content=json.dumps(body), media_type="application/json") except json.JSONDecodeError: