-
Notifications
You must be signed in to change notification settings - Fork 127
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add LLM analyzers #508
Open
rbs-jacob
wants to merge
24
commits into
redballoonsecurity:master
Choose a base branch
from
rbs-jacob:feature/llm-analyzer
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Add LLM analyzers #508
Changes from all commits
Commits
Show all changes
24 commits
Select commit
Hold shift + click to select a range
2c067d0
Add initial version of LLM analyzer
rbs-jacob d9dcc13
Fix bugs when using in the UI
rbs-jacob e883681
Fix component configs default field=None bug
rbs-jacob 24c1124
Update CHANGELOG
rbs-jacob 9c34469
Merge branch 'bugfix/default-none-ui-components' into feature/llm-ana…
rbs-jacob 57ef1ba
Merge branch 'master' into feature/llm-analyzer
rbs-jacob 9642b36
Add first-pass AI analyzer UI
rbs-jacob 9776803
Fix frontend folder copying
rbs-jacob 87d9e7a
Add simple LLM test using Ollama
rbs-jacob 3ae9c3c
Remove print statement
rbs-jacob fdb4789
Fix type errors
rbs-jacob cd629fd
Add LLM function analyzer
rbs-jacob 486ee86
Improve function analyzer output quality
rbs-jacob 6b3d750
Add whole-program LLM analyzer
rbs-jacob cc6ca8d
Add tests for LLM program analyzer
rbs-jacob d2ea6af
Merge branch 'master' into feature/llm-analyzer
rbs-jacob 9ceb04e
Fix failing tests
rbs-jacob 546eaf8
Update CHANGELOG
rbs-jacob 025a694
Use the correct analyzer in the UI
rbs-jacob 76672e4
Add explanatory docstring with install instructions
rbs-jacob 3e1a1c1
Add LlmFunctionAnalyzer test case
rbs-jacob 12cd0bb
Assert existence of attributes.description in tests
rbs-jacob f60e905
Show system prompt in textarea
rbs-jacob af754e1
Lint (oops)
rbs-jacob File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,211 @@ | ||
| <style> | ||
| .container { | ||
| min-height: 100%; | ||
| max-height: 100%; | ||
| display: flex; | ||
| flex-direction: column; | ||
| flex-wrap: nowrap; | ||
| justify-content: center; | ||
| align-items: stretch; | ||
| align-content: center; | ||
| overflow: auto; | ||
| } | ||
|
|
||
| .inputs *:first-child { | ||
| margin-top: 0; | ||
| } | ||
|
|
||
| .output { | ||
| flex-grow: 1; | ||
| } | ||
|
|
||
| pre { | ||
| white-space: pre-wrap; | ||
| } | ||
|
|
||
| .actions { | ||
| margin-top: 2em; | ||
| display: flex; | ||
| flex-direction: row; | ||
| flex-wrap: nowrap; | ||
| justify-content: space-evenly; | ||
| align-items: center; | ||
| align-content: center; | ||
| } | ||
|
|
||
| input { | ||
| background: inherit; | ||
| color: inherit; | ||
| border: none; | ||
| border-bottom: 1px solid white; | ||
| flex-grow: 1; | ||
| margin-left: 1ch; | ||
| } | ||
|
|
||
| input:focus, | ||
| textarea:focus { | ||
| outline: none; | ||
| box-shadow: inset 0 -1px 0 var(--main-fg-color); | ||
| } | ||
|
|
||
| textarea { | ||
| background: inherit; | ||
| color: inherit; | ||
| border: 1px solid white; | ||
| flex-grow: 1; | ||
| margin-left: 1ch; | ||
| resize: vertical; | ||
| font-size: inherit; | ||
| } | ||
|
|
||
| select { | ||
| margin-bottom: 1em; | ||
| background-color: inherit; | ||
| color: inherit; | ||
| border: 1px solid var(--main-fg-color); | ||
| border-radius: 0; | ||
| font-size: inherit; | ||
| font-family: var(--font); | ||
| box-shadow: none; | ||
| } | ||
|
|
||
| select:focus { | ||
| outline: none; | ||
| box-shadow: inset 0 -1px 0 var(--main-fg-color); | ||
| } | ||
|
|
||
| label { | ||
| margin-bottom: 1em; | ||
| display: flex; | ||
| flex-direction: row; | ||
| flex-wrap: nowrap; | ||
| justify-content: space-evenly; | ||
| align-items: baseline; | ||
| align-content: center; | ||
| white-space: nowrap; | ||
| } | ||
| </style> | ||
|
|
||
| <script> | ||
| import Button from "../utils/Button.svelte"; | ||
| import LoadingText from "../utils/LoadingText.svelte"; | ||
|
|
||
| import { selectedResource } from "../stores.js"; | ||
| import { onMount } from "svelte"; | ||
|
|
||
| export let modifierView; | ||
| let apiUrl = JSON.parse(window.localStorage.getItem("aiApiUrl") ?? "null"), | ||
| model = JSON.parse(window.localStorage.getItem("aiApiModel") ?? "null"), | ||
| key = JSON.parse(window.localStorage.getItem("aiApiKey") ?? "null"), | ||
| prompt; | ||
| $: window.localStorage.setItem("aiApiUrl", JSON.stringify(apiUrl)); | ||
| $: window.localStorage.setItem("aiApiModel", JSON.stringify(model)); | ||
| $: window.localStorage.setItem("aiApiKey", JSON.stringify(key)); | ||
|
|
||
| function getSystemPrompt(r) { | ||
| return r.get_config_for_component("LlmAnalyzer").then(({ fields }) => { | ||
| for (const field of fields) { | ||
| if (field.name == "system_prompt") { | ||
| prompt = field.default; | ||
| return field.default; | ||
| } | ||
| } | ||
| }); | ||
| } | ||
| $: promptPromise = getSystemPrompt($selectedResource); | ||
|
|
||
| function getAnalyzer() { | ||
| if ($selectedResource.tags.includes("ofrak.core.program.Program")) { | ||
| return "LlmProgramAnalyzer"; | ||
| } | ||
| if ( | ||
| $selectedResource.tags.includes("ofrak.core.complex_block.ComplexBlock") | ||
| ) { | ||
| return "LlmFunctionAnalyzer"; | ||
| } | ||
| return "LlmAnalyzer"; | ||
| } | ||
| $: analyzer = getAnalyzer(); | ||
|
|
||
| let resultPromise = undefined; | ||
| function llmAnalyzer() { | ||
| resultPromise = fetch( | ||
| `${$selectedResource.uri}/run_component?component=${analyzer}`, | ||
| { | ||
| credentials: "omit", | ||
| body: JSON.stringify([ | ||
| "ofrak.core.llm.LlmAnalyzerConfig", | ||
| { | ||
| api_url: apiUrl, | ||
| model: model, | ||
| api_key: key || undefined, | ||
| ...(analyzer == "LlmAnalyzer" ? { system_prompt: prompt } : {}), | ||
| }, | ||
| ]), | ||
| method: "POST", | ||
| mode: "cors", | ||
| } | ||
| ) | ||
| .then(async (r) => { | ||
| if (!r.ok) { | ||
| throw Error(JSON.stringify(await r.json(), undefined, 2)); | ||
| } | ||
| return await r.json(); | ||
| }) | ||
| .then((r) => { | ||
| $selectedResource.attributes = $selectedResource.attributes; | ||
| return r; | ||
| }) | ||
| .then(({ modified }) => { | ||
| const { attributes } = modified.find( | ||
| ({ id }) => id == $selectedResource.resource_id | ||
| ); | ||
| for (const [type, [_, { description }]] of attributes) { | ||
| if (type == "ofrak.core.llm.LlmAttributes") { | ||
| return description; | ||
| } | ||
| } | ||
| }); | ||
| } | ||
| </script> | ||
|
|
||
| <div class="container"> | ||
| {#await promptPromise} | ||
| <LoadingText /> | ||
| {:then _} | ||
| <div class="inputs"> | ||
| <select bind:value="{analyzer}"> | ||
| <option value="LlmAnalyzer">LlmAnalyzer</option> | ||
| <option value="LlmFunctionAnalyzer">LlmFunctionAnalyzer</option> | ||
| <option value="LlmProgramAnalyzer">LlmProgramAnalyzer</option> | ||
| </select> | ||
| <label> | ||
| AI API URL | ||
| <input type="text" bind:value="{apiUrl}" /> | ||
| </label> | ||
| <label> | ||
| AI Model | ||
| <input type="text" bind:value="{model}" /> | ||
| </label> | ||
| <label> | ||
| AI API Key (Optional) | ||
| <input type="password" bind:value="{key}" /> | ||
| </label> | ||
| {#if analyzer == "LlmAnalyzer"} | ||
| <label> | ||
| System Prompt (Optional) | ||
| <textarea value="{prompt}"></textarea> | ||
| </label> | ||
| {/if} | ||
| </div> | ||
| {:catch e} | ||
| <p>{e}</p> | ||
| {/await} | ||
| <pre class="output">{#await resultPromise}<LoadingText | ||
| />{:then result}{#if result}{result}{/if}{:catch e}Try re-running the analyzer. | ||
| {e}{/await}</pre> | ||
| <div class="actions"> | ||
| <Button on:click="{llmAnalyzer}">Analyze</Button> | ||
| <Button on:click="{() => (modifierView = undefined)}">Cancel</Button> | ||
| </div> | ||
| </div> |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Are we sure we want to add this to the Docker image?
This feels like an incomplete install step -- we're adding a dependency that requires the user to download/install more things to actually use it. The impact is a larger image size for all users, plus more work for users who actually want to use the LLM feature.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
How would you propose doing it instead?
It would be ideal to have this whole thing in a separate package so that users could conditionally include the AI components. Then we could install Ollama and pull a model without worrying about unnecessarily inflating the core image. But one of the original constraints was that this should be in OFRAK core.
If we don't have Ollama installed, then we can't run the tests without having the tests themselves pull down and install the binary. As far as I know, in other OFRAK tests, there is no precedent for installing a dependency like this. It would be doable, but seems a little weird.
On the other hand, we could pull a model in the OFRAK core Dockerstub after installing Ollama, but this seems like a bad idea. We don't want to add another several-hundred-MB dependency to the OFRAK core images. We especially don't want to add models this large to the base image if the user is just going to use OpenAI anyway.
Installing Ollama, but not pulling a model, seemed like a reasonable tradeoff. If it seems wrong to you, I would appreciate guidance on a better way to do it.