Update app.py

app.py

@@ -1,151 +1,142 @@
 import concurrent.futures
 import gradio as gr
-from huggingface_hub import InferenceClient
 import json
-import re
 import traceback
 
-class AgentPlanner:
+class FunctionManager:
     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
-4. Format the response as a JSON with the following structure:
-{
-    "task_analysis": "Brief analysis of the main task",
-    "steps": [
-        {
-            "step_number": 1,
-            "description": "Step description",
-            "estimated_time": "Time estimate",
-            "considerations": ["List of considerations"]
-        }
-    ],
-    "potential_challenges": ["List of potential challenges"],
-    "resources_needed": ["List of required resources"]
-}
+        self.functions = {}
 
-Keep your responses focused and practical."""
-
-    def parse_json_response(self, response_text):
-        """Extract JSON from the response text."""
+    def validate_function_json(self, json_str):
+        """Validate function JSON format."""
         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
+            function_data = json.loads(json_str)
+            required_fields = ['name', 'description', 'parameters']
+            if not all(field in function_data for field in required_fields):
+                return False, "Missing required fields: name, description, or parameters"
+            
+            if not isinstance(function_data['parameters'], dict):
+                return False, "Parameters must be a JSON object"
+                
+            return True, function_data
+        except json.JSONDecodeError as e:
+            return False, f"Invalid JSON format: {str(e)}"
+        except Exception as e:
+            return False, f"Validation error: {str(e)}"
 
-    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 add_function(self, function_json):
+        """Add a new function definition."""
+        try:
+            is_valid, result = self.validate_function_json(function_json)
+            if not is_valid:
+                return f"Error: {result}"
+            
+            function_name = result['name']
+            self.functions[function_name] = result
+            return f"Successfully added function: {function_name}"
+        except Exception as e:
+            return f"Error adding function: {str(e)}\n{traceback.format_exc()}"
 
-    def generate_plan(self, task, max_tokens=1024, temperature=0.7, top_p=0.95):
+    def list_functions(self):
+        """List all registered functions."""
         try:
-            messages = [
-                {"role": "system", "content": self.system_prompt},
-                {"role": "user", "content": task}
-            ]
+            if not self.functions:
+                return "No functions registered."
+            
+            output = ["# Registered Functions"]
+            for name, func in self.functions.items():
+                output.append(f"\n## {name}")
+                output.append(f"Description: {func['description']}")
+                output.append("\nParameters:")
+                params = func['parameters']
+                if 'properties' in params:
+                    for param_name, param_info in params['properties'].items():
+                        output.append(f"- {param_name}: {param_info.get('description', 'No description')}")
+                        if param_info.get('type'):
+                            output.append(f"  Type: {param_info['type']}")
+                        if param_info.get('required', False):
+                            output.append("  Required: Yes")
+            
+            return "\n".join(output)
+        except Exception as e:
+            return f"Error listing functions: {str(e)}\n{traceback.format_exc()}"
 
-            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
+    def format_function_call(self, function_name, parameters):
+        """Format a function call into the HuggingChat format."""
+        try:
+            if function_name not in self.functions:
+                return False, f"Function '{function_name}' not found"
+            
+            function_def = self.functions[function_name]
+            function_call = {
+                "name": function_name,
+                "arguments": parameters
+            }
+            
+            return True, json.dumps(function_call, indent=2)
         except Exception as e:
-            yield f"Error generating plan: {str(e)}\n{traceback.format_exc()}"
+            return False, f"Error formatting function call: {str(e)}"
 
     def __del__(self):
         self.executor.shutdown(wait=False)
 
 def create_interface():
-    planner = AgentPlanner()
+    manager = FunctionManager()
     
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# AI Agent Planner")
-        gr.Markdown("Get organized, structured plans for any task")
-
-        task_input = gr.Textbox(
-            label="Task Description",
-            placeholder="Describe the task you need help planning...",
-            lines=3
-        )
-
-        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"
+        gr.Markdown("# HuggingChat Function Manager")
+        
+        with gr.Tab("Add Function"):
+            function_json = gr.Textbox(
+                label="Function Definition (JSON)",
+                placeholder='''{
+    "name": "example_function",
+    "description": "An example function",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "param1": {
+                "type": "string",
+                "description": "First parameter"
+            }
+        },
+        "required": ["param1"]
+    }
+}''',
+                lines=10
             )
-            top_p = gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.95,
-                step=0.05,
-                label="Top-p"
+            add_output = gr.Textbox(label="Result", lines=2)
+            add_btn = gr.Button("Add Function", variant="primary")
+            add_btn.click(
+                manager.add_function,
+                inputs=[function_json],
+                outputs=[add_output]
             )
 
-        output = gr.Markdown(label="Generated Plan")
+        with gr.Tab("List Functions"):
+            list_output = gr.Markdown(label="Registered Functions")
+            list_btn = gr.Button("List Functions")
+            list_btn.click(
+                manager.list_functions,
+                outputs=[list_output]
+            )
 
-        generate_btn = gr.Button("Generate Plan", variant="primary")
-        generate_btn.click(
-            planner.generate_plan,
-            inputs=[task_input, max_tokens, temperature, top_p],
-            outputs=output
-        )
+        with gr.Tab("Format Function Call"):
+            with gr.Row():
+                call_name = gr.Textbox(label="Function Name", lines=1)
+                call_params = gr.Textbox(
+                    label="Parameters (JSON)",
+                    placeholder='{"param1": "value1"}',
+                    lines=5
+                )
+            format_output = gr.Textbox(label="Formatted Call", lines=5)
+            format_btn = gr.Button("Format Call")
+            format_btn.click(
+                lambda n, p: manager.format_function_call(n, json.loads(p))[1],
+                inputs=[call_name, call_params],
+                outputs=[format_output]
+            )
 
     return demo