feature: add Fireworks AI provider to KB generator - #345
feature: add Fireworks AI provider to KB generator#345ticket-fixer[bot] wants to merge 1 commit into
Conversation
Support the fireworks LLM provider via the --provider CLI option, backed by the FIREWORKS_API_KEY environment variable and the OpenAI-compatible Fireworks inference endpoint.
![ostorlab-ai-pr-review[bot]](https://avatars.githubusercontent.com/in/1072062?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review Summary
Total Issues Found: 3
Critical Issues (Bugs): 0
Security Issues: 0
Suggestions: 3
Estimated Review Effort: 2/5
Key Findings:
The PR contains 3 review findings, all non-critical. The most notable issue is in the tests: patching tenacity.nap.sleep does not affect the retry controller's stored sleep callable, so the missing-key test performs real 2-second sleeps between retries, making the suite unnecessarily slow. Relatedly, the retry decorator currently wraps a deterministic missing-API-key ValueError into a tenacity.RetryError, masking a clear configuration error as a transient-style retry failure. Finally, LLMProvider should be marked with @enum.unique to enforce distinct values. No critical, security, or data-corruption issues were identified.
| ) -> None: | ||
| """generate_kb with the fireworks provider and no API key raises a ValueError.""" | ||
| mocker.patch("tools.kb_generator.OpenAI") | ||
| mocker.patch("tenacity.nap.sleep") |
There was a problem hiding this comment.
Patching tenacity.nap.sleep has no effect on the retry loop. The Retrying controller created by @tenacity.retry stores the original tenacity.nap.sleep function as its sleep callable when generate_kb is decorated, so replacing the module attribute later does not change the sleep it performs. Patch time.sleep (which tenacity.nap.sleep calls dynamically) or kb_generator.generate_kb.retry.sleep instead, so the missing-key test does not perform two real 2-second sleeps between its three attempts.
| "Sql Injection", kb_generator.RiskRating.HIGH, kb_generator.Platform.WEB | ||
| ) | ||
|
|
||
| with pytest.raises(tenacity.RetryError) as excinfo: |
There was a problem hiding this comment.
This test codifies that a missing Fireworks API key surfaces as tenacity.RetryError after three attempts instead of an immediate ValueError. Since generate_kb is decorated with retry_if_exception_type() with no exception types, the deterministic missing-key ValueError is retried and wrapped. Validate the selected provider's API key before entering the retried function, or restrict the retry predicate to transient API failures, so callers receive the clear configuration error immediately.
| OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1" | ||
|
|
||
|
|
||
| class LLMProvider(enum.Enum): |
There was a problem hiding this comment.
Decorate LLMProvider with @enum.unique to enforce distinct provider values and follow the Python coding convention for enum classes.
1 participant
Summary
Adds the Fireworks AI inference provider to the KB generator so vulnerability descriptions, code examples and metadata can be generated with Fireworks-hosted models (e.g. Llama 3.3 70B) through their OpenAI-compatible endpoint, selected via the new
--provider fireworksCLI option and authenticated with theFIREWORKS_API_KEYenvironment variable.Changes
LLMProviderenum with the existingopenrouterdefault and the newfireworksprovider.fireworksrequests throughhttps://api.fireworks.ai/inference/v1with the modelaccounts/fireworks/models/llama-v3p3-70b-instruct, readingFIREWORKS_API_KEYat call time.--providerCLI option (defaults toopenrouter, preserving current behaviour).Fixes ticket https://report.ostorlab.co/o/os/remediation/tickets/os-36783