Spaces:
Build error
Build error
| import asyncio | |
| import gradio as gr | |
| from autogen.runtime_logging import start, stop | |
| from autogen_agentchat.agents import AssistantAgent | |
| from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination | |
| from autogen_agentchat.teams import RoundRobinGroupChat | |
| from autogen_ext.models.openai import OpenAIChatCompletionClient | |
| from autogen_agentchat.base import TaskResult | |
| # Configuration | |
| LOG_FILE = "team_runtime.log" | |
| config_list_primary = [ | |
| {"model": "lm-broca", "api_type": "openai", "max_tokens": 4096, "api_key": "sk-", "base_url": "<base-url>", "tags": ["lm-broca", "openai"]}, | |
| ] | |
| config_list_critic = [ | |
| {"model": "groq-mixtral-8x7b-32768", "api_type": "openai", "max_tokens": 16192, "api_key": "sk-", "base_url": "<base-url>", "tags": ["groq-mixtral-8x7b-32768", "openai"]}, | |
| ] | |
| llm_config_primary = { | |
| "config_list": config_list_primary, | |
| } | |
| llm_config_critic = { | |
| "config_list": config_list_critic, | |
| } | |
| #def create_llm_config_critic(): | |
| # return { | |
| # "model": "groq-mixtral-8x7b-32768", | |
| # "api_key": "sk-BhYjxpcKH_4w4H9jduTVwA", | |
| # "base_url": "https://litellm.j78.org/v1", | |
| # "cache_seed": None | |
| # } | |
| # Create the team with primary and critic agents | |
| def create_team(llm_config_primary, primary_system_message, critic_system_message): | |
| model_client = OpenAIChatCompletionClient(**llm_config) | |
| primary_agent = AssistantAgent( | |
| "primary", | |
| llm_config_primary={"config_list": config_list_primary}) | |
| system_message=primary_system_message, | |
| ) | |
| critic_agent = AssistantAgent( | |
| "critic", | |
| llm_config_critic={"config_list": config_list_critic}) | |
| system_message=critic_system_message | |
| ) | |
| # Set termination conditions (10-message cap OR "APPROVE" detected) | |
| max_message_termination = MaxMessageTermination(max_messages=10) | |
| text_termination = TextMentionTermination("APPROVE") | |
| combined_termination = max_message_termination | text_termination | |
| team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=combined_termination) | |
| return team, model_client | |
| # Function to stream the task through the workflow | |
| async def async_stream_task(task_message, api_key, primary_system_message, critic_system_message, documentation_system_message): | |
| # Start logging | |
| logging_session_id = start(logger_type="file", config={"filename": LOG_FILE}) | |
| print(f"Logging session ID: {logging_session_id}") | |
| llm_config_primary={"config_list": config_list_primary}) | |
| llm_config={"config_list": config_list_primary}) | |
| team, model_client = create_team(llm_config, primary_system_message, critic_system_message) | |
| documentation_triggered = False # Track if documentation agent was triggered | |
| final_output = None # Store the final approved output | |
| try: | |
| async for message in team.run_stream(task=task_message): | |
| if hasattr(message, "source") and hasattr(message, "content"): | |
| # Handle critic's approval | |
| if message.source == "critic" and "APPROVE" in message.content: | |
| print("Critic approved the response. Handing off to Documentation Agent...") | |
| documentation_triggered = True | |
| final_output = task_message # Capture the final approved output | |
| break | |
| yield message.source, message.content | |
| # Trigger Documentation Agent if approved | |
| if documentation_triggered and final_output: | |
| documentation_agent = AssistantAgent( | |
| "documentation", | |
| model_client=model_client, | |
| system_message=documentation_system_message, | |
| ) | |
| doc_task = f"Generate a '--help' message for the following code:\n\n{final_output}" | |
| async for doc_message in documentation_agent.run_stream(task=doc_task): | |
| if isinstance(doc_message, TaskResult): | |
| # Extract messages from TaskResult | |
| for msg in doc_message.messages: | |
| yield msg.source, msg.content | |
| else: | |
| yield doc_message.source, doc_message.content | |
| finally: | |
| # Stop logging | |
| stop() | |
| # Gradio interface function | |
| async def chat_interface(api_key, primary_system_message, critic_system_message, documentation_system_message, task_message): | |
| primary_messages = [] | |
| critic_messages = [] | |
| documentation_messages = [] | |
| # Append new messages while streaming | |
| async for source, output in async_stream_task(task_message, api_key, primary_system_message, critic_system_message, documentation_system_message): | |
| if source == "primary": | |
| primary_messages.append(output) | |
| elif source == "critic": | |
| critic_messages.append(output) | |
| elif source == "documentation": | |
| documentation_messages.append(output) | |
| # Return all outputs | |
| yield ( | |
| "\n".join(primary_messages), | |
| "\n".join(critic_messages), | |
| "\n".join(documentation_messages), | |
| ) | |
| # Gradio interface | |
| iface = gr.Interface( | |
| fn=chat_interface, | |
| inputs=[ | |
| gr.Textbox(label="OpenAI API Key", type="password", placeholder="Enter your OpenAI API Key"), | |
| gr.Textbox(label="Primary Agent System Message", placeholder="Enter the system message for the primary agent", value="You are a creative assistant focused on producing high-quality code."), | |
| gr.Textbox(label="Critic Agent System Message", placeholder="Enter the system message for the critic agent (requires APPROVAL tag!)", value="You are a critic assistant highly skilled in evaluating the quality of a given code or response. Provide constructive feedback and respond with 'APPROVE' once the feedback is addressed. Do not produce any content or code yourself, only provide feedback!"), | |
| gr.Textbox(label="Documentation Agent System Message", placeholder="Enter the system message for the documentation agent", value="You are a documentation assistant. Write a short and concise '--help' message for the provided code."), | |
| gr.Textbox(label="Task Message", placeholder="Code a random password generator using python."), | |
| ], | |
| outputs=[ | |
| gr.Textbox(label="The Primary Assistant Messages"), | |
| gr.Textbox(label="The Critics Assistant Messages"), | |
| gr.Textbox(label="The Documentation Assistant Message"), | |
| ], | |
| title="Team Workflow with Documentation Agent and Hard Cap", | |
| description="""Collaborative workflow between Primary, Critic, and Documentation agents. | |
| 1. The user can send a prompt to the primary agent. | |
| 2. The response will then be evaluated by the critic, which either sends feedback back to the primary agent or gives the APPROVAL sign. | |
| 3. If the APPROVAL sign is given, the documentation agent is asked to write a short documentation for the code (that has been approved by the critic and generated by the priamry agent. | |
| 4. (Note: There is a hard cap of 10 messages for the critic to approve the output of the primary agent. If it fails to do so the workflow is interrupted to prevent long loops)""" | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| iface.launch(share=True) |