feat: Update Issue Label Generator to Issue Analysis Generator

Changelog:
- Changed class name to `IssueAnalysisGenerator`
- Add new dataclass for issue description
- Update prompt for issue label generation
- Add prompt for issue description generation

Signed-off-by: Maharshi Basu <[email protected]>

kaizen/generator/issue_analysis.py

@@ -0,0 +1,133 @@
+from typing import Optional, List, Dict, Generator
+import logging
+from dataclasses import dataclass
+import json
+
+from kaizen.helpers import output, parser
+from kaizen.llms.provider import LLMProvider
+from kaizen.llms.prompts.issue_analysis_prompts import (
+    ISSUE_LABEL_PROMPT,
+    ISSUE_DESC_PROMPT,
+    ISSUE_LABEL_SYSTEM_PROMPT,
+)
+
+@dataclass
+class IssueLabelOutput:
+    labels: List[str]
+    usage: Dict[str, int]
+    model_name: str
+    cost: Dict[str, float]
+
+@dataclass
+class IssueDescOutput:
+    desc: str
+    usage: Dict[str, int]
+    model_name: str
+    cost: Dict[str, float]
+
+class IssueAnalysisGenerator:
+    def __init__(self, llm_provider: LLMProvider):
+        self.logger = logging.getLogger(__name__)
+        self.provider = llm_provider
+        self.provider.system_prompt = ISSUE_LABEL_SYSTEM_PROMPT
+        self.total_usage = {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0,
+        }
+
+    def generate_issue_labels(
+        self,
+        issue_label_list: List[str],
+        issue_title: str,
+        issue_desc: str,
+        user: Optional[str] = None,
+    ) -> IssueLabelOutput:
+        prompt = ISSUE_LABEL_PROMPT.format(
+            ISSUE_LABEL_LIST=issue_label_list,
+            ISSUE_TITLE=issue_title,
+            ISSUE_DESCRIPTION=issue_desc,
+        )
+        if not issue_label_list and not issue_desc:
+            raise Exception("Issue labels missing for this repository. Create labels to ensure issue categorization.")
+
+        if issue_label_list and issue_desc and self.provider.is_inside_token_limit(PROMPT=prompt):
+            labels = self._process_issue_for_labels(
+                issue_title,
+                issue_desc,
+                user,
+            )
+
+        prompt_cost, completion_cost = self.provider.get_usage_cost(
+            total_usage=self.total_usage
+        )
+
+        return IssueLabelOutput(
+            labels=labels,
+            usage=self.total_usage,
+            model_name=self.provider.model,
+            cost={"prompt_cost": prompt_cost, "completion_cost": completion_cost}
+        )
+
+    def generate_issue_desc(
+        self,
+        issue_title,
+        issue_desc,
+        user: Optional[str] = None,
+    ) -> IssueDescOutput:
+        prompt = ISSUE_DESC_PROMPT.format(
+            ISSUE_TITLE=issue_title,
+            ISSUE_DESC=issue_desc
+        )
+        if not issue_desc:
+            raise Exception("Original issue description is empty!")
+
+        if issue_desc and self.provider.is_inside_token_limit(PROMPT=prompt):
+            desc = self._process_issue_for_desc(
+                issue_title,
+                issue_desc,
+                user,
+            )
+
+        body = output.create_issue_description(desc, issue_desc)
+        prompt_cost, completion_cost = self.provider.get_usage_cost(
+            total_usage=self.total_usage
+        )
+
+        return IssueDescOutput(
+            desc=body,
+            usage=self.total_usage,
+            model_name=self.provider.model,
+            cost={"prompt_cost": prompt_cost, "completion_cost": completion_cost},
+        )
+
+    # TODO: Convert `labels` to a format suitable for the github handler
+    def _process_issue_for_labels(
+        self,
+        prompt: str,
+        user: Optional[str],
+    ) -> List[str]:
+        self.logger.debug("Processing Issue for labels")
+        resp, usage = self.provider.chat_completion(
+            prompt=prompt,
+            user=user,
+        )
+        labels = parser.extract_code_from_markdown(resp)
+        self.total_usage = self.provider.update_usage(self.total_usage, usage)
+
+        return labels
+
+    def _process_issue_for_desc(
+        self,
+        prompt: str,
+        user: Optional[str]
+    ) -> str:
+        self.logger.debug("Processing issue for description")
+        resp, usage = self.provider.chat_completion(
+            prompt=prompt,
+            user=user,
+        )
+        desc = parser.extract_code_from_markdown(resp)
+        self.total_usage = self.provider.update_usage(self.total_usage, usage)
+
+        return desc

kaizen/helpers/output.py

@@ -42,6 +42,8 @@ def create_pr_description(desc, original_desc):
     markdown_output += "\n\n" + DESC_COLLAPSIBLE_TEMPLATE.format(desc=original_desc)
     return markdown_output
 
+def create_issue_description(desc, original_desc):
+    pass
 
 async def get_html(url):
     async with async_playwright() as p:

kaizen/llms/prompts/issue_analysis_prompts.py

@@ -0,0 +1,53 @@
+ISSUE_LABEL_SYSTEM_PROMPT = """
+As a senior project manager and a senior software developer reviewing raised GitHub issues, feature requests, and bug reports, make decisions by reviewing the issues and give constructive feedback and suggestions for improvement. Understand the nature of the issue and the problem being reported. Use your expertise to provide comprehensive feedback without asking clarifying questions.
+"""
+
+ISSUE_LABEL_PROMPT = """
+Understand the type of issue and choose labels from a given list to be added to the issue. The existing labels will be provided as a string in the following format:
+
+<label>: (Optional)<description>
+
+Output Format:
+```markdown
+- {{label_1}}
+- {{label_2}}
+...
+- {{label_n}}
+```
+
+Instructions:
+- Understand the labels from the list and their description.
+- Choose labels only from the provided list that best fit the issue.
+
+Based on:
+Title: {ISSUE_TITLE}
+Description: {ISSUE_DESCRIPTION}
+Exisiting labels: {ISSUE_LABEL_LIST}
+"""
+
+ISSUE_DESC_PROMPT="""
+Understand the issue and summarize the main purpose, scope of changes, modifications and approach required to tackle the issue. Also provide a brief summary of why the change proposed in the issue is necessary.
+
+Output Format:
+```markdown
+# {{Generated PR Title}}
+
+## Overview
+{{Brief summary of the issue}}
+
+## Suggestions
+- Primary actions: {{List the most critical steps to address the issue}}
+- Potential enhancements: {{List additional improvements or optimizations}}
+- Risk mitigation: {{List any risks and how to mitigate them}}
+```
+
+Instructions:
+- Create a concise summary of the issue's main purpose.
+- Use markdown formatting for readability.
+- Focus of significant suggestions and avoid repetition.
+
+Based on:
+Title: {ISSUE_TITLE}
+Description: {ISSUE_DESCRIPTION}
+"""
+