Architecture Proposal: Enhancing Extensibility and Multi-LLM Support

[pi-2r]

Hi team !

I just discovered your work this week during a technical workshop on LLM security at Devoxx France, led by Xavier and Joseph. First of all, thank you for open-sourcing this framework. The approach to vulnerability triage using AI and MCPs is fantastic, and it’s great to see the project gaining so much traction!

As the community starts experimenting with the agent (especially those of us looking to run it with local LLMs or custom MCP transports), I’ve been studying the codebase and identified a few architectural tweaks that could make the framework even more extensible and contributor-friendly.

I would love to help implement some of these structural improvements if you are open to it. Here are a few proposals:

1. Modularizing the Execution Engine (runner.py)

runner.py currently handles a lot of responsibilities (model resolution, Jinja2 templating, session state, and async execution).

• Proposal: Extracting the Jinja2 prompt building and the model resolution into their own dedicated modules (e.g., prompt_builder.py). This would make unit testing prompts much easier without needing to mock the entire OpenAI execution loop.

2. Facilitating Custom LLM Providers (Dependency Injection)

Currently, in agent.py, the initialization of the OpenAI client is tightly coupled with GitHub's infrastructure (via capi.py).

• Proposal: We could decouple this by injecting an already configured AsyncOpenAI client (or an abstract interface) into the TaskAgent. This would make it incredibly easy for the community to plug in local models (like Ollama) or other cloud providers without altering the core agent logic.

3. Dynamic Registry for MCP Transports (mcp_lifecycle.py)

The build_mcp_servers function currently uses a static match/case to handle transport types (stdio, sse).

• Proposal: Implementing a simple Registry or Factory pattern here. This way, new custom MCP transports could just "register" themselves without requiring modifications to the core mcp_lifecycle.py file.

4. Streamlining API Retries

The retry and backoff logic in runner.py is currently handled via manual while loops.

• Proposal: Delegating this to a robust library like tenacity. Using decorators like @retry would remove dozens of lines of state-management code and make the network calls much cleaner.

I understand the team might have its own roadmap, so I don't want to drop a massive PR out of nowhere.

Would you be open to a Draft PR starting with point 1 ?

Thanks for your time and the great work !