mirror of
https://github.com/KeygraphHQ/shannon.git
synced 2026-05-24 17:44:17 +02:00
feat: add three-tier model system with Bedrock support
Introduce small/medium/large model tiers so agents use the appropriate model for their task complexity. Pre-recon uses Opus (large) for deep source code analysis, most agents use Sonnet (medium), and report uses Haiku (small) for summarization. - Add src/ai/models.ts with ModelTier type and resolveModel() - Add modelTier field to AgentDefinition - Refactor claude-executor env var passthrough into loop - Add Bedrock credential validation in preflight and CLI - Pass through Bedrock and model env vars in docker-compose
This commit is contained in:
@@ -26,6 +26,28 @@ ANTHROPIC_API_KEY=your-api-key-here
|
||||
# OPENROUTER_API_KEY=sk-or-your-openrouter-key
|
||||
# ROUTER_DEFAULT=openrouter,google/gemini-3-flash-preview
|
||||
|
||||
# =============================================================================
|
||||
# Model Tier Overrides (Anthropic API / OAuth / Bedrock)
|
||||
# =============================================================================
|
||||
# Override which model is used for each tier. Defaults are used if not set.
|
||||
# ANTHROPIC_SMALL_MODEL=... # Small tier (default: claude-haiku-4-5-20251001)
|
||||
# ANTHROPIC_MEDIUM_MODEL=... # Medium tier (default: claude-sonnet-4-6)
|
||||
# ANTHROPIC_LARGE_MODEL=... # Large tier (default: claude-opus-4-6)
|
||||
|
||||
# =============================================================================
|
||||
# OPTION 3: AWS Bedrock
|
||||
# =============================================================================
|
||||
# https://aws.amazon.com/blogs/machine-learning/accelerate-ai-development-with-amazon-bedrock-api-keys/
|
||||
# Requires the model tier overrides above to be set with Bedrock-specific model IDs.
|
||||
# Example Bedrock model IDs for us-east-1:
|
||||
# ANTHROPIC_SMALL_MODEL=us.anthropic.claude-haiku-4-5-20251001-v1:0
|
||||
# ANTHROPIC_MEDIUM_MODEL=us.anthropic.claude-sonnet-4-6
|
||||
# ANTHROPIC_LARGE_MODEL=us.anthropic.claude-opus-4-6
|
||||
|
||||
# CLAUDE_CODE_USE_BEDROCK=1
|
||||
# AWS_REGION=us-east-1
|
||||
# AWS_BEARER_TOKEN_BEDROCK=your-bearer-token
|
||||
|
||||
# =============================================================================
|
||||
# Available Models
|
||||
# =============================================================================
|
||||
|
||||
@@ -87,6 +87,7 @@ Shannon is available in two editions:
|
||||
- [Usage Examples](#usage-examples)
|
||||
- [Workspaces and Resuming](#workspaces-and-resuming)
|
||||
- [Configuration (Optional)](#configuration-optional)
|
||||
- [AWS Bedrock](#aws-bedrock)
|
||||
- [[EXPERIMENTAL - UNSUPPORTED] Router Mode (Alternative Providers)](#experimental---unsupported-router-mode-alternative-providers)
|
||||
- [Output and Results](#output-and-results)
|
||||
- [Sample Reports](#-sample-reports)
|
||||
@@ -107,6 +108,7 @@ Shannon is available in two editions:
|
||||
- **AI Provider Credentials** (choose one):
|
||||
- **Anthropic API key** (recommended) - Get from [Anthropic Console](https://console.anthropic.com)
|
||||
- **Claude Code OAuth token**
|
||||
- **AWS Bedrock** - Route through Amazon Bedrock with AWS credentials (see [AWS Bedrock](#aws-bedrock))
|
||||
- **[EXPERIMENTAL - UNSUPPORTED] Alternative providers via Router Mode** - OpenAI or Google Gemini via OpenRouter (see [Router Mode](#experimental---unsupported-router-mode-alternative-providers))
|
||||
|
||||
### Quick Start
|
||||
@@ -348,6 +350,33 @@ pipeline:
|
||||
|
||||
`max_concurrent_pipelines` controls how many vulnerability pipelines run simultaneously (1-5, default: 5). Lower values reduce the chance of hitting rate limits but increase wall-clock time.
|
||||
|
||||
### AWS Bedrock
|
||||
|
||||
Shannon also supports [Amazon Bedrock](https://aws.amazon.com/bedrock/) instead of using an Anthropic API key.
|
||||
|
||||
#### Quick Setup
|
||||
|
||||
1. Add your AWS credentials to `.env`:
|
||||
|
||||
```bash
|
||||
CLAUDE_CODE_USE_BEDROCK=1
|
||||
AWS_REGION=us-east-1
|
||||
AWS_BEARER_TOKEN_BEDROCK=your-bearer-token
|
||||
|
||||
# Set models with Bedrock-specific IDs for your region
|
||||
ANTHROPIC_SMALL_MODEL=us.anthropic.claude-haiku-4-5-20251001-v1:0
|
||||
ANTHROPIC_MEDIUM_MODEL=us.anthropic.claude-sonnet-4-6
|
||||
ANTHROPIC_LARGE_MODEL=us.anthropic.claude-opus-4-6
|
||||
```
|
||||
|
||||
2. Run Shannon as usual:
|
||||
|
||||
```bash
|
||||
./shannon start URL=https://example.com REPO=repo-name
|
||||
```
|
||||
|
||||
Shannon uses three model tiers: **small** (`claude-haiku-4-5-20251001`) for summarization, **medium** (`claude-sonnet-4-6`) for security analysis, and **large** (`claude-opus-4-6`) for deep reasoning. Set `ANTHROPIC_SMALL_MODEL`, `ANTHROPIC_MEDIUM_MODEL`, and `ANTHROPIC_LARGE_MODEL` to the Bedrock model IDs for your region.
|
||||
|
||||
### [EXPERIMENTAL - UNSUPPORTED] Router Mode (Alternative Providers)
|
||||
|
||||
Shannon can experimentally route requests through alternative AI providers using claude-code-router. This mode is not officially supported and is intended primarily for:
|
||||
|
||||
@@ -24,6 +24,12 @@ services:
|
||||
- ANTHROPIC_AUTH_TOKEN=${ANTHROPIC_AUTH_TOKEN:-} # Auth token for router
|
||||
- ROUTER_DEFAULT=${ROUTER_DEFAULT:-} # Model name when using router (e.g., "gemini,gemini-2.5-pro")
|
||||
- CLAUDE_CODE_OAUTH_TOKEN=${CLAUDE_CODE_OAUTH_TOKEN:-}
|
||||
- CLAUDE_CODE_USE_BEDROCK=${CLAUDE_CODE_USE_BEDROCK:-}
|
||||
- AWS_REGION=${AWS_REGION:-}
|
||||
- AWS_BEARER_TOKEN_BEDROCK=${AWS_BEARER_TOKEN_BEDROCK:-}
|
||||
- ANTHROPIC_SMALL_MODEL=${ANTHROPIC_SMALL_MODEL:-}
|
||||
- ANTHROPIC_MEDIUM_MODEL=${ANTHROPIC_MEDIUM_MODEL:-}
|
||||
- ANTHROPIC_LARGE_MODEL=${ANTHROPIC_LARGE_MODEL:-}
|
||||
- CLAUDE_CODE_MAX_OUTPUT_TOKENS=${CLAUDE_CODE_MAX_OUTPUT_TOKENS:-64000}
|
||||
depends_on:
|
||||
temporal:
|
||||
|
||||
@@ -142,14 +142,27 @@ cmd_start() {
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check for API key (router mode can use alternative provider API keys)
|
||||
# Check for API key (Bedrock and router modes can bypass this)
|
||||
if [ -z "$ANTHROPIC_API_KEY" ] && [ -z "$CLAUDE_CODE_OAUTH_TOKEN" ]; then
|
||||
if [ "$ROUTER" = "true" ] && { [ -n "$OPENAI_API_KEY" ] || [ -n "$OPENROUTER_API_KEY" ]; }; then
|
||||
if [ "$CLAUDE_CODE_USE_BEDROCK" = "1" ]; then
|
||||
# Bedrock mode — validate required AWS credentials
|
||||
MISSING=""
|
||||
[ -z "$AWS_REGION" ] && MISSING="$MISSING AWS_REGION"
|
||||
[ -z "$AWS_BEARER_TOKEN_BEDROCK" ] && MISSING="$MISSING AWS_BEARER_TOKEN_BEDROCK"
|
||||
[ -z "$ANTHROPIC_SMALL_MODEL" ] && MISSING="$MISSING ANTHROPIC_SMALL_MODEL"
|
||||
[ -z "$ANTHROPIC_MEDIUM_MODEL" ] && MISSING="$MISSING ANTHROPIC_MEDIUM_MODEL"
|
||||
[ -z "$ANTHROPIC_LARGE_MODEL" ] && MISSING="$MISSING ANTHROPIC_LARGE_MODEL"
|
||||
if [ -n "$MISSING" ]; then
|
||||
echo "ERROR: Bedrock mode requires the following env vars in .env:$MISSING"
|
||||
exit 1
|
||||
fi
|
||||
elif [ "$ROUTER" = "true" ] && { [ -n "$OPENAI_API_KEY" ] || [ -n "$OPENROUTER_API_KEY" ]; }; then
|
||||
# Router mode with alternative provider - set a placeholder for SDK init
|
||||
export ANTHROPIC_API_KEY="router-mode"
|
||||
else
|
||||
echo "ERROR: Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env"
|
||||
echo " (or use ROUTER=true with OPENAI_API_KEY or OPENROUTER_API_KEY)"
|
||||
echo " (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock,"
|
||||
echo " or ROUTER=true with OPENAI_API_KEY or OPENROUTER_API_KEY)"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
+20
-13
@@ -24,6 +24,7 @@ import { detectExecutionContext, formatErrorOutput, formatCompletionMessage } fr
|
||||
import { createProgressManager } from './progress-manager.js';
|
||||
import { createAuditLogger } from './audit-logger.js';
|
||||
import { getActualModelName } from './router-utils.js';
|
||||
import { resolveModel, type ModelTier } from './models.js';
|
||||
import type { ActivityLogger } from '../types/activity-logger.js';
|
||||
|
||||
declare global {
|
||||
@@ -202,7 +203,8 @@ export async function runClaudePrompt(
|
||||
description: string = 'Claude analysis',
|
||||
agentName: string | null = null,
|
||||
auditSession: AuditSession | null = null,
|
||||
logger: ActivityLogger
|
||||
logger: ActivityLogger,
|
||||
modelTier: ModelTier = 'medium'
|
||||
): Promise<ClaudePromptResult> {
|
||||
// 1. Initialize timing and prompt
|
||||
const timer = new Timer(`agent-${description.toLowerCase().replace(/\s+/g, '-')}`);
|
||||
@@ -225,22 +227,27 @@ export async function runClaudePrompt(
|
||||
const sdkEnv: Record<string, string> = {
|
||||
CLAUDE_CODE_MAX_OUTPUT_TOKENS: process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS || '64000',
|
||||
};
|
||||
if (process.env.ANTHROPIC_API_KEY) {
|
||||
sdkEnv.ANTHROPIC_API_KEY = process.env.ANTHROPIC_API_KEY;
|
||||
}
|
||||
if (process.env.CLAUDE_CODE_OAUTH_TOKEN) {
|
||||
sdkEnv.CLAUDE_CODE_OAUTH_TOKEN = process.env.CLAUDE_CODE_OAUTH_TOKEN;
|
||||
}
|
||||
if (process.env.ANTHROPIC_BASE_URL) {
|
||||
sdkEnv.ANTHROPIC_BASE_URL = process.env.ANTHROPIC_BASE_URL;
|
||||
}
|
||||
if (process.env.ANTHROPIC_AUTH_TOKEN) {
|
||||
sdkEnv.ANTHROPIC_AUTH_TOKEN = process.env.ANTHROPIC_AUTH_TOKEN;
|
||||
const passthroughVars = [
|
||||
'ANTHROPIC_API_KEY',
|
||||
'CLAUDE_CODE_OAUTH_TOKEN',
|
||||
'ANTHROPIC_BASE_URL',
|
||||
'ANTHROPIC_AUTH_TOKEN',
|
||||
'CLAUDE_CODE_USE_BEDROCK',
|
||||
'AWS_REGION',
|
||||
'AWS_BEARER_TOKEN_BEDROCK',
|
||||
'ANTHROPIC_SMALL_MODEL',
|
||||
'ANTHROPIC_MEDIUM_MODEL',
|
||||
'ANTHROPIC_LARGE_MODEL',
|
||||
];
|
||||
for (const name of passthroughVars) {
|
||||
if (process.env[name]) {
|
||||
sdkEnv[name] = process.env[name]!;
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Configure SDK options
|
||||
const options = {
|
||||
model: 'claude-sonnet-4-5-20250929',
|
||||
model: resolveModel(modelTier),
|
||||
maxTurns: 10_000,
|
||||
cwd: sourceDir,
|
||||
permissionMode: 'bypassPermissions' as const,
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
// Copyright (C) 2025 Keygraph, Inc.
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License version 3
|
||||
// as published by the Free Software Foundation.
|
||||
|
||||
/**
|
||||
* Model tier definitions and resolution.
|
||||
*
|
||||
* Three tiers mapped to capability levels:
|
||||
* - "small" (Haiku — summarization, structured extraction)
|
||||
* - "medium" (Sonnet — tool use, general analysis)
|
||||
* - "large" (Opus — deep reasoning, complex analysis)
|
||||
*
|
||||
* Users override via ANTHROPIC_SMALL_MODEL / ANTHROPIC_MEDIUM_MODEL / ANTHROPIC_LARGE_MODEL,
|
||||
* which works across all providers (direct, Bedrock, Vertex).
|
||||
*/
|
||||
|
||||
export type ModelTier = 'small' | 'medium' | 'large';
|
||||
|
||||
const DEFAULT_MODELS: Readonly<Record<ModelTier, string>> = {
|
||||
small: 'claude-haiku-4-5-20251001',
|
||||
medium: 'claude-sonnet-4-6',
|
||||
large: 'claude-opus-4-6',
|
||||
};
|
||||
|
||||
/** Resolve a model tier to a concrete model ID. */
|
||||
export function resolveModel(tier: ModelTier = 'medium'): string {
|
||||
switch (tier) {
|
||||
case 'small':
|
||||
return process.env.ANTHROPIC_SMALL_MODEL || DEFAULT_MODELS.small;
|
||||
case 'large':
|
||||
return process.env.ANTHROPIC_LARGE_MODEL || DEFAULT_MODELS.large;
|
||||
default:
|
||||
return process.env.ANTHROPIC_MEDIUM_MODEL || DEFAULT_MODELS.medium;
|
||||
}
|
||||
}
|
||||
@@ -156,7 +156,8 @@ export class AgentExecutionService {
|
||||
agentName, // description
|
||||
agentName,
|
||||
auditSession,
|
||||
logger
|
||||
logger,
|
||||
AGENTS[agentName].modelTier
|
||||
);
|
||||
|
||||
// 6. Spending cap check - defense-in-depth
|
||||
|
||||
@@ -24,6 +24,7 @@ import { PentestError, isRetryableError } from './error-handling.js';
|
||||
import { ErrorCode } from '../types/errors.js';
|
||||
import { type Result, ok, err } from '../types/result.js';
|
||||
import { parseConfig } from '../config-parser.js';
|
||||
import { resolveModel } from '../ai/models.js';
|
||||
import type { ActivityLogger } from '../types/activity-logger.js';
|
||||
|
||||
// === Repository Validation ===
|
||||
@@ -165,11 +166,30 @@ async function validateCredentials(
|
||||
return ok(undefined);
|
||||
}
|
||||
|
||||
// 2. Check that at least one credential is present
|
||||
// 2. Bedrock mode — validate required AWS credentials are present
|
||||
if (process.env.CLAUDE_CODE_USE_BEDROCK === '1') {
|
||||
const required = ['AWS_REGION', 'AWS_BEARER_TOKEN_BEDROCK', 'ANTHROPIC_SMALL_MODEL', 'ANTHROPIC_MEDIUM_MODEL', 'ANTHROPIC_LARGE_MODEL'];
|
||||
const missing = required.filter(v => !process.env[v]);
|
||||
if (missing.length > 0) {
|
||||
return err(
|
||||
new PentestError(
|
||||
`Bedrock mode requires the following env vars in .env: ${missing.join(', ')}`,
|
||||
'config',
|
||||
false,
|
||||
{ missing },
|
||||
ErrorCode.AUTH_FAILED
|
||||
)
|
||||
);
|
||||
}
|
||||
logger.info('Bedrock credentials OK');
|
||||
return ok(undefined);
|
||||
}
|
||||
|
||||
// 3. Check that at least one credential is present
|
||||
if (!process.env.ANTHROPIC_API_KEY && !process.env.CLAUDE_CODE_OAUTH_TOKEN) {
|
||||
return err(
|
||||
new PentestError(
|
||||
'No API credentials found. Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env',
|
||||
'No API credentials found. Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock)',
|
||||
'config',
|
||||
false,
|
||||
{},
|
||||
@@ -178,12 +198,12 @@ async function validateCredentials(
|
||||
);
|
||||
}
|
||||
|
||||
// 3. Validate via SDK query
|
||||
// 4. Validate via SDK query
|
||||
const authType = process.env.CLAUDE_CODE_OAUTH_TOKEN ? 'OAuth token' : 'API key';
|
||||
logger.info(`Validating ${authType} via SDK...`);
|
||||
|
||||
try {
|
||||
for await (const message of query({ prompt: 'hi', options: { model: 'claude-haiku-4-5-20251001', maxTurns: 1 } })) {
|
||||
for await (const message of query({ prompt: 'hi', options: { model: resolveModel('small'), maxTurns: 1 } })) {
|
||||
if (message.type === 'assistant' && message.error) {
|
||||
return classifySdkError(message.error, authType);
|
||||
}
|
||||
|
||||
@@ -18,6 +18,7 @@ export const AGENTS: Readonly<Record<AgentName, AgentDefinition>> = Object.freez
|
||||
prerequisites: [],
|
||||
promptTemplate: 'pre-recon-code',
|
||||
deliverableFilename: 'code_analysis_deliverable.md',
|
||||
modelTier: 'large',
|
||||
},
|
||||
'recon': {
|
||||
name: 'recon',
|
||||
@@ -102,6 +103,7 @@ export const AGENTS: Readonly<Record<AgentName, AgentDefinition>> = Object.freez
|
||||
prerequisites: ['injection-exploit', 'xss-exploit', 'auth-exploit', 'ssrf-exploit', 'authz-exploit'],
|
||||
promptTemplate: 'report-executive',
|
||||
deliverableFilename: 'comprehensive_security_assessment_report.md',
|
||||
modelTier: 'small',
|
||||
},
|
||||
});
|
||||
|
||||
|
||||
@@ -58,6 +58,7 @@ export interface AgentDefinition {
|
||||
prerequisites: AgentName[];
|
||||
promptTemplate: string;
|
||||
deliverableFilename: string;
|
||||
modelTier?: 'small' | 'medium' | 'large';
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
Reference in New Issue
Block a user