importasynciofrompydanticimportBaseModelfromagentsimport (Agent,GuardrailFunctionOutput,InputGuardrailTripwireTriggered,RunContextWrapper,Runner,TResponseInputItem,input_guardrail,)classMathHomeworkOutput(BaseModel):is_math_homework:boolreasoning:strguardrail_agent=Agent(name="Guardrail check",instructions="Check if the user is asking you to do their math homework.",output_type=MathHomeworkOutput,)@input_guardrailasyncdefmath_guardrail(ctx:RunContextWrapper[None],agent:Agent,input:str|list[TResponseInputItem])->GuardrailFunctionOutput:result=awaitRunner.run(guardrail_agent,input,context=ctx.context)returnGuardrailFunctionOutput(output_info=result.final_output,tripwire_triggered=result.final_output.is_math_homework,    )math_agent=Agent(name="Customer support agent",instructions="You are a customer support agent. You help customers with their questions.",input_guardrails=[math_guardrail],)

from .agentimportmath_agentclassMyChatKitServer(ChatKitServer):def__init__(self,data_store:Store,attachment_store:AttachmentStore|None=None    ):super().__init__(data_store,attachment_store)assistant_agent=math_agentasyncdefrespond(self,thread:ThreadMetadata,input:UserMessageItem|None,context:Any,    )->AsyncIterator[ThreadStreamEvent]:context=AgentContext(thread=thread,store=self.store,request_context=context,        )result=Runner.run_streamed(self.assistant_agent,awaitsimple_to_agent_input(input)ifinputelse [],context=context,        )asyncforeventinstream_agent_response(context,result,        ):yieldevent# ...