Skip to content

feat: adds linear tools #499

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Feb 14, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ dependencies = [
"numpy>=2.2.2",
"mcp[cli]",
"neo4j",
"modal>=0.73.45",
]

license = { text = "Apache-2.0" }
Expand Down
2 changes: 1 addition & 1 deletion src/codegen/extensions/events/app.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging

import modal
import modal # deptry: ignore

from codegen.extensions.events.linear import Linear

Expand Down
2 changes: 1 addition & 1 deletion src/codegen/extensions/events/interface.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from typing import Protocol

import modal
import modal # deptry: ignore


class EventHandlerManagerProtocol(Protocol):
Expand Down
2 changes: 1 addition & 1 deletion src/codegen/extensions/events/linear.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import os
from typing import Callable

import modal
import modal # deptry: ignore
from anthropic import BaseModel

from codegen.extensions.clients.linear import LinearClient
Expand All @@ -23,15 +23,15 @@
self.access_token = os.environ["LINEAR_ACCESS_TOKEN"] # move to extensions config.
self.signing_secret = os.environ["LINEAR_SIGNING_SECRET"]
self.linear_team_id = os.environ["LINEAR_TEAM_ID"]
self.registered_handlers = {}

Check failure on line 26 in src/codegen/extensions/events/linear.py

View workflow job for this annotation

GitHub Actions / mypy

error: Need type annotation for "registered_handlers" (hint: "registered_handlers: dict[<type>, <type>] = ...") [var-annotated]

def subscribe_handler_to_webhook(self, web_url: str, event_name: str):

Check failure on line 28 in src/codegen/extensions/events/linear.py

View workflow job for this annotation

GitHub Actions / mypy

error: Signature of "subscribe_handler_to_webhook" incompatible with supertype "EventHandlerManagerProtocol" [override]
client = LinearClient(access_token=self.access_token)

result = client.register_webhook(team_id=self.linear_team_id, webhook_url=web_url, enabled=True, resource_types=[event_name], secret=self.signing_secret)
return result

def unsubscribe_handler_to_webhook(self, registered_handler: RegisteredWebhookHandler):

Check failure on line 34 in src/codegen/extensions/events/linear.py

View workflow job for this annotation

GitHub Actions / mypy

error: Signature of "unsubscribe_handler_to_webhook" incompatible with supertype "EventHandlerManagerProtocol" [override]
webhook_id = registered_handler.webhook_id

client = LinearClient(access_token=self.access_token)
Expand Down
86 changes: 84 additions & 2 deletions src/codegen/extensions/langchain/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -37,9 +43,9 @@
class ViewFileTool(BaseTool):
"""Tool for viewing file contents and metadata."""

name: ClassVar[str] = "view_file"

Check failure on line 46 in src/codegen/extensions/langchain/tools.py

View workflow job for this annotation

GitHub Actions / mypy

error: Cannot override instance variable (previously declared on base class "BaseTool") with class variable [misc]
description: ClassVar[str] = "View the contents and metadata of a file in the codebase"

Check failure on line 47 in src/codegen/extensions/langchain/tools.py

View workflow job for this annotation

GitHub Actions / mypy

error: Cannot override instance variable (previously declared on base class "BaseTool") with class variable [misc]
args_schema: ClassVar[type[BaseModel]] = ViewFileInput

Check failure on line 48 in src/codegen/extensions/langchain/tools.py

View workflow job for this annotation

GitHub Actions / mypy

error: Cannot override instance variable (previously declared on base class "BaseTool") with class variable [misc]
codebase: Codebase = Field(exclude=True)

def __init__(self, codebase: Codebase) -> None:
Expand All @@ -60,8 +66,8 @@
class ListDirectoryTool(BaseTool):
"""Tool for listing directory contents."""

name: ClassVar[str] = "list_directory"

Check failure on line 69 in src/codegen/extensions/langchain/tools.py

View workflow job for this annotation

GitHub Actions / mypy

error: Cannot override instance variable (previously declared on base class "BaseTool") with class variable [misc]
description: ClassVar[str] = "List contents of a directory in the codebase"

Check failure on line 70 in src/codegen/extensions/langchain/tools.py

View workflow job for this annotation

GitHub Actions / mypy

error: Cannot override instance variable (previously declared on base class "BaseTool") with class variable [misc]
args_schema: ClassVar[type[BaseModel]] = ListDirectoryInput
codebase: Codebase = Field(exclude=True)

Expand Down Expand Up @@ -184,7 +190,10 @@

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")

Expand Down Expand Up @@ -281,7 +290,10 @@
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")


Expand Down Expand Up @@ -453,6 +465,73 @@
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.

Expand All @@ -479,4 +558,7 @@
SemanticEditTool(codebase),
SemanticSearchTool(codebase),
ViewFileTool(codebase),
LinearGetIssueTool(codebase),
LinearGetIssueCommentsTool(codebase),
LinearCommentOnIssueTool(codebase),
]
3 changes: 3 additions & 0 deletions src/codegen/extensions/linear/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .linear_client import LinearClient

__all__ = ["LinearClient"]
196 changes: 196 additions & 0 deletions src/codegen/extensions/linear/linear_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,196 @@
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

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)
11 changes: 11 additions & 0 deletions src/codegen/extensions/tools/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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",
Expand Down
Loading