Update app.py

app.py

@@ -45,158 +45,27 @@ class InferenceClient:
 
 class MyChatbot(gr.Chatbot):
     """Custom Chatbot class for enhanced functionality."""
-    def __init__(self, fn, **kwargs):
-        super().__init__(fn, **kwargs)
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
         self.issues = []  # Store fetched issues
         self.current_issue = None  # Store the currently selected issue
 
     def postprocess(self, y):
         """Post-processes the response to handle commands and display results."""
-        # Extract the response from the dictionary
-        assistant_response = y['assistant_response']
-
-        # Handle commands
-        if y['command'] == "/github":
-            if not y['github_api_token']:
-                return "Please enter your GitHub API token first."
-            else:
-                try:
-                    self.issues = fetch_github_issues(y['github_api_token'], y['github_username'], y['github_repository'])
-                    issue_list = "\n".join([f"{i+1}. {issue['title']}" for i, issue in enumerate(self.issues)])
-                    return f"Available GitHub Issues:\n{issue_list}\n\nEnter the issue number to analyze:"
-                except Exception as e:
-                    return f"Error fetching GitHub issues: {e}"
-        elif y['command'] == "/help":
-            return """Available commands:
-                - `/github`: Analyze a GitHub issue
-                - `/help`: Show this help message
-                - `/generate_code [code description]`: Generate code based on the description
-                - `/explain_concept [concept]`: Explain a concept
-                - `/write_documentation [topic]`: Write documentation for a given topic
-                - `/translate_code [code] to [target language]`: Translate code to another language"""
-        elif y['command'].isdigit() and self.issues:
-            try:
-                issue_number = int(y['command']) - 1
-                self.current_issue = self.issues[issue_number]  # Store the selected issue
-                issue_text = self.current_issue['title'] + "\n\n" + self.current_issue['body']
-                resolution = analyze_issues(issue_text, y['selected_model'], y['severity'], y['programming_language'])
-                related_issues = find_related_issues(issue_text, self.issues)
-                related_issue_text = "\n".join(
-                    [f"- {issue['title']} (Similarity: {similarity:.2f})" for issue, similarity in related_issues]
-                )
-                return f"Resolution for Issue '{self.current_issue['title']}':\n{resolution['assistant_response']}\n\nRelated Issues:\n{related_issue_text}"
-            except Exception as e:
-                return f"Error analyzing issue: {e}"
-        elif y['command'].startswith("/"):
-            # Handle commands like `/generate_code`, `/explain_concept`, etc.
-            if self.current_issue:
-                # Use the current issue's context for these commands
-                issue_text = self.current_issue['title'] + "\n\n" + self.current_issue['body']
-                return analyze_issues(issue_text, y['selected_model'], y['severity'], y['programming_language'])['assistant_response']
-            else:
-                return "Please select an issue first using `/github`."
+        if not y:
+            return []
+            
+        # Handle different types of input
+        if isinstance(y, dict):
+            assistant_response = y.get('assistant_response', '')
+            command = y.get('command', '')
         else:
-            # For free-form text, simply display the assistant's response
-            return assistant_response
-
-def analyze_issues(issue_text: str, model_name: str, severity: str = None, programming_language: str = None) -> str:
-    """Analyzes issues and provides solutions using a specified language model."""
-    model = pipeline("text-generation", model=model_name)
-    response = model(
-        f"{system_message}\n{issue_text}\nAssistant: ",
-        max_length=max_tokens,
-        do_sample=True,
-        temperature=temperature,
-        top_k=top_p,
-    )
-    assistant_response = response[0]['generated_text'].strip()
-
-    # Extract severity and programming language from the response
-    if "Severity" in assistant_response:
-        severity = assistant_response.split(":")[1].strip()
-    if "Programming Language" in assistant_response:
-        programming_language = assistant_response.split(":")[1].strip()
-
-    return {
-        'assistant_response': assistant_response,
-        'severity': severity,
-        'programming_language': programming_language,
-    }
+            assistant_response = str(y)
+            command = ''
 
-def find_related_issues(issue_text: str, issues: list) -> list:
-    """Finds semantically related issues from a list of issues based on the input issue text."""
-    similarity_model = SentenceTransformer('all-mpnet-base-v2')
-    issue_embedding = similarity_model.encode(issue_text)
-    similarities = [util.cos_sim(issue_embedding, similarity_model.encode(issue['title'])) for issue in issues]
-    sorted_issues = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)
-    related_issues = [issues[i] for i, similarity in sorted_issues[:5]]
-    return related_issues
-
-def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str) -> list:
-    """Fetches issues from a specified GitHub repository using the GitHub API."""
-    headers = {'Authorization': f'token {github_api_token}'}
-    url = f"https://api.github.com/repos/{github_username}/{github_repository}/issues"
-    response = requests.get(url, headers=headers)
-    issues = response.json()
-    return issues
-
-def push_to_repo(github_api_token: str, github_username: str, github_repository: str, commit_message: str, commit_file: str):
-    """Pushes changes to a GitHub repository."""
-    try:
-        repo = git.Repo.clone_from(f"https://github.com/{github_username}/{github_repository}.git", f"{github_repository}")
-        repo.git.add(commit_file)
-        repo.git.commit(m=commit_message)
-        repo.git.push()
-        return f"Changes pushed to {github_repository} successfully."
-    except gitdb.exc.InvalidGitRepositoryError:
-        return f"Invalid repository: {github_repository}"
-    except Exception as e:
-        return f"Error pushing to repository: {e}"
-
-def resolve_issue(github_api_token: str, github_username: str, github_repository: str, issue_number: int, resolution: str):
-    """Resolves an issue by pushing a commit with the resolution."""
-    try:
-        issue_text = fetch_github_issues(github_api_token, github_username, github_repository)[issue_number]['body']
-        commit_message = f"Resolved issue {issue_number}: {issue_text}"
-        commit_file = "resolution.txt"
-        with open(commit_file, "w") as f:
-            f.write(resolution)
-        return push_to_repo(github_api_token, github_username, github_repository, commit_message, commit_file)
-    except Exception as e:
-        return f"Error resolving issue: {e}"
-
-def respond(
-    command,
-    history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    github_api_token,
-    github_username,
-    github_repository,
-    selected_model,
-    severity,
-    programming_language,
-    *args,
-    **kwargs,
-) -> dict:
-    """Handles user commands and generates responses using the selected language model."""
-    model = pipeline("text-generation", model="OpenBMB/multilingual-codeparrot")
-    response = model(
-        f"{system_message}\n{command}\n{history}\n{github_username}/{github_repository}\n{severity}\n{programming_language}\nAssistant: ",
-        max_length=max_tokens,
-        do_sample=True,
-        temperature=temperature,
-        top_k=top_p,
-    )
-    assistant_response = response[0]['generated_text'].strip()
-    return {
-        'assistant_response': assistant_response,
-        'severity': severity,
-        'programming_language': programming_language,
-    }
+        return [(None, assistant_response)]
 
+# Update the Blocks implementation
 with gr.Blocks() as demo:
     with gr.Row():
         github_api_token = gr.Textbox(label="GitHub API Token", type="password")
@@ -226,45 +95,126 @@ with gr.Blocks() as demo:
 
     programming_language_textbox = gr.Textbox(label="Programming Language")
 
-    chatbot = MyChatbot(
-        respond,
-        inputs=[
-            gr.Textbox(label="Command"),
-            gr.Textbox(label="History"),
+    # Create the chatbot instance
+    chatbot = MyChatbot()
+    
+    # Create input textbox for user messages
+    msg = gr.Textbox(label="Message")
+    
+    # Create state for storing conversation history
+    state = gr.State([])
+
+    def user(user_message, history):
+        """Handle user messages"""
+        return "", history + [[user_message, None]]
+
+    def bot(history, system_msg, max_tokens, temp, top_p, api_token, username, repo, model, severity, lang):
+        """Generate bot response"""
+        if len(history) == 0:
+            return history
+            
+        user_message = history[-1][0]
+        response = respond(
+            command=user_message,
+            history=history[:-1],
+            system_message=system_msg,
+            max_tokens=max_tokens,
+            temperature=temp,
+            top_p=top_p,
+            github_api_token=api_token,
+            github_username=username,
+            github_repository=repo,
+            selected_model=model,
+            severity=severity,
+            programming_language=lang
+        )
+        
+        history[-1][1] = response['assistant_response']
+        return history
+
+    # Connect the components
+    msg.submit(
+        user,
+        [msg, chatbot],
+        [msg, chatbot]
+    ).then(
+        bot,
+        [
+            chatbot,
             system_message,
             gr.Slider(minimum=1, maximum=8192, value=2048, step=1, label="Max new tokens"),
             gr.Slider(minimum=0.1, maximum=4.0, value=0.71, step=0.1, label="Temperature"),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.95,
-                step=0.01,
-                label="Top-p (nucleus sampling)",
-            ),
+            gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.01, label="Top-p"),
             github_api_token,
             github_username,
             github_repository,
             model_dropdown,
             severity_dropdown,
-            programming_language_textbox,
+            programming_language_textbox
         ],
-        outputs=[gr.Textbox(label="Assistant Response")],
+        [chatbot]
     )
 
     # Add a button to fetch GitHub issues
     fetch_issues_button = gr.Button(label="Fetch Issues")
-    fetch_issues_button.click(fn=lambda github_api_token, github_username, github_repository: chatbot.issues, inputs=[github_api_token, github_username, github_repository], outputs=[chatbot])
+    
+    def fetch_issues(api_token, username, repo):
+        """Fetch GitHub issues and update the chatbot"""
+        issues = fetch_github_issues(api_token, username, repo)
+        chatbot.issues = issues
+        return gr.Dropdown(choices=[f"{i+1}. {issue['title']}" for i, issue in enumerate(issues)])
+
+    fetch_issues_button.click(
+        fetch_issues,
+        inputs=[github_api_token, github_username, github_repository],
+        outputs=[issue_dropdown]
+    )
 
     # Add a dropdown to select an issue
     issue_dropdown = gr.Dropdown(label="Select Issue", choices=[], interactive=True)
-    issue_dropdown.change(fn=lambda issue_number, chatbot: chatbot.postprocess(issue_number), inputs=[issue_dropdown, chatbot], outputs=[chatbot])
+    
+    def select_issue(selected, chatbot):
+        """Handle issue selection"""
+        if selected:
+            issue_num = int(selected.split('.')[0]) - 1
+            chatbot.current_issue = chatbot.issues[issue_num]
+            return chatbot.postprocess({
+                'command': str(issue_num + 1),
+                'github_api_token': github_api_token.value,
+                'github_username': github_username.value,
+                'github_repository': github_repository.value,
+                'selected_model': model_dropdown.value,
+                'severity': severity_dropdown.value,
+                'programming_language': programming_language_textbox.value
+            })
+        return []
+
+    issue_dropdown.change(
+        select_issue,
+        inputs=[issue_dropdown, chatbot],
+        outputs=[chatbot]
+    )
 
     # Add a button to resolve an issue
     resolve_issue_button = gr.Button(label="Resolve Issue")
-    resolve_issue_button.click(fn=lambda github_api_token, github_username, github_repository, issue_number, resolution: resolve_issue(github_api_token, github_username, github_repository, issue_number, resolution), inputs=[github_api_token, github_username, github_repository, gr.Number(label="Issue Number"), chatbot], outputs=[chatbot])
-
-    # Connect the chatbot input to the issue dropdown
-    chatbot.input.change(fn=lambda chatbot, github_api_token, github_username, github_repository: chatbot.postprocess("/github"), inputs=[chatbot, github_api_token, github_username, github_repository], outputs=[chatbot])
+    
+    def resolve_selected_issue(api_token, username, repo, resolution):
+        """Resolve the currently selected issue"""
+        if chatbot.current_issue:
+            issue_number = chatbot.issues.index(chatbot.current_issue)
+            return resolve_issue(api_token, username, repo, issue_number, resolution)
+        return "No issue selected"
+
+    resolve_issue_button.click(
+        resolve_selected_issue,
+        inputs=[
+            github_api_token,
+            github_username,
+            github_repository,
+            chatbot
+        ],
+        outputs=[chatbot]
+    )
 
 if __name__ == "__main__":
     demo.queue().launch(