Spaces:
Sleeping
Sleeping
File size: 6,078 Bytes
92c4ae6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | import asyncio
import json
import logging
import os
from typing import Any, Dict, List, Optional
from ..core.models import Finding, Severity
from core.llm_service import LLMService
logger = logging.getLogger(__name__)
class LLMAnalyzer:
"""
Open-source LLM analyzer for security scanning.
Supports:
- BYOK (Bring Your Own Key) for OpenAI and Anthropic.
- Local CPU inference for privacy and offline use.
"""
def __init__(
self,
mode: str = "local",
model: Optional[str] = None,
api_key: Optional[str] = None,
provider: Optional[str] = None
):
"""
Initialize the LLM Analyzer.
Args:
mode: "local" or "byok"
model: Model name (e.g., "gpt-4o", "claude-3-5-sonnet", or local model path)
api_key: API key for the provider
provider: "openai" or "anthropic" (for byok mode)
"""
self.mode = mode
self.model = model or ("Qwen/Qwen2.5-1.5B-Instruct" if mode == "local" else "gpt-4o")
self.api_key = api_key or os.getenv("ATOM_SECURITY_LLM_API_KEY")
self.provider = provider or os.getenv("ATOM_SECURITY_LLM_PROVIDER", "openai")
self.pipeline = None
# Initialize LLMService for unified LLM interactions (replaces direct clients)
self.llm_service = LLMService(workspace_id="default")
if self.mode == "local":
self._init_local()
else:
self._init_byok()
def _init_local(self):
"""Initialize local transformers pipeline."""
try:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
logger.info(f"Loading local model: {self.model}...")
self.tokenizer = AutoTokenizer.from_pretrained(self.model, trust_remote_code=True)
self.model_obj = AutoModelForCausalLM.from_pretrained(
self.model,
device_map="cpu",
torch_dtype=torch.float32,
trust_remote_code=True
)
self.pipeline = pipeline(
"text-generation",
model=self.model_obj,
tokenizer=self.tokenizer,
max_new_tokens=512,
temperature=0.1
)
except Exception as e:
logger.error(f"Failed to load local model: {e}")
raise
def _init_byok(self):
"""
Initialize BYOK mode using LLMService.
LLMService handles provider selection, API key resolution,
and client creation internally via BYOKHandler.
"""
# LLMService initialized in __init__ handles all BYOK configuration
# No direct client creation needed
pass
async def analyze(self, skill_name: str, content: str) -> List[Finding]:
"""Run analysis on skill content."""
system_prompt = (
"You are a security expert. Analyze the AI agent skill for:\n"
"1. Prompt Injection\n2. Code Injection\n3. Data Exfiltration\n\n"
"Return JSON: {\"findings\": [{\"category\": \"...\", \"severity\": \"...\", \"description\": \"...\"}]}"
)
user_prompt = f"Skill: {skill_name}\n\nContent:\n{content[:4000]}"
if self.mode == "local":
return await self._analyze_local(system_prompt, user_prompt)
else:
return await self._analyze_byok(system_prompt, user_prompt)
async def _analyze_local(self, system_prompt: str, user_prompt: str) -> List[Finding]:
"""Local inference."""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
outputs = await asyncio.to_thread(self.pipeline, prompt)
text = outputs[0]["generated_text"].replace(prompt, "")
return self._parse_json(text)
async def _analyze_byok(self, system_prompt: str, user_prompt: str) -> List[Finding]:
"""
BYOK API call via LLMService.
Uses unified LLMService interface for all providers (OpenAI, Anthropic).
LLMService handles provider selection, API key resolution, and cost tracking.
"""
# Build messages in OpenAI format
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
# Use model parameter (gpt-4o, claude-3-5-sonnet, etc.)
# Note: response_format not supported yet, JSON mode requested in system prompt instead
response = await self.llm_service.generate_completion(
messages=messages,
model=self.model,
temperature=0.1,
max_tokens=1024
)
# Extract content from LLMService response format
text = response.get("content", "")
return self._parse_json(text)
def _parse_json(self, text: str) -> List[Finding]:
"""Parse findings from LLM output."""
try:
# Simple cleanup for markdown
if "```json" in text:
text = text.split("```json")[1].split("```")[0]
elif "```" in text:
text = text.split("```")[1].split("```")[0]
data = json.loads(text)
findings = []
for f in data.get("findings", []):
findings.append(Finding(
rule_id=f.get("category", "LLM_DETECTED"),
category=f.get("category", "OTHER"),
severity=Severity(f.get("severity", "MEDIUM").upper()),
title=f.get("category", "Security issue"),
description=f.get("description", ""),
analyzer="llm"
))
return findings
except Exception as e:
logger.warning(f"Failed to parse LLM response: {e}")
return []
|