From 8ef74461a525e57b411fa2b65caffbcc03ee4b7a Mon Sep 17 00:00:00 2001 From: jayhack Date: Fri, 14 Feb 2025 16:10:24 -0800 Subject: [PATCH 1/5] . --- src/codegen/extensions/langchain/tools.py | 51 +++++++++++++++++++ .../extensions/linear/linear_client.py | 50 ++++++++++++++++++ src/codegen/extensions/tools/linear_tools.py | 18 +++++++ 3 files changed, 119 insertions(+) diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index cbce0c1cc..797920958 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -10,8 +10,10 @@ from codegen.extensions.linear.linear_client import LinearClient from codegen.extensions.tools.linear_tools import ( linear_comment_on_issue_tool, + linear_create_issue_tool, linear_get_issue_comments_tool, linear_get_issue_tool, + linear_search_issues_tool, ) from ..tools import ( @@ -532,6 +534,53 @@ def _run(self, issue_id: str, body: str) -> str: return json.dumps(result, indent=2) +class LinearSearchIssuesInput(BaseModel): + """Input for searching Linear issues.""" + + query: str = Field(..., description="Search query string") + limit: int = Field(default=10, description="Maximum number of issues to return") + + +class LinearSearchIssuesTool(BaseTool): + """Tool for searching Linear issues.""" + + name: ClassVar[str] = "linear_search_issues" + description: ClassVar[str] = "Search for Linear issues using a query string" + args_schema: ClassVar[type[BaseModel]] = LinearSearchIssuesInput + client: LinearClient = Field(exclude=True) + + def __init__(self, client: LinearClient) -> None: + super().__init__(client=client) + + def _run(self, query: str, limit: int = 10) -> str: + result = linear_search_issues_tool(self.client, query, limit) + return json.dumps(result, indent=2) + + +class LinearCreateIssueInput(BaseModel): + """Input for creating a Linear issue.""" + + team_id: str = Field(..., description="ID of the team to create the issue in") + title: str = Field(..., description="Title of the issue") + description: str | None = Field(None, description="Optional description of the issue") + + +class LinearCreateIssueTool(BaseTool): + """Tool for creating Linear issues.""" + + name: ClassVar[str] = "linear_create_issue" + description: ClassVar[str] = "Create a new Linear issue" + args_schema: ClassVar[type[BaseModel]] = LinearCreateIssueInput + client: LinearClient = Field(exclude=True) + + def __init__(self, client: LinearClient) -> None: + super().__init__(client=client) + + def _run(self, team_id: str, title: str, description: str | None = None) -> str: + result = linear_create_issue_tool(self.client, team_id, title, description) + return json.dumps(result, indent=2) + + def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: """Get all workspace tools initialized with a codebase. @@ -561,4 +610,6 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: LinearGetIssueTool(codebase), LinearGetIssueCommentsTool(codebase), LinearCommentOnIssueTool(codebase), + LinearSearchIssuesTool(codebase), + LinearCreateIssueTool(codebase), ] diff --git a/src/codegen/extensions/linear/linear_client.py b/src/codegen/extensions/linear/linear_client.py index ef0e1b96a..ae7adea7a 100644 --- a/src/codegen/extensions/linear/linear_client.py +++ b/src/codegen/extensions/linear/linear_client.py @@ -194,3 +194,53 @@ def search_issues(self, query: str, limit: int = 10) -> list[LinearIssue]: except Exception as e: msg = f"Error searching issues\n{data}\n{e}" raise Exception(msg) + + def create_issue(self, team_id: str, title: str, description: str | None = None) -> LinearIssue: + """Create a new issue. + + Args: + team_id: ID of the team to create the issue in + title: Title of the issue + description: Optional description of the issue + + Returns: + The created LinearIssue object + """ + mutation = """ + mutation createIssue($input: IssueCreateInput!) { + issueCreate(input: $input) { + success + issue { + id + title + description + } + } + } + """ + + variables = { + "input": { + "teamId": team_id, + "title": title, + "description": description, + } + } + + response = requests.post( + self.api_endpoint, + headers=self.api_headers, + json={"query": mutation, "variables": variables}, + ) + data = response.json() + + try: + issue_data = data["data"]["issueCreate"]["issue"] + return LinearIssue( + id=issue_data["id"], + title=issue_data["title"], + description=issue_data["description"], + ) + except Exception as e: + msg = f"Error creating issue\n{data}\n{e}" + raise Exception(msg) diff --git a/src/codegen/extensions/tools/linear_tools.py b/src/codegen/extensions/tools/linear_tools.py index 9bc02f5fe..c6c5a6ddc 100644 --- a/src/codegen/extensions/tools/linear_tools.py +++ b/src/codegen/extensions/tools/linear_tools.py @@ -37,3 +37,21 @@ def linear_register_webhook_tool(client: LinearClient, webhook_url: str, team_id return {"status": "success", "response": response} except Exception as e: return {"error": f"Failed to register webhook: {e!s}"} + + +def linear_search_issues_tool(client: LinearClient, query: str, limit: int = 10) -> dict[str, Any]: + """Search for issues using a query string.""" + try: + issues = client.search_issues(query, limit) + return {"status": "success", "issues": [issue.dict() for issue in issues]} + except Exception as e: + return {"error": f"Failed to search issues: {e!s}"} + + +def linear_create_issue_tool(client: LinearClient, team_id: str, title: str, description: str | None = None) -> dict[str, Any]: + """Create a new issue.""" + try: + issue = client.create_issue(team_id, title, description) + return {"status": "success", "issue": issue.dict()} + except Exception as e: + return {"error": f"Failed to create issue: {e!s}"} From e1c74a7096f7535d1e3e5d0e733dd242c9f877c2 Mon Sep 17 00:00:00 2001 From: jayhack Date: Fri, 14 Feb 2025 16:42:21 -0800 Subject: [PATCH 2/5] . --- src/codegen/extensions/langchain/tools.py | 23 ++++++- .../extensions/linear/linear_client.py | 67 ++++++++++++++++++- src/codegen/extensions/tools/linear_tools.py | 13 +++- tests/integration/extension/test_linear.py | 49 +++++++++++++- 4 files changed, 142 insertions(+), 10 deletions(-) diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index 797920958..4007f28fa 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -13,6 +13,7 @@ linear_create_issue_tool, linear_get_issue_comments_tool, linear_get_issue_tool, + linear_get_teams_tool, linear_search_issues_tool, ) @@ -560,9 +561,9 @@ def _run(self, query: str, limit: int = 10) -> str: class LinearCreateIssueInput(BaseModel): """Input for creating a Linear issue.""" - team_id: str = Field(..., description="ID of the team to create the issue in") title: str = Field(..., description="Title of the issue") description: str | None = Field(None, description="Optional description of the issue") + team_id: str | None = Field(None, description="Optional team ID. If not provided, uses the client's configured team_id") class LinearCreateIssueTool(BaseTool): @@ -576,8 +577,23 @@ class LinearCreateIssueTool(BaseTool): def __init__(self, client: LinearClient) -> None: super().__init__(client=client) - def _run(self, team_id: str, title: str, description: str | None = None) -> str: - result = linear_create_issue_tool(self.client, team_id, title, description) + def _run(self, title: str, description: str | None = None, team_id: str | None = None) -> str: + result = linear_create_issue_tool(self.client, title, description, team_id) + return json.dumps(result, indent=2) + + +class LinearGetTeamsTool(BaseTool): + """Tool for getting Linear teams.""" + + name: ClassVar[str] = "linear_get_teams" + description: ClassVar[str] = "Get all Linear teams the authenticated user has access to" + client: LinearClient = Field(exclude=True) + + def __init__(self, client: LinearClient) -> None: + super().__init__(client=client) + + def _run(self) -> str: + result = linear_get_teams_tool(self.client) return json.dumps(result, indent=2) @@ -612,4 +628,5 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: LinearCommentOnIssueTool(codebase), LinearSearchIssuesTool(codebase), LinearCreateIssueTool(codebase), + LinearGetTeamsTool(codebase), ] diff --git a/src/codegen/extensions/linear/linear_client.py b/src/codegen/extensions/linear/linear_client.py index ae7adea7a..e27867dca 100644 --- a/src/codegen/extensions/linear/linear_client.py +++ b/src/codegen/extensions/linear/linear_client.py @@ -17,6 +17,14 @@ class LinearUser(BaseModel): name: str +class LinearTeam(BaseModel): + """Represents a Linear team.""" + + id: str + name: str + key: str + + class LinearComment(BaseModel): id: str body: str @@ -33,13 +41,18 @@ class LinearClient: api_headers: dict api_endpoint = "https://api.linear.app/graphql" - def __init__(self, access_token: Optional[str] = None): + def __init__(self, access_token: Optional[str] = None, team_id: Optional[str] = None): if not access_token: access_token = os.getenv("LINEAR_ACCESS_TOKEN") if not access_token: msg = "access_token is required" raise ValueError(msg) self.access_token = access_token + + if not team_id: + team_id = os.getenv("LINEAR_TEAM_ID") + self.team_id = team_id + self.api_headers = { "Content-Type": "application/json", "Authorization": self.access_token, @@ -195,17 +208,26 @@ def search_issues(self, query: str, limit: int = 10) -> list[LinearIssue]: msg = f"Error searching issues\n{data}\n{e}" raise Exception(msg) - def create_issue(self, team_id: str, title: str, description: str | None = None) -> LinearIssue: + def create_issue(self, title: str, description: str | None = None, team_id: str | None = None) -> LinearIssue: """Create a new issue. Args: - team_id: ID of the team to create the issue in title: Title of the issue description: Optional description of the issue + team_id: Optional team ID. If not provided, uses the client's configured team_id Returns: The created LinearIssue object + + Raises: + ValueError: If no team_id is provided or configured """ + if not team_id: + team_id = self.team_id + if not team_id: + msg = "team_id must be provided either during client initialization or in the create_issue call" + raise ValueError(msg) + mutation = """ mutation createIssue($input: IssueCreateInput!) { issueCreate(input: $input) { @@ -244,3 +266,42 @@ def create_issue(self, team_id: str, title: str, description: str | None = None) except Exception as e: msg = f"Error creating issue\n{data}\n{e}" raise Exception(msg) + + def get_teams(self) -> list[LinearTeam]: + """Get all teams the authenticated user has access to. + + Returns: + List of LinearTeam objects + """ + query = """ + query { + teams { + nodes { + id + name + key + } + } + } + """ + + response = requests.post( + self.api_endpoint, + headers=self.api_headers, + json={"query": query}, + ) + data = response.json() + + try: + teams_data = data["data"]["teams"]["nodes"] + return [ + LinearTeam( + id=team["id"], + name=team["name"], + key=team["key"], + ) + for team in teams_data + ] + except Exception as e: + msg = f"Error getting teams\n{data}\n{e}" + raise Exception(msg) diff --git a/src/codegen/extensions/tools/linear_tools.py b/src/codegen/extensions/tools/linear_tools.py index c6c5a6ddc..c15cb5d10 100644 --- a/src/codegen/extensions/tools/linear_tools.py +++ b/src/codegen/extensions/tools/linear_tools.py @@ -48,10 +48,19 @@ def linear_search_issues_tool(client: LinearClient, query: str, limit: int = 10) return {"error": f"Failed to search issues: {e!s}"} -def linear_create_issue_tool(client: LinearClient, team_id: str, title: str, description: str | None = None) -> dict[str, Any]: +def linear_create_issue_tool(client: LinearClient, title: str, description: str | None = None, team_id: str | None = None) -> dict[str, Any]: """Create a new issue.""" try: - issue = client.create_issue(team_id, title, description) + issue = client.create_issue(title, description, team_id) return {"status": "success", "issue": issue.dict()} except Exception as e: return {"error": f"Failed to create issue: {e!s}"} + + +def linear_get_teams_tool(client: LinearClient) -> dict[str, Any]: + """Get all teams the authenticated user has access to.""" + try: + teams = client.get_teams() + return {"status": "success", "teams": [team.dict() for team in teams]} + except Exception as e: + return {"error": f"Failed to get teams: {e!s}"} diff --git a/tests/integration/extension/test_linear.py b/tests/integration/extension/test_linear.py index c6eb3606d..5a758f2cb 100644 --- a/tests/integration/extension/test_linear.py +++ b/tests/integration/extension/test_linear.py @@ -7,8 +7,11 @@ from codegen.extensions.linear.linear_client import LinearClient from codegen.extensions.tools.linear_tools import ( linear_comment_on_issue_tool, + linear_create_issue_tool, linear_get_issue_comments_tool, linear_get_issue_tool, + linear_get_teams_tool, + linear_search_issues_tool, ) @@ -18,7 +21,10 @@ def client() -> LinearClient: token = os.getenv("LINEAR_ACCESS_TOKEN") if not token: pytest.skip("LINEAR_ACCESS_TOKEN environment variable not set") - return LinearClient(token) + team_id = os.getenv("LINEAR_TEAM_ID") + if not team_id: + pytest.skip("LINEAR_TEAM_ID environment variable not set") + return LinearClient(token, team_id) def test_linear_get_issue(client: LinearClient) -> None: @@ -45,6 +51,45 @@ def test_linear_comment_on_issue(client: LinearClient) -> None: def test_search_issues(client: LinearClient) -> None: """Test searching for issues in Linear.""" - issues = client.search_issues("REVEAL_SYMBOL") + issues = linear_search_issues_tool(client, "REVEAL_SYMBOL") assert issues["status"] == "success" assert len(issues["issues"]) > 0 + + +def test_create_issue(client: LinearClient) -> None: + """Test creating an issue in Linear.""" + # Test creating an issue with explicit team_id + title = "Test Issue - Automated Testing (Explicit Team)" + description = "This is a test issue created by automated testing with explicit team_id" + + issue = client.create_issue(title, description) + assert issue.title == title + assert issue.description == description + + # Test creating an issue using default team_id from environment + title2 = "Test Issue - Automated Testing (Default Team)" + description2 = "This is a test issue created by automated testing with default team_id" + + issue2 = client.create_issue(title2, description2) + assert issue2.title == title2 + assert issue2.description == description2 + + # Test the tool wrapper with default team_id + result = linear_create_issue_tool(client, "Test Tool Issue", "Test description from tool") + assert result["status"] == "success" + assert result["issue"]["title"] == "Test Tool Issue" + assert result["issue"]["description"] == "Test description from tool" + + +def test_get_teams(client: LinearClient) -> None: + """Test getting teams from Linear.""" + result = linear_get_teams_tool(client) + assert result["status"] == "success" + assert len(result["teams"]) > 0 + + # Verify team structure + team = result["teams"][0] + print(result) + assert "id" in team + assert "name" in team + assert "key" in team From f88a002de6623d348ed73cbd91faa74a1ff8c441 Mon Sep 17 00:00:00 2001 From: jayhack Date: Fri, 14 Feb 2025 16:44:57 -0800 Subject: [PATCH 3/5] . --- src/codegen/extensions/langchain/tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index 4007f28fa..8ab951b2d 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -563,7 +563,7 @@ class LinearCreateIssueInput(BaseModel): title: str = Field(..., description="Title of the issue") description: str | None = Field(None, description="Optional description of the issue") - team_id: str | None = Field(None, description="Optional team ID. If not provided, uses the client's configured team_id") + team_id: str | None = Field(None, description="Optional team ID. If not provided, uses the default team_id (recommended)") class LinearCreateIssueTool(BaseTool): From 2109e714629b1516371c1a71c9a6835d8ad3c7e5 Mon Sep 17 00:00:00 2001 From: jayhack Date: Fri, 14 Feb 2025 19:16:15 -0800 Subject: [PATCH 4/5] . --- src/codegen/extensions/langchain/tools.py | 11 +++++++---- src/codegen/extensions/tools/linear/__init__.py | 0 .../tools/{linear_tools.py => linear/linear.py} | 0 tests/integration/__init__.py | 0 tests/integration/extension/__init__.py | 0 tests/integration/extension/test_github.py | 0 tests/integration/extension/test_linear.py | 2 +- 7 files changed, 8 insertions(+), 5 deletions(-) create mode 100644 src/codegen/extensions/tools/linear/__init__.py rename src/codegen/extensions/tools/{linear_tools.py => linear/linear.py} (100%) create mode 100644 tests/integration/__init__.py create mode 100644 tests/integration/extension/__init__.py create mode 100644 tests/integration/extension/test_github.py diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index 8ab951b2d..a48310593 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -8,7 +8,7 @@ from codegen import Codebase from codegen.extensions.linear.linear_client import LinearClient -from codegen.extensions.tools.linear_tools import ( +from codegen.extensions.tools.linear.linear import ( linear_comment_on_issue_tool, linear_create_issue_tool, linear_get_issue_comments_tool, @@ -609,9 +609,6 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: return [ CommitTool(codebase), CreateFileTool(codebase), - CreatePRTool(codebase), - CreatePRCommentTool(codebase), - CreatePRReviewCommentTool(codebase), DeleteFileTool(codebase), EditFileTool(codebase), GetPRcontentsTool(codebase), @@ -623,6 +620,12 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: SemanticEditTool(codebase), SemanticSearchTool(codebase), ViewFileTool(codebase), + # Github + CreatePRTool(codebase), + CreatePRCommentTool(codebase), + CreatePRReviewCommentTool(codebase), + GetPRcontentsTool(codebase), + # Linear LinearGetIssueTool(codebase), LinearGetIssueCommentsTool(codebase), LinearCommentOnIssueTool(codebase), diff --git a/src/codegen/extensions/tools/linear/__init__.py b/src/codegen/extensions/tools/linear/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/codegen/extensions/tools/linear_tools.py b/src/codegen/extensions/tools/linear/linear.py similarity index 100% rename from src/codegen/extensions/tools/linear_tools.py rename to src/codegen/extensions/tools/linear/linear.py diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integration/extension/__init__.py b/tests/integration/extension/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integration/extension/test_github.py b/tests/integration/extension/test_github.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integration/extension/test_linear.py b/tests/integration/extension/test_linear.py index 5a758f2cb..a962fa366 100644 --- a/tests/integration/extension/test_linear.py +++ b/tests/integration/extension/test_linear.py @@ -5,7 +5,7 @@ import pytest from codegen.extensions.linear.linear_client import LinearClient -from codegen.extensions.tools.linear_tools import ( +from codegen.extensions.tools.linear import ( linear_comment_on_issue_tool, linear_create_issue_tool, linear_get_issue_comments_tool, From 11e8740fc22e6813a9ecbb1a7d246fa43e02cc27 Mon Sep 17 00:00:00 2001 From: jayhack Date: Fri, 14 Feb 2025 19:30:25 -0800 Subject: [PATCH 5/5] . --- src/codegen/extensions/langchain/agent.py | 16 +++--- src/codegen/extensions/langchain/tools.py | 53 ++++++++++++------- src/codegen/extensions/tools/__init__.py | 2 +- .../extensions/tools/github/__init__.py | 11 ++++ .../extensions/tools/github/view_pr.py | 2 +- .../extensions/tools/linear/__init__.py | 19 +++++++ tests/integration/extension/test_github.py | 26 +++++++++ tests/integration/extension/test_linear.py | 2 +- 8 files changed, 101 insertions(+), 30 deletions(-) create mode 100644 src/codegen/extensions/tools/github/__init__.py diff --git a/src/codegen/extensions/langchain/agent.py b/src/codegen/extensions/langchain/agent.py index 8909e02d7..22964c5ac 100644 --- a/src/codegen/extensions/langchain/agent.py +++ b/src/codegen/extensions/langchain/agent.py @@ -13,12 +13,12 @@ from .tools import ( CommitTool, CreateFileTool, - CreatePRCommentTool, - CreatePRReviewCommentTool, - CreatePRTool, DeleteFileTool, EditFileTool, - GetPRcontentsTool, + GithubCreatePRCommentTool, + GithubCreatePRReviewCommentTool, + GithubCreatePRTool, + GithubViewPRTool, ListDirectoryTool, MoveSymbolTool, RenameFileTool, @@ -68,10 +68,10 @@ def create_codebase_agent( SemanticEditTool(codebase), SemanticSearchTool(codebase), CommitTool(codebase), - CreatePRTool(codebase), - GetPRcontentsTool(codebase), - CreatePRCommentTool(codebase), - CreatePRReviewCommentTool(codebase), + GithubCreatePRTool(codebase), + GithubViewPRTool(codebase), + GithubCreatePRCommentTool(codebase), + GithubCreatePRReviewCommentTool(codebase), ] # Get the prompt to use diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index a48310593..2aa76f25d 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -354,19 +354,24 @@ def _run(self, query: str, k: int = 5, preview_length: int = 200) -> str: return json.dumps(result, indent=2) -class CreatePRInput(BaseModel): +######################################################################################################################## +# GITHUB +######################################################################################################################## + + +class GithubCreatePRInput(BaseModel): """Input for creating a PR""" title: str = Field(..., description="The title of the PR") body: str = Field(..., description="The body of the PR") -class CreatePRTool(BaseTool): +class GithubCreatePRTool(BaseTool): """Tool for creating a PR.""" name: ClassVar[str] = "create_pr" description: ClassVar[str] = "Create a PR for the current branch" - args_schema: ClassVar[type[BaseModel]] = CreatePRInput + args_schema: ClassVar[type[BaseModel]] = GithubCreatePRInput codebase: Codebase = Field(exclude=True) def __init__(self, codebase: Codebase) -> None: @@ -377,18 +382,18 @@ def _run(self, title: str, body: str) -> str: return json.dumps(result, indent=2) -class GetPRContentsInput(BaseModel): +class GithubViewPRInput(BaseModel): """Input for getting PR contents.""" pr_id: int = Field(..., description="Number of the PR to get the contents for") -class GetPRcontentsTool(BaseTool): +class GithubViewPRTool(BaseTool): """Tool for getting PR data.""" - name: ClassVar[str] = "get_pr_contents" - description: ClassVar[str] = "Get the diff and modified symbols of a PR along with the dependencies of the modified symbols" - args_schema: ClassVar[type[BaseModel]] = GetPRContentsInput + name: ClassVar[str] = "view_pr" + description: ClassVar[str] = "View the diff and associated context for a pull request" + args_schema: ClassVar[type[BaseModel]] = GithubViewPRInput codebase: Codebase = Field(exclude=True) def __init__(self, codebase: Codebase) -> None: @@ -399,19 +404,19 @@ def _run(self, pr_id: int) -> str: return json.dumps(result, indent=2) -class CreatePRCommentInput(BaseModel): +class GithubCreatePRCommentInput(BaseModel): """Input for creating a PR comment""" pr_number: int = Field(..., description="The PR number to comment on") body: str = Field(..., description="The comment text") -class CreatePRCommentTool(BaseTool): +class GithubCreatePRCommentTool(BaseTool): """Tool for creating a general PR comment.""" name: ClassVar[str] = "create_pr_comment" description: ClassVar[str] = "Create a general comment on a pull request" - args_schema: ClassVar[type[BaseModel]] = CreatePRCommentInput + args_schema: ClassVar[type[BaseModel]] = GithubCreatePRCommentInput codebase: Codebase = Field(exclude=True) def __init__(self, codebase: Codebase) -> None: @@ -422,7 +427,7 @@ def _run(self, pr_number: int, body: str) -> str: return json.dumps(result, indent=2) -class CreatePRReviewCommentInput(BaseModel): +class GithubCreatePRReviewCommentInput(BaseModel): """Input for creating an inline PR review comment""" pr_number: int = Field(..., description="The PR number to comment on") @@ -434,12 +439,12 @@ class CreatePRReviewCommentInput(BaseModel): start_line: int | None = Field(None, description="For multi-line comments, the starting line") -class CreatePRReviewCommentTool(BaseTool): +class GithubCreatePRReviewCommentTool(BaseTool): """Tool for creating inline PR review comments.""" name: ClassVar[str] = "create_pr_review_comment" description: ClassVar[str] = "Create an inline review comment on a specific line in a pull request" - args_schema: ClassVar[type[BaseModel]] = CreatePRReviewCommentInput + args_schema: ClassVar[type[BaseModel]] = GithubCreatePRReviewCommentInput codebase: Codebase = Field(exclude=True) def __init__(self, codebase: Codebase) -> None: @@ -468,6 +473,11 @@ def _run( return json.dumps(result, indent=2) +######################################################################################################################## +# LINEAR +######################################################################################################################## + + class LinearGetIssueInput(BaseModel): """Input for getting a Linear issue.""" @@ -597,6 +607,11 @@ def _run(self) -> str: return json.dumps(result, indent=2) +######################################################################################################################## +# EXPORT +######################################################################################################################## + + def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: """Get all workspace tools initialized with a codebase. @@ -611,7 +626,7 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: CreateFileTool(codebase), DeleteFileTool(codebase), EditFileTool(codebase), - GetPRcontentsTool(codebase), + GithubViewPRTool(codebase), ListDirectoryTool(codebase), MoveSymbolTool(codebase), RenameFileTool(codebase), @@ -621,10 +636,10 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: SemanticSearchTool(codebase), ViewFileTool(codebase), # Github - CreatePRTool(codebase), - CreatePRCommentTool(codebase), - CreatePRReviewCommentTool(codebase), - GetPRcontentsTool(codebase), + GithubCreatePRTool(codebase), + GithubCreatePRCommentTool(codebase), + GithubCreatePRReviewCommentTool(codebase), + GithubViewPRTool(codebase), # Linear LinearGetIssueTool(codebase), LinearGetIssueCommentsTool(codebase), diff --git a/src/codegen/extensions/tools/__init__.py b/src/codegen/extensions/tools/__init__.py index 06ededaf8..70375f9b8 100644 --- a/src/codegen/extensions/tools/__init__.py +++ b/src/codegen/extensions/tools/__init__.py @@ -8,7 +8,7 @@ from .github.create_pr_comment import create_pr_comment from .github.create_pr_review_comment import create_pr_review_comment from .github.view_pr import view_pr -from .linear_tools import ( +from .linear import ( linear_comment_on_issue_tool, linear_get_issue_comments_tool, linear_get_issue_tool, diff --git a/src/codegen/extensions/tools/github/__init__.py b/src/codegen/extensions/tools/github/__init__.py new file mode 100644 index 000000000..a59669dd2 --- /dev/null +++ b/src/codegen/extensions/tools/github/__init__.py @@ -0,0 +1,11 @@ +from .create_pr import create_pr +from .create_pr_comment import create_pr_comment +from .create_pr_review_comment import create_pr_review_comment +from .view_pr import view_pr + +__all__ = [ + "create_pr", + "create_pr_comment", + "create_pr_review_comment", + "view_pr", +] diff --git a/src/codegen/extensions/tools/github/view_pr.py b/src/codegen/extensions/tools/github/view_pr.py index 13b90a0f5..c077b56a1 100644 --- a/src/codegen/extensions/tools/github/view_pr.py +++ b/src/codegen/extensions/tools/github/view_pr.py @@ -18,4 +18,4 @@ def view_pr(codebase: Codebase, pr_id: int) -> dict[str, Any]: modified_symbols, patch = codebase.get_modified_symbols_in_pr(pr_id) # Convert modified_symbols set to list for JSON serialization - return {"status": "success", "modified_symbols": list(modified_symbols), "patch": patch} + return {"status": "success", "patch": patch} diff --git a/src/codegen/extensions/tools/linear/__init__.py b/src/codegen/extensions/tools/linear/__init__.py index e69de29bb..8d15e1b72 100644 --- a/src/codegen/extensions/tools/linear/__init__.py +++ b/src/codegen/extensions/tools/linear/__init__.py @@ -0,0 +1,19 @@ +from .linear import ( + linear_comment_on_issue_tool, + linear_create_issue_tool, + linear_get_issue_comments_tool, + linear_get_issue_tool, + linear_get_teams_tool, + linear_register_webhook_tool, + linear_search_issues_tool, +) + +__all__ = [ + "linear_comment_on_issue_tool", + "linear_create_issue_tool", + "linear_get_issue_comments_tool", + "linear_get_issue_tool", + "linear_get_teams_tool", + "linear_register_webhook_tool", + "linear_search_issues_tool", +] diff --git a/tests/integration/extension/test_github.py b/tests/integration/extension/test_github.py index e69de29bb..a7ced0a3e 100644 --- a/tests/integration/extension/test_github.py +++ b/tests/integration/extension/test_github.py @@ -0,0 +1,26 @@ +"""Tests for Linear tools.""" + +import os + +import pytest + +from codegen import Codebase +from codegen.extensions.linear.linear_client import LinearClient +from codegen.extensions.tools.github import view_pr + + +@pytest.fixture +def client() -> LinearClient: + """Create a Linear client for testing.""" + token = os.getenv("CODEGEN_SECRETS__GITHUB_TOKEN") + if not token: + pytest.skip("CODEGEN_SECRETS__GITHUB_TOKEN environment variable not set") + codebase = Codebase.from_repo("codegen-sh/Kevin-s-Adventure-Game") + return codebase + + +def test_github_view_pr(client: LinearClient) -> None: + """Test getting an issue from Linear.""" + # Link to PR: https://github.com/codegen-sh/Kevin-s-Adventure-Game/pull/419 + pr = view_pr(client, 419) + print(pr) diff --git a/tests/integration/extension/test_linear.py b/tests/integration/extension/test_linear.py index a962fa366..4d39c9788 100644 --- a/tests/integration/extension/test_linear.py +++ b/tests/integration/extension/test_linear.py @@ -5,7 +5,7 @@ import pytest from codegen.extensions.linear.linear_client import LinearClient -from codegen.extensions.tools.linear import ( +from codegen.extensions.tools.linear.linear import ( linear_comment_on_issue_tool, linear_create_issue_tool, linear_get_issue_comments_tool,