From 5da3395d0830527a86a00bd7cef60a0c25f344f4 Mon Sep 17 00:00:00 2001
From: yingying <yingying.xwy@antgroup.com>
Date: Mon, 8 Apr 2024 19:04:59 +0800
Subject: [PATCH] feat: add tools: search_issues, search_code

---
 server/agent/stream.py |  5 ++--
 server/tools/issue.py  | 53 ++++++++++++++++++++++++++++++++----------
 2 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/server/agent/stream.py b/server/agent/stream.py
index 55ef4e2e..2f257193 100644
--- a/server/agent/stream.py
+++ b/server/agent/stream.py
@@ -47,10 +47,9 @@ def get_datetime() -> datetime:
 TOOL_MAPPING = {
     "get_datetime": get_datetime,
     "create_issue": issue.create_issue,
-    "get_issues_list": issue.get_issues_list,
-    "get_issues_by_number": issue.get_issues_by_number
+    "search_issues": issue.search_issues,
 }
-TOOLS = ["get_datetime", "create_issue", "get_issues_list", "get_issues_by_number"]
+TOOLS = ["get_datetime", "create_issue", "search_issues"]
 
 
 def _create_agent_with_tools(openai_api_key: str ) -> AgentExecutor:
diff --git a/server/tools/issue.py b/server/tools/issue.py
index 08480487..683750ec 100644
--- a/server/tools/issue.py
+++ b/server/tools/issue.py
@@ -5,6 +5,8 @@
 
 GITHUB_TOKEN = os.getenv('GITHUB_TOKEN')
 
+DEFAULT_REPOSITORY = "ant-design/ant-design"
+
 g = Github(GITHUB_TOKEN)
 
 @tool
@@ -30,47 +32,74 @@ def create_issue(repo_name, title, body):
         return None
     
 @tool
-def get_issues_list(repo_name, state="open"):
+def get_issues(
+        max_num=5, 
+        repo_name=DEFAULT_REPOSITORY,
+        state="all",
+        sort="created",
+        order="desc",
+    ):
         """
         Fetches issues from the configured repository
         
+        :param max_num: The maximum number of issues to fetch
         :param repo_name: The name of the repository, e.g., "octocat/Hello-World"
-        :state: The state of the issue, e.g: open, closed, all
+        :param state: The state of the issue, e.g: open, closed, all
+        :param sort: The sorting method, e.g: created, updated, comments
+        :param order: The order of the sorting, e.g: asc, desc
         """
         try:
             # Obtain the repository object
             repo = g.get_repo(repo_name)
             
-            # Retrieve a list of open issues from the repository
-            open_issues = repo.get_issues(state=state)
+            # Retrieve a list of issues from the repository
+            issues = repo.get_issues(state=state, sort=sort, direction=order)[:max_num]
             
             issues_list = [
                 {
                     'issue_name': f'Issue #{issue.number} - {issue.title}',
                     'issue_url': issue.html_url
                 }
-                for issue in open_issues
+                for issue in issues
             ]
             return json.dumps(issues_list)
         except Exception as e:
             print(f"An error occurred: {e}")
             return json.dumps([])  
 
-@tool
-def get_issues_by_number(repo_name, number):
+def search_issues(
+        keyword, 
+        repo_name=DEFAULT_REPOSITORY,
+        max_num=5,
+        sort="created",
+        order="desc",
+    ):
         """
-        Match issue by the issue number
+        Fetches issues from the configured repository
         
+        :param keyword: The keyword to search for in the issues
         :param repo_name: The name of the repository, e.g., "octocat/Hello-World"
-        :number: The number of the issue
+        :param max_num: The maximum number of issues to fetch
+        :param sort: The sorting method, e.g: created, updated, comments
+        :param order: The order of the sorting, e.g: asc, desc
         """
         try:
             # Obtain the repository object
             repo = g.get_repo(repo_name)
             
-            issues = repo.get_issue(number=number)
-            print(f"issues: {issues}")
-            return issues
+            search_query = f'repo:{repo_name} {keyword} in:title,body,comments'
+
+            # Retrieve a list of open issues from the repository
+            issues = repo.search_issues(query=search_query, sort=sort, order=order)[:max_num]
+         
+            issues_list = [
+                {
+                    'issue_name': f'Issue #{issue.number} - {issue.title}',
+                    'issue_url': issue.html_url
+                }
+                for issue in issues
+            ]
+            return json.dumps(issues_list)
         except Exception as e:
             print(f"An error occurred: {e}")
             return json.dumps([])