mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-10-01 00:23:17 +00:00
refactor: extract shared utils from server.py into agent/utils/ package
This commit is contained in:
parent
a1a6f81015
commit
8c9bc67335
8 changed files with 214 additions and 198 deletions
|
|
@ -1,10 +1,9 @@
|
|||
from .check_message_queue import LinearNotifyState, check_message_queue_before_model
|
||||
from .check_message_queue import check_message_queue_before_model
|
||||
from .open_pr import open_pr_if_needed
|
||||
from .post_to_linear import post_to_linear_after_model
|
||||
from .tool_error_handler import ToolErrorMiddleware
|
||||
|
||||
__all__ = [
|
||||
"LinearNotifyState",
|
||||
"ToolErrorMiddleware",
|
||||
"check_message_queue_before_model",
|
||||
"open_pr_if_needed",
|
||||
|
|
|
|||
|
|
@ -17,6 +17,11 @@ from langchain.agents.middleware import AgentState, after_agent
|
|||
from langgraph.config import get_config
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from ..encryption import decrypt_token
|
||||
from ..utils.github import create_github_pr, get_github_default_branch
|
||||
from ..utils.linear import comment_on_linear_issue
|
||||
from ..utils.sandbox_state import SANDBOX_BACKENDS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
|
@ -41,19 +46,11 @@ def _extract_pr_params_from_messages(messages: list) -> dict[str, str] | None:
|
|||
|
||||
|
||||
@after_agent
|
||||
async def open_pr_if_needed( # noqa: PLR0912, PLR0915
|
||||
async def open_pr_if_needed(
|
||||
state: AgentState,
|
||||
runtime: Runtime, # noqa: ARG001
|
||||
runtime: Runtime,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Middleware that commits/pushes changes and comments on Linear after agent runs."""
|
||||
from ..encryption import decrypt_token
|
||||
from ..server import (
|
||||
_SANDBOX_BACKENDS,
|
||||
comment_on_linear_issue,
|
||||
create_github_pr,
|
||||
get_github_default_branch,
|
||||
)
|
||||
|
||||
logger.info("After-agent middleware started")
|
||||
pr_url = None
|
||||
pr_number = None
|
||||
|
|
@ -103,7 +100,7 @@ async def open_pr_if_needed( # noqa: PLR0912, PLR0915
|
|||
repo_owner = repo_config.get("owner")
|
||||
repo_name = repo_config.get("name")
|
||||
|
||||
sandbox_backend = _SANDBOX_BACKENDS.get(thread_id)
|
||||
sandbox_backend = SANDBOX_BACKENDS.get(thread_id)
|
||||
|
||||
repo_dir = f"/workspace/{repo_name}"
|
||||
|
||||
|
|
@ -214,7 +211,7 @@ async def open_pr_if_needed( # noqa: PLR0912, PLR0915
|
|||
|
||||
if linear_issue_id and last_message_content:
|
||||
if pr_url:
|
||||
comment = f"""✅ **Pull Request Created**
|
||||
comment = f"""**Pull Request Created**
|
||||
|
||||
I've created a pull request to address this issue:
|
||||
|
||||
|
|
@ -226,7 +223,7 @@ I've created a pull request to address this issue:
|
|||
|
||||
{last_message_content}"""
|
||||
else:
|
||||
comment = f"""🤖 **Agent Response**
|
||||
comment = f""" **Agent Response**
|
||||
|
||||
{last_message_content}"""
|
||||
await comment_on_linear_issue(linear_issue_id, comment)
|
||||
|
|
@ -241,7 +238,7 @@ I've created a pull request to address this issue:
|
|||
linear_issue = configurable.get("linear_issue", {})
|
||||
linear_issue_id = linear_issue.get("id")
|
||||
if linear_issue_id:
|
||||
error_comment = f"""❌ **Agent Error**
|
||||
error_comment = f""" **Agent Error**
|
||||
|
||||
An error occurred while processing this issue:
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from langchain.agents.middleware import after_model
|
|||
from langgraph.config import get_config
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from ..utils.linear import comment_on_linear_issue
|
||||
from .check_message_queue import LinearNotifyState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -34,8 +35,6 @@ async def post_to_linear_after_model( # noqa: PLR0911, PLR0912
|
|||
- The AI response has text content (not just tool calls)
|
||||
- The message hasn't already been sent (tracked via linear_messages_sent_count)
|
||||
"""
|
||||
from ..server import comment_on_linear_issue
|
||||
|
||||
try:
|
||||
config = get_config()
|
||||
configurable = config.get("configurable", {})
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@
|
|||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -95,184 +94,7 @@ SANDBOX_CREATING = "__creating__"
|
|||
SANDBOX_CREATION_TIMEOUT = 180
|
||||
SANDBOX_POLL_INTERVAL = 1.0
|
||||
|
||||
# HTTP status codes
|
||||
HTTP_CREATED = 201
|
||||
HTTP_UNPROCESSABLE_ENTITY = 422
|
||||
|
||||
# Message count thresholds
|
||||
_SANDBOX_BACKENDS: dict[str, Any] = {}
|
||||
|
||||
import httpx
|
||||
|
||||
LINEAR_API_KEY = os.environ.get("LINEAR_API_KEY", "")
|
||||
|
||||
|
||||
async def create_github_pr(
|
||||
repo_owner: str,
|
||||
repo_name: str,
|
||||
github_token: str,
|
||||
title: str,
|
||||
head_branch: str,
|
||||
base_branch: str,
|
||||
body: str,
|
||||
) -> tuple[str | None, int | None]:
|
||||
"""Create a GitHub pull request via the API.
|
||||
|
||||
Args:
|
||||
repo_owner: Repository owner (e.g., "langchain-ai")
|
||||
repo_name: Repository name (e.g., "deepagents")
|
||||
github_token: GitHub access token
|
||||
title: PR title
|
||||
head_branch: Source branch name
|
||||
base_branch: Target branch name
|
||||
body: PR description
|
||||
|
||||
Returns:
|
||||
Tuple of (pr_url, pr_number) if successful, (None, None) otherwise
|
||||
"""
|
||||
pr_payload = {
|
||||
"title": title,
|
||||
"head": head_branch,
|
||||
"base": base_branch,
|
||||
"body": body,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Creating PR: head=%s, base=%s, repo=%s/%s",
|
||||
head_branch,
|
||||
base_branch,
|
||||
repo_owner,
|
||||
repo_name,
|
||||
)
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient() as http_client:
|
||||
pr_response = await http_client.post(
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}/pulls",
|
||||
headers={
|
||||
"Authorization": f"Bearer {github_token}",
|
||||
"Accept": "application/vnd.github+json",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
},
|
||||
json=pr_payload,
|
||||
)
|
||||
|
||||
pr_data = pr_response.json()
|
||||
|
||||
if pr_response.status_code == HTTP_CREATED:
|
||||
pr_url = pr_data.get("html_url")
|
||||
pr_number = pr_data.get("number")
|
||||
logger.info("PR created successfully: %s", pr_url)
|
||||
return pr_url, pr_number
|
||||
|
||||
if pr_response.status_code == HTTP_UNPROCESSABLE_ENTITY:
|
||||
logger.error("GitHub API validation error (422): %s", pr_data.get("message"))
|
||||
else:
|
||||
logger.error(
|
||||
"GitHub API error (%s): %s",
|
||||
pr_response.status_code,
|
||||
pr_data.get("message"),
|
||||
)
|
||||
|
||||
if "errors" in pr_data:
|
||||
logger.error("GitHub API errors detail: %s", pr_data.get("errors"))
|
||||
|
||||
return None, None
|
||||
|
||||
except httpx.HTTPError:
|
||||
logger.exception("Failed to create PR via GitHub API")
|
||||
return None, None
|
||||
|
||||
|
||||
async def get_github_default_branch(
|
||||
repo_owner: str,
|
||||
repo_name: str,
|
||||
github_token: str,
|
||||
) -> str:
|
||||
"""Get the default branch of a GitHub repository via the API.
|
||||
|
||||
Args:
|
||||
repo_owner: Repository owner (e.g., "langchain-ai")
|
||||
repo_name: Repository name (e.g., "deepagents")
|
||||
github_token: GitHub access token
|
||||
|
||||
Returns:
|
||||
The default branch name (e.g., "main" or "master")
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient() as http_client:
|
||||
response = await http_client.get(
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}",
|
||||
headers={
|
||||
"Authorization": f"Bearer {github_token}",
|
||||
"Accept": "application/vnd.github+json",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
},
|
||||
)
|
||||
|
||||
if response.status_code == 200: # noqa: PLR2004
|
||||
repo_data = response.json()
|
||||
default_branch = repo_data.get("default_branch", "main")
|
||||
logger.debug("Got default branch from GitHub API: %s", default_branch)
|
||||
return default_branch
|
||||
|
||||
logger.warning(
|
||||
"Failed to get repo info from GitHub API (%s), falling back to 'main'",
|
||||
response.status_code,
|
||||
)
|
||||
return "main"
|
||||
|
||||
except httpx.HTTPError:
|
||||
logger.exception("Failed to get default branch from GitHub API, falling back to 'main'")
|
||||
return "main"
|
||||
|
||||
|
||||
async def comment_on_linear_issue(issue_id: str, comment_body: str) -> bool:
|
||||
"""Add a comment to a Linear issue.
|
||||
|
||||
Args:
|
||||
issue_id: The Linear issue ID
|
||||
comment_body: The comment text
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if not LINEAR_API_KEY:
|
||||
return False
|
||||
|
||||
import httpx
|
||||
|
||||
url = "https://api.linear.app/graphql"
|
||||
|
||||
mutation = """
|
||||
mutation CommentCreate($issueId: String!, $body: String!) {
|
||||
commentCreate(input: { issueId: $issueId, body: $body }) {
|
||||
success
|
||||
comment {
|
||||
id
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
async with httpx.AsyncClient() as http_client:
|
||||
try:
|
||||
response = await http_client.post(
|
||||
url,
|
||||
headers={
|
||||
"Authorization": LINEAR_API_KEY,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json={
|
||||
"query": mutation,
|
||||
"variables": {"issueId": issue_id, "body": comment_body},
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
return bool(result.get("data", {}).get("commentCreate", {}).get("success"))
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
from .utils.sandbox_state import SANDBOX_BACKENDS
|
||||
|
||||
|
||||
async def _clone_or_pull_repo_in_sandbox( # noqa: PLR0915
|
||||
|
|
@ -532,7 +354,7 @@ async def get_agent(config: RunnableConfig) -> Pregel: # noqa: PLR0915
|
|||
logger.exception("Failed to pull repo in existing sandbox")
|
||||
raise
|
||||
|
||||
_SANDBOX_BACKENDS[thread_id] = sandbox_backend
|
||||
SANDBOX_BACKENDS[thread_id] = sandbox_backend
|
||||
|
||||
linear_issue = config["configurable"].get("linear_issue", {})
|
||||
linear_project_id = linear_issue.get("linear_project_id", "")
|
||||
|
|
|
|||
0
apps/agent/agent/utils/__init__.py
Normal file
0
apps/agent/agent/utils/__init__.py
Normal file
133
apps/agent/agent/utils/github.py
Normal file
133
apps/agent/agent/utils/github.py
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
"""GitHub API utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# HTTP status codes
|
||||
HTTP_CREATED = 201
|
||||
HTTP_UNPROCESSABLE_ENTITY = 422
|
||||
|
||||
|
||||
async def create_github_pr(
|
||||
repo_owner: str,
|
||||
repo_name: str,
|
||||
github_token: str,
|
||||
title: str,
|
||||
head_branch: str,
|
||||
base_branch: str,
|
||||
body: str,
|
||||
) -> tuple[str | None, int | None]:
|
||||
"""Create a GitHub pull request via the API.
|
||||
|
||||
Args:
|
||||
repo_owner: Repository owner (e.g., "langchain-ai")
|
||||
repo_name: Repository name (e.g., "deepagents")
|
||||
github_token: GitHub access token
|
||||
title: PR title
|
||||
head_branch: Source branch name
|
||||
base_branch: Target branch name
|
||||
body: PR description
|
||||
|
||||
Returns:
|
||||
Tuple of (pr_url, pr_number) if successful, (None, None) otherwise
|
||||
"""
|
||||
pr_payload = {
|
||||
"title": title,
|
||||
"head": head_branch,
|
||||
"base": base_branch,
|
||||
"body": body,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Creating PR: head=%s, base=%s, repo=%s/%s",
|
||||
head_branch,
|
||||
base_branch,
|
||||
repo_owner,
|
||||
repo_name,
|
||||
)
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient() as http_client:
|
||||
pr_response = await http_client.post(
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}/pulls",
|
||||
headers={
|
||||
"Authorization": f"Bearer {github_token}",
|
||||
"Accept": "application/vnd.github+json",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
},
|
||||
json=pr_payload,
|
||||
)
|
||||
|
||||
pr_data = pr_response.json()
|
||||
|
||||
if pr_response.status_code == HTTP_CREATED:
|
||||
pr_url = pr_data.get("html_url")
|
||||
pr_number = pr_data.get("number")
|
||||
logger.info("PR created successfully: %s", pr_url)
|
||||
return pr_url, pr_number
|
||||
|
||||
if pr_response.status_code == HTTP_UNPROCESSABLE_ENTITY:
|
||||
logger.error("GitHub API validation error (422): %s", pr_data.get("message"))
|
||||
else:
|
||||
logger.error(
|
||||
"GitHub API error (%s): %s",
|
||||
pr_response.status_code,
|
||||
pr_data.get("message"),
|
||||
)
|
||||
|
||||
if "errors" in pr_data:
|
||||
logger.error("GitHub API errors detail: %s", pr_data.get("errors"))
|
||||
|
||||
return None, None
|
||||
|
||||
except httpx.HTTPError:
|
||||
logger.exception("Failed to create PR via GitHub API")
|
||||
return None, None
|
||||
|
||||
|
||||
async def get_github_default_branch(
|
||||
repo_owner: str,
|
||||
repo_name: str,
|
||||
github_token: str,
|
||||
) -> str:
|
||||
"""Get the default branch of a GitHub repository via the API.
|
||||
|
||||
Args:
|
||||
repo_owner: Repository owner (e.g., "langchain-ai")
|
||||
repo_name: Repository name (e.g., "deepagents")
|
||||
github_token: GitHub access token
|
||||
|
||||
Returns:
|
||||
The default branch name (e.g., "main" or "master")
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient() as http_client:
|
||||
response = await http_client.get(
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}",
|
||||
headers={
|
||||
"Authorization": f"Bearer {github_token}",
|
||||
"Accept": "application/vnd.github+json",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
},
|
||||
)
|
||||
|
||||
if response.status_code == 200: # noqa: PLR2004
|
||||
repo_data = response.json()
|
||||
default_branch = repo_data.get("default_branch", "main")
|
||||
logger.debug("Got default branch from GitHub API: %s", default_branch)
|
||||
return default_branch
|
||||
|
||||
logger.warning(
|
||||
"Failed to get repo info from GitHub API (%s), falling back to 'main'",
|
||||
response.status_code,
|
||||
)
|
||||
return "main"
|
||||
|
||||
except httpx.HTTPError:
|
||||
logger.exception("Failed to get default branch from GitHub API, falling back to 'main'")
|
||||
return "main"
|
||||
58
apps/agent/agent/utils/linear.py
Normal file
58
apps/agent/agent/utils/linear.py
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
"""Linear API utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LINEAR_API_KEY = os.environ.get("LINEAR_API_KEY", "")
|
||||
|
||||
|
||||
async def comment_on_linear_issue(issue_id: str, comment_body: str) -> bool:
|
||||
"""Add a comment to a Linear issue.
|
||||
|
||||
Args:
|
||||
issue_id: The Linear issue ID
|
||||
comment_body: The comment text
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if not LINEAR_API_KEY:
|
||||
return False
|
||||
|
||||
url = "https://api.linear.app/graphql"
|
||||
|
||||
mutation = """
|
||||
mutation CommentCreate($issueId: String!, $body: String!) {
|
||||
commentCreate(input: { issueId: $issueId, body: $body }) {
|
||||
success
|
||||
comment {
|
||||
id
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
async with httpx.AsyncClient() as http_client:
|
||||
try:
|
||||
response = await http_client.post(
|
||||
url,
|
||||
headers={
|
||||
"Authorization": LINEAR_API_KEY,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json={
|
||||
"query": mutation,
|
||||
"variables": {"issueId": issue_id, "body": comment_body},
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
return bool(result.get("data", {}).get("commentCreate", {}).get("success"))
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
8
apps/agent/agent/utils/sandbox_state.py
Normal file
8
apps/agent/agent/utils/sandbox_state.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
"""Shared sandbox state used by server and middleware."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
# Thread ID -> SandboxBackend mapping, shared between server.py and middleware
|
||||
SANDBOX_BACKENDS: dict[str, Any] = {}
|
||||
Loading…
Add table
Reference in a new issue