Skip to content

Commit c44d98e

Browse files
authored
feat: adds SearchGithubIssuesTool (#747)
# Motivation <!-- Why is this change necessary? --> # Content <!-- Please include a summary of the change --> # Testing <!-- How was the change tested? --> # Please check the following before marking your PR as ready for review - [ ] I have added tests for my changes - [ ] I have updated the documentation or added new documentation as needed
1 parent 26b5ca4 commit c44d98e

File tree

6 files changed

+200
-16
lines changed

6 files changed

+200
-16
lines changed

src/codegen/agents/code_agent.py

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,9 @@
88
from langsmith import Client
99

1010
from codegen.extensions.langchain.agent import create_codebase_agent
11-
from codegen.extensions.langchain.utils.get_langsmith_url import find_and_print_langsmith_run_url
11+
from codegen.extensions.langchain.utils.get_langsmith_url import (
12+
find_and_print_langsmith_run_url,
13+
)
1214

1315
if TYPE_CHECKING:
1416
from codegen import Codebase
@@ -17,6 +19,13 @@
1719
class CodeAgent:
1820
"""Agent for interacting with a codebase."""
1921

22+
codebase: "Codebase"
23+
agent: any
24+
langsmith_client: Client
25+
project_name: str
26+
thread_id: str | None = None
27+
config: dict = {}
28+
2029
def __init__(
2130
self,
2231
codebase: "Codebase",
@@ -43,7 +52,14 @@ def __init__(
4352
- max_tokens: Maximum number of tokens to generate
4453
"""
4554
self.codebase = codebase
46-
self.agent = create_codebase_agent(self.codebase, model_provider=model_provider, model_name=model_name, memory=memory, additional_tools=tools, **kwargs)
55+
self.agent = create_codebase_agent(
56+
self.codebase,
57+
model_provider=model_provider,
58+
model_name=model_name,
59+
memory=memory,
60+
additional_tools=tools,
61+
**kwargs,
62+
)
4763
self.langsmith_client = Client()
4864
self.run_id = run_id
4965
self.instance_id = instance_id
@@ -64,6 +80,14 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str:
6480
"""
6581
if thread_id is None:
6682
thread_id = str(uuid4())
83+
self.thread_id = thread_id
84+
self.config = {
85+
"configurable": {
86+
"thread_id": thread_id,
87+
"metadata": {"project": self.project_name},
88+
},
89+
"recursion_limit": 100,
90+
}
6791

6892
# this message has a reducer which appends the current message to the existing history
6993
# see more https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers
@@ -134,3 +158,9 @@ def get_agent_trace_url(self) -> str | None:
134158
print(traceback.format_exc())
135159
print(separator)
136160
return None
161+
162+
def get_tools(self) -> list[BaseTool]:
163+
return list(self.agent.get_graph().nodes["tools"].data.tools_by_name.values())
164+
165+
def get_state(self) -> dict:
166+
return self.agent.get_state(self.config)

src/codegen/extensions/langchain/tools.py

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
from codegen.extensions.linear.linear_client import LinearClient
99
from codegen.extensions.tools.bash import run_bash_command
10+
from codegen.extensions.tools.github.search import search
1011
from codegen.extensions.tools.linear.linear import (
1112
linear_comment_on_issue_tool,
1213
linear_create_issue_tool,
@@ -20,7 +21,6 @@
2021
from codegen.extensions.tools.relace_edit import relace_edit
2122
from codegen.extensions.tools.replacement_edit import replacement_edit
2223
from codegen.extensions.tools.reveal_symbol import reveal_symbol
23-
from codegen.extensions.tools.search import search
2424
from codegen.extensions.tools.semantic_edit import semantic_edit
2525
from codegen.extensions.tools.semantic_search import semantic_search
2626
from codegen.sdk.core.codebase import Codebase
@@ -560,6 +560,28 @@ def _run(self, title: str, body: str) -> str:
560560
return result.render()
561561

562562

563+
class GithubSearchIssuesInput(BaseModel):
564+
"""Input for searching GitHub issues."""
565+
566+
query: str = Field(..., description="Search query string to find issues")
567+
568+
569+
class GithubSearchIssuesTool(BaseTool):
570+
"""Tool for searching GitHub issues."""
571+
572+
name: ClassVar[str] = "search_issues"
573+
description: ClassVar[str] = "Search for GitHub issues/PRs using a query string from pygithub, e.g. 'is:pr is:open test_query'"
574+
args_schema: ClassVar[type[BaseModel]] = GithubSearchIssuesInput
575+
codebase: Codebase = Field(exclude=True)
576+
577+
def __init__(self, codebase: Codebase) -> None:
578+
super().__init__(codebase=codebase)
579+
580+
def _run(self, query: str) -> str:
581+
result = search(self.codebase, query)
582+
return result.render()
583+
584+
563585
class GithubViewPRInput(BaseModel):
564586
"""Input for getting PR contents."""
565587

@@ -856,6 +878,7 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]:
856878
GithubCreatePRCommentTool(codebase),
857879
GithubCreatePRReviewCommentTool(codebase),
858880
GithubViewPRTool(codebase),
881+
GithubSearchIssuesTool(codebase),
859882
# Linear
860883
LinearGetIssueTool(codebase),
861884
LinearGetIssueCommentsTool(codebase),
@@ -870,22 +893,28 @@ class ReplacementEditInput(BaseModel):
870893
filepath: str = Field(..., description="Path to the file to edit relative to the workspace root. The file must exist and be a text file.")
871894
pattern: str = Field(
872895
...,
873-
description="Regular expression pattern to match text that should be replaced. Supports all Python regex syntax including capture groups (\1, \2, etc). The pattern is compiled with re.MULTILINE flag by default.",
896+
description="""Regular expression pattern to match text that should be replaced.
897+
Supports all Python regex syntax including capture groups (\1, \2, etc). The pattern is compiled with re.MULTILINE flag by default.""",
874898
)
875899
replacement: str = Field(
876900
...,
877-
description="Text to replace matched patterns with. Can reference regex capture groups using \1, \2, etc. If using regex groups in pattern, make sure to preserve them in replacement if needed.",
901+
description="""Text to replace matched patterns with.
902+
Can reference regex capture groups using \1, \2, etc. If using regex groups in pattern, make sure to preserve them in replacement if needed.""",
878903
)
879904
start: int = Field(
880-
default=1, description="Starting line number (1-indexed, inclusive) to begin replacements from. Use this with 'end' to limit changes to a specific region. Default is 1 (start of file)."
905+
default=1,
906+
description="""Starting line number (1-indexed, inclusive) to begin replacements from.
907+
Use this with 'end' to limit changes to a specific region. Default is 1 (start of file).""",
881908
)
882909
end: int = Field(
883910
default=-1,
884-
description="Ending line number (1-indexed, inclusive) to stop replacements at. Use -1 to indicate end of file. Use this with 'start' to limit changes to a specific region. Default is -1 (end of file).",
911+
description="""Ending line number (1-indexed, inclusive) to stop replacements at.
912+
Use -1 to indicate end of file. Use this with 'start' to limit changes to a specific region. Default is -1 (end of file).""",
885913
)
886914
count: Optional[int] = Field(
887915
default=None,
888-
description="Maximum number of replacements to make. Use None to replace all occurrences (default), or specify a number to limit replacements. Useful when you only want to replace the first N occurrences.",
916+
description="""Maximum number of replacements to make. Use None to replace all occurrences (default), or specify a number to limit replacements.
917+
Useful when you only want to replace the first N occurrences.""",
889918
)
890919

891920

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
from .create_pr import create_pr
22
from .create_pr_comment import create_pr_comment
33
from .create_pr_review_comment import create_pr_review_comment
4+
from .search import search
45
from .view_pr import view_pr
56

67
__all__ = [
78
"create_pr",
89
"create_pr_comment",
910
"create_pr_review_comment",
11+
"search",
1012
"view_pr",
1113
]
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
"""Tools for searching GitHub issues and pull requests."""
2+
3+
from typing import ClassVar
4+
5+
from pydantic import Field
6+
7+
from codegen.sdk.core.codebase import Codebase
8+
9+
from ..observation import Observation
10+
11+
12+
class SearchResultObservation(Observation):
13+
"""Response from searching issues and pull requests."""
14+
15+
query: str = Field(
16+
description="The search query that was used",
17+
)
18+
results: list[dict] = Field(
19+
description="List of matching issues/PRs with their details. Use is:pr in query to search for PRs, is:issue for issues.",
20+
)
21+
22+
str_template: ClassVar[str] = "Found {total} results matching query: {query}"
23+
24+
@property
25+
def total(self) -> int:
26+
return len(self.results)
27+
28+
29+
def search(
30+
codebase: Codebase,
31+
query: str,
32+
max_results: int = 20,
33+
) -> SearchResultObservation:
34+
"""Search for GitHub issues and pull requests using the provided query.
35+
36+
To search for pull requests specifically, include 'is:pr' in your query.
37+
To search for issues specifically, include 'is:issue' in your query.
38+
If neither is specified, both issues and PRs will be included in results.
39+
40+
Args:
41+
codebase: The codebase to operate on
42+
query: Search query string (e.g. "is:pr label:bug", "is:issue is:open")
43+
state: Filter by state ("open", "closed", or "all")
44+
max_results: Maximum number of results to return
45+
"""
46+
try:
47+
# Get the GitHub repo object
48+
repo = codebase._op.remote_git_repo
49+
50+
# Search using PyGitHub's search_issues (which searches both issues and PRs)
51+
results = []
52+
for item in repo.search_issues(query)[:max_results]:
53+
result = {
54+
"title": item.title,
55+
"number": item.number,
56+
"state": item.state,
57+
"labels": [label.name for label in item.labels],
58+
"created_at": item.created_at.isoformat(),
59+
"updated_at": item.updated_at.isoformat(),
60+
"url": item.html_url,
61+
"is_pr": item.pull_request is not None,
62+
}
63+
results.append(result)
64+
65+
return SearchResultObservation(
66+
status="success",
67+
query=query,
68+
results=results,
69+
)
70+
71+
except Exception as e:
72+
return SearchResultObservation(
73+
status="error",
74+
error=f"Failed to search: {e!s}",
75+
query=query,
76+
results=[],
77+
)

src/codegen/git/clients/git_repo_client.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from github.Commit import Commit
88
from github.GithubException import GithubException, UnknownObjectException
99
from github.GithubObject import NotSet, Opt
10+
from github.Issue import Issue
1011
from github.IssueComment import IssueComment
1112
from github.Label import Label
1213
from github.PullRequest import PullRequest
@@ -431,3 +432,13 @@ def merge_upstream(self, branch_name: str) -> bool:
431432
post_parameters = {"branch": branch_name}
432433
status, _, _ = self.repo._requester.requestJson("POST", f"{self.repo.url}/merge-upstream", input=post_parameters)
433434
return status == 200
435+
436+
####################################################################################################################
437+
# SEARCH
438+
####################################################################################################################
439+
440+
def search_issues(self, query: str, **kwargs) -> list[Issue]:
441+
return self.gh_client.client.search_issues(query, **kwargs)
442+
443+
def search_prs(self, query: str, **kwargs) -> list[PullRequest]:
444+
return self.gh_client.client.search_issues(query, **kwargs)
Lines changed: 43 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,61 @@
1-
"""Tests for Linear tools."""
1+
"""Tests for GitHub tools."""
22

33
import os
44

55
import pytest
66

7-
from codegen.extensions.linear.linear_client import LinearClient
8-
from codegen.extensions.tools.github import view_pr
7+
from codegen.extensions.tools.github import search, view_pr
98
from codegen.sdk.core.codebase import Codebase
109

1110

1211
@pytest.fixture
13-
def client() -> LinearClient:
14-
"""Create a Linear client for testing."""
12+
def codebase() -> Codebase:
13+
"""Create a Codebase instance for testing."""
1514
token = os.getenv("GITHUB_TOKEN")
1615
if not token:
1716
pytest.skip("GITHUB_TOKEN environment variable not set")
1817
codebase = Codebase.from_repo("codegen-sh/Kevin-s-Adventure-Game")
1918
return codebase
2019

2120

22-
def test_github_view_pr(client: LinearClient) -> None:
23-
"""Test getting an issue from Linear."""
21+
def test_github_view_pr(codebase: Codebase) -> None:
22+
"""Test viewing a PR from GitHub."""
2423
# Link to PR: https://github.com/codegen-sh/Kevin-s-Adventure-Game/pull/419
25-
pr = view_pr(client, 419)
24+
pr = view_pr(codebase, 419)
2625
print(pr)
26+
27+
28+
def test_github_search_issues(codebase: Codebase) -> None:
29+
"""Test searching GitHub issues."""
30+
# Search for closed issues with the 'bug' label
31+
result = search(codebase, query="is:issue is:closed")
32+
assert result.status == "success"
33+
assert len(result.results) > 0
34+
assert "is:issue is:closed" in result.query
35+
36+
# Verify issue structure
37+
if result.results:
38+
issue = result.results[0]
39+
assert "title" in issue
40+
assert "number" in issue
41+
assert "state" in issue
42+
assert issue["state"] == "closed"
43+
assert not issue["is_pr"] # Should be an issue, not a PR
44+
45+
46+
def test_github_search_prs(codebase: Codebase) -> None:
47+
"""Test searching GitHub pull requests."""
48+
# Search for merged PRs
49+
result = search(codebase, query="is:pr is:merged")
50+
assert result.status == "success"
51+
assert len(result.results) > 0
52+
assert "is:pr is:merged" in result.query
53+
54+
# Verify PR structure
55+
if result.results:
56+
pr = result.results[0]
57+
assert "title" in pr
58+
assert "number" in pr
59+
assert "state" in pr
60+
assert pr["state"] == "closed"
61+
assert pr["is_pr"] # Should be a PR

0 commit comments

Comments
 (0)