fix CI: black formatting and update test to use input_filters

This commit is contained in:
Adil Hafeez 2026-03-17 17:10:14 -07:00
parent a00a102f86
commit 6d58469b82
No known key found for this signature in database
GPG key ID: 9B18EF7691369645
6 changed files with 61 additions and 16 deletions

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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,

View file

@ -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)}
(
{
**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)}
(
{
**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:

View file

@ -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: