mirror of
https://github.com/Rarebuffalo/securelens-backend.git
synced 2026-06-19 07:00:30 +00:00
make securelens show usage when run with no subcommand and pass api_base in scan triage and analysis
This commit is contained in:
@@ -70,9 +70,10 @@ async def call_ai_json(
|
|||||||
api_key: str,
|
api_key: str,
|
||||||
model: str,
|
model: str,
|
||||||
temperature: float = 0.2,
|
temperature: float = 0.2,
|
||||||
|
api_base: Optional[str] = None,
|
||||||
) -> Optional[dict]:
|
) -> Optional[dict]:
|
||||||
"""Convenience wrapper — calls AI in JSON mode and parses the result."""
|
"""Convenience wrapper — calls AI in JSON mode and parses the result."""
|
||||||
raw = await call_ai(prompt, api_key, model, temperature=temperature, json_mode=True)
|
raw = await call_ai(prompt, api_key, model, temperature=temperature, json_mode=True, api_base=api_base)
|
||||||
if not raw:
|
if not raw:
|
||||||
return None
|
return None
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -52,7 +52,7 @@ def main(ctx):
|
|||||||
Scan codebases, URLs and get instant security reports.
|
Scan codebases, URLs and get instant security reports.
|
||||||
"""
|
"""
|
||||||
if ctx.invoked_subcommand is None:
|
if ctx.invoked_subcommand is None:
|
||||||
ctx.invoke(scan, path=".", model=None, output=None, max_files=None, ci=False, fail_on=None, no_ai=False, sync=False)
|
click.echo(ctx.get_help())
|
||||||
|
|
||||||
|
|
||||||
# ── configure ─────────────────────────────────────────────────────────────────
|
# ── configure ─────────────────────────────────────────────────────────────────
|
||||||
@@ -269,7 +269,7 @@ async def _scan_async(path, model, output, max_files, ci, fail_on, no_ai, sync):
|
|||||||
for v in vulnerabilities
|
for v in vulnerabilities
|
||||||
]
|
]
|
||||||
prompt = summary_prompt(str(root), _json.dumps(issues_data, indent=2))
|
prompt = summary_prompt(str(root), _json.dumps(issues_data, indent=2))
|
||||||
ai_summary = await call_ai(prompt, cfg.api_key, cfg.default_model, temperature=0.4)
|
ai_summary = await call_ai(prompt, cfg.api_key, cfg.default_model, temperature=0.4, api_base=cfg.api_base)
|
||||||
progress.update(task_summary, completed=100, total=100, detail="Done")
|
progress.update(task_summary, completed=100, total=100, detail="Done")
|
||||||
|
|
||||||
# ── Build result ─────────────────────────────────────────────────────────
|
# ── Build result ─────────────────────────────────────────────────────────
|
||||||
@@ -389,7 +389,7 @@ async def _web_async(url, model, output, ci, fail_on, no_ai):
|
|||||||
]
|
]
|
||||||
prompt = web_summary_prompt(url, _json.dumps(issues_data, indent=2),
|
prompt = web_summary_prompt(url, _json.dumps(issues_data, indent=2),
|
||||||
result.score, result.grade)
|
result.score, result.grade)
|
||||||
result.ai_summary = await call_ai(prompt, cfg.api_key, cfg.default_model, temperature=0.4)
|
result.ai_summary = await call_ai(prompt, cfg.api_key, cfg.default_model, temperature=0.4, api_base=cfg.api_base)
|
||||||
progress.update(task2, completed=100, total=100, detail="Done")
|
progress.update(task2, completed=100, total=100, detail="Done")
|
||||||
|
|
||||||
fmt = cfg.output_format
|
fmt = cfg.output_format
|
||||||
|
|||||||
@@ -195,7 +195,7 @@ async def triage_files(
|
|||||||
if rel_paths and remaining_budget > 0 and cfg.api_key:
|
if rel_paths and remaining_budget > 0 and cfg.api_key:
|
||||||
file_list_str = "\n".join(rel_paths[:300]) # cap to ~300 paths for token budget
|
file_list_str = "\n".join(rel_paths[:300]) # cap to ~300 paths for token budget
|
||||||
prompt = triage_prompt(file_list_str, remaining_budget)
|
prompt = triage_prompt(file_list_str, remaining_budget)
|
||||||
result = await call_ai_json(prompt, cfg.api_key, cfg.default_model, temperature=0.1)
|
result = await call_ai_json(prompt, cfg.api_key, cfg.default_model, temperature=0.1, api_base=cfg.api_base)
|
||||||
if result and "critical_files" in result:
|
if result and "critical_files" in result:
|
||||||
for rel in result["critical_files"]:
|
for rel in result["critical_files"]:
|
||||||
abs_path = root / rel
|
abs_path = root / rel
|
||||||
@@ -233,7 +233,7 @@ async def analyze_file(
|
|||||||
content = content[:30_000] + "\n... (truncated)"
|
content = content[:30_000] + "\n... (truncated)"
|
||||||
|
|
||||||
prompt = analysis_prompt(rel, content)
|
prompt = analysis_prompt(rel, content)
|
||||||
result = await call_ai_json(prompt, cfg.api_key, cfg.default_model, temperature=0.2)
|
result = await call_ai_json(prompt, cfg.api_key, cfg.default_model, temperature=0.2, api_base=cfg.api_base)
|
||||||
if not result:
|
if not result:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
|||||||
@@ -52,3 +52,45 @@ async def test_backend_call_ai_passes_api_base():
|
|||||||
finally:
|
finally:
|
||||||
settings.ai_api_key = original_key
|
settings.ai_api_key = original_key
|
||||||
settings.ai_api_base = original_base
|
settings.ai_api_base = original_base
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_triage_files_passes_api_base():
|
||||||
|
from securelens.scanners import triage_files
|
||||||
|
from securelens.config import CLIConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
cfg = CLIConfig()
|
||||||
|
cfg.api_key = "mock_key"
|
||||||
|
cfg.api_base = "https://agentrouter.org/v1"
|
||||||
|
cfg.default_model = "openai/deepseek-chat"
|
||||||
|
|
||||||
|
with patch("securelens.scanners.call_ai_json", new_callable=AsyncMock) as mock_call_ai_json:
|
||||||
|
mock_call_ai_json.return_value = {"critical_files": []}
|
||||||
|
|
||||||
|
await triage_files([Path("test.py")], Path("."), cfg)
|
||||||
|
|
||||||
|
mock_call_ai_json.assert_called_once()
|
||||||
|
called_kwargs = mock_call_ai_json.call_args[1]
|
||||||
|
assert called_kwargs["api_base"] == "https://agentrouter.org/v1"
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_analyze_file_passes_api_base():
|
||||||
|
from securelens.scanners import analyze_file
|
||||||
|
from securelens.config import CLIConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
cfg = CLIConfig()
|
||||||
|
cfg.api_key = "mock_key"
|
||||||
|
cfg.api_base = "https://agentrouter.org/v1"
|
||||||
|
cfg.default_model = "openai/deepseek-chat"
|
||||||
|
|
||||||
|
mock_file = Path("test.py")
|
||||||
|
with patch("pathlib.Path.read_text", return_value="print('hello')"):
|
||||||
|
with patch("securelens.scanners.call_ai_json", new_callable=AsyncMock) as mock_call_ai_json:
|
||||||
|
mock_call_ai_json.return_value = {"vulnerabilities": []}
|
||||||
|
|
||||||
|
await analyze_file(mock_file, Path("."), cfg)
|
||||||
|
|
||||||
|
mock_call_ai_json.assert_called_once()
|
||||||
|
called_kwargs = mock_call_ai_json.call_args[1]
|
||||||
|
assert called_kwargs["api_base"] == "https://agentrouter.org/v1"
|
||||||
|
|||||||
Reference in New Issue
Block a user