feat(wren-ai-service): create Streamlit UI for configuring LLM models #1690
Conversation
…_state.py; combine init and override behaviors (import yaml files)
1. Rename custom_llm_ui.py to app.py for better naming convention 2. Extract Embedder and Document Store configuration from app.py to ui_components.py to improve code organization 3. Fix data type issue: ensure "timeout" value is properly converted to integer 4. Optimize alias handling: only include alias field in dictionary when it has a value
- Implemented pipeline configuration UI. - Set init_pipeline in session_state.py to manage pipeline state.
…and embedding models
…guration handling - Removed the document store configuration section and its associated rendering logic. - Updated the embedder configuration to include dimension input for embedding models. - Simplified the rendering of pipeline configurations and improved the selection process for LLM models. - Enhanced the preview and generation functionality for configuration files. - Adjusted the layout and organization of UI components for better clarity and usability.
… component integration - Refactored configuration path retrieval by introducing dedicated functions for fetching paths to config.yaml, config.done, and .env files. - Updated UI components to utilize the new configuration path functions, enhancing maintainability and clarity. - Simplified file operations in the API key management section to ensure consistent access to the environment file.
…prove error handling - Implemented validation for required configuration blocks (LLM, Embedder, Document Store, Pipeline) to ensure they are present and not empty. - Enhanced error handling by replacing print statements with Streamlit error messages for better user feedback during configuration loading failures. - Improved user experience by warning users about missing configuration blocks and indicating that default values will be used.
…ontainer launch logging - Updated the config extraction logic to return the first block of each type, improving clarity and efficiency. - Modified the launch command to dynamically find and log the available port for the Streamlit UI, enhancing user feedback during startup.
…tecture support and digest handling
…ing configurations fix(ui): update save button key to include original index for correct LLM updates
WalkthroughAdds a Streamlit-based provider setup UI, related Python modules and Dockerfile, a GitHub Actions workflow to build/publish its multi-arch image, and launcher/docker utilities to run the UI container and integrate the generated config into the launch flow. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant CLI as Launcher (CLI)
participant FS as Local FS (~/.wrenai)
participant Docker
participant UI as Streamlit UI
User->>CLI: run launch with provider="custom"
CLI->>FS: TryGetWrenAIDir()
alt wrenai dir exists
CLI->>Docker: RemoveContainerIfExists("wrenai-providers-setup")
CLI->>Docker: buildStreamlitImage()
CLI->>FS: ensure config.yaml and .env (download if missing)
CLI->>Docker: runStreamlitContainer(port 8501, mounts -> FS)
User->>UI: configure providers via web UI
UI->>FS: write config.yaml/.env, set config.done = "true"
CLI->>FS: poll IsCustomConfigReady() until true or timeout
CLI->>Docker: RemoveContainerIfExists("wrenai-providers-setup")
else no dir
CLI->>CLI: log warning and skip UI
end
CLI->>Docker: continue launching main services
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (9)
wren-ai-service/tools/providers-setup/constants.py (1)
2-2: Remove unused import.The
osmodule is imported but never used in this file.-import os🧰 Tools
🪛 Ruff (0.11.9)
2-2:
osimported but unusedRemove unused import:
os(F401)
wren-launcher/commands/launch.go (1)
251-251: Fix the typo in the comment.There's a spelling error in the comment.
- // chech if ~/.wrenai is exist + // check if ~/.wrenai exists.github/workflows/setup-provider-image.yaml (2)
29-29: Remove trailing spaces.Static analysis detected trailing spaces that should be removed for cleaner code.
- - runner: ubuntu-latest + - runner: ubuntu-latest platform: linux/amd64 - - runner: ubuntu-24.04-arm + - runner: ubuntu-24.04-arm platform: linux/arm64Also applies to: 31-31
🧰 Tools
🪛 YAMLlint (1.37.1)
[error] 29-29: trailing spaces
(trailing-spaces)
130-131: Remove extra blank line.There's an unnecessary blank line at the end of the file.
$(printf '${{ env.REGISTRY }}/${{ env.PROVIDER_UI_IMAGE }}@sha256:%s ' *) \ $TAGS -🧰 Tools
🪛 YAMLlint (1.37.1)
[warning] 130-130: too many blank lines
(1 > 0) (empty-lines)
wren-ai-service/tools/providers-setup/config_loader.py (2)
88-99: Clarify the fallback logic in group_blocks function.The logic for handling blocks without a "type" field could be more explicit and documented.
def group_blocks(blocks: List[Dict[str, Any]]) -> Dict[str, Any]: """ Group YAML blocks by their 'type' field. If multiple blocks share the same type, they are stored as a list. + Falls back to 'settings' key if 'type' field is missing. """ save_blocks = {} for block in blocks: - key = block.get("type") or ("settings" if "settings" in block else None) + key = block.get("type") + if not key and "settings" in block: + key = "settings" if not key: continue
47-48: Improve error message consistency.The error message format should be consistent with other error messages in the function.
if not config_list: - raise ValueError(f"⚠️ Received empty YAML content from: {url}") + raise ValueError(f"Received empty YAML content from: {url}")wren-ai-service/tools/providers-setup/session_state.py (1)
167-167: Consider making the .env path configurable.The hardcoded path
/app/data/.envassumes a containerized environment. Consider making this configurable or using a constant to improve portability.Consider importing the path from the constants module:
+from constants import get_env_path +env_path = get_env_path() -env_path = Path("/app/data/.env")wren-ai-service/tools/providers-setup/ui_components.py (2)
219-220: Remove unnecessary f-string prefix.The f-string on line 220 doesn't contain any placeholders.
- with st.expander(f"{default_name}", expanded=False): + with st.expander(default_name, expanded=False):🧰 Tools
🪛 Ruff (0.11.9)
220-220: f-string without any placeholders
Remove extraneous
fprefix(F541)
294-342: Remove commented code block.This large block of commented code should be removed to keep the codebase clean. If this functionality might be needed later, consider tracking it in version control history or a separate feature branch.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
wren-ai-service/tools/providers-setup/poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (11)
.github/workflows/setup-provider-image.yaml(1 hunks)wren-ai-service/tools/providers-setup/Dockerfile(1 hunks)wren-ai-service/tools/providers-setup/app.py(1 hunks)wren-ai-service/tools/providers-setup/config_loader.py(1 hunks)wren-ai-service/tools/providers-setup/constants.py(1 hunks)wren-ai-service/tools/providers-setup/dry_run_test.py(1 hunks)wren-ai-service/tools/providers-setup/pyproject.toml(1 hunks)wren-ai-service/tools/providers-setup/session_state.py(1 hunks)wren-ai-service/tools/providers-setup/ui_components.py(1 hunks)wren-launcher/commands/launch.go(1 hunks)wren-launcher/utils/docker.go(4 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (4)
wren-launcher/commands/launch.go (2)
wren-launcher/utils/docker.go (4)
TryGetWrenAIDir(552-564)RemoveContainerIfExists(657-674)RunStreamlitUIContainer(576-606)IsCustomConfigReady(678-693)wren-launcher/utils/network.go (1)
FindAvailablePort(16-34)
wren-ai-service/tools/providers-setup/app.py (3)
wren-ai-service/tools/providers-setup/config_loader.py (2)
load_config_yaml_blocks(9-22)group_blocks(82-99)wren-ai-service/tools/providers-setup/session_state.py (2)
ConfigState(5-206)init(17-27)wren-ai-service/tools/providers-setup/ui_components.py (7)
render_llm_config(143-206)render_embedder_config(210-291)render_import_yaml(9-32)render_pipeline_config(343-379)render_preview(383-390)render_apikey(35-141)render_generate_button(392-412)
wren-ai-service/tools/providers-setup/dry_run_test.py (1)
wren-ai-service/tools/providers-setup/session_state.py (1)
ConfigState(5-206)
wren-ai-service/tools/providers-setup/session_state.py (2)
wren-ai-service/src/config.py (1)
config_loader(87-97)wren-ai-service/tools/providers-setup/config_loader.py (2)
fetch_yaml_from_url(37-54)fetch_example_yaml_filenames(101-113)
🪛 actionlint (1.7.7)
.github/workflows/setup-provider-image.yaml
61-61: got unexpected character '/' while lexing expression, expecting 'a'..'z', 'A'..'Z', '_', '0'..'9', ''', '}', '(', ')', '[', ']', '.', '!', '<', '>', '=', '&', '|', '*', ',', ' '
(expression)
🪛 YAMLlint (1.37.1)
.github/workflows/setup-provider-image.yaml
[error] 29-29: trailing spaces
(trailing-spaces)
[error] 31-31: trailing spaces
(trailing-spaces)
[warning] 130-130: too many blank lines
(1 > 0) (empty-lines)
🪛 Ruff (0.11.9)
wren-ai-service/tools/providers-setup/ui_components.py
220-220: f-string without any placeholders
Remove extraneous f prefix
(F541)
wren-ai-service/tools/providers-setup/constants.py
2-2: os imported but unused
Remove unused import: os
(F401)
🔇 Additional comments (16)
wren-ai-service/tools/providers-setup/constants.py (1)
9-28: LGTM! Well-implemented GitHub API function.The function correctly handles the GitHub API request with proper timeout, error handling, and fallback logic. The print statements provide useful debugging information.
wren-ai-service/tools/providers-setup/pyproject.toml (1)
1-18: LGTM! Well-configured Poetry project.The project configuration is appropriate for a Streamlit-based provider setup tool. The dependency version ranges are reasonable and the build system configuration is correct.
wren-ai-service/tools/providers-setup/Dockerfile (2)
1-58: LGTM! Well-structured Dockerfile with good practices.The Dockerfile follows best practices with proper layer caching, minimal base image, and appropriate environment configuration for Streamlit. The multi-stage approach and dependency installation strategy are well-implemented.
34-35:❌ Incorrect review comment
Verify poetry.lock file exists.
The Dockerfile copies
poetry.lockbut this file isn't included in the PR. This could cause build failures if the lock file doesn't exist.
🏁 Script executed:
#!/bin/bash # Check if poetry.lock exists in the providers-setup directory fd -t f "poetry.lock" wren-ai-service/tools/providers-setup/Length of output: 110
poetry.lock exists – no action needed
The
poetry.lockfile is present atwren-ai-service/tools/providers-setup/poetry.lock, so theCOPY poetry.lockinstruction in the Dockerfile is valid. No changes required.Likely an incorrect or invalid review comment.
wren-ai-service/tools/providers-setup/app.py (3)
1-12: LGTM! Clean and well-organized imports.The imports are well-structured, importing only the necessary functions from each module. The modular approach promotes good separation of concerns.
33-49: Excellent validation logic for configuration blocks.The validation properly checks both type and content of required configuration blocks, providing clear user feedback when blocks are missing. This helps users understand what needs to be configured.
59-88: Well-designed UI layout with logical organization.The two-column layout effectively separates input controls from preview functionality. The component rendering is modular and follows a logical flow from API keys through configuration to final generation.
wren-ai-service/tools/providers-setup/dry_run_test.py (1)
6-31: LGTM! Well-structured LLM completion test function.The function properly handles the completion API call with good error handling and response validation.
wren-ai-service/tools/providers-setup/config_loader.py (2)
9-23: LGTM! Excellent configuration loading structure.The functions provide robust configuration loading with proper error handling, fallback mechanisms, and clear separation of concerns.
Also applies to: 24-36, 37-55
115-127: LGTM! Clean integration with session state.The
apply_config_blocksfunction properly integrates the loaded configuration with the Streamlit session state management.wren-ai-service/tools/providers-setup/session_state.py (3)
16-28: LGTM!The initialization orchestration method is well-structured and provides appropriate control over state overwriting.
30-68: Well-implemented form initialization logic.The method properly handles form creation with unique IDs and correctly flattens kwargs for the model list.
121-130:✅ Verification successful
Consider handling empty response from remote config fetch.
When
fetch_yaml_from_urlreturns an empty list (indicating a fetch error), the code continues as if there's no latest pipeline configuration. Consider adding explicit error handling to distinguish between fetch failures and legitimate empty configurations.
🏁 Script executed:
#!/bin/bash # Check how fetch_yaml_from_url handles errors ast-grep --pattern 'def fetch_yaml_from_url($_) -> $_: $$$ 'Length of output: 1745
No additional empty-response handling needed.
The fetch_yaml_from_url function already raises a ValueError on empty YAML content, logs the error via st.error, and returns an empty list only on failure. In init_pipeline, receiving an empty list (and thus latest_pipeline remaining None) correctly reflects a fetch/parsing error and skips the update as intended.
wren-launcher/utils/docker.go (2)
298-344: Great improvement to custom provider setup!The automatic downloading of missing configuration files and USER_UUID injection significantly improves the user experience compared to the previous strict error returns.
576-693: Well-implemented container management functions.The new functions for Streamlit UI container management are properly structured with good error handling, appropriate volume mounts, and clear separation of concerns.
wren-ai-service/tools/providers-setup/ui_components.py (1)
35-142: Well-implemented API key management.The API key configuration section properly handles CRUD operations on the .env file with appropriate error handling and file format preservation.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
wren-ai-service/tools/providers-setup/ui_components.py (1)
431-437:⚠️ Potential issueCritical: Remove eval() to prevent code injection.
Using
eval()on user input is a serious security vulnerability that could allow arbitrary code execution.
🧹 Nitpick comments (4)
wren-ai-service/tools/providers-setup/ui_components.py (4)
25-33: Add success message after YAML upload.The YAML upload section is missing a success message after successfully processing the uploaded file, unlike the GitHub import section.
uploaded_file = st.file_uploader("Upload your own YAML file", type=["yaml", "yml"]) if uploaded_file is not None and st.button("Import.yaml", key="import_own_yaml"): try: blocks = list(yaml.safe_load_all(uploaded_file)) apply_config_blocks(blocks) + st.success("YAML import succeeded. Settings updated.") except Exception as e: st.error(f"Failed to import YAML file: {e}")
221-221: Remove unnecessary f-string prefix.The f-string prefix is not needed here since there are no placeholders.
- with st.expander(f"{default_name}", expanded=False): + with st.expander(default_name, expanded=False):
233-233: Fix typo in comment.There's a typo in the comment: "doucument_store" should be "document_store".
- document_store_dim = st.text_input("Embedding_model_dim (placed in doucument_store)", key="embedding_model_dim", value="3072") + document_store_dim = st.text_input("Embedding_model_dim (placed in document_store)", key="embedding_model_dim", value="3072")
296-344: Remove commented-out code.This entire function is commented out and marked for deletion. Dead code should be removed to maintain code cleanliness.
Remove lines 296-344 entirely instead of keeping them as comments.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
wren-ai-service/tools/providers-setup/session_state.py(1 hunks)wren-ai-service/tools/providers-setup/ui_components.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/tools/providers-setup/session_state.py
🧰 Additional context used
🧬 Code Graph Analysis (1)
wren-ai-service/tools/providers-setup/ui_components.py (5)
wren-ai-service/tools/providers-setup/session_state.py (1)
ConfigState(5-208)wren-ai-service/src/config.py (1)
config_loader(87-97)wren-ai-service/tools/providers-setup/config_loader.py (2)
load_selected_example_yaml(24-35)apply_config_blocks(115-127)wren-ai-service/tools/providers-setup/dry_run_test.py (2)
llm_completion_test(6-31)llm_embedding_test(34-63)wren-ai-service/tools/providers-setup/constants.py (3)
get_env_path(80-87)get_config_path(69-77)get_config_done_path(59-67)
🪛 Ruff (0.11.9)
wren-ai-service/tools/providers-setup/ui_components.py
222-222: f-string without any placeholders
Remove extraneous f prefix
(F541)
… timeout for user configuration
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (7)
wren-ai-service/tools/providers-setup/constants.py (1)
59-87: The directory naming is now consistent.All path utility functions correctly use the
.wrenaidirectory name (lowercase), which resolves the previous naming inconsistency concern.🧰 Tools
🪛 Pylint (3.3.7)
[refactor] 64-67: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
[refactor] 74-77: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
[refactor] 84-87: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
wren-ai-service/tools/providers-setup/dry_run_test.py (2)
44-48: Excellent bounds checking implementation.The code now properly checks for empty models list before accessing the first element, addressing the previous IndexError concern.
74-77: Good error handling for session state.The function now uses
.get()with a default value and checks for non-empty API keys, addressing the previous KeyError concern..github/workflows/setup-provider-image.yaml (2)
11-13: Permissions are correctly configured.The permissions section is now properly uncommented and configured for GitHub Container Registry access.
62-62: String replacement syntax is now correct.The expression now uses the proper
trcommand syntax instead of the invalid//operator.wren-launcher/utils/docker.go (1)
613-615: Path validation is now properly implemented.The build path is now validated before attempting to build, addressing the previous robustness concern.
wren-ai-service/tools/providers-setup/ui_components.py (1)
435-438: Security improvement confirmed - eval() replaced with ast.literal_eval().Great work addressing the previous security concern! The code now safely evaluates literals using
ast.literal_eval()instead of the dangerouseval()function, preventing potential code injection attacks.
🧹 Nitpick comments (10)
wren-ai-service/tools/providers-setup/constants.py (2)
2-2: Remove unused import.The
osmodule is imported but never used in this file.-import os🧰 Tools
🪛 Ruff (0.11.9)
2-2:
osimported but unusedRemove unused import:
os(F401)
19-28: Consider simplifying control flow.The
elseclause afterreturnstatements can be removed to reduce nesting levels.- if response.status_code == 200: - data = response.json() - return data["tag_name"] - else: - print(f"Failed to get latest release: {response.status_code}") + if response.status_code == 200: + data = response.json() + return data["tag_name"] + print(f"Failed to get latest release: {response.status_code}")🧰 Tools
🪛 Pylint (3.3.7)
[refactor] 20-24: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
wren-ai-service/tools/providers-setup/dry_run_test.py (1)
25-28: Consider simplifying control flow.The
elseclauses afterreturnstatements can be removed to reduce nesting levels.- if choices and choices[0].get("message", {}).get("content"): - return True, response_completion["choices"][0]["message"]["content"] - else: - return False, "No valid response content" + if choices and choices[0].get("message", {}).get("content"): + return True, response_completion["choices"][0]["message"]["content"] + return False, "No valid response content"- if isinstance(embedding_vector, list) and len(embedding_vector) > 0: - return True, f"Embedding length: {len(embedding_vector)}" - else: - return False, "Embedding vector is empty or invalid" + if isinstance(embedding_vector, list) and len(embedding_vector) > 0: + return True, f"Embedding length: {len(embedding_vector)}" + return False, "Embedding vector is empty or invalid"Also applies to: 58-61
🧰 Tools
🪛 Pylint (3.3.7)
[refactor] 25-28: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
.github/workflows/setup-provider-image.yaml (1)
29-29: Fix formatting issues.Remove trailing spaces and extra blank lines to comply with YAML formatting standards.
- - runner: ubuntu-latest + - runner: ubuntu-latest - - runner: ubuntu-24.04-arm + - runner: ubuntu-24.04-arm - +Also applies to: 31-31, 130-130
🧰 Tools
🪛 YAMLlint (1.37.1)
[error] 29-29: trailing spaces
(trailing-spaces)
wren-launcher/utils/docker.go (1)
662-681: Consider adding logging for container removal operations.The function properly handles container existence checking, but consider adding more detailed logging for troubleshooting.
func RemoveContainerIfExists(name string) error { // Check if the container exists (inspect will fail if not) err := exec.Command("docker", "inspect", name).Run() if err != nil { - pterm.Info.Println("🔍 Container does not exist, skipping:", name) + pterm.Debug.Println("🔍 Container does not exist, skipping:", name) return nil } + pterm.Info.Println("🧹 Removing existing container:", name) // Remove container forcefully cmd := exec.Command("docker", "rm", "-f", name) out, err := cmd.CombinedOutput() if err != nil { return fmt.Errorf("❌ Failed to force-remove container: %v\n%s", err, string(out)) } pterm.Info.Println("🧹 Container forcibly removed:", name) return nil }wren-ai-service/tools/providers-setup/ui_components.py (5)
30-34: Add success message for custom YAML import.The success message is missing when importing a custom YAML file, unlike the example YAML import which shows a success message on line 24.
try: blocks = list(yaml.safe_load_all(uploaded_file)) apply_config_blocks(blocks) - + st.success("YAML import succeeded. Settings updated.") except Exception as e: st.error(f"Failed to import YAML file: {e}")
37-143: Consider refactoring this function to reduce complexity.This function handles too many responsibilities (adding, editing, saving, deleting API keys) and has high complexity according to static analysis (21 local variables, 18 branches, 62 statements).
Consider breaking this into smaller functions:
_render_api_key_form()for form rendering_save_api_keys()for save logic_delete_api_key()for deletion logicThis would improve maintainability and testability.
🧰 Tools
🪛 Pylint (3.3.7)
[refactor] 37-37: Too many local variables (21/15)
(R0914)
[refactor] 37-37: Too many branches (18/12)
(R0912)
[refactor] 37-37: Too many statements (62/50)
(R0915)
224-224: Remove unnecessary f-string prefix.The f-string has no placeholders, making the
fprefix unnecessary.- st.markdown(f"**type:** `embedder`") + st.markdown("**type:** `embedder`")🧰 Tools
🪛 Ruff (0.11.9)
224-224: f-string without any placeholders
Remove extraneous
fprefix(F541)
298-346: Remove commented-out code.This large block of commented code should be removed to keep the codebase clean, unless there's a specific reason to keep it for reference.
-# DELETE -# def render_document_store_config(): -# """ -# Render the document store configuration section. -# Displays current settings and allows updating index location, dimensions, timeout, etc. -# """ -# with st.expander("Document Store Configuration", expanded=False): - [... rest of commented code ...]
460-466: Remove unnecessary else clause after return.The else clause is unnecessary since the previous block contains a return statement.
existing_ids = [m.get("id") for m in st.session_state[ConfigState.LLM_MODELS_KEY]] if form_id in existing_ids: index = existing_ids.index(form_id) st.session_state[ConfigState.LLM_MODELS_KEY][index] = {**saved_entry, "id": form_id} return True, f"Updated model: {form['model']}" - else: - st.session_state[ConfigState.LLM_MODELS_KEY].append({**saved_entry, "id": form_id}) - return True, f"Added new model: {form['model']}" + + st.session_state[ConfigState.LLM_MODELS_KEY].append({**saved_entry, "id": form_id}) + return True, f"Added new model: {form['model']}"🧰 Tools
🪛 Pylint (3.3.7)
[refactor] 460-466: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
.github/workflows/setup-provider-image.yaml(1 hunks)wren-ai-service/tools/providers-setup/constants.py(1 hunks)wren-ai-service/tools/providers-setup/dry_run_test.py(1 hunks)wren-ai-service/tools/providers-setup/ui_components.py(1 hunks)wren-launcher/commands/launch.go(1 hunks)wren-launcher/utils/docker.go(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-launcher/commands/launch.go
🧰 Additional context used
🪛 Ruff (0.11.9)
wren-ai-service/tools/providers-setup/ui_components.py
224-224: f-string without any placeholders
Remove extraneous f prefix
(F541)
wren-ai-service/tools/providers-setup/constants.py
2-2: os imported but unused
Remove unused import: os
(F401)
🪛 Pylint (3.3.7)
wren-ai-service/tools/providers-setup/ui_components.py
[refactor] 37-37: Too many local variables (21/15)
(R0914)
[refactor] 37-37: Too many branches (18/12)
(R0912)
[refactor] 37-37: Too many statements (62/50)
(R0915)
[refactor] 145-145: Too many branches (14/12)
(R0912)
[refactor] 214-214: Too many branches (15/12)
(R0912)
[error] 382-382: Possibly using variable 'selected_llm' before assignment
(E0606)
[refactor] 460-466: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
wren-ai-service/tools/providers-setup/constants.py
[refactor] 20-24: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
[refactor] 64-67: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
[refactor] 74-77: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
[refactor] 84-87: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
wren-ai-service/tools/providers-setup/dry_run_test.py
[refactor] 25-28: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
[refactor] 58-61: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
🪛 YAMLlint (1.37.1)
.github/workflows/setup-provider-image.yaml
[error] 29-29: trailing spaces
(trailing-spaces)
[error] 31-31: trailing spaces
(trailing-spaces)
[warning] 130-130: too many blank lines
(1 > 0) (empty-lines)
🔇 Additional comments (3)
wren-launcher/utils/docker.go (3)
297-344: Excellent improvement to custom provider handling.The logic now gracefully handles missing configuration files by downloading defaults instead of failing with errors. This provides a much better user experience for the custom provider setup flow.
573-606: Well-structured Streamlit UI container management.The
RunStreamlitUIContainerfunction properly orchestrates the container lifecycle with good error handling and clear separation of concerns.
683-700: Robust configuration readiness checking.The
IsCustomConfigReadyfunction implements proper error handling and case-insensitive comparison for the configuration state check.
The selected_llm variable may be used before assignment if the pipeline has no llm key or if the current LLM value is not in the options list. Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
wren-ai-service/tools/providers-setup/ui_components.py (1)
298-346: Remove commented-out code to improve maintainability.Large blocks of commented-out code should be removed rather than left in the codebase. If this functionality might be needed in the future, consider documenting the decision or using version control history instead.
-# DELETE -# def render_document_store_config(): -# """ -# Render the document store configuration section. -# Displays current settings and allows updating index location, dimensions, timeout, etc. -# """ -# with st.expander("Document Store Configuration", expanded=False): -# st.markdown(f"**type:** `document_store`") -# st.markdown(f"**provider:** `{st.session_state[ConfigState.DOC_STORE_KEY].get('provider')}`") -# st.markdown(f"**location:** `{st.session_state[ConfigState.DOC_STORE_KEY].get('location')}`") - -# document_store_timeout = st.text_input("Timeout (default: 120)", key="document_store_timeout", value="120") -# st.markdown(f"**timeout:** `120`") -# st.markdown(f"**recreate_index:** `{st.session_state[ConfigState.DOC_STORE_KEY].get('recreate_index')}`") - -# document_store_dim = st.text_input("Embedding_model_dim", value="3072") - -# if st.button("💾 save", key="save_document_store"): -# errors = [] -# if not document_store_dim: -# errors.append("Embedding model dim is required.") -# else: -# try: -# int(document_store_dim) -# except ValueError: -# errors.append("Embedding model dim must be an integer.") - -# if not document_store_timeout: -# errors.append("Timeout is required.") -# else: -# try: -# int(document_store_timeout) -# except ValueError: -# errors.append("Timeout must be an integer.") - -# if errors: -# for error in errors: -# st.error(error) -# else: -# st.session_state.document_store = { -# "type": "document_store", -# "provider": st.session_state[ConfigState.DOC_STORE_KEY].get("provider"), -# "location": st.session_state[ConfigState.DOC_STORE_KEY].get("location"), -# "embedding_model_dim": document_store_dim, -# "timeout": document_store_timeout, -# "recreate_index": st.session_state[ConfigState.DOC_STORE_KEY].get("recreate_index") -# } -# st.success("Updated document store configuration")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/tools/providers-setup/ui_components.py(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
wren-ai-service/tools/providers-setup/ui_components.py (5)
wren-ai-service/tools/providers-setup/session_state.py (1)
ConfigState(5-208)wren-ai-service/src/config.py (1)
config_loader(87-97)wren-ai-service/tools/providers-setup/config_loader.py (2)
load_selected_example_yaml(24-35)apply_config_blocks(115-127)wren-ai-service/tools/providers-setup/dry_run_test.py (2)
llm_completion_test(6-31)llm_embedding_test(34-66)wren-ai-service/tools/providers-setup/constants.py (3)
get_env_path(80-87)get_config_path(69-77)get_config_done_path(59-67)
🪛 Ruff (0.11.9)
wren-ai-service/tools/providers-setup/ui_components.py
370-371: SyntaxError: Expected an indented block after if statement
395-395: SyntaxError: Unexpected indentation
399-399: SyntaxError: Expected a statement
🪛 Pylint (3.3.7)
wren-ai-service/tools/providers-setup/ui_components.py
[error] 371-371: Parsing failed: 'expected an indented block after 'if' statement on line 370 (ui_components, line 371)'
(E0001)
🔇 Additional comments (7)
wren-ai-service/tools/providers-setup/ui_components.py (7)
1-9: LGTM: Clean and appropriate imports.The imports are well-organized and include the necessary modules for the Streamlit UI functionality. The
astimport indicates proper security practices for safe evaluation.
10-35: LGTM: Secure YAML loading with proper error handling.The function correctly uses
yaml.safe_load_all()for secure YAML parsing and includes appropriate error handling for both GitHub examples and user uploads.
37-143: LGTM: Comprehensive API key management with proper validation.The function provides robust API key management with appropriate validation, duplicate checking, and secure .env file handling. The logic for updating existing keys and adding new ones is well-implemented.
145-211: LGTM: Well-structured LLM configuration interface.The function provides a comprehensive interface for managing LLM models with dynamic form handling, validation, and testing capabilities. The separation of concerns with validation in a separate function is good practice.
214-296: LGTM: Robust embedding model configuration with validation.The function provides comprehensive embedding model configuration with proper input validation, type checking, and testing capabilities. The error handling for integer conversion is well-implemented.
399-429: LGTM: Clean preview and configuration generation logic.Both functions are well-implemented with proper error handling and clear separation of concerns between preview display and configuration saving.
🧰 Tools
🪛 Ruff (0.11.9)
399-399: SyntaxError: Expected a statement
432-563: LGTM: Security issue resolved and comprehensive validation.The utility functions are well-implemented with proper validation and error handling. Notably, the security concern from previous reviews has been addressed by using
ast.literal_eval()instead ofeval()on line 448, which safely evaluates string literals without the risk of code injection.
| def render_pipeline_config(): | ||
| """ | ||
| Render the pipeline configuration section. | ||
| Allows selecting LLM models for each defined pipeline step. | ||
| """ | ||
| pipeline_llm_options = [] | ||
| all_pipelines = st.session_state[ConfigState.PIPELINE_KEY].get("pipes", []) | ||
|
|
||
| # set all LLM models options | ||
| for model in st.session_state[ConfigState.LLM_MODELS_KEY]: | ||
| if model.get("alias"): | ||
| pipeline_llm_options.append("litellm_llm." + model["alias"]) | ||
| elif model.get("model"): | ||
| pipeline_llm_options.append("litellm_llm." + model["model"]) | ||
|
|
||
| # enumerate all pipelines, record original index | ||
| for original_idx, pipe in enumerate(all_pipelines): | ||
| if not pipe.get("llm"): | ||
| continue | ||
|
|
||
| pipe_name = pipe.get("name", f"Unnamed Pipeline {original_idx}") | ||
| with st.expander(f"🔧 Pipeline: {pipe_name}", expanded=False): | ||
| for key, value in pipe.items(): | ||
| if key == "llm": | ||
| for key, value in pipe.items(): | ||
| if key == "llm": | ||
| # Ensure we have a default value | ||
| default_index = 0 | ||
| if value in pipeline_llm_options: | ||
| default_index = pipeline_llm_options.index(value) | ||
|
|
||
| selected_llm = st.selectbox( | ||
| "LLM Model", | ||
| options=pipeline_llm_options, | ||
| index=default_index, | ||
| key=f"llm_{original_idx}" | ||
| ) | ||
| else: | ||
| st.markdown(f"**{key}:** `{value}`") | ||
|
|
||
| if st.button("💾 Save this llm", key=f"save_{pipe_name}_{original_idx}"): | ||
| # Only proceed if we have a selected LLM | ||
| if 'selected_llm' in locals(): | ||
| # ✅ use original index to update llm | ||
| st.session_state[ConfigState.PIPELINE_KEY]["pipes"][original_idx]["llm"] = selected_llm | ||
| st.success(f"✅ Updated pipeline `{pipe_name}` LLM to `{selected_llm}`") | ||
| else: | ||
| st.error("No LLM selected for this pipeline.") | ||
| st.success(f"✅ Updated pipeline `{pipe_name}` LLM to `{selected_llm}`") | ||
|
|
There was a problem hiding this comment.
Critical: Fix syntax errors and restructure pipeline configuration logic.
This section contains multiple critical syntax errors that prevent the code from running:
- Syntax Error: The for loop at line 369 has no body (missing indented block)
- Duplicate Code: The for loop is duplicated at lines 369 and 371
- Incorrect Indentation: Lines 371-396 have wrong indentation levels
- Duplicate Code: Line 395 duplicates the success message from line 392
Apply this diff to fix the syntax errors and logic:
- for key, value in pipe.items():
- if key == "llm":
-for key, value in pipe.items():
- if key == "llm":
- # Ensure we have a default value
- default_index = 0
- if value in pipeline_llm_options:
- default_index = pipeline_llm_options.index(value)
-
- selected_llm = st.selectbox(
- "LLM Model",
- options=pipeline_llm_options,
- index=default_index,
- key=f"llm_{original_idx}"
- )
- else:
- st.markdown(f"**{key}:** `{value}`")
-
-if st.button("💾 Save this llm", key=f"save_{pipe_name}_{original_idx}"):
- # Only proceed if we have a selected LLM
- if 'selected_llm' in locals():
- # ✅ use original index to update llm
- st.session_state[ConfigState.PIPELINE_KEY]["pipes"][original_idx]["llm"] = selected_llm
- st.success(f"✅ Updated pipeline `{pipe_name}` LLM to `{selected_llm}`")
- else:
- st.error("No LLM selected for this pipeline.")
- st.success(f"✅ Updated pipeline `{pipe_name}` LLM to `{selected_llm}`")
+ selected_llm = None
+ for key, value in pipe.items():
+ if key == "llm":
+ # Ensure we have a default value
+ default_index = 0
+ if value in pipeline_llm_options:
+ default_index = pipeline_llm_options.index(value)
+
+ selected_llm = st.selectbox(
+ "LLM Model",
+ options=pipeline_llm_options,
+ index=default_index,
+ key=f"llm_{original_idx}"
+ )
+ else:
+ st.markdown(f"**{key}:** `{value}`")
+
+ if st.button("💾 Save this llm", key=f"save_{pipe_name}_{original_idx}"):
+ # Only proceed if we have a selected LLM
+ if selected_llm is not None:
+ # ✅ use original index to update llm
+ st.session_state[ConfigState.PIPELINE_KEY]["pipes"][original_idx]["llm"] = selected_llm
+ st.success(f"✅ Updated pipeline `{pipe_name}` LLM to `{selected_llm}`")
+ else:
+ st.error("No LLM selected for this pipeline.")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def render_pipeline_config(): | |
| """ | |
| Render the pipeline configuration section. | |
| Allows selecting LLM models for each defined pipeline step. | |
| """ | |
| pipeline_llm_options = [] | |
| all_pipelines = st.session_state[ConfigState.PIPELINE_KEY].get("pipes", []) | |
| # set all LLM models options | |
| for model in st.session_state[ConfigState.LLM_MODELS_KEY]: | |
| if model.get("alias"): | |
| pipeline_llm_options.append("litellm_llm." + model["alias"]) | |
| elif model.get("model"): | |
| pipeline_llm_options.append("litellm_llm." + model["model"]) | |
| # enumerate all pipelines, record original index | |
| for original_idx, pipe in enumerate(all_pipelines): | |
| if not pipe.get("llm"): | |
| continue | |
| pipe_name = pipe.get("name", f"Unnamed Pipeline {original_idx}") | |
| with st.expander(f"🔧 Pipeline: {pipe_name}", expanded=False): | |
| for key, value in pipe.items(): | |
| if key == "llm": | |
| for key, value in pipe.items(): | |
| if key == "llm": | |
| # Ensure we have a default value | |
| default_index = 0 | |
| if value in pipeline_llm_options: | |
| default_index = pipeline_llm_options.index(value) | |
| selected_llm = st.selectbox( | |
| "LLM Model", | |
| options=pipeline_llm_options, | |
| index=default_index, | |
| key=f"llm_{original_idx}" | |
| ) | |
| else: | |
| st.markdown(f"**{key}:** `{value}`") | |
| if st.button("💾 Save this llm", key=f"save_{pipe_name}_{original_idx}"): | |
| # Only proceed if we have a selected LLM | |
| if 'selected_llm' in locals(): | |
| # ✅ use original index to update llm | |
| st.session_state[ConfigState.PIPELINE_KEY]["pipes"][original_idx]["llm"] = selected_llm | |
| st.success(f"✅ Updated pipeline `{pipe_name}` LLM to `{selected_llm}`") | |
| else: | |
| st.error("No LLM selected for this pipeline.") | |
| st.success(f"✅ Updated pipeline `{pipe_name}` LLM to `{selected_llm}`") | |
| def render_pipeline_config(): | |
| """ | |
| Render the pipeline configuration section. | |
| Allows selecting LLM models for each defined pipeline step. | |
| """ | |
| pipeline_llm_options = [] | |
| all_pipelines = st.session_state[ConfigState.PIPELINE_KEY].get("pipes", []) | |
| # set all LLM models options | |
| for model in st.session_state[ConfigState.LLM_MODELS_KEY]: | |
| if model.get("alias"): | |
| pipeline_llm_options.append("litellm_llm." + model["alias"]) | |
| elif model.get("model"): | |
| pipeline_llm_options.append("litellm_llm." + model["model"]) | |
| # enumerate all pipelines, record original index | |
| for original_idx, pipe in enumerate(all_pipelines): | |
| if not pipe.get("llm"): | |
| continue | |
| pipe_name = pipe.get("name", f"Unnamed Pipeline {original_idx}") | |
| with st.expander(f"🔧 Pipeline: {pipe_name}", expanded=False): | |
| # initialize selection | |
| selected_llm = None | |
| for key, value in pipe.items(): | |
| if key == "llm": | |
| # Ensure we have a default value | |
| default_index = 0 | |
| if value in pipeline_llm_options: | |
| default_index = pipeline_llm_options.index(value) | |
| selected_llm = st.selectbox( | |
| "LLM Model", | |
| options=pipeline_llm_options, | |
| index=default_index, | |
| key=f"llm_{original_idx}" | |
| ) | |
| else: | |
| st.markdown(f"**{key}:** `{value}`") | |
| if st.button("💾 Save this llm", key=f"save_{pipe_name}_{original_idx}"): | |
| # Only proceed if we have a selected LLM | |
| if selected_llm is not None: | |
| st.session_state[ConfigState.PIPELINE_KEY]["pipes"][original_idx]["llm"] = selected_llm | |
| st.success(f"✅ Updated pipeline `{pipe_name}` LLM to `{selected_llm}`") | |
| else: | |
| st.error("No LLM selected for this pipeline.") |
🧰 Tools
🪛 Ruff (0.11.9)
370-371: SyntaxError: Expected an indented block after if statement
395-395: SyntaxError: Unexpected indentation
🪛 Pylint (3.3.7)
[error] 371-371: Parsing failed: 'expected an indented block after 'if' statement on line 370 (ui_components, line 371)'
(E0001)
🤖 Prompt for AI Agents
In wren-ai-service/tools/providers-setup/ui_components.py around lines 347 to
396, fix the syntax errors by removing the duplicated for loop at line 371 and
ensuring the for loop at line 369 has a proper indented body. Correct the
indentation of lines 371 to 396 so they are nested correctly inside the pipeline
expander and the for loop. Remove the duplicate success message on line 395.
This will resolve the missing block error, eliminate duplicate code, and
properly structure the pipeline configuration rendering logic.
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (1)
wren-launcher/utils/docker.go (1)
617-631: Make Docker build path robust and absolute.
This revisits the prior feedback and resolves fragility when the working directory varies by using the executable’s dir + filepath.Abs.-func buildStreamlitImage() error { - // Get the absolute path relative to the executable or use a configuration - buildPath := "../wren-ai-service/tools/providers-setup" - // Verify the path exists before attempting to build - if _, err := os.Stat(buildPath); os.IsNotExist(err) { - return fmt.Errorf("build path does not exist: %s", buildPath) - } - cmd := exec.Command("docker", "build", "-t", "wrenai-providers-setup", buildPath) +func buildStreamlitImage() error { + exe, err := os.Executable() + if err != nil { + return fmt.Errorf("cannot resolve executable path: %v", err) + } + base := filepath.Dir(exe) + buildPath := filepath.Join(base, "..", "wren-ai-service", "tools", "providers-setup") + buildPath, err = filepath.Abs(buildPath) + if err != nil { + return fmt.Errorf("cannot resolve build path: %v", err) + } + if _, err := os.Stat(buildPath); os.IsNotExist(err) { + return fmt.Errorf("build path does not exist: %s", buildPath) + } + cmd := exec.Command("docker", "build", "-t", "wrenai-providers-setup", buildPath) output, err := cmd.CombinedOutput() if err != nil { return fmt.Errorf("build failed: %v\n%s", err, output) } return nil }Add this import:
// at the top import block import ( // ... "path/filepath" )
🧹 Nitpick comments (3)
wren-launcher/utils/docker.go (3)
633-643: Use 0600 for config.done for consistency.
Not strictly sensitive, but consistent restrictive perms are safer by default.- err := os.WriteFile(donePath, []byte("false"), 0644) + err := os.WriteFile(donePath, []byte("false"), 0600)
645-649: Prefer filepath.Join for OS portability.
Windows paths can break with path.Join; new code can use filepath.Join going forward.-func getMountPaths(wrenDir string) (string, string, string) { - return path.Join(wrenDir, "config.yaml"), - path.Join(wrenDir, ".env"), - path.Join(wrenDir, "config.done") +func getMountPaths(wrenDir string) (string, string, string) { + return filepath.Join(wrenDir, "config.yaml"), + filepath.Join(wrenDir, ".env"), + filepath.Join(wrenDir, "config.done") }Add to imports if not already present:
import "path/filepath"
651-669: Minor: consider --mount for clarity and better path handling.
-v works, but --mount is more explicit and resilient with odd paths.- cmd := exec.Command("docker", "run", "--rm", "-d", + cmd := exec.Command("docker", "run", "--rm", "-d", "-p", fmt.Sprintf("%d:8501", port), "--name", "wrenai-providers-setup", - "-v", configPath+":/app/data/config.yaml", - "-v", envPath+":/app/data/.env", - "-v", donePath+":/app/data/config.done", + "--mount", fmt.Sprintf("type=bind,src=%s,dst=/app/data/config.yaml", configPath), + "--mount", fmt.Sprintf("type=bind,src=%s,dst=/app/data/.env", envPath), + "--mount", fmt.Sprintf("type=bind,src=%s,dst=/app/data/config.done", donePath), "wrenai-providers-setup", )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
wren-launcher/commands/launch.go(1 hunks)wren-launcher/utils/docker.go(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-launcher/commands/launch.go
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: fmt-and-test
- GitHub Check: lint
🔇 Additional comments (2)
wren-launcher/utils/docker.go (2)
9-9: Import of os/exec looks fine.
Used appropriately for docker build/run helpers.
671-690: LGTM: pragmatic forced removal helper.
Clear UX and safe early-exit if container is absent.
| // Generate or load existing USER_UUID | ||
| userUUID, err := prepareUserUUID(projectDir) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| // Ensure .env exists (download if missing) | ||
| envFilePath := getEnvFilePath(projectDir) | ||
| if _, err := os.Stat(envFilePath); os.IsNotExist(err) { | ||
| pterm.Println(".env file not found, downloading from default URL...") | ||
| err = downloadFile(envFilePath, DOCKER_COMPOSE_ENV_URL) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to download .env: %v", err) | ||
| } | ||
|
|
||
| // Read downloaded .env content | ||
| contentBytes, err := os.ReadFile(envFilePath) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to read .env: %v", err) | ||
| } | ||
| str := string(contentBytes) | ||
|
|
||
| // Replace or append USER_UUID | ||
| reg := regexp.MustCompile(`(?m)^USER_UUID=.*$`) | ||
| if reg.MatchString(str) { | ||
| str = reg.ReplaceAllString(str, "USER_UUID="+userUUID) | ||
| } else { | ||
| str = str + "\nUSER_UUID=" + userUUID + "\n" | ||
| } | ||
|
|
||
| // Save updated .env file | ||
| err = os.WriteFile(envFilePath, []byte(str), 0644) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to write updated .env: %v", err) | ||
| } | ||
| } | ||
|
|
||
| // if config.yaml file does not exist, return error | ||
| if _, err := os.Stat(getConfigFilePath(projectDir)); os.IsNotExist(err) { | ||
| return fmt.Errorf("config.yaml file does not exist, please download the config.yaml file from %s to ~/.wrenai, rename it to config.yaml and fill in the required information", AI_SERVICE_CONFIG_URL) | ||
| // Ensure config.yaml exists (download if missing) | ||
| configFilePath := getConfigFilePath(projectDir) | ||
| if _, err := os.Stat(configFilePath); os.IsNotExist(err) { | ||
| pterm.Println("config.yaml not found, downloading from default URL...") | ||
| err = downloadFile(configFilePath, AI_SERVICE_CONFIG_URL) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to download config.yaml: %v", err) | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Harden .env/config handling for “custom” provider (permissions + robustness).
- .env will carry secrets; write it with 0600 and chmod the downloaded files.
- After downloads, chmod config.yaml to 0600 as well.
- Consider making downloadFile fail on non-2xx to avoid writing error pages.
Apply these diffs in this hunk:
@@
- err = downloadFile(envFilePath, DOCKER_COMPOSE_ENV_URL)
+ err = downloadFile(envFilePath, DOCKER_COMPOSE_ENV_URL)
if err != nil {
return fmt.Errorf("failed to download .env: %v", err)
}
+ // Ensure secret-friendly perms
+ _ = os.Chmod(envFilePath, 0600)
@@
- err = os.WriteFile(envFilePath, []byte(str), 0644)
+ err = os.WriteFile(envFilePath, []byte(str), 0600)
if err != nil {
return fmt.Errorf("failed to write updated .env: %v", err)
}
@@
- err = downloadFile(configFilePath, AI_SERVICE_CONFIG_URL)
+ err = downloadFile(configFilePath, AI_SERVICE_CONFIG_URL)
if err != nil {
return fmt.Errorf("failed to download config.yaml: %v", err)
}
+ // Lock down config.yaml as it may later contain sensitive info
+ _ = os.Chmod(configFilePath, 0600)Additionally, improve downloadFile to fail on HTTP errors:
func downloadFile(filepath string, url string) error {
// Get the data
- resp, err := http.Get(url) // #nosec G107 -- URL is from trusted source constants
+ resp, err := http.Get(url) // #nosec G107 -- URL is from trusted source constants
if err != nil {
return err
}
defer func() { _ = resp.Body.Close() }()
+ if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+ return fmt.Errorf("download failed: %s returned status %d", url, resp.StatusCode)
+ }🤖 Prompt for AI Agents
In wren-launcher/utils/docker.go around lines 307 to 353, the .env and
config.yaml handling needs hardened permissions and more robust downloads:
update the code so downloaded files are written with 0600 (use os.WriteFile with
mode 0600) and explicitly os.Chmod the downloaded file paths to 0600 after a
successful download or write; when updating the .env ensure you still write the
final file with 0600; and modify downloadFile to treat non-2xx HTTP responses as
errors (return a non-nil error and do not write the response body) so we never
save error pages as files.
| func TryGetWrenAIDir() (string, error) { | ||
| homeDir, err := os.UserHomeDir() | ||
| if err != nil { | ||
| return "", err | ||
| } | ||
| wrenDir := path.Join(homeDir, ".wrenai") | ||
|
|
||
| info, err := os.Stat(wrenDir) | ||
| if err != nil || !info.IsDir() { | ||
| return "", nil | ||
| } | ||
| return wrenDir, nil | ||
| } | ||
|
|
There was a problem hiding this comment.
Bug: empty return path can cause files to be created in the CWD.
TryGetWrenAIDir returns "" when ~/.wrenai is missing. Callers (e.g., RunStreamlitUIContainer, IsCustomConfigReady) then join against "", leading to ./config.done etc. Write or read ends up in the working dir instead of ~/.wrenai.
I’ve proposed fixes in the RunStreamlitUIContainer and IsCustomConfigReady comments to create/guard the directory. Optionally, consider changing TryGetWrenAIDir to create the dir or return an error when missing.
| if _, err := os.Stat(filePath); os.IsNotExist(err) { | ||
| return os.WriteFile(filePath, defaultContent, 0644) | ||
| } | ||
| return nil | ||
| } | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Default to secret-friendly perms when creating files.
Use 0600 instead of 0644; this function is used for .env and config.yaml.
func ensureFileExists(filePath string, defaultContent []byte) error {
if _, err := os.Stat(filePath); os.IsNotExist(err) {
- return os.WriteFile(filePath, defaultContent, 0644)
+ return os.WriteFile(filePath, defaultContent, 0600)
}
return nil
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func ensureFileExists(filePath string, defaultContent []byte) error { | |
| if _, err := os.Stat(filePath); os.IsNotExist(err) { | |
| return os.WriteFile(filePath, defaultContent, 0644) | |
| } | |
| return nil | |
| } | |
| func ensureFileExists(filePath string, defaultContent []byte) error { | |
| if _, err := os.Stat(filePath); os.IsNotExist(err) { | |
| return os.WriteFile(filePath, defaultContent, 0600) | |
| } | |
| return nil | |
| } |
🤖 Prompt for AI Agents
In wren-launcher/utils/docker.go around lines 575 to 580, ensureFileExists
currently writes new files with mode 0644 which is too permissive for secrets
like .env and config.yaml; change the os.WriteFile call to use file mode 0600 so
created files are only readable/writable by the owner, preserving the rest of
the function behavior.
| // It ensures that config.yaml, .env, and config.done are mounted, | ||
| // and initializes config.done with 'false' for setup flow control. | ||
| func RunStreamlitUIContainer(port int) error { | ||
|
|
||
| // Build the Docker image for the Streamlit UI | ||
| if err := buildStreamlitImage(); err != nil { | ||
| return err | ||
| } | ||
|
|
||
| // Get ~/.wrenai directory | ||
| wrenAIDir, err := TryGetWrenAIDir() | ||
| if err != nil { | ||
| return fmt.Errorf("failed to get ~/.wrenai: %v", err) | ||
| } | ||
|
|
||
| // Initialize config.done with 'false' | ||
| donePath, err := prepareConfigDoneFile(wrenAIDir) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to write to config.done: %v", err) | ||
| } | ||
|
|
||
| // Mount user config.yaml and .env for the UI to read/write | ||
| configPath, envPath, _ := getMountPaths(wrenAIDir) | ||
| _ = ensureFileExists(configPath, []byte("# Create a temporary yaml file")) | ||
| _ = ensureFileExists(envPath, []byte("# Put your API keys here\n")) | ||
|
|
||
| // run docker and mount volume | ||
| if err := runStreamlitContainer(configPath, envPath, donePath, port); err != nil { | ||
| return err | ||
| } | ||
|
|
||
| return nil | ||
| } | ||
|
|
There was a problem hiding this comment.
Fix incorrect handling of missing ~/.wrenai, validate port, and avoid container-name collisions.
- Create ~/.wrenai if TryGetWrenAIDir returns "" to avoid writing config.done into CWD.
- Validate port range early.
- Remove existing container before starting to prevent “Conflict. The container name is already in use”.
func RunStreamlitUIContainer(port int) error {
+ // Validate port
+ if port < 1 || port > 65535 {
+ return fmt.Errorf("invalid port: %d", port)
+ }
+
// Build the Docker image for the Streamlit UI
if err := buildStreamlitImage(); err != nil {
return err
}
// Get ~/.wrenai directory
wrenAIDir, err := TryGetWrenAIDir()
if err != nil {
return fmt.Errorf("failed to get ~/.wrenai: %v", err)
}
+ if wrenAIDir == "" {
+ homeDir, err := os.UserHomeDir()
+ if err != nil {
+ return fmt.Errorf("failed to resolve user home dir: %v", err)
+ }
+ wrenAIDir = path.Join(homeDir, ".wrenai")
+ if err := os.MkdirAll(wrenAIDir, 0700); err != nil {
+ return fmt.Errorf("failed to create %s: %v", wrenAIDir, err)
+ }
+ }
// Initialize config.done with 'false'
donePath, err := prepareConfigDoneFile(wrenAIDir)
if err != nil {
return fmt.Errorf("failed to write to config.done: %v", err)
}
// Mount user config.yaml and .env for the UI to read/write
configPath, envPath, _ := getMountPaths(wrenAIDir)
_ = ensureFileExists(configPath, []byte("# Create a temporary yaml file"))
_ = ensureFileExists(envPath, []byte("# Put your API keys here\n"))
// run docker and mount volume
+ // Ensure previous container is gone to avoid name conflicts
+ _ = RemoveContainerIfExists("wrenai-providers-setup")
if err := runStreamlitContainer(configPath, envPath, donePath, port); err != nil {
return err
}
return nil
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // RunStreamlitUIContainer builds and runs the Streamlit UI container. | |
| // It ensures that config.yaml, .env, and config.done are mounted, | |
| // and initializes config.done with 'false' for setup flow control. | |
| func RunStreamlitUIContainer(port int) error { | |
| // Build the Docker image for the Streamlit UI | |
| if err := buildStreamlitImage(); err != nil { | |
| return err | |
| } | |
| // Get ~/.wrenai directory | |
| wrenAIDir, err := TryGetWrenAIDir() | |
| if err != nil { | |
| return fmt.Errorf("failed to get ~/.wrenai: %v", err) | |
| } | |
| // Initialize config.done with 'false' | |
| donePath, err := prepareConfigDoneFile(wrenAIDir) | |
| if err != nil { | |
| return fmt.Errorf("failed to write to config.done: %v", err) | |
| } | |
| // Mount user config.yaml and .env for the UI to read/write | |
| configPath, envPath, _ := getMountPaths(wrenAIDir) | |
| _ = ensureFileExists(configPath, []byte("# Create a temporary yaml file")) | |
| _ = ensureFileExists(envPath, []byte("# Put your API keys here\n")) | |
| // run docker and mount volume | |
| if err := runStreamlitContainer(configPath, envPath, donePath, port); err != nil { | |
| return err | |
| } | |
| return nil | |
| } | |
| func RunStreamlitUIContainer(port int) error { | |
| // Validate port | |
| if port < 1 || port > 65535 { | |
| return fmt.Errorf("invalid port: %d", port) | |
| } | |
| // Build the Docker image for the Streamlit UI | |
| if err := buildStreamlitImage(); err != nil { | |
| return err | |
| } | |
| // Get ~/.wrenai directory | |
| wrenAIDir, err := TryGetWrenAIDir() | |
| if err != nil { | |
| return fmt.Errorf("failed to get ~/.wrenai: %v", err) | |
| } | |
| if wrenAIDir == "" { | |
| homeDir, err := os.UserHomeDir() | |
| if err != nil { | |
| return fmt.Errorf("failed to resolve user home dir: %v", err) | |
| } | |
| wrenAIDir = path.Join(homeDir, ".wrenai") | |
| if err := os.MkdirAll(wrenAIDir, 0700); err != nil { | |
| return fmt.Errorf("failed to create %s: %v", wrenAIDir, err) | |
| } | |
| } | |
| // Initialize config.done with 'false' | |
| donePath, err := prepareConfigDoneFile(wrenAIDir) | |
| if err != nil { | |
| return fmt.Errorf("failed to write to config.done: %v", err) | |
| } | |
| // Mount user config.yaml and .env for the UI to read/write | |
| configPath, envPath, _ := getMountPaths(wrenAIDir) | |
| _ = ensureFileExists(configPath, []byte("# Create a temporary yaml file")) | |
| _ = ensureFileExists(envPath, []byte("# Put your API keys here\n")) | |
| // Ensure previous container is gone to avoid name conflicts | |
| _ = RemoveContainerIfExists("wrenai-providers-setup") | |
| // run docker and mount volume | |
| if err := runStreamlitContainer(configPath, envPath, donePath, port); err != nil { | |
| return err | |
| } | |
| return nil | |
| } |
🤖 Prompt for AI Agents
In wren-launcher/utils/docker.go around lines 582 to 615, handle three issues:
if TryGetWrenAIDir returns an empty string create the ~/.wrenai directory (or
return an explicit error) before attempting to write config.done so we don't
accidentally write to the CWD; validate the provided port immediately (ensure
it's an integer within 1–65535 and return a clear error if not) before building
images or creating files; and before running the Streamlit container check for
an existing container with the intended name and remove/stop it (or use a unique
name) to avoid "container name already in use" conflicts, then proceed to build
image, prepare files, and start the container.
wren-launcher/utils/docker.go
Outdated
There was a problem hiding this comment.
🛠️ Refactor suggestion
Guard against missing ~/.wrenai in readiness check.
Prevent reading ./config.done implicitly when the dir is absent.
func IsCustomConfigReady() bool {
wrenAIDir, err := TryGetWrenAIDir()
if err != nil {
return false
}
+ if wrenAIDir == "" {
+ return false
+ }
configDonePath := path.Join(wrenAIDir, "config.done")
data, err := os.ReadFile(configDonePath)
if err != nil {
return false
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // IsCustomConfigReady checks whether the config.done file contains 'true', | |
| // indicating that the Streamlit configuration process is complete. | |
| func IsCustomConfigReady() bool { | |
| wrenAIDir, err := TryGetWrenAIDir() | |
| if err != nil { | |
| return false | |
| } | |
| configDonePath := path.Join(wrenAIDir, "config.done") | |
| data, err := os.ReadFile(configDonePath) | |
| if err != nil { | |
| return false | |
| } | |
| // Trim whitespace and compare case-insensitively | |
| trimmed := strings.TrimSpace(string(data)) | |
| return strings.EqualFold(trimmed, "true") | |
| } | |
| // IsCustomConfigReady checks whether the config.done file contains 'true', | |
| // indicating that the Streamlit configuration process is complete. | |
| func IsCustomConfigReady() bool { | |
| wrenAIDir, err := TryGetWrenAIDir() | |
| if err != nil { | |
| return false | |
| } | |
| if wrenAIDir == "" { | |
| return false | |
| } | |
| configDonePath := path.Join(wrenAIDir, "config.done") | |
| data, err := os.ReadFile(configDonePath) | |
| if err != nil { | |
| return false | |
| } | |
| // Trim whitespace and compare case-insensitively | |
| trimmed := strings.TrimSpace(string(data)) | |
| return strings.EqualFold(trimmed, "true") | |
| } |
🤖 Prompt for AI Agents
In wren-launcher/utils/docker.go around lines 692–709, the readiness check can
end up reading ./config.done if the WrenAI dir is missing; after obtaining
wrenAIDir from TryGetWrenAIDir, verify that wrenAIDir is non-empty and that the
path exists and is a directory (use os.Stat and check for os.IsNotExist /
FileInfo.IsDir), return false if it does not exist or isn’t a directory, and use
filepath.Join to build configDonePath before attempting to read the file.
2 participants
Create a UI to help users set up the model.
Summary by CodeRabbit
New Features
Chores