|
| 1 | +package de.tuda.stg.securecoder.engine.guardian |
| 2 | + |
| 3 | +import de.tuda.stg.securecoder.engine.llm.ChatMessage |
| 4 | +import de.tuda.stg.securecoder.engine.llm.LlmClient |
| 5 | +import de.tuda.stg.securecoder.engine.llm.chatStructured |
| 6 | +import de.tuda.stg.securecoder.guardian.AnalyzeRequest |
| 7 | +import de.tuda.stg.securecoder.guardian.AnalyzeResponse |
| 8 | +import de.tuda.stg.securecoder.guardian.Guardian |
| 9 | + |
| 10 | +class LlmGuardian( |
| 11 | + private val client: LlmClient, |
| 12 | + private val systemPrompt: String = DEFAULT_SYSTEM_PROMPT, |
| 13 | +) : Guardian { |
| 14 | + override suspend fun run(req: AnalyzeRequest): AnalyzeResponse { |
| 15 | + val messages = buildMessages(req) |
| 16 | + val llmResp = client.chatStructured<LlmAnalyzeResponse>( |
| 17 | + messages = messages, |
| 18 | + params = LlmClient.GenerationParams( |
| 19 | + temperature = 0.0 |
| 20 | + ) |
| 21 | + ) |
| 22 | + return llmResp.toApi() |
| 23 | + } |
| 24 | + |
| 25 | + private fun buildMessages(req: AnalyzeRequest): List<ChatMessage> { |
| 26 | + return listOf( |
| 27 | + ChatMessage(ChatMessage.Role.System, systemPrompt), |
| 28 | + ChatMessage(ChatMessage.Role.User, buildString { |
| 29 | + appendLine("You are given a set of source files to analyze for security issues.") |
| 30 | + appendLine("Only consider the provided files; do not assume hidden context.") |
| 31 | + appendLine() |
| 32 | + req.files.forEach { file -> |
| 33 | + appendLine("===== FILE: ${file.name} =====") |
| 34 | + appendLine(withLineNumbers(file.content)) |
| 35 | + appendLine("===== END FILE: ${file.name} =====") |
| 36 | + appendLine() |
| 37 | + } |
| 38 | + appendLine("Return your analysis strictly using the structured schema provided by the tool.") |
| 39 | + }) |
| 40 | + ) |
| 41 | + } |
| 42 | + |
| 43 | + private fun withLineNumbers(text: String): String = buildString { |
| 44 | + text.lineSequence().forEachIndexed { idx, line -> |
| 45 | + append(idx + 1) |
| 46 | + append(": ") |
| 47 | + append(line) |
| 48 | + append('\n') |
| 49 | + } |
| 50 | + }.removeSuffix("\n") |
| 51 | + |
| 52 | + companion object { |
| 53 | + private const val DEFAULT_SYSTEM_PROMPT: String = |
| 54 | + """ |
| 55 | + You are SecureCoder Guardian. Analyze code for security vulnerabilities. |
| 56 | + Use conservative judgment; highlight clear issues or suspicious patterns. |
| 57 | + Provide precise file and line locations when possible. If unsure, leave |
| 58 | + optional fields null. Do not include any prose outside the structured result. |
| 59 | + """ |
| 60 | + } |
| 61 | +} |
0 commit comments