Merge pull request #148 from Multi-Agent-LLMs/judge

Judge Agent

Author: jonas-becker

README.md

@@ -120,14 +120,18 @@ use_ablation: bool = False
 shuffle_input_samples: bool = False
 all_agents_generate_first_draft: bool = False
 all_agents_generate_draft: bool = False
-policy: Optional[str] = None
 voting_protocols_with_alterations: bool = False
 calculate_persona_diversity: bool = False
 challenge_final_results: bool = False
+judge_intervention: Optional[str] = None
+judge_metric: Optional[str] = None
+judge_endpoint_url: Optional[str] = None
+judge_api_key: str = "-"
+judge_always_intervene: bool = False
 ```
 
 ### Discussion Parameters:
-Response Generators: `freetext`, `json`, `simple`, `splitfreetext`
+Response Generators: `freetext`, `simple`, `splitfreetext`
 
 Decision Protocols: `approval_voting`, `consensus_voting`, `cumulative_voting`, `hybrid_consensus`, `majority_consensus`, `ranked_voting`, `simple_voting`, `summary`, `supermajority_consensus`, `unanimity_consensus`
 

mallm/agents/agent.py

@@ -232,6 +232,12 @@ def get_memories(
 
         return context_memory, memory_ids, current_draft
 
+    def forget_memories(self, turn: int) -> None:
+        keys_to_delete = [key for key, memory in self.memory.items() if memory.turn == turn]
+        for key in keys_to_delete:
+            del self.memory[key]
+        logger.debug(f"Forgot memories {keys_to_delete} from turn {turn} from agent {self.id}")
+
     def get_own_messages(self, context_length: Optional[int] = None) -> list[str]:
         """
         Retrieves memory from the agents memory bucket as a string

mallm/agents/judge.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import httpx
+
+from mallm.models.Chat import Chat
+from mallm.models.discussion.ResponseGenerator import ResponseGenerator
+
+if TYPE_CHECKING:
+    from mallm.coordinator import Coordinator
+
+from mallm.agents.agent import Agent
+from mallm.evaluation.evaluator import Evaluator
+from mallm.utils.types import Memory, TemplateFilling
+
+logger = logging.getLogger("mallm")
+
+
+class Judge(Agent):
+    def __init__(
+        self,
+        llm: Chat,
+        client: httpx.Client,
+        coordinator: Coordinator,
+        response_generator: ResponseGenerator,
+        persona: str,
+        persona_description: str,
+        metric: str,
+        chain_of_thought: bool = False,
+        drafting_agent: bool = False,
+        intervention_type: str = "regenerate",
+        references: Optional[list[str]] = None,
+    ):
+        if references is None:
+            references = []
+        super().__init__(
+            llm,
+            client,
+            coordinator,
+            response_generator,
+            persona,
+            persona_description,
+            chain_of_thought,
+            drafting_agent,
+        )
+        self.metric = Evaluator._initialize_metrics([metric])[0]
+        self.judgements: list[Optional[bool]] = []
+        self.performances: list[float] = []
+        self.judged_solutions: list[str] = []
+        self.intervention_type = intervention_type
+        self.coordinator = coordinator
+        self.references = references
+
+    def llm_as_a_judge(self, template_filling: TemplateFilling) -> Optional[bool]:
+        repeats = 0
+        while repeats < 3:
+            # check for drift
+            response = self.response_generator.generate_judgement(
+                template_filling, self.judged_solutions[-2], self.judged_solutions[-1]
+            )
+            if "[[A]]" in response.message:
+                return True     # answer_before is better
+            if "[[B]]" in response.message:
+                return False    # answer_after is better (problem drift)
+            logger.warning(f"Judge verdict is not valid: {response.message}. Retry number {repeats + 1}.")
+            repeats += 1
+        logger.warning(f"Judge verdict is not valid: {response.message}. All retries failed. The verdict will be saved as None.")
+        return None
+
+    def intervention(self,
+        unique_id: int,
+        turn: int,
+        memory_ids: list[int],
+        template_filling: TemplateFilling,
+        answer: str,
+        threshold: float = 0,
+        always_intervene: bool = False,
+        ) -> tuple[int, int]:
+        self.judged_solutions.append(answer)
+
+        if self.coordinator.judge_llm is not None:
+            if len(self.judged_solutions) < 2:
+                logger.debug("Judge skipped this turn because there are not enough solutions to judge.")
+                return unique_id, turn
+            on_track = self.llm_as_a_judge(template_filling)
+        else:
+            self.performances.append(Evaluator.calculate_score(answer, self.references, self.metric)["value"])
+            on_track = len(self.performances) > 1 and self.performances[-1] + threshold < self.performances[-2]
+        self.judgements.append(on_track)
+
+        if on_track is False or always_intervene:  # regenerates at most once per turn
+            if self.intervention_type == "regenerate":
+                # delete and restart the turn
+                logger.debug("Judge decided to regenerate the turn.")
+                self.coordinator.forget_memories(turn)
+                return unique_id - len(self.coordinator.agents) + 1, turn - 1
+            if self.intervention_type == "policy":
+                # Give the agents tips on how to improve their policy
+                logger.debug("Judge decided to give policy feedback.")
+                response = self.response_generator.generate_policy_intervention(
+                    template_filling,
+                    provide_labels=False
+                )
+                memory = Memory(
+                    message_id=unique_id,
+                    turn=turn,
+                    agent_id=self.id,
+                    persona=self.persona,
+                    contribution="judge",
+                    message=response.message,
+                    agreement=None,
+                    solution=None,
+                    memory_ids=memory_ids,
+                    additional_args={},
+                )
+                self.coordinator.update_memories([memory], self.coordinator.agents)
+                self.coordinator.memory.append(memory)
+                return unique_id + 1, turn
+        return unique_id, turn

mallm/agents/policyFeedback.py

@@ -7,12 +7,12 @@
 from typing import Optional
 
 import httpx
-from rich.progress import Console  # type: ignore
+from rich.progress import Console
 
 from mallm.agents.agent import Agent
 from mallm.agents.draftProposer import DraftProposer
+from mallm.agents.judge import Judge
 from mallm.agents.panelist import Panelist
-from mallm.agents.policyFeedback import PolicyFeedback
 from mallm.decision_protocol.protocol import DecisionProtocol
 from mallm.discourse_policy.policy import DiscoursePolicy
 from mallm.models.Chat import Chat
@@ -50,9 +50,9 @@ def __init__(
         model: Chat,
         client: httpx.Client,
         agent_generators: Optional[list[str]] = None,
-        policy: Optional[str] = None,
         num_neutral_agents: int = 0,
         console: Optional[Console] = None,
+        judge_model: Optional[Chat] = None,
     ):
         if agent_generators is None:
             agent_generators = ["expert", "expert", "expert"]
@@ -68,9 +68,9 @@ def __init__(
         self.response_generator: ResponseGenerator = SimpleResponseGenerator(self.llm)
         self.client = client
         self.agent_generators = agent_generators
-        self.policy = policy
         self.memory: list[Memory] = []
         self.console = console or Console()
+        self.judge_llm = judge_model
 
     def init_agents(
         self,
@@ -80,14 +80,16 @@ def init_agents(
         num_agents: int,
         chain_of_thought: bool,
         sample: InputExample,
+        judge_intervention: Optional[str] = None,
+        judge_metric: Optional[str] = None,
     ) -> None:
         """
         Instantiates the agents by
         1) identify helpful personas depending on the agent_generator
         2) create agents with the personas
         """
         logger.debug(
-            f"Coordinator {self.id} creates {num_agents} agents ({self.agent_generators}). Policy: {self.policy}"
+            f"Coordinator {self.id} creates {num_agents} agents ({self.agent_generators})."
         )
         self.panelists = []
         self.agents = []
@@ -142,16 +144,22 @@ def init_agents(
                 "Created only 1 agent. The discussion will be replaced by a self-improvement mechanism."
             )
 
-        if self.policy:
-            policyFeedback = PolicyFeedback(
-                self.llm,
+        self.judge = None
+        if judge_intervention and self.judge_llm:
+            self.judge = Judge(
+                self.judge_llm,
                 self.client,
                 self,
                 response_generator=self.response_generator,
-                persona="Policy Moderator",
-                policy=self.policy,
+                persona="Judge",
+                persona_description="Responsible for evaluating the solutions and providing feedback to the agents.",
+                metric=str(judge_metric),
+                chain_of_thought=False,
+                drafting_agent=False,
+                intervention_type=judge_intervention,
+                references=sample.references,
             )
-            self.agents.append(policyFeedback)
+            self.agents.append(self.judge)
 
     def get_agents(
         self, config: Config, worker_functions: WorkerFunctions
@@ -202,6 +210,8 @@ def discuss(
         bool,
         dict[int, Optional[VotingResultList]],
         ChallengeResult,
+        Optional[list[Optional[bool]]],
+        Optional[list[str]],
     ]:
         """
         The routine responsible for the discussion between agents to solve a task.
@@ -241,6 +251,8 @@ def discuss(
             num_agents=config.num_agents,
             chain_of_thought=config.use_chain_of_thought,
             sample=sample,
+            judge_intervention=config.judge_intervention,
+            judge_metric=config.judge_metric,
         )
 
         if config.decision_protocol not in DECISION_PROTOCOLS:
@@ -349,6 +361,8 @@ def discuss(
             decision_success,
             voting_results_per_turn,
             challenged_answers,
+            self.judge.judgements if self.judge else None,
+            self.judge.judged_solutions if self.judge else None,
         )
 
     def challenge_solution(
@@ -432,6 +446,10 @@ def get_memories(
 
         return context_memory, memory_ids, current_draft
 
+    def forget_memories(self, turn: int) -> None:
+        self.memory = [memory for memory in self.memory if memory.turn != turn]
+        logger.debug(f"Memories from turn {turn} have been removed from global memory.")
+
     def get_discussion_history(
         self,
         context_length: Optional[int] = None,

mallm/coordinator.py

@@ -42,6 +42,7 @@ def make_decision(
     ) -> tuple[str, bool, list[Agreement], str, Optional[VotingResultList]]:
         if len(agreements) > self.total_agents:
             agreements = agreements[-self.total_agents :]
+
         if agent_index != self.total_agents - 1:
             return "", False, agreements, "", None
 
@@ -166,6 +167,10 @@ def process_votes(
     ) -> tuple[str, Any, bool, str]:
         success = False
         vote_int = int("".join([x for x in vote_str if x.isnumeric()]))
+
+        # if len(final_answers) == 1:    # TODO: Add this in a future PR
+        #    vote_int = 0    # If there is only one answer, the agent must vote for it
+
         if 0 <= vote_int < len(final_answers):
             vote.append(vote_int)
             logger.info(