Add support for PersonaAgent in PersonaDialogGenerator and tests

Author: sergioburdisso

src/sdialog/generators.py

@@ -17,7 +17,7 @@
 from langchain_core.messages import HumanMessage, SystemMessage
 
 from . import Dialog, Turn
-from .personas import Persona
+from .personas import Persona, PersonaAgent
 
 
 class LLMDialogOutput(BaseModel):
@@ -152,10 +152,13 @@ class PersonaDialogGenerator(DialogGenerator):
     :ivar persona_b: The second persona.
     :vartype persona_b: Persona
     """
+    _agent_a = None
+    _agent_b = None
+
     def __init__(self,
                  model: Union[ChatOllama, str],
-                 persona_a: Persona,
-                 persona_b: Persona,
+                 persona_a: Union[Persona, PersonaAgent],
+                 persona_b: Union[Persona, PersonaAgent],
                  dialogue_details: str = "",
                  response_details: str = "responses SHOULD NOT be too long and wordy, should be "
                                          "approximately one utterance long",
@@ -166,9 +169,9 @@ def __init__(self,
         :param model: The LLM or model name to use.
         :type model: Union[ChatOllama, str]
         :param persona_a: The first persona.
-        :type persona_a: Persona
+        :type persona_a: Persona (or PersonaAgent)
         :param persona_b: The second persona.
-        :type persona_b: Persona
+        :type persona_b: Persona (or PersonaAgent)
         :param dialogue_details: Additional dialogue instructions.
         :type dialogue_details: str
         :param response_details: Instructions for response style.
@@ -177,6 +180,10 @@ def __init__(self,
         :type scenario: dict
         """
 
+        if isinstance(persona_a, PersonaAgent) and isinstance(persona_b, PersonaAgent):
+            self._agent_a = persona_a
+            self._agent_b = persona_b
+
         dialogue_details = f"""Role play as the following two characters having a conversations. The characters are defined by the personas in the following lines. You always stay in character.
 [[ ## BEGING FIRST PERSONA ## ]]
 {persona_a}
@@ -194,3 +201,12 @@ def __init__(self,
         super().__init__(model=model,
                          dialogue_details=dialogue_details,
                          scenario=scenario)
+
+    def generate(self, seed: int = None, id: int = None, max_iterations: int = 20):
+        if self._agent_a and self._agent_b:
+            return self._agent_a.dialog_with(self._agent_b,
+                                             max_iterations=max_iterations,
+                                             id=id,
+                                             seed=seed)
+        else:
+            return super().generate(seed=seed, id=id)

tests/test_generators.py

@@ -1,5 +1,5 @@
 from sdialog.generators import DialogGenerator, PersonaDialogGenerator, LLMDialogOutput, Turn
-from sdialog.personas import Persona
+from sdialog.personas import Persona, PersonaAgent
 
 
 MODEL = "smollm:135m"
@@ -38,3 +38,12 @@ def test_persona_dialog_generator(monkeypatch):
     gen = PersonaDialogGenerator(MODEL, persona_a, persona_b)
     dialog = gen()
     assert hasattr(dialog, "turns")
+
+
+def test_persona_dialog_generator_with_agents(monkeypatch):
+    monkeypatch.setattr("sdialog.generators.ChatOllama", DummyLLM)
+    persona_a = PersonaAgent(DummyLLM(), name="A")
+    persona_b = PersonaAgent(DummyLLM(), name="B")
+    gen = PersonaDialogGenerator(MODEL, persona_a, persona_b)
+    dialog = gen()
+    assert hasattr(dialog, "turns")

tests/test_personas.py

@@ -45,7 +45,7 @@ def test_persona_agent_init(monkeypatch):
     persona = Persona(name="Alice")
     agent = PersonaAgent(DummyLLM(), persona=persona, name="Alice")
     assert agent.get_name() == "Alice"
-    assert "Role play" in agent.get_prompt()
+    assert "role play" in agent.get_prompt().lower()
     agent.set_first_utterances("Hi!")
     assert agent.first_utterances == "Hi!"
     agent.clear_orchestrators()