├── .env.example ├── .gitignore ├── LICENSE.md ├── README.md ├── install.sh ├── requirements.txt ├── screenshots ├── 01.png ├── 01_up.png ├── 02.png ├── 03.png └── 04.png ├── setup.py └── src └── shellsage ├── __init__.py ├── cli.py ├── command_generator.py ├── error_interceptor.py ├── helpers.py ├── llm_handler.py └── model_manager.py /.env.example: -------------------------------------------------------------------------------- 1 | # Shell Sage Configuration 2 | # ------------------------ 3 | 4 | # Operation Mode (local/api) 5 | MODE=local # Options: 'local' (Ollama) | 'api' (Cloud providers) 6 | OLLAMA_HOST=http://localhost:11434 7 | 8 | # Local Configuration 9 | LOCAL_MODEL=llama3:8b-instruct-q4_1 # Ollama model name for local mode 10 | 11 | # API Configuration 12 | ACTIVE_API_PROVIDER=groq # Current provider: groq, openai, anthropic, fireworks, openrouter, deepseek 13 | API_MODEL=mixtral-8x7b-32768 # Provider-specific model name, you can add any model supported by your provider 14 | 15 | # Provider API Keys (only set for your active provider) 16 | GROQ_API_KEY= # For Groq Cloud (https://console.groq.com) 17 | OPENAI_API_KEY= # For OpenAI (https://platform.openai.com) 18 | ANTHROPIC_API_KEY= # For Anthropic Claude (https://console.anthropic.com) 19 | FIREWORKS_API_KEY= # For Fireworks AI (https://app.fireworks.ai) 20 | OPENROUTER_API_KEY= # For OpenRouter (https://openrouter.ai) 21 | DEEPSEEK_API_KEY= # For Deepseek (https://platform.deepseek.com) -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | shellsage_env/ 2 | 3 | # Python 4 | __pycache__/ 5 | *.egg-info/ 6 | .installed.cfg 7 | *.egg 8 | 9 | # Virtual Environments 10 | .env 11 | .venv 12 | env/ 13 | venv/ 14 | ENV/ 15 | env.bak/ 16 | venv.bak/ 17 | .nox/ 18 | .tox/ 19 | 20 | # IDE 21 | .idea/ 22 | .vscode/ 23 | *.swp 24 | *.swo 25 | *~ 26 | 27 | # Version control 28 | .git/ 29 | .hg/ 30 | .svn/ 31 | CVS/ 32 | _darcs/ 33 | .bzr/ 34 | 35 | # Project specific 36 | .cache/ 37 | 38 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [2025] [Dheeraj C L] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Shell Sage 🐚✨ 2 | 3 | **Intelligent Terminal Companion | AI-Powered Terminal Assistant** 4 | *(Development Preview - v0.2.0)* 5 | 6 | --- 7 | 8 | ## Features 9 | 10 | ### 🌟 Next-Gen Terminal Experience 11 | - 🏠 Local AI Support (Ollama) & Cloud AI (Groq) 12 | - 🔍 Context-aware error diagnosis 13 | - 🪄 Natural language to command translation 14 | - ⚡ Safe command execution workflows 15 | 16 | ## 🔧 Core Capabilities 17 | 18 | ### Error Diagnosis 19 | 20 | ```bash 21 | # Error analysis example 22 | $ rm -rf /important-folder 23 | 🔎 Analysis → 🛠️ Fix: `rm -rf ./important-folder` 24 | ``` 25 | ![Error Analysis](screenshots/01_up.png) 26 | 27 | ### Natural Language to Commands 28 | 29 | ```bash 30 | # Command generation 31 | $ shellsage ask "find large files over 1GB" 32 | # → find / -type f -size +1G -exec ls -lh {} \; 33 | ``` 34 | ![Command generation](screenshots/02.png) 35 | 36 | ### ⚡ Interactive Workflows 37 | - Confirm before executing generated commands 38 | - Step-by-step complex operations 39 | - Safety checks for destructive commands 40 | 41 | 42 | ### 🌐 Supported API Providers 43 | - Groq 44 | - OpenAI 45 | - Anthropic 46 | - Fireworks.ai 47 | - OpenRouter 48 | - Deepseek 49 | 50 | *Switch providers with `shellsage config --provider `* 51 | 52 | --- 53 | 54 | ## Installation 55 | 56 | ### Prerequisites 57 | - Python 3.8+ 58 | - (4GB+ recommended for local models) 59 | 60 | ```bash 61 | # 1. Clone & install Shell Sage 62 | git clone https://github.com/dheerajcl/Terminal_assistant.git 63 | cd Terminal_assistant 64 | ./install.sh 65 | 66 | # 2. Install Ollama for local AI 67 | curl -fsSL https://ollama.com/install.sh | sh 68 | 69 | # 3. Get base model (3.8GB) 70 | #for example 71 | ollama pull llama3:8b-instruct-q4_1 72 | 73 | # or API key (Currently supports Groq, OpenAI, Anthropic, Fireworks, OpenRouter, Deepseek) 74 | # put your desired provider api in .env file 75 | shellsage config --mode api --provider groq 76 | 77 | 78 | ``` 79 | 80 | ### Configuration Notes 81 | - Rename `.env.example` → `.env` and populate required values 82 | - API performance varies by provider (Groq fastest, Anthropic most capable) 83 | - Local models need 4GB+ RAM (llama3:8b) to 16GB+ (llama3:70b) 84 | - Response quality depends on selected model capabilities 85 | 86 | 87 | ### Custom Model Selection 88 | 89 | While we provide common defaults for each AI provider, many services offer hundreds of models. To use a specific model: 90 | 91 | - Check your provider's documentation for available models 92 | - Set in .env: 93 | ``` 94 | API_PROVIDER=openrouter 95 | API_MODEL=your-model-name-here # e.g. google/gemini-2.0-pro-exp-02-05:free 96 | 97 | ``` 98 | 99 | --- 100 | 101 | 102 | 103 | ## Configuration 104 | 105 | ### First-Time Setup 106 | ```bash 107 | # Interactive configuration wizard 108 | shellsage setup 109 | 110 | ? Select operation mode: 111 | ▸ Local (Privacy-first, needs 4GB+ RAM) 112 | API (Faster but requires internet) 113 | 114 | ? Choose local model: 115 | ▸ llama3:8b-instruct-q4_1 (Recommended) 116 | mistral:7b-instruct-v0.3 117 | phi3:mini-128k-instruct 118 | 119 | # If API mode selected: 120 | ? Choose API provider: 121 | ▸ Groq 122 | OpenAI 123 | Anthropic 124 | Fireworks 125 | Deepseek 126 | 127 | ? Enter Groq API key: [hidden input] 128 | 129 | ? Select Groq model: 130 | ▸ mixtral-8x7b-32768 131 | llama3-70b-8192 # It isn't necessary to select models from the shown list, you can add any model of your choice supported by your provider in your .env `API_MODEL=` 132 | 133 | ✅ API configuration updated! 134 | 135 | ``` 136 | 137 | ### Runtime Control 138 | 139 | ```bash 140 | # Switch modes 141 | shellsage config --mode api # or 'local' 142 | 143 | 144 | # Switch to specific model 145 | shellsage config --mode local --model 146 | 147 | # Interactive switch 148 | shellsage config --mode local 149 | ? Select local model: 150 | ▸ llama3:8b-instruct-q4_1 151 | mistral:7b-instruct-v0.3 152 | phi3:mini-128k-instruct 153 | ``` 154 | 155 | ![interactive_flow1](screenshots/03.png) 156 | 157 | ![interactive_flow2](screenshots/04.png) 158 | 159 | --- 160 | 161 | ## Development Status 🚧 162 | 163 | Shell Sage is currently in **alpha development**. 164 | 165 | **Known Limitations**: 166 | - Limited Windows support 167 | - Compatibility issues with zsh, fish 168 | - Occasional false positives in error detection 169 | - API mode requires provider-specific key 170 | 171 | **Roadmap**: 172 | - [x] Local LLM support 173 | - [x] Hybrid cloud(api)/local mode switching 174 | - [x] Model configuration wizard 175 | - [ ] Better Context Aware 176 | - [ ] Windows PowerShell integration 177 | - [ ] Tmux Integration 178 | - [ ] CI/CD error pattern database 179 | 180 | --- 181 | 182 | ## Contributing 183 | 184 | We welcome contributions! Please follow these steps: 185 | 186 | 1. Fork the repository 187 | 2. Create feature branch (`git checkout -b feat/amazing-feature`) 188 | 3. Commit changes (`git commit -m 'Add amazing feature'`) 189 | 4. Push to branch (`git push origin feat/amazing-feature`) 190 | 5. Open Pull Request 191 | 192 | --- 193 | 194 | 195 | > **Note**: This project is not affiliated with any API or model providers. 196 | > Local models require adequate system resources. 197 | > Internet required for initial setup and API mode. 198 | > Use at your own risk with critical operations. 199 | > Always verify commands before execution 200 | -------------------------------------------------------------------------------- /install.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Shell Sage Installation Script 4 | set -e 5 | 6 | # Colors for output 7 | RED='\033[0;31m' 8 | GREEN='\033[0;32m' 9 | YELLOW='\033[1;33m' 10 | NC='\033[0m' # No Color 11 | 12 | echo -e "${YELLOW}🚀 Starting Shell Sage Installation...${NC}" 13 | 14 | # # Check for Python 3.8+ 15 | # if ! python3 -c 'import sys; exit(1) if sys.version_info < (3,8) else exit(0)' &>/dev/null; then 16 | # echo -e "${RED}❌ Python 3.8 or newer is required${NC}" 17 | # exit 1 18 | # fi 19 | 20 | # Create virtual environment 21 | if [ ! -d "shellsage_env" ]; then 22 | echo -e "${YELLOW}⚙️ Creating virtual environment...${NC}" 23 | python3 -m venv shellsage_env 24 | else 25 | echo -e "${YELLOW}⚙️ Using existing virtual environment...${NC}" 26 | fi 27 | 28 | # Activate virtual environment 29 | source shellsage_env/bin/activate 30 | 31 | # Install system dependencies 32 | echo -e "${YELLOW}⚙️ Checking for system dependencies...${NC}" 33 | if ! command -v ollama &>/dev/null; then 34 | echo -e "${YELLOW}⚠️ Ollama not found. Install from https://ollama.ai/ for local models${NC}" 35 | fi 36 | 37 | # Install Python dependencies 38 | echo -e "${YELLOW}⚙️ Installing Python dependencies...${NC}" 39 | pip install -U pip 40 | pip install -r requirements.txt 41 | 42 | # Install in editable mode 43 | echo -e "${YELLOW}⚙️ Installing Shell Sage...${NC}" 44 | pip install -e . 45 | 46 | # Post-install setup 47 | echo -e "${YELLOW}⚙️ Running initial configuration...${NC}" 48 | if [ ! -f .env ]; then 49 | echo -e "${YELLOW}⚙️ Creating .env file from example...${NC}" 50 | cp .env.example .env 51 | fi 52 | shellsage setup 53 | 54 | # Install shell hook 55 | echo -e "${YELLOW}⚙️ Installing shell hook...${NC}" 56 | HOOK=$'shell_sage_prompt() {\n local EXIT=$?\n local CMD=$(fc -ln -1 | awk \'{$1=$1}1\' | sed \'s/\\\\/\\\\\\\\/g\')\n [ $EXIT -ne 0 ] && shellsage run --analyze "$CMD" --exit-code $EXIT\n history -s "$CMD" # Force into session history\n}\nPROMPT_COMMAND="shell_sage_prompt"' 57 | 58 | if [ -f ~/.bashrc ]; then 59 | echo -e "\n# Shell Sage Hook\n$HOOK" >> ~/.bashrc 60 | echo -e "${GREEN}✅ Added to ~/.bashrc${NC}" 61 | # Refresh bash if we're in bash 62 | if [ -n "$BASH" ]; then 63 | source ~/.bashrc 64 | fi 65 | fi 66 | 67 | if [ -f ~/.zshrc ]; then 68 | echo -e "\n# Shell Sage Hook\n$HOOK" >> ~/.zshrc 69 | echo -e "${GREEN}✅ Added to ~/.zshrc${NC}" 70 | # Refresh zsh if we're in zsh 71 | if [ -n "$ZSH_VERSION" ]; then 72 | source ~/.zshrc 73 | fi 74 | fi 75 | 76 | echo -e "\n${GREEN}✅ Installation Complete!${NC}" 77 | echo -e "To start using Shell Sage:" 78 | echo -e "1. Activate environment: ${YELLOW}source shellsage_env/bin/activate${NC}" 79 | echo -e "2. Add api keys and your desired model supported by the listed providers manually in .env if the model you intend to use is not listed${NC}" 80 | echo -e "3. Test installation: ${YELLOW}shellsage ask 'update packages'${NC}" 81 | echo -e "4. For local models: ${YELLOW}ollama pull llama3:8b-instruct-q4_1${NC}" 82 | 83 | exit 0 -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | click>=8.1.0 2 | requests>=2.31.0 3 | openai>=1.12.0 4 | pyyaml>=6.0.1 5 | inquirer>=3.1.3 6 | ctransformers>=0.2.27 7 | python-dotenv>=1.0.0 8 | anthropic>=0.25.0 9 | rich>=13.7.1 10 | -------------------------------------------------------------------------------- /screenshots/01.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dheerajcl/Shellsage/17651dcb7399bdd9b3207fed64ebc306fcf10c78/screenshots/01.png -------------------------------------------------------------------------------- /screenshots/01_up.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dheerajcl/Shellsage/17651dcb7399bdd9b3207fed64ebc306fcf10c78/screenshots/01_up.png -------------------------------------------------------------------------------- /screenshots/02.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dheerajcl/Shellsage/17651dcb7399bdd9b3207fed64ebc306fcf10c78/screenshots/02.png -------------------------------------------------------------------------------- /screenshots/03.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dheerajcl/Shellsage/17651dcb7399bdd9b3207fed64ebc306fcf10c78/screenshots/03.png -------------------------------------------------------------------------------- /screenshots/04.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dheerajcl/Shellsage/17651dcb7399bdd9b3207fed64ebc306fcf10c78/screenshots/04.png -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup( 4 | name='shellsage', 5 | version='1.0.0', 6 | packages=find_packages(where='src'), 7 | package_dir={'': 'src'}, 8 | install_requires=[ 9 | 'click>=8.1.0', 10 | 'requests>=2.31.0', 11 | 'openai>=1.12.0', 12 | 'pyyaml>=6.0.1', 13 | 'inquirer>=3.1.3', 14 | 'ctransformers>=0.2.27', 15 | 'python-dotenv>=1.0.0' 16 | ], 17 | entry_points={ 18 | 'console_scripts': [ 19 | 'shellsage=shellsage.cli:cli', 20 | ], 21 | }, 22 | include_package_data=True, 23 | ) -------------------------------------------------------------------------------- /src/shellsage/__init__.py: -------------------------------------------------------------------------------- 1 | from .cli import cli 2 | from .error_interceptor import ErrorInterceptor 3 | 4 | __all__ = ['cli', 'ErrorInterceptor'] -------------------------------------------------------------------------------- /src/shellsage/cli.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import click 3 | import os 4 | import subprocess 5 | import inquirer 6 | from .error_interceptor import ErrorInterceptor 7 | from .command_generator import CommandGenerator 8 | from .model_manager import ModelManager, PROVIDERS 9 | from .helpers import update_env_file, update_env_variable 10 | from dotenv import load_dotenv 11 | import re 12 | from rich.console import Console 13 | from rich.panel import Panel 14 | from rich.markdown import Markdown 15 | from rich.syntax import Syntax 16 | from rich.columns import Columns 17 | 18 | @click.group() 19 | def cli(): 20 | """Terminal Assistant - Error Analysis & Command Generation""" 21 | 22 | @cli.command(context_settings={"ignore_unknown_options": True}) 23 | @click.argument('command', nargs=-1) 24 | @click.option('--analyze', is_flag=True, hidden=True) 25 | @click.option('--exit-code', type=int, hidden=True) 26 | def run(command, analyze, exit_code): 27 | """Execute command with error analysis""" 28 | interceptor = ErrorInterceptor() 29 | if analyze: 30 | interceptor.auto_analyze(' '.join(command), exit_code) 31 | else: 32 | interceptor.run_command(command) 33 | 34 | # cli.py - update the ask command 35 | 36 | @cli.command() 37 | @click.argument('query', nargs=-1, required=True) 38 | @click.option('--execute', is_flag=True, help='Execute commands with safety checks') 39 | def ask(query, execute): 40 | """Generate and execute commands with safety checks""" 41 | console = Console() 42 | generator = CommandGenerator() 43 | interceptor = ErrorInterceptor() 44 | 45 | # Add distro detection at the top of the ask command 46 | try: 47 | with open('/etc/os-release') as f: 48 | dist_info = {k.lower(): v.strip('"') for k,v in 49 | [line.split('=') for line in f if '=' in line]} 50 | dist_name = dist_info.get('pretty_name', 'Linux') 51 | except FileNotFoundError: 52 | import platform 53 | dist_name = f"{platform.system()} {platform.release()}" 54 | 55 | context = { 56 | 'os': dist_name, 57 | 'cwd': os.getcwd(), 58 | 'git': os.path.exists('.git'), 59 | 'history': interceptor.command_history 60 | } 61 | 62 | results = generator.generate_commands(' '.join(query), context) 63 | 64 | # Command Analysis Display 65 | console.print(Panel.fit("[bold cyan]COMMAND ANALYSIS[/]", style="cyan")) 66 | 67 | # Thinking Process 68 | thinking_items = [item for item in results if item['type'] == 'thinking'] 69 | if thinking_items: 70 | console.print(Panel.fit( 71 | "\n".join(f"[dim]› {item['content']}[/dim]" for item in thinking_items), 72 | title="[gold1]Thinking Process[/]", 73 | border_style="gold1", 74 | padding=(0, 2) 75 | )) 76 | 77 | # Main Results 78 | command_item = next((i for i in results if i['type'] == 'command'), None) 79 | 80 | analysis_col = [] 81 | details_col = [] 82 | 83 | for item in results: 84 | if item['type'] == 'warning': 85 | analysis_col.append(f"[red]⚠ {item['content']}[/]") 86 | elif item['type'] == 'analysis': 87 | analysis_col.append(f"[cyan]ⓘ {item['content']}[/]") 88 | elif item['type'] == 'details': 89 | details_col.append(f"[dim]{item['content']}[/]") 90 | 91 | console.print(Columns([ 92 | Panel.fit("\n".join(analysis_col), title="[blue]Analysis[/]", padding=(0, 1)), 93 | Panel.fit("\n".join(details_col), title="[grey70]Technical Details[/]", padding=(0, 1)) 94 | ], equal=True, expand=False)) 95 | 96 | if command_item and command_item['content']: 97 | # Clean markdown backticks before display 98 | clean_command = command_item['content'].strip('`').strip() 99 | console.print(Panel.fit( 100 | Syntax(clean_command, "bash", theme="monokai", line_numbers=False), 101 | title="[green]Generated Command[/]", 102 | border_style="green", 103 | padding=0 104 | )) 105 | 106 | if execute: 107 | if console.input("\n[bold gold1]› Execute command?[/] [[y]/n]: ").lower() != 'n': 108 | subprocess.run(command_item['content'], shell=True) 109 | else: 110 | console.print(Panel.fit( 111 | "No valid command generated", 112 | style="red" 113 | )) 114 | 115 | @cli.command() 116 | def install(): 117 | """Install automatic error handling""" 118 | hook = r""" 119 | shell_sage_prompt() { 120 | local EXIT=$? 121 | local CMD=$(fc -ln -1 | awk '{$1=$1}1' | sed 's/\\/\\\\/g') 122 | [ $EXIT -ne 0 ] && shellsage run --analyze "$CMD" --exit-code $EXIT 123 | history -s "$CMD" # Force into session history 124 | } 125 | PROMPT_COMMAND="shell_sage_prompt" 126 | """ 127 | click.echo("# Add this to your shell config:") 128 | click.echo(hook) 129 | click.echo("\n# Then run: source ~/.bashrc") 130 | 131 | @cli.command() 132 | def setup(): 133 | """Interactive configuration setup""" 134 | if not os.path.exists('.env'): 135 | click.echo("❌ Missing .env file - clone the repository properly") 136 | return 137 | 138 | # Mode selection 139 | mode_q = inquirer.List( 140 | 'mode', 141 | message="Select operation mode:", 142 | choices=['local', 'api'], 143 | default=os.getenv('MODE', 'local') 144 | ) 145 | answers = inquirer.prompt([mode_q]) 146 | mode = answers['mode'] 147 | 148 | if mode == 'local': 149 | # Local mode configuration 150 | manager = ModelManager() 151 | models = manager.get_ollama_models() 152 | 153 | if not models: 154 | click.echo("❌ No local models found. Install Ollama first.") 155 | return 156 | 157 | model_q = inquirer.List( 158 | 'model', 159 | message="Select local model:", 160 | choices=models, 161 | default=os.getenv('LOCAL_MODEL') 162 | ) 163 | answers = inquirer.prompt([model_q]) 164 | update_env_variable('LOCAL_MODEL', answers['model']) 165 | update_env_variable('MODE', 'local') 166 | click.echo(f"✅ Local mode configured with model: {answers['model']}") 167 | 168 | elif mode == 'api': 169 | # API provider selection 170 | update_env_variable('MODE', 'api') 171 | provider_q = inquirer.List( 172 | 'provider', 173 | message="Select API Provider:", 174 | choices=list(PROVIDERS.keys()) 175 | ) 176 | answers = inquirer.prompt([provider_q]) 177 | provider = answers['provider'] 178 | 179 | # Key entry for any provider 180 | with open('.env') as f: 181 | env_content = f.read() 182 | existing_key = re.search(f"{provider.upper()}_API_KEY=(.*)", env_content) 183 | 184 | if not existing_key or not existing_key.group(1).strip(): 185 | key_q = inquirer.Text( 186 | 'key', 187 | message=f"Enter {provider} API key:" 188 | ) 189 | key_answers = inquirer.prompt([key_q]) 190 | update_env_file(provider, key_answers['key']) 191 | load_dotenv(override=True) 192 | 193 | # Model selection for chosen provider 194 | models = PROVIDERS[provider]['models'] 195 | model_q = inquirer.List( 196 | 'model', 197 | message=f"Select {provider} model:", 198 | choices=models 199 | ) 200 | model_answers = inquirer.prompt([model_q]) 201 | update_env_variable('ACTIVE_API_PROVIDER', provider) 202 | update_env_variable('API_MODEL', model_answers['model']) 203 | click.echo(f"✅ API mode configured with {provider}/{model_answers['model']}") 204 | 205 | @cli.command() 206 | @click.option('--mode', type=click.Choice(['local', 'api'])) 207 | @click.option('--provider', type=click.Choice(list(PROVIDERS.keys()))) 208 | @click.option('--model', help="Specify model name") 209 | def config(mode, provider, model): 210 | """Configure operation mode and models""" 211 | manager = ModelManager() 212 | 213 | if mode == 'local': 214 | models = manager.get_ollama_models() 215 | question = [ 216 | inquirer.List('model', 217 | message="Select local model:", 218 | choices=models, 219 | default=os.getenv('LOCAL_MODEL') 220 | ) 221 | ] 222 | answers = inquirer.prompt(question) 223 | update_env_variable('LOCAL_MODEL', answers['model']) 224 | update_env_variable('MODE', 'local') 225 | click.echo(f"✅ Switched to local mode using {answers['model']}") 226 | 227 | elif mode == 'api': 228 | if not provider: 229 | # Interactive provider selection 230 | provider_q = inquirer.List( 231 | 'provider', 232 | message="Select API Provider:", 233 | choices=list(PROVIDERS.keys()) 234 | ) 235 | provider = inquirer.prompt([provider_q])['provider'] 236 | 237 | # Update provider FIRST before checking key 238 | update_env_variable('ACTIVE_API_PROVIDER', provider) 239 | load_dotenv(override=True) 240 | 241 | # Now check key in updated environment 242 | key = os.getenv(f"{provider.upper()}_API_KEY") 243 | if not key: 244 | key_q = inquirer.Text( 245 | 'key', 246 | message=f"Enter {provider} API key:" 247 | ) 248 | key_answers = inquirer.prompt([key_q]) 249 | # Update .env 250 | update_env_variable(f"{provider.upper()}_API_KEY", key_answers['key']) 251 | 252 | # Model selection 253 | models = PROVIDERS[provider]['models'] 254 | if not model: 255 | model_q = inquirer.List( 256 | 'model', 257 | message=f"Select {provider} model:", 258 | choices=models 259 | ) 260 | model = inquirer.prompt([model_q])['model'] 261 | 262 | # Update config 263 | manager.switch_mode('api', model_name=model) 264 | click.echo(f"✅ Switched to API mode using {provider}/{model}") 265 | 266 | else: 267 | current_mode = manager.config['mode'] 268 | click.echo(f"Current mode: {current_mode}") 269 | if current_mode == 'local': 270 | click.echo(f"Local model: {manager.config['local']['model']}") 271 | else: 272 | provider = os.getenv('ACTIVE_API_PROVIDER', 'groq') 273 | click.echo(f"API Provider: {provider}") 274 | 275 | @cli.command() 276 | @click.option('--provider', type=click.Choice(['ollama', 'huggingface'])) 277 | def models(provider): 278 | """Manage local models""" 279 | manager = ModelManager() 280 | 281 | if provider: 282 | manager.config['local']['provider'] = provider 283 | manager._save_config() 284 | 285 | click.echo(f"Local provider: {manager.config['local']['provider']}") 286 | 287 | if manager.config['local']['provider'] == 'ollama': 288 | models = manager.get_ollama_models() 289 | click.echo("\nInstalled Ollama models:") 290 | else: 291 | models = manager._get_hf_models() 292 | click.echo("\nInstalled HuggingFace models:") 293 | 294 | for model in models: 295 | click.echo(f"- {model}") 296 | 297 | 298 | if __name__ == "__main__": 299 | cli() -------------------------------------------------------------------------------- /src/shellsage/command_generator.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from .model_manager import ModelManager 4 | 5 | 6 | class CommandGenerator: 7 | def __init__(self): 8 | self.manager = ModelManager() 9 | 10 | def generate_commands(self, query, context=None): 11 | try: 12 | prompt = self._build_prompt(query, context) 13 | response = self.manager.generate(prompt) 14 | 15 | # Check if response contains thinking tokens 16 | has_thinking = '' in response and '' in response 17 | 18 | if has_thinking: 19 | thoughts = [] 20 | remaining_response = response 21 | while '' in remaining_response and '' in remaining_response: 22 | think_start = remaining_response.find('') + len('') 23 | think_end = remaining_response.find('') 24 | if think_start > -1 and think_end > -1: 25 | thought = remaining_response[think_start:think_end].strip() 26 | thoughts.append(thought) 27 | remaining_response = remaining_response[think_end + len(''):] 28 | 29 | final_response = remaining_response.strip() 30 | return self._format_thinking_response(thoughts, final_response) 31 | else: 32 | return self._parse_response(response) 33 | 34 | except Exception as e: 35 | return [{ 36 | 'type': 'warning', 37 | 'content': f"Error: {str(e)}" 38 | }, { 39 | 'type': 'command', 40 | 'content': None, 41 | 'details': None 42 | }] 43 | 44 | def _build_prompt(self, query, context): 45 | # Determine the primary context based on the query and environment 46 | system_context = f"""SYSTEM: You are a Linux terminal expert. Generate exactly ONE command or command sequence. 47 | Primary focus is on system-level operations (package management, system updates, file operations). 48 | Only consider Git operations if the query explicitly mentions Git/repository operations. 49 | 50 | USER QUERY: {query} 51 | 52 | RESPONSE FORMAT: 53 | 🧠 Analysis: [1-line explanation] 54 | 🛠️ Command: ```[executable command(s)]``` 55 | 📝 Details: [technical specifics] 56 | ⚠️ Warning: [if dangerous] 57 | 58 | CURRENT CONTEXT: 59 | - OS: {context.get('os', 'Linux')} 60 | - Directory: {context.get('cwd', 'Unknown')} 61 | {f'- Git repo: Yes (only relevant for Git-specific queries)' if context.get('git') else ''} 62 | 63 | PRIORITY ORDER: 64 | 1. System-level operations (apt, dnf, pacman, etc.) 65 | 2. File system operations 66 | 3. Repository operations (only if explicitly requested) 67 | 68 | EXAMPLES: 69 | Query: "update packages" 70 | 🧠 Analysis: Update system packages using the appropriate package manager 71 | 🛠️ Command: ```sudo apt update && sudo apt upgrade -y``` 72 | 📝 Details: Updates package lists and upgrades all installed packages 73 | ⚠️ Warning: System may require restart after certain updates 74 | 75 | Query: "update git repo" 76 | 🧠 Analysis: Update local Git repository with remote changes 77 | 🛠️ Command: ```git pull origin main``` 78 | 📝 Details: Fetches and merges changes from the remote repository 79 | ⚠️ Warning: Ensure working directory is clean before updating 80 | """ 81 | return system_context 82 | 83 | 84 | def _format_thinking_response(self, thoughts, final_response): 85 | results = [] 86 | 87 | # Add thinking process to results 88 | for thought in thoughts: 89 | results.append({ 90 | 'type': 'thinking', 91 | 'content': thought 92 | }) 93 | 94 | # Parse the final response 95 | components = self._parse_response(final_response) 96 | 97 | # Clean up any duplicate sections caused by thinking process 98 | seen_types = set() 99 | cleaned_components = [] 100 | 101 | for comp in components: 102 | if comp['type'] not in seen_types and comp['content']: 103 | seen_types.add(comp['type']) 104 | # Remove any duplicate content within the same component 105 | if isinstance(comp['content'], str): 106 | comp['content'] = '\n'.join(dict.fromkeys(comp['content'].split('\n'))) 107 | cleaned_components.append(comp) 108 | 109 | results.extend(cleaned_components) 110 | return results 111 | 112 | 113 | def _parse_response(self, response): 114 | # Clean up response by removing any remaining XML-like tags 115 | cleaned = re.sub(r'<[^>]+>', '', response) 116 | 117 | components = { 118 | 'analysis': None, 119 | 'command': None, 120 | 'details': None, 121 | 'warning': None 122 | } 123 | 124 | markers = { 125 | 'analysis': ['🧠', 'Analysis:'], 126 | 'command': ['🛠️', 'Command:'], 127 | 'details': ['📝', 'Details:'], 128 | 'warning': ['⚠️', 'Warning:'] 129 | } 130 | 131 | lines = cleaned.split('\n') 132 | current_section = None 133 | 134 | for line in lines: 135 | line = line.strip() 136 | if not line: 137 | continue 138 | 139 | for section, section_markers in markers.items(): 140 | if any(marker in line for marker in section_markers): 141 | current_section = section 142 | for marker in section_markers: 143 | line = line.replace(marker, '').strip() 144 | components[section] = line 145 | break 146 | 147 | if current_section and not any( 148 | marker in line for markers_list in markers.values() 149 | for marker in markers_list 150 | ): 151 | if components[current_section]: 152 | components[current_section] += '\n' + line 153 | else: 154 | components[current_section] = line 155 | 156 | return [{ 157 | 'type': key, 158 | 'content': '\n'.join(dict.fromkeys(value.strip().split('\n'))) if value else None 159 | } for key, value in components.items()] -------------------------------------------------------------------------------- /src/shellsage/error_interceptor.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import sys 3 | import os 4 | import re 5 | import yaml 6 | import click 7 | from collections import deque 8 | from .llm_handler import DeepSeekLLMHandler 9 | from rich.console import Console 10 | from rich.panel import Panel 11 | from rich.text import Text 12 | from rich.syntax import Syntax 13 | from rich.columns import Columns 14 | from rich.rule import Rule 15 | from rich.markdown import Markdown 16 | from rich.console import Group 17 | 18 | class ErrorInterceptor: 19 | def __init__(self): 20 | self.llm_handler = DeepSeekLLMHandler() 21 | self.command_history = deque(maxlen=20) # Increased history depth 22 | self.last_command = "" 23 | self.context_cache = {} 24 | 25 | def run_command(self, command): 26 | """Execute command with error interception""" 27 | try: 28 | full_cmd = ' '.join(command) 29 | self.last_command = full_cmd 30 | # Maintain full session history while respecting maxlen 31 | if full_cmd != self.command_history[-1] if self.command_history else True: 32 | self.command_history.append(full_cmd) 33 | 34 | # Execute with live terminal interaction 35 | result = subprocess.run( 36 | full_cmd, 37 | shell=True, 38 | check=False, 39 | stdin=sys.stdin, 40 | capture_output=True, 41 | text=True 42 | ) 43 | 44 | if result.returncode != 0: 45 | self.context_cache = self._get_additional_context() # Cache context 46 | self._handle_error(result, self.context_cache) 47 | 48 | sys.exit(result.returncode) 49 | 50 | except Exception as e: 51 | print(f"\n\033[91mExecution Error: {e}\033[0m") 52 | sys.exit(1) 53 | 54 | def auto_analyze(self, command, exit_code): 55 | """Automatically analyze failed commands from shell hook""" 56 | self.last_command = command 57 | self.command_history.append(command) 58 | result = subprocess.CompletedProcess( 59 | args=command, 60 | returncode=exit_code, 61 | stdout='', 62 | stderr=self._get_native_error(command) 63 | ) 64 | self._handle_error(result, self.context_cache) 65 | 66 | def _handle_error(self, result, context): 67 | """Process and analyze command errors""" 68 | # Get relevant files from command history 69 | relevant_files = self._get_relevant_files_from_history() 70 | 71 | error_context = { 72 | 'command': self.last_command, 73 | 'error_output': self._get_full_error_output(result), 74 | 'cwd': os.getcwd(), 75 | 'exit_code': result.returncode, 76 | 'history': list(self.command_history), 77 | 'relevant_files': relevant_files, 78 | **context 79 | } 80 | 81 | # Enhanced context for file operations 82 | parts = self.last_command.split() 83 | if len(parts) > 0: 84 | base_cmd = parts[0] 85 | error_context['man_excerpt'] = self._get_man_page(base_cmd) 86 | 87 | if os.getenv('SHELLSAGE_DEBUG'): 88 | print("\n\033[90m[DEBUG] Error Context:") 89 | print(yaml.dump(error_context, allow_unicode=True) + "\033[0m") 90 | 91 | print("\n\033[90m🔎 Analyzing error...\033[0m") 92 | solution = self.llm_handler.get_error_solution(error_context) 93 | 94 | if solution: 95 | self._show_analysis(solution, error_context) 96 | else: 97 | print("\n\033[91mError: Could not get analysis\033[0m") 98 | 99 | def _get_relevant_files_from_history(self): 100 | """Extract recently referenced files from command history""" 101 | files = [] 102 | git_operations = ['add', 'commit', 'push', 'pull'] 103 | 104 | for cmd in reversed(list(self.command_history)[:-1]): 105 | parts = cmd.split() 106 | if parts and parts[0] == 'git' and len(parts) > 1: 107 | if parts[1] in git_operations and len(parts) > 2: 108 | files.append(parts[-1]) 109 | elif parts and parts[0] in ['touch', 'mkdir', 'cp', 'mv', 'vim', 'nano']: 110 | files.append(parts[-1]) 111 | 112 | if len(files) >= 3: 113 | break 114 | return files 115 | 116 | def _get_man_page(self, command): 117 | """Get relevant sections from man page""" 118 | try: 119 | # Special case for git 120 | if command == 'git': 121 | result = subprocess.run( 122 | 'git status --porcelain', 123 | shell=True, 124 | capture_output=True, 125 | text=True 126 | ) 127 | if result.returncode == 0 and not result.stdout.strip(): 128 | return "Git status: No changes to commit (working directory clean)" 129 | 130 | result = subprocess.run( 131 | f'man {command} 2>/dev/null | col -b', 132 | shell=True, 133 | capture_output=True, 134 | text=True 135 | ) 136 | 137 | if result.returncode == 0: 138 | content = result.stdout 139 | 140 | # Extract relevant sections 141 | sections = [] 142 | current_section = None 143 | for line in content.split('\n'): 144 | if line.upper() in ['NAME', 'SYNOPSIS', 'DESCRIPTION']: 145 | current_section = line 146 | sections.append(line) 147 | elif current_section and line.startswith(' '): 148 | sections.append(line.strip()) 149 | if len(sections) > 10: # Limit size 150 | break 151 | 152 | return '\n'.join(sections) 153 | return "No manual entry available" 154 | except Exception: 155 | return "Error retrieving manual page" 156 | 157 | # Replace _get_full_error_output in ErrorInterceptor 158 | def _get_full_error_output(self, result): 159 | """Combine stderr/stdout and sanitize with context enhancement""" 160 | error_output = '' 161 | if hasattr(result, 'stderr') and result.stderr: 162 | error_output += result.stderr if isinstance(result.stderr, str) else result.stderr.decode() 163 | if hasattr(result, 'stdout') and result.stdout: 164 | error_output += '\n' + (result.stdout if isinstance(result.stdout, str) else result.stdout.decode()) 165 | 166 | # Clean ANSI color codes 167 | clean_error = re.sub(r'\x1B\[[0-?]*[ -/]*[@-~]', '', error_output).strip() 168 | 169 | # Try to enhance error with additional context 170 | enhanced_error = clean_error 171 | 172 | # Check for common error patterns and add hints 173 | if "permission denied" in clean_error.lower(): 174 | enhanced_error += "\nHint: This may be a permissions issue. Current user: " + os.getenv('USER', 'unknown') 175 | elif "command not found" in clean_error.lower(): 176 | enhanced_error += "\nHint: Command may not be installed or not in PATH" 177 | elif "no such file" in clean_error.lower(): 178 | enhanced_error += "\nHint: File or directory does not exist in the current context" 179 | 180 | # Check for git commit without add 181 | if ("git commit" in self.last_command.lower() and 182 | not any("git add" in cmd.lower() for cmd in self.command_history) and 183 | "no changes added to commit" in clean_error.lower()): 184 | enhanced_error += "\nHint: No files staged for commit. Did you forget 'git add'?" 185 | 186 | return enhanced_error 187 | 188 | def _show_analysis(self, solution, context): 189 | """Display analysis with thinking process""" 190 | console = Console() 191 | 192 | # Extract thinking blocks first 193 | thoughts = [] 194 | remaining = solution 195 | while '' in remaining and '' in remaining: 196 | think_start = remaining.find('') + len('') 197 | think_end = remaining.find('') 198 | if think_start > -1 and think_end > -1: 199 | thoughts.append(remaining[think_start:think_end].strip()) 200 | remaining = remaining[think_end + len(''):] 201 | 202 | console.print("\n[bold cyan]Error Analysis[/bold cyan]") 203 | 204 | # Display thinking process if any 205 | if thoughts: 206 | console.print(Panel( 207 | "\n".join(f"[dim]› {thought}[/dim]" for thought in thoughts), 208 | title="[gold1]Cognitive Process[/]", 209 | border_style="gold1", 210 | padding=(0, 2) 211 | )) 212 | 213 | # Context information 214 | context_content = [] 215 | if context['history']: 216 | context_content.append( 217 | Panel("\n".join(f"[dim]› {cmd}[/dim]" for cmd in context['history'][-3:]), 218 | title="[grey70]Recent Commands[/]", 219 | border_style="grey58") 220 | ) 221 | 222 | if context.get('relevant_files'): 223 | context_content.append( 224 | Panel("\n".join(f"[dim]› {file}[/dim]" for file in context['relevant_files']), 225 | title="[grey70]Related Files[/]", 226 | border_style="grey58") 227 | ) 228 | 229 | if context.get('man_excerpt') and "No manual entry" not in context['man_excerpt']: 230 | context_content.append( 231 | Panel( 232 | Syntax(context['man_excerpt'], "man", theme="ansi_light", line_numbers=False), 233 | title="[bold medium_blue]📘 MANUAL REFERENCE[/]", 234 | border_style="bright_blue", 235 | padding=(0, 1), 236 | # subtitle=f"for {os.path.basename(context['command'].split()[0])}" 237 | ) 238 | ) 239 | 240 | if context_content: 241 | console.print(Columns(context_content, equal=True, expand=False)) 242 | 243 | # Error Components 244 | components = { 245 | 'cause': re.search(r'🔍 Root Cause: (.+?)(?=\n🛠️|\n📚|\n⚠️|\n🔒|$)', remaining, re.DOTALL), 246 | 'fix': re.search(r'🛠️ Fix: (`{1,3}(.*?)`{1,3}|([^\n]+))', remaining, re.DOTALL), 247 | 'explanation': re.search(r'📚 Technical Explanation: (.+?)(?=\n⚠️|\n🔒|$)', remaining, re.DOTALL), 248 | 'risk': re.search(r'⚠️ Potential Risks: (.+?)(?=\n🔒|$)', remaining, re.DOTALL), 249 | 'prevention': re.search(r'🔒 Prevention Tip: (.+?)(?=\n|$)', remaining, re.DOTALL) 250 | } 251 | 252 | # Main Analysis Content 253 | analysis_blocks = [] 254 | if components['cause']: 255 | analysis_blocks.append(Markdown(f"**Root Cause**\n{components['cause'].group(1)}")) 256 | if components['explanation']: 257 | analysis_blocks.append(Markdown(f"**Technical Explanation**\n{components['explanation'].group(1)}")) 258 | 259 | if analysis_blocks: 260 | console.print(Panel( 261 | Group(*analysis_blocks), 262 | title="[cyan]Diagnosis[/]", 263 | border_style="cyan", 264 | padding=(0, 2) 265 | )) 266 | 267 | # Recommended Fix 268 | if components['fix']: 269 | fix_command = components['fix'].group(1).strip('`') 270 | console.print(Panel( 271 | Syntax(fix_command, "bash", theme="ansi_light", line_numbers=False), 272 | title="[bold bright_green]⚡ RECOMMENDED FIX[/]", 273 | border_style="bright_green", 274 | padding=(1, 2), 275 | # subtitle="Copy-paste ready solution" 276 | )) 277 | 278 | # Additional Information 279 | info_blocks = [] 280 | if components['risk']: 281 | info_blocks.append(Markdown(f"**Potential Risks**\n{components['risk'].group(1)}")) 282 | if components['prevention']: 283 | info_blocks.append(Markdown(f"**Prevention Tip**\n{components['prevention'].group(1)}")) 284 | 285 | if info_blocks: 286 | console.print(Panel( 287 | Group(*info_blocks), 288 | title="[yellow]Additional Information[/]", 289 | border_style="yellow", 290 | padding=(0, 2) 291 | )) 292 | 293 | def _print_component(self, match, color, label): 294 | """Enhanced component display""" 295 | if match: 296 | cleaned = match.group(1).replace('\n', ' ').strip() 297 | print(f"{color}▸ {label}:\n {cleaned}\033[0m") 298 | 299 | def _prompt_fix(self, command, relevant_files): 300 | """Smart fix suggestion using context""" 301 | clean_cmd = re.sub(r'^\s*\[.*?\]\s*', '', command).strip() 302 | 303 | # If the command contains 'filename' or similar placeholder and we have relevant files 304 | if ('filename' in clean_cmd.lower() or 'file' in clean_cmd.lower()) and relevant_files: 305 | clean_cmd = clean_cmd.replace('filename', relevant_files[0]) 306 | clean_cmd = clean_cmd.replace('file', relevant_files[0]) 307 | 308 | print(f"\n\033[95m💡 Recommended fix command:\033[0m \033[92m{clean_cmd}\033[0m") 309 | 310 | def _get_native_error(self, command): 311 | """Get error output directly from command""" 312 | try: 313 | result = subprocess.run( 314 | command, 315 | shell=True, 316 | capture_output=True, 317 | text=True 318 | ) 319 | return result.stderr.strip() 320 | except Exception: 321 | return "Command execution failed" 322 | 323 | def _get_additional_context(self): 324 | """Enhanced context gathering for error analysis""" 325 | context = { 326 | 'env_vars': self._get_relevant_env_vars(), 327 | 'process_tree': self._get_process_tree(), 328 | 'file_context': self._get_file_context(), 329 | 'network_state': self._get_network_state() 330 | } 331 | 332 | context['command_history'] = self._enhance_command_history() 333 | 334 | context.update(self._get_specialized_context()) 335 | 336 | return context 337 | 338 | def _get_relevant_env_vars(self): 339 | return { 340 | 'PATH': os.getenv('PATH', ''), 341 | 'SHELL': os.getenv('SHELL', ''), 342 | 'USER': os.getenv('USER', ''), 343 | 'HOME': os.getenv('HOME', ''), 344 | 'PWD': os.getenv('PWD', ''), 345 | 'OLDPWD': os.getenv('OLDPWD', '') 346 | } 347 | 348 | def _get_process_tree(self): 349 | try: 350 | ps_output = subprocess.check_output( 351 | ['ps', '-ef', '--forest'], 352 | stderr=subprocess.DEVNULL, 353 | text=True 354 | ).strip() 355 | return ps_output.split('\n')[-10:] # Last 10 processes 356 | except Exception: 357 | return [] 358 | 359 | 360 | def _get_file_context(self): 361 | cwd = os.getcwd() 362 | context = { 363 | 'files': [f for f in os.listdir(cwd) if os.path.isfile(f)][:10], 364 | 'dirs': [d for d in os.listdir(cwd) if os.path.isdir(d)][:5] 365 | } 366 | 367 | 368 | cmd_parts = self.last_command.split() 369 | if not cmd_parts: 370 | return context 371 | 372 | 373 | potential_files = [p for p in cmd_parts if os.path.exists(p) and os.path.isfile(p)] 374 | 375 | 376 | file_contents = {} 377 | for f in potential_files[:2]: # Limit to 2 most relevant files 378 | try: 379 | with open(f, 'r') as file: 380 | content = "".join(file.readlines()[:20]) 381 | file_contents[f] = content 382 | except Exception: 383 | file_contents[f] = "Unable to read file content" 384 | 385 | context['file_contents'] = file_contents 386 | return context 387 | 388 | def _get_network_state(self): 389 | try: 390 | return subprocess.check_output( 391 | ['ss', '-tulpn'], 392 | stderr=subprocess.DEVNULL, 393 | text=True 394 | ).strip().split('\n')[:5] 395 | except Exception: 396 | return [] 397 | 398 | def _get_git_context(self): 399 | try: 400 | git_status = subprocess.run( 401 | 'git status --porcelain', 402 | shell=True, 403 | capture_output=True, 404 | text=True 405 | ) 406 | git_remotes = subprocess.run( 407 | 'git remote -v', 408 | shell=True, 409 | capture_output=True, 410 | text=True 411 | ).stdout 412 | return { 413 | 'git_status': git_status.stdout, 414 | 'git_remotes': git_remotes 415 | } 416 | except Exception: 417 | return {} 418 | 419 | 420 | def _enhance_command_history(self): 421 | """Track both commands and their outputs""" 422 | history_dict = {} 423 | for cmd in self.command_history: 424 | if cmd not in history_dict: 425 | try: 426 | result = subprocess.run( 427 | cmd, 428 | shell=True, 429 | capture_output=True, 430 | text=True, 431 | timeout=2 # Set timeout to avoid hanging 432 | ) 433 | output = result.stdout[:200] + ("..." if len(result.stdout) > 200 else "") 434 | history_dict[cmd] = output 435 | except Exception: 436 | history_dict[cmd] = "Error capturing output" 437 | 438 | return history_dict 439 | 440 | def _get_specialized_context(self): 441 | """Get command-specific context based on command type""" 442 | cmd_parts = self.last_command.split() 443 | if not cmd_parts: 444 | return {} 445 | 446 | base_cmd = cmd_parts[0] 447 | context = {} 448 | 449 | # Git command context 450 | if base_cmd == 'git': 451 | context.update(self._get_git_context()) 452 | 453 | # Docker command context 454 | elif base_cmd == 'docker' or base_cmd == 'docker-compose': 455 | context.update(self._get_docker_context()) 456 | 457 | # Package manager context 458 | elif base_cmd in ['apt', 'apt-get', 'yum', 'dnf', 'pacman', 'brew']: 459 | context.update(self._get_package_context(base_cmd)) 460 | 461 | # Server/service context 462 | elif base_cmd in ['systemctl', 'service', 'nginx', 'apache2']: 463 | context.update(self._get_service_context(base_cmd)) 464 | 465 | return context 466 | 467 | def _get_docker_context(self): 468 | """Get Docker-specific context""" 469 | try: 470 | containers = subprocess.run( 471 | 'docker ps --format "{{.Names}} ({{.Status}})"', 472 | shell=True, 473 | capture_output=True, 474 | text=True 475 | ).stdout.strip() 476 | 477 | compose_files = [] 478 | for file in ['docker-compose.yml', 'docker-compose.yaml']: 479 | if os.path.exists(file): 480 | compose_files.append(file) 481 | 482 | return { 483 | 'docker_containers': containers.split('\n') if containers else [], 484 | 'compose_files': compose_files 485 | } 486 | except Exception: 487 | return {} 488 | 489 | def _get_package_context(self, manager): 490 | """Get package manager context""" 491 | try: 492 | if manager in ['apt', 'apt-get']: 493 | updates = subprocess.run( 494 | 'apt list --upgradable 2>/dev/null | head -n 5', 495 | shell=True, 496 | capture_output=True, 497 | text=True 498 | ).stdout.strip() 499 | return {'available_updates': updates.split('\n') if updates else []} 500 | return {} 501 | except Exception: 502 | return {} 503 | 504 | def _get_service_context(self, service_manager): 505 | """Get service/systemd context""" 506 | try: 507 | if service_manager in ['systemctl', 'service']: 508 | # Get failed services 509 | failed = subprocess.run( 510 | 'systemctl list-units --state=failed --no-legend | head -n 3', 511 | shell=True, 512 | capture_output=True, 513 | text=True 514 | ).stdout.strip() 515 | 516 | return {'failed_services': failed.split('\n') if failed else []} 517 | return {} 518 | except Exception: 519 | return {} -------------------------------------------------------------------------------- /src/shellsage/helpers.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | def update_env_file(provider, key): 4 | """Update provider key in .env without duplicates""" 5 | env_path = Path('.env') 6 | 7 | if not env_path.exists(): 8 | env_path.touch() 9 | 10 | lines = env_path.read_text().splitlines() 11 | key_line = f"{provider.upper()}_API_KEY={key}" 12 | 13 | # Remove existing entries 14 | new_lines = [line for line in lines if not line.startswith(f"{provider.upper()}_API_KEY=")] 15 | 16 | # Add new key at the end 17 | new_lines.append(key_line) 18 | 19 | # Write back to file 20 | env_path.write_text("\n".join(new_lines)) 21 | 22 | def update_env_variable(variable, value): 23 | """Update any .env variable without duplicates""" 24 | env_path = Path('.env') 25 | 26 | if not env_path.exists(): 27 | env_path.touch() 28 | 29 | lines = env_path.read_text().splitlines() 30 | 31 | # Remove existing entries 32 | new_lines = [line for line in lines if not line.startswith(f"{variable}=")] 33 | 34 | # Add new value 35 | new_lines.append(f"{variable}={value}") 36 | 37 | # Write back to file 38 | env_path.write_text("\n".join(new_lines)) -------------------------------------------------------------------------------- /src/shellsage/llm_handler.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from .model_manager import ModelManager 4 | 5 | class DeepSeekLLMHandler: 6 | def __init__(self): 7 | self.manager = ModelManager() 8 | 9 | def get_error_solution(self, error_context): 10 | prompt = self._build_prompt(error_context) 11 | try: 12 | response = self.manager.generate(prompt, max_tokens=1024) 13 | return self._format_response(response) 14 | except Exception as e: 15 | return f"Error: {str(e)}" 16 | 17 | # Update _build_prompt in DeepSeekLLMHandler 18 | def _build_prompt(self, context): 19 | # Extract files mentioned in error if any 20 | error_files = [] 21 | if context.get('error_output'): 22 | # Simple regex to find file paths in error messages 23 | file_matches = re.findall(r'\'(.*?)\'|\"(.*?)\"|\b([\/\w\.-]+\.\w+)\b', context.get('error_output', '')) 24 | for match in file_matches: 25 | for group in match: 26 | if group and os.path.exists(group) and os.path.isfile(group): 27 | error_files.append(group) 28 | 29 | # Gather command-specific context details 30 | specialized_context = "" 31 | if context.get('git_status'): 32 | specialized_context += f"\n**Git Status**: {context['git_status'][:200]}" 33 | if context.get('docker_containers'): 34 | specialized_context += f"\n**Docker Containers**: {', '.join(context['docker_containers'][:3])}" 35 | if context.get('failed_services'): 36 | specialized_context += f"\n**Failed Services**: {', '.join(context['failed_services'])}" 37 | 38 | # File content context 39 | file_context = "" 40 | if context.get('file_context', {}).get('file_contents'): 41 | for file, content in context['file_context']['file_contents'].items(): 42 | if len(content) > 300: 43 | content = content[:300] + "..." 44 | file_context += f"\n**File {file}**: ```\n{content}\n```" 45 | 46 | # Build the enhanced prompt 47 | prompt = f"""**[Terminal Context Analysis]** 48 | **System Environment**: {context.get('env_vars', {}).get('SHELL', 'Unknown')} on {context.get('os', 'Linux')} 49 | **Working Directory**: {context['cwd']} ({len(context.get('file_context', {}).get('files', []))} files) 50 | **Recent Commands**: {', '.join(context.get('history', [])[-3:])} 51 | **Failed Command**: `{context['command']}` 52 | **Error Message**: {context['error_output']} 53 | **Exit Code**: {context['exit_code']} 54 | **Referenced Files**: {', '.join(error_files) if error_files else 'None detected'} 55 | **Man Page Excerpt**: {context.get('man_excerpt', 'N/A')} 56 | {specialized_context} 57 | {file_context} 58 | 59 | **Required Analysis Format:** 60 | 61 | Step 1: Identify the exact error message and command that failed 62 | Step 2: Analyze why the command failed (syntax, missing files, permissions, etc.) 63 | Step 3: Find the correct command or fix based on context 64 | Step 4: Consider any potential risks 65 | 66 | 67 | Root Cause: <1-line diagnosis> 68 | Fix: `[executable command]` 69 | Technical Explanation: 70 | Potential Risks: 71 | Prevention Tip: """ 72 | 73 | return prompt 74 | 75 | def _format_response(self, raw): 76 | # Detect reasoning model response 77 | is_reasoning_model = any(x in self.manager.local_model.lower() 78 | for x in ['deepseek', 'r1', 'think', 'expert']) 79 | 80 | if is_reasoning_model and '' in raw: 81 | # Extract all thinking blocks and final response 82 | thoughts = [] 83 | remaining = raw 84 | while '' in remaining and '' in remaining: 85 | think_start = remaining.find('') + len('') 86 | think_end = remaining.find('') 87 | if think_start > -1 and think_end > -1: 88 | thoughts.append(remaining[think_start:think_end].strip()) 89 | remaining = remaining[think_end + len(''):] 90 | raw = remaining.strip() 91 | 92 | # Existing cleaning logic 93 | cleaned = re.sub(r'\n+', '\n', raw) 94 | cleaned = re.sub(r'(\d\.\s|\*\*)', '', cleaned) 95 | 96 | return re.sub( 97 | r'(Root Cause|Fix|Technical Explanation|Potential Risks|Prevention Tip):?', 98 | lambda m: f"🔍 {m.group(1)}:" if m.group(1) == "Root Cause" else 99 | f"🛠️ {m.group(1)}:" if m.group(1) == "Fix" else 100 | f"📚 {m.group(1)}:" if m.group(1) == "Technical Explanation" else 101 | f"⚠️ {m.group(1)}:" if m.group(1) == "Potential Risks" else 102 | f"🔒 {m.group(1)}:", 103 | cleaned 104 | ) -------------------------------------------------------------------------------- /src/shellsage/model_manager.py: -------------------------------------------------------------------------------- 1 | from .helpers import update_env_variable 2 | import os 3 | import yaml 4 | import requests 5 | from pathlib import Path 6 | from openai import OpenAI 7 | import inquirer 8 | from anthropic import Anthropic 9 | from dotenv import load_dotenv 10 | 11 | 12 | # Define providers at module level 13 | PROVIDERS = { 14 | 'groq': { 15 | 'client': OpenAI, 16 | 'base_url': 'https://api.groq.com/openai/v1', 17 | 'models': ['llama-3.1-8b-instant', 'deepseek-r1-distill-llama-70b', 'gemma2-9b-it', 'llama-3.3-70b-versatile', 'llama3-70b-8192', 'llama3-8b-8192', 'mixtral-8x7b-32768'] 18 | }, 19 | 'openai': { 20 | 'client': OpenAI, 21 | 'base_url': 'https://api.openai.com/v1', 22 | 'models': ['gpt-4o', 'chatgpt-4o-latest', 'o1', 'o1-mini', 'o1-preview', 'gpt-4o-2024-08-06', 'gpt-4o-mini-2024-07-18', 'gpt-4-turbo', 'gpt-3.5-turbo'] 23 | }, 24 | 'anthropic': { 25 | 'client': Anthropic, 26 | 'models': ['claude-3-5-sonnet-20241022', 'claude-3-opus-20240229', 'claude-3-sonnet-20240229'] 27 | }, 28 | 'fireworks': { 29 | 'client': OpenAI, 30 | 'base_url': 'https://api.fireworks.ai/inference/v1', 31 | 'models': ['accounts/fireworks/models/llama-v3p1-405b-instruct', 'accounts/fireworks/models/deepseek-v3', 'accounts/fireworks/models/llama-v3p1-8b-instruct', 'accounts/fireworks/models/llama-v3p3-70b-instruct'] 32 | }, 33 | 'openrouter': { 34 | 'client': OpenAI, 35 | 'base_url': 'https://openrouter.ai/api/v1', 36 | 'models': ['deepseek/deepseek-r1-distill-llama-70b:free', 'deepseek/deepseek-r1-distill-qwen-32b', 'mistralai/mistral-small-24b-instruct-2501', 'openai/gpt-3.5-turbo-instruct', 'microsoft/phi-4', 'google/gemini-2.0-flash-thinking-exp:free', 'google/gemini-2.0-pro-exp-02-05:free', 'deepseek/deepseek-r1:free', 'qwen/qwen-vl-plus:free'] 37 | }, 38 | 'deepseek': { 39 | 'client': OpenAI, 40 | 'base_url': 'https://api.deepseek.com/v1', 41 | 'models': ['deepseek-chat'] 42 | } 43 | } 44 | 45 | class ModelManager: 46 | PROVIDERS = PROVIDERS # Add this line to expose the module-level PROVIDERS 47 | 48 | def __init__(self): 49 | load_dotenv(override=True) 50 | self.mode = os.getenv('MODE', 'local') 51 | self.local_model = os.getenv('LOCAL_MODEL', 'llama3:8b-instruct-q4_1') 52 | self.client = None 53 | self._init_client() 54 | 55 | def _init_client(self): 56 | """Initialize active client based on config""" 57 | if self.mode == 'api': 58 | provider = os.getenv('ACTIVE_API_PROVIDER', 'groq') 59 | api_key = os.environ.get(f"{provider.upper()}_API_KEY") 60 | 61 | if not api_key: 62 | raise ValueError(f"API key for {provider} not set. Run 'shellsage setup'") 63 | 64 | if self.PROVIDERS[provider]['client'] == OpenAI: 65 | self.client = OpenAI( 66 | api_key=api_key, 67 | base_url=self.PROVIDERS[provider].get('base_url') 68 | ) 69 | # Special case for Anthropic 70 | elif provider == 'anthropic': 71 | self.client = Anthropic(api_key=api_key) 72 | else: 73 | raise ValueError(f"Unsupported provider: {provider}") 74 | else: 75 | # Initialize local client if needed 76 | self.client = "ollama" # Just a flag for local mode 77 | 78 | def switch_mode(self, new_mode, model_name=None): 79 | """Change mode with optional model selection""" 80 | update_env_variable('MODE', new_mode) 81 | 82 | if new_mode == 'local' and model_name: 83 | update_env_variable('LOCAL_MODEL', model_name) 84 | elif new_mode == 'api' and model_name: 85 | provider = next(p for p in self.PROVIDERS if model_name in self.PROVIDERS[p]['models']) 86 | update_env_variable('ACTIVE_API_PROVIDER', provider) 87 | update_env_variable('API_MODEL', model_name) 88 | 89 | load_dotenv(override=True) 90 | self._init_client() 91 | 92 | def get_ollama_models(self): 93 | """List installed Ollama models""" 94 | try: 95 | ollama_host = os.getenv('OLLAMA_HOST', 'http://localhost:11434') 96 | response = requests.get(f"{ollama_host}/api/tags") 97 | return [m['name'] for m in response.json().get('models', [])] 98 | except requests.ConnectionError: 99 | return [] 100 | 101 | def interactive_setup(self): 102 | """Guide user through configuration""" 103 | questions = [ 104 | inquirer.List('mode', 105 | message="Select operation mode:", 106 | choices=['local', 'api'], 107 | default=self.mode 108 | ), 109 | inquirer.List('local_model', 110 | message="Select local model:", 111 | choices=self.get_ollama_models(), 112 | default=self.local_model 113 | ), 114 | inquirer.Text('api_key', 115 | message="Enter Groq API key:", 116 | default=os.getenv(f"GROQ_API_KEY", ''), 117 | ignore=lambda x: x['mode'] != 'api' 118 | ) 119 | ] 120 | 121 | answers = inquirer.prompt(questions) 122 | self._update_config(answers) 123 | self._init_client() 124 | 125 | def _update_config(self, answers): 126 | """Update configuration from answers""" 127 | self.mode = answers['mode'] 128 | self.local_model = answers['local_model'] 129 | os.environ["ACTIVE_API_PROVIDER"] = "groq" if self.mode == 'api' else "" # Fixed provider assignment 130 | os.environ["GROQ_API_KEY"] = answers['api_key'] 131 | load_dotenv(override=True) 132 | 133 | def list_local_models(self): 134 | """Get all available local models""" 135 | models = [] 136 | if self.mode == 'local': 137 | models = self.get_ollama_models() 138 | return models 139 | 140 | def generate(self, prompt, max_tokens=512): 141 | """Unified generation interface""" 142 | try: 143 | if self.mode == 'api': 144 | return self._api_generate(prompt, max_tokens) 145 | return self._local_generate(prompt) 146 | except Exception as e: 147 | raise RuntimeError(f"Generation failed: {str(e)}") 148 | 149 | def _api_generate(self, prompt, max_tokens): 150 | """Generate using selected API provider""" 151 | provider = os.getenv('ACTIVE_API_PROVIDER', 'groq') 152 | model = os.getenv('API_MODEL') # New environment variable 153 | 154 | try: 155 | if self.PROVIDERS[provider]['client'] == OpenAI: 156 | response = self.client.chat.completions.create( 157 | model=model, 158 | messages=[{"role": "user", "content": prompt}], 159 | temperature=0.1, 160 | max_tokens=max_tokens 161 | ) 162 | return response.choices[0].message.content 163 | elif provider == 'anthropic': 164 | response = self.client.messages.create( 165 | model=model, 166 | max_tokens=max_tokens, 167 | messages=[{"role": "user", "content": prompt}] 168 | ) 169 | return response.content[0].text 170 | except Exception as e: 171 | raise RuntimeError(f"API Error ({provider}): {str(e)}") 172 | 173 | def _local_generate(self, prompt): 174 | """Generate using local provider""" 175 | if self.mode == 'local': 176 | return self._ollama_generate(prompt) 177 | return self._hf_generate(prompt) 178 | 179 | # model_manager.py 180 | 181 | def _ollama_generate(self, prompt): 182 | try: 183 | ollama_host = os.getenv('OLLAMA_HOST', 'http://localhost:11434') 184 | # Detect if it's a reasoning model based on model name 185 | is_reasoning_model = any(x in self.local_model.lower() for x in ['deepseek', 'r1', 'think', 'expert']) 186 | 187 | options = { 188 | "temperature": 0.1, 189 | "num_predict": 200048 190 | } 191 | 192 | # Only set stop tokens for non-reasoning models 193 | if not is_reasoning_model: 194 | options["stop"] = ["\n\n\n", "USER QUERY:"] 195 | 196 | response = requests.post( 197 | f"{ollama_host}/api/generate", 198 | json={ 199 | "model": self.local_model, 200 | "prompt": prompt, 201 | "stream": False, 202 | "options": options 203 | } 204 | ) 205 | response.raise_for_status() 206 | return response.json()['response'] 207 | except Exception as e: 208 | raise RuntimeError(f"Ollama error: {str(e)}") 209 | 210 | def _hf_generate(self, prompt): 211 | """Generate using HuggingFace model""" 212 | from ctransformers import AutoModelForCausalLM 213 | 214 | try: 215 | model = AutoModelForCausalLM.from_pretrained( 216 | model_path=self.local_model, 217 | model_type='llama' 218 | ) 219 | return model(prompt) 220 | except Exception as e: 221 | raise RuntimeError(f"HuggingFace error: {str(e)}") --------------------------------------------------------------------------------