@@ -301,41 +301,45 @@ class TestLLMNotTriggeredByInjectedInput:
301301 when its required inputs (e.g. `query`) never arrive.
302302 """
303303
304- USER_PROMPT = '{% message role="user" %}{{ query }}{% endmessage %}'
305-
306304 def test_llm_not_triggered_by_injected_streaming_callback (self ):
305+
307306 @component
308307 class Planner :
309- @component .output_types (query = str , last_role = str )
308+ @component .output_types (messages = list [ ChatMessage ] , last_role = str )
310309 def run (self ) -> dict :
311- return {"query " : "What is 2+2?" , "last_role" : "assistant" }
310+ return {"messages " : [ ChatMessage . from_user ( "hello" )] , "last_role" : "assistant" }
312311
313312 chat_generator = MockChatGenerator ()
314- llm = LLM (chat_generator = chat_generator , user_prompt = self . USER_PROMPT )
313+ llm = LLM (chat_generator = chat_generator )
315314 chat_generator .run = MagicMock (return_value = {"replies" : [ChatMessage .from_assistant ("x" )]})
316315
317316 router = ConditionalRouter (
318317 routes = [
319318 {
320319 "condition" : "{{ last_role == 'tool' }}" ,
321- "output" : "{{ query }}" ,
320+ "output" : "{{ messages }}" ,
322321 "output_name" : "processing" ,
323- "output_type" : str ,
322+ "output_type" : list [ChatMessage ],
323+ },
324+ {
325+ "condition" : "{{ True }}" ,
326+ "output" : "{{ messages }}" ,
327+ "output_name" : "planning" ,
328+ "output_type" : list [ChatMessage ],
324329 },
325- {"condition" : "{{ True }}" , "output" : "{{ query }}" , "output_name" : "planning" , "output_type" : str },
326330 ],
327331 unsafe = True ,
328332 )
329333
330334 pipeline = Pipeline ()
331335 pipeline .add_component ("planner" , Planner ())
332336 pipeline .add_component ("router" , router )
333- pipeline .add_component ("branch_joiner" , BranchJoiner (type_ = str ))
337+ pipeline .add_component ("branch_joiner" , BranchJoiner (type_ = list [ ChatMessage ] ))
334338 pipeline .add_component ("llm" , llm )
335- pipeline .connect ("planner.query " , "router.query " )
339+ pipeline .connect ("planner.messages " , "router.messages " )
336340 pipeline .connect ("planner.last_role" , "router.last_role" )
337341 pipeline .connect ("router.processing" , "branch_joiner.value" )
338- pipeline .connect ("branch_joiner.value" , "llm.query " )
342+ pipeline .connect ("branch_joiner.value" , "llm.messages " )
339343
340344 result = pipeline .run (data = {"llm" : {"streaming_callback" : sync_streaming_callback }})
341345
0 commit comments