From a895aad1600c679138b46c2bf90b6e0d81157ab5 Mon Sep 17 00:00:00 2001 From: jayhack Date: Thu, 13 Feb 2025 19:24:08 -0800 Subject: [PATCH 1/3] . --- src/codegen/extensions/langchain/tools.py | 76 +++++++++ src/codegen/extensions/linear/__init__.py | 3 + .../extensions/linear/linear_client.py | 153 ++++++++++++++++++ src/codegen/extensions/tools/__init__.py | 11 ++ src/codegen/extensions/tools/linear_tools.py | 39 +++++ tests/integration/extension/test_linear.py | 43 +++++ 6 files changed, 325 insertions(+) create mode 100644 src/codegen/extensions/linear/__init__.py create mode 100644 src/codegen/extensions/linear/linear_client.py create mode 100644 src/codegen/extensions/tools/linear_tools.py create mode 100644 tests/integration/extension/test_linear.py diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index 6299cc556..b09ea7105 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -7,6 +7,12 @@ from pydantic import BaseModel, Field from codegen import Codebase +from codegen.extensions.linear.linear_client import LinearClient +from codegen.extensions.tools.linear_tools import ( + linear_comment_on_issue_tool, + linear_get_issue_comments_tool, + linear_get_issue_tool, +) from ..tools import ( commit, @@ -453,6 +459,73 @@ def _run( return json.dumps(result, indent=2) +class LinearGetIssueInput(BaseModel): + """Input for getting a Linear issue.""" + + issue_id: str = Field(..., description="ID of the Linear issue to retrieve") + + +class LinearGetIssueTool(BaseTool): + """Tool for getting Linear issue details.""" + + name: ClassVar[str] = "linear_get_issue" + description: ClassVar[str] = "Get details of a Linear issue by its ID" + args_schema: ClassVar[type[BaseModel]] = LinearGetIssueInput + client: LinearClient = Field(exclude=True) + + def __init__(self, client: LinearClient) -> None: + super().__init__(client=client) + + def _run(self, issue_id: str) -> str: + result = linear_get_issue_tool(self.client, issue_id) + return json.dumps(result, indent=2) + + +class LinearGetIssueCommentsInput(BaseModel): + """Input for getting Linear issue comments.""" + + issue_id: str = Field(..., description="ID of the Linear issue to get comments for") + + +class LinearGetIssueCommentsTool(BaseTool): + """Tool for getting Linear issue comments.""" + + name: ClassVar[str] = "linear_get_issue_comments" + description: ClassVar[str] = "Get all comments on a Linear issue" + args_schema: ClassVar[type[BaseModel]] = LinearGetIssueCommentsInput + client: LinearClient = Field(exclude=True) + + def __init__(self, client: LinearClient) -> None: + super().__init__(client=client) + + def _run(self, issue_id: str) -> str: + result = linear_get_issue_comments_tool(self.client, issue_id) + return json.dumps(result, indent=2) + + +class LinearCommentOnIssueInput(BaseModel): + """Input for commenting on a Linear issue.""" + + issue_id: str = Field(..., description="ID of the Linear issue to comment on") + body: str = Field(..., description="The comment text") + + +class LinearCommentOnIssueTool(BaseTool): + """Tool for commenting on Linear issues.""" + + name: ClassVar[str] = "linear_comment_on_issue" + description: ClassVar[str] = "Add a comment to a Linear issue" + args_schema: ClassVar[type[BaseModel]] = LinearCommentOnIssueInput + client: LinearClient = Field(exclude=True) + + def __init__(self, client: LinearClient) -> None: + super().__init__(client=client) + + def _run(self, issue_id: str, body: str) -> str: + result = linear_comment_on_issue_tool(self.client, issue_id, body) + return json.dumps(result, indent=2) + + def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: """Get all workspace tools initialized with a codebase. @@ -479,4 +552,7 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: SemanticEditTool(codebase), SemanticSearchTool(codebase), ViewFileTool(codebase), + LinearGetIssueTool(codebase), + LinearGetIssueCommentsTool(codebase), + LinearCommentOnIssueTool(codebase), ] diff --git a/src/codegen/extensions/linear/__init__.py b/src/codegen/extensions/linear/__init__.py new file mode 100644 index 000000000..8ba060245 --- /dev/null +++ b/src/codegen/extensions/linear/__init__.py @@ -0,0 +1,3 @@ +from .linear_client import LinearClient + +__all__ = ["LinearClient"] diff --git a/src/codegen/extensions/linear/linear_client.py b/src/codegen/extensions/linear/linear_client.py new file mode 100644 index 000000000..b7d1a66d2 --- /dev/null +++ b/src/codegen/extensions/linear/linear_client.py @@ -0,0 +1,153 @@ +import json +import logging +import os +from typing import Optional + +import requests +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + + +# --- TYPES + + +class LinearUser(BaseModel): + id: str + name: str + + +class LinearComment(BaseModel): + id: str + body: str + user: LinearUser | None = None + + +class LinearIssue(BaseModel): + id: str + title: str + description: str | None = None + + +class LinearClient: + api_headers: dict + api_endpoint = "https://api.linear.app/graphql" + + def __init__(self, access_token: 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 + self.api_headers = { + "Content-Type": "application/json", + "Authorization": self.access_token, + } + + def get_issue(self, issue_id: str) -> LinearIssue: + query = """ + query getIssue($issueId: String!) { + issue(id: $issueId) { + id + title + description + } + } + """ + variables = {"issueId": issue_id} + response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": query, "variables": variables}) + data = response.json() + issue_data = data["data"]["issue"] + return LinearIssue(id=issue_data["id"], title=issue_data["title"], description=issue_data["description"]) + + def get_issue_comments(self, issue_id: str) -> list[LinearComment]: + query = """ + query getIssueComments($issueId: String!) { + issue(id: $issueId) { + comments { + nodes { + id + body + user { + id + name + } + } + + } + } + } + """ + variables = {"issueId": issue_id} + response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": query, "variables": variables}) + data = response.json() + comments = data["data"]["issue"]["comments"]["nodes"] + + # Parse comments into list of LinearComment objects + parsed_comments = [] + for comment in comments: + user = comment.get("user", None) + parsed_comment = LinearComment(id=comment["id"], body=comment["body"], user=LinearUser(id=user.get("id"), name=user.get("name")) if user else None) + parsed_comments.append(parsed_comment) + + # Convert raw comments to LinearComment objects + return parsed_comments + + def comment_on_issue(self, issue_id: str, body: str) -> dict: + """issue_id is our internal issue ID""" + query = """mutation makeComment($issueId: String!, $body: String!) { + commentCreate(input: {issueId: $issueId, body: $body}) { + comment { + id + body + url + user { + id + name + } + } + } + } + """ + variables = {"issueId": issue_id, "body": body} + response = requests.post( + self.api_endpoint, + headers=self.api_headers, + data=json.dumps({"query": query, "variables": variables}), + ) + data = response.json() + try: + comment_data = data["data"]["commentCreate"]["comment"] + + return comment_data + except Exception as e: + msg = f"Error creating comment\n{data}\n{e}" + raise Exception(msg) + + def register_webhook(self, webhook_url: str, team_id: str, secret: str, enabled: bool, resource_types: list[str]): + mutation = """ + mutation createWebhook($input: WebhookCreateInput!) { + webhookCreate(input: $input) { + success + webhook { + id + enabled + } + } + } + """ + + variables = { + "input": { + "url": webhook_url, + "teamId": team_id, + "resourceTypes": resource_types, + "enabled": enabled, + "secret": secret, + } + } + + response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": mutation, "variables": variables}) + body = response.json() + return body diff --git a/src/codegen/extensions/tools/__init__.py b/src/codegen/extensions/tools/__init__.py index 74f8ba83c..06ededaf8 100644 --- a/src/codegen/extensions/tools/__init__.py +++ b/src/codegen/extensions/tools/__init__.py @@ -8,6 +8,12 @@ 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 ( + linear_comment_on_issue_tool, + linear_get_issue_comments_tool, + linear_get_issue_tool, + linear_register_webhook_tool, +) from .list_directory import list_directory from .move_symbol import move_symbol from .rename_file import rename_file @@ -27,6 +33,11 @@ "create_pr_review_comment", "delete_file", "edit_file", + # Linear operations + "linear_comment_on_issue_tool", + "linear_get_issue_comments_tool", + "linear_get_issue_tool", + "linear_register_webhook_tool", "list_directory", # Symbol operations "move_symbol", diff --git a/src/codegen/extensions/tools/linear_tools.py b/src/codegen/extensions/tools/linear_tools.py new file mode 100644 index 000000000..9bc02f5fe --- /dev/null +++ b/src/codegen/extensions/tools/linear_tools.py @@ -0,0 +1,39 @@ +from typing import Any + +from codegen.extensions.linear.linear_client import LinearClient + + +def linear_get_issue_tool(client: LinearClient, issue_id: str) -> dict[str, Any]: + """Get an issue by its ID.""" + try: + issue = client.get_issue(issue_id) + return {"status": "success", "issue": issue.dict()} + except Exception as e: + return {"error": f"Failed to get issue: {e!s}"} + + +def linear_get_issue_comments_tool(client: LinearClient, issue_id: str) -> dict[str, Any]: + """Get comments for a specific issue.""" + try: + comments = client.get_issue_comments(issue_id) + return {"status": "success", "comments": [comment.dict() for comment in comments]} + except Exception as e: + return {"error": f"Failed to get issue comments: {e!s}"} + + +def linear_comment_on_issue_tool(client: LinearClient, issue_id: str, body: str) -> dict[str, Any]: + """Add a comment to an issue.""" + try: + comment = client.comment_on_issue(issue_id, body) + return {"status": "success", "comment": comment} + except Exception as e: + return {"error": f"Failed to comment on issue: {e!s}"} + + +def linear_register_webhook_tool(client: LinearClient, webhook_url: str, team_id: str, secret: str, enabled: bool, resource_types: list[str]) -> dict[str, Any]: + """Register a webhook with Linear.""" + try: + response = client.register_webhook(webhook_url, team_id, secret, enabled, resource_types) + return {"status": "success", "response": response} + except Exception as e: + return {"error": f"Failed to register webhook: {e!s}"} diff --git a/tests/integration/extension/test_linear.py b/tests/integration/extension/test_linear.py new file mode 100644 index 000000000..1c0742121 --- /dev/null +++ b/tests/integration/extension/test_linear.py @@ -0,0 +1,43 @@ +"""Tests for Linear tools.""" + +import os + +import pytest + +from codegen.extensions.linear.linear_client import LinearClient +from codegen.extensions.tools.linear_tools import ( + linear_comment_on_issue_tool, + linear_get_issue_comments_tool, + linear_get_issue_tool, +) + + +@pytest.fixture +def client() -> LinearClient: + """Create a Linear client for testing.""" + token = os.getenv("LINEAR_ACCESS_TOKEN") + if not token: + pytest.skip("LINEAR_ACCESS_TOKEN environment variable not set") + return LinearClient(token) + + +def test_linear_get_issue(client: LinearClient) -> None: + """Test getting an issue from Linear.""" + # Link to issue: https://linear.app/codegen-sh/issue/CG-10775/read-file-and-reveal-symbol-tool-size-limits + issue = linear_get_issue_tool(client, "CG-10775") + assert issue["status"] == "success" + assert issue["issue"]["id"] == "d5a7d6db-e20d-4d67-98f8-acedef6d3536" + + +def test_linear_get_issue_comments(client: LinearClient) -> None: + """Test getting comments for an issue from Linear.""" + comments = linear_get_issue_comments_tool(client, "CG-10775") + assert comments["status"] == "success" + assert len(comments["comments"]) > 1 + + +def test_linear_comment_on_issue(client: LinearClient) -> None: + """Test commenting on a Linear issue.""" + test_comment = "Test comment from automated testing" + result = linear_comment_on_issue_tool(client, "CG-10775", test_comment) + assert result["status"] == "success" From 6044663bb0e8708f81fe35f1bbae11ceaa488f98 Mon Sep 17 00:00:00 2001 From: jayhack Date: Thu, 13 Feb 2025 19:47:43 -0800 Subject: [PATCH 2/3] . --- src/codegen/extensions/langchain/tools.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index b09ea7105..cbce0c1cc 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -190,7 +190,10 @@ class RevealSymbolInput(BaseModel): symbol_name: str = Field(..., description="Name of the symbol to analyze") degree: int = Field(default=1, description="How many degrees of separation to traverse") - max_tokens: Optional[int] = Field(default=None, description="Optional maximum number of tokens for all source code combined") + max_tokens: Optional[int] = Field( + default=None, + description="Optional maximum number of tokens for all source code combined", + ) collect_dependencies: bool = Field(default=True, description="Whether to collect dependencies") collect_usages: bool = Field(default=True, description="Whether to collect usages") @@ -287,7 +290,10 @@ class MoveSymbolInput(BaseModel): source_file: str = Field(..., description="Path to the file containing the symbol") symbol_name: str = Field(..., description="Name of the symbol to move") target_file: str = Field(..., description="Path to the destination file") - strategy: Literal["update_all_imports", "add_back_edge"] = Field(default="update_all_imports", description="Strategy for handling imports: 'update_all_imports' (default) or 'add_back_edge'") + strategy: Literal["update_all_imports", "add_back_edge"] = Field( + default="update_all_imports", + description="Strategy for handling imports: 'update_all_imports' (default) or 'add_back_edge'", + ) include_dependencies: bool = Field(default=True, description="Whether to move dependencies along with the symbol") From ebff424ff987878b61ef9e004b7f5defa89e9004 Mon Sep 17 00:00:00 2001 From: jayhack Date: Thu, 13 Feb 2025 19:54:03 -0800 Subject: [PATCH 3/3] . --- src/codegen/cli/commands/create/main.py | 2 - .../extensions/linear/linear_client.py | 43 +++++++++++++++++++ tests/integration/extension/test_linear.py | 7 +++ 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/src/codegen/cli/commands/create/main.py b/src/codegen/cli/commands/create/main.py index 551d85e85..bbed5c3eb 100644 --- a/src/codegen/cli/commands/create/main.py +++ b/src/codegen/cli/commands/create/main.py @@ -5,7 +5,6 @@ from codegen.cli.api.client import RestAPI from codegen.cli.auth.constants import PROMPTS_DIR -from codegen.cli.auth.decorators import requires_auth from codegen.cli.auth.session import CodegenSession from codegen.cli.auth.token_manager import get_current_token from codegen.cli.codemod.convert import convert_to_cli @@ -66,7 +65,6 @@ def make_relative(path: Path) -> str: @click.command(name="create") -@requires_auth @click.argument("name", type=str) @click.argument("path", type=click.Path(path_type=Path), default=Path.cwd()) @click.option("--description", "-d", default=None, help="Description of what this codemod does.") diff --git a/src/codegen/extensions/linear/linear_client.py b/src/codegen/extensions/linear/linear_client.py index b7d1a66d2..ef0e1b96a 100644 --- a/src/codegen/extensions/linear/linear_client.py +++ b/src/codegen/extensions/linear/linear_client.py @@ -151,3 +151,46 @@ def register_webhook(self, webhook_url: str, team_id: str, secret: str, enabled: response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": mutation, "variables": variables}) body = response.json() return body + + def search_issues(self, query: str, limit: int = 10) -> list[LinearIssue]: + """Search for issues using a query string. + + Args: + query: Search query string + limit: Maximum number of issues to return (default: 10) + + Returns: + List of LinearIssue objects matching the search query + """ + graphql_query = """ + query searchIssues($query: String!, $limit: Int!) { + issueSearch(query: $query, first: $limit) { + nodes { + id + title + description + } + } + } + """ + variables = {"query": query, "limit": limit} + response = requests.post( + self.api_endpoint, + headers=self.api_headers, + json={"query": graphql_query, "variables": variables}, + ) + data = response.json() + + try: + issues_data = data["data"]["issueSearch"]["nodes"] + return [ + LinearIssue( + id=issue["id"], + title=issue["title"], + description=issue["description"], + ) + for issue in issues_data + ] + except Exception as e: + msg = f"Error searching issues\n{data}\n{e}" + raise Exception(msg) diff --git a/tests/integration/extension/test_linear.py b/tests/integration/extension/test_linear.py index 1c0742121..c6eb3606d 100644 --- a/tests/integration/extension/test_linear.py +++ b/tests/integration/extension/test_linear.py @@ -41,3 +41,10 @@ def test_linear_comment_on_issue(client: LinearClient) -> None: test_comment = "Test comment from automated testing" result = linear_comment_on_issue_tool(client, "CG-10775", test_comment) assert result["status"] == "success" + + +def test_search_issues(client: LinearClient) -> None: + """Test searching for issues in Linear.""" + issues = client.search_issues("REVEAL_SYMBOL") + assert issues["status"] == "success" + assert len(issues["issues"]) > 0