Update app.py

app.py

@@ -1,64 +1,128 @@
+import asyncio
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+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"
+
+def create_llm_config(api_key):
+    return {
+        "model": "gpt-4o",
+        "api_key": api_key,
+        "cache_seed": None
+    }
+
+# Create the team with primary and critic agents
+def create_team(llm_config, primary_system_message, critic_system_message):
+    model_client = OpenAIChatCompletionClient(**llm_config)
+
+    primary_agent = AssistantAgent(
+        "primary",
+        model_client=model_client,
+        system_message=primary_system_message,
+    )
+
+    critic_agent = AssistantAgent(
+        "critic",
+        model_client=model_client,
+        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):
+    # Start logging
+    logging_session_id = start(logger_type="file", config={"filename": LOG_FILE})
+    print(f"Logging session ID: {logging_session_id}")
+
+    llm_config = create_llm_config(api_key)
+    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="You are a documentation assistant. Write a short and concise '--help' message for the provided code.",
+            )
+            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, 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):
+        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", value="Critic. You are a helpful 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."),
+        gr.Textbox(label="Task Message", placeholder="Enter your task message"),
     ],
+    outputs=[
+        gr.Textbox(label="Primary Assistant Messages"),
+        gr.Textbox(label="Critic Messages"),
+        gr.Textbox(label="Documentation Messages"),
+    ],
+    title="Team Workflow with Documentation Agent and Hard Cap",
+    description="Collaborative workflow between Primary, Critic, and Documentation agents with a hard cap on messages."
 )
 
-
+# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch(share=True)