Reflexión (Reflection)

Reflexión (Reflection) es un patrón de diseño en el que una generación de un modelo LLM es seguida por una reflexión, que a su vez es otra generación de LLM condicionada por la salida de la primera. Por ejemplo, dado un encargo para escribir código, el primer LLM puede generar un fragmento de código, y el segundo LLM puede generar una crítica sobre dicho código.

En el contexto de Saptiva-Agents y sus agentes, la reflexión puede implementarse como un par de agentes, donde el primer agente genera un mensaje y el segundo genera una respuesta a dicho mensaje. Estos dos agentes continúan interactuando hasta alcanzar una condición de parada, como un número máximo de iteraciones o una aprobación explícita del segundo agente.

🛠 Vamos a implementar este patrón de diseño de reflexión usando agentes de Saptiva-Agents.

Tendremos dos agentes:

• 🧑‍💻 Agente Coder: genera un fragmento de código.

• 🧑‍🏫 Agente Reviewer: genera una crítica del fragmento de código generado.

---

Protocolo de Mensajes

Antes de definir los agentes, primero necesitamos definir el protocolo de mensajes que usarán los agentes para comunicarse entre sí.

from dataclasses import dataclass


@dataclass
class CodeWritingTask:
    task: str


@dataclass
class CodeWritingResult:
    task: str
    code: str
    review: str


@dataclass
class CodeReviewTask:
    session_id: str
    code_writing_task: str
    code_writing_scratchpad: str
    code: str


@dataclass
class CodeReviewResult:
    review: str
    session_id: str
    approved: bool

El conjunto de mensajes anterior define el protocolo para nuestro ejemplo de patrón de diseño de reflexión:

• La aplicación envía un mensaje CodeWritingTask al agente coder

• El agente coder genera un mensaje CodeReviewTask, que se envía al agente reviewer

• El agente reviewer genera un mensaje CodeReviewResult, que se envía de regreso al agente coder

• Dependiendo del mensaje CodeReviewResult, si el código es aprobado, el agente coder envía un mensaje CodeWritingResult de vuelta a la aplicación; de lo contrario, el agente coder envía otro mensaje CodeReviewTask al agente reviewer, y el proceso continúa.

Podemos visualizar este protocolo de mensajes utilizando un diagrama de flujo de datos:

coder-reviewer data flow

---

Agentes

Ahora, definamos los agentes para el patrón de diseño de reflexión.

import json
import re
import uuid
from typing import Dict, List, Union

from saptiva_agents.core import MessageContext, RoutedAgent, TopicId, default_subscription, message_handler
from saptiva_agents.models import (
    AssistantMessage,
    LLMMessage,
    SystemMessage,
    UserMessage,
)
from saptiva_agents.base import SaptivaAIChatCompletionClient

Usamos la API de difusión (Broadcast API) para implementar el patrón de diseño. Los agentes implementan el modelo de publicación/suscripción (pub/sub). El agente coder se suscribe a los mensajes CodeWritingTask y CodeReviewResult, y publica los mensajes CodeReviewTask y CodeWritingResult.

@default_subscription
class CoderAgent(RoutedAgent):
    """An agent that performs code writing tasks."""

    def __init__(self, model_client: SaptivaAIChatCompletionClient) -> None:
        super().__init__("A code writing agent.")
        self._system_messages: List[LLMMessage] = [
            SystemMessage(
                content="""You are a proficient coder. You write code to solve problems.
Work with the reviewer to improve your code.
Always put all finished code in a single Markdown code block.
For example:
```python
def hello_world():
    print("Hello, World!")
```

Respond using the following format:

Thoughts: <Your comments>
Code: <Your code>
""",
            )
        ]
        self._model_client = model_client
        self._session_memory: Dict[str, List[CodeWritingTask | CodeReviewTask | CodeReviewResult]] = {}

    @message_handler
    async def handle_code_writing_task(self, message: CodeWritingTask, ctx: MessageContext) -> None:
        # Almacenar los mensajes en una memoria temporal solo para esta solicitud.
        session_id = str(uuid.uuid4())
        self._session_memory.setdefault(session_id, []).append(message)
        # Generar una respuesta usando la API de completado de chat.
        response = await self._model_client.create(
            self._system_messages + [UserMessage(content=message.task, source=self.metadata["type"])],
            cancellation_token=ctx.cancellation_token,
        )
        assert isinstance(response.content, str)
        # Extraer el bloque de código de la respuesta.
        code_block = self._extract_code_block(response.content)
        if code_block is None:
            raise ValueError("No se encontró el bloque de código.")
        # Crear una tarea de revisión de código.
        code_review_task = CodeReviewTask(
            session_id=session_id,
            code_writing_task=message.task,
            code_writing_scratchpad=response.content,
            code=code_block,
        )
        # Almacenar la tarea de revisión de código en la memoria de sesión.
        self._session_memory[session_id].append(code_review_task)
        # Publicar una tarea de revisión de código.
        await self.publish_message(code_review_task, topic_id=TopicId("default", self.id.key))
    
    @message_handler
    async def handle_code_review_result(self, message: CodeReviewResult, ctx: MessageContext) -> None:
        # Almacenar el resultado de la revisión en la memoria de sesión.
        self._session_memory[message.session_id].append(message)
        # Obtener la solicitud a partir de los mensajes previos.
        review_request = next(
            m for m in reversed(self._session_memory[message.session_id]) if isinstance(m, CodeReviewTask)
        )
        assert review_request is not None
        # Verificar si el código fue aprobado.
        if message.approved:
            # Publicar el resultado de la escritura de código.
            await self.publish_message(
                CodeWritingResult(
                    code=review_request.code,
                    task=review_request.code_writing_task,
                    review=message.review,
                ),
                topic_id=TopicId("default", self.id.key),
            )
            print("Code Writing Result:")
            print("-" * 80)
            print(f"Task:\n{review_request.code_writing_task}")
            print("-" * 80)
            print(f"Code:\n{review_request.code}")
            print("-" * 80)
            print(f"Review:\n{message.review}")
            print("-" * 80)
        else:
            # Crear una lista de mensajes LLM para enviar al modelo.
            messages: List[LLMMessage] = [*self._system_messages]
            for m in self._session_memory[message.session_id]:
                if isinstance(m, CodeReviewResult):
                    messages.append(UserMessage(content=m.review, source="Reviewer"))
                elif isinstance(m, CodeReviewTask):
                    messages.append(AssistantMessage(content=m.code_writing_scratchpad, source="Coder"))
                elif isinstance(m, CodeWritingTask):
                    messages.append(UserMessage(content=m.task, source="User"))
                else:
                    raise ValueError(f"Unexpected message type: {m}")
            # Generar una revisión usando la API de completado de chat.
            response = await self._model_client.create(messages, cancellation_token=ctx.cancellation_token)
            assert isinstance(response.content, str)
            # Extraer el bloque de código de la respuesta.
            code_block = self._extract_code_block(response.content)
            if code_block is None:
                raise ValueError("Code block not found.")
            # Crear una nueva tarea de revisión de código.
            code_review_task = CodeReviewTask(
                session_id=message.session_id,
                code_writing_task=review_request.code_writing_task,
                code_writing_scratchpad=response.content,
                code=code_block,
            )
            # Almacenar la tarea de revisión de código en la memoria de sesión.
            self._session_memory[message.session_id].append(code_review_task)
            # Publicar una nueva tarea de revisión de código.
            await self.publish_message(code_review_task, topic_id=TopicId("default", self.id.key))

    def _extract_code_block(self, markdown_text: str) -> Union[str, None]:
        pattern = r"```(\w+)\n(.*?)\n```"
        # Buscar el patrón en el texto markdown.
        match = re.search(pattern, markdown_text, re.DOTALL)
        # Extraer el lenguaje y el bloque de código si se encuentra una coincidencia.
        if match:
            return match.group(2)
        return None

Algunas cosas a tener en cuenta sobre CoderAgent:

• Utiliza chain-of-thought prompting (encadenamiento de pensamientos) en su mensaje del sistema.

• Almacena historiales de mensajes para diferentes CodeWritingTask en un diccionario, de modo que cada tarea tiene su propio historial.

• Al hacer una solicitud de inferencia al LLM usando su cliente de modelo, transforma el historial de mensajes en una lista de objetos saptiva_agents.models.LLMMessage para pasarlos al cliente del modelo.

El agente revisor se suscribe al mensaje CodeReviewTask y publica el mensaje CodeReviewResult.

@default_subscription
class ReviewerAgent(RoutedAgent):
    """An agent that performs code review tasks."""

    def __init__(self, model_client: SaptivaAIChatCompletionClient) -> None:
        super().__init__("A code reviewer agent.")
        self._system_messages: List[LLMMessage] = [
            SystemMessage(
                content="""You are a code reviewer. You focus on correctness, efficiency and safety of the code.
Respond using the following JSON format:
{
    "correctness": "<Your comments>",
    "efficiency": "<Your comments>",
    "safety": "<Your comments>",
    "approval": "<APPROVE or REVISE>",
    "suggested_changes": "<Your comments>"
}
""",
            )
        ]
        self._session_memory: Dict[str, List[CodeReviewTask | CodeReviewResult]] = {}
        self._model_client = model_client

    @message_handler
    async def handle_code_review_task(self, message: CodeReviewTask, ctx: MessageContext) -> None:
        # Formatear el prompt para la revisión de código.
        # Recopilar los comentarios previos si están disponibles.
        previous_feedback = ""
        if message.session_id in self._session_memory:
            previous_review = next(
                (m for m in reversed(self._session_memory[message.session_id]) if isinstance(m, CodeReviewResult)),
                None,
            )
            if previous_review is not None:
                previous_feedback = previous_review.review
        # Almacenar los mensajes en una memoria temporal solo para esta solicitud.
        self._session_memory.setdefault(message.session_id, []).append(message)
        prompt = f"""The problem statement is: {message.code_writing_task}
The code is:
```
{message.code}
```

Previous feedback:
{previous_feedback}

Please review the code. If previous feedback was provided, see if it was addressed.
"""
        # Generar una respuesta usando la API de finalización de chat.
        response = await self._model_client.create(
            self._system_messages + [UserMessage(content=prompt, source=self.metadata["type"])],
            cancellation_token=ctx.cancellation_token,
            json_output=True,
        )
        assert isinstance(response.content, str)
        # TODO: usar una librería de generación estructurada, por ejemplo, guidance, para asegurar que la respuesta tenga el formato esperado.
        # Analizar el JSON de la respuesta.
        review = json.loads(response.content)
        # Construir el texto de la revisión.
        review_text = "Code review:\n" + "\n".join([f"{k}: {v}" for k, v in review.items()])
        approved = review["approval"].lower().strip() == "approve"
        result = CodeReviewResult(
            review=review_text,
            session_id=message.session_id,
            approved=approved,
        )
        # Almacenar el resultado de la revisión en la memoria de sesión.
        self._session_memory[message.session_id].append(result)
        # Publicar el resultado de la revisión.
        await self.publish_message(result, topic_id=TopicId("default", self.id.key))

El ReviewerAgent utiliza el modo JSON (JSON-mode) al hacer una solicitud de inferencia al LLM, y también emplea chain-of-thought prompting (encadenamiento de pensamientos) en su mensaje del sistema.

---

Registro (Logging)

Activa el registro (logging) para ver los mensajes intercambiados entre los agentes.

import logging

logging.basicConfig(level=logging.INFO)
logging.getLogger("saptiva_agents.core").setLevel(logging.DEBUG)

---

Ejecutando el Patrón de Diseño

Probemos el patrón de diseño con una tarea de programación. Dado que todos los agentes están decorados con el decorador de clase default_subscription(), los agentes, al ser creados, se suscriben automáticamente al topic por defecto.

Publicamos un mensaje CodeWritingTask al topic por defecto para iniciar el proceso de reflexión.

from saptiva_agents import SAPTIVA_LEGACY
from saptiva_agents.core import DefaultTopicId, SingleThreadedAgentRuntime
from saptiva_agents.base import SaptivaAIChatCompletionClient

runtime = SingleThreadedAgentRuntime()
model_client = SaptivaAIChatCompletionClient(model=SAPTIVA_LEGACY, api_key="TU_SAPTIVA_API_KEY")
await ReviewerAgent.register(runtime, "ReviewerAgent", lambda: ReviewerAgent(model_client=model_client))
await CoderAgent.register(runtime, "CoderAgent", lambda: CoderAgent(model_client=model_client))
runtime.start()
await runtime.publish_message(
    message=CodeWritingTask(task="Write a function to find the sum of all even numbers in a list."),
    topic_id=DefaultTopicId(),
)

# Sigue procesando mensajes hasta que esté inactivo.
await runtime.stop_when_idle()
# Cierra el cliente del modelo.
await model_client.close()

INFO:autogen_core:Publishing message of type CodeWritingTask to all subscribers: {'task': 'Write a function to find the sum of all even numbers in a list.'}
INFO:autogen_core:Calling message handler for ReviewerAgent with message type CodeWritingTask published by Unknown
INFO:autogen_core:Calling message handler for CoderAgent with message type CodeWritingTask published by Unknown
INFO:autogen_core:Unhandled message: CodeWritingTask(task='Write a function to find the sum of all even numbers in a list.')
INFO:autogen_core.events:{"prompt_tokens": 101, "completion_tokens": 88, "type": "LLMCall"}
INFO:autogen_core:Publishing message of type CodeReviewTask to all subscribers: {'session_id': '51db93d5-3e29-4b7f-9f96-77be7bb02a5e', 'code_writing_task': 'Write a function to find the sum of all even numbers in a list.', 'code_writing_scratchpad': 'Thoughts: To find the sum of all even numbers in a list, we can use a list comprehension to filter out the even numbers and then use the `sum()` function to calculate their total. The implementation should handle edge cases like an empty list or a list with no even numbers.\n\nCode:\n```python\ndef sum_of_even_numbers(numbers):\n    return sum(num for num in numbers if num % 2 == 0)\n```', 'code': 'def sum_of_even_numbers(numbers):\n    return sum(num for num in numbers if num % 2 == 0)'}
INFO:autogen_core:Calling message handler for ReviewerAgent with message type CodeReviewTask published by CoderAgent:default
INFO:autogen_core.events:{"prompt_tokens": 163, "completion_tokens": 235, "type": "LLMCall"}
INFO:autogen_core:Publishing message of type CodeReviewResult to all subscribers: {'review': "Code review:\ncorrectness: The function correctly identifies and sums all even numbers in the provided list. The use of a generator expression ensures that only even numbers are processed, which is correct.\nefficiency: The function is efficient as it utilizes a generator expression that avoids creating an intermediate list, therefore using less memory. The time complexity is O(n) where n is the number of elements in the input list, which is optimal for this task.\nsafety: The function does not include checks for input types. If a non-iterable or a list containing non-integer types is passed, it could lead to unexpected behavior or errors. It’s advisable to handle such cases.\napproval: REVISE\nsuggested_changes: Consider adding input validation to ensure that 'numbers' is a list and contains only integers. You could raise a ValueError if the input is invalid. Example: 'if not isinstance(numbers, list) or not all(isinstance(num, int) for num in numbers): raise ValueError('Input must be a list of integers')'. This will make the function more robust.", 'session_id': '51db93d5-3e29-4b7f-9f96-77be7bb02a5e', 'approved': False}
INFO:autogen_core:Calling message handler for CoderAgent with message type CodeReviewResult published by ReviewerAgent:default
INFO:autogen_core.events:{"prompt_tokens": 421, "completion_tokens": 119, "type": "LLMCall"}
INFO:autogen_core:Publishing message of type CodeReviewTask to all subscribers: {'session_id': '51db93d5-3e29-4b7f-9f96-77be7bb02a5e', 'code_writing_task': 'Write a function to find the sum of all even numbers in a list.', 'code_writing_scratchpad': "Thoughts: I appreciate the reviewer's feedback on input validation. Adding type checks ensures that the function can handle unexpected inputs gracefully. I will implement the suggested changes and include checks for both the input type and the elements within the list to confirm that they are integers.\n\nCode:\n```python\ndef sum_of_even_numbers(numbers):\n    if not isinstance(numbers, list) or not all(isinstance(num, int) for num in numbers):\n        raise ValueError('Input must be a list of integers')\n    \n    return sum(num for num in numbers if num % 2 == 0)\n```", 'code': "def sum_of_even_numbers(numbers):\n    if not isinstance(numbers, list) or not all(isinstance(num, int) for num in numbers):\n        raise ValueError('Input must be a list of integers')\n    \n    return sum(num for num in numbers if num % 2 == 0)"}
INFO:autogen_core:Calling message handler for ReviewerAgent with message type CodeReviewTask published by CoderAgent:default
INFO:autogen_core.events:{"prompt_tokens": 420, "completion_tokens": 153, "type": "LLMCall"}
INFO:autogen_core:Publishing message of type CodeReviewResult to all subscribers: {'review': 'Code review:\ncorrectness: The function correctly sums all even numbers in the provided list. It raises a ValueError if the input is not a list of integers, which is a necessary check for correctness.\nefficiency: The function remains efficient with a time complexity of O(n) due to the use of a generator expression. There are no unnecessary intermediate lists created, so memory usage is optimal.\nsafety: The function includes input validation, which enhances safety by preventing incorrect input types. It raises a ValueError for invalid inputs, making the function more robust against unexpected data.\napproval: APPROVE\nsuggested_changes: No further changes are necessary as the previous feedback has been adequately addressed.', 'session_id': '51db93d5-3e29-4b7f-9f96-77be7bb02a5e', 'approved': True}
INFO:autogen_core:Calling message handler for CoderAgent with message type CodeReviewResult published by ReviewerAgent:default
INFO:autogen_core:Publishing message of type CodeWritingResult to all subscribers: {'task': 'Write a function to find the sum of all even numbers in a list.', 'code': "def sum_of_even_numbers(numbers):\n    if not isinstance(numbers, list) or not all(isinstance(num, int) for num in numbers):\n        raise ValueError('Input must be a list of integers')\n    \n    return sum(num for num in numbers if num % 2 == 0)", 'review': 'Code review:\ncorrectness: The function correctly sums all even numbers in the provided list. It raises a ValueError if the input is not a list of integers, which is a necessary check for correctness.\nefficiency: The function remains efficient with a time complexity of O(n) due to the use of a generator expression. There are no unnecessary intermediate lists created, so memory usage is optimal.\nsafety: The function includes input validation, which enhances safety by preventing incorrect input types. It raises a ValueError for invalid inputs, making the function more robust against unexpected data.\napproval: APPROVE\nsuggested_changes: No further changes are necessary as the previous feedback has been adequately addressed.'}
INFO:autogen_core:Calling message handler for ReviewerAgent with message type CodeWritingResult published by CoderAgent:default
INFO:autogen_core:Unhandled message: CodeWritingResult(task='Write a function to find the sum of all even numbers in a list.', code="def sum_of_even_numbers(numbers):\n    if not isinstance(numbers, list) or not all(isinstance(num, int) for num in numbers):\n        raise ValueError('Input must be a list of integers')\n    \n    return sum(num for num in numbers if num % 2 == 0)", review='Code review:\ncorrectness: The function correctly sums all even numbers in the provided list. It raises a ValueError if the input is not a list of integers, which is a necessary check for correctness.\nefficiency: The function remains efficient with a time complexity of O(n) due to the use of a generator expression. There are no unnecessary intermediate lists created, so memory usage is optimal.\nsafety: The function includes input validation, which enhances safety by preventing incorrect input types. It raises a ValueError for invalid inputs, making the function more robust against unexpected data.\napproval: APPROVE\nsuggested_changes: No further changes are necessary as the previous feedback has been adequately addressed.')

Code Writing Result:
--------------------------------------------------------------------------------
Task:
Write a function to find the sum of all even numbers in a list.
--------------------------------------------------------------------------------
Code:
def sum_of_even_numbers(numbers):
    if not isinstance(numbers, list) or not all(isinstance(num, int) for num in numbers):
        raise ValueError('Input must be a list of integers')
    
    return sum(num for num in numbers if num % 2 == 0)
--------------------------------------------------------------------------------
Review:
Code review:
correctness: The function correctly sums all even numbers in the provided list. It raises a ValueError if the input is not a list of integers, which is a necessary check for correctness.
efficiency: The function remains efficient with a time complexity of O(n) due to the use of a generator expression. There are no unnecessary intermediate lists created, so memory usage is optimal.
safety: The function includes input validation, which enhances safety by preventing incorrect input types. It raises a ValueError for invalid inputs, making the function more robust against unexpected data.
approval: APPROVE
suggested_changes: No further changes are necessary as the previous feedback has been adequately addressed.
--------------------------------------------------------------------------------

Los mensajes de registro muestran la interacción entre los agentes coder (programador) y reviewer (revisor). La salida final muestra el fragmento de código generado por el agente coder y la crítica generada por el agente reviewer.