mirror of
https://github.com/katanemo/plano.git
synced 2026-05-21 13:55:15 +02:00
fix CI: black formatting and update test to use input_filters
This commit is contained in:
parent
a00a102f86
commit
6d58469b82
6 changed files with 61 additions and 16 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue