Update app.py

app.py

@@ -1,11 +1,15 @@
+import concurrent.futures
 import gradio as gr
 from huggingface_hub import InferenceClient
 import json
 import re
+import traceback
 
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-SYSTEM_PROMPT = """You are an AI agent planner that helps break down tasks into clear, actionable steps. For each task, you will:
+class AgentPlanner:
+    def __init__(self):
+        self.client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+        self.system_prompt = """You are an AI agent planner that helps break down tasks into clear, actionable steps. For each task, you will:
 1. Analyze the task and break it down into specific sub-tasks
 2. Create a structured plan with numbered steps
 3. Include any relevant considerations or potential challenges
@@ -26,103 +30,125 @@ SYSTEM_PROMPT = """You are an AI agent planner that helps break down tasks into
 
 Keep your responses focused and practical."""
 
-def parse_json_response(response_text):
-    """Extract JSON from the response text."""
-    try:
-        # Find JSON pattern in the text
-        json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
-        if json_match:
-            json_str = json_match.group()
-            return json.loads(json_str)
-        return None
-    except json.JSONDecodeError:
-        return None
+    def parse_json_response(self, response_text):
+        """Extract JSON from the response text."""
+        try:
+            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
+            if json_match:
+                json_str = json_match.group()
+                return json.loads(json_str)
+            return None
+        except json.JSONDecodeError:
+            return None
 
-def format_plan(plan_json):
-    """Format the JSON plan into a readable markdown string."""
-    if not plan_json:
-        return "Error: Could not parse the plan. Please try again."
-    
-    output = []
-    output.append("# Task Analysis")
-    output.append(plan_json.get("task_analysis", ""))
-    output.append("\n## Detailed Steps")
-    
-    for step in plan_json.get("steps", []):
-        output.append(f"\n### Step {step.get('step_number')}: {step.get('description')}")
-        output.append(f"- Estimated time: {step.get('estimated_time')}")
-        if step.get('considerations'):
-            output.append("\nConsiderations:")
-            for consideration in step['considerations']:
-                output.append(f"- {consideration}")
-    
-    if plan_json.get("potential_challenges"):
-        output.append("\n## Potential Challenges")
-        for challenge in plan_json["potential_challenges"]:
-            output.append(f"- {challenge}")
-    
-    if plan_json.get("resources_needed"):
-        output.append("\n## Required Resources")
-        for resource in plan_json["resources_needed"]:
-            output.append(f"- {resource}")
+    def format_plan(self, plan_json):
+        """Format the JSON plan into a readable markdown string."""
+        if not plan_json:
+            return "Error: Could not parse the plan. Please try again."
+        
+        output = []
+        output.append("# Task Analysis")
+        output.append(plan_json.get("task_analysis", ""))
+        output.append("\n## Detailed Steps")
+        
+        for step in plan_json.get("steps", []):
+            output.append(f"\n### Step {step.get('step_number')}: {step.get('description')}")
+            output.append(f"- Estimated time: {step.get('estimated_time')}")
+            if step.get('considerations'):
+                output.append("\nConsiderations:")
+                for consideration in step['considerations']:
+                    output.append(f"- {consideration}")
+        
+        if plan_json.get("potential_challenges"):
+            output.append("\n## Potential Challenges")
+            for challenge in plan_json["potential_challenges"]:
+                output.append(f"- {challenge}")
+        
+        if plan_json.get("resources_needed"):
+            output.append("\n## Required Resources")
+            for resource in plan_json["resources_needed"]:
+                output.append(f"- {resource}")
+        
+        return "\n".join(output)
+
+    def generate_plan(self, task, max_tokens=1024, temperature=0.7, top_p=0.95):
+        try:
+            messages = [
+                {"role": "system", "content": self.system_prompt},
+                {"role": "user", "content": task}
+            ]
+
+            response = ""
+            for message in self.client.chat_completion(
+                messages,
+                max_tokens=max_tokens,
+                stream=True,
+                temperature=temperature,
+                top_p=top_p,
+            ):
+                token = message.choices[0].delta.content
+                response += token
+                
+                # Try to parse and format JSON as it comes in
+                plan_json = self.parse_json_response(response)
+                if plan_json:
+                    formatted_response = self.format_plan(plan_json)
+                    yield formatted_response
+                else:
+                    yield response
+        except Exception as e:
+            yield f"Error generating plan: {str(e)}\n{traceback.format_exc()}"
+
+    def __del__(self):
+        self.executor.shutdown(wait=False)
+
+def create_interface():
+    planner = AgentPlanner()
     
-    return "\n".join(output)
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# AI Agent Planner")
+        gr.Markdown("Get organized, structured plans for any task")
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message=SYSTEM_PROMPT,
-    max_tokens=1024,
-    temperature=0.7,
-    top_p=0.95,
-):
-    messages = [{"role": "system", "content": system_message}]
+        task_input = gr.Textbox(
+            label="Task Description",
+            placeholder="Describe the task you need help planning...",
+            lines=3
+        )
 
-    for user_msg, assistant_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
+        with gr.Row():
+            max_tokens = gr.Slider(
+                minimum=1,
+                maximum=2048,
+                value=1024,
+                step=1,
+                label="Max Tokens"
+            )
+            temperature = gr.Slider(
+                minimum=0.1,
+                maximum=4.0,
+                value=0.7,
+                step=0.1,
+                label="Temperature"
+            )
+            top_p = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.95,
+                step=0.05,
+                label="Top-p"
+            )
 
-    messages.append({"role": "user", "content": message})
+        output = gr.Markdown(label="Generated Plan")
 
-    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
-        
-        # Try to parse and format JSON as it comes in
-        plan_json = parse_json_response(response)
-        if plan_json:
-            formatted_response = format_plan(plan_json)
-            yield formatted_response
-        else:
-            yield response
+        generate_btn = gr.Button("Generate Plan", variant="primary")
+        generate_btn.click(
+            planner.generate_plan,
+            inputs=[task_input, max_tokens, temperature, top_p],
+            outputs=output
+        )
 
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value=SYSTEM_PROMPT, label="System message", lines=5),
-        gr.Slider(minimum=1, maximum=2048, value=1024, 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)",
-        ),
-    ],
-    title="AI Agent Planner",
-    description="I help break down tasks into clear, actionable steps. Describe your task, and I'll create a detailed plan.",
-    theme=gr.themes.Soft(),
-)
+    return demo
 
 if __name__ == "__main__":
+    demo = create_interface()
     demo.launch()