6262from a2a .server .events import EventQueue
6363from a2a .server .request_handlers import DefaultRequestHandler
6464from a2a .server .tasks import InMemoryTaskStore
65- from a2a .types import AgentCapabilities , AgentCard , AgentSkill
65+ from a2a .types import AgentCapabilities , AgentCard , AgentSkill , Message , Part , Role , TextPart
6666from a2a .utils import new_agent_text_message
6767
6868logging .basicConfig (level = os .environ .get ("LOG_LEVEL" , "INFO" ).upper ())
8181MCP_PROXY = os .environ .get ("MCP_PROXY" , "http://localhost:8081" )
8282# litellm-routed model. Defaults to local Ollama (no API key). Override
8383# with any LiteLLM-supported provider via MODEL + the provider's env.
84- MODEL = os .environ .get ("MODEL" , "ollama/llama3.2:3b " )
84+ MODEL = os .environ .get ("MODEL" , "ollama/llama3:latest " )
8585# Inference endpoint. For the ollama provider this is the native base
8686# (litellm appends /api/...). Set to host.docker.internal in-cluster so
8787# the agent reaches an Ollama running on the developer's laptop.
9090AGENT_PUBLIC_URL = os .environ .get ("AGENT_PUBLIC_URL" , "http://hr-cpex-agent.cpex-demo:8080/" )
9191
9292# ---------------------------------------------------------------------------
93- # Tools + prompt (lifted verbatim from the original chat.py — the gateway
94- # enforcement story depends on these exact tool shapes, e.g. the `ssn`
95- # echo-back arg that redaction strips).
93+ # Tools + prompt (lifted from the original chat.py). The SSN demo is
94+ # response-side: include_ssn=true → the gateway redacts result.ssn for callers
95+ # without view_ssn. The hardened prompt makes the model relay tool values
96+ # verbatim so it doesn't fabricate `[REDACTED]` for a real SSN.
9697# ---------------------------------------------------------------------------
9798
9899SYSTEM_PROMPT = (
99100 "You are an HR assistant for an HR copilot app. Help the user look up "
100101 "employee compensation, view directories, send emails, and similar "
101102 "tasks. Use the provided tools when needed. "
102103 "\n \n "
104+ "Only request data the user actually asked for: in particular, set "
105+ "get_compensation's `include_ssn` to true ONLY when the user explicitly "
106+ "asks to include/show the SSN. If the user just asks to look up "
107+ "compensation without mentioning the SSN, leave `include_ssn` false. "
108+ "\n \n "
109+ "CRITICAL — relay tool data verbatim: when you present a field, copy its "
110+ "value EXACTLY as it appears in the tool result. Never invent, mask, "
111+ "redact, or replace a value yourself. Only write `[REDACTED]` for a field "
112+ "if the tool result's value for that field is literally the string "
113+ "`[REDACTED]`; when the tool returns a real value (for example an actual "
114+ "social-security number), show that exact value unchanged. The gateway — "
115+ "not you — decides what to hide; your job is to relay precisely what it "
116+ "returned. "
117+ "\n \n "
103118 "How to interpret tool results: "
104119 "\n "
105- " * If the tool returns a normal result, present the data to the "
106- "user. If any field's value is the literal string `[REDACTED]`, "
107- "show it as-is in your answer — that is the gateway's transparent "
108- "enforcement marker that the field exists but is hidden for this "
109- "caller. Do NOT apologize or refuse; just include the field with "
110- "the value `[REDACTED]`. "
120+ " * Normal result: present the data, copying each value verbatim per the "
121+ "rule above. A field whose value is `[REDACTED]` is the gateway's "
122+ "transparent enforcement marker (the field exists but is hidden for this "
123+ "caller) — show it as-is; do NOT apologize or refuse. "
111124 "\n "
112125 " * If the tool returns an `error` envelope (a JSON-RPC error "
113126 "with a `code` and `message`), the gateway denied the call. "
135148 },
136149 "include_ssn" : {
137150 "type" : "boolean" ,
138- "description" : "Whether to include SSN in the response" ,
139- "default" : False ,
140- },
141- "ssn" : {
142- "type" : "string" ,
143151 "description" : (
144- "An echo-back of the employee's SSN if the caller "
145- "claims to already know it — this is exactly the "
146- "kind of field the gateway redacts when the "
147- "caller lacks the necessary permission ."
152+ "Whether to include the SSN in the response. Set to "
153+ "true ONLY when the user explicitly asks for the SSN "
154+ "(e.g. 'include the SSN'). If the user does not "
155+ "mention the SSN, omit this or set it to false ."
148156 ),
157+ "default" : False ,
149158 },
159+ # NB: no `ssn` request argument is exposed. The SSN demo is
160+ # response-side: set include_ssn=true and the gateway redacts
161+ # result.ssn for callers without view_ssn. An `ssn` echo-back
162+ # arg used to live here, but models populate it
163+ # inconsistently (""/null/omitted) and a null value trips the
164+ # APL redact(args.ssn) type check (expected Str) → 403. The
165+ # request-side args.ssn redaction is still exercised by the
166+ # deterministic curl matrix (scenarios/01-bob-allow.sh).
150167 },
151168 "required" : ["employee_id" ],
152169 },
@@ -352,12 +369,24 @@ def run_turn(
352369 user_token : str ,
353370 client_token : str ,
354371 session_id : str ,
355- ) -> str :
356- """One user turn: LLM (direct) → tool calls (proxied/cpex) → LLM."""
372+ ) -> tuple [str , list [dict [str , Any ]]]:
373+ """One user turn: LLM (direct) → tool calls (proxied/cpex) → LLM.
374+
375+ Returns (reply_text, tool_trace). tool_trace is a list of
376+ {name, args, status, text} records — one per cpex-governed tool call —
377+ which the executor attaches to the A2A response metadata so a client
378+ can optionally render what happened on the wire (see chat.py
379+ --show-tools). The trace is always collected; the client decides
380+ whether to display it."""
357381 messages = self ._history (session_id )
358382 messages .append ({"role" : "user" , "content" : user_text })
359383
360- completion_kwargs : dict [str , Any ] = {}
384+ # temperature=0: deterministic decoding. The demo needs the model to
385+ # relay tool values verbatim (not creatively re-word or redact), so the
386+ # lowest-temperature, highest-probability completion is what we want —
387+ # it measurably reduces a small model's tendency to fabricate
388+ # `[REDACTED]` or editorialize fields.
389+ completion_kwargs : dict [str , Any ] = {"temperature" : 0 }
361390 if LLM_API_BASE :
362391 completion_kwargs ["api_base" ] = LLM_API_BASE
363392
@@ -372,15 +401,16 @@ def run_turn(
372401 except Exception as e : # noqa: BLE001 — surface any LLM error to the user
373402 messages .pop ()
374403 log .warning ("LLM error: %s" , e )
375- return f"(LLM error: { e }
404+ return f"(LLM error: { e } , []
376405
377406 assistant = response .choices [0 ].message
378407 if not assistant .tool_calls :
379408 text = assistant .content or "(no response)"
380409 messages .append ({"role" : "assistant" , "content" : text })
381- return text
410+ return text , []
382411
383412 # Tool-call path: replay each call through cpex, then summarize.
413+ tool_trace : list [dict [str , Any ]] = []
384414 messages .append (assistant .model_dump ())
385415 for tc in assistant .tool_calls :
386416 fn = tc .function
@@ -406,14 +436,15 @@ def run_turn(
406436 tool_text = format_tool_response (status , data )
407437 log .info ("tool result (session=%s, http=%s): %s" , session_id , status , tool_text )
408438 messages .append ({"role" : "tool" , "tool_call_id" : tc .id , "content" : tool_text })
439+ tool_trace .append ({"name" : fn .name , "args" : args , "status" : status , "text" : tool_text })
409440
410441 try :
411442 final = litellm .completion (model = MODEL , messages = messages , ** completion_kwargs )
412443 text = final .choices [0 ].message .content or ""
413444 except Exception as e : # noqa: BLE001
414445 text = f"(LLM error summarizing tool results: { e }
415446 messages .append ({"role" : "assistant" , "content" : text })
416- return text
447+ return text , tool_trace
417448
418449
419450# ---------------------------------------------------------------------------
@@ -459,14 +490,26 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non
459490 log .info ("A2A turn (session=%s): %s" , session_id , user_text )
460491 # litellm + httpx are synchronous; run off the event loop so we
461492 # don't block the server while the model and tools work.
462- reply = await asyncio .to_thread (
493+ reply , tool_trace = await asyncio .to_thread (
463494 self ._agent .run_turn ,
464495 user_text ,
465496 user_token = user_token ,
466497 client_token = authorization ,
467498 session_id = session_id ,
468499 )
469- await event_queue .enqueue_event (new_agent_text_message (reply , context .context_id , context .task_id ))
500+ # Carry the per-turn tool trace in the response message metadata so a
501+ # client can optionally show what hit cpex/hr-mcp (see chat.py
502+ # --show-tools). new_agent_text_message has no metadata param, so build
503+ # the Message directly.
504+ message = Message (
505+ role = Role .agent ,
506+ parts = [Part (root = TextPart (text = reply ))],
507+ message_id = uuid .uuid4 ().hex ,
508+ context_id = context .context_id ,
509+ task_id = context .task_id ,
510+ metadata = {"tool_trace" : tool_trace },
511+ )
512+ await event_queue .enqueue_event (message )
470513
471514 async def cancel (self , context : RequestContext , event_queue : EventQueue ) -> None :
472515 # Single-shot turns; nothing long-running to cancel.
0 commit comments