
    ni                        U d Z ddlmZ ddlZddlmc mZ ddl	Z	ddl
Z
ddlmZmZmZ ddlmZmZ  G d d          Z edd	d
           eddd          dZded<    G d d          Z G d d          ZdS )zTests for the PM Agent router.    )annotationsN)AgentConfigMeetingState
PMDecision)build_pm_system_promptrun_pmc                  "    e Zd ZdZddZdd
ZdS )MockLLMClientz1Mock LLM client returning a canned JSON response.responsestrreturnNonec                0    || _         d| _        d| _        d S )Nr    )r   
call_countlast_system)selfr   s     4/root/projects/multi-agents-meeting/tests/test_pm.py__init__zMockLLMClient.__init__   s      "    systemmessageslist[dict[str, str]]c                >    | xj         dz  c_         || _        | j        S )N   )r   r   r   )r   r   r   s      r   chatzMockLLMClient.chat   s"    1!}r   N)r   r   r   r   r   r   r   r   r   r   )__name__
__module____qualname____doc__r   r    r   r   r
   r
      sB        ;;# # # #
     r   r
   	architectSoftware ArchitectzYou are an architect.)namerolesystem_promptdevopsDevOps EngineerzYou are DevOps.)r#   r(   zdict[str, AgentConfig]SAMPLE_AGENTSc                  *    e Zd ZdZddZddZddZdS )	TestBuildPMSystemPromptz!Tests for build_pm_system_prompt.r   r   c                F   t          t                    }d}||v }|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j	        |                    d x}}d}||v }|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j	        |                    d x}}d	}||v }|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j	        |                    d x}}d
}||v }|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j	        |                    d x}}d S )Nr#   inz%(py1)s in %(py3)spromptpy1py3assert %(py5)spy5r$   r(   r)   
r   r*   
@pytest_ar_call_reprcompare	_saferepr@py_builtinslocals_should_repr_global_nameAssertionError_format_explanationr   r1   @py_assert0@py_assert2@py_format4@py_format6s         r   #test_includes_agent_names_and_rolesz;TestBuildPMSystemPrompt.test_includes_agent_names_and_roles$   s#   '66${f$$$$$$$$$${f$$$$${$$$$$$$$$$$f$$$$$f$$$$$$$$$$$$$$$$$$$$$$$$$$$#-#v----------#v-----#-----------v-----v---------------------------!x6!!!!!!!!!!x6!!!!!x!!!!!!!!!!!6!!!!!6!!!!!!!!!!!!!!!!!!!!!!!!!!! * F********** F***** ***********F*****F*****************************r   c                    t          t                    }d}||v }|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j	        |                    d x}}d}||v }|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j	        |                    d x}}d	}||v }|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j	        |                    d x}}d S )
NJSONr.   r0   r1   r2   r5   r6   
CALL_AGENTFINISHr7   r@   s         r   %test_includes_json_schema_instructionz=TestBuildPMSystemPrompt.test_includes_json_schema_instruction+   s[   '66vvv%|v%%%%%%%%%%|v%%%%%|%%%%%%%%%%%v%%%%%v%%%%%%%%%%%%%%%%%%%%%%%%%%%!x6!!!!!!!!!!x6!!!!!x!!!!!!!!!!!6!!!!!6!!!!!!!!!!!!!!!!!!!!!!!!!!!!!r   c                   t          t                    }g }d}|j        } |            }||v }|}|sd}|j        }	 |	            }
||
v }|}|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          ndt          j        |          t          j        |          dz  }dd|iz  }|	                    |           |st          j        d|fd	||
f          t          j        |          dt          j                    v st          j        |          rt          j        |          ndt          j        |	          t          j        |
          d
z  }dd|iz  }|	                    |           t          j
        |d          i z  }dd|iz  }t          t          j        |                    dx}x}x}x}x}x}x}x}x}	}
g }d}|j        } |            }||v }|}|sd}|j        }	 |	            }
||
v }|}|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          ndt          j        |          t          j        |          dz  }dd|iz  }|	                    |           |st          j        d|fd	||
f          t          j        |          dt          j                    v st          j        |          rt          j        |          ndt          j        |	          t          j        |
          d
z  }dd|iz  }|	                    |           t          j
        |d          i z  }dd|iz  }t          t          j        |                    dx}x}x}x}x}x}x}x}x}	}
dS )z?PM system prompt includes keywords encouraging deep discussion.z	follow-upz	follow upr.   )zD%(py3)s in %(py9)s
{%(py9)s = %(py7)s
{%(py7)s = %(py5)s.lower
}()
}r1   )r4   r6   py7py9z%(py11)spy11)zJ%(py14)s in %(py20)s
{%(py20)s = %(py18)s
{%(py18)s = %(py16)s.lower
}()
})py14py16py18py20z%(py22)spy22r   zassert %(py25)spy25Ndeeperdeep)r   r*   lowerr8   r9   r:   r;   r<   r=   append_format_boolopr>   r?   )r   r1   @py_assert1rB   @py_assert6@py_assert8@py_assert4rA   @py_assert13@py_assert17@py_assert19@py_assert15@py_format10@py_format12@py_format21@py_format23@py_format24@py_format26s                     r   test_pm_prompt_encourages_depthz7TestBuildPMSystemPrompt.test_pm_prompt_encourages_depth1   s   '66MM{MflMllnnM{n,MMMMMMv|M||~~M~0MMMMMMMMMMMM{nMMMMM{MMMMMMMMMMMfMMMMMfMMMMMMlMMMMnMMMMMMMMMMMMMMMMMMMMMMMMMMM~MMMMMMMMMMMMMMMMvMMMMMvMMMMMM|MMMM~MMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMEExE6<E<<>>Ex>)EEEEEVEv|E||~~EV~-EEEEEEEEEEEEx>EEEEExEEEEEEEEEEE6EEEEE6EEEEEE<EEEE>EEEEEEEEEEEEEEEEEEEEEEEEEEEV~EEEEEVEEEEEEEEEEEvEEEEEvEEEEEE|EEEE~EEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEr   Nr   r   )r   r   r    r!   rE   rJ   rh   r"   r   r   r,   r,   !   s\        +++ + + +" " " "F F F F F Fr   r,   c                  Z    e Zd ZdZddZddZddZddZddZdd	Z	dd
Z
ddZddZdS )	TestRunPMzTests for run_pm function.r   r   c                   t          j        ddddd          }t          |          }t          d          }t	          |t
          |          }t          |t                    }|sd	d
t          j	                    v st          j        t                    rt          j        t                    nd
dt          j	                    v st          j        |          rt          j        |          nddt          j	                    v st          j        t                    rt          j        t                    ndt          j        |          dz  }t          t          j        |                    d }|j        }d}||k    }|st          j        d|fd||f          dt          j	                    v st          j        |          rt          j        |          ndt          j        |          t          j        |          dz  }	dd|	iz  }
t          t          j        |
                    d x}x}}|j        }d}||k    }|st          j        d|fd||f          dt          j	                    v st          j        |          rt          j        |          ndt          j        |          t          j        |          dz  }	dd|	iz  }
t          t          j        |
                    d x}x}}d S )NzNeed architect input.rH   r#   zWhat about scalability?analysisnext_actiontarget_agentprompt_for_agentr   zDatabase migrationtopic5assert %(py4)s
{%(py4)s = %(py0)s(%(py1)s, %(py2)s)
}
isinstancedecisionr   py0r3   py2py4==z3%(py2)s
{%(py2)s = %(py0)s.next_action
} == %(py5)sry   rz   r6   assert %(py7)srL   z4%(py2)s
{%(py2)s = %(py0)s.target_agent
} == %(py5)s)jsondumpsr
   r   r   r*   rv   r   r;   r<   r8   r=   r:   r>   r?   ro   r9   rp   )r   decision_jsonmock_clientstaterw   @py_assert3@py_format5rZ   r]   rD   @py_format8s              r   'test_run_pm_returns_call_agent_decisionz1TestRunPM.test_run_pm_returns_call_agent_decision;   s(   
/'' 9	$
 $
   $];;;#7888%<<(J///////////////z/////z/////////////(/////(/////////////J/////J////////////////////////#3|3#|3333333333#|333333333333x33333x333333#3333|333333333333333333333333333$33$3333333333$333333333333x33333x333333$333333333333333333333333333333333r   c                j  	 t          j        dddd          	 G 	fdd          } |            }t          dd	          }t          |t          |           |j        d
         d         }d}||v }|st          j        d|fd||f          t          j        |          dt          j
                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j        |                    dx}}d}||v }|st          j        d|fd||f          t          j        |          dt          j
                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j        |                    dx}}dS )z>When MeetingState has context, PM receives it in user content.OK.rI   # Donern   ro   final_reportc                  $    e Zd Zd
dZd fdZd	S )JTestRunPM.test_pm_user_content_includes_context.<locals>.CaptureMockClientr   r   c                    g | _         d S Nlast_messagesr   s    r   r   zSTestRunPM.test_pm_user_content_includes_context.<locals>.CaptureMockClient.__init__R       ;="""r   r   r   r   r   c                    || _         S r   r   r   r   r   r   s      r   r   zOTestRunPM.test_pm_user_content_includes_context.<locals>.CaptureMockClient.chatU       %-"$$r   Nri   r   r   r   r    r   r   r   s   r   CaptureMockClientr   Q   G        > > > >% % % % % % % %r   r   TestzImportant background infort   contextr   contentr.   r0   user_msgr2   r5   r6   NzBackground Context)r   r   r   r   r*   r   r8   r9   r:   r;   r<   r=   r>   r?   )
r   r   r   r   r   rA   rB   rC   rD   r   s
            @r   %test_pm_user_content_includes_contextz/TestRunPM.test_pm_user_content_includes_contextI   s0   
#$$
 $
  	% 	% 	% 	% 	% 	% 	% 	% 	% 	% ('))63NOOOum[111,Q/	:*6*h6666666666*h66666*66666666666h66666h666666666666666666666666666#/#x//////////#x/////#///////////x/////x/////////////////////////////r   c                   t          j        dddd          }t          |          }t          d          }t	          |t
          |          }|j        }d}||k    }|st          j        d|fd	||f          d
t          j
                    v st          j        |          rt          j        |          nd
t          j        |          t          j        |          dz  }dd|iz  }	t          t          j        |	                    d x}x}}d}
|j        }|
|v }|st          j        d|fd|
|f          t          j        |
          d
t          j
                    v st          j        |          rt          j        |          nd
t          j        |          dz  }dd|iz  }	t          t          j        |	                    d x}
x}}d S )NzAll covered.rI   z# Report

Done.r   rr   r   rs   r|   r~   rw   r   r   rL   Reportr.   )z4%(py1)s in %(py5)s
{%(py5)s = %(py3)s.final_report
})r3   r4   r6   )r   r   r
   r   r   r*   ro   r8   r9   r;   r<   r=   r:   r>   r?   r   )r   r   r   r   rw   rZ   r]   r   rD   r   rA   rB   s               r   #test_run_pm_returns_finish_decisionz-TestRunPM.test_run_pm_returns_finish_decision`   s!   
&#/$
 $
  
 $];;;6***%<<#/x/#x//////////#x////////////x/////x//////#////x///////////////////////////0800x00000000000x000000x000000000008000008000000000000000000000000000000000000r   c                R   d G fdd          }t          d          }t          |t           |                      }|j        }d}||k    }|st	          j        d|fd||f          d	t          j                    v st	          j        |          rt	          j	        |          nd	t	          j	        |          t	          j	        |          d
z  }dd|iz  }t          t	          j        |                    d x}x}}d}	|	k    }|st	          j        d|fd|	f          dt          j                    v st	          j                  rt	          j	                  ndt	          j	        |	          dz  }
dd|
iz  }t          t	          j        |                    d x}}	d S )Nr   c                      e Zd Zd fdZdS )	FTestRunPM.test_run_pm_retries_on_invalid_json.<locals>.RetryMockClientr   r   r   r   r   c                N    dz  dk     rdS t          j        dddd          S )Nr      zThis is not JSONOKrI   r   r   r   r   r   r   r   r   s      r   r   zKTestRunPM.test_run_pm_retries_on_invalid_json.<locals>.RetryMockClient.chatp   sB    a
>>--z $#+$,# #   r   Nr   r   r   r    r   r   s   r   RetryMockClientr   o   s3        	 	 	 	 	 	 	 	r   r   r   rs   rI   r|   r~   rw   r   r   rL   r   z%(py0)s == %(py3)sr   ry   r4   r5   r6   )r   r   r*   ro   r8   r9   r;   r<   r=   r:   r>   r?   )r   r   r   rw   rZ   r]   r   rD   r   rB   rC   r   s              @r   #test_run_pm_retries_on_invalid_jsonz-TestRunPM.test_run_pm_retries_on_invalid_jsonl   s   

	 
	 
	 
	 
	 
	 
	 
	 
	 
	 6***%0A0ABB#/x/#x//////////#x////////////x/////x//////#////x///////////////////////////zQzQzzQr   c                    t          d          }t          d          }t          j        t          d          5  t          |t          |           d d d            d S # 1 swxY w Y   d S )Nznot json at allrr   r   rs   zFailed to get valid PM decision)match)r
   r   pytestraisesRuntimeErrorr   r*   )r   r   r   s      r   $test_run_pm_raises_after_max_retriesz.TestRunPM.test_run_pm_raises_after_max_retries   s    #->???6***]</PQQQ 	6 	65-555	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6s   A  A$'A$c                R   d G fdd          }t          d          }t          |t           |                      }|j        }d}||k    }|st	          j        d|fd||f          d	t          j                    v st	          j        |          rt	          j	        |          nd	t	          j	        |          t	          j	        |          d
z  }dd|iz  }t          t	          j        |                    d x}x}}d}	|	k    }|st	          j        d|fd|	f          dt          j                    v st	          j                  rt	          j	                  ndt	          j	        |	          dz  }
dd|
iz  }t          t	          j        |                    d x}}	d S )Nr   c                      e Zd Zd fdZdS )	VTestRunPM.test_run_pm_validates_target_agent_exists.<locals>.AgentValidationMockClientr   r   r   r   r   c                ~    dz  dk    rt          j        ddddd          S t          j        dddd	d          S )
Nr   zNeed input.rH   nonexistent_agentzHello?rm   r   r#   zWhat do you think?r   r   s      r   r   z[TestRunPM.test_run_pm_validates_target_agent_exists.<locals>.AgentValidationMockClient.chat   sk    a
??:$1'3(;,4	' '    z %#/$/(<	# #   r   Nr   r   r   s   r   AgentValidationMockClientr      s3               r   r   r   rs   r#   r|   r   rw   r   r   rL      r   r   r   r5   r6   )r   r   r*   rp   r8   r9   r;   r<   r=   r:   r>   r?   )r   r   r   rw   rZ   r]   r   rD   r   rB   rC   r   s              @r   )test_run_pm_validates_target_agent_existsz3TestRunPM.test_run_pm_validates_target_agent_exists   s   
	 	 	 	 	 	 	 	 	 	$ 6***%0I0I0K0KLL$33$3333333333$333333333333x33333x333333$3333333333333333333333333333333zQzQzzQr   c                t   t          t                    }d}||v }|st          j        d|fd||f          t          j        |          dt          j                    v st          j        |          rt          j        |          nddz  }dd|iz  }t          t          j	        |                    dx}}dS )	zBPM system prompt includes instructions for handling HUMAN entries.HUMANr.   r0   r1   r2   r5   r6   Nr7   r@   s         r   ,test_pm_system_prompt_handles_human_feedbackz6TestRunPM.test_pm_system_prompt_handles_human_feedback   s    '66 w&          w&     w           &     &                             r   c                   ddl m} t          j        dddd           G fdd          } |            }t	          d	d
d          } |d          }t          |t          ||          }t          |t                    }|sddt          j
                    v st          j        t                    rt          j        t                    nddt          j
                    v st          j        |          rt          j        |          nddt          j
                    v st          j        t                    rt          j        t                    ndt          j        |          dz  }t          t          j        |                    d}|j        d         d         }	d	}
|
|	v }|st          j        d|fd|
|	f          t          j        |
          dt          j
                    v st          j        |	          rt          j        |	          nddz  }dd|iz  }t          t          j        |                    dx}
}dS )zDWhen context_manager is provided, PM receives compressed whiteboard.r   )ContextManagerr   rI   r   r   c                  $    e Zd Zd
dZd fdZd	S )]TestRunPM.test_run_pm_with_context_manager_uses_compressed_context.<locals>.CaptureMockClientr   r   c                    g | _         d S r   r   r   s    r   r   zfTestRunPM.test_run_pm_with_context_manager_uses_compressed_context.<locals>.CaptureMockClient.__init__   r   r   r   r   r   r   c                    || _         S r   r   r   s      r   r   zbTestRunPM.test_run_pm_with_context_manager_uses_compressed_context.<locals>.CaptureMockClient.chat   r   r   Nri   r   r   r   s   r   r   r      r   r   r   r   zSome contextr   )rt   r   current_roundip  )budget_tokens)context_managerru   rv   rw   r   rx   Nr   r.   r0   r   r2   r5   r6   )src.context_managerr   r   r   r   r   r*   rv   r   r;   r<   r8   r=   r:   r>   r?   r   r9   )r   r   r   r   r   cmrw   r   r   r   rA   rB   rC   rD   r   s                 @r   8test_run_pm_with_context_manager_uses_compressed_contextzBTestRunPM.test_run_pm_with_context_manager_uses_compressed_context   s   666666
#$$
 $
  	% 	% 	% 	% 	% 	% 	% 	% 	% 	% ('))6>QRSSS^$///%RPPP(J///////////////z/////z/////////////(/////(/////////////J/////J////////////////////////,Q/	:!v!!!!!!!!!!v!!!!!v!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!r   c                &   t          j        dddd          }t          |          }t          dd          }t	          |t
          |          }|j        }d}||k    }|st          j        d	|fd
||f          dt          j
                    v st          j        |          rt          j        |          ndt          j        |          t          j        |          dz  }dd|iz  }	t          t          j        |	                    dx}x}}dS )z?When context_manager is None, run_pm behaves exactly as before.r   rI   r   r   rr   r   
Backgroundr   r|   r~   rw   r   r   rL   N)r   r   r
   r   r   r*   ro   r8   r9   r;   r<   r=   r:   r>   r?   )
r   r   r   r   rw   rZ   r]   r   rD   r   s
             r   7test_run_pm_without_context_manager_backward_compatiblezATestRunPM.test_run_pm_without_context_manager_backward_compatible   sB   
#$$
 $
  
 $];;;6<@@@%<<#/x/#x//////////#x////////////x/////x//////#////x/////////////////////////////r   Nri   )r   r   r    r!   r   r   r   r   r   r   r   r   r   r"   r   r   rk   rk   8   s        $$4 4 4 40 0 0 0.
1 
1 
1 
1   (6 6 6 6   4! ! ! !
" " " "6
0 
0 
0 
0 
0 
0r   rk   )r!   
__future__r   builtinsr;   _pytest.assertion.rewrite	assertionrewriter8   r   r   
src.modelsr   r   r   src.pmr   r   r
   r*   __annotations__r,   rk   r"   r   r   <module>r      sd   $ $ $ " " " " " "                < < < < < < < < < < 1 1 1 1 1 1 1 1        +4HXopppkx.?O`aaa) )    F F F F F F F F.R0 R0 R0 R0 R0 R0 R0 R0 R0 R0r   