Manejo De Estados

Hasta ahora, hemos discutido cómo construir componentes en una aplicación multi-agente: agentes, equipos, condiciones de terminación. En muchos casos, es útil guardar el estado de estos componentes en el disco y cargarlos más tarde. Esto es particularmente útil en una aplicación web donde los endpoints sin estado responden a solicitudes y necesitan cargar el estado de la aplicación desde un almacenamiento persistente.

En este cuaderno, discutiremos cómo guardar y cargar el estado de agentes, equipos y terminación.

---

Agentes de Ahorro y Carga

Podemos obtener el estado de un agente llamando al método save_state() en un AssistantAgent.

import asyncio

from saptiva_agents import SAPTIVA_LEGACY
from saptiva_agents.agents import AssistantAgent
from saptiva_agents.conditions import MaxMessageTermination
from saptiva_agents.messages import TextMessage
from saptiva_agents.teams import RoundRobinGroupChat
from saptiva_agents.ui import Console
from saptiva_agents.core import CancellationToken
from saptiva_agents.base import SaptivaAIChatCompletionClient

model_client = SaptivaAIChatCompletionClient(
        model=SAPTIVA_LEGACY,
        api_key="TU_SAPTIVA_API_KEY",
)
assistant_agent = AssistantAgent(
    name="assistant_agent",
    system_message="You are a helpful assistant",
    model_client=model_client
)

async def main() -> None:
    # Utilice `asyncio.run(...)` cuando lo ejecute en un script
    response = await assistant_agent.on_messages(
        [TextMessage(content="Write a 3 line poem on lake tangayika", source="user")], CancellationToken()
    )
    print(response.chat_message.content)
    await model_client.close()
    
asyncio.run(main())

In Tanganyika's embrace so wide and deep,  
Ancient waters cradle secrets they keep,  
Echoes of time where horizons sleep.  

agent_state = await assistant_agent.save_state()
print(agent_state)

{'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_messages': [{'content': 'Write a 3 line poem on lake tangayika', 'source': 'user', 'type': 'UserMessage'}, {'content': "In Tanganyika's embrace so wide and deep,  \nAncient waters cradle secrets they keep,  \nEchoes of time where horizons sleep.  ", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}

new_assistant_agent = AssistantAgent(
    name="assistant_agent",
    system_message="You are a helpful assistant",
    model_client=SaptivaAIChatCompletionClient(
        model=SAPTIVA_LEGACY,
    ),
)
await new_assistant_agent.load_state(agent_state)

# Utilice `asyncio.run(...)` cuando lo ejecute en un script
response = await new_assistant_agent.on_messages(
    [TextMessage(content="What was the last line of the previous poem you wrote", source="user")], CancellationToken()
)
print(response.chat_message.content)

The last line of the poem was: "Echoes of time where horizons sleep."

Nota

Para AssistantAgent, su estado consiste en el modelo_context. Si escribes tu propio agente personalizado, considera sobrescribir los métodos save_state() y load_state() para personalizar el comportamiento. Las implementaciones predeterminadas guardan y cargan un estado vacío.

---

Equipos de Ahorro y Carga

Podemos obtener el estado de un equipo llamando al método save_state en el equipo y cargarlo de nuevo llamando al método load_state en el equipo.

Cuando llamamos a save_state en un equipo, guarda el estado de todos los agentes en el equipo.

Comenzaremos creando un equipo simple de RoundRobinGroupChat con un solo agente y le pediremos que escriba un poema.

# Definir un equipo.
assistant_agent = AssistantAgent(
    name="assistant_agent",
    system_message="You are a helpful assistant",
    model_client = SaptivaAIChatCompletionClient(
        model=SAPTIVA_LEG,
        api_key="TU_SAPTIVA_API_KEY",
    ),
)
agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))

# Ejecutar el equipo y transmitir mensajes a la consola.
stream = agent_team.run_stream(task="Write a beautiful poem 3-line about lake tangayika")

# Utilice `asyncio.run(...)` cuando lo ejecute en un script
await Console(stream)

# Guardar el estado del equipo del agente.
team_state = await agent_team.save_state()

---------- user ----------
Write a beautiful poem 3-line about lake tangayika
---------- assistant_agent ----------
In Tanganyika's gleam, beneath the azure skies,  
Whispers of ancient waters, in tranquil guise,  
Nature's mirror, where dreams and serenity lie.
[Prompt tokens: 29, Completion tokens: 34]
---------- Summary ----------
Number of messages: 2
Finish reason: Maximum number of messages 2 reached, current message count: 2
Total prompt tokens: 29
Total completion tokens: 34
Duration: 0.71 seconds

Si reiniciamos el equipo (simulando la creación del equipo) y hacemos la pregunta ¿What was the last line of the poem you wrote?, vemos que el equipo no puede lograrlo ya que no hay referencia a la ejecución anterior.

await agent_team.reset()
stream = agent_team.run_stream(task="What was the last line of the poem you wrote?")
await Console(stream)

---------- user ----------
What was the last line of the poem you wrote?
---------- assistant_agent ----------
I'm sorry, but I am unable to recall or access previous interactions, including any specific poem I may have composed in our past conversations. If you like, I can write a new poem for you.
[Prompt tokens: 28, Completion tokens: 40]
---------- Summary ----------
Number of messages: 2
Finish reason: Maximum number of messages 2 reached, current message count: 2
Total prompt tokens: 28
Total completion tokens: 40
Duration: 0.70 seconds

TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=28, completion_tokens=40), content="I'm sorry, but I am unable to recall or access previous interactions, including any specific poem I may have composed in our past conversations. If you like, I can write a new poem for you.", type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')

A continuación, cargamos el estado del equipo y hacemos la misma pregunta. Vemos que el equipo es capaz de devolver con precisión la última línea del poema que escribió.

print(team_state)

# Cargar estado del equipo
await agent_team.load_state(team_state)
stream = agent_team.run_stream(task="What was the last line of the poem you wrote?")
await Console(stream)

{'type': 'TeamState', 'version': '1.0.0', 'agent_states': {'group_chat_manager/a55364ad-86fd-46ab-9449-dcb5260b1e06': {'type': 'RoundRobinManagerState', 'version': '1.0.0', 'message_thread': [{'source': 'user', 'models_usage': None, 'content': 'Write a beautiful poem 3-line about lake tangayika', 'type': 'TextMessage'}, {'source': 'assistant_agent', 'models_usage': {'prompt_tokens': 29, 'completion_tokens': 34}, 'content': "In Tanganyika's gleam, beneath the azure skies,  \nWhispers of ancient waters, in tranquil guise,  \nNature's mirror, where dreams and serenity lie.", 'type': 'TextMessage'}], 'current_turn': 0, 'next_speaker_index': 0}, 'collect_output_messages/a55364ad-86fd-46ab-9449-dcb5260b1e06': {}, 'assistant_agent/a55364ad-86fd-46ab-9449-dcb5260b1e06': {'type': 'ChatAgentContainerState', 'version': '1.0.0', 'agent_state': {'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_messages': [{'content': 'Write a beautiful poem 3-line about lake tangayika', 'source': 'user', 'type': 'UserMessage'}, {'content': "In Tanganyika's gleam, beneath the azure skies,  \nWhispers of ancient waters, in tranquil guise,  \nNature's mirror, where dreams and serenity lie.", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}, 'message_buffer': []}}, 'team_id': 'a55364ad-86fd-46ab-9449-dcb5260b1e06'}
---------- user ----------
What was the last line of the poem you wrote?
---------- assistant_agent ----------
The last line of the poem I wrote is:  
"Nature's mirror, where dreams and serenity lie."
[Prompt tokens: 86, Completion tokens: 22]
---------- Summary ----------
Number of messages: 2
Finish reason: Maximum number of messages 2 reached, current message count: 2
Total prompt tokens: 86
Total completion tokens: 22
Duration: 0.96 seconds

TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=22), content='The last line of the poem I wrote is:  \n"Nature\'s mirror, where dreams and serenity lie."', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')

---

Estado Persistente (Archivo o Base de Datos)

En muchos casos, es posible que deseemos persistir el estado del equipo en disco (o en una base de datos) y cargarlo posteriormente. El estado es un diccionario que puede ser serializado a un archivo o escrito en una base de datos.

import json

## Guardar el estado en el disco.

with open("coding/team_state.json", "w") as f:
    json.dump(team_state, f)

## Cargar el estado desde el disco.
with open("coding/team_state.json", "r") as f:
    team_state = json.load(f)

new_agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))
await new_agent_team.load_state(team_state)
stream = new_agent_team.run_stream(task="What was the last line of the poem you wrote?")
await Console(stream)
await model_client.close()

---------- user ----------
What was the last line of the poem you wrote?
---------- assistant_agent ----------
The last line of the poem I wrote is:  
"Nature's mirror, where dreams and serenity lie."
[Prompt tokens: 86, Completion tokens: 22]
---------- Summary ----------
Number of messages: 2
Finish reason: Maximum number of messages 2 reached, current message count: 2
Total prompt tokens: 86
Total completion tokens: 22
Duration: 0.72 seconds

TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=22), content='The last line of the poem I wrote is:  \n"Nature\'s mirror, where dreams and serenity lie."', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')