add agent mode

anycoder_app/agent.py

@@ -0,0 +1,234 @@
+"""
+Agent functionality for interactive code generation with follow-up questions and task planning.
+"""
+import os
+from typing import Dict, List, Optional, Tuple, Generator
+import gradio as gr
+
+from .models import (
+    get_inference_client, get_real_model_id, history_to_messages,
+    history_to_chatbot_messages, strip_thinking_tags
+)
+from .deploy import generation_code
+
+
+def agent_generate_with_questions(
+    query: Optional[str],
+    setting: Dict[str, str],
+    history: List,
+    current_model: Dict,
+    language: str,
+    provider: str,
+    profile: Optional[gr.OAuthProfile] = None,
+    token: Optional[gr.OAuthToken] = None,
+    max_questions: int = 3
+) -> Generator[Tuple[List, List], None, None]:
+    """
+    Agent that asks follow-up questions, creates a task list, and generates code.
+    
+    Args:
+        query: Initial user request
+        setting: System settings
+        history: Conversation history
+        current_model: Selected model configuration
+        language: Target programming language/framework
+        provider: Model provider
+        profile: User OAuth profile
+        token: User OAuth token
+        max_questions: Maximum number of follow-up questions to ask
+        
+    Yields:
+        Tuples of (history, chatbot_messages) at each step
+    """
+    if not query or not query.strip():
+        return
+    
+    # Initialize history with user's initial query
+    current_history = history + [[query, ""]]
+    
+    # Step 1: Agent analyzes the request and asks follow-up questions
+    agent_system_prompt = """You are a helpful coding assistant that helps users clarify their requirements before generating code.
+
+Your task is to:
+1. Analyze the user's request
+2. Ask 1-3 clarifying questions to better understand their needs
+3. Focus on important details like:
+   - Target audience and use case
+   - Specific features or functionality needed
+   - Design preferences (colors, layout, style)
+   - Data sources or APIs to integrate
+   - Performance or scalability requirements
+
+Output ONLY the questions, numbered 1, 2, 3, etc. Keep questions concise and focused.
+Do not generate code yet - just ask the questions."""
+
+    # Get LLM client
+    client = get_inference_client(current_model.get('model_id', 'Qwen/Qwen2.5-Coder-32B-Instruct'), provider)
+    model_id = get_real_model_id(current_model.get('model_id', 'Qwen/Qwen2.5-Coder-32B-Instruct'))
+    
+    # Prepare messages for follow-up questions
+    messages = [
+        {'role': 'system', 'content': agent_system_prompt},
+        {'role': 'user', 'content': f"User wants to create: {query}\n\nLanguage/Framework: {language}\n\nAsk clarifying questions."}
+    ]
+    
+    # Generate follow-up questions
+    questions_response = ""
+    try:
+        # Try to use the client (works for both InferenceClient and OpenAI-compatible clients)
+        stream = client.chat.completions.create(
+            model=model_id,
+            messages=messages,
+            temperature=0.7,
+            max_tokens=500,
+            stream=True
+        )
+        for chunk in stream:
+            if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
+                questions_response += chunk.choices[0].delta.content
+                # Update display in real-time
+                temp_history = current_history[:-1] + [[query, f"🤔 **Analyzing your request...**\n\n{questions_response}"]]
+                yield (temp_history, history_to_chatbot_messages(temp_history))
+    except Exception as e:
+        error_msg = f"❌ Error asking follow-up questions: {str(e)}"
+        temp_history = current_history[:-1] + [[query, error_msg]]
+        yield (temp_history, history_to_chatbot_messages(temp_history))
+        return
+    
+    # Update history with agent's questions
+    current_history[-1][1] = f"🤔 **Let me ask you a few questions to better understand your needs:**\n\n{questions_response}\n\n💬 Please answer these questions in your next message."
+    yield (current_history, history_to_chatbot_messages(current_history))
+    
+    # Wait for user response (this will be handled by the UI)
+    # For now, we'll return and let the user respond, then continue in the next call
+    return
+
+
+def agent_process_answers_and_generate(
+    user_answers: str,
+    original_query: str,
+    questions: str,
+    setting: Dict[str, str],
+    history: List,
+    current_model: Dict,
+    language: str,
+    provider: str,
+    profile: Optional[gr.OAuthProfile] = None,
+    token: Optional[gr.OAuthToken] = None,
+    code_output=None,
+    history_output=None,
+    history_state=None
+) -> Generator:
+    """
+    Process user's answers, create task list, and generate code.
+    
+    Args:
+        user_answers: User's responses to the questions
+        original_query: Original user request
+        questions: Agent's questions
+        setting: System settings
+        history: Conversation history
+        current_model: Selected model configuration
+        language: Target programming language/framework
+        provider: Model provider
+        profile: User OAuth profile
+        token: User OAuth token
+        code_output: Code output component
+        history_output: History output component
+        history_state: History state
+        
+    Yields:
+        Updates to code output and history
+    """
+    # Step 2: Create task list based on answers
+    task_planning_prompt = f"""Based on the user's request and their answers, create a detailed task list for implementing the solution.
+
+Original Request: {original_query}
+
+Questions Asked:
+{questions}
+
+User's Answers:
+{user_answers}
+
+Create a numbered task list with 5-8 specific, actionable tasks. Each task should be clear and focused.
+Start with "📋 **Task List:**" and then list the tasks."""
+
+    client = get_inference_client(current_model.get('model_id', 'Qwen/Qwen2.5-Coder-32B-Instruct'), provider)
+    model_id = get_real_model_id(current_model.get('model_id', 'Qwen/Qwen2.5-Coder-32B-Instruct'))
+    
+    messages = [
+        {'role': 'system', 'content': 'You are a helpful coding assistant creating a task plan.'},
+        {'role': 'user', 'content': task_planning_prompt}
+    ]
+    
+    # Generate task list
+    task_list = ""
+    try:
+        stream = client.chat.completions.create(
+            model=model_id,
+            messages=messages,
+            temperature=0.7,
+            max_tokens=800,
+            stream=True
+        )
+        for chunk in stream:
+            if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
+                task_list += chunk.choices[0].delta.content
+                # Update display
+                temp_history = history + [[user_answers, f"📋 **Creating task list...**\n\n{task_list}"]]
+                yield {
+                    history_state: temp_history,
+                    history_output: history_to_chatbot_messages(temp_history)
+                }
+    except Exception as e:
+        error_msg = f"❌ Error creating task list: {str(e)}"
+        temp_history = history + [[user_answers, error_msg]]
+        yield {
+            history_state: temp_history,
+            history_output: history_to_chatbot_messages(temp_history)
+        }
+        return
+    
+    # Update history with task list
+    updated_history = history + [[user_answers, task_list]]
+    yield {
+        history_state: updated_history,
+        history_output: history_to_chatbot_messages(updated_history)
+    }
+    
+    # Step 3: Generate code based on refined requirements
+    refined_query = f"""{original_query}
+
+Additional Requirements (based on follow-up):
+{user_answers}
+
+Task List:
+{task_list}
+
+Please implement the above requirements following the task list."""
+    
+    # Add a message indicating code generation is starting
+    code_gen_start_history = updated_history + [["[System]", "🚀 **Starting code generation based on your requirements...**"]]
+    yield {
+        history_state: code_gen_start_history,
+        history_output: history_to_chatbot_messages(code_gen_start_history)
+    }
+    
+    # Use the existing generation_code function for actual code generation
+    # We need to pass the refined query and updated history
+    for result in generation_code(
+        refined_query,
+        setting,
+        updated_history,
+        current_model,
+        language,
+        provider,
+        profile,
+        token,
+        code_output,
+        history_output,
+        history_state
+    ):
+        yield result
+

anycoder_app/ui.py

@@ -29,6 +29,9 @@ from .deploy import (
     generate_requirements_txt_with_llm, prettify_comfyui_json_for_html,
     get_trending_models, import_model_from_hf, get_trending_spaces, import_space_from_hf
 )
+from .agent import (
+    agent_generate_with_questions, agent_process_answers_and_generate
+)
 
 # Main application with proper Gradio theming
 with gr.Blocks(
@@ -91,6 +94,12 @@ with gr.Blocks(
     last_login_state = gr.State(None)
     models_first_change = gr.State(True)
     spaces_first_change = gr.State(True)
+    agent_mode_enabled = gr.State(False)
+    agent_conversation_state = gr.State({
+        "stage": "initial",  # initial, waiting_for_answers, generating
+        "original_query": "",
+        "questions": ""
+    })
 
     with gr.Sidebar() as sidebar:
         login_button = gr.LoginButton()
@@ -166,6 +175,15 @@ with gr.Blocks(
             label="Code Language",
             visible=True
         )
+        
+        # Agent mode checkbox
+        agent_mode_checkbox = gr.Checkbox(
+            label="🤖 Enable Agent Mode",
+            value=False,
+            info="Agent will ask follow-up questions and create a task list before coding",
+            visible=True
+        )
+        
         # Removed image generation components
         with gr.Row():
             btn = gr.Button("Generate", variant="secondary", size="lg", scale=2, visible=True, interactive=False)
@@ -861,35 +879,107 @@ with gr.Blocks(
 
 
 
-    def begin_generation_ui():
-        # Collapse the sidebar when generation starts; keep status hidden
-        return [gr.update(open=False), gr.update(visible=False)]
+    def begin_generation_ui(agent_enabled):
+        # In agent mode, keep sidebar open during question/task planning phase
+        # Only close it when actual code generation starts
+        if agent_enabled:
+            return [gr.update(), gr.update(visible=False)]  # Keep sidebar as-is
+        else:
+            # Normal mode: collapse sidebar immediately
+            return [gr.update(open=False), gr.update(visible=False)]
 
     def end_generation_ui():
         # Open sidebar after generation; hide the status
         return [gr.update(open=True), gr.update(visible=False)]
+    
+    def close_sidebar_for_coding():
+        # Close sidebar when transitioning to actual code generation
+        return gr.update(open=False)
 
-    def generation_code_wrapper(inp, sett, hist, model, lang, prov, profile: Optional[gr.OAuthProfile] = None, token: Optional[gr.OAuthToken] = None):
-        """Wrapper to call generation_code and pass component references"""
-        # Generate code and update both history and chat_history
-        for result in generation_code(inp, sett, hist, model, lang, prov, profile, token, code_output, history_output, history):
-            # generation_code yields dictionaries with component keys
-            # Extract the values and yield them for our outputs
-            code_val = result.get(code_output, "")
-            hist_val = result.get(history, hist)
-            history_output_val = result.get(history_output, [])
-            # Yield for: code_output, history, history_output, chat_history
-            yield code_val, hist_val, history_output_val, history_output_val
-
+    def generation_code_wrapper(inp, sett, hist, model, lang, prov, agent_enabled, agent_state, profile: Optional[gr.OAuthProfile] = None, token: Optional[gr.OAuthToken] = None):
+        """Wrapper to call generation_code or agent mode based on settings"""
+        
+        # Check if agent mode is enabled
+        if agent_enabled and agent_state["stage"] == "initial":
+            # Agent mode - first interaction, ask questions
+            # Sidebar stays open during this phase
+            for updated_hist, chatbot_msgs in agent_generate_with_questions(
+                inp, sett, hist, model, lang, prov, profile, token
+            ):
+                # Update agent state to track that we're waiting for answers
+                new_agent_state = {
+                    "stage": "waiting_for_answers",
+                    "original_query": inp,
+                    "questions": updated_hist[-1][1] if updated_hist else ""
+                }
+                # Yield: code_output, history, history_output, chat_history, agent_conversation_state, sidebar
+                yield "", updated_hist, chatbot_msgs, chatbot_msgs, new_agent_state, gr.update()
+            return
+        
+        elif agent_enabled and agent_state["stage"] == "waiting_for_answers":
+            # Agent mode - user has answered questions, now create task list and generate code
+            original_query = agent_state.get("original_query", "")
+            questions = agent_state.get("questions", "")
+            
+            # Track whether we've started code generation to close sidebar
+            started_code_generation = False
+            
+            # Process answers and generate code
+            for result in agent_process_answers_and_generate(
+                inp, original_query, questions, sett, hist, model, lang, prov,
+                profile, token, code_output, history_output, history
+            ):
+                # Extract values from result dict
+                code_val = result.get(code_output, "")
+                hist_val = result.get(history, hist)
+                history_output_val = result.get(history_output, [])
+                
+                # Reset agent state after generation
+                reset_agent_state = {
+                    "stage": "initial",
+                    "original_query": "",
+                    "questions": ""
+                }
+                
+                # Close sidebar when we start generating code (when code_output has content)
+                if code_val and not started_code_generation:
+                    sidebar_update = gr.update(open=False)
+                    started_code_generation = True
+                else:
+                    sidebar_update = gr.update()
+                
+                # Yield: code_output, history, history_output, chat_history, agent_conversation_state, sidebar
+                yield code_val, hist_val, history_output_val, history_output_val, reset_agent_state, sidebar_update
+            return
+        
+        else:
+            # Normal mode - direct code generation
+            # Sidebar was already closed by begin_generation_ui
+            for result in generation_code(inp, sett, hist, model, lang, prov, profile, token, code_output, history_output, history):
+                # generation_code yields dictionaries with component keys
+                # Extract the values and yield them for our outputs
+                code_val = result.get(code_output, "")
+                hist_val = result.get(history, hist)
+                history_output_val = result.get(history_output, [])
+                # Yield for: code_output, history, history_output, chat_history, agent_conversation_state, sidebar
+                yield code_val, hist_val, history_output_val, history_output_val, agent_state, gr.update()
+
+    # Update agent_mode_enabled state when checkbox changes
+    agent_mode_checkbox.change(
+        lambda enabled: enabled,
+        inputs=[agent_mode_checkbox],
+        outputs=[agent_mode_enabled]
+    )
+    
     btn.click(
         begin_generation_ui,
-        inputs=None,
+        inputs=[agent_mode_enabled],
         outputs=[sidebar, generating_status],
         show_progress="hidden",
     ).then(
         generation_code_wrapper,
-        inputs=[input, setting, history, current_model, language_dropdown, provider_state],
-        outputs=[code_output, history, history_output, chat_history]
+        inputs=[input, setting, history, current_model, language_dropdown, provider_state, agent_mode_enabled, agent_conversation_state],
+        outputs=[code_output, history, history_output, chat_history, agent_conversation_state, sidebar]
     ).then(
         end_generation_ui,
         inputs=None,
@@ -931,13 +1021,13 @@ with gr.Blocks(
     # Pressing Enter in the main input should trigger generation and collapse the sidebar
     input.submit(
         begin_generation_ui,
-        inputs=None,
+        inputs=[agent_mode_enabled],
         outputs=[sidebar, generating_status],
         show_progress="hidden",
     ).then(
         generation_code_wrapper,
-        inputs=[input, setting, history, current_model, language_dropdown, provider_state],
-        outputs=[code_output, history, history_output, chat_history]
+        inputs=[input, setting, history, current_model, language_dropdown, provider_state, agent_mode_enabled, agent_conversation_state],
+        outputs=[code_output, history, history_output, chat_history, agent_conversation_state, sidebar]
     ).then(
         end_generation_ui,
         inputs=None,
@@ -1001,8 +1091,17 @@ with gr.Blocks(
         </div>
         """
     
+    def reset_agent_state():
+        """Reset agent conversation state when clearing history"""
+        return {
+            "stage": "initial",
+            "original_query": "",
+            "questions": ""
+        }
+    
     clear_btn.click(clear_history, outputs=[history, history_output, chat_history])
     clear_btn.click(hide_deploy_components, None, [deploy_btn])
+    clear_btn.click(reset_agent_state, outputs=[agent_conversation_state])
     # Reset button text when clearing
     clear_btn.click(
         lambda: gr.update(value="Publish"),