Skip to content

Commit d7fa98c

Browse files
authored
Merge pull request #138 from kevinbackhouse/hatch-fmt-linter
hatch fmt --linter
2 parents 78c2c85 + ebe62e5 commit d7fa98c

26 files changed

+179
-188
lines changed

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ Instead of generating CodeQL queries itself, the CodeQL MCP Server is used to pr
3030

3131
## Requirements
3232

33-
Python >= 3.9 or Docker
33+
Python >= 3.10 or Docker
3434

3535
## Configuration
3636

pyproject.toml

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ name = "seclab-taskflow-agent"
77
dynamic = ["version"]
88
description = "A taskflow agent for the SecLab project, enabling secure and automated workflow execution."
99
readme = "README.md"
10-
requires-python = ">=3.9"
10+
requires-python = ">=3.10"
1111
license = "MIT"
1212
keywords = []
1313
authors = [
@@ -16,8 +16,6 @@ authors = [
1616
classifiers = [
1717
"Development Status :: 4 - Beta",
1818
"Programming Language :: Python",
19-
"Programming Language :: Python :: 3.8",
20-
"Programming Language :: Python :: 3.9",
2119
"Programming Language :: Python :: 3.10",
2220
"Programming Language :: Python :: 3.11",
2321
"Programming Language :: Python :: 3.12",

release_tools/copy_files.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,16 @@
33

44
import os
55
import shutil
6-
import sys
76
import subprocess
7+
import sys
88

99

1010
def read_file_list(list_path):
1111
"""
1212
Reads a file containing file paths, ignoring empty lines and lines starting with '#'.
1313
Returns a list of relative file paths.
1414
"""
15-
with open(list_path, "r") as f:
15+
with open(list_path) as f:
1616
lines = [line.strip() for line in f]
1717
return [line for line in lines if line and not line.startswith("#")]
1818

release_tools/publish_docker.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
# SPDX-FileCopyrightText: 2025 GitHub
22
# SPDX-License-Identifier: MIT
33

4-
import os
5-
import shutil
64
import subprocess
75
import sys
86

src/seclab_taskflow_agent/__main__.py

Lines changed: 39 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,46 @@
11
# SPDX-FileCopyrightText: 2025 GitHub
22
# SPDX-License-Identifier: MIT
33

4-
import asyncio
5-
from threading import Thread
64
import argparse
7-
import os
8-
import sys
9-
from dotenv import load_dotenv, find_dotenv
5+
import asyncio
6+
import json
107
import logging
11-
from logging.handlers import RotatingFileHandler
12-
from pprint import pprint, pformat
8+
import os
9+
import pathlib
1310
import re
14-
import json
11+
import sys
1512
import uuid
16-
import pathlib
13+
from collections.abc import Callable
14+
from logging.handlers import RotatingFileHandler
15+
from pprint import pformat
1716

18-
from .agent import DEFAULT_MODEL, TaskRunHooks, TaskAgentHooks
17+
from agents import Agent, RunContextWrapper, TContext, Tool
18+
from agents.agent import ModelSettings
1919

2020
# from agents.run import DEFAULT_MAX_TURNS # XXX: this is 10, we need more than that
21-
from agents.exceptions import MaxTurnsExceeded, AgentsException
22-
from agents.agent import ModelSettings
23-
from agents.mcp import MCPServer, MCPServerStdio, MCPServerSse, MCPServerStreamableHttp, create_static_tool_filter
21+
from agents.exceptions import AgentsException, MaxTurnsExceeded
2422
from agents.extensions.handoff_prompt import prompt_with_handoff_instructions
25-
from agents import Tool, RunContextWrapper, TContext, Agent
26-
from openai import BadRequestError, APITimeoutError, RateLimitError
23+
from agents.mcp import MCPServerSse, MCPServerStdio, MCPServerStreamableHttp, create_static_tool_filter
24+
from dotenv import find_dotenv, load_dotenv
25+
from openai import APITimeoutError, BadRequestError, RateLimitError
2726
from openai.types.responses import ResponseTextDeltaEvent
28-
from typing import Callable
2927

30-
from .shell_utils import shell_tool_call
28+
from .agent import DEFAULT_MODEL, TaskAgent, TaskAgentHooks, TaskRunHooks
29+
from .available_tools import AvailableTools
30+
from .capi import get_AI_token, list_tool_call_models
31+
from .env_utils import TmpEnv
3132
from .mcp_utils import (
3233
DEFAULT_MCP_CLIENT_SESSION_TIMEOUT,
33-
ReconnectingMCPServerStdio,
3434
MCPNamespaceWrap,
35-
mcp_client_params,
36-
mcp_system_prompt,
35+
ReconnectingMCPServerStdio,
3736
StreamableMCPThread,
3837
compress_name,
38+
mcp_client_params,
39+
mcp_system_prompt,
3940
)
40-
from .render_utils import render_model_output, flush_async_output
41-
from .env_utils import TmpEnv
42-
from .agent import TaskAgent
43-
from .capi import list_tool_call_models, get_AI_token
44-
from .available_tools import AvailableTools
4541
from .path_utils import log_file_name
42+
from .render_utils import flush_async_output, render_model_output
43+
from .shell_utils import shell_tool_call
4644

4745
load_dotenv(find_dotenv(usecwd=True))
4846

@@ -91,7 +89,7 @@ def parse_prompt_args(available_tools: AvailableTools, user_prompt: str | None =
9189
args = parser.parse_known_args(user_prompt.split(" ") if user_prompt else None)
9290
except SystemExit as e:
9391
if e.code == 2:
94-
logging.error(f"User provided incomplete prompt: {user_prompt}")
92+
logging.exception(f"User provided incomplete prompt: {user_prompt}")
9593
return None, None, None, None, help_msg
9694
p = args[0].p.strip() if args[0].p else None
9795
t = args[0].t.strip() if args[0].t else None
@@ -258,14 +256,13 @@ async def mcp_session_task(mcp_servers: list, connected: asyncio.Event, cleanup:
258256
except Exception as e:
259257
print(f"Streamable mcp server process exception: {e}")
260258
except asyncio.CancelledError:
261-
logging.error(f"Timeout on cleanup for mcp server: {server._name}")
259+
logging.exception(f"Timeout on cleanup for mcp server: {server._name}")
262260
finally:
263261
mcp_servers.remove(s)
264262
except RuntimeError as e:
265-
logging.error(f"RuntimeError in mcp session task: {e}")
263+
logging.exception("RuntimeError in mcp session task")
266264
except asyncio.CancelledError as e:
267-
logging.error(f"Timeout on main session task: {e}")
268-
pass
265+
logging.exception("Timeout on main session task")
269266
finally:
270267
mcp_servers.clear()
271268

@@ -353,17 +350,17 @@ async def _run_streamed():
353350
return
354351
except APITimeoutError:
355352
if not max_retry:
356-
logging.error(f"Max retries for APITimeoutError reached")
353+
logging.exception("Max retries for APITimeoutError reached")
357354
raise
358355
max_retry -= 1
359356
except RateLimitError:
360357
if rate_limit_backoff == MAX_RATE_LIMIT_BACKOFF:
361-
raise APITimeoutError(f"Max rate limit backoff reached")
358+
raise APITimeoutError("Max rate limit backoff reached")
362359
if rate_limit_backoff > MAX_RATE_LIMIT_BACKOFF:
363360
rate_limit_backoff = MAX_RATE_LIMIT_BACKOFF
364361
else:
365362
rate_limit_backoff += rate_limit_backoff
366-
logging.error(f"Hit rate limit ... holding for {rate_limit_backoff}")
363+
logging.exception(f"Hit rate limit ... holding for {rate_limit_backoff}")
367364
await asyncio.sleep(rate_limit_backoff)
368365

369366
await _run_streamed()
@@ -372,16 +369,16 @@ async def _run_streamed():
372369
# raise exceptions up to here for anything that indicates a task failure
373370
except MaxTurnsExceeded as e:
374371
await render_model_output(f"** 🤖❗ Max Turns Reached: {e}\n", async_task=async_task, task_id=task_id)
375-
logging.error(f"Exceeded max_turns: {max_turns}")
372+
logging.exception(f"Exceeded max_turns: {max_turns}")
376373
except AgentsException as e:
377374
await render_model_output(f"** 🤖❗ Agent Exception: {e}\n", async_task=async_task, task_id=task_id)
378-
logging.error(f"Agent Exception: {e}")
375+
logging.exception("Agent Exception")
379376
except BadRequestError as e:
380377
await render_model_output(f"** 🤖❗ Request Error: {e}\n", async_task=async_task, task_id=task_id)
381-
logging.error(f"Bad Request: {e}")
378+
logging.exception("Bad Request")
382379
except APITimeoutError as e:
383380
await render_model_output(f"** 🤖❗ Timeout Error: {e}\n", async_task=async_task, task_id=task_id)
384-
logging.error(f"Bad Request: {e}")
381+
logging.exception("Bad Request")
385382

386383
if async_task:
387384
await flush_async_output(task_id)
@@ -392,14 +389,14 @@ async def _run_streamed():
392389
# signal mcp sessions task that it can disconnect our servers
393390
start_cleanup.set()
394391
cleanup_attempts_left = len(mcp_servers)
395-
while cleanup_attempts_left and len(mcp_servers):
392+
while cleanup_attempts_left and mcp_servers:
396393
try:
397394
cleanup_attempts_left -= 1
398395
await asyncio.wait_for(mcp_sessions, timeout=MCP_CLEANUP_TIMEOUT)
399-
except asyncio.TimeoutError as e:
396+
except asyncio.TimeoutError:
400397
continue
401398
except Exception as e:
402-
logging.error(f"Exception in mcp server cleanup task: {e}")
399+
logging.exception("Exception in mcp server cleanup task")
403400

404401

405402
async def main(available_tools: AvailableTools, p: str | None, t: str | None, cli_globals: dict, prompt: str | None):
@@ -581,15 +578,15 @@ def preprocess_prompt(prompt: str, tag: str, kv: Callable[[str], dict], kv_subke
581578
async def run_prompts(async_task=False, max_concurrent_tasks=5):
582579
# if this is a shell task, execute that and append the results
583580
if run:
584-
await render_model_output(f"** 🤖🐚 Executing Shell Task\n")
581+
await render_model_output("** 🤖🐚 Executing Shell Task\n")
585582
# this allows e.g. shell based jq output to become available for repeat prompts
586583
try:
587584
result = shell_tool_call(run).content[0].model_dump_json()
588585
last_mcp_tool_results.append(result)
589586
return True
590587
except RuntimeError as e:
591588
await render_model_output(f"** 🤖❗ Shell Task Exception: {e}\n")
592-
logging.error(f"Shell task error: {e}")
589+
logging.exception("Shell task error")
593590
return False
594591

595592
tasks = []

src/seclab_taskflow_agent/agent.py

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,34 +2,32 @@
22
# SPDX-License-Identifier: MIT
33

44
# https://openai.github.io/openai-agents-python/agents/
5-
import os
65
import logging
7-
from dotenv import load_dotenv, find_dotenv
6+
import os
87
from collections.abc import Callable
98
from typing import Any
109
from urllib.parse import urlparse
1110

12-
from openai import AsyncOpenAI
13-
from agents.agent import ModelSettings, ToolsToFinalOutputResult, FunctionToolResult
14-
from agents.run import DEFAULT_MAX_TURNS
15-
from agents.run import RunHooks
1611
from agents import (
1712
Agent,
18-
Runner,
1913
AgentHooks,
20-
RunHooks,
21-
result,
22-
function_tool,
23-
Tool,
14+
OpenAIChatCompletionsModel,
2415
RunContextWrapper,
16+
RunHooks,
17+
Runner,
2518
TContext,
26-
OpenAIChatCompletionsModel,
27-
set_default_openai_client,
19+
Tool,
20+
result,
2821
set_default_openai_api,
22+
set_default_openai_client,
2923
set_tracing_disabled,
3024
)
25+
from agents.agent import FunctionToolResult, ModelSettings, ToolsToFinalOutputResult
26+
from agents.run import DEFAULT_MAX_TURNS, RunHooks
27+
from dotenv import find_dotenv, load_dotenv
28+
from openai import AsyncOpenAI
3129

32-
from .capi import COPILOT_INTEGRATION_ID, get_AI_endpoint, get_AI_token, AI_API_ENDPOINT_ENUM
30+
from .capi import AI_API_ENDPOINT_ENUM, COPILOT_INTEGRATION_ID, get_AI_endpoint, get_AI_token
3331

3432
# grab our secrets from .env, this must be in .gitignore
3533
load_dotenv(find_dotenv(usecwd=True))

src/seclab_taskflow_agent/available_tools.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
# SPDX-FileCopyrightText: 2025 GitHub
22
# SPDX-License-Identifier: MIT
33

4-
from enum import Enum
5-
import logging
64
import importlib.resources
5+
from enum import Enum
6+
77
import yaml
88

99

src/seclab_taskflow_agent/capi.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,14 @@
22
# SPDX-License-Identifier: MIT
33

44
# CAPI specific interactions
5-
import httpx
65
import json
76
import logging
87
import os
9-
from strenum import StrEnum
108
from urllib.parse import urlparse
119

10+
import httpx
11+
from strenum import StrEnum
12+
1213

1314
# Enumeration of currently supported API endpoints.
1415
class AI_API_ENDPOINT_ENUM(StrEnum):
@@ -96,11 +97,11 @@ def list_capi_models(token: str) -> dict[str, dict]:
9697
for model in models_list:
9798
models[model.get("id")] = dict(model)
9899
except httpx.RequestError as e:
99-
logging.error(f"Request error: {e}")
100+
logging.exception("Request error")
100101
except json.JSONDecodeError as e:
101-
logging.error(f"JSON error: {e}")
102+
logging.exception("JSON error")
102103
except httpx.HTTPStatusError as e:
103-
logging.error(f"HTTP error: {e}")
104+
logging.exception("HTTP error")
104105
return models
105106

106107

src/seclab_taskflow_agent/env_utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
# SPDX-FileCopyrightText: 2025 GitHub
22
# SPDX-License-Identifier: MIT
33

4-
import re
54
import os
5+
import re
66

77

88
def swap_env(s):

0 commit comments

Comments
 (0)