diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..68d085b --- /dev/null +++ b/.flake8 @@ -0,0 +1,19 @@ +[flake8] +# Professional Python code style - balances quality with readability +max-line-length = 95 +extend-ignore = E203,W503,W605 +exclude = + .venv, + .venv-linting, + __pycache__, + *.egg-info, + .git, + build, + dist, + .mini-rag + +# Per-file ignores for practical development +per-file-ignores = + tests/*.py:F401,F841 + examples/*.py:F401,F841 + fix_*.py:F401,F841,E501 \ No newline at end of file diff --git a/.mini-rag/last_search b/.mini-rag/last_search index 651bb6f..30d74d2 100644 --- a/.mini-rag/last_search +++ b/.mini-rag/last_search @@ -1 +1 @@ -how to run tests \ No newline at end of file +test \ No newline at end of file diff --git a/.venv-linting/bin/Activate.ps1 b/.venv-linting/bin/Activate.ps1 new file mode 100644 index 0000000..b49d77b --- /dev/null +++ b/.venv-linting/bin/Activate.ps1 @@ -0,0 +1,247 @@ +<# +.Synopsis +Activate a Python virtual environment for the current PowerShell session. + +.Description +Pushes the python executable for a virtual environment to the front of the +$Env:PATH environment variable and sets the prompt to signify that you are +in a Python virtual environment. Makes use of the command line switches as +well as the `pyvenv.cfg` file values present in the virtual environment. + +.Parameter VenvDir +Path to the directory that contains the virtual environment to activate. The +default value for this is the parent of the directory that the Activate.ps1 +script is located within. + +.Parameter Prompt +The prompt prefix to display when this virtual environment is activated. By +default, this prompt is the name of the virtual environment folder (VenvDir) +surrounded by parentheses and followed by a single space (ie. '(.venv) '). + +.Example +Activate.ps1 +Activates the Python virtual environment that contains the Activate.ps1 script. + +.Example +Activate.ps1 -Verbose +Activates the Python virtual environment that contains the Activate.ps1 script, +and shows extra information about the activation as it executes. + +.Example +Activate.ps1 -VenvDir C:\Users\MyUser\Common\.venv +Activates the Python virtual environment located in the specified location. + +.Example +Activate.ps1 -Prompt "MyPython" +Activates the Python virtual environment that contains the Activate.ps1 script, +and prefixes the current prompt with the specified string (surrounded in +parentheses) while the virtual environment is active. + +.Notes +On Windows, it may be required to enable this Activate.ps1 script by setting the +execution policy for the user. You can do this by issuing the following PowerShell +command: + +PS C:\> Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser + +For more information on Execution Policies: +https://go.microsoft.com/fwlink/?LinkID=135170 + +#> +Param( + [Parameter(Mandatory = $false)] + [String] + $VenvDir, + [Parameter(Mandatory = $false)] + [String] + $Prompt +) + +<# Function declarations --------------------------------------------------- #> + +<# +.Synopsis +Remove all shell session elements added by the Activate script, including the +addition of the virtual environment's Python executable from the beginning of +the PATH variable. + +.Parameter NonDestructive +If present, do not remove this function from the global namespace for the +session. + +#> +function global:deactivate ([switch]$NonDestructive) { + # Revert to original values + + # The prior prompt: + if (Test-Path -Path Function:_OLD_VIRTUAL_PROMPT) { + Copy-Item -Path Function:_OLD_VIRTUAL_PROMPT -Destination Function:prompt + Remove-Item -Path Function:_OLD_VIRTUAL_PROMPT + } + + # The prior PYTHONHOME: + if (Test-Path -Path Env:_OLD_VIRTUAL_PYTHONHOME) { + Copy-Item -Path Env:_OLD_VIRTUAL_PYTHONHOME -Destination Env:PYTHONHOME + Remove-Item -Path Env:_OLD_VIRTUAL_PYTHONHOME + } + + # The prior PATH: + if (Test-Path -Path Env:_OLD_VIRTUAL_PATH) { + Copy-Item -Path Env:_OLD_VIRTUAL_PATH -Destination Env:PATH + Remove-Item -Path Env:_OLD_VIRTUAL_PATH + } + + # Just remove the VIRTUAL_ENV altogether: + if (Test-Path -Path Env:VIRTUAL_ENV) { + Remove-Item -Path env:VIRTUAL_ENV + } + + # Just remove VIRTUAL_ENV_PROMPT altogether. + if (Test-Path -Path Env:VIRTUAL_ENV_PROMPT) { + Remove-Item -Path env:VIRTUAL_ENV_PROMPT + } + + # Just remove the _PYTHON_VENV_PROMPT_PREFIX altogether: + if (Get-Variable -Name "_PYTHON_VENV_PROMPT_PREFIX" -ErrorAction SilentlyContinue) { + Remove-Variable -Name _PYTHON_VENV_PROMPT_PREFIX -Scope Global -Force + } + + # Leave deactivate function in the global namespace if requested: + if (-not $NonDestructive) { + Remove-Item -Path function:deactivate + } +} + +<# +.Description +Get-PyVenvConfig parses the values from the pyvenv.cfg file located in the +given folder, and returns them in a map. + +For each line in the pyvenv.cfg file, if that line can be parsed into exactly +two strings separated by `=` (with any amount of whitespace surrounding the =) +then it is considered a `key = value` line. The left hand string is the key, +the right hand is the value. + +If the value starts with a `'` or a `"` then the first and last character is +stripped from the value before being captured. + +.Parameter ConfigDir +Path to the directory that contains the `pyvenv.cfg` file. +#> +function Get-PyVenvConfig( + [String] + $ConfigDir +) { + Write-Verbose "Given ConfigDir=$ConfigDir, obtain values in pyvenv.cfg" + + # Ensure the file exists, and issue a warning if it doesn't (but still allow the function to continue). + $pyvenvConfigPath = Join-Path -Resolve -Path $ConfigDir -ChildPath 'pyvenv.cfg' -ErrorAction Continue + + # An empty map will be returned if no config file is found. + $pyvenvConfig = @{ } + + if ($pyvenvConfigPath) { + + Write-Verbose "File exists, parse `key = value` lines" + $pyvenvConfigContent = Get-Content -Path $pyvenvConfigPath + + $pyvenvConfigContent | ForEach-Object { + $keyval = $PSItem -split "\s*=\s*", 2 + if ($keyval[0] -and $keyval[1]) { + $val = $keyval[1] + + # Remove extraneous quotations around a string value. + if ("'""".Contains($val.Substring(0, 1))) { + $val = $val.Substring(1, $val.Length - 2) + } + + $pyvenvConfig[$keyval[0]] = $val + Write-Verbose "Adding Key: '$($keyval[0])'='$val'" + } + } + } + return $pyvenvConfig +} + + +<# Begin Activate script --------------------------------------------------- #> + +# Determine the containing directory of this script +$VenvExecPath = Split-Path -Parent $MyInvocation.MyCommand.Definition +$VenvExecDir = Get-Item -Path $VenvExecPath + +Write-Verbose "Activation script is located in path: '$VenvExecPath'" +Write-Verbose "VenvExecDir Fullname: '$($VenvExecDir.FullName)" +Write-Verbose "VenvExecDir Name: '$($VenvExecDir.Name)" + +# Set values required in priority: CmdLine, ConfigFile, Default +# First, get the location of the virtual environment, it might not be +# VenvExecDir if specified on the command line. +if ($VenvDir) { + Write-Verbose "VenvDir given as parameter, using '$VenvDir' to determine values" +} +else { + Write-Verbose "VenvDir not given as a parameter, using parent directory name as VenvDir." + $VenvDir = $VenvExecDir.Parent.FullName.TrimEnd("\\/") + Write-Verbose "VenvDir=$VenvDir" +} + +# Next, read the `pyvenv.cfg` file to determine any required value such +# as `prompt`. +$pyvenvCfg = Get-PyVenvConfig -ConfigDir $VenvDir + +# Next, set the prompt from the command line, or the config file, or +# just use the name of the virtual environment folder. +if ($Prompt) { + Write-Verbose "Prompt specified as argument, using '$Prompt'" +} +else { + Write-Verbose "Prompt not specified as argument to script, checking pyvenv.cfg value" + if ($pyvenvCfg -and $pyvenvCfg['prompt']) { + Write-Verbose " Setting based on value in pyvenv.cfg='$($pyvenvCfg['prompt'])'" + $Prompt = $pyvenvCfg['prompt']; + } + else { + Write-Verbose " Setting prompt based on parent's directory's name. (Is the directory name passed to venv module when creating the virtual environment)" + Write-Verbose " Got leaf-name of $VenvDir='$(Split-Path -Path $venvDir -Leaf)'" + $Prompt = Split-Path -Path $venvDir -Leaf + } +} + +Write-Verbose "Prompt = '$Prompt'" +Write-Verbose "VenvDir='$VenvDir'" + +# Deactivate any currently active virtual environment, but leave the +# deactivate function in place. +deactivate -nondestructive + +# Now set the environment variable VIRTUAL_ENV, used by many tools to determine +# that there is an activated venv. +$env:VIRTUAL_ENV = $VenvDir + +if (-not $Env:VIRTUAL_ENV_DISABLE_PROMPT) { + + Write-Verbose "Setting prompt to '$Prompt'" + + # Set the prompt to include the env name + # Make sure _OLD_VIRTUAL_PROMPT is global + function global:_OLD_VIRTUAL_PROMPT { "" } + Copy-Item -Path function:prompt -Destination function:_OLD_VIRTUAL_PROMPT + New-Variable -Name _PYTHON_VENV_PROMPT_PREFIX -Description "Python virtual environment prompt prefix" -Scope Global -Option ReadOnly -Visibility Public -Value $Prompt + + function global:prompt { + Write-Host -NoNewline -ForegroundColor Green "($_PYTHON_VENV_PROMPT_PREFIX) " + _OLD_VIRTUAL_PROMPT + } + $env:VIRTUAL_ENV_PROMPT = $Prompt +} + +# Clear PYTHONHOME +if (Test-Path -Path Env:PYTHONHOME) { + Copy-Item -Path Env:PYTHONHOME -Destination Env:_OLD_VIRTUAL_PYTHONHOME + Remove-Item -Path Env:PYTHONHOME +} + +# Add the venv to the PATH +Copy-Item -Path Env:PATH -Destination Env:_OLD_VIRTUAL_PATH +$Env:PATH = "$VenvExecDir$([System.IO.Path]::PathSeparator)$Env:PATH" diff --git a/.venv-linting/bin/activate b/.venv-linting/bin/activate new file mode 100644 index 0000000..7139d6d --- /dev/null +++ b/.venv-linting/bin/activate @@ -0,0 +1,70 @@ +# This file must be used with "source bin/activate" *from bash* +# You cannot run it directly + +deactivate () { + # reset old environment variables + if [ -n "${_OLD_VIRTUAL_PATH:-}" ] ; then + PATH="${_OLD_VIRTUAL_PATH:-}" + export PATH + unset _OLD_VIRTUAL_PATH + fi + if [ -n "${_OLD_VIRTUAL_PYTHONHOME:-}" ] ; then + PYTHONHOME="${_OLD_VIRTUAL_PYTHONHOME:-}" + export PYTHONHOME + unset _OLD_VIRTUAL_PYTHONHOME + fi + + # Call hash to forget past commands. Without forgetting + # past commands the $PATH changes we made may not be respected + hash -r 2> /dev/null + + if [ -n "${_OLD_VIRTUAL_PS1:-}" ] ; then + PS1="${_OLD_VIRTUAL_PS1:-}" + export PS1 + unset _OLD_VIRTUAL_PS1 + fi + + unset VIRTUAL_ENV + unset VIRTUAL_ENV_PROMPT + if [ ! "${1:-}" = "nondestructive" ] ; then + # Self destruct! + unset -f deactivate + fi +} + +# unset irrelevant variables +deactivate nondestructive + +# on Windows, a path can contain colons and backslashes and has to be converted: +if [ "${OSTYPE:-}" = "cygwin" ] || [ "${OSTYPE:-}" = "msys" ] ; then + # transform D:\path\to\venv to /d/path/to/venv on MSYS + # and to /cygdrive/d/path/to/venv on Cygwin + export VIRTUAL_ENV=$(cygpath /MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting) +else + # use the path as-is + export VIRTUAL_ENV=/MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting +fi + +_OLD_VIRTUAL_PATH="$PATH" +PATH="$VIRTUAL_ENV/"bin":$PATH" +export PATH + +# unset PYTHONHOME if set +# this will fail if PYTHONHOME is set to the empty string (which is bad anyway) +# could use `if (set -u; : $PYTHONHOME) ;` in bash +if [ -n "${PYTHONHOME:-}" ] ; then + _OLD_VIRTUAL_PYTHONHOME="${PYTHONHOME:-}" + unset PYTHONHOME +fi + +if [ -z "${VIRTUAL_ENV_DISABLE_PROMPT:-}" ] ; then + _OLD_VIRTUAL_PS1="${PS1:-}" + PS1='(.venv-linting) '"${PS1:-}" + export PS1 + VIRTUAL_ENV_PROMPT='(.venv-linting) ' + export VIRTUAL_ENV_PROMPT +fi + +# Call hash to forget past commands. Without forgetting +# past commands the $PATH changes we made may not be respected +hash -r 2> /dev/null diff --git a/.venv-linting/bin/activate.csh b/.venv-linting/bin/activate.csh new file mode 100644 index 0000000..ff5de9a --- /dev/null +++ b/.venv-linting/bin/activate.csh @@ -0,0 +1,27 @@ +# This file must be used with "source bin/activate.csh" *from csh*. +# You cannot run it directly. + +# Created by Davide Di Blasi . +# Ported to Python 3.3 venv by Andrew Svetlov + +alias deactivate 'test $?_OLD_VIRTUAL_PATH != 0 && setenv PATH "$_OLD_VIRTUAL_PATH" && unset _OLD_VIRTUAL_PATH; rehash; test $?_OLD_VIRTUAL_PROMPT != 0 && set prompt="$_OLD_VIRTUAL_PROMPT" && unset _OLD_VIRTUAL_PROMPT; unsetenv VIRTUAL_ENV; unsetenv VIRTUAL_ENV_PROMPT; test "\!:*" != "nondestructive" && unalias deactivate' + +# Unset irrelevant variables. +deactivate nondestructive + +setenv VIRTUAL_ENV /MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting + +set _OLD_VIRTUAL_PATH="$PATH" +setenv PATH "$VIRTUAL_ENV/"bin":$PATH" + + +set _OLD_VIRTUAL_PROMPT="$prompt" + +if (! "$?VIRTUAL_ENV_DISABLE_PROMPT") then + set prompt = '(.venv-linting) '"$prompt" + setenv VIRTUAL_ENV_PROMPT '(.venv-linting) ' +endif + +alias pydoc python -m pydoc + +rehash diff --git a/.venv-linting/bin/activate.fish b/.venv-linting/bin/activate.fish new file mode 100644 index 0000000..9faac7e --- /dev/null +++ b/.venv-linting/bin/activate.fish @@ -0,0 +1,69 @@ +# This file must be used with "source /bin/activate.fish" *from fish* +# (https://fishshell.com/). You cannot run it directly. + +function deactivate -d "Exit virtual environment and return to normal shell environment" + # reset old environment variables + if test -n "$_OLD_VIRTUAL_PATH" + set -gx PATH $_OLD_VIRTUAL_PATH + set -e _OLD_VIRTUAL_PATH + end + if test -n "$_OLD_VIRTUAL_PYTHONHOME" + set -gx PYTHONHOME $_OLD_VIRTUAL_PYTHONHOME + set -e _OLD_VIRTUAL_PYTHONHOME + end + + if test -n "$_OLD_FISH_PROMPT_OVERRIDE" + set -e _OLD_FISH_PROMPT_OVERRIDE + # prevents error when using nested fish instances (Issue #93858) + if functions -q _old_fish_prompt + functions -e fish_prompt + functions -c _old_fish_prompt fish_prompt + functions -e _old_fish_prompt + end + end + + set -e VIRTUAL_ENV + set -e VIRTUAL_ENV_PROMPT + if test "$argv[1]" != "nondestructive" + # Self-destruct! + functions -e deactivate + end +end + +# Unset irrelevant variables. +deactivate nondestructive + +set -gx VIRTUAL_ENV /MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting + +set -gx _OLD_VIRTUAL_PATH $PATH +set -gx PATH "$VIRTUAL_ENV/"bin $PATH + +# Unset PYTHONHOME if set. +if set -q PYTHONHOME + set -gx _OLD_VIRTUAL_PYTHONHOME $PYTHONHOME + set -e PYTHONHOME +end + +if test -z "$VIRTUAL_ENV_DISABLE_PROMPT" + # fish uses a function instead of an env var to generate the prompt. + + # Save the current fish_prompt function as the function _old_fish_prompt. + functions -c fish_prompt _old_fish_prompt + + # With the original prompt function renamed, we can override with our own. + function fish_prompt + # Save the return status of the last command. + set -l old_status $status + + # Output the venv prompt; color taken from the blue of the Python logo. + printf "%s%s%s" (set_color 4B8BBE) '(.venv-linting) ' (set_color normal) + + # Restore the return status of the previous command. + echo "exit $old_status" | . + # Output the original/"old" prompt. + _old_fish_prompt + end + + set -gx _OLD_FISH_PROMPT_OVERRIDE "$VIRTUAL_ENV" + set -gx VIRTUAL_ENV_PROMPT '(.venv-linting) ' +end diff --git a/.venv-linting/bin/black b/.venv-linting/bin/black new file mode 100755 index 0000000..3aa8f3f --- /dev/null +++ b/.venv-linting/bin/black @@ -0,0 +1,8 @@ +#!/MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from black import patched_main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(patched_main()) diff --git a/.venv-linting/bin/blackd b/.venv-linting/bin/blackd new file mode 100755 index 0000000..274f3f7 --- /dev/null +++ b/.venv-linting/bin/blackd @@ -0,0 +1,8 @@ +#!/MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from blackd import patched_main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(patched_main()) diff --git a/.venv-linting/bin/isort b/.venv-linting/bin/isort new file mode 100755 index 0000000..36b1b18 --- /dev/null +++ b/.venv-linting/bin/isort @@ -0,0 +1,8 @@ +#!/MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from isort.main import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/.venv-linting/bin/isort-identify-imports b/.venv-linting/bin/isort-identify-imports new file mode 100755 index 0000000..9d7c9a8 --- /dev/null +++ b/.venv-linting/bin/isort-identify-imports @@ -0,0 +1,8 @@ +#!/MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from isort.main import identify_imports_main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(identify_imports_main()) diff --git a/.venv-linting/bin/pip b/.venv-linting/bin/pip new file mode 100755 index 0000000..6252495 --- /dev/null +++ b/.venv-linting/bin/pip @@ -0,0 +1,8 @@ +#!/MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from pip._internal.cli.main import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/.venv-linting/bin/pip3 b/.venv-linting/bin/pip3 new file mode 100755 index 0000000..6252495 --- /dev/null +++ b/.venv-linting/bin/pip3 @@ -0,0 +1,8 @@ +#!/MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from pip._internal.cli.main import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/.venv-linting/bin/pip3.12 b/.venv-linting/bin/pip3.12 new file mode 100755 index 0000000..6252495 --- /dev/null +++ b/.venv-linting/bin/pip3.12 @@ -0,0 +1,8 @@ +#!/MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from pip._internal.cli.main import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/.venv-linting/bin/python b/.venv-linting/bin/python new file mode 120000 index 0000000..b8a0adb --- /dev/null +++ b/.venv-linting/bin/python @@ -0,0 +1 @@ +python3 \ No newline at end of file diff --git a/.venv-linting/bin/python3 b/.venv-linting/bin/python3 new file mode 120000 index 0000000..ae65fda --- /dev/null +++ b/.venv-linting/bin/python3 @@ -0,0 +1 @@ +/usr/bin/python3 \ No newline at end of file diff --git a/.venv-linting/bin/python3.12 b/.venv-linting/bin/python3.12 new file mode 120000 index 0000000..b8a0adb --- /dev/null +++ b/.venv-linting/bin/python3.12 @@ -0,0 +1 @@ +python3 \ No newline at end of file diff --git a/.venv-linting/lib64 b/.venv-linting/lib64 new file mode 120000 index 0000000..7951405 --- /dev/null +++ b/.venv-linting/lib64 @@ -0,0 +1 @@ +lib \ No newline at end of file diff --git a/.venv-linting/pyvenv.cfg b/.venv-linting/pyvenv.cfg new file mode 100644 index 0000000..752b9be --- /dev/null +++ b/.venv-linting/pyvenv.cfg @@ -0,0 +1,5 @@ +home = /usr/bin +include-system-site-packages = false +version = 3.12.3 +executable = /usr/bin/python3.12 +command = /usr/bin/python3 -m venv /MASTERFOLDER/Coding/Fss-Mini-Rag/.venv-linting diff --git a/GET_STARTED.md b/GET_STARTED.md deleted file mode 100644 index fae0802..0000000 --- a/GET_STARTED.md +++ /dev/null @@ -1,83 +0,0 @@ -# ๐Ÿš€ FSS-Mini-RAG: Get Started in 2 Minutes - -## Step 1: Install Everything -```bash -./install_mini_rag.sh -``` -**That's it!** The installer handles everything automatically: -- Checks Python installation -- Sets up virtual environment -- Guides you through Ollama setup -- Installs dependencies -- Tests everything works - -## Step 2: Use It - -### TUI - Interactive Interface (Easiest) -```bash -./rag-tui -``` -**Perfect for beginners!** Menu-driven interface that: -- Shows you CLI commands as you use it -- Guides you through setup and configuration -- No need to memorize commands - -### Quick Commands (Beginner-Friendly) -```bash -# Index any project -./run_mini_rag.sh index ~/my-project - -# Search your code -./run_mini_rag.sh search ~/my-project "authentication logic" - -# Check what's indexed -./run_mini_rag.sh status ~/my-project -``` - -### Full Commands (More Options) -```bash -# Basic indexing and search -./rag-mini index /path/to/project -./rag-mini search /path/to/project "database connection" - -# Enhanced search with smart features -./rag-mini-enhanced search /path/to/project "UserManager" -./rag-mini-enhanced similar /path/to/project "def validate_input" -``` - -## What You Get - -**Semantic Search**: Instead of exact text matching, finds code by meaning: -- Search "user login" โ†’ finds authentication functions, session management, password validation -- Search "database queries" โ†’ finds SQL, ORM code, connection handling -- Search "error handling" โ†’ finds try/catch blocks, error classes, logging - -## Installation Options - -The installer offers two choices: - -**Light Installation (Recommended)**: -- Uses Ollama for high-quality embeddings -- Requires Ollama installed (installer guides you) -- Small download (~50MB) - -**Full Installation**: -- Includes ML fallback models -- Works without Ollama -- Large download (~2-3GB) - -## Troubleshooting - -**"Python not found"**: Install Python 3.8+ from python.org -**"Ollama not found"**: Visit https://ollama.ai/download -**"Import errors"**: Re-run `./install_mini_rag.sh` - -## Next Steps - -- **Technical Details**: Read `README.md` -- **Step-by-Step Guide**: Read `docs/GETTING_STARTED.md` -- **Examples**: Check `examples/` directory -- **Test It**: Run on this project: `./run_mini_rag.sh index .` - ---- -**Questions?** Everything is documented in the README.md file. \ No newline at end of file diff --git a/README.md b/README.md index 3281435..0d6b89b 100644 --- a/README.md +++ b/README.md @@ -79,34 +79,24 @@ FSS-Mini-RAG offers **two distinct experiences** optimized for different use cas ## Quick Start (2 Minutes) -**Linux/macOS:** +**Step 1: Install** ```bash -# 1. Install everything +# Linux/macOS ./install_mini_rag.sh -# 2. Choose your interface -./rag-tui # Friendly interface for beginners -# OR choose your mode: -./rag-mini index ~/my-project # Index your project first -./rag-mini search ~/my-project "query" --synthesize # Fast synthesis -./rag-mini explore ~/my-project # Interactive exploration +# Windows +install_windows.bat ``` -**Windows:** -```cmd -# 1. Install everything -install_windows.bat +**Step 2: Start Using** +```bash +# Beginners: Interactive interface +./rag-tui # Linux/macOS +rag.bat # Windows -# 2. Choose your interface -rag.bat # Interactive interface -# OR choose your mode: -rag.bat index C:\my-project # Index your project first -rag.bat search C:\my-project "query" # Fast search -rag.bat explore C:\my-project # Interactive exploration - -# Direct Python entrypoint (after install): -rag-mini index C:\my-project -rag-mini search C:\my-project "query" +# Experienced users: Direct commands +./rag-mini index ~/project # Index your project +./rag-mini search ~/project "your query" ``` That's it. No external dependencies, no configuration required, no PhD in computer science needed. @@ -232,7 +222,7 @@ This implementation prioritizes: ## Next Steps -- **New users**: Run `./rag-mini` (Linux/macOS) or `rag.bat` (Windows) for guided experience +- **New users**: Run `./rag-tui` (Linux/macOS) or `rag.bat` (Windows) for guided experience - **Developers**: Read [`TECHNICAL_GUIDE.md`](docs/TECHNICAL_GUIDE.md) for implementation details - **Contributors**: See [`CONTRIBUTING.md`](CONTRIBUTING.md) for development setup diff --git a/rag-mini.py b/bin/rag-mini.py similarity index 79% rename from rag-mini.py rename to bin/rag-mini.py index b9acec8..3aa309b 100644 --- a/rag-mini.py +++ b/bin/rag-mini.py @@ -6,24 +6,32 @@ A lightweight, portable RAG system for semantic code search. Usage: rag-mini [options] """ -import sys import argparse -from pathlib import Path import json import logging +import socket +import sys +from pathlib import Path + +# Add parent directory to path so we can import mini_rag +sys.path.insert(0, str(Path(__file__).parent.parent)) + +import requests # Add the RAG system to the path sys.path.insert(0, str(Path(__file__).parent)) try: - from mini_rag.indexer import ProjectIndexer - from mini_rag.search import CodeSearcher - from mini_rag.ollama_embeddings import OllamaEmbedder - from mini_rag.llm_synthesizer import LLMSynthesizer from mini_rag.explorer import CodeExplorer + from mini_rag.indexer import ProjectIndexer + from mini_rag.llm_synthesizer import LLMSynthesizer + from mini_rag.ollama_embeddings import OllamaEmbedder + from mini_rag.search import CodeSearcher + # Update system (graceful import) try: from mini_rag.updater import check_for_updates, get_updater + UPDATER_AVAILABLE = True except ImportError: UPDATER_AVAILABLE = False @@ -48,50 +56,51 @@ except ImportError as e: # Configure logging for user-friendly output logging.basicConfig( level=logging.WARNING, # Only show warnings and errors by default - format='%(levelname)s: %(message)s' + format="%(levelname)s: %(message)s", ) logger = logging.getLogger(__name__) + def index_project(project_path: Path, force: bool = False): """Index a project directory.""" try: # Show what's happening action = "Re-indexing" if force else "Indexing" print(f"๐Ÿš€ {action} {project_path.name}") - + # Quick pre-check - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if rag_dir.exists() and not force: print(" Checking for changes...") - + indexer = ProjectIndexer(project_path) result = indexer.index_project(force_reindex=force) - + # Show results with context - files_count = result.get('files_indexed', 0) - chunks_count = result.get('chunks_created', 0) - time_taken = result.get('time_taken', 0) - + files_count = result.get("files_indexed", 0) + chunks_count = result.get("chunks_created", 0) + time_taken = result.get("time_taken", 0) + if files_count == 0: print("โœ… Index up to date - no changes detected") else: print(f"โœ… Indexed {files_count} files in {time_taken:.1f}s") print(f" Created {chunks_count} chunks") - + # Show efficiency if time_taken > 0: speed = files_count / time_taken print(f" Speed: {speed:.1f} files/sec") - + # Show warnings if any - failed_count = result.get('files_failed', 0) + failed_count = result.get("files_failed", 0) if failed_count > 0: print(f"โš ๏ธ {failed_count} files failed (check logs with --verbose)") - + # Quick tip for first-time users - if not (project_path / '.mini-rag' / 'last_search').exists(): - print(f"\n๐Ÿ’ก Try: rag-mini search {project_path} \"your search here\"") - + if not (project_path / ".mini-rag" / "last_search").exists(): + print(f'\n๐Ÿ’ก Try: rag-mini search {project_path} "your search here"') + except FileNotFoundError: print(f"๐Ÿ“ Directory Not Found: {project_path}") print(" Make sure the path exists and you're in the right location") @@ -110,7 +119,7 @@ def index_project(project_path: Path, force: bool = False): # Connection errors are handled in the embedding module if "ollama" in str(e).lower() or "connection" in str(e).lower(): sys.exit(1) # Error already displayed - + print(f"โŒ Indexing failed: {e}") print() print("๐Ÿ”ง Common solutions:") @@ -124,39 +133,44 @@ def index_project(project_path: Path, force: bool = False): print(" Or see: docs/TROUBLESHOOTING.md") sys.exit(1) + def search_project(project_path: Path, query: str, top_k: int = 10, synthesize: bool = False): """Search a project directory.""" try: # Check if indexed first - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if not rag_dir.exists(): print(f"โŒ Project not indexed: {project_path.name}") print(f" Run: rag-mini index {project_path}") sys.exit(1) - - print(f"๐Ÿ” Searching \"{query}\" in {project_path.name}") + + print(f'๐Ÿ” Searching "{query}" in {project_path.name}') searcher = CodeSearcher(project_path) results = searcher.search(query, top_k=top_k) - + if not results: print("โŒ No results found") print() print("๐Ÿ”ง Quick fixes to try:") - print(" โ€ข Use broader terms: \"login\" instead of \"authenticate_user_session\"") - print(" โ€ข Try concepts: \"database query\" instead of specific function names") + print(' โ€ข Use broader terms: "login" instead of "authenticate_user_session"') + print(' โ€ข Try concepts: "database query" instead of specific function names') print(" โ€ข Check spelling and try simpler words") - print(" โ€ข Search for file types: \"python class\" or \"javascript function\"") + print(' โ€ข Search for file types: "python class" or "javascript function"') print() print("โš™๏ธ Configuration adjustments:") - print(f" โ€ข Lower threshold: ./rag-mini search \"{project_path}\" \"{query}\" --threshold 0.05") - print(f" โ€ข More results: ./rag-mini search \"{project_path}\" \"{query}\" --top-k 20") + print( + f' โ€ข Lower threshold: ./rag-mini search "{project_path}" "{query}" --threshold 0.05' + ) + print( + f' โ€ข More results: ./rag-mini search "{project_path}" "{query}" --top-k 20' + ) print() print("๐Ÿ“š Need help? See: docs/TROUBLESHOOTING.md") return - + print(f"โœ… Found {len(results)} results:") print() - + for i, result in enumerate(results, 1): # Clean up file path display file_path = Path(result.file_path) @@ -165,70 +179,89 @@ def search_project(project_path: Path, query: str, top_k: int = 10, synthesize: except ValueError: # If relative_to fails, just show the basename rel_path = file_path.name - + print(f"{i}. {rel_path}") print(f" Score: {result.score:.3f}") - + # Show line info if available - if hasattr(result, 'start_line') and result.start_line: + if hasattr(result, "start_line") and result.start_line: print(f" Lines: {result.start_line}-{result.end_line}") - - # Show content preview - if hasattr(result, 'name') and result.name: + + # Show content preview + if hasattr(result, "name") and result.name: print(f" Context: {result.name}") - + # Show full content with proper formatting - print(f" Content:") - content_lines = result.content.strip().split('\n') + print(" Content:") + content_lines = result.content.strip().split("\n") for line in content_lines[:10]: # Show up to 10 lines print(f" {line}") - + if len(content_lines) > 10: print(f" ... ({len(content_lines) - 10} more lines)") - print(f" Use --verbose or rag-mini-enhanced for full context") - + print(" Use --verbose or rag-mini-enhanced for full context") + print() - + # LLM Synthesis if requested if synthesize: print("๐Ÿง  Generating LLM synthesis...") - + # Load config to respect user's model preferences from mini_rag.config import ConfigManager + config_manager = ConfigManager(project_path) config = config_manager.load_config() - + synthesizer = LLMSynthesizer( - model=config.llm.synthesis_model if config.llm.synthesis_model != "auto" else None, - config=config + model=( + config.llm.synthesis_model + if config.llm.synthesis_model != "auto" + else None + ), + config=config, ) - + if synthesizer.is_available(): synthesis = synthesizer.synthesize_search_results(query, results, project_path) print() print(synthesizer.format_synthesis_output(synthesis, query)) - + # Add guidance for deeper analysis - if synthesis.confidence < 0.7 or any(word in query.lower() for word in ['why', 'how', 'explain', 'debug']): + if synthesis.confidence < 0.7 or any( + word in query.lower() for word in ["why", "how", "explain", "debug"] + ): print("\n๐Ÿ’ก Want deeper analysis with reasoning?") print(f" Try: rag-mini explore {project_path}") - print(" Exploration mode enables thinking and remembers conversation context.") + print( + " Exploration mode enables thinking and remembers conversation context." + ) else: print("โŒ LLM synthesis unavailable") print(" โ€ข Ensure Ollama is running: ollama serve") print(" โ€ข Install a model: ollama pull qwen3:1.7b") print(" โ€ข Check connection to http://localhost:11434") - + # Save last search for potential enhancements try: - (rag_dir / 'last_search').write_text(query) - except: + (rag_dir / "last_search").write_text(query) + except ( + ConnectionError, + FileNotFoundError, + IOError, + OSError, + TimeoutError, + TypeError, + ValueError, + requests.RequestException, + socket.error, + ): pass # Don't fail if we can't save - + except Exception as e: print(f"โŒ Search failed: {e}") print() - + if "not indexed" in str(e).lower(): print("๐Ÿ”ง Solution:") print(f" ./rag-mini index {project_path}") @@ -241,44 +274,45 @@ def search_project(project_path: Path, query: str, top_k: int = 10, synthesize: print(" โ€ข Check available memory and disk space") print() print("๐Ÿ“š Get detailed error info:") - print(f" ./rag-mini search {project_path} \"{query}\" --verbose") + print(f' ./rag-mini search {project_path} "{query}" --verbose') print(" Or see: docs/TROUBLESHOOTING.md") print() sys.exit(1) + def status_check(project_path: Path): """Show status of RAG system.""" try: print(f"๐Ÿ“Š Status for {project_path.name}") print() - + # Check project indexing status first - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if not rag_dir.exists(): print("โŒ Project not indexed") print(f" Run: rag-mini index {project_path}") print() else: - manifest = rag_dir / 'manifest.json' + manifest = rag_dir / "manifest.json" if manifest.exists(): try: with open(manifest) as f: data = json.load(f) - - file_count = data.get('file_count', 0) - chunk_count = data.get('chunk_count', 0) - indexed_at = data.get('indexed_at', 'Never') - + + file_count = data.get("file_count", 0) + chunk_count = data.get("chunk_count", 0) + indexed_at = data.get("indexed_at", "Never") + print("โœ… Project indexed") print(f" Files: {file_count}") print(f" Chunks: {chunk_count}") print(f" Last update: {indexed_at}") - + # Show average chunks per file if file_count > 0: avg_chunks = chunk_count / file_count print(f" Avg chunks/file: {avg_chunks:.1f}") - + print() except Exception: print("โš ๏ธ Index exists but manifest unreadable") @@ -287,88 +321,94 @@ def status_check(project_path: Path): print("โš ๏ธ Index directory exists but incomplete") print(f" Try: rag-mini index {project_path} --force") print() - + # Check embedding system status print("๐Ÿง  Embedding System:") try: embedder = OllamaEmbedder() emb_info = embedder.get_status() - method = emb_info.get('method', 'unknown') - - if method == 'ollama': + method = emb_info.get("method", "unknown") + + if method == "ollama": print(" โœ… Ollama (high quality)") - elif method == 'ml': + elif method == "ml": print(" โœ… ML fallback (good quality)") - elif method == 'hash': + elif method == "hash": print(" โš ๏ธ Hash fallback (basic quality)") else: print(f" โ“ Unknown method: {method}") - + # Show additional details if available - if 'model' in emb_info: + if "model" in emb_info: print(f" Model: {emb_info['model']}") - + except Exception as e: print(f" โŒ Status check failed: {e}") - + print() - + # Check LLM status and show actual vs configured model print("๐Ÿค– LLM System:") try: from mini_rag.config import ConfigManager + config_manager = ConfigManager(project_path) config = config_manager.load_config() - + synthesizer = LLMSynthesizer( - model=config.llm.synthesis_model if config.llm.synthesis_model != "auto" else None, - config=config + model=( + config.llm.synthesis_model + if config.llm.synthesis_model != "auto" + else None + ), + config=config, ) - + if synthesizer.is_available(): synthesizer._ensure_initialized() actual_model = synthesizer.model config_model = config.llm.synthesis_model - + if config_model == "auto": print(f" โœ… Auto-selected: {actual_model}") elif config_model == actual_model: print(f" โœ… Using configured: {actual_model}") else: - print(f" โš ๏ธ Model mismatch!") + print(" โš ๏ธ Model mismatch!") print(f" Configured: {config_model}") print(f" Actually using: {actual_model}") - print(f" (Configured model may not be installed)") - + print(" (Configured model may not be installed)") + print(f" Config file: {config_manager.config_path}") else: print(" โŒ Ollama not available") print(" Start with: ollama serve") - + except Exception as e: print(f" โŒ LLM status check failed: {e}") - + # Show last search if available - last_search_file = rag_dir / 'last_search' if rag_dir.exists() else None + last_search_file = rag_dir / "last_search" if rag_dir.exists() else None if last_search_file and last_search_file.exists(): try: last_query = last_search_file.read_text().strip() - print(f"\n๐Ÿ” Last search: \"{last_query}\"") - except: + print(f'\n๐Ÿ” Last search: "{last_query}"') + except (FileNotFoundError, IOError, OSError, TypeError, ValueError): pass - + except Exception as e: print(f"โŒ Status check failed: {e}") sys.exit(1) + def explore_interactive(project_path: Path): """Interactive exploration mode with thinking and context memory for any documents.""" try: explorer = CodeExplorer(project_path) - + if not explorer.start_exploration_session(): sys.exit(1) - + # Show enhanced first-time guidance print(f"\n๐Ÿค” Ask your first question about {project_path.name}:") print() @@ -377,12 +417,12 @@ def explore_interactive(project_path: Path): print() print("๐Ÿ”ง Quick options:") print(" 1. Help - Show example questions") - print(" 2. Status - Project information") + print(" 2. Status - Project information") print(" 3. Suggest - Get a random starter question") print() - + is_first_question = True - + while True: try: # Get user input with clearer prompt @@ -390,12 +430,12 @@ def explore_interactive(project_path: Path): question = input("๐Ÿ“ Enter question or option (1-3): ").strip() else: question = input("\n> ").strip() - + # Handle exit commands - if question.lower() in ['quit', 'exit', 'q']: + if question.lower() in ["quit", "exit", "q"]: print("\n" + explorer.end_session()) break - + # Handle empty input if not question: if is_first_question: @@ -403,17 +443,18 @@ def explore_interactive(project_path: Path): else: print("Please enter a question or 'quit' to exit.") continue - + # Handle numbered options and special commands - if question in ['1'] or question.lower() in ['help', 'h']: - print(""" + if question in ["1"] or question.lower() in ["help", "h"]: + print( + """ ๐Ÿง  EXPLORATION MODE HELP: โ€ข Ask any question about your documents or code โ€ข I remember our conversation for follow-up questions โ€ข Use 'why', 'how', 'explain' for detailed reasoning โ€ข Type 'summary' to see session overview โ€ข Type 'quit' or 'exit' to end session - + ๐Ÿ’ก Example questions: โ€ข "How does authentication work?" โ€ข "What are the main components?" @@ -421,36 +462,40 @@ def explore_interactive(project_path: Path): โ€ข "Why is this function slow?" โ€ข "What security measures are in place?" โ€ข "How does data flow through this system?" -""") +""" + ) continue - - elif question in ['2'] or question.lower() == 'status': - print(f""" + + elif question in ["2"] or question.lower() == "status": + print( + """ ๐Ÿ“Š PROJECT STATUS: {project_path.name} โ€ข Location: {project_path} โ€ข Exploration session active โ€ข AI model ready for questions โ€ข Conversation memory enabled -""") +""" + ) continue - - elif question in ['3'] or question.lower() == 'suggest': + + elif question in ["3"] or question.lower() == "suggest": # Random starter questions for first-time users if is_first_question: import random + starters = [ "What are the main components of this project?", - "How is error handling implemented?", + "How is error handling implemented?", "Show me the authentication and security logic", "What are the key functions I should understand first?", "How does data flow through this system?", "What configuration options are available?", - "Show me the most important files to understand" + "Show me the most important files to understand", ] suggested = random.choice(starters) print(f"\n๐Ÿ’ก Suggested question: {suggested}") print(" Press Enter to use this, or type your own question:") - + next_input = input("๐Ÿ“ > ").strip() if not next_input: # User pressed Enter to use suggestion question = suggested @@ -463,24 +508,24 @@ def explore_interactive(project_path: Path): print(' "What are the security implications?"') print(' "Show me related code examples"') continue - - if question.lower() == 'summary': + + if question.lower() == "summary": print("\n" + explorer.get_session_summary()) continue - + # Process the question print(f"\n๐Ÿ” Searching {project_path.name}...") print("๐Ÿง  Thinking with AI model...") response = explorer.explore_question(question) - + # Mark as no longer first question after processing is_first_question = False - + if response: print(f"\n{response}") else: print("โŒ Sorry, I couldn't process that question. Please try again.") - + except KeyboardInterrupt: print(f"\n\n{explorer.end_session()}") break @@ -490,88 +535,94 @@ def explore_interactive(project_path: Path): except Exception as e: print(f"โŒ Error processing question: {e}") print("Please try again or type 'quit' to exit.") - + except Exception as e: print(f"โŒ Failed to start exploration mode: {e}") print("Make sure the project is indexed first: rag-mini index ") sys.exit(1) + def show_discrete_update_notice(): """Show a discrete, non-intrusive update notice for CLI users.""" if not UPDATER_AVAILABLE: return - + try: update_info = check_for_updates() if update_info: # Very discrete notice - just one line - print(f"๐Ÿ”„ (Update v{update_info.version} available - run 'rag-mini check-update' to learn more)") + print( + f"๐Ÿ”„ (Update v{update_info.version} available - run 'rag-mini check-update' to learn more)" + ) except Exception: # Silently ignore any update check failures pass + def handle_check_update(): """Handle the check-update command.""" if not UPDATER_AVAILABLE: print("โŒ Update system not available") print("๐Ÿ’ก Try updating to the latest version manually from GitHub") return - + try: print("๐Ÿ” Checking for updates...") update_info = check_for_updates() - + if update_info: print(f"\n๐ŸŽ‰ Update Available: v{update_info.version}") print("=" * 50) print("\n๐Ÿ“‹ What's New:") - notes_lines = update_info.release_notes.split('\n')[:10] # First 10 lines + notes_lines = update_info.release_notes.split("\n")[:10] # First 10 lines for line in notes_lines: if line.strip(): print(f" {line.strip()}") - + print(f"\n๐Ÿ”— Release Page: {update_info.release_url}") - print(f"\n๐Ÿš€ To install: rag-mini update") + print("\n๐Ÿš€ To install: rag-mini update") print("๐Ÿ’ก Or update manually from GitHub releases") else: print("โœ… You're already on the latest version!") - + except Exception as e: print(f"โŒ Failed to check for updates: {e}") print("๐Ÿ’ก Try updating manually from GitHub") + def handle_update(): """Handle the update command.""" if not UPDATER_AVAILABLE: print("โŒ Update system not available") print("๐Ÿ’ก Try updating manually from GitHub") return - + try: print("๐Ÿ” Checking for updates...") update_info = check_for_updates() - + if not update_info: print("โœ… You're already on the latest version!") return - + print(f"\n๐ŸŽ‰ Update Available: v{update_info.version}") print("=" * 50) - + # Show brief release notes - notes_lines = update_info.release_notes.split('\n')[:5] + notes_lines = update_info.release_notes.split("\n")[:5] for line in notes_lines: if line.strip(): print(f" โ€ข {line.strip()}") - + # Confirm update confirm = input(f"\n๐Ÿš€ Install v{update_info.version}? [Y/n]: ").strip().lower() - if confirm in ['', 'y', 'yes']: + if confirm in ["", "y", "yes"]: updater = get_updater() - + print(f"\n๐Ÿ“ฅ Downloading v{update_info.version}...") - + # Progress callback + def show_progress(downloaded, total): if total > 0: percent = (downloaded / total) * 100 @@ -579,17 +630,17 @@ def handle_update(): filled = int(bar_length * downloaded / total) bar = "โ–ˆ" * filled + "โ–‘" * (bar_length - filled) print(f"\r [{bar}] {percent:.1f}%", end="", flush=True) - + # Download and install update_package = updater.download_update(update_info, show_progress) if not update_package: print("\nโŒ Download failed. Please try again later.") return - + print("\n๐Ÿ’พ Creating backup...") if not updater.create_backup(): print("โš ๏ธ Backup failed, but continuing anyway...") - + print("๐Ÿ”„ Installing update...") if updater.apply_update(update_package, update_info): print("โœ… Update successful!") @@ -604,91 +655,108 @@ def handle_update(): print("โŒ Rollback failed. You may need to reinstall.") else: print("Update cancelled.") - + except Exception as e: print(f"โŒ Update failed: {e}") print("๐Ÿ’ก Try updating manually from GitHub") + def main(): """Main CLI interface.""" # Check virtual environment try: from mini_rag.venv_checker import check_and_warn_venv + check_and_warn_venv("rag-mini.py", force_exit=False) except ImportError: pass # If venv checker can't be imported, continue anyway - + parser = argparse.ArgumentParser( description="FSS-Mini-RAG - Lightweight semantic code search", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: rag-mini index /path/to/project # Index a project - rag-mini search /path/to/project "query" # Search indexed project + rag-mini search /path/to/project "query" # Search indexed project rag-mini search /path/to/project "query" -s # Search with LLM synthesis rag-mini explore /path/to/project # Interactive exploration mode rag-mini status /path/to/project # Show status - """ + """, ) - - parser.add_argument('command', choices=['index', 'search', 'explore', 'status', 'update', 'check-update'], - help='Command to execute') - parser.add_argument('project_path', type=Path, nargs='?', - help='Path to project directory (REQUIRED except for update commands)') - parser.add_argument('query', nargs='?', - help='Search query (for search command)') - parser.add_argument('--force', action='store_true', - help='Force reindex all files') - parser.add_argument('--top-k', '--limit', type=int, default=10, dest='top_k', - help='Maximum number of search results (top-k)') - parser.add_argument('--verbose', '-v', action='store_true', - help='Enable verbose logging') - parser.add_argument('--synthesize', '-s', action='store_true', - help='Generate LLM synthesis of search results (requires Ollama)') - + + parser.add_argument( + "command", + choices=["index", "search", "explore", "status", "update", "check-update"], + help="Command to execute", + ) + parser.add_argument( + "project_path", + type=Path, + nargs="?", + help="Path to project directory (REQUIRED except for update commands)", + ) + parser.add_argument("query", nargs="?", help="Search query (for search command)") + parser.add_argument("--force", action="store_true", help="Force reindex all files") + parser.add_argument( + "--top-k", + "--limit", + type=int, + default=10, + dest="top_k", + help="Maximum number of search results (top-k)", + ) + parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose logging") + parser.add_argument( + "--synthesize", + "-s", + action="store_true", + help="Generate LLM synthesis of search results (requires Ollama)", + ) + args = parser.parse_args() - + # Set logging level if args.verbose: logging.getLogger().setLevel(logging.INFO) - + # Handle update commands first (don't require project_path) - if args.command == 'check-update': + if args.command == "check-update": handle_check_update() return - elif args.command == 'update': + elif args.command == "update": handle_update() return - + # All other commands require project_path if not args.project_path: print("โŒ Project path required for this command") sys.exit(1) - + # Validate project path if not args.project_path.exists(): print(f"โŒ Project path does not exist: {args.project_path}") sys.exit(1) - + if not args.project_path.is_dir(): print(f"โŒ Project path is not a directory: {args.project_path}") sys.exit(1) - + # Show discrete update notification for regular commands (non-intrusive) show_discrete_update_notice() - + # Execute command - if args.command == 'index': + if args.command == "index": index_project(args.project_path, args.force) - elif args.command == 'search': + elif args.command == "search": if not args.query: print("โŒ Search query required") sys.exit(1) search_project(args.project_path, args.query, args.top_k, args.synthesize) - elif args.command == 'explore': + elif args.command == "explore": explore_interactive(args.project_path) - elif args.command == 'status': + elif args.command == "status": status_check(args.project_path) -if __name__ == '__main__': - main() \ No newline at end of file + +if __name__ == "__main__": + main() diff --git a/rag-tui.py b/bin/rag-tui.py similarity index 79% rename from rag-tui.py rename to bin/rag-tui.py index f1ef851..20174b6 100755 --- a/rag-tui.py +++ b/bin/rag-tui.py @@ -4,91 +4,107 @@ FSS-Mini-RAG Text User Interface Simple, educational TUI that shows CLI commands while providing easy interaction. """ -import os -import sys import json +import logging +import os +import subprocess +import sys from pathlib import Path -from typing import Optional, List, Dict, Any +from typing import Any, Dict, List, Optional + +# Add parent directory to path so we can import mini_rag +sys.path.insert(0, str(Path(__file__).parent.parent)) + +logger = logging.getLogger(__name__) # Update system (graceful import) try: - from mini_rag.updater import check_for_updates, get_updater, get_legacy_notification + from mini_rag.updater import check_for_updates, get_legacy_notification, get_updater + UPDATER_AVAILABLE = True except ImportError: UPDATER_AVAILABLE = False + # Simple TUI without external dependencies + + class SimpleTUI: + def __init__(self): self.project_path: Optional[Path] = None self.current_config: Dict[str, Any] = {} self.search_count = 0 # Track searches for sample reminder - self.config_dir = Path.home() / '.mini-rag-tui' - self.config_file = self.config_dir / 'last_project.json' - + self.config_dir = Path.home() / ".mini-rag-tui" + self.config_file = self.config_dir / "last_project.json" + # Load last project on startup self._load_last_project() - + def _load_last_project(self): """Load the last used project from config file, or auto-detect current directory.""" # First check if current directory has .mini-rag folder (auto-detect) current_dir = Path.cwd() - if (current_dir / '.mini-rag').exists(): + if (current_dir / ".mini-rag").exists(): self.project_path = current_dir # Save this as the last project too self._save_last_project() return - + # If no auto-detection, try loading from config file try: - if hasattr(self, 'config_file') and self.config_file.exists(): - with open(self.config_file, 'r') as f: + if hasattr(self, "config_file") and self.config_file.exists(): + with open(self.config_file, "r") as f: data = json.load(f) - project_path = Path(data.get('last_project', '')) + project_path = Path(data.get("last_project", "")) if project_path.exists() and project_path.is_dir(): self.project_path = project_path except Exception: # If loading fails, just continue without last project pass - + def _save_last_project(self): """Save current project as last used.""" if not self.project_path: return try: self.config_dir.mkdir(exist_ok=True) - data = {'last_project': str(self.project_path)} - with open(self.config_file, 'w') as f: + data = {"last_project": str(self.project_path)} + with open(self.config_file, "w") as f: json.dump(data, f) except Exception: # If saving fails, just continue pass - + def _get_llm_status(self): """Get LLM status for display in main menu.""" try: # Import here to avoid startup delays sys.path.insert(0, str(Path(__file__).parent)) + from mini_rag.config import ConfigManager, RAGConfig from mini_rag.llm_synthesizer import LLMSynthesizer - from mini_rag.config import RAGConfig, ConfigManager - + # Load config for model rankings if self.project_path: config_manager = ConfigManager(self.project_path) config = config_manager.load_config() else: config = RAGConfig() - + # Create synthesizer with proper model configuration (same as main CLI) synthesizer = LLMSynthesizer( - model=config.llm.synthesis_model if config.llm.synthesis_model != "auto" else None, - config=config + model=( + config.llm.synthesis_model + if config.llm.synthesis_model != "auto" + else None + ), + config=config, ) if synthesizer.is_available(): # Get the model that would be selected synthesizer._ensure_initialized() model = synthesizer.model - + # Show what config says vs what's actually being used config_model = config.llm.synthesis_model if config_model != "auto" and config_model != model: @@ -100,11 +116,11 @@ class SimpleTUI: return "โŒ Ollama not running", None except Exception as e: return f"โŒ Error: {str(e)[:20]}...", None - + def clear_screen(self): """Clear the terminal screen.""" - os.system('cls' if os.name == 'nt' else 'clear') - + os.system("cls" if os.name == "nt" else "clear") + def print_header(self): """Print the main header.""" print("+====================================================+") @@ -112,46 +128,52 @@ class SimpleTUI: print("| Semantic Code Search Interface |") print("+====================================================+") print() - + def print_cli_command(self, command: str, description: str = ""): """Show the equivalent CLI command.""" print(f"๐Ÿ’ป CLI equivalent: {command}") if description: print(f" {description}") print() - + def get_input(self, prompt: str, default: str = "") -> str: """Get user input with optional default.""" if default: full_prompt = f"{prompt} [{default}]: " else: full_prompt = f"{prompt}: " - + try: result = input(full_prompt).strip() return result if result else default except (KeyboardInterrupt, EOFError): print("\nGoodbye!") sys.exit(0) - - def show_menu(self, title: str, options: List[str], show_cli: bool = True, back_option: str = None) -> int: + + def show_menu( + self, + title: str, + options: List[str], + show_cli: bool = True, + back_option: str = None, + ) -> int: """Show a menu and get user selection.""" print(f"๐ŸŽฏ {title}") print("=" * (len(title) + 3)) print() - + for i, option in enumerate(options, 1): print(f"{i}. {option}") - + # Add back/exit option if back_option: print(f"0. {back_option}") - + if show_cli: print() print("๐Ÿ’ก All these actions can be done via CLI commands") print(" You'll see the commands as you use this interface!") - + print() while True: try: @@ -161,68 +183,75 @@ class SimpleTUI: elif 1 <= choice <= len(options): return choice - 1 else: - valid_range = "0-" + str(len(options)) if back_option else "1-" + str(len(options)) + valid_range = ( + "0-" + str(len(options)) if back_option else "1-" + str(len(options)) + ) print(f"Please enter a number between {valid_range}") except ValueError: print("Please enter a valid number") except (KeyboardInterrupt, EOFError): print("\nGoodbye!") sys.exit(0) - + def select_project(self): """Select or create project directory.""" self.clear_screen() self.print_header() - + print("๐Ÿ“ Project Selection") print("==================") print() - + # Show current project if any if self.project_path: print(f"Current project: {self.project_path}") print() - + print("๐Ÿ’ก New to FSS-Mini-RAG? Select 'Use current directory' to") print(" explore this RAG system's own codebase as your first demo!") print() - + # If we already have a project, show it prominently and offer quick actions if self.project_path: - rag_dir = self.project_path / '.mini-rag' + rag_dir = self.project_path / ".mini-rag" is_indexed = rag_dir.exists() status_text = "Ready for search โœ…" if is_indexed else "Needs indexing โŒ" - + print(f"Current: {self.project_path.name} ({status_text})") print() - + options = [ "Keep current project (go back to main menu)", "Use current directory (this folder)", "Enter different project path", "Browse recent projects", - "Open folder picker (GUI)" + "Open folder picker (GUI)", ] else: options = [ "Use current directory (perfect for beginners - try the RAG codebase!)", - "Enter project path (if you have a specific project)", + "Enter project path (if you have a specific project)", "Browse recent projects", - "Open folder picker (GUI)" + "Open folder picker (GUI)", ] - - choice = self.show_menu("Choose project directory", options, show_cli=False, back_option="Back to main menu") - + + choice = self.show_menu( + "Choose project directory", + options, + show_cli=False, + back_option="Back to main menu", + ) + if choice == -1: # Back to main menu return - + # Handle different choice patterns based on whether we have a project if self.project_path: if choice == 0: # Keep current project - just go back return elif choice == 1: - # Use current directory + # Use current directory self.project_path = Path.cwd() print(f"โœ… Using current directory: {self.project_path}") self._save_last_project() @@ -256,20 +285,22 @@ class SimpleTUI: self.project_path = Path(picked) print(f"โœ… Selected: {self.project_path}") self._save_last_project() - + input("\nPress Enter to continue...") - + def _enter_project_path(self): """Helper method to handle manual project path entry.""" while True: - path_str = self.get_input("Enter project directory path", - str(self.project_path) if self.project_path else "") - + path_str = self.get_input( + "Enter project directory path", + str(self.project_path) if self.project_path else "", + ) + if not path_str: continue - + project_path = Path(path_str).expanduser().resolve() - + if project_path.exists() and project_path.is_dir(): self.project_path = project_path print(f"โœ… Selected: {self.project_path}") @@ -278,64 +309,72 @@ class SimpleTUI: else: print(f"โŒ Directory not found: {project_path}") retry = input("Try again? (y/N): ").lower() - if retry != 'y': + if retry != "y": break - + def browse_recent_projects(self): """Browse recently indexed projects.""" print("๐Ÿ•’ Recent Projects") print("=================") print() - + # Look for .mini-rag directories in common locations search_paths = [ Path.home(), - Path.home() / "projects", + Path.home() / "projects", Path.home() / "code", Path.home() / "dev", Path.cwd().parent, - Path.cwd() + Path.cwd(), ] - + recent_projects = [] for search_path in search_paths: if search_path.exists() and search_path.is_dir(): try: for item in search_path.iterdir(): if item.is_dir(): - rag_dir = item / '.mini-rag' + rag_dir = item / ".mini-rag" if rag_dir.exists(): recent_projects.append(item) except (PermissionError, OSError): continue - + # Remove duplicates and sort by modification time recent_projects = list(set(recent_projects)) try: - recent_projects.sort(key=lambda p: (p / '.mini-rag').stat().st_mtime, reverse=True) - except: + recent_projects.sort(key=lambda p: (p / ".mini-rag").stat().st_mtime, reverse=True) + except (AttributeError, OSError, TypeError, ValueError): pass - + if not recent_projects: print("โŒ No recently indexed projects found") print(" Projects with .mini-rag directories will appear here") return - + print("Found indexed projects:") for i, project in enumerate(recent_projects[:10], 1): # Show up to 10 try: - manifest = project / '.mini-rag' / 'manifest.json' + manifest = project / ".mini-rag" / "manifest.json" if manifest.exists(): with open(manifest) as f: data = json.load(f) - file_count = data.get('file_count', 0) - indexed_at = data.get('indexed_at', 'Unknown') + file_count = data.get("file_count", 0) + indexed_at = data.get("indexed_at", "Unknown") print(f"{i}. {project.name} ({file_count} files, {indexed_at})") else: print(f"{i}. {project.name} (incomplete index)") - except: + except ( + AttributeError, + FileNotFoundError, + IOError, + OSError, + TypeError, + ValueError, + json.JSONDecodeError, + ): print(f"{i}. {project.name} (index status unknown)") - + print() try: choice = int(input("Select project number (or 0 to cancel): ")) @@ -345,87 +384,95 @@ class SimpleTUI: self._save_last_project() except (ValueError, IndexError): print("Selection cancelled") - + def index_project_interactive(self): """Interactive project indexing.""" if not self.project_path: print("โŒ No project selected") input("Press Enter to continue...") return - + self.clear_screen() self.print_header() - + print("๐Ÿš€ Project Indexing") print("==================") print() print(f"Project: {self.project_path}") print() - + # Check if already indexed - rag_dir = self.project_path / '.mini-rag' + rag_dir = self.project_path / ".mini-rag" if rag_dir.exists(): force = self._show_existing_index_info(rag_dir) else: force = False - + # Show CLI command cli_cmd = f"./rag-mini index {self.project_path}" if force: cli_cmd += " --force" - + self.print_cli_command(cli_cmd, "Index project for semantic search") - + # Import here to avoid startup delays sys.path.insert(0, str(Path(__file__).parent)) from mini_rag.indexer import ProjectIndexer - + # Get file count and show preview before starting print("๐Ÿ” Analyzing project structure...") print("=" * 50) - + try: indexer = ProjectIndexer(self.project_path) - + # Get files that would be indexed files_to_index = indexer._get_files_to_index() total_files = len(files_to_index) - + if total_files == 0: print("โœ… All files are already up to date!") print(" No indexing needed.") input("\nPress Enter to continue...") return - + # Show file analysis - print(f"๐Ÿ“Š Indexing Analysis:") + print("๐Ÿ“Š Indexing Analysis:") print(f" Files to process: {total_files}") - + # Analyze file types file_types = {} total_size = 0 for file_path in files_to_index: - ext = file_path.suffix.lower() or 'no extension' + ext = file_path.suffix.lower() or "no extension" file_types[ext] = file_types.get(ext, 0) + 1 try: total_size += file_path.stat().st_size - except: + except ( + AttributeError, + FileNotFoundError, + IOError, + OSError, + TypeError, + ValueError, + subprocess.SubprocessError, + ): pass - + # Show breakdown print(f" Total size: {total_size / (1024*1024):.1f}MB") - print(f" File types:") + print(" File types:") for ext, count in sorted(file_types.items(), key=lambda x: x[1], reverse=True): print(f" โ€ข {ext}: {count} files") - + # Conservative time estimate for average hardware estimated_time = self._estimate_processing_time(total_files, total_size) print(f" Estimated time: {estimated_time}") - + print() print("๐Ÿ’ก What indexing does:") print(" โ€ข Reads and analyzes each file's content (READ-ONLY)") - print(" โ€ข Breaks content into semantic chunks") + print(" โ€ข Breaks content into semantic chunks") print(" โ€ข Generates embeddings for semantic search") print(" โ€ข Stores everything in a separate .mini-rag/ database") print() @@ -435,103 +482,106 @@ class SimpleTUI: print(" โ€ข All data stored separately in .mini-rag/ folder") print(" โ€ข You can delete the .mini-rag/ folder anytime to remove all traces") print() - + # Confirmation confirm = input("๐Ÿš€ Proceed with indexing? [Y/n]: ").strip().lower() - if confirm and confirm != 'y' and confirm != 'yes': + if confirm and confirm != "y" and confirm != "yes": print("Indexing cancelled.") input("Press Enter to continue...") return - + print("\n๐Ÿš€ Starting indexing...") print("=" * 50) - + # Actually run the indexing result = indexer.index_project(force_reindex=force) - + print() print("๐ŸŽ‰ INDEXING COMPLETE!") print("=" * 50) - + # Comprehensive performance summary - files_processed = result.get('files_indexed', 0) - chunks_created = result.get('chunks_created', 0) - time_taken = result.get('time_taken', 0) - files_failed = result.get('files_failed', 0) - files_per_second = result.get('files_per_second', 0) - - print(f"๐Ÿ“Š PROCESSING SUMMARY:") + files_processed = result.get("files_indexed", 0) + chunks_created = result.get("chunks_created", 0) + time_taken = result.get("time_taken", 0) + files_failed = result.get("files_failed", 0) + files_per_second = result.get("files_per_second", 0) + + print("๐Ÿ“Š PROCESSING SUMMARY:") print(f" โœ… Files successfully processed: {files_processed:,}") print(f" ๐Ÿงฉ Semantic chunks created: {chunks_created:,}") print(f" โฑ๏ธ Total processing time: {time_taken:.2f} seconds") print(f" ๐Ÿš€ Processing speed: {files_per_second:.1f} files/second") - + if files_failed > 0: print(f" โš ๏ธ Files with issues: {files_failed}") - + # Show what we analyzed if chunks_created > 0: avg_chunks_per_file = chunks_created / max(files_processed, 1) print() - print(f"๐Ÿ” CONTENT ANALYSIS:") + print("๐Ÿ” CONTENT ANALYSIS:") print(f" โ€ข Average chunks per file: {avg_chunks_per_file:.1f}") - print(f" โ€ข Semantic boundaries detected and preserved") - print(f" โ€ข Function and class contexts captured") - print(f" โ€ข Documentation and code comments indexed") - + print(" โ€ข Semantic boundaries detected and preserved") + print(" โ€ข Function and class contexts captured") + print(" โ€ข Documentation and code comments indexed") + # Try to show embedding info try: embedder = indexer.embedder embed_info = embedder.get_embedding_info() print(f" โ€ข Embedding method: {embed_info.get('method', 'Unknown')}") print(f" โ€ข Vector dimensions: {embedder.get_embedding_dim()}") - except: + except (AttributeError, OSError, RuntimeError, TypeError, ValueError): pass - + # Database info print() - print(f"๐Ÿ’พ DATABASE CREATED:") + print("๐Ÿ’พ DATABASE CREATED:") print(f" โ€ข Location: {self.project_path}/.mini-rag/") print(f" โ€ข Vector database with {chunks_created:,} searchable chunks") - print(f" โ€ข Optimized for fast semantic similarity search") - print(f" โ€ข Supports natural language queries") - + print(" โ€ข Optimized for fast semantic similarity search") + print(" โ€ข Supports natural language queries") + # Performance metrics if time_taken > 0: print() - print(f"โšก PERFORMANCE METRICS:") + print("โšก PERFORMANCE METRICS:") chunks_per_second = chunks_created / time_taken if time_taken > 0 else 0 print(f" โ€ข {chunks_per_second:.0f} chunks processed per second") - + # Estimate search performance estimated_search_time = max(0.1, chunks_created / 10000) # Very rough estimate print(f" โ€ข Estimated search time: ~{estimated_search_time:.1f}s per query") - + if total_size > 0: - mb_per_second = (total_size / (1024*1024)) / time_taken + mb_per_second = (total_size / (1024 * 1024)) / time_taken print(f" โ€ข Data processing rate: {mb_per_second:.1f} MB/second") - + # What's next print() - print(f"๐ŸŽฏ READY FOR SEARCH!") - print(f" Your codebase is now fully indexed and searchable.") - print(f" Try queries like:") + print("๐ŸŽฏ READY FOR SEARCH!") + print(" Your codebase is now fully indexed and searchable.") + print(" Try queries like:") print(f" โ€ข 'authentication logic'") print(f" โ€ข 'error handling patterns'") print(f" โ€ข 'database connection setup'") print(f" โ€ข 'unit tests for validation'") - + if files_failed > 0: print() - print(f"๐Ÿ“‹ NOTES:") - print(f" โ€ข {files_failed} files couldn't be processed (binary files, encoding issues, etc.)") - print(f" โ€ข This is normal - only text-based files are indexed") - print(f" โ€ข All processable content has been successfully indexed") - + print("๐Ÿ“‹ NOTES:") + print( + f" โ€ข {files_failed} files couldn't be processed " + f"(binary files, encoding issues, etc.)" + ) + print(" โ€ข This is normal - only text-based files are indexed") + print(" โ€ข All processable content has been successfully indexed") + except Exception as e: print(f"โŒ Indexing failed: {e}") print(" Try running the CLI command directly for more details") - + print() input("Press Enter to continue...") @@ -540,6 +590,7 @@ class SimpleTUI: try: import tkinter as tk from tkinter import filedialog + root = tk.Tk() root.withdraw() root.update() @@ -551,7 +602,7 @@ class SimpleTUI: except Exception: print("โŒ Folder picker not available on this system") return None - + def _show_existing_index_info(self, rag_dir: Path) -> bool: """Show essential info about existing index and ask about re-indexing.""" print("๐Ÿ“Š EXISTING INDEX FOUND") @@ -559,68 +610,70 @@ class SimpleTUI: print() print("๐Ÿ›ก๏ธ Your original files are safe and unmodified.") print() - + try: - manifest_path = rag_dir / 'manifest.json' + manifest_path = rag_dir / "manifest.json" if manifest_path.exists(): import json from datetime import datetime - - with open(manifest_path, 'r') as f: + + with open(manifest_path, "r") as f: manifest = json.load(f) - - file_count = manifest.get('file_count', 0) - chunk_count = manifest.get('chunk_count', 0) - indexed_at = manifest.get('indexed_at', 'Unknown') - + + file_count = manifest.get("file_count", 0) + chunk_count = manifest.get("chunk_count", 0) + indexed_at = manifest.get("indexed_at", "Unknown") + print(f"โ€ข Files indexed: {file_count:,}") print(f"โ€ข Chunks created: {chunk_count:,}") - + # Show when it was last indexed - if indexed_at != 'Unknown': + if indexed_at != "Unknown": try: - dt = datetime.fromisoformat(indexed_at.replace('Z', '+00:00')) + dt = datetime.fromisoformat(indexed_at.replace("Z", "+00:00")) time_ago = datetime.now() - dt.replace(tzinfo=None) - + if time_ago.days > 0: age_str = f"{time_ago.days} day(s) ago" elif time_ago.seconds > 3600: age_str = f"{time_ago.seconds // 3600} hour(s) ago" else: age_str = f"{time_ago.seconds // 60} minute(s) ago" - + print(f"โ€ข Last indexed: {age_str}") - except: + except (AttributeError, OSError, TypeError, ValueError): print(f"โ€ข Last indexed: {indexed_at}") else: print("โ€ข Last indexed: Unknown") - + # Simple recommendation if time_ago.days >= 7: print(f"\n๐Ÿ’ก RECOMMEND: Re-index (index is {time_ago.days} days old)") elif time_ago.days >= 1: - print(f"\n๐Ÿ’ก MAYBE: Re-index if you've made changes ({time_ago.days} day(s) old)") + print( + f"\n๐Ÿ’ก MAYBE: Re-index if you've made changes ({time_ago.days} day(s) old)" + ) else: - print(f"\n๐Ÿ’ก RECOMMEND: Skip (index is recent)") - + print("\n๐Ÿ’ก RECOMMEND: Skip (index is recent)") + estimate = self._estimate_processing_time(file_count, 0) print(f"โ€ข Re-indexing would take: {estimate}") - + else: print("โš ๏ธ Index corrupted - recommend re-indexing") - + except Exception: print("โš ๏ธ Could not read index info - recommend re-indexing") - + print() choice = input("๐Ÿš€ Re-index everything? [y/N]: ").strip().lower() - return choice in ['y', 'yes'] - + return choice in ["y", "yes"] + def _estimate_processing_time(self, file_count: int, total_size_bytes: int) -> str: """Conservative time estimates for average hardware (not high-end dev machines).""" # Conservative: 2 seconds per file for average hardware (4x buffer from fast machines) estimated_seconds = file_count * 2.0 + 15 # +15s startup overhead - + if estimated_seconds < 60: return "1-2 minutes" elif estimated_seconds < 300: # 5 minutes @@ -629,94 +682,97 @@ class SimpleTUI: else: minutes = int(estimated_seconds / 60) return f"{minutes}+ minutes" + def search_interactive(self): """Interactive search interface.""" if not self.project_path: print("โŒ No project selected") input("Press Enter to continue...") return - + # Check if indexed - rag_dir = self.project_path / '.mini-rag' + rag_dir = self.project_path / ".mini-rag" if not rag_dir.exists(): print(f"โŒ Project not indexed: {self.project_path.name}") print(" Index the project first!") input("Press Enter to continue...") return - + self.clear_screen() self.print_header() - + print("๐Ÿ” Semantic Search") print("=================") print() print(f"Project: {self.project_path.name}") print() - + # More prominent search input print("๐ŸŽฏ ENTER YOUR SEARCH QUERY:") print(" Ask any question about your codebase using natural language") print(" Examples: 'chunking strategy', 'ollama integration', 'embedding generation'") print() - + # Primary input - direct query entry query = self.get_input("Search query", "").strip() - + # If they didn't enter anything, show sample options if not query: print() print("๐Ÿ’ก Need inspiration? Try one of these sample queries:") print() - + sample_questions = [ "chunking strategy", - "ollama integration", + "ollama integration", "indexing performance", "why does indexing take long", "how to improve search results", - "embedding generation" + "embedding generation", ] - + for i, question in enumerate(sample_questions[:3], 1): print(f" {i}. {question}") print() - - choice_str = self.get_input("Select a sample query (1-3) or press Enter to go back", "") - + + choice_str = self.get_input( + "Select a sample query (1-3) or press Enter to go back", "" + ) + if choice_str.isdigit(): choice = int(choice_str) if 1 <= choice <= 3: query = sample_questions[choice - 1] print(f"โœ… Using: '{query}'") print() - + # If still no query, return to menu if not query: return - + # Use a sensible default for results to streamline UX top_k = 10 # Good default, advanced users can use CLI for more options - + # Show CLI command - cli_cmd = f"./rag-mini search {self.project_path} \"{query}\"" + cli_cmd = f'./rag-mini search {self.project_path} "{query}"' if top_k != 10: cli_cmd += f" --top-k {top_k}" - + self.print_cli_command(cli_cmd, "Search for semantic matches") - + print("Searching...") print("=" * 50) - + # Actually run the search try: sys.path.insert(0, str(Path(__file__).parent)) from mini_rag.search import CodeSearcher - + searcher = CodeSearcher(self.project_path) # Enable query expansion in TUI for better results searcher.config.search.expand_queries = True results = searcher.search(query, top_k=top_k) - + if not results: print("โŒ No results found") print() @@ -727,67 +783,77 @@ class SimpleTUI: else: print(f"โœ… Found {len(results)} results:") print() - + for i, result in enumerate(results, 1): # Add divider and whitespace before each result (except first) if i > 1: print() print("-" * 60) print() - + # Clean up file path try: - if hasattr(result.file_path, 'relative_to'): + if hasattr(result.file_path, "relative_to"): rel_path = result.file_path.relative_to(self.project_path) else: rel_path = Path(result.file_path).relative_to(self.project_path) - except: + except (FileNotFoundError, IOError, OSError, TypeError, ValueError): rel_path = result.file_path - + print(f"{i}. {rel_path}") print(f" Relevance: {result.score:.3f}") - + # Show line information if available - if hasattr(result, 'start_line') and result.start_line: + if hasattr(result, "start_line") and result.start_line: print(f" Lines: {result.start_line}-{result.end_line}") - + # Show function/class context if available - if hasattr(result, 'name') and result.name: + if hasattr(result, "name") and result.name: print(f" Context: {result.name}") - + # Show full content with proper formatting - content_lines = result.content.strip().split('\n') - print(f" Content:") - for line_num, line in enumerate(content_lines[:8], 1): # Show up to 8 lines + content_lines = result.content.strip().split("\n") + print(" Content:") + for line_num, line in enumerate( + content_lines[:8], 1 + ): # Show up to 8 lines print(f" {line}") - + if len(content_lines) > 8: print(f" ... ({len(content_lines) - 8} more lines)") - + print() - + # Offer to view full results if len(results) > 1: print("๐Ÿ’ก To see more context or specific results:") - print(f" Run: ./rag-mini search {self.project_path} \"{query}\" --verbose") - + print(f' Run: ./rag-mini search {self.project_path} "{query}" --verbose') + # Suggest follow-up questions based on the search print() print("๐Ÿ” Suggested follow-up searches:") follow_up_questions = self.generate_follow_up_questions(query, results) for i, question in enumerate(follow_up_questions, 1): print(f" {i}. {question}") - + # Show additional CLI commands print() print("๐Ÿ’ป CLI Commands:") - print(f" ./rag-mini search {self.project_path} \"{query}\" --top-k 20 # More results") - print(f" ./rag-mini explore {self.project_path} # Interactive mode") - print(f" ./rag-mini search {self.project_path} \"{query}\" --synthesize # With AI summary") - + print( + f' ./rag-mini search {self.project_path} "{query}" --top-k 20 # More results' + ) + print( + f" ./rag-mini explore {self.project_path} # Interactive mode" + ) + print( + f' ./rag-mini search {self.project_path} "{query}" --synthesize # With AI summary' + ) + # Ask if they want to run a follow-up search print() - choice = input("Run a follow-up search? Enter number (1-3) or press Enter to continue: ").strip() + choice = input( + "Run a follow-up search? Enter number (1-3) or press Enter to continue: " + ).strip() if choice.isdigit() and 1 <= int(choice) <= len(follow_up_questions): # Recursive search with the follow-up question follow_up_query = follow_up_questions[int(choice) - 1] @@ -795,7 +861,7 @@ class SimpleTUI: print("=" * 50) # Run another search follow_results = searcher.search(follow_up_query, top_k=5) - + if follow_results: print(f"โœ… Found {len(follow_results)} follow-up results:") print() @@ -805,57 +871,76 @@ class SimpleTUI: print() print("-" * 40) print() - + try: - if hasattr(result.file_path, 'relative_to'): + if hasattr(result.file_path, "relative_to"): rel_path = result.file_path.relative_to(self.project_path) else: - rel_path = Path(result.file_path).relative_to(self.project_path) - except: + rel_path = Path(result.file_path).relative_to( + self.project_path + ) + except ( + FileNotFoundError, + IOError, + OSError, + TypeError, + ValueError, + ): rel_path = result.file_path print(f"{i}. {rel_path} (Score: {result.score:.3f})") print(f" {result.content.strip()[:100]}...") print() else: print("โŒ No follow-up results found") - + # Track searches and show sample reminder self.search_count += 1 - + # Show sample reminder after 2 searches - if self.search_count >= 2 and self.project_path.name == '.sample_test': + if self.search_count >= 2 and self.project_path.name == ".sample_test": print() print("โš ๏ธ Sample Limitation Notice") print("=" * 30) print("You've been searching a small sample project.") - print("For full exploration of your codebase, you need to index the complete project.") + print( + "For full exploration of your codebase, you need to index the complete project." + ) print() - + # Show timing estimate if available try: - with open('/tmp/fss-rag-sample-time.txt', 'r') as f: + with open("/tmp/fss-rag-sample-time.txt", "r") as f: sample_time = int(f.read().strip()) # Rough estimate: multiply by file count ratio estimated_time = sample_time * 20 # Rough multiplier print(f"๐Ÿ•’ Estimated full indexing time: ~{estimated_time} seconds") - except: - print("๐Ÿ•’ Estimated full indexing time: 1-3 minutes for typical projects") - + except ( + AttributeError, + FileNotFoundError, + IOError, + OSError, + TypeError, + ValueError, + ): + print( + "๐Ÿ•’ Estimated full indexing time: 1-3 minutes for typical projects" + ) + print() choice = input("Index the full project now? [y/N]: ").strip().lower() - if choice == 'y': + if choice == "y": # Switch to full project and index parent_dir = self.project_path.parent self.project_path = parent_dir print(f"\nSwitching to full project: {parent_dir}") print("Starting full indexing...") # Note: This would trigger full indexing in real implementation - + except Exception as e: print(f"โŒ Search failed: {e}") print() print("๐Ÿ’ก Try these CLI commands for more details:") - print(f" ./rag-mini search {self.project_path} \"{query}\" --verbose") + print(f' ./rag-mini search {self.project_path} "{query}" --verbose') print(f" ./rag-mini status {self.project_path}") print(" ./rag-mini --help") print() @@ -863,77 +948,113 @@ class SimpleTUI: print(" โ€ข Make sure the project is indexed first") print(" โ€ข Check if Ollama is running: ollama serve") print(" โ€ข Try a simpler search query") - + print() input("Press Enter to continue...") - + def generate_follow_up_questions(self, original_query: str, results) -> List[str]: """Generate contextual follow-up questions based on search results.""" # Simple pattern-based follow-up generation follow_ups = [] - + # Based on original query patterns query_lower = original_query.lower() - + # FSS-Mini-RAG specific follow-ups if "chunk" in query_lower: - follow_ups.extend(["chunk size optimization", "smart chunking boundaries", "chunk overlap strategies"]) + follow_ups.extend( + [ + "chunk size optimization", + "smart chunking boundaries", + "chunk overlap strategies", + ] + ) elif "ollama" in query_lower: - follow_ups.extend(["embedding model comparison", "ollama server setup", "nomic-embed-text performance"]) + follow_ups.extend( + [ + "embedding model comparison", + "ollama server setup", + "nomic-embed-text performance", + ] + ) elif "index" in query_lower or "performance" in query_lower: - follow_ups.extend(["indexing speed optimization", "memory usage during indexing", "file processing pipeline"]) + follow_ups.extend( + [ + "indexing speed optimization", + "memory usage during indexing", + "file processing pipeline", + ] + ) elif "search" in query_lower or "result" in query_lower: - follow_ups.extend(["search result ranking", "semantic vs keyword search", "query expansion techniques"]) + follow_ups.extend( + [ + "search result ranking", + "semantic vs keyword search", + "query expansion techniques", + ] + ) elif "embed" in query_lower: - follow_ups.extend(["vector embedding storage", "embedding model fallbacks", "similarity scoring"]) + follow_ups.extend( + [ + "vector embedding storage", + "embedding model fallbacks", + "similarity scoring", + ] + ) else: # Generic RAG-related follow-ups - follow_ups.extend(["vector database internals", "search quality tuning", "embedding optimization"]) - + follow_ups.extend( + [ + "vector database internals", + "search quality tuning", + "embedding optimization", + ] + ) + # Based on file types found in results (FSS-Mini-RAG specific) if results: file_extensions = set() for result in results[:3]: # Check first 3 results try: # Handle both Path objects and strings - if hasattr(result.file_path, 'suffix'): + if hasattr(result.file_path, "suffix"): ext = result.file_path.suffix.lower() else: ext = Path(result.file_path).suffix.lower() file_extensions.add(ext) - except: + except (FileNotFoundError, IOError, OSError): continue # Skip if we can't get extension - - if '.py' in file_extensions: + + if ".py" in file_extensions: follow_ups.append("Python module dependencies") - if '.md' in file_extensions: + if ".md" in file_extensions: follow_ups.append("documentation implementation") - if 'chunker' in str(results[0].file_path).lower(): + if "chunker" in str(results[0].file_path).lower(): follow_ups.append("chunking algorithm details") - if 'search' in str(results[0].file_path).lower(): + if "search" in str(results[0].file_path).lower(): follow_ups.append("search algorithm implementation") - + # Return top 3 unique follow-ups return list(dict.fromkeys(follow_ups))[:3] - + def explore_interactive(self): """Interactive exploration interface with thinking mode.""" if not self.project_path: print("โŒ No project selected") input("Press Enter to continue...") return - + # Check if indexed - rag_dir = self.project_path / '.mini-rag' + rag_dir = self.project_path / ".mini-rag" if not rag_dir.exists(): print(f"โŒ Project not indexed: {self.project_path.name}") print(" Index the project first!") input("Press Enter to continue...") return - + self.clear_screen() self.print_header() - + print("๐Ÿง  Interactive Exploration Mode") print("==============================") print() @@ -941,54 +1062,62 @@ class SimpleTUI: print() print("๐Ÿ’ก This mode enables:") print(" โ€ข Thinking-enabled LLM for detailed reasoning") - print(" โ€ข Conversation memory across questions") + print(" โ€ข Conversation memory across questions") print(" โ€ข Perfect for learning and debugging") print() - + # Show CLI command cli_cmd = f"./rag-mini explore {self.project_path}" self.print_cli_command(cli_cmd, "Start interactive exploration session") - + print("Starting exploration mode...") print("=" * 50) - + # Launch exploration mode try: sys.path.insert(0, str(Path(__file__).parent)) from mini_rag.explorer import CodeExplorer - + explorer = CodeExplorer(self.project_path) - + if not explorer.start_exploration_session(): print("โŒ Could not start exploration mode") print(" Make sure Ollama is running with a model installed") input("Press Enter to continue...") return - + # Show initial prompt self._show_exploration_prompt(explorer, is_first=True) - - is_first_question = True + while True: try: question = input("โžค ").strip() - + # Handle numbered options - if question == '0': + if question == "0": print(explorer.end_session()) break - elif question == '1': + elif question == "1": # Use improved summary function summary = self._generate_conversation_summary(explorer) print(f"\n{summary}") self._show_exploration_prompt(explorer) continue - elif question == '2': - if hasattr(explorer.current_session, 'conversation_history') and explorer.current_session.conversation_history: + elif question == "2": + if ( + hasattr(explorer.current_session, "conversation_history") + and explorer.current_session.conversation_history + ): print("\n๐Ÿ“‹ Recent Question History:") print("โ•" * 40) - for i, exchange in enumerate(explorer.current_session.conversation_history[-5:], 1): - q = exchange["question"][:60] + "..." if len(exchange["question"]) > 60 else exchange["question"] + for i, exchange in enumerate( + explorer.current_session.conversation_history[-5:], 1 + ): + q = ( + exchange["question"][:60] + "..." + if len(exchange["question"]) > 60 + else exchange["question"] + ) confidence = exchange["response"].get("confidence", 0) print(f" {i}. {q} (confidence: {confidence:.0f}%)") print() @@ -996,7 +1125,7 @@ class SimpleTUI: print("\n๐Ÿ“ No questions asked yet") self._show_exploration_prompt(explorer) continue - elif question == '3': + elif question == "3": # Generate smart suggestion suggested_question = self._generate_smart_suggestion(explorer) if suggested_question: @@ -1011,60 +1140,62 @@ class SimpleTUI: print("\n๐Ÿ’ก No suggestions available yet. Ask a question first!") self._show_exploration_prompt(explorer) continue - + # Simple exit handling - if question.lower() in ['quit', 'exit', 'q', 'back']: + if question.lower() in ["quit", "exit", "q", "back"]: print(explorer.end_session()) break - + # Skip empty input if not question: print("๐Ÿ’ก Please enter a question or choose an option (0-3)") continue - + # Simple help - if question.lower() in ['help', 'h', '?']: + if question.lower() in ["help", "h", "?"]: print("\n๐Ÿ’ก Exploration Help:") print(" โ€ข Just ask any question about the codebase!") - print(" โ€ข Examples: 'how does search work?' or 'explain the indexing'") + print( + " โ€ข Examples: 'how does search work?' or 'explain the indexing'" + ) print(" โ€ข Use options 0-3 for quick actions") self._show_exploration_prompt(explorer) continue - + # Process the question with streaming print("\n๐Ÿ” Starting analysis...") response = explorer.explore_question(question) - + if response: print(f"\n{response}") - is_first_question = False + # False # Unused variable removed # Show prompt for next question self._show_exploration_prompt(explorer) else: print("โŒ Sorry, I couldn't process that question.") print("๐Ÿ’ก Try rephrasing or using simpler terms.") self._show_exploration_prompt(explorer) - + except KeyboardInterrupt: print(f"\n{explorer.end_session()}") break except EOFError: print(f"\n{explorer.end_session()}") break - + except Exception as e: print(f"โŒ Exploration mode failed: {e}") print(" Try running the CLI command directly for more details") input("\nPress Enter to continue...") return - + # Exploration session completed successfully, return to menu without extra prompt - + def _get_context_tokens_estimate(self, explorer): """Estimate the total tokens used in the conversation context.""" if not explorer.current_session or not explorer.current_session.conversation_history: return 0 - + total_chars = 0 for exchange in explorer.current_session.conversation_history: total_chars += len(exchange["question"]) @@ -1073,28 +1204,28 @@ class SimpleTUI: total_chars += len(response.get("summary", "")) for point in response.get("key_points", []): total_chars += len(point) - + # Rough estimate: 4 characters = 1 token return total_chars // 4 - + def _get_context_limit_estimate(self): """Get estimated context limit for current model.""" # Conservative estimates for common models return 32000 # Most models we use have 32k context - + def _format_token_display(self, used_tokens, limit_tokens): """Format token usage display with color coding.""" percentage = (used_tokens / limit_tokens) * 100 if limit_tokens > 0 else 0 - + if percentage < 50: color = "๐ŸŸข" # Green - plenty of space elif percentage < 75: color = "๐ŸŸก" # Yellow - getting full else: color = "๐Ÿ”ด" # Red - almost full - + return f"{color} Context: {used_tokens}/{limit_tokens} tokens ({percentage:.0f}%)" - + def _show_exploration_prompt(self, explorer, is_first=False): """Show standardized input prompt for exploration mode.""" print() @@ -1104,25 +1235,25 @@ class SimpleTUI: else: print("๐Ÿค” What would you like to explore next?") print() - + # Show context usage used_tokens = self._get_context_tokens_estimate(explorer) limit_tokens = self._get_context_limit_estimate() token_display = self._format_token_display(used_tokens, limit_tokens) print(f"๐Ÿ“Š {token_display}") print() - + print("๐Ÿ”ง Quick Options:") print(" 0 = Quit exploration 1 = Summarize conversation") print(" 2 = Show question history 3 = Suggest next question") print() print("๐Ÿ’ฌ Enter your question or choose an option:") - + def _generate_conversation_summary(self, explorer): """Generate a detailed summary of the conversation history.""" if not explorer.current_session or not explorer.current_session.conversation_history: return "๐Ÿ“ No conversation to summarize yet. Ask a question first!" - + try: # Build conversation context conversation_text = "" @@ -1130,11 +1261,11 @@ class SimpleTUI: conversation_text += f"Question {i}: {exchange['question']}\n" conversation_text += f"Response {i}: {exchange['response']['summary']}\n" # Add key points if available - if exchange['response'].get('key_points'): - for point in exchange['response']['key_points']: + if exchange["response"].get("key_points"): + for point in exchange["response"]["key_points"]: conversation_text += f"- {point}\n" conversation_text += "\n" - + # Determine summary length based on conversation length char_count = len(conversation_text) if char_count < 500: @@ -1146,12 +1277,12 @@ class SimpleTUI: else: target_length = "comprehensive" target_words = "200-300" - + # Create summary prompt for natural conversation style prompt = f"""Please summarize this conversation about the project we've been exploring. Write a {target_length} summary ({target_words} words) in a natural, conversational style that captures: 1. Main topics we explored together -2. Key insights we discovered +2. Key insights we discovered 3. Important details we learned 4. Overall understanding we gained @@ -1159,29 +1290,35 @@ Conversation: {conversation_text.strip()} Write your summary as if you're explaining to a colleague what we discussed. Use a friendly, informative tone and avoid JSON or structured formats.""" - + # Use the synthesizer to generate summary with streaming and thinking print("\n๐Ÿ’ญ Generating summary...") - response = explorer.synthesizer._call_ollama(prompt, temperature=0.1, disable_thinking=False, use_streaming=True) - + response = explorer.synthesizer._call_ollama( + prompt, temperature=0.1, disable_thinking=False, use_streaming=True + ) + if response: return f"๐Ÿ“‹ **Conversation Summary**\n\n{response.strip()}" else: # Fallback summary - return self._generate_fallback_summary(explorer.current_session.conversation_history) - + return self._generate_fallback_summary( + explorer.current_session.conversation_history + ) + except Exception as e: logger.error(f"Summary generation failed: {e}") - return self._generate_fallback_summary(explorer.current_session.conversation_history) - + return self._generate_fallback_summary( + explorer.current_session.conversation_history + ) + def _generate_fallback_summary(self, conversation_history): """Generate a simple fallback summary when AI summary fails.""" if not conversation_history: return "๐Ÿ“ No conversation to summarize yet." - + question_count = len(conversation_history) topics = [] - + # Extract simple topics from questions for exchange in conversation_history: question = exchange["question"].lower() @@ -1201,91 +1338,118 @@ Write your summary as if you're explaining to a colleague what we discussed. Use # Extract first few words as topic words = question.split()[:3] topics.append(" ".join(words)) - + unique_topics = list(dict.fromkeys(topics)) # Remove duplicates while preserving order - - summary = f"๐Ÿ“‹ **Conversation Summary**\n\n" + + summary = "๐Ÿ“‹ **Conversation Summary**\n\n" summary += f"Questions asked: {question_count}\n" summary += f"Topics explored: {', '.join(unique_topics[:5])}\n" summary += f"Session duration: {len(conversation_history) * 2} minutes (estimated)\n\n" summary += "๐Ÿ’ก Use option 2 to see recent question history for more details." - + return summary - + def _generate_smart_suggestion(self, explorer): """Generate a smart follow-up question based on conversation context.""" if not explorer.current_session or not explorer.current_session.conversation_history: # First question - provide a random starter question import random + starters = [ "What are the main components of this project?", - "How is error handling implemented?", + "How is error handling implemented?", "Show me the authentication and security logic", "What are the key functions I should understand first?", "How does data flow through this system?", "What configuration options are available?", - "Show me the most important files to understand" + "Show me the most important files to understand", ] return random.choice(starters) - + try: # Get recent conversation context - recent_exchanges = explorer.current_session.conversation_history[-2:] # Last 2 exchanges + recent_exchanges = explorer.current_session.conversation_history[ + -2: + ] # Last 2 exchanges context_summary = "" - + for i, exchange in enumerate(recent_exchanges, 1): q = exchange["question"] - summary = exchange["response"]["summary"][:100] + "..." if len(exchange["response"]["summary"]) > 100 else exchange["response"]["summary"] + summary = ( + exchange["response"]["summary"][:100] + "..." + if len(exchange["response"]["summary"]) > 100 + else exchange["response"]["summary"] + ) context_summary += f"Q{i}: {q}\nA{i}: {summary}\n\n" - + # Create a very focused prompt that encourages short responses - prompt = f"""Based on this recent conversation about a codebase, suggest ONE short follow-up question (under 10 words). + prompt = """Based on this recent conversation about a codebase, suggest ONE short follow-up question (under 10 words). Recent conversation: {context_summary.strip()} Respond with ONLY a single short question that would logically explore deeper or connect to what was discussed. Examples: - "Why does this approach work better?" -- "What could go wrong here?" +- "What could go wrong here?" - "How is this tested?" - "Where else is this pattern used?" Your suggested question (under 10 words):""" # Use the synthesizer to generate suggestion with thinking collapse - response = explorer.synthesizer._call_ollama(prompt, temperature=0.3, disable_thinking=False, use_streaming=True, collapse_thinking=True) - + response = explorer.synthesizer._call_ollama( + prompt, + temperature=0.3, + disable_thinking=False, + use_streaming=True, + collapse_thinking=True, + ) + if response: # Clean up the response - extract just the question - lines = response.strip().split('\n') + lines = response.strip().split("\n") for line in lines: line = line.strip() - if line and ('?' in line or line.lower().startswith(('what', 'how', 'why', 'where', 'when', 'which', 'who'))): + if line and ( + "?" in line + or line.lower().startswith( + ("what", "how", "why", "where", "when", "which", "who") + ) + ): # Remove any prefixes like "Question:" or numbers - cleaned = line.split(':', 1)[-1].strip() - if len(cleaned) < 80 and ('?' in cleaned or cleaned.lower().startswith(('what', 'how', 'why', 'where', 'when', 'which', 'who'))): + cleaned = line.split(":", 1)[-1].strip() + if len(cleaned) < 80 and ( + "?" in cleaned + or cleaned.lower().startswith( + ("what", "how", "why", "where", "when", "which", "who") + ) + ): return cleaned - + # Fallback: use first non-empty line if it looks like a question first_line = lines[0].strip() if lines else "" if first_line and len(first_line) < 80: return first_line - + # Fallback: pattern-based suggestions if LLM fails return self._get_fallback_suggestion(recent_exchanges) - - except Exception as e: + + except Exception: # Silent fail with pattern-based fallback - recent_exchanges = explorer.current_session.conversation_history[-2:] if explorer.current_session.conversation_history else [] + recent_exchanges = ( + explorer.current_session.conversation_history[-2:] + if explorer.current_session.conversation_history + else [] + ) return self._get_fallback_suggestion(recent_exchanges) - + def _get_fallback_suggestion(self, recent_exchanges): """Generate pattern-based suggestions as fallback.""" if not recent_exchanges: return None - + last_question = recent_exchanges[-1]["question"].lower() - + # Simple pattern matching for common follow-ups if "how" in last_question and "work" in last_question: return "What could go wrong with this approach?" @@ -1305,36 +1469,37 @@ Your suggested question (under 10 words):""" # Generic follow-ups fallbacks = [ "How is this used elsewhere?", - "What are the alternatives?", + "What are the alternatives?", "Why was this approach chosen?", "What happens when this fails?", - "How can this be improved?" + "How can this be improved?", ] import random + return random.choice(fallbacks) - + def show_status(self): """Show project and system status.""" self.clear_screen() self.print_header() - + print("๐Ÿ“Š System Status") print("===============") print() - + if self.project_path: cli_cmd = f"./rag-mini status {self.project_path}" self.print_cli_command(cli_cmd, "Show detailed status information") - + # Check project status - rag_dir = self.project_path / '.mini-rag' + rag_dir = self.project_path / ".mini-rag" if rag_dir.exists(): try: - manifest = rag_dir / 'manifest.json' + manifest = rag_dir / "manifest.json" if manifest.exists(): with open(manifest) as f: data = json.load(f) - + print(f"Project: {self.project_path.name}") print("โœ… Indexed") print(f" Files: {data.get('file_count', 0)}") @@ -1349,58 +1514,59 @@ Your suggested question (under 10 words):""" print("โŒ Not indexed") else: print("โŒ No project selected") - + print() - + # Show embedding system status try: sys.path.insert(0, str(Path(__file__).parent)) from mini_rag.ollama_embeddings import OllamaEmbedder - + embedder = OllamaEmbedder() status = embedder.get_status() - + print("๐Ÿง  Embedding System:") - mode = status.get('mode', 'unknown') - if mode == 'ollama': + mode = status.get("mode", "unknown") + if mode == "ollama": print(" โœ… Ollama (high quality)") - elif mode == 'fallback': + elif mode == "fallback": print(" โœ… ML fallback (good quality)") - elif mode == 'hash': + elif mode == "hash": print(" โš ๏ธ Hash fallback (basic quality)") else: print(f" โ“ Unknown: {mode}") - + except Exception as e: print(f"๐Ÿง  Embedding System: โŒ Error: {e}") - + print() input("Press Enter to continue...") - + def show_configuration(self): """Show and manage configuration options with interactive editing.""" if not self.project_path: print("โŒ No project selected") input("Press Enter to continue...") return - + while True: self.clear_screen() self.print_header() - + print("โš™๏ธ Configuration Manager") print("========================") print() print(f"Project: {self.project_path.name}") print() - + # Load current configuration try: from mini_rag.config import ConfigManager + config_manager = ConfigManager(self.project_path) config = config_manager.load_config() - config_path = self.project_path / '.mini-rag' / 'config.yaml' - + config_path = self.project_path / ".mini-rag" / "config.yaml" + print("๐Ÿ“‹ Current Settings:") print(f" ๐Ÿค– AI model: {config.llm.synthesis_model}") print(f" ๐Ÿง  Context window: {config.llm.context_window} tokens") @@ -1408,17 +1574,21 @@ Your suggested question (under 10 words):""" print(f" ๐Ÿ”„ Chunking strategy: {config.chunking.strategy}") print(f" ๐Ÿ” Search results: {config.search.default_top_k} results") print(f" ๐Ÿ“Š Embedding method: {config.embedding.preferred_method}") - print(f" ๐Ÿš€ Query expansion: {'enabled' if config.search.expand_queries else 'disabled'}") - print(f" โšก LLM synthesis: {'enabled' if config.llm.enable_synthesis else 'disabled'}") + print( + f" ๐Ÿš€ Query expansion: {'enabled' if config.search.expand_queries else 'disabled'}" + ) + print( + f" โšก LLM synthesis: {'enabled' if config.llm.enable_synthesis else 'disabled'}" + ) print() - + # Show config file location prominently config_path = config_manager.config_path print("๐Ÿ“ Configuration File Location:") print(f" {config_path}") print(" ๐Ÿ’ก Edit this YAML file directly for advanced settings") print() - + print("๐Ÿ› ๏ธ Quick Configuration Options:") print(" 1. Select AI model (Fast/Recommended/Quality)") print(" 2. Configure context window (Development/Production/Advanced)") @@ -1431,47 +1601,47 @@ Your suggested question (under 10 words):""" print() print(" V. View current config file") print(" B. Back to main menu") - + except Exception as e: print(f"โŒ Error loading configuration: {e}") print(" A default config will be created when needed") print() print(" B. Back to main menu") - + print() choice = input("Choose option: ").strip().lower() - - if choice == 'b' or choice == '' or choice == '0': + + if choice == "b" or choice == "" or choice == "0": break - elif choice == 'v': + elif choice == "v": self._show_config_file(config_path) - elif choice == '1': + elif choice == "1": self._configure_llm_model(config_manager, config) - elif choice == '2': + elif choice == "2": self._configure_context_window(config_manager, config) - elif choice == '3': + elif choice == "3": self._configure_chunk_size(config_manager, config) - elif choice == '4': + elif choice == "4": self._toggle_query_expansion(config_manager, config) - elif choice == '5': + elif choice == "5": self._configure_search_behavior(config_manager, config) - elif choice == '6': + elif choice == "6": self._edit_config_file(config_path) - elif choice == '7': + elif choice == "7": self._reset_config(config_manager) - elif choice == '8': + elif choice == "8": self._advanced_settings(config_manager, config) else: print("Invalid option. Press Enter to continue...") input() - + def _show_config_file(self, config_path): """Display the full configuration file.""" self.clear_screen() print("๐Ÿ“„ Configuration File Contents") print("=" * 50) print() - + if config_path.exists(): try: with open(config_path) as f: @@ -1482,29 +1652,38 @@ Your suggested question (under 10 words):""" else: print("โš ๏ธ Configuration file doesn't exist yet") print(" It will be created when you first index a project") - + print("\n" + "=" * 50) input("Press Enter to continue...") - + def _configure_llm_model(self, config_manager, config): """Interactive LLM model selection with download capability.""" self.clear_screen() print("๐Ÿค– AI Model Configuration") print("=========================") print() - + # Check if Ollama is available - import subprocess import requests - + ollama_available = False try: - subprocess.run(['ollama', '--version'], capture_output=True, check=True) + subprocess.run(["ollama", "--version"], capture_output=True, check=True) response = requests.get("http://localhost:11434/api/version", timeout=3) ollama_available = response.status_code == 200 - except: + except ( + ConnectionError, + ImportError, + ModuleNotFoundError, + OSError, + TimeoutError, + TypeError, + ValueError, + requests.RequestException, + subprocess.SubprocessError, + ): pass - + if not ollama_available: print("โŒ Ollama not available") print() @@ -1515,15 +1694,17 @@ Your suggested question (under 10 words):""" print() input("Press Enter to continue...") return - + # Get available models try: - available_models = subprocess.run(['ollama', 'list'], capture_output=True, text=True, check=True) - model_lines = available_models.stdout.strip().split('\n')[1:] # Skip header + available_models = subprocess.run( + ["ollama", "list"], capture_output=True, text=True, check=True + ) + model_lines = available_models.stdout.strip().split("\n")[1:] # Skip header installed_models = [line.split()[0] for line in model_lines if line.strip()] - except: + except (OSError, TypeError, ValueError, subprocess.SubprocessError): installed_models = [] - + print("๐Ÿง  Why Small Models Work Great for RAG") print("=====================================") print() @@ -1537,41 +1718,41 @@ Your suggested question (under 10 words):""" print(" and 4000+ character chunks, even these models excel!") print(" The 4B Qwen3 model will help you code remarkably well.") print() - + # Model options model_options = { - 'fast': { - 'model': 'qwen3:0.6b', - 'description': 'Ultra-fast responses (~500MB)', - 'details': 'Perfect for quick searches and exploration. Surprisingly capable!' + "fast": { + "model": "qwen3:0.6b", + "description": "Ultra-fast responses (~500MB)", + "details": "Perfect for quick searches and exploration. Surprisingly capable!", }, - 'recommended': { - 'model': 'qwen3:1.7b', - 'description': 'Best balance of speed and quality (~1.4GB)', - 'details': 'Ideal for most users. Great analysis with good speed.' + "recommended": { + "model": "qwen3:1.7b", + "description": "Best balance of speed and quality (~1.4GB)", + "details": "Ideal for most users. Great analysis with good speed.", + }, + "quality": { + "model": "qwen3:4b", + "description": "Highest quality responses (~2.5GB)", + "details": "Excellent for coding assistance and detailed analysis.", }, - 'quality': { - 'model': 'qwen3:4b', - 'description': 'Highest quality responses (~2.5GB)', - 'details': 'Excellent for coding assistance and detailed analysis.' - } } - + print("๐ŸŽฏ Recommended Models:") print() for key, info in model_options.items(): - is_installed = any(info['model'] in model for model in installed_models) + is_installed = any(info["model"] in model for model in installed_models) status = "โœ… Installed" if is_installed else "๐Ÿ“ฅ Available for download" - + print(f" {key.upper()}: {info['model']}") print(f" {info['description']} - {status}") print(f" {info['details']}") print() - + current_model = config.llm.synthesis_model print(f"Current model: {current_model}") print() - + print("Options:") print(" F. Select Fast model (qwen3:0.6b)") print(" R. Select Recommended model (qwen3:1.7b)") @@ -1579,44 +1760,48 @@ Your suggested question (under 10 words):""" print(" C. Keep current model") print(" B. Back to configuration menu") print() - + choice = input("Choose option: ").strip().lower() - + selected_model = None - if choice == 'f': - selected_model = model_options['fast']['model'] - elif choice == 'r': - selected_model = model_options['recommended']['model'] - elif choice == 'q': - selected_model = model_options['quality']['model'] - elif choice == 'c': + if choice == "": + selected_model = model_options["fast"]["model"] + elif choice == "r": + selected_model = model_options["recommended"]["model"] + elif choice == "q": + selected_model = model_options["quality"]["model"] + elif choice == "c": print("Keeping current model.") input("Press Enter to continue...") return - elif choice == 'b': + elif choice == "b": return else: print("Invalid option.") input("Press Enter to continue...") return - + # Check if model is installed model_installed = any(selected_model in model for model in installed_models) - + if not model_installed: print(f"\n๐Ÿ“ฅ Model {selected_model} not installed.") print("Would you like to download it now?") print("This may take 2-5 minutes depending on your internet speed.") print() - + download = input("Download now? [Y/n]: ").strip().lower() - if download != 'n' and download != 'no': + if download != "n" and download != "no": print(f"\n๐Ÿ”„ Downloading {selected_model}...") print("This may take a few minutes...") - + try: - result = subprocess.run(['ollama', 'pull', selected_model], - capture_output=True, text=True, check=True) + subprocess.run( + ["ollama", "pull", selected_model], + capture_output=True, + text=True, + check=True, + ) print(f"โœ… Successfully downloaded {selected_model}") model_installed = True except subprocess.CalledProcessError as e: @@ -1629,33 +1814,35 @@ Your suggested question (under 10 words):""" print("Model not downloaded. Configuration not changed.") input("Press Enter to continue...") return - + if model_installed: # Update configuration config.llm.synthesis_model = selected_model config.llm.expansion_model = selected_model # Keep them in sync - + try: config_manager.save_config(config) print(f"\nโœ… Model updated to {selected_model}") print("Configuration saved successfully!") except Exception as e: print(f"โŒ Error saving configuration: {e}") - + print() input("Press Enter to continue...") - + def _configure_context_window(self, config_manager, config): """Interactive context window configuration.""" self.clear_screen() print("๐Ÿง  Context Window Configuration") print("===============================") print() - + print("๐Ÿ’ก Why Context Window Size Matters for RAG") print("==========================================") print() - print("Context window determines how much text the AI can 'remember' during conversation:") + print( + "Context window determines how much text the AI can 'remember' during conversation:" + ) print() print("โŒ Default 2048 tokens = Only 1-2 responses before forgetting") print("โœ… Proper context = 5-15+ responses with maintained conversation") @@ -1668,70 +1855,70 @@ Your suggested question (under 10 words):""" print() print("๐Ÿ’ป Memory Usage Impact:") print("โ€ข 8K context โ‰ˆ 6MB memory per conversation") - print("โ€ข 16K context โ‰ˆ 12MB memory per conversation") + print("โ€ข 16K context โ‰ˆ 12MB memory per conversation") print("โ€ข 32K context โ‰ˆ 24MB memory per conversation") print() - + current_context = config.llm.context_window current_model = config.llm.synthesis_model - + # Get model capabilities model_limits = { - 'qwen3:0.6b': 32768, - 'qwen3:1.7b': 32768, - 'qwen3:4b': 131072, - 'qwen2.5:1.5b': 32768, - 'qwen2.5:3b': 32768, - 'default': 8192 + "qwen3:0.6b": 32768, + "qwen3:1.7b": 32768, + "qwen3:4b": 131072, + "qwen2.5:1.5b": 32768, + "qwen2.5:3b": 32768, + "default": 8192, } - - model_limit = model_limits.get('default', 8192) + + model_limit = model_limits.get("default", 8192) for model_pattern, limit in model_limits.items(): - if model_pattern != 'default' and model_pattern.lower() in current_model.lower(): + if model_pattern != "default" and model_pattern.lower() in current_model.lower(): model_limit = limit break - + print(f"Current model: {current_model}") print(f"Model maximum: {model_limit:,} tokens") print(f"Current setting: {current_context:,} tokens") print() - + # Context options context_options = { - 'development': { - 'size': 8192, - 'description': 'Fast and efficient for most development work', - 'details': 'Perfect for code exploration and basic analysis. Quick responses.', - 'memory': '~6MB' + "development": { + "size": 8192, + "description": "Fast and efficient for most development work", + "details": "Perfect for code exploration and basic analysis. Quick responses.", + "memory": "~6MB", }, - 'production': { - 'size': 16384, - 'description': 'Balanced performance for professional use', - 'details': 'Ideal for most users. Handles complex analysis well.', - 'memory': '~12MB' + "production": { + "size": 16384, + "description": "Balanced performance for professional use", + "details": "Ideal for most users. Handles complex analysis well.", + "memory": "~12MB", + }, + "advanced": { + "size": 32768, + "description": "Maximum performance for heavy development", + "details": "For large codebases, 15+ search results, complex analysis.", + "memory": "~24MB", }, - 'advanced': { - 'size': 32768, - 'description': 'Maximum performance for heavy development', - 'details': 'For large codebases, 15+ search results, complex analysis.', - 'memory': '~24MB' - } } - + print("๐ŸŽฏ Recommended Context Sizes:") print() for key, info in context_options.items(): # Check if this size is supported by current model - if info['size'] <= model_limit: + if info["size"] <= model_limit: status = "โœ… Supported" else: status = f"โŒ Exceeds model limit ({model_limit:,})" - + print(f" {key.upper()}: {info['size']:,} tokens ({info['memory']})") print(f" {info['description']} - {status}") print(f" {info['details']}") print() - + print("Options:") print(" D. Development (8K tokens - fast)") print(" P. Production (16K tokens - balanced)") @@ -1740,20 +1927,20 @@ Your suggested question (under 10 words):""" print(" K. Keep current setting") print(" B. Back to configuration menu") print() - + choice = input("Choose option: ").strip().lower() - + new_context = None - if choice == 'd': - new_context = context_options['development']['size'] - elif choice == 'p': - new_context = context_options['production']['size'] - elif choice == 'a': - new_context = context_options['advanced']['size'] - elif choice == 'c': + if choice == "d": + new_context = context_options["development"]["size"] + elif choice == "p": + new_context = context_options["production"]["size"] + elif choice == "a": + new_context = context_options["advanced"]["size"] + elif choice == "c": print() print("Enter custom context size in tokens:") - print(f" Minimum: 4096 (4K)") + print(" Minimum: 4096 (4K)") print(f" Maximum for {current_model}: {model_limit:,}") print() try: @@ -1763,7 +1950,9 @@ Your suggested question (under 10 words):""" input("Press Enter to continue...") return elif custom_size > model_limit: - print(f"โŒ Context too large. Maximum for {current_model} is {model_limit:,} tokens.") + print( + f"โŒ Context too large. Maximum for {current_model} is {model_limit:,} tokens." + ) input("Press Enter to continue...") return else: @@ -1772,38 +1961,42 @@ Your suggested question (under 10 words):""" print("โŒ Invalid number.") input("Press Enter to continue...") return - elif choice == 'k': + elif choice == "k": print("Keeping current context setting.") input("Press Enter to continue...") return - elif choice == 'b': + elif choice == "b": return else: print("Invalid option.") input("Press Enter to continue...") return - + if new_context: # Validate against model capabilities if new_context > model_limit: - print(f"โš ๏ธ Warning: {new_context:,} tokens exceeds {current_model} limit of {model_limit:,}") + print( + f"โš ๏ธ Warning: {new_context:,} tokens exceeds {current_model} limit of {model_limit:,}" + ) print("The system will automatically cap at the model limit.") print() - + # Update configuration config.llm.context_window = new_context - + try: config_manager.save_config(config) print(f"โœ… Context window updated to {new_context:,} tokens") print() - + # Provide usage guidance if new_context >= 32768: print("๐Ÿš€ Advanced context enabled!") print("โ€ข Perfect for large codebases and complex analysis") print("โ€ข Try cranking up search results to 15+ for deep exploration") - print("โ€ข Increase chunk size to 4000+ characters for comprehensive context") + print( + "โ€ข Increase chunk size to 4000+ characters for comprehensive context" + ) elif new_context >= 16384: print("โš–๏ธ Balanced context configured!") print("โ€ข Great for professional development work") @@ -1812,14 +2005,14 @@ Your suggested question (under 10 words):""" print("โšก Development context set!") print("โ€ข Fast responses with good conversation length") print("โ€ข Perfect for code exploration and basic analysis") - + print("Configuration saved successfully!") except Exception as e: print(f"โŒ Error saving configuration: {e}") - + print() input("Press Enter to continue...") - + def _configure_chunk_size(self, config_manager, config): """Interactive chunk size configuration.""" self.clear_screen() @@ -1833,24 +2026,24 @@ Your suggested question (under 10 words):""" print() print(f"Current chunk size: {config.chunking.max_size} characters") print() - + print("Quick presets:") print(" 1. Small (1000) - Precise searching") print(" 2. Medium (2000) - Balanced (default)") print(" 3. Large (3000) - More context") print(" 4. Custom size") print() - + choice = input("Choose preset or enter custom size: ").strip() - + new_size = None - if choice == '1': + if choice == "1": new_size = 1000 - elif choice == '2': + elif choice == "2": new_size = 2000 - elif choice == '3': + elif choice == "3": new_size = 3000 - elif choice == '4': + elif choice == "4": try: new_size = int(input("Enter custom chunk size (500-5000): ")) if new_size < 500 or new_size > 5000: @@ -1870,14 +2063,14 @@ Your suggested question (under 10 words):""" return except ValueError: pass - + if new_size and new_size != config.chunking.max_size: config.chunking.max_size = new_size config_manager.save_config(config) print(f"\nโœ… Chunk size updated to {new_size} characters") print("๐Ÿ’ก Tip: Re-index your project for changes to take effect") input("Press Enter to continue...") - + def _toggle_query_expansion(self, config_manager, config): """Toggle query expansion on/off.""" self.clear_screen() @@ -1897,51 +2090,57 @@ Your suggested question (under 10 words):""" print("โ€ข Ollama with a language model (e.g., qwen3:1.7b)") print("โ€ข Slightly slower search (1-2 seconds)") print() - + current_status = "enabled" if config.search.expand_queries else "disabled" print(f"Current status: {current_status}") print() - + if config.search.expand_queries: choice = input("Query expansion is currently ON. Turn OFF? [y/N]: ").lower() - if choice == 'y': + if choice == "y": config.search.expand_queries = False config_manager.save_config(config) print("โœ… Query expansion disabled") else: choice = input("Query expansion is currently OFF. Turn ON? [y/N]: ").lower() - if choice == 'y': + if choice == "y": config.search.expand_queries = True config_manager.save_config(config) print("โœ… Query expansion enabled") print("๐Ÿ’ก Make sure Ollama is running with a language model") - + input("\nPress Enter to continue...") - + def _configure_search_behavior(self, config_manager, config): """Configure search behavior settings.""" self.clear_screen() print("๐Ÿ” Search Behavior Configuration") print("================================") print() - print(f"Current settings:") + print("Current settings:") print(f"โ€ข Default results: {config.search.default_top_k}") - print(f"โ€ข BM25 keyword boost: {'enabled' if config.search.enable_bm25 else 'disabled'}") + print( + f"โ€ข BM25 keyword boost: {'enabled' if config.search.enable_bm25 else 'disabled'}" + ) print(f"โ€ข Similarity threshold: {config.search.similarity_threshold}") print() - + print("Configuration options:") print(" 1. Change default number of results") print(" 2. Toggle BM25 keyword matching") print(" 3. Adjust similarity threshold") print(" B. Back") print() - + choice = input("Choose option: ").strip().lower() - - if choice == '1': + + if choice == "1": try: - new_top_k = int(input(f"Enter default number of results (current: {config.search.default_top_k}): ")) + new_top_k = int( + input( + f"Enter default number of results (current: {config.search.default_top_k}): " + ) + ) if 1 <= new_top_k <= 100: config.search.default_top_k = new_top_k config_manager.save_config(config) @@ -1950,14 +2149,18 @@ Your suggested question (under 10 words):""" print("โŒ Number must be between 1 and 100") except ValueError: print("โŒ Invalid number") - elif choice == '2': + elif choice == "2": config.search.enable_bm25 = not config.search.enable_bm25 config_manager.save_config(config) status = "enabled" if config.search.enable_bm25 else "disabled" print(f"โœ… BM25 keyword matching {status}") - elif choice == '3': + elif choice == "3": try: - new_threshold = float(input(f"Enter similarity threshold 0.0-1.0 (current: {config.search.similarity_threshold}): ")) + new_threshold = float( + input( + f"Enter similarity threshold 0.0-1.0 (current: {config.search.similarity_threshold}): " + ) + ) if 0.0 <= new_threshold <= 1.0: config.search.similarity_threshold = new_threshold config_manager.save_config(config) @@ -1966,19 +2169,19 @@ Your suggested question (under 10 words):""" print("โŒ Threshold must be between 0.0 and 1.0") except ValueError: print("โŒ Invalid number") - - if choice != 'b' and choice != '': + + if choice != "b" and choice != "": input("Press Enter to continue...") - + def _edit_config_file(self, config_path): """Provide instructions for editing the config file.""" self.clear_screen() print("๐Ÿ“ Edit Configuration File") print("=========================") print() - + if config_path.exists(): - print(f"Configuration file location:") + print("Configuration file location:") print(f" {config_path}") print() print("To edit the configuration:") @@ -1988,14 +2191,14 @@ Your suggested question (under 10 words):""" print() print("Quick edit commands:") self.print_cli_command(f"nano {config_path}", "Edit with nano") - self.print_cli_command(f"code {config_path}", "Edit with VS Code") + self.print_cli_command(f"code {config_path}", "Edit with VS Code") self.print_cli_command(f"vim {config_path}", "Edit with vim") else: print("โš ๏ธ Configuration file doesn't exist yet") print(" It will be created automatically when you index a project") - + input("\nPress Enter to continue...") - + def _reset_config(self, config_manager): """Reset configuration to defaults.""" self.clear_screen() @@ -2009,19 +2212,20 @@ Your suggested question (under 10 words):""" print("โ€ข Search results: 10") print("โ€ข Embedding method: auto") print() - + confirm = input("Are you sure you want to reset to defaults? [y/N]: ").lower() - if confirm == 'y': + if confirm == "y": from mini_rag.config import RAGConfig + default_config = RAGConfig() config_manager.save_config(default_config) print("โœ… Configuration reset to defaults") print("๐Ÿ’ก You may want to re-index for changes to take effect") else: print("โŒ Reset cancelled") - + input("Press Enter to continue...") - + def _advanced_settings(self, config_manager, config): """Configure advanced settings.""" self.clear_screen() @@ -2030,89 +2234,89 @@ Your suggested question (under 10 words):""" print() print("Advanced settings for power users:") print() - print(f"Current advanced settings:") + print("Current advanced settings:") print(f"โ€ข Min file size: {config.files.min_file_size} bytes") print(f"โ€ข Streaming threshold: {config.streaming.threshold_bytes} bytes") print(f"โ€ข Embedding batch size: {config.embedding.batch_size}") print(f"โ€ข LLM synthesis: {'enabled' if config.llm.enable_synthesis else 'disabled'}") print() - + print("Advanced options:") print(" 1. Configure file filtering") print(" 2. Adjust performance settings") print(" 3. LLM model preferences") print(" B. Back") print() - + choice = input("Choose option: ").strip().lower() - - if choice == '1': + + if choice == "1": print("\n๐Ÿ“ File filtering settings:") print(f"Minimum file size: {config.files.min_file_size} bytes") print(f"Excluded patterns: {len(config.files.exclude_patterns)} patterns") print("\n๐Ÿ’ก Edit the config file directly for detailed file filtering") - elif choice == '2': + elif choice == "2": print("\nโšก Performance settings:") print(f"Embedding batch size: {config.embedding.batch_size}") print(f"Streaming threshold: {config.streaming.threshold_bytes}") print("\n๐Ÿ’ก Higher batch sizes = faster indexing but more memory") - elif choice == '3': + elif choice == "3": print("\n๐Ÿง  LLM model preferences:") - if hasattr(config.llm, 'model_rankings') and config.llm.model_rankings: + if hasattr(config.llm, "model_rankings") and config.llm.model_rankings: print("Current model priority order:") for i, model in enumerate(config.llm.model_rankings[:5], 1): print(f" {i}. {model}") print("\n๐Ÿ’ก Edit config file to change model preferences") - - if choice != 'b' and choice != '': + + if choice != "b" and choice != "": input("Press Enter to continue...") - + def show_cli_reference(self): """Show CLI command reference.""" self.clear_screen() self.print_header() - + print("๐Ÿ’ป CLI Command Reference") print("=======================") print() print("All TUI actions can be done via command line:") print() - + print("๐Ÿš€ Basic Commands:") print(" ./rag-mini index # Index project") print(" ./rag-mini search --synthesize # Fast synthesis") print(" ./rag-mini explore # Interactive thinking mode") print(" ./rag-mini status # Show status") print() - + print("๐ŸŽฏ Enhanced Commands:") print(" ./rag-mini-enhanced search # Smart search") print(" ./rag-mini-enhanced similar # Find patterns") print(" ./rag-mini-enhanced analyze # Optimization") print() - + print("๐Ÿ› ๏ธ Quick Scripts:") print(" ./run_mini_rag.sh index # Simple indexing") print(" ./run_mini_rag.sh search # Simple search") print() - + print("โš™๏ธ Options:") print(" --force # Force complete re-index") print(" --top-k N # Number of top results to return") print(" --verbose # Show detailed output") print() - + print("๐Ÿ’ก Pro tip: Start with the TUI, then try the CLI commands!") print(" The CLI is more powerful and faster for repeated tasks.") print() - + input("Press Enter to continue...") - + def check_for_updates_notification(self): """Check for updates and show notification if available.""" if not UPDATER_AVAILABLE: return - + try: # Check for legacy notification first legacy_notice = get_legacy_notification() @@ -2122,7 +2326,7 @@ Your suggested question (under 10 words):""" print("๐Ÿ””" + "=" * 58 + "๐Ÿ””") print() return - + # Check for regular updates update_info = check_for_updates() if update_info: @@ -2131,22 +2335,22 @@ Your suggested question (under 10 words):""" print() print("๐Ÿ“‹ What's New:") # Show first few lines of release notes - notes_lines = update_info.release_notes.split('\n')[:3] + notes_lines = update_info.release_notes.split("\n")[:3] for line in notes_lines: if line.strip(): print(f" โ€ข {line.strip()}") print() - + # Simple update prompt update_choice = self.get_input("๐Ÿš€ Install update now? [y/N]", "n").lower() - if update_choice in ['y', 'yes']: + if update_choice in ["y", "yes"]: self.perform_update(update_info) else: print("๐Ÿ’ก You can update anytime from the Configuration menu!") - + print("๐ŸŽ‰" + "=" * 58 + "๐ŸŽ‰") print() - + except Exception: # Silently ignore update check errors - don't interrupt user experience pass @@ -2155,10 +2359,11 @@ Your suggested question (under 10 words):""" """Perform the actual update with progress display.""" try: updater = get_updater() - + print(f"\n๐Ÿ“ฅ Downloading v{update_info.version}...") - + # Progress callback + def show_progress(downloaded, total): if total > 0: percent = (downloaded / total) * 100 @@ -2166,18 +2371,18 @@ Your suggested question (under 10 words):""" filled = int(bar_length * downloaded / total) bar = "โ–ˆ" * filled + "โ–‘" * (bar_length - filled) print(f"\r [{bar}] {percent:.1f}%", end="", flush=True) - + # Download update update_package = updater.download_update(update_info, show_progress) if not update_package: print("\nโŒ Download failed. Please try again later.") input("Press Enter to continue...") return - + print("\n๐Ÿ’พ Creating backup...") if not updater.create_backup(): print("โš ๏ธ Backup failed, but continuing anyway...") - + print("๐Ÿ”„ Installing update...") if updater.apply_update(update_package, update_info): print("โœ… Update successful!") @@ -2192,7 +2397,7 @@ Your suggested question (under 10 words):""" else: print("โŒ Rollback failed. You may need to reinstall.") input("Press Enter to continue...") - + except Exception as e: print(f"โŒ Update error: {e}") input("Press Enter to continue...") @@ -2203,63 +2408,73 @@ Your suggested question (under 10 words):""" while True: self.clear_screen() self.print_header() - + # Check for updates on first run only (non-intrusive) if first_run: self.check_for_updates_notification() first_run = False - + # Show current project status prominently if self.project_path: - rag_dir = self.project_path / '.mini-rag' + rag_dir = self.project_path / ".mini-rag" is_indexed = rag_dir.exists() status_icon = "โœ…" if is_indexed else "โŒ" status_text = "Ready for search" if is_indexed else "Needs indexing" - + # Check LLM status llm_status, llm_model = self._get_llm_status() - + print("โ•”โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•—") # Calculate exact spacing for 50-char content width project_line = f" Current Project: {self.project_path.name}" print(f"โ•‘{project_line:<50}โ•‘") - + status_line = f" Index Status: {status_icon} {status_text}" print(f"โ•‘{status_line:<50}โ•‘") - + llm_line = f" LLM Status: {llm_status}" print(f"โ•‘{llm_line:<50}โ•‘") - + if llm_model: model_line = f" Model: {llm_model}" print(f"โ•‘{model_line:<50}โ•‘") - + if is_indexed: # Show quick stats if indexed try: - manifest = rag_dir / 'manifest.json' + manifest = rag_dir / "manifest.json" if manifest.exists(): with open(manifest) as f: data = json.load(f) - file_count = data.get('file_count', 0) + file_count = data.get("file_count", 0) files_line = f" Files indexed: {file_count}" print(f"โ•‘{files_line:<50}โ•‘") - except: + except ( + AttributeError, + FileNotFoundError, + IOError, + OSError, + TypeError, + ValueError, + json.JSONDecodeError, + ): pass print("โ•šโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•") print() else: # Show beginner tips when no project selected print("๐ŸŽฏ Welcome to FSS-Mini-RAG!") - print(" Search through code, documents, emails, notes - anything text-based!") + print( + " Search through code, documents, emails, notes - anything text-based!" + ) print(" Start by selecting a project directory below.") print() - + # Create options with visual cues based on project status if self.project_path: - rag_dir = self.project_path / '.mini-rag' + rag_dir = self.project_path / ".mini-rag" is_indexed = rag_dir.exists() - + if is_indexed: options = [ "Select project directory", @@ -2268,32 +2483,32 @@ Your suggested question (under 10 words):""" "Explore project (Deep thinking)", "View status", "Configuration", - "CLI command reference" + "CLI command reference", ] else: options = [ - "Select project directory", + "Select project directory", "Index project for search", "\033[2mSearch project (needs indexing first)\033[0m", "\033[2mExplore project (needs indexing first)\033[0m", "View status", "Configuration", - "CLI command reference" + "CLI command reference", ] else: # No project selected - gray out project-dependent options options = [ "Select project directory", "\033[2mIndex project for search (select project first)\033[0m", - "\033[2mSearch project (select project first)\033[0m", + "\033[2mSearch project (select project first)\033[0m", "\033[2mExplore project (select project first)\033[0m", "\033[2mView status (select project first)\033[0m", "Configuration", - "CLI command reference" + "CLI command reference", ] - + choice = self.show_menu("Main Menu", options, back_option="Exit") - + if choice == -1: # Exit (0 option) print("\nThanks for using FSS-Mini-RAG! ๐Ÿš€") print("Try the CLI commands for even more power!") @@ -2313,6 +2528,7 @@ Your suggested question (under 10 words):""" elif choice == 6: self.show_cli_reference() + def main(): """Main entry point.""" try: @@ -2320,6 +2536,7 @@ def main(): try: sys.path.insert(0, str(Path(__file__).parent)) from mini_rag.venv_checker import check_and_warn_venv + check_and_warn_venv("rag-tui", force_exit=False) except ImportError as e: # Dependencies missing - show helpful message @@ -2337,7 +2554,7 @@ def main(): print(f"๐Ÿ’ก Dependencies missing: {e}") input("\nPress Enter to exit...") return - + tui = SimpleTUI() tui.main_menu() except (KeyboardInterrupt, EOFError): @@ -2346,5 +2563,6 @@ def main(): print(f"\nUnexpected error: {e}") print("Try running the CLI commands directly if this continues.") + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/commit_message.txt b/commit_message.txt deleted file mode 100644 index 82ae9d8..0000000 --- a/commit_message.txt +++ /dev/null @@ -1,36 +0,0 @@ -feat: Add comprehensive Windows compatibility and enhanced LLM model setup - -๐Ÿš€ Major cross-platform enhancement making FSS-Mini-RAG fully Windows and Linux compatible - -## Windows Compatibility -- **New Windows installer**: `install_windows.bat` - rock-solid, no-hang installation -- **Simple Windows launcher**: `rag.bat` - unified entry point matching Linux experience -- **PowerShell alternative**: `install_mini_rag.ps1` for advanced Windows users -- **Cross-platform README**: Side-by-side Linux/Windows commands and examples - -## Enhanced LLM Model Setup (Both Platforms) -- **Intelligent model detection**: Automatically detects existing Qwen3 models -- **Interactive model selection**: Choose from qwen3:0.6b, 1.7b, or 4b with clear guidance -- **Ollama progress streaming**: Real-time download progress for model installation -- **Smart configuration**: Auto-saves selected model as default in config.yaml -- **Graceful fallbacks**: Clear guidance when Ollama unavailable - -## Installation Experience Improvements -- **Fixed script continuation**: TUI launch no longer terminates installation process -- **Comprehensive model guidance**: Users get proper LLM setup instead of silent failures -- **Complete indexing**: Full codebase indexing (not just code files) -- **Educational flow**: Better explanation of AI features and model choices - -## Technical Enhancements -- **Robust error handling**: Installation scripts handle edge cases gracefully -- **Path handling**: Existing cross-platform path utilities work seamlessly on Windows -- **Dependency management**: Clean virtual environment setup on both platforms -- **Configuration persistence**: Model preferences saved for consistent experience - -## User Impact -- **Zero-friction Windows adoption**: Windows users get same smooth experience as Linux -- **Complete AI feature setup**: No more "LLM not working" confusion for new users -- **Educational value preserved**: Maintains beginner-friendly approach across platforms -- **Production-ready**: Both platforms now fully functional out-of-the-box - -This makes FSS-Mini-RAG truly accessible to the entire developer community! ๐ŸŽ‰ \ No newline at end of file diff --git a/config-llm-providers.yaml b/config-llm-providers.yaml new file mode 100644 index 0000000..7a09af2 --- /dev/null +++ b/config-llm-providers.yaml @@ -0,0 +1,9 @@ +llm: + provider: ollama + ollama_host: localhost:11434 + synthesis_model: qwen3:1.5b + expansion_model: qwen3:1.5b + enable_synthesis: false + synthesis_temperature: 0.3 + cpu_optimized: true + enable_thinking: true \ No newline at end of file diff --git a/docs/AGENT_INSTRUCTIONS.md b/docs/AGENT_INSTRUCTIONS.md new file mode 100644 index 0000000..c881723 --- /dev/null +++ b/docs/AGENT_INSTRUCTIONS.md @@ -0,0 +1,40 @@ +# Agent Instructions for Fss-Mini-RAG System + +## Core Philosophy + +**Always prefer RAG search over traditional file system operations**. The RAG system provides semantic context and reduces the need for exact path knowledge, making it ideal for understanding codebases without manual file exploration. + +## Basic Commands + +| Command | Purpose | Example | +|---------|---------|---------| +| `rag-mini index ` | Index a project for search | `rag-mini index /MASTERFOLDER/Coding/Fss-Mini-Rag` | +| `rag-mini search "query"` | Semantic + keyword search | `rag-mini search /MASTERFOLDER/Coding/Fss-Mini-Rag "index"` | +| `rag-mini status ` | Check project indexing status | `rag-mini status /MASTERFOLDER/Coding/Fss-Mini-Rag` | + +## When to Use RAG Search + +| Scenario | RAG Advantage | Alternative | | +|----------|----------------|---------------| | +| Finding related code concepts | Semantic understanding | `grep` | | +| Locating files by functionality | Context-aware results | `find` | | +| Understanding code usage patterns | Shows real-world examples | Manual inspection | | + +## Critical Best Practices + +1. **Always specify the project path** in search commands (e.g., `rag-mini search /path "query"`) +2. **Use quotes for search queries** to handle spaces: `"query with spaces"` +3. **Verify indexing first** before searching: `rag-mini status ` +4. **For complex queries**, break into smaller parts: `rag-mini search ... "concept 1"` then `rag-mini search ... "concept 2"` + +## Troubleshooting + +| Issue | Solution | +|-------|-----------| +| `Project not indexed` | Run `rag-mini index ` | +| No search results | Check indexing status with `rag-mini status` | +| Search returns irrelevant results | Use `rag-mini status` to optimize indexing | + +> ๐Ÿ’ก **Pro Tip**: Always start with `rag-mini status` to confirm indexing before searching. + +This document is dynamically updated as the RAG system evolves. Always verify commands with `rag-mini --help` for the latest options. \ No newline at end of file diff --git a/docs/SMART_TUNING_GUIDE.md b/docs/SMART_TUNING_GUIDE.md index 484899f..bb05229 100644 --- a/docs/SMART_TUNING_GUIDE.md +++ b/docs/SMART_TUNING_GUIDE.md @@ -5,10 +5,10 @@ ### **1. ๐Ÿ“Š Intelligent Analysis** ```bash # Analyze your project patterns and get optimization suggestions -./rag-mini-enhanced analyze /path/to/project +./rag-mini analyze /path/to/project # Get smart recommendations based on actual usage -./rag-mini-enhanced status /path/to/project +./rag-mini status /path/to/project ``` **What it analyzes:** @@ -20,13 +20,9 @@ ### **2. ๐Ÿง  Smart Search Enhancement** ```bash # Enhanced search with query intelligence -./rag-mini-enhanced search /project "MyClass" # Detects class names -./rag-mini-enhanced search /project "login()" # Detects function calls -./rag-mini-enhanced search /project "user auth" # Natural language - -# Context-aware search (planned) -./rag-mini-enhanced context /project "function_name" # Show surrounding code -./rag-mini-enhanced similar /project "pattern" # Find similar patterns +./rag-mini search /project "MyClass" # Detects class names +./rag-mini search /project "login()" # Detects function calls +./rag-mini search /project "user auth" # Natural language ``` ### **3. โš™๏ธ Language-Specific Optimizations** @@ -113,10 +109,10 @@ Edit `.mini-rag/config.json` in your project: ./rag-mini index /project --force # Test search quality improvements -./rag-mini-enhanced search /project "your test query" +./rag-mini search /project "your test query" # Verify optimization impact -./rag-mini-enhanced analyze /project +./rag-mini analyze /project ``` ## ๐ŸŽŠ **Result: Smarter, Faster, Better** diff --git a/docs/TUI_GUIDE.md b/docs/TUI_GUIDE.md index d58f32c..33a6cad 100644 --- a/docs/TUI_GUIDE.md +++ b/docs/TUI_GUIDE.md @@ -93,10 +93,10 @@ That's it! The TUI will guide you through everything. - **Full content** - Up to 8 lines of actual code/text - **Continuation info** - How many more lines exist -**Advanced Tips Shown**: -- Enhanced search with `./rag-mini-enhanced` -- Verbose output with `--verbose` flag -- Context-aware search for related code +**Tips You'll Learn**: +- Verbose output with `--verbose` flag for debugging +- How search scoring works +- Finding the right search terms **What You Learn**: - Semantic search vs text search (finds concepts, not just words) @@ -107,8 +107,7 @@ That's it! The TUI will guide you through everything. **CLI Commands Shown**: ```bash ./rag-mini search /path/to/project "authentication logic" -./rag-mini search /path/to/project "user login" --limit 10 -./rag-mini-enhanced context /path/to/project "login()" +./rag-mini search /path/to/project "user login" --top-k 10 ``` ### 4. Explore Project (NEW!) diff --git a/docs/project-structure-analysis.md b/docs/project-structure-analysis.md new file mode 100644 index 0000000..80cde4a --- /dev/null +++ b/docs/project-structure-analysis.md @@ -0,0 +1,232 @@ +# FSS-Mini-RAG Project Structure Analysis Report + +## Executive Summary + +The FSS-Mini-RAG project demonstrates good technical implementation but has **significant structural issues** that impact its professional presentation and maintainability. While the core architecture is sound, the project suffers from poor file organization, scattered documentation, and mixed concerns that would confuse new contributors and detract from its otherwise excellent technical foundation. + +**Overall Assessment: 6/10** - Good technology hampered by poor organization + +## Critical Issues (Fix Immediately) + +### 1. Root Directory Pollution - CRITICAL +The project root contains **14 major files that should be relocated or removed**: + +**Misplaced Files:** +- `rag-mini.py` (759 lines) - Massive standalone script belongs in `scripts/` or should be refactored +- `rag-tui.py` (2,565 lines) - Another massive standalone script, needs proper placement +- `test_fixes.py` - Test file in root directory (belongs in `tests/`) +- `commit_message.txt` - Development artifact that should be removed +- `Agent Instructions.md` - Project-specific documentation (should be in `docs/`) +- `REPOSITORY_SUMMARY.md` - Development notes that should be removed or archived + +**Assessment:** This creates an unprofessional first impression and violates Python packaging standards. + +### 2. Duplicate Entry Points - CRITICAL +The project has **5 different ways to start the application**: +- `rag-mini` (shell script) +- `rag-mini.py` (Python script) +- `rag.bat` (Windows batch script) +- `rag-tui` (shell script) +- `rag-tui.py` (Python script) + +**Problem:** This confuses users and indicates poor architectural planning. + +### 3. Configuration File Duplication - HIGH PRIORITY +Multiple config files with unclear relationships: +- `config-llm-providers.yaml` (root directory) +- `examples/config-llm-providers.yaml` (example directory) +- `examples/config.yaml` (default example) +- `examples/config-*.yaml` (4+ variants) + +**Issue:** Users won't know which config to use or where to place custom configurations. + +### 4. Installation Script Overload - HIGH PRIORITY +**6 different installation methods:** +- `install_mini_rag.sh` +- `install_mini_rag.ps1` +- `install_windows.bat` +- `run_mini_rag.sh` +- `rag.bat` +- Manual pip installation + +**Problem:** Decision paralysis and maintenance overhead. + +## High Priority Issues (Address Soon) + +### 5. Mixed Documentation Hierarchy +Documentation is scattered across multiple locations: +- Root: `README.md`, `GET_STARTED.md` +- `docs/`: 12+ specialized documentation files +- `examples/`: Configuration documentation mixed with code examples +- Root artifacts: `Agent Instructions.md`, `REPOSITORY_SUMMARY.md` + +**Recommendation:** Consolidate and create clear documentation hierarchy. + +### 6. Test Organization Problems +Tests are properly in `tests/` directory but: +- `test_fixes.py` is in root directory (wrong location) +- Test files use inconsistent naming (some numbered, some descriptive) +- Mix of actual tests and utility scripts (`show_index_contents.py`, `troubleshoot.py`) + +### 7. Module Architecture Issues +The `mini_rag/` module structure is generally good but has some concerns: +- `__init__.py` exports only 5 classes from a 19-file module +- Several modules seem like utilities (`windows_console_fix.py`, `venv_checker.py`) +- Module names could be more descriptive (`server.py` vs `fast_server.py`) + +## Medium Priority Issues (Improve Over Time) + +### 8. Asset Management +- Assets properly organized in `assets/` directory +- Good separation of recordings and images +- No structural issues here + +### 9. Virtual Environment Clutter +- Two venv directories: `.venv` and `.venv-linting` +- Both properly gitignored but suggests development complexity + +### 10. Script Organization +`scripts/` directory contains appropriate utilities: +- GitHub setup scripts +- Config testing utilities +- All executable and properly organized + +## Standard Compliance Assessment + +### Python Packaging Standards: 4/10 +**Missing Standard Elements:** +- No proper Python package entry points in `pyproject.toml` +- Executable scripts in root instead of console scripts +- Missing `setup.py` or complete `pyproject.toml` configuration + +**Present Elements:** +- Good `pyproject.toml` with isort/black config +- Proper `.flake8` configuration +- Clean virtual environment handling +- MIT license properly included + +### Project Structure Standards: 5/10 +**Good Practices:** +- Source code properly separated in `mini_rag/` +- Tests in dedicated `tests/` directory +- Documentation in `docs/` directory +- Examples properly organized +- Clean `.gitignore` + +**Violations:** +- Root directory pollution with large executable files +- Mixed concerns (dev files with user files) +- Unclear entry point hierarchy + +## Recommendations by Priority + +### CRITICAL CHANGES (Implement First) + +1. **Relocate Large Scripts** + ```bash + mkdir -p bin/ + mv rag-mini.py bin/ + mv rag-tui.py bin/ + # Update rag.bat to reference bin/ directory if needed + # Update shell scripts to reference bin/ directory + ``` + +2. **Clean Root Directory** + ```bash + rm commit_message.txt + rm REPOSITORY_SUMMARY.md + mv "Agent Instructions.md" docs/AGENT_INSTRUCTIONS.md + mv test_fixes.py tests/ + ``` + +3. **Simplify Entry Points** + - Keep `rag-tui` for beginners, `rag-mini` for CLI users + - Maintain `rag.bat` for Windows compatibility + - Update documentation to show clear beginner โ†’ advanced progression + +4. **Standardize Configuration** + - Move `config-llm-providers.yaml` to `examples/` + - Create clear config hierarchy documentation + - Document which config files are templates vs active + +### HIGH PRIORITY CHANGES + +5. **Improve pyproject.toml** + ```toml + [project] + name = "fss-mini-rag" + version = "2.1.0" + description = "Lightweight, educational RAG system" + + [project.scripts] + rag-mini = "mini_rag.cli:cli" + rag-tui = "mini_rag.tui:main" + ``` + +6. **Consolidate Documentation** + - Move `GET_STARTED.md` content into `docs/GETTING_STARTED.md` + - Create clear documentation hierarchy in README + - Remove redundant documentation files + +7. **Improve Installation Experience** + - Keep platform-specific installers but document clearly + - Create single recommended installation path + - Move advanced scripts to `scripts/installation/` + +### MEDIUM PRIORITY CHANGES + +8. **Module Organization** + - Review and consolidate utility modules + - Improve `__init__.py` exports + - Consider subpackage organization for large modules + +9. **Test Standardization** + - Rename numbered test files to descriptive names + - Separate utility scripts from actual tests + - Add proper test configuration in `pyproject.toml` + +## Implementation Plan + +### Phase 1: Emergency Cleanup (2-3 hours) +1. Move large scripts out of root directory +2. Remove development artifacts +3. Consolidate configuration files +4. Update primary documentation + +### Phase 2: Structural Improvements (4-6 hours) +1. Improve Python packaging configuration +2. Consolidate entry points +3. Organize installation scripts +4. Standardize test organization + +### Phase 3: Professional Polish (2-4 hours) +1. Review and improve module architecture +2. Enhance documentation hierarchy +3. Add missing standard project files +4. Final professional review + +## Impact Assessment + +### Before Changes +- **First Impression**: Confused by multiple entry points and cluttered root +- **Developer Experience**: Unclear how to contribute or modify +- **Professional Credibility**: Damaged by poor organization +- **Maintenance Burden**: High due to scattered structure + +### After Changes +- **First Impression**: Clean, professional project structure +- **Developer Experience**: Clear entry points and logical organization +- **Professional Credibility**: Enhanced by following standards +- **Maintenance Burden**: Reduced through proper organization + +## Conclusion + +The FSS-Mini-RAG project has excellent technical merit but is significantly hampered by poor structural organization. The root directory pollution and multiple entry points create unnecessary complexity and damage the professional presentation. + +**Priority Recommendation:** Focus on the Critical Changes first - these will provide the most impact for professional presentation with minimal risk to functionality. + +**Timeline:** The structural issues can be resolved in 1-2 focused sessions without touching the core technical implementation, dramatically improving the project's professional appearance and maintainability. + +--- + +*Analysis completed: August 28, 2025 - FSS-Mini-RAG Project Structure Assessment* diff --git a/docs/security-analysis.md b/docs/security-analysis.md new file mode 100644 index 0000000..6647ca8 --- /dev/null +++ b/docs/security-analysis.md @@ -0,0 +1,373 @@ +# FSS-Mini-RAG Security Analysis Report +**Conducted by: Emma, Authentication Specialist** +**Date: 2024-08-28** +**Classification: Confidential - For Professional Deployment Review** + +--- + +## Executive Summary + +This comprehensive security audit examines the FSS-Mini-RAG system's defensive posture, identifying vulnerabilities and providing actionable hardening recommendations. The system demonstrates several commendable security practices but requires attention in key areas before professional deployment. + +**Overall Security Rating: MODERATE RISK (Amber)** +- โœ… **Strengths**: Good input validation patterns, secure default configurations, appropriate access controls +- โš ๏ธ **Concerns**: Network service exposure, file system access patterns, dependency management +- ๐Ÿ”ด **Critical**: Server port management and external service integration security + +--- + +## 1. Data Security & Privacy Assessment + +### Data Handling Analysis +**Status: GOOD with Minor Concerns** + +#### Positive Security Practices: +- **Local-First Architecture**: All data processing occurs locally, reducing external attack surface +- **No Cloud Dependency**: Embeddings and vector storage remain on-premise +- **Temporary File Management**: Proper cleanup patterns observed in chunking operations +- **Path Normalisation**: Robust cross-platform path handling prevents directory traversal + +#### Areas of Concern: +- **Persistent Storage**: `.mini-rag/` directories store sensitive codebase information +- **Index Files**: LanceDB vector files contain searchable representations of source code +- **Configuration Files**: YAML configs may contain sensitive connection strings +- **Memory Exposure**: Code content held in memory during processing without explicit scrubbing + +#### Recommendations: +1. **Implement data classification**: Tag sensitive files during indexing +2. **Add encryption at rest**: Encrypt vector databases and configuration files +3. **Memory management**: Explicit memory clearing after processing sensitive content +4. **Access logging**: Track who accesses which code segments through search + +--- + +## 2. Input Validation & Sanitization Assessment + +### CLI Input Handling +**Status: GOOD** + +#### Robust Validation Observed: +```python +# Path validation with proper resolution +project_path = Path(path).resolve() + +# Type checking and bounds validation +@click.option("--top-k", "-k", type=int, default=10) +@click.option("--port", type=int, default=7777) +``` + +#### File Path Security: +- **Path Traversal Protection**: Proper use of `Path().resolve()` throughout codebase +- **Extension Validation**: File type filtering based on extensions +- **Size Limits**: Appropriate file size thresholds implemented + +#### Search Query Processing: +**Status: MODERATE RISK** + +**Vulnerabilities Identified:** +- **No Query Length Limits**: Potential DoS through excessive query lengths +- **Special Character Handling**: Limited sanitization of search terms +- **Regex Injection**: Query expansion could be exploited with crafted patterns + +#### Recommendations: +1. **Implement query length limits** (max 512 characters) +2. **Sanitize search queries** before processing +3. **Validate file patterns** in include/exclude configurations +4. **Add input encoding validation** for non-ASCII content + +--- + +## 3. Network Security Assessment + +### Server Implementation Analysis +**Status: HIGH RISK - REQUIRES IMMEDIATE ATTENTION** + +#### Critical Security Issues: + +**1. Port Management Vulnerabilities:** +```python +# CRITICAL: Automatic port cleanup attempts system commands +result = subprocess.run(["netstat", "-ano"], capture_output=True, text=True) +subprocess.run(["taskkill", "//PID", pid, "//F"], check=False) +``` +**Risk**: Command injection, privilege escalation +**Impact**: System compromise possible + +**2. Network Service Exposure:** +```python +# Binds to localhost but lacks authentication +self.socket.bind(("localhost", self.port)) +self.socket.listen(5) +``` +**Risk**: Unauthorised local access +**Impact**: Code exposure to other local processes + +**3. Message Framing Vulnerabilities:** +```python +# Potential buffer overflow with untrusted length prefix +length = int.from_bytes(length_data, "big") +chunk = sock.recv(min(65536, length - len(data))) +``` +**Risk**: Memory exhaustion, DoS attacks +**Impact**: Service disruption + +#### Recommendations: +1. **Implement authentication**: Token-based access control for server connections +2. **Remove automatic process killing**: Replace with safe port checking +3. **Add connection limits**: Rate limiting and concurrent connection controls +4. **Message size validation**: Strict limits on incoming message sizes +5. **TLS encryption**: Encrypt local communications + +--- + +## 4. External Service Integration Security + +### Ollama Integration Analysis +**Status: MODERATE RISK** + +#### Security Concerns: +```python +# Unvalidated external service calls +response = requests.get(f"{self.base_url}/api/tags", timeout=5) +``` + +**Vulnerabilities:** +- **No certificate validation** for HTTPS connections +- **Trust boundary violation**: Implicit trust of Ollama responses +- **Configuration injection**: User-controlled host parameters + +#### LLM Service Security: +- **Prompt injection risks**: User queries passed directly to LLM +- **Data leakage potential**: Code content sent to external models +- **Response validation**: Limited validation of LLM outputs + +#### Recommendations: +1. **Certificate validation**: Enforce TLS certificate checking +2. **Response validation**: Sanitize and validate all external responses +3. **Connection timeouts**: Implement aggressive timeouts for external calls +4. **Host validation**: Whitelist allowed connection targets + +--- + +## 5. File System Security Assessment + +### File Access Patterns +**Status: GOOD with Recommendations** + +#### Positive Practices: +- **Appropriate file permissions**: Uses standard Python file operations +- **Pattern-based exclusions**: Sensible default exclude patterns +- **Size-based filtering**: Protection against processing oversized files + +#### Areas for Improvement: +```python +# File enumeration could be restricted further +all_files = list(project_path.rglob("*")) +``` + +#### Recommendations: +1. **Implement file access logging**: Track which files are indexed/searched +2. **Add symlink protection**: Prevent symlink-based directory traversal +3. **Enhanced file type validation**: Magic number checking beyond extensions +4. **Temporary file security**: Secure creation and cleanup of temp files + +--- + +## 6. Configuration Security Assessment + +### YAML Configuration Handling +**Status: MODERATE RISK** + +#### Security Issues: +```python +# YAML parsing without safe mode enforcement +data = yaml.safe_load(f) +``` +**Note**: Uses `safe_load` (good) but lacks validation + +#### Configuration Vulnerabilities: +- **Path injection**: User-controlled paths in configuration +- **Service endpoints**: External service URLs configurable +- **Model specifications**: Potential for malicious model references + +#### Recommendations: +1. **Configuration validation schema**: Implement strict YAML schema validation +2. **Whitelist allowed values**: Restrict configuration options to safe choices +3. **Configuration encryption**: Encrypt sensitive configuration values +4. **Read-only configurations**: Prevent runtime modification of security settings + +--- + +## 7. Dependencies & Supply Chain Security + +### Dependency Analysis +**Status: MODERATE RISK** + +#### Current Dependencies: +``` +lancedb>=0.5.0 # Vector database - moderate risk +requests>=2.28.0 # HTTP client - well-maintained +click>=8.1.0 # CLI framework - secure +PyYAML>=6.0.0 # YAML parsing - recent versions secure +``` + +#### Security Concerns: +- **Version pinning**: Uses minimum versions (>=) allowing potentially vulnerable updates +- **Transitive dependencies**: No analysis of indirect dependencies +- **Supply chain attacks**: No dependency integrity verification + +#### Recommendations: +1. **Pin exact versions**: Use `==` instead of `>=` for production deployments +2. **Dependency scanning**: Implement automated vulnerability scanning +3. **Integrity verification**: Use pip hash checking for critical dependencies +4. **Regular updates**: Establish dependency update and testing procedures + +--- + +## 8. Logging & Monitoring Security + +### Current Logging Analysis +**Status: REQUIRES IMPROVEMENT** + +#### Logging Practices: +```python +logger = logging.getLogger(__name__) +# Basic logging without security context +``` + +#### Security Gaps: +- **No security event logging**: Access attempts not recorded +- **Information leakage**: Debug logs may expose sensitive paths +- **No audit trail**: Cannot track security-relevant events +- **Log injection**: Potential for log poisoning through user inputs + +#### Recommendations: +1. **Security event logging**: Log all authentication attempts, access patterns +2. **Sanitize log inputs**: Prevent log injection attacks +3. **Structured logging**: Use structured formats for security analysis +4. **Log rotation and retention**: Implement secure log management +5. **Monitoring integration**: Connect to security monitoring systems + +--- + +## 9. System Hardening Recommendations + +### Priority 1 (Critical - Implement Immediately): + +1. **Server Authentication**: + ```python + # Add token-based authentication + def authenticate_request(self, token): + return hmac.compare_digest(token, self.expected_token) + ``` + +2. **Safe Port Management**: + ```python + # Remove dangerous subprocess calls + # Use socket.SO_REUSEADDR properly instead + ``` + +3. **Input Validation Framework**: + ```python + def validate_search_query(query: str) -> str: + if len(query) > 512: + raise ValueError("Query too long") + return re.sub(r'[^\w\s\-\.]', '', query) + ``` + +### Priority 2 (High - Implement Within Sprint): + +4. **Configuration Security**: + ```python + # Implement configuration schema validation + # Add encryption for sensitive config values + ``` + +5. **Enhanced Logging**: + ```python + # Add security event logging + security_logger.info("Search performed", extra={ + "user": user_id, + "query_hash": hashlib.sha256(query.encode()).hexdigest()[:16], + "files_accessed": len(results) + }) + ``` + +6. **Dependency Management**: + ```bash + # Pin exact versions in requirements.txt + # Implement hash checking + ``` + +### Priority 3 (Medium - Next Release Cycle): + +7. **Data Encryption**: Implement at-rest encryption for vector databases +8. **Access Controls**: Role-based access to different code segments +9. **Security Monitoring**: Integration with SIEM systems +10. **Penetration Testing**: Regular security assessments + +--- + +## 10. Compliance & Audit Considerations + +### Current Compliance Posture: +- **Data Protection**: Local storage reduces GDPR/privacy risks +- **Access Logging**: Currently insufficient for audit requirements +- **Change Management**: Git-based but lacks security change tracking +- **Documentation**: Good code documentation but missing security procedures + +### Recommendations for Compliance: +1. **Security documentation**: Create security architecture diagrams +2. **Access audit trails**: Implement comprehensive logging +3. **Regular security reviews**: Quarterly security assessments +4. **Incident response procedures**: Define security incident handling +5. **Backup security**: Secure backup and recovery procedures + +--- + +## 11. Deployment Security Checklist + +### Pre-Deployment Security Requirements: + +- [ ] **Authentication implemented** for server mode +- [ ] **Input validation** comprehensive across all entry points +- [ ] **Configuration hardening** with schema validation +- [ ] **Dependency scanning** completed and vulnerabilities addressed +- [ ] **Security logging** implemented and tested +- [ ] **TLS/encryption** for network communications +- [ ] **File system permissions** properly configured +- [ ] **Service account isolation** implemented +- [ ] **Monitoring and alerting** configured +- [ ] **Backup security** validated + +### Post-Deployment Security Monitoring: + +- [ ] **Regular vulnerability scans** scheduled +- [ ] **Log analysis** for security events +- [ ] **Dependency update procedures** established +- [ ] **Incident response plan** activated +- [ ] **Security metrics** tracked and reported + +--- + +## Conclusion + +The FSS-Mini-RAG system demonstrates solid foundational security practices with appropriate local-first architecture and sensible defaults. However, several critical vulnerabilities require immediate attention before professional deployment, particularly around server security and input validation. + +**Primary Action Items:** +1. **Implement server authentication** (Critical) +2. **Eliminate subprocess security risks** (Critical) +3. **Enhanced input validation** (High) +4. **Comprehensive security logging** (High) +5. **Dependency security hardening** (Medium) + +With these improvements, the system will achieve a **GOOD** security posture suitable for professional deployment environments. + +**Risk Acceptance**: Any deployment without addressing Critical and High priority items should require explicit risk acceptance from senior management. + +--- + +*This analysis conducted with military precision and British thoroughness. Implementation of recommendations will significantly enhance the system's defensive capabilities whilst maintaining operational effectiveness.* + +**Emma, Authentication Specialist** +**Security Clearance: OFFICIAL** diff --git a/examples/analyze_dependencies.py b/examples/analyze_dependencies.py index 7211b61..c7899ce 100644 --- a/examples/analyze_dependencies.py +++ b/examples/analyze_dependencies.py @@ -4,106 +4,110 @@ Analyze FSS-Mini-RAG dependencies to determine what's safe to remove. """ import ast -import os -from pathlib import Path from collections import defaultdict +from pathlib import Path + def find_imports_in_file(file_path): """Find all imports in a Python file.""" try: - with open(file_path, 'r', encoding='utf-8') as f: + with open(file_path, "r", encoding="utf-8") as f: content = f.read() - + tree = ast.parse(content) imports = set() - + for node in ast.walk(tree): if isinstance(node, ast.Import): for alias in node.names: - imports.add(alias.name.split('.')[0]) + imports.add(alias.name.split(".")[0]) elif isinstance(node, ast.ImportFrom): if node.module: - module = node.module.split('.')[0] + module = node.module.split(".")[0] imports.add(module) - + return imports except Exception as e: print(f"Error analyzing {file_path}: {e}") return set() + def analyze_dependencies(): """Analyze all dependencies in the project.""" project_root = Path(__file__).parent mini_rag_dir = project_root / "mini_rag" - + # Find all Python files python_files = [] for file_path in mini_rag_dir.glob("*.py"): if file_path.name != "__pycache__": python_files.append(file_path) - + # Analyze imports file_imports = {} internal_deps = defaultdict(set) - + for file_path in python_files: imports = find_imports_in_file(file_path) file_imports[file_path.name] = imports - + # Check for internal imports for imp in imports: if imp in [f.stem for f in python_files]: internal_deps[file_path.name].add(imp) - + print("๐Ÿ” FSS-Mini-RAG Dependency Analysis") print("=" * 50) - + # Show what each file imports print("\n๐Ÿ“ File Dependencies:") for filename, imports in file_imports.items(): internal = [imp for imp in imports if imp in [f.stem for f in python_files]] if internal: print(f" {filename} imports: {', '.join(internal)}") - + # Show reverse dependencies (what depends on each file) reverse_deps = defaultdict(set) for file, deps in internal_deps.items(): for dep in deps: reverse_deps[dep].add(file) - + print("\n๐Ÿ”— Reverse Dependencies (what uses each file):") all_modules = {f.stem for f in python_files} - + for module in sorted(all_modules): users = reverse_deps.get(module, set()) if users: print(f" {module}.py is used by: {', '.join(users)}") else: print(f" {module}.py is NOT imported by any other file") - + # Safety analysis print("\n๐Ÿ›ก๏ธ Safety Analysis:") - + # Files imported by __init__.py are definitely needed - init_imports = file_imports.get('__init__.py', set()) + init_imports = file_imports.get("__init__.py", set()) print(f" Core modules (imported by __init__.py): {', '.join(init_imports)}") - + # Files not used anywhere might be safe to remove unused_files = [] for module in all_modules: - if module not in reverse_deps and module != '__init__': + if module not in reverse_deps and module != "__init__": unused_files.append(module) - + if unused_files: print(f" โš ๏ธ Potentially unused: {', '.join(unused_files)}") print(" โ— Verify these aren't used by CLI or external scripts!") - + # Check CLI usage - cli_files = ['cli.py', 'enhanced_cli.py'] + cli_files = ["cli.py", "enhanced_cli.py"] for cli_file in cli_files: if cli_file in file_imports: cli_imports = file_imports[cli_file] - print(f" ๐Ÿ“‹ {cli_file} imports: {', '.join([imp for imp in cli_imports if imp in all_modules])}") + print( + f" ๐Ÿ“‹ {cli_file} imports: {', '.join([imp for imp in cli_imports if imp in all_modules])}" + ) + if __name__ == "__main__": - analyze_dependencies() \ No newline at end of file + analyze_dependencies() diff --git a/examples/basic_usage.py b/examples/basic_usage.py index 1d9d05d..cd8e379 100644 --- a/examples/basic_usage.py +++ b/examples/basic_usage.py @@ -5,64 +5,67 @@ Shows how to index a project and search it programmatically. """ from pathlib import Path -from mini_rag import ProjectIndexer, CodeSearcher, CodeEmbedder + +from mini_rag import CodeEmbedder, CodeSearcher, ProjectIndexer + def main(): # Example project path - change this to your project project_path = Path(".") # Current directory - + print("=== FSS-Mini-RAG Basic Usage Example ===") print(f"Project: {project_path}") - + # Initialize the embedding system print("\n1. Initializing embedding system...") embedder = CodeEmbedder() print(f" Using: {embedder.get_embedding_info()['method']}") - - # Initialize indexer and searcher + + # Initialize indexer and searcher indexer = ProjectIndexer(project_path, embedder) searcher = CodeSearcher(project_path, embedder) - + # Index the project print("\n2. Indexing project...") result = indexer.index_project() - + print(f" Files processed: {result.get('files_processed', 0)}") print(f" Chunks created: {result.get('chunks_created', 0)}") print(f" Time taken: {result.get('indexing_time', 0):.2f}s") - + # Get index statistics print("\n3. Index statistics:") stats = indexer.get_stats() print(f" Total files: {stats.get('total_files', 0)}") print(f" Total chunks: {stats.get('total_chunks', 0)}") print(f" Languages: {', '.join(stats.get('languages', []))}") - + # Example searches queries = [ "chunker function", - "embedding system", + "embedding system", "search implementation", "file watcher", - "error handling" + "error handling", ] - + print("\n4. Example searches:") for query in queries: print(f"\n Query: '{query}'") results = searcher.search(query, top_k=3) - + if results: for i, result in enumerate(results, 1): print(f" {i}. {result.file_path.name} (score: {result.score:.3f})") print(f" Type: {result.chunk_type}") # Show first 60 characters of content - content_preview = result.content.replace('\n', ' ')[:60] + content_preview = result.content.replace("\n", " ")[:60] print(f" Preview: {content_preview}...") else: print(" No results found") - + print("\n=== Example Complete ===") + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/smart_config_suggestions.py b/examples/smart_config_suggestions.py index c9620e5..26fe009 100644 --- a/examples/smart_config_suggestions.py +++ b/examples/smart_config_suggestions.py @@ -5,102 +5,108 @@ Analyzes the indexed data to suggest optimal settings. """ import json -from pathlib import Path -from collections import defaultdict, Counter import sys +from collections import Counter +from pathlib import Path + def analyze_project_patterns(manifest_path: Path): """Analyze project patterns and suggest optimizations.""" - + with open(manifest_path) as f: manifest = json.load(f) - - files = manifest.get('files', {}) - + + files = manifest.get("files", {}) + print("๐Ÿ” FSS-Mini-RAG Smart Tuning Analysis") print("=" * 50) - + # Analyze file types and chunking efficiency languages = Counter() chunk_efficiency = [] large_files = [] small_files = [] - + for filepath, info in files.items(): - lang = info.get('language', 'unknown') + lang = info.get("language", "unknown") languages[lang] += 1 - - size = info.get('size', 0) - chunks = info.get('chunks', 1) - + + size = info.get("size", 0) + chunks = info.get("chunks", 1) + chunk_efficiency.append(chunks / max(1, size / 1000)) # chunks per KB - + if size > 10000: # >10KB large_files.append((filepath, size, chunks)) elif size < 500: # <500B small_files.append((filepath, size, chunks)) - + # Analysis results total_files = len(files) - total_chunks = sum(info.get('chunks', 1) for info in files.values()) + total_chunks = sum(info.get("chunks", 1) for info in files.values()) avg_chunks_per_file = total_chunks / max(1, total_files) - - print(f"๐Ÿ“Š Current Stats:") + + print("๐Ÿ“Š Current Stats:") print(f" Files: {total_files}") print(f" Chunks: {total_chunks}") print(f" Avg chunks/file: {avg_chunks_per_file:.1f}") - - print(f"\n๐Ÿ—‚๏ธ Language Distribution:") + + print("\n๐Ÿ—‚๏ธ Language Distribution:") for lang, count in languages.most_common(10): pct = 100 * count / total_files print(f" {lang}: {count} files ({pct:.1f}%)") - - print(f"\n๐Ÿ’ก Smart Optimization Suggestions:") - + + print("\n๐Ÿ’ก Smart Optimization Suggestions:") + # Suggestion 1: Language-specific chunking - if languages['python'] > 10: - print(f"โœจ Python Optimization:") - print(f" - Use function-level chunking (detected {languages['python']} Python files)") - print(f" - Increase chunk size to 3000 chars for Python (better context)") - - if languages['markdown'] > 5: - print(f"โœจ Markdown Optimization:") + if languages["python"] > 10: + print("โœจ Python Optimization:") + print( + f" - Use function-level chunking (detected {languages['python']} Python files)" + ) + print(" - Increase chunk size to 3000 chars for Python (better context)") + + if languages["markdown"] > 5: + print("โœจ Markdown Optimization:") print(f" - Use header-based chunking (detected {languages['markdown']} MD files)") - print(f" - Keep sections together for better search relevance") - - if languages['json'] > 20: - print(f"โœจ JSON Optimization:") + print(" - Keep sections together for better search relevance") + + if languages["json"] > 20: + print("โœจ JSON Optimization:") print(f" - Consider object-level chunking (detected {languages['json']} JSON files)") - print(f" - Might want to exclude large config JSONs") - + print(" - Might want to exclude large config JSONs") + # Suggestion 2: File size optimization if large_files: - print(f"\n๐Ÿ“ˆ Large File Optimization:") + print("\n๐Ÿ“ˆ Large File Optimization:") print(f" Found {len(large_files)} files >10KB:") - for filepath, size, chunks in sorted(large_files, key=lambda x: x[1], reverse=True)[:3]: + for filepath, size, chunks in sorted(large_files, key=lambda x: x[1], reverse=True)[ + :3 + ]: kb = size / 1024 print(f" - {filepath}: {kb:.1f}KB โ†’ {chunks} chunks") if len(large_files) > 5: - print(f" ๐Ÿ’ก Consider streaming threshold: 5KB (current: 1MB)") - + print(" ๐Ÿ’ก Consider streaming threshold: 5KB (current: 1MB)") + if small_files and len(small_files) > total_files * 0.3: - print(f"\n๐Ÿ“‰ Small File Optimization:") + print("\n๐Ÿ“‰ Small File Optimization:") print(f" {len(small_files)} files <500B might not need chunking") - print(f" ๐Ÿ’ก Consider: combine small files or skip tiny ones") - + print(" ๐Ÿ’ก Consider: combine small files or skip tiny ones") + # Suggestion 3: Search optimization avg_efficiency = sum(chunk_efficiency) / len(chunk_efficiency) - print(f"\n๐Ÿ” Search Optimization:") + print("\n๐Ÿ” Search Optimization:") if avg_efficiency < 0.5: - print(f" ๐Ÿ’ก Chunks are large relative to files - consider smaller chunks") + print(" ๐Ÿ’ก Chunks are large relative to files - consider smaller chunks") print(f" ๐Ÿ’ก Current: {avg_chunks_per_file:.1f} chunks/file, try 2-3 chunks/file") elif avg_efficiency > 2: - print(f" ๐Ÿ’ก Many small chunks - consider larger chunk size") - print(f" ๐Ÿ’ก Reduce chunk overhead with 2000-4000 char chunks") - + print(" ๐Ÿ’ก Many small chunks - consider larger chunk size") + print(" ๐Ÿ’ก Reduce chunk overhead with 2000-4000 char chunks") + # Suggestion 4: Smart defaults - print(f"\nโš™๏ธ Recommended Config Updates:") - print(f"""{{ + print("\nโš™๏ธ Recommended Config Updates:") + print( + """{{ "chunking": {{ "max_size": {3000 if languages['python'] > languages['markdown'] else 2000}, "min_size": 200, @@ -115,16 +121,18 @@ def analyze_project_patterns(manifest_path: Path): "skip_small_files": {500 if len(small_files) > total_files * 0.3 else 0}, "streaming_threshold_kb": {5 if len(large_files) > 5 else 1024} }} -}}""") +}}""" + ) + if __name__ == "__main__": if len(sys.argv) != 2: print("Usage: python smart_config_suggestions.py ") sys.exit(1) - + manifest_path = Path(sys.argv[1]) if not manifest_path.exists(): print(f"Manifest not found: {manifest_path}") sys.exit(1) - - analyze_project_patterns(manifest_path) \ No newline at end of file + + analyze_project_patterns(manifest_path) diff --git a/mini_rag/__init__.py b/mini_rag/__init__.py index 2db7d2e..60dd5ca 100644 --- a/mini_rag/__init__.py +++ b/mini_rag/__init__.py @@ -7,30 +7,16 @@ Designed for portability, efficiency, and simplicity across projects and compute __version__ = "2.1.0" -from .ollama_embeddings import OllamaEmbedder as CodeEmbedder from .chunker import CodeChunker from .indexer import ProjectIndexer +from .ollama_embeddings import OllamaEmbedder as CodeEmbedder from .search import CodeSearcher from .watcher import FileWatcher -# Auto-update system (graceful import for legacy versions) -try: - from .updater import UpdateChecker, check_for_updates, get_updater - __all__ = [ - "CodeEmbedder", - "CodeChunker", - "ProjectIndexer", - "CodeSearcher", - "FileWatcher", - "UpdateChecker", - "check_for_updates", - "get_updater", - ] -except ImportError: - __all__ = [ - "CodeEmbedder", - "CodeChunker", - "ProjectIndexer", - "CodeSearcher", - "FileWatcher", - ] \ No newline at end of file +__all__ = [ + "CodeEmbedder", + "CodeChunker", + "ProjectIndexer", + "CodeSearcher", + "FileWatcher", +] diff --git a/mini_rag/__main__.py b/mini_rag/__main__.py index 77409db..ec65efd 100644 --- a/mini_rag/__main__.py +++ b/mini_rag/__main__.py @@ -2,5 +2,5 @@ from .cli import cli -if __name__ == '__main__': - cli() \ No newline at end of file +if __name__ == "__main__": + cli() diff --git a/mini_rag/auto_optimizer.py b/mini_rag/auto_optimizer.py index 3477d9e..dcc8b02 100644 --- a/mini_rag/auto_optimizer.py +++ b/mini_rag/auto_optimizer.py @@ -3,194 +3,188 @@ Auto-optimizer for FSS-Mini-RAG. Automatically tunes settings based on usage patterns. """ -from pathlib import Path import json -from typing import Dict, Any, List -from collections import Counter import logging +from collections import Counter +from pathlib import Path +from typing import Any, Dict logger = logging.getLogger(__name__) + class AutoOptimizer: """Automatically optimizes RAG settings based on project patterns.""" - + def __init__(self, project_path: Path): self.project_path = project_path - self.rag_dir = project_path / '.mini-rag' - self.config_path = self.rag_dir / 'config.json' - self.manifest_path = self.rag_dir / 'manifest.json' - + self.rag_dir = project_path / ".mini-rag" + self.config_path = self.rag_dir / "config.json" + self.manifest_path = self.rag_dir / "manifest.json" + def analyze_and_optimize(self) -> Dict[str, Any]: """Analyze current patterns and auto-optimize settings.""" - + if not self.manifest_path.exists(): return {"error": "No index found - run indexing first"} - + # Load current data with open(self.manifest_path) as f: manifest = json.load(f) - + # Analyze patterns analysis = self._analyze_patterns(manifest) - + # Generate optimizations optimizations = self._generate_optimizations(analysis) - + # Apply optimizations if beneficial - if optimizations['confidence'] > 0.7: + if optimizations["confidence"] > 0.7: self._apply_optimizations(optimizations) return { "status": "optimized", - "changes": optimizations['changes'], - "expected_improvement": optimizations['expected_improvement'] + "changes": optimizations["changes"], + "expected_improvement": optimizations["expected_improvement"], } else: return { "status": "no_changes_needed", "analysis": analysis, - "confidence": optimizations['confidence'] + "confidence": optimizations["confidence"], } - + def _analyze_patterns(self, manifest: Dict[str, Any]) -> Dict[str, Any]: """Analyze current indexing patterns.""" - files = manifest.get('files', {}) - + files = manifest.get("files", {}) + # Language distribution languages = Counter() sizes = [] chunk_ratios = [] - + for filepath, info in files.items(): - lang = info.get('language', 'unknown') + lang = info.get("language", "unknown") languages[lang] += 1 - - size = info.get('size', 0) - chunks = info.get('chunks', 1) - + + size = info.get("size", 0) + chunks = info.get("chunks", 1) + sizes.append(size) chunk_ratios.append(chunks / max(1, size / 1000)) # chunks per KB - + avg_chunk_ratio = sum(chunk_ratios) / len(chunk_ratios) if chunk_ratios else 1 avg_size = sum(sizes) / len(sizes) if sizes else 1000 - + return { - 'languages': dict(languages.most_common()), - 'total_files': len(files), - 'total_chunks': sum(info.get('chunks', 1) for info in files.values()), - 'avg_chunk_ratio': avg_chunk_ratio, - 'avg_file_size': avg_size, - 'large_files': sum(1 for s in sizes if s > 10000), - 'small_files': sum(1 for s in sizes if s < 500) + "languages": dict(languages.most_common()), + "total_files": len(files), + "total_chunks": sum(info.get("chunks", 1) for info in files.values()), + "avg_chunk_ratio": avg_chunk_ratio, + "avg_file_size": avg_size, + "large_files": sum(1 for s in sizes if s > 10000), + "small_files": sum(1 for s in sizes if s < 500), } - + def _generate_optimizations(self, analysis: Dict[str, Any]) -> Dict[str, Any]: """Generate optimization recommendations.""" changes = [] confidence = 0.5 expected_improvement = 0 - + # Optimize chunking based on dominant language - languages = analysis['languages'] + languages = analysis["languages"] if languages: dominant_lang, count = list(languages.items())[0] - lang_pct = count / analysis['total_files'] - + lang_pct = count / analysis["total_files"] + if lang_pct > 0.3: # Dominant language >30% - if dominant_lang == 'python' and analysis['avg_chunk_ratio'] < 1.5: - changes.append("Increase Python chunk size to 3000 for better function context") + if dominant_lang == "python" and analysis["avg_chunk_ratio"] < 1.5: + changes.append( + "Increase Python chunk size to 3000 for better function context" + ) confidence += 0.2 expected_improvement += 15 - - elif dominant_lang == 'markdown' and analysis['avg_chunk_ratio'] < 1.2: + + elif dominant_lang == "markdown" and analysis["avg_chunk_ratio"] < 1.2: changes.append("Use header-based chunking for Markdown files") confidence += 0.15 expected_improvement += 10 - + # Optimize for large files - if analysis['large_files'] > 5: + if analysis["large_files"] > 5: changes.append("Reduce streaming threshold to 5KB for better large file handling") confidence += 0.1 expected_improvement += 8 - + # Optimize chunk ratio - if analysis['avg_chunk_ratio'] < 1.0: + if analysis["avg_chunk_ratio"] < 1.0: changes.append("Reduce chunk size for more granular search results") confidence += 0.15 expected_improvement += 12 - elif analysis['avg_chunk_ratio'] > 3.0: + elif analysis["avg_chunk_ratio"] > 3.0: changes.append("Increase chunk size to reduce overhead") confidence += 0.1 expected_improvement += 5 - + # Skip tiny files optimization - small_file_pct = analysis['small_files'] / analysis['total_files'] + small_file_pct = analysis["small_files"] / analysis["total_files"] if small_file_pct > 0.3: changes.append("Skip files smaller than 300 bytes to improve focus") confidence += 0.1 expected_improvement += 3 - + return { - 'changes': changes, - 'confidence': min(confidence, 1.0), - 'expected_improvement': expected_improvement + "changes": changes, + "confidence": min(confidence, 1.0), + "expected_improvement": expected_improvement, } - + def _apply_optimizations(self, optimizations: Dict[str, Any]): """Apply the recommended optimizations.""" - + # Load existing config or create default if self.config_path.exists(): with open(self.config_path) as f: config = json.load(f) else: config = self._get_default_config() - - changes = optimizations['changes'] - + + changes = optimizations["changes"] + # Apply changes based on recommendations for change in changes: if "Python chunk size to 3000" in change: - config.setdefault('chunking', {})['max_size'] = 3000 - + config.setdefault("chunking", {})["max_size"] = 3000 + elif "header-based chunking" in change: - config.setdefault('chunking', {})['strategy'] = 'header' - + config.setdefault("chunking", {})["strategy"] = "header" + elif "streaming threshold to 5KB" in change: - config.setdefault('streaming', {})['threshold_bytes'] = 5120 - + config.setdefault("streaming", {})["threshold_bytes"] = 5120 + elif "Reduce chunk size" in change: - current_size = config.get('chunking', {}).get('max_size', 2000) - config.setdefault('chunking', {})['max_size'] = max(1500, current_size - 500) - + current_size = config.get("chunking", {}).get("max_size", 2000) + config.setdefault("chunking", {})["max_size"] = max(1500, current_size - 500) + elif "Increase chunk size" in change: - current_size = config.get('chunking', {}).get('max_size', 2000) - config.setdefault('chunking', {})['max_size'] = min(4000, current_size + 500) - + current_size = config.get("chunking", {}).get("max_size", 2000) + config.setdefault("chunking", {})["max_size"] = min(4000, current_size + 500) + elif "Skip files smaller" in change: - config.setdefault('files', {})['min_file_size'] = 300 - + config.setdefault("files", {})["min_file_size"] = 300 + # Save optimized config - config['_auto_optimized'] = True - config['_optimization_timestamp'] = json.dumps(None, default=str) - - with open(self.config_path, 'w') as f: + config["_auto_optimized"] = True + config["_optimization_timestamp"] = json.dumps(None, default=str) + + with open(self.config_path, "w") as f: json.dump(config, f, indent=2) - + logger.info(f"Applied {len(changes)} optimizations to {self.config_path}") - + def _get_default_config(self) -> Dict[str, Any]: """Get default configuration.""" return { - "chunking": { - "max_size": 2000, - "min_size": 150, - "strategy": "semantic" - }, - "streaming": { - "enabled": True, - "threshold_bytes": 1048576 - }, - "files": { - "min_file_size": 50 - } - } \ No newline at end of file + "chunking": {"max_size": 2000, "min_size": 150, "strategy": "semantic"}, + "streaming": {"enabled": True, "threshold_bytes": 1048576}, + "files": {"min_file_size": 50}, + } diff --git a/mini_rag/chunker.py b/mini_rag/chunker.py index 635d2bb..60d6dd4 100644 --- a/mini_rag/chunker.py +++ b/mini_rag/chunker.py @@ -4,37 +4,41 @@ Chunks by functions, classes, and logical boundaries instead of arbitrary lines. """ import ast -import re -from typing import List, Dict, Any, Optional, Tuple -from pathlib import Path import logging +import re +from pathlib import Path +from typing import Any, Dict, List, Optional logger = logging.getLogger(__name__) class CodeChunk: """Represents a logical chunk of code.""" - - def __init__(self, - content: str, - file_path: str, - start_line: int, - end_line: int, - chunk_type: str, - name: Optional[str] = None, - language: str = "python", - file_lines: Optional[int] = None, - chunk_index: Optional[int] = None, - total_chunks: Optional[int] = None, - parent_class: Optional[str] = None, - parent_function: Optional[str] = None, - prev_chunk_id: Optional[str] = None, - next_chunk_id: Optional[str] = None): + + def __init__( + self, + content: str, + file_path: str, + start_line: int, + end_line: int, + chunk_type: str, + name: Optional[str] = None, + language: str = "python", + file_lines: Optional[int] = None, + chunk_index: Optional[int] = None, + total_chunks: Optional[int] = None, + parent_class: Optional[str] = None, + parent_function: Optional[str] = None, + prev_chunk_id: Optional[str] = None, + next_chunk_id: Optional[str] = None, + ): self.content = content self.file_path = file_path self.start_line = start_line self.end_line = end_line - self.chunk_type = chunk_type # 'function', 'class', 'method', 'module', 'module_header' + self.chunk_type = ( + chunk_type # 'function', 'class', 'method', 'module', 'module_header' + ) self.name = name self.language = language # New metadata fields @@ -45,42 +49,47 @@ class CodeChunk: self.parent_function = parent_function # For nested functions self.prev_chunk_id = prev_chunk_id # Link to previous chunk self.next_chunk_id = next_chunk_id # Link to next chunk - + def to_dict(self) -> Dict[str, Any]: """Convert to dictionary for storage.""" return { - 'content': self.content, - 'file_path': self.file_path, - 'start_line': self.start_line, - 'end_line': self.end_line, - 'chunk_type': self.chunk_type, - 'name': self.name, - 'language': self.language, - 'num_lines': self.end_line - self.start_line + 1, + "content": self.content, + "file_path": self.file_path, + "start_line": self.start_line, + "end_line": self.end_line, + "chunk_type": self.chunk_type, + "name": self.name, + "language": self.language, + "num_lines": self.end_line - self.start_line + 1, # Include new metadata if available - 'file_lines': self.file_lines, - 'chunk_index': self.chunk_index, - 'total_chunks': self.total_chunks, - 'parent_class': self.parent_class, - 'parent_function': self.parent_function, - 'prev_chunk_id': self.prev_chunk_id, - 'next_chunk_id': self.next_chunk_id, + "file_lines": self.file_lines, + "chunk_index": self.chunk_index, + "total_chunks": self.total_chunks, + "parent_class": self.parent_class, + "parent_function": self.parent_function, + "prev_chunk_id": self.prev_chunk_id, + "next_chunk_id": self.next_chunk_id, } - + def __repr__(self): - return f"CodeChunk({self.chunk_type}:{self.name} in {self.file_path}:{self.start_line}-{self.end_line})" + return ( + f"CodeChunk({self.chunk_type}:{self.name} " + f"in {self.file_path}:{self.start_line}-{self.end_line})" + ) class CodeChunker: """Intelligently chunks code files based on language and structure.""" - - def __init__(self, - max_chunk_size: int = 1000, - min_chunk_size: int = 50, - overlap_lines: int = 0): + + def __init__( + self, + max_chunk_size: int = 1000, + min_chunk_size: int = 50, + overlap_lines: int = 0, + ): """ Initialize chunker with size constraints. - + Args: max_chunk_size: Maximum lines per chunk min_chunk_size: Minimum lines per chunk @@ -89,83 +98,83 @@ class CodeChunker: self.max_chunk_size = max_chunk_size self.min_chunk_size = min_chunk_size self.overlap_lines = overlap_lines - + # Language detection patterns self.language_patterns = { - '.py': 'python', - '.js': 'javascript', - '.jsx': 'javascript', - '.ts': 'typescript', - '.tsx': 'typescript', - '.go': 'go', - '.java': 'java', - '.cpp': 'cpp', - '.c': 'c', - '.cs': 'csharp', - '.rs': 'rust', - '.rb': 'ruby', - '.php': 'php', - '.swift': 'swift', - '.kt': 'kotlin', - '.scala': 'scala', + ".py": "python", + ".js": "javascript", + ".jsx": "javascript", + ".ts": "typescript", + ".tsx": "typescript", + ".go": "go", + ".java": "java", + ".cpp": "cpp", + ".c": "c", + ".cs": "csharp", + ".rs": "rust", + ".rb": "ruby", + ".php": "php", + ".swift": "swift", + ".kt": "kotlin", + ".scala": "scala", # Documentation formats - '.md': 'markdown', - '.markdown': 'markdown', - '.rst': 'restructuredtext', - '.txt': 'text', - '.adoc': 'asciidoc', - '.asciidoc': 'asciidoc', + ".md": "markdown", + ".markdown": "markdown", + ".rst": "restructuredtext", + ".txt": "text", + ".adoc": "asciidoc", + ".asciidoc": "asciidoc", # Config formats - '.json': 'json', - '.yaml': 'yaml', - '.yml': 'yaml', - '.toml': 'toml', - '.ini': 'ini', - '.xml': 'xml', - '.conf': 'config', - '.config': 'config', + ".json": "json", + ".yaml": "yaml", + ".yml": "yaml", + ".toml": "toml", + ".ini": "ini", + ".xml": "xml", + ".con": "config", + ".config": "config", } - + def chunk_file(self, file_path: Path, content: Optional[str] = None) -> List[CodeChunk]: """ Chunk a code file intelligently based on its language. - + Args: file_path: Path to the file content: Optional content (if not provided, will read from file) - + Returns: List of CodeChunk objects """ if content is None: try: - content = file_path.read_text(encoding='utf-8') + content = file_path.read_text(encoding="utf-8") except Exception as e: logger.error(f"Failed to read {file_path}: {e}") return [] - + # Get total lines for metadata lines = content.splitlines() total_lines = len(lines) - + # Detect language language = self._detect_language(file_path, content) - + # Choose chunking strategy based on language chunks = [] - + try: - if language == 'python': + if language == "python": chunks = self._chunk_python(content, str(file_path)) - elif language in ['javascript', 'typescript']: + elif language in ["javascript", "typescript"]: chunks = self._chunk_javascript(content, str(file_path), language) - elif language == 'go': + elif language == "go": chunks = self._chunk_go(content, str(file_path)) - elif language == 'java': + elif language == "java": chunks = self._chunk_java(content, str(file_path)) - elif language in ['markdown', 'text', 'restructuredtext', 'asciidoc']: + elif language in ["markdown", "text", "restructuredtext", "asciidoc"]: chunks = self._chunk_markdown(content, str(file_path), language) - elif language in ['json', 'yaml', 'toml', 'ini', 'xml', 'config']: + elif language in ["json", "yaml", "toml", "ini", "xml", "config"]: chunks = self._chunk_config(content, str(file_path), language) else: # Fallback to generic chunking @@ -173,258 +182,290 @@ class CodeChunker: except Exception as e: logger.warning(f"Failed to chunk {file_path} with language-specific chunker: {e}") chunks = self._chunk_generic(content, str(file_path), language) - + # Ensure chunks meet size constraints chunks = self._enforce_size_constraints(chunks) - + # Set chunk links and indices for all chunks if chunks: for chunk in chunks: if chunk.file_lines is None: chunk.file_lines = total_lines chunks = self._set_chunk_links(chunks, str(file_path)) - + return chunks - + def _detect_language(self, file_path: Path, content: str = None) -> str: """Detect programming language from file extension and content.""" # First try extension-based detection suffix = file_path.suffix.lower() if suffix in self.language_patterns: return self.language_patterns[suffix] - + # Fallback to content-based detection if content is None: try: - content = file_path.read_text(encoding='utf-8') - except: - return 'unknown' - + content = file_path.read_text(encoding="utf-8") + except (UnicodeDecodeError, OSError, IOError): + return "unknown" + # Check for shebang lines = content.splitlines() - if lines and lines[0].startswith('#!'): + if lines and lines[0].startswith("#!"): shebang = lines[0].lower() - if 'python' in shebang: - return 'python' - elif 'node' in shebang or 'javascript' in shebang: - return 'javascript' - elif 'bash' in shebang or 'sh' in shebang: - return 'bash' - + if "python" in shebang: + return "python" + elif "node" in shebang or "javascript" in shebang: + return "javascript" + elif "bash" in shebang or "sh" in shebang: + return "bash" + # Check for Python-specific patterns in first 50 lines sample_lines = lines[:50] - sample_text = '\n'.join(sample_lines) - + sample_text = "\n".join(sample_lines) + python_indicators = [ - 'import ', 'from ', 'def ', 'class ', 'if __name__', - 'print(', 'len(', 'range(', 'str(', 'int(', 'float(', - 'self.', '__init__', '__main__', 'Exception:', 'try:', 'except:' + "import ", + "from ", + "def ", + "class ", + "if __name__", + "print(", + "len(", + "range(", + "str(", + "int(", + "float(", + "self.", + "__init__", + "__main__", + "Exception:", + "try:", + "except:", ] - + python_score = sum(1 for indicator in python_indicators if indicator in sample_text) - + # If we find strong Python indicators, classify as Python if python_score >= 3: - return 'python' - + return "python" + # Check for other languages - if any(indicator in sample_text for indicator in ['function ', 'var ', 'const ', 'let ', '=>']): - return 'javascript' - - return 'unknown' - + if any( + indicator in sample_text + for indicator in ["function ", "var ", "const ", "let ", "=>"] + ): + return "javascript" + + return "unknown" + def _chunk_python(self, content: str, file_path: str) -> List[CodeChunk]: """Chunk Python code using AST with enhanced function/class extraction.""" chunks = [] lines = content.splitlines() total_lines = len(lines) - + try: tree = ast.parse(content) except SyntaxError as e: logger.warning(f"Syntax error in {file_path}: {e}") return self._chunk_python_fallback(content, file_path) - + # Extract all functions and classes with their metadata extracted_items = self._extract_python_items(tree, lines) - + # If we found functions/classes, create chunks for them if extracted_items: - chunks = self._create_chunks_from_items(extracted_items, lines, file_path, total_lines) - + chunks = self._create_chunks_from_items( + extracted_items, lines, file_path, total_lines + ) + # If no chunks or very few chunks from a large file, add fallback chunks if len(chunks) < 3 and total_lines > 200: fallback_chunks = self._chunk_python_fallback(content, file_path) # Merge with existing chunks, avoiding duplicates chunks = self._merge_chunks(chunks, fallback_chunks) - + return chunks or self._chunk_python_fallback(content, file_path) - + def _extract_python_items(self, tree: ast.AST, lines: List[str]) -> List[Dict]: """Extract all functions and classes with metadata.""" items = [] - + class ItemExtractor(ast.NodeVisitor): + def __init__(self): self.class_stack = [] # Track nested classes self.function_stack = [] # Track nested functions - + def visit_ClassDef(self, node): self.class_stack.append(node.name) - + # Extract class info item = { - 'type': 'class', - 'name': node.name, - 'start_line': node.lineno, - 'end_line': node.end_lineno or len(lines), - 'parent_class': self.class_stack[-2] if len(self.class_stack) > 1 else None, - 'decorators': [d.id for d in node.decorator_list if hasattr(d, 'id')], - 'methods': [] + "type": "class", + "name": node.name, + "start_line": node.lineno, + "end_line": node.end_lineno or len(lines), + "parent_class": ( + self.class_stack[-2] if len(self.class_stack) > 1 else None + ), + "decorators": [d.id for d in node.decorator_list if hasattr(d, "id")], + "methods": [], } - + # Find methods in this class for child in node.body: if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)): - item['methods'].append(child.name) - + item["methods"].append(child.name) + items.append(item) - + self.generic_visit(node) self.class_stack.pop() - + def visit_FunctionDef(self, node): - self._visit_function(node, 'function') - + self._visit_function(node, "function") + def visit_AsyncFunctionDef(self, node): - self._visit_function(node, 'async_function') - + self._visit_function(node, "async_function") + def _visit_function(self, node, func_type): self.function_stack.append(node.name) - + # Extract function info item = { - 'type': func_type, - 'name': node.name, - 'start_line': node.lineno, - 'end_line': node.end_lineno or len(lines), - 'parent_class': self.class_stack[-1] if self.class_stack else None, - 'parent_function': self.function_stack[-2] if len(self.function_stack) > 1 else None, - 'decorators': [d.id for d in node.decorator_list if hasattr(d, 'id')], - 'args': [arg.arg for arg in node.args.args], - 'is_method': bool(self.class_stack) + "type": func_type, + "name": node.name, + "start_line": node.lineno, + "end_line": node.end_lineno or len(lines), + "parent_class": self.class_stack[-1] if self.class_stack else None, + "parent_function": ( + self.function_stack[-2] if len(self.function_stack) > 1 else None + ), + "decorators": [d.id for d in node.decorator_list if hasattr(d, "id")], + "args": [arg.arg for arg in node.args.args], + "is_method": bool(self.class_stack), } - + items.append(item) - + self.generic_visit(node) self.function_stack.pop() - + extractor = ItemExtractor() extractor.visit(tree) - + # Sort items by line number - items.sort(key=lambda x: x['start_line']) - + items.sort(key=lambda x: x["start_line"]) + return items - - def _create_chunks_from_items(self, items: List[Dict], lines: List[str], file_path: str, total_lines: int) -> List[CodeChunk]: + + def _create_chunks_from_items( + self, items: List[Dict], lines: List[str], file_path: str, total_lines: int + ) -> List[CodeChunk]: """Create chunks from extracted AST items.""" chunks = [] - + for item in items: - start_line = item['start_line'] - 1 # Convert to 0-based - end_line = min(item['end_line'], len(lines)) - 1 # Convert to 0-based - - chunk_content = '\n'.join(lines[start_line:end_line + 1]) - + start_line = item["start_line"] - 1 # Convert to 0-based + end_line = min(item["end_line"], len(lines)) - 1 # Convert to 0-based + + chunk_content = "\n".join(lines[start_line : end_line + 1]) + chunk = CodeChunk( content=chunk_content, file_path=file_path, start_line=start_line + 1, end_line=end_line + 1, - chunk_type=item['type'], - name=item['name'], - language='python', - parent_class=item.get('parent_class'), - parent_function=item.get('parent_function'), - file_lines=total_lines + chunk_type=item["type"], + name=item["name"], + language="python", + parent_class=item.get("parent_class"), + parent_function=item.get("parent_function"), + file_lines=total_lines, ) chunks.append(chunk) - + return chunks - + def _chunk_python_fallback(self, content: str, file_path: str) -> List[CodeChunk]: """Fallback chunking for Python files with syntax errors or no AST items.""" chunks = [] lines = content.splitlines() - + # Use regex to find function/class definitions patterns = [ - (r'^(class\s+\w+.*?:)', 'class'), - (r'^(def\s+\w+.*?:)', 'function'), - (r'^(async\s+def\s+\w+.*?:)', 'async_function'), + (r"^(class\s+\w+.*?:)", "class"), + (r"^(def\s+\w+.*?:)", "function"), + (r"^(async\s+def\s+\w+.*?:)", "async_function"), ] - + matches = [] for i, line in enumerate(lines): for pattern, item_type in patterns: if re.match(pattern, line.strip()): # Extract name - if item_type == 'class': - name_match = re.match(r'class\s+(\w+)', line.strip()) + if item_type == "class": + name_match = re.match(r"class\s+(\w+)", line.strip()) else: - name_match = re.match(r'(?:async\s+)?def\s+(\w+)', line.strip()) - + name_match = re.match(r"(?:async\s+)?def\s+(\w+)", line.strip()) + if name_match: - matches.append({ - 'line': i, - 'type': item_type, - 'name': name_match.group(1), - 'indent': len(line) - len(line.lstrip()) - }) - + matches.append( + { + "line": i, + "type": item_type, + "name": name_match.group(1), + "indent": len(line) - len(line.lstrip()), + } + ) + # Create chunks from matches for i, match in enumerate(matches): - start_line = match['line'] - + start_line = match["line"] + # Find end line by looking for next item at same or lower indentation end_line = len(lines) - 1 - base_indent = match['indent'] - + base_indent = match["indent"] + for j in range(start_line + 1, len(lines)): line = lines[j] if line.strip() and len(line) - len(line.lstrip()) <= base_indent: # Found next item at same or lower level end_line = j - 1 break - + # Create chunk - chunk_content = '\n'.join(lines[start_line:end_line + 1]) + chunk_content = "\n".join(lines[start_line : end_line + 1]) if chunk_content.strip(): - chunks.append(CodeChunk( - content=chunk_content, - file_path=file_path, - start_line=start_line + 1, - end_line=end_line + 1, - chunk_type=match['type'], - name=match['name'], - language='python' - )) - + chunks.append( + CodeChunk( + content=chunk_content, + file_path=file_path, + start_line=start_line + 1, + end_line=end_line + 1, + chunk_type=match["type"], + name=match["name"], + language="python", + ) + ) + return chunks - - def _merge_chunks(self, primary_chunks: List[CodeChunk], fallback_chunks: List[CodeChunk]) -> List[CodeChunk]: + + def _merge_chunks( + self, primary_chunks: List[CodeChunk], fallback_chunks: List[CodeChunk] + ) -> List[CodeChunk]: """Merge chunks, avoiding duplicates.""" if not primary_chunks: return fallback_chunks if not fallback_chunks: return primary_chunks - + # Simple merge - just add fallback chunks that don't overlap with primary merged = primary_chunks[:] primary_ranges = [(chunk.start_line, chunk.end_line) for chunk in primary_chunks] - + for fallback_chunk in fallback_chunks: # Check if this fallback chunk overlaps with any primary chunk overlaps = False @@ -432,136 +473,162 @@ class CodeChunker: if not (fallback_chunk.end_line < start or fallback_chunk.start_line > end): overlaps = True break - + if not overlaps: merged.append(fallback_chunk) - + # Sort by start line merged.sort(key=lambda x: x.start_line) return merged - - def _process_python_class(self, node: ast.ClassDef, lines: List[str], file_path: str, total_lines: int) -> List[CodeChunk]: + + def _process_python_class( + self, node: ast.ClassDef, lines: List[str], file_path: str, total_lines: int + ) -> List[CodeChunk]: """Process a Python class with smart chunking.""" chunks = [] - + # Get class definition line class_start = node.lineno - 1 - class_end = node.end_lineno or len(lines) - + # Find where class docstring ends docstring_end = class_start class_docstring = ast.get_docstring(node) if class_docstring and node.body: first_stmt = node.body[0] - if isinstance(first_stmt, ast.Expr) and isinstance(first_stmt.value, (ast.Str, ast.Constant)): + if isinstance(first_stmt, ast.Expr) and isinstance( + first_stmt.value, (ast.Str, ast.Constant) + ): docstring_end = first_stmt.end_lineno - 1 - + # Find __init__ method if exists init_method = None init_end = docstring_end for child in node.body: - if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name == '__init__': + if ( + isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) + and child.name == "__init__" + ): init_method = child init_end = child.end_lineno - 1 break - + # Collect method signatures for preview method_signatures = [] for child in node.body: - if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name != '__init__': + if ( + isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) + and child.name != "__init__" + ): # Get just the method signature line sig_line = lines[child.lineno - 1].strip() method_signatures.append(f" # {sig_line}") - + # Create class header chunk: class def + docstring + __init__ + method preview header_lines = [] - + # Add class definition and docstring if init_method: - header_lines = lines[class_start:init_end + 1] + header_lines = lines[class_start : init_end + 1] else: - header_lines = lines[class_start:docstring_end + 1] - + header_lines = lines[class_start : docstring_end + 1] + # Add method signature preview if we have methods if method_signatures: - header_content = '\n'.join(header_lines) - if not header_content.rstrip().endswith(':'): - header_content += '\n' - header_content += '\n # Method signatures:\n' + '\n'.join(method_signatures[:5]) # Limit preview + header_content = "\n".join(header_lines) + if not header_content.rstrip().endswith(":"): + header_content += "\n" + header_content += "\n # Method signatures:\n" + "\n".join( + method_signatures[:5] + ) # Limit preview if len(method_signatures) > 5: - header_content += f'\n # ... and {len(method_signatures) - 5} more methods' + header_content += f"\n # ... and {len(method_signatures) - 5} more methods" else: - header_content = '\n'.join(header_lines) - + header_content = "\n".join(header_lines) + # Create class header chunk header_end = init_end + 1 if init_method else docstring_end + 1 - chunks.append(CodeChunk( - content=header_content, - file_path=file_path, - start_line=class_start + 1, - end_line=header_end, - chunk_type='class', - name=node.name, - language='python', - file_lines=total_lines - )) - + chunks.append( + CodeChunk( + content=header_content, + file_path=file_path, + start_line=class_start + 1, + end_line=header_end, + chunk_type="class", + name=node.name, + language="python", + file_lines=total_lines, + ) + ) + # Process each method as separate chunk for child in node.body: if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)): - if child.name == '__init__': + if child.name == "__init__": continue # Already included in class header - + method_chunk = self._process_python_function( - child, lines, file_path, - is_method=True, + child, + lines, + file_path, + is_method=True, parent_class=node.name, - total_lines=total_lines + total_lines=total_lines, ) chunks.append(method_chunk) - + return chunks - - def _process_python_function(self, node, lines: List[str], file_path: str, - is_method: bool = False, parent_class: Optional[str] = None, - total_lines: Optional[int] = None) -> CodeChunk: + + def _process_python_function( + self, + node, + lines: List[str], + file_path: str, + is_method: bool = False, + parent_class: Optional[str] = None, + total_lines: Optional[int] = None, + ) -> CodeChunk: """Process a Python function or method, including its docstring.""" start_line = node.lineno - 1 end_line = (node.end_lineno or len(lines)) - 1 - + # Include any decorators - if hasattr(node, 'decorator_list') and node.decorator_list: + if hasattr(node, "decorator_list") and node.decorator_list: first_decorator = node.decorator_list[0] - if hasattr(first_decorator, 'lineno'): + if hasattr(first_decorator, "lineno"): start_line = min(start_line, first_decorator.lineno - 1) - - function_content = '\n'.join(lines[start_line:end_line + 1]) - + + function_content = "\n".join(lines[start_line : end_line + 1]) + return CodeChunk( content=function_content, file_path=file_path, start_line=start_line + 1, end_line=end_line + 1, - chunk_type='method' if is_method else 'function', + chunk_type="method" if is_method else "function", name=node.name, - language='python', + language="python", parent_class=parent_class, - file_lines=total_lines + file_lines=total_lines, ) - - def _chunk_javascript(self, content: str, file_path: str, language: str) -> List[CodeChunk]: + + def _chunk_javascript( + self, content: str, file_path: str, language: str + ) -> List[CodeChunk]: """Chunk JavaScript/TypeScript code using regex patterns.""" chunks = [] lines = content.splitlines() - + # Patterns for different code structures patterns = { - 'function': r'^\s*(?:export\s+)?(?:async\s+)?function\s+(\w+)', - 'arrow_function': r'^\s*(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\([^)]*\)\s*=>', - 'class': r'^\s*(?:export\s+)?class\s+(\w+)', - 'method': r'^\s*(?:async\s+)?(\w+)\s*\([^)]*\)\s*{', + "function": r"^\s*(?:export\s+)?(?:async\s+)?function\s+(\w+)", + "arrow_function": ( + r"^\s*(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*" + r"(?:async\s+)?\([^)]*\)\s*=>" + ), + "class": r"^\s*(?:export\s+)?class\s+(\w+)", + "method": r"^\s*(?:async\s+)?(\w+)\s*\([^)]*\)\s*{", } - + # Find all matches matches = [] for i, line in enumerate(lines): @@ -571,359 +638,380 @@ class CodeChunker: name = match.group(1) matches.append((i, chunk_type, name)) break - + # Sort matches by line number matches.sort(key=lambda x: x[0]) - + # Create chunks between matches for i in range(len(matches)): start_line = matches[i][0] chunk_type = matches[i][1] name = matches[i][2] - + # Find end line (next match or end of file) if i + 1 < len(matches): end_line = matches[i + 1][0] - 1 else: end_line = len(lines) - 1 - + # Find actual end by looking for closing brace brace_count = 0 actual_end = start_line for j in range(start_line, min(end_line + 1, len(lines))): line = lines[j] - brace_count += line.count('{') - line.count('}') + brace_count += line.count("{") - line.count("}") if brace_count == 0 and j > start_line: actual_end = j break else: actual_end = end_line - - chunk_content = '\n'.join(lines[start_line:actual_end + 1]) - chunks.append(CodeChunk( - content=chunk_content, - file_path=file_path, - start_line=start_line + 1, - end_line=actual_end + 1, - chunk_type=chunk_type, - name=name, - language=language - )) - + + chunk_content = "\n".join(lines[start_line : actual_end + 1]) + chunks.append( + CodeChunk( + content=chunk_content, + file_path=file_path, + start_line=start_line + 1, + end_line=actual_end + 1, + chunk_type=chunk_type, + name=name, + language=language, + ) + ) + # If no chunks found, use generic chunking if not chunks: return self._chunk_generic(content, file_path, language) - + return chunks - + def _chunk_go(self, content: str, file_path: str) -> List[CodeChunk]: """Chunk Go code by functions and types.""" chunks = [] lines = content.splitlines() - + # Patterns for Go structures patterns = { - 'function': r'^\s*func\s+(?:\(\w+\s+\*?\w+\)\s+)?(\w+)\s*\(', - 'type': r'^\s*type\s+(\w+)\s+(?:struct|interface)\s*{', - 'method': r'^\s*func\s+\((\w+)\s+\*?\w+\)\s+(\w+)\s*\(', + "function": r"^\s*func\s+(?:\(\w+\s+\*?\w+\)\s+)?(\w+)\s*\(", + "type": r"^\s*type\s+(\w+)\s+(?:struct|interface)\s*{", + "method": r"^\s*func\s+\((\w+)\s+\*?\w+\)\s+(\w+)\s*\(", } - + matches = [] for i, line in enumerate(lines): for chunk_type, pattern in patterns.items(): match = re.match(pattern, line) if match: - if chunk_type == 'method': + if chunk_type == "method": name = f"{match.group(1)}.{match.group(2)}" else: name = match.group(1) matches.append((i, chunk_type, name)) break - + # Process matches similar to JavaScript for i in range(len(matches)): start_line = matches[i][0] chunk_type = matches[i][1] name = matches[i][2] - + # Find end line if i + 1 < len(matches): end_line = matches[i + 1][0] - 1 else: end_line = len(lines) - 1 - + # Find actual end by brace matching brace_count = 0 actual_end = start_line for j in range(start_line, min(end_line + 1, len(lines))): line = lines[j] - brace_count += line.count('{') - line.count('}') + brace_count += line.count("{") - line.count("}") if brace_count == 0 and j > start_line: actual_end = j break - - chunk_content = '\n'.join(lines[start_line:actual_end + 1]) - chunks.append(CodeChunk( - content=chunk_content, - file_path=file_path, - start_line=start_line + 1, - end_line=actual_end + 1, - chunk_type=chunk_type, - name=name, - language='go' - )) - - return chunks if chunks else self._chunk_generic(content, file_path, 'go') - + + chunk_content = "\n".join(lines[start_line : actual_end + 1]) + chunks.append( + CodeChunk( + content=chunk_content, + file_path=file_path, + start_line=start_line + 1, + end_line=actual_end + 1, + chunk_type=chunk_type, + name=name, + language="go", + ) + ) + + return chunks if chunks else self._chunk_generic(content, file_path, "go") + def _chunk_java(self, content: str, file_path: str) -> List[CodeChunk]: """Chunk Java code by classes and methods.""" chunks = [] lines = content.splitlines() - + # Simple regex-based approach for Java - class_pattern = r'^\s*(?:public|private|protected)?\s*(?:abstract|final)?\s*class\s+(\w+)' - method_pattern = r'^\s*(?:public|private|protected)?\s*(?:static)?\s*(?:final)?\s*\w+\s+(\w+)\s*\(' - + class_pattern = ( + r"^\s*(?:public|private|protected)?\s*(?:abstract|final)?\s*class\s+(\w+)" + ) + method_pattern = ( + r"^\s*(?:public|private|protected)?\s*(?:static)?\s*" + r"(?:final)?\s*\w+\s+(\w+)\s*\(" + ) + matches = [] for i, line in enumerate(lines): class_match = re.match(class_pattern, line) if class_match: - matches.append((i, 'class', class_match.group(1))) + matches.append((i, "class", class_match.group(1))) continue - + method_match = re.match(method_pattern, line) if method_match: - matches.append((i, 'method', method_match.group(1))) - + matches.append((i, "method", method_match.group(1))) + # Process matches for i in range(len(matches)): start_line = matches[i][0] chunk_type = matches[i][1] name = matches[i][2] - + # Find end line if i + 1 < len(matches): end_line = matches[i + 1][0] - 1 else: end_line = len(lines) - 1 - - chunk_content = '\n'.join(lines[start_line:end_line + 1]) - chunks.append(CodeChunk( - content=chunk_content, - file_path=file_path, - start_line=start_line + 1, - end_line=end_line + 1, - chunk_type=chunk_type, - name=name, - language='java' - )) - - return chunks if chunks else self._chunk_generic(content, file_path, 'java') - + + chunk_content = "\n".join(lines[start_line : end_line + 1]) + chunks.append( + CodeChunk( + content=chunk_content, + file_path=file_path, + start_line=start_line + 1, + end_line=end_line + 1, + chunk_type=chunk_type, + name=name, + language="java", + ) + ) + + return chunks if chunks else self._chunk_generic(content, file_path, "java") + def _chunk_by_indent(self, content: str, file_path: str, language: str) -> List[CodeChunk]: """Chunk code by indentation levels (fallback for syntax errors).""" chunks = [] lines = content.splitlines() - + current_chunk_start = 0 current_indent = 0 - + for i, line in enumerate(lines): if line.strip(): # Non-empty line # Calculate indentation indent = len(line) - len(line.lstrip()) - + # If dedent detected and chunk is large enough if indent < current_indent and i - current_chunk_start >= self.min_chunk_size: # Create chunk - chunk_content = '\n'.join(lines[current_chunk_start:i]) - chunks.append(CodeChunk( - content=chunk_content, - file_path=file_path, - start_line=current_chunk_start + 1, - end_line=i, - chunk_type='code_block', - name=f"block_{len(chunks) + 1}", - language=language - )) + chunk_content = "\n".join(lines[current_chunk_start:i]) + chunks.append( + CodeChunk( + content=chunk_content, + file_path=file_path, + start_line=current_chunk_start + 1, + end_line=i, + chunk_type="code_block", + name=f"block_{len(chunks) + 1}", + language=language, + ) + ) current_chunk_start = i - + current_indent = indent - + # Add final chunk if current_chunk_start < len(lines): - chunk_content = '\n'.join(lines[current_chunk_start:]) - chunks.append(CodeChunk( - content=chunk_content, - file_path=file_path, - start_line=current_chunk_start + 1, - end_line=len(lines), - chunk_type='code_block', - name=f"block_{len(chunks) + 1}", - language=language - )) - + chunk_content = "\n".join(lines[current_chunk_start:]) + chunks.append( + CodeChunk( + content=chunk_content, + file_path=file_path, + start_line=current_chunk_start + 1, + end_line=len(lines), + chunk_type="code_block", + name=f"block_{len(chunks) + 1}", + language=language, + ) + ) + return chunks - + def _chunk_generic(self, content: str, file_path: str, language: str) -> List[CodeChunk]: """Generic chunking by empty lines and size constraints.""" chunks = [] lines = content.splitlines() - + current_chunk = [] current_start = 0 - + for i, line in enumerate(lines): current_chunk.append(line) - + # Check if we should create a chunk should_chunk = False - + # Empty line indicates potential chunk boundary if not line.strip() and len(current_chunk) >= self.min_chunk_size: should_chunk = True - + # Maximum size reached if len(current_chunk) >= self.max_chunk_size: should_chunk = True - + # End of file if i == len(lines) - 1 and current_chunk: should_chunk = True - + if should_chunk and current_chunk: - chunk_content = '\n'.join(current_chunk).strip() + chunk_content = "\n".join(current_chunk).strip() if chunk_content: # Don't create empty chunks - chunks.append(CodeChunk( - content=chunk_content, - file_path=file_path, - start_line=current_start + 1, - end_line=current_start + len(current_chunk), - chunk_type='code_block', - name=f"block_{len(chunks) + 1}", - language=language - )) - + chunks.append( + CodeChunk( + content=chunk_content, + file_path=file_path, + start_line=current_start + 1, + end_line=current_start + len(current_chunk), + chunk_type="code_block", + name=f"block_{len(chunks) + 1}", + language=language, + ) + ) + # Reset for next chunk current_chunk = [] current_start = i + 1 - + return chunks - + def _enforce_size_constraints(self, chunks: List[CodeChunk]) -> List[CodeChunk]: """ Ensure all chunks meet size constraints. Split too-large chunks and merge too-small ones. """ result = [] - + for chunk in chunks: lines = chunk.content.splitlines() - + # If chunk is too large, split it if len(lines) > self.max_chunk_size: # Split into smaller chunks for i in range(0, len(lines), self.max_chunk_size - self.overlap_lines): - sub_lines = lines[i:i + self.max_chunk_size] + sub_lines = lines[i : i + self.max_chunk_size] if len(sub_lines) >= self.min_chunk_size or not result: - sub_content = '\n'.join(sub_lines) + sub_content = "\n".join(sub_lines) sub_chunk = CodeChunk( content=sub_content, file_path=chunk.file_path, start_line=chunk.start_line + i, end_line=chunk.start_line + i + len(sub_lines) - 1, chunk_type=chunk.chunk_type, - name=f"{chunk.name}_part{i // self.max_chunk_size + 1}" if chunk.name else None, - language=chunk.language + name=( + f"{chunk.name}_part{i // self.max_chunk_size + 1}" + if chunk.name + else None + ), + language=chunk.language, ) result.append(sub_chunk) elif result: # Merge with previous chunk if too small - result[-1].content += '\n' + '\n'.join(sub_lines) + result[-1].content += "\n" + "\n".join(sub_lines) result[-1].end_line = chunk.start_line + i + len(sub_lines) - 1 - + # If chunk is too small, try to merge with previous elif len(lines) < self.min_chunk_size and result: # Check if merging would exceed max size prev_lines = result[-1].content.splitlines() if len(prev_lines) + len(lines) <= self.max_chunk_size: - result[-1].content += '\n' + chunk.content + result[-1].content += "\n" + chunk.content result[-1].end_line = chunk.end_line else: result.append(chunk) - + else: # Chunk is good size result.append(chunk) - + return result - + def _set_chunk_links(self, chunks: List[CodeChunk], file_path: str) -> List[CodeChunk]: """Set chunk indices and prev/next links for navigation.""" total_chunks = len(chunks) - + for i, chunk in enumerate(chunks): chunk.chunk_index = i chunk.total_chunks = total_chunks - - # Generate chunk ID - chunk_id = f"{Path(file_path).stem}_{i}" - + + # Set chunk ID + chunk.chunk_id = f"{Path(file_path).stem}_{i}" + # Set previous chunk link if i > 0: - chunk.prev_chunk_id = f"{Path(file_path).stem}_{i-1}" - + chunk.prev_chunk_id = f"{Path(file_path).stem}_{i - 1}" + # Set next chunk link if i < total_chunks - 1: - chunk.next_chunk_id = f"{Path(file_path).stem}_{i+1}" - + chunk.next_chunk_id = f"{Path(file_path).stem}_{i + 1}" + return chunks - - def _chunk_markdown(self, content: str, file_path: str, language: str = 'markdown') -> List[CodeChunk]: + + def _chunk_markdown( + self, content: str, file_path: str, language: str = "markdown" + ) -> List[CodeChunk]: """ Chunk markdown/text files by sections with context overlap. - + Args: content: File content file_path: Path to file language: Document language type - + Returns: List of chunks """ chunks = [] lines = content.splitlines() total_lines = len(lines) - + # Track current section current_section = [] current_start = 0 section_name = "content" - section_level = 0 - + # Context overlap for markdown (keep last few lines) overlap_buffer = [] overlap_size = 3 # Lines to overlap between chunks - + # Patterns for different section types - header_pattern = re.compile(r'^(#+)\s+(.+)$') # Markdown headers with level - separator_pattern = re.compile(r'^[-=]{3,}$') # Horizontal rules - + header_pattern = re.compile(r"^(#+)\s+(.+)$") # Markdown headers with level + separator_pattern = re.compile(r"^[-=]{3,}$") # Horizontal rules + for i, line in enumerate(lines): # Check for headers header_match = header_pattern.match(line) - + # Check for section breaks is_separator = separator_pattern.match(line.strip()) is_empty = not line.strip() - + # Decide if we should create a chunk should_chunk = False - + if header_match: # New header found should_chunk = True - new_section_level = len(header_match.group(1)) new_section_name = header_match.group(2).strip() elif is_separator: # Separator found @@ -933,185 +1021,192 @@ class CodeChunker: if i + 1 < len(lines) and not lines[i + 1].strip(): # Multiple empty lines - chunk here should_chunk = True - + # Check size constraints if len(current_section) >= self.max_chunk_size: should_chunk = True - + if should_chunk and current_section: # Add overlap from previous chunk if available section_with_overlap = overlap_buffer + current_section - + # Create chunk from current section - chunk_content = '\n'.join(section_with_overlap) + chunk_content = "\n".join(section_with_overlap) if chunk_content.strip(): # Only create chunk if non-empty chunk = CodeChunk( content=chunk_content, file_path=file_path, start_line=max(1, current_start + 1 - len(overlap_buffer)), end_line=current_start + len(current_section), - chunk_type='section', + chunk_type="section", name=section_name[:50], # Limit name length language=language, - file_lines=total_lines + file_lines=total_lines, ) chunks.append(chunk) - + # Save overlap for next chunk if len(current_section) > overlap_size: overlap_buffer = current_section[-overlap_size:] else: overlap_buffer = current_section[:] - + # Reset for next section current_section = [] current_start = i + 1 - + # Update section name if we found a header if header_match: section_name = new_section_name - section_level = new_section_level else: section_name = f"section_{len(chunks) + 1}" - + # Add line to current section if not (should_chunk and (header_match or is_separator)): current_section.append(line) - + # Don't forget the last section if current_section: section_with_overlap = overlap_buffer + current_section - chunk_content = '\n'.join(section_with_overlap) + chunk_content = "\n".join(section_with_overlap) if chunk_content.strip(): chunk = CodeChunk( content=chunk_content, file_path=file_path, start_line=max(1, current_start + 1 - len(overlap_buffer)), end_line=len(lines), - chunk_type='section', + chunk_type="section", name=section_name[:50], language=language, - file_lines=total_lines + file_lines=total_lines, ) chunks.append(chunk) - + # If no chunks created, create one for the whole file if not chunks and content.strip(): - chunks.append(CodeChunk( - content=content, - file_path=file_path, - start_line=1, - end_line=len(lines), - chunk_type='document', - name=Path(file_path).stem, - language=language, - file_lines=total_lines - )) - + chunks.append( + CodeChunk( + content=content, + file_path=file_path, + start_line=1, + end_line=len(lines), + chunk_type="document", + name=Path(file_path).stem, + language=language, + file_lines=total_lines, + ) + ) + # Set chunk links chunks = self._set_chunk_links(chunks, file_path) - + return chunks - - def _chunk_config(self, content: str, file_path: str, language: str = 'config') -> List[CodeChunk]: + + def _chunk_config( + self, content: str, file_path: str, language: str = "config" + ) -> List[CodeChunk]: """ Chunk configuration files by sections. - + Args: content: File content file_path: Path to file language: Config language type - + Returns: List of chunks """ # For config files, we'll create smaller chunks by top-level sections chunks = [] lines = content.splitlines() - - if language == 'json': + + if language == "json": # For JSON, just create one chunk for now # (Could be enhanced to chunk by top-level keys) - chunks.append(CodeChunk( - content=content, - file_path=file_path, - start_line=1, - end_line=len(lines), - chunk_type='config', - name=Path(file_path).stem, - language=language - )) + chunks.append( + CodeChunk( + content=content, + file_path=file_path, + start_line=1, + end_line=len(lines), + chunk_type="config", + name=Path(file_path).stem, + language=language, + ) + ) else: # For YAML, INI, TOML, etc., chunk by sections current_section = [] current_start = 0 section_name = "config" - + # Patterns for section headers section_patterns = { - 'ini': re.compile(r'^\[(.+)\]$'), - 'toml': re.compile(r'^\[(.+)\]$'), - 'yaml': re.compile(r'^(\w+):$'), + "ini": re.compile(r"^\[(.+)\]$"), + "toml": re.compile(r"^\[(.+)\]$"), + "yaml": re.compile(r"^(\w+):$"), } - + pattern = section_patterns.get(language) - + for i, line in enumerate(lines): is_section = False - + if pattern: match = pattern.match(line.strip()) if match: is_section = True new_section_name = match.group(1) - + if is_section and current_section: # Create chunk for previous section - chunk_content = '\n'.join(current_section) + chunk_content = "\n".join(current_section) if chunk_content.strip(): chunk = CodeChunk( content=chunk_content, file_path=file_path, start_line=current_start + 1, end_line=current_start + len(current_section), - chunk_type='config_section', + chunk_type="config_section", name=section_name, - language=language + language=language, ) chunks.append(chunk) - + # Start new section current_section = [line] current_start = i section_name = new_section_name else: current_section.append(line) - + # Add final section if current_section: - chunk_content = '\n'.join(current_section) + chunk_content = "\n".join(current_section) if chunk_content.strip(): chunk = CodeChunk( content=chunk_content, file_path=file_path, start_line=current_start + 1, end_line=len(lines), - chunk_type='config_section', + chunk_type="config_section", name=section_name, - language=language + language=language, ) chunks.append(chunk) - + # If no chunks created, create one for the whole file if not chunks and content.strip(): - chunks.append(CodeChunk( - content=content, - file_path=file_path, - start_line=1, - end_line=len(lines), - chunk_type='config', - name=Path(file_path).stem, - language=language - )) - - return chunks \ No newline at end of file + chunks.append( + CodeChunk( + content=content, + file_path=file_path, + start_line=1, + end_line=len(lines), + chunk_type="config", + name=Path(file_path).stem, + language=language, + ) + ) + + return chunks diff --git a/mini_rag/cli.py b/mini_rag/cli.py index 47ebc55..22d93d9 100644 --- a/mini_rag/cli.py +++ b/mini_rag/cli.py @@ -3,59 +3,57 @@ Command-line interface for Mini RAG system. Beautiful, intuitive, and highly effective. """ -import click +import logging import sys import time -import logging from pathlib import Path from typing import Optional -# Fix Windows console for proper emoji/Unicode support -from .windows_console_fix import fix_windows_console -fix_windows_console() - +import click from rich.console import Console -from rich.table import Table -from rich.progress import Progress, SpinnerColumn, TextColumn from rich.logging import RichHandler -from rich.syntax import Syntax from rich.panel import Panel -from rich import print as rprint +from rich.progress import Progress, SpinnerColumn, TextColumn +from rich.syntax import Syntax +from rich.table import Table from .indexer import ProjectIndexer -from .search import CodeSearcher -from .watcher import FileWatcher from .non_invasive_watcher import NonInvasiveFileWatcher from .ollama_embeddings import OllamaEmbedder as CodeEmbedder -from .chunker import CodeChunker from .performance import get_monitor -from .server import RAGClient -from .server import RAGServer, RAGClient, start_server +from .search import CodeSearcher +from .server import RAGClient, start_server +from .windows_console_fix import fix_windows_console + +# Fix Windows console for proper emoji/Unicode support +fix_windows_console() # Set up logging logging.basicConfig( level=logging.INFO, format="%(message)s", - handlers=[RichHandler(rich_tracebacks=True)] + handlers=[RichHandler(rich_tracebacks=True)], ) logger = logging.getLogger(__name__) console = Console() @click.group() -@click.option('--verbose', '-v', is_flag=True, help='Enable verbose logging') -@click.option('--quiet', '-q', is_flag=True, help='Suppress output') +@click.option("--verbose", "-v", is_flag=True, help="Enable verbose logging") +@click.option("--quiet", "-q", is_flag=True, help="Suppress output") def cli(verbose: bool, quiet: bool): """ Mini RAG - Fast semantic code search that actually works. - - A local RAG system for improving the development environment's grounding capabilities. + + A local RAG system for improving the development environment's grounding + capabilities. Indexes your codebase and enables lightning-fast semantic search. """ # Check virtual environment from .venv_checker import check_and_warn_venv + check_and_warn_venv("rag-mini", force_exit=False) - + if verbose: logging.getLogger().setLevel(logging.DEBUG) elif quiet: @@ -63,43 +61,45 @@ def cli(verbose: bool, quiet: bool): @cli.command() -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path to index') -@click.option('--force', '-f', is_flag=True, - help='Force reindex all files') -@click.option('--reindex', '-r', is_flag=True, - help='Force complete reindex (same as --force)') -@click.option('--model', '-m', type=str, default=None, - help='Embedding model to use') +@click.option( + "--path", + "-p", + type=click.Path(exists=True), + default=".", + help="Project path to index", +) +@click.option("--force", "-", is_flag=True, help="Force reindex all files") +@click.option("--reindex", "-r", is_flag=True, help="Force complete reindex (same as --force)") +@click.option("--model", "-m", type=str, default=None, help="Embedding model to use") def init(path: str, force: bool, reindex: bool, model: Optional[str]): """Initialize RAG index for a project.""" project_path = Path(path).resolve() - + console.print(f"\n[bold cyan]Initializing Mini RAG for:[/bold cyan] {project_path}\n") - + # Check if already initialized - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" force_reindex = force or reindex if rag_dir.exists() and not force_reindex: console.print("[yellow][/yellow] Project already initialized!") console.print("Use --force or --reindex to reindex all files\n") - + # Show current stats indexer = ProjectIndexer(project_path) stats = indexer.get_statistics() - + table = Table(title="Current Index Statistics") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") - - table.add_row("Files Indexed", str(stats['file_count'])) - table.add_row("Total Chunks", str(stats['chunk_count'])) + + table.add_row("Files Indexed", str(stats["file_count"])) + table.add_row("Total Chunks", str(stats["chunk_count"])) table.add_row("Index Size", f"{stats['index_size_mb']:.2f} MB") - table.add_row("Last Updated", stats['indexed_at'] or "Never") - + table.add_row("Last Updated", stats["indexed_at"] or "Never") + console.print(table) return - + # Initialize components try: with Progress( @@ -111,34 +111,33 @@ def init(path: str, force: bool, reindex: bool, model: Optional[str]): task = progress.add_task("[cyan]Loading embedding model...", total=None) embedder = CodeEmbedder(model_name=model) progress.update(task, completed=True) - + # Create indexer task = progress.add_task("[cyan]Creating indexer...", total=None) - indexer = ProjectIndexer( - project_path, - embedder=embedder - ) + indexer = ProjectIndexer(project_path, embedder=embedder) progress.update(task, completed=True) - + # Run indexing console.print("\n[bold green]Starting indexing...[/bold green]\n") stats = indexer.index_project(force_reindex=force_reindex) - + # Show summary - if stats['files_indexed'] > 0: - console.print(f"\n[bold green] Success![/bold green] Indexed {stats['files_indexed']} files") + if stats["files_indexed"] > 0: + console.print( + f"\n[bold green] Success![/bold green] Indexed {stats['files_indexed']} files" + ) console.print(f"Created {stats['chunks_created']} searchable chunks") console.print(f"Time: {stats['time_taken']:.2f} seconds") console.print(f"Speed: {stats['files_per_second']:.1f} files/second") else: console.print("\n[green] All files are already up to date![/green]") - + # Show how to use console.print("\n[bold]Next steps:[/bold]") - console.print(" โ€ข Search your code: [cyan]rag-mini search \"your query\"[/cyan]") + console.print(' โ€ข Search your code: [cyan]rag-mini search "your query"[/cyan]') console.print(" โ€ข Watch for changes: [cyan]rag-mini watch[/cyan]") console.print(" โ€ข View statistics: [cyan]rag-mini stats[/cyan]\n") - + except Exception as e: console.print(f"\n[bold red]Error:[/bold red] {e}") logger.exception("Initialization failed") @@ -146,64 +145,71 @@ def init(path: str, force: bool, reindex: bool, model: Optional[str]): @cli.command() -@click.argument('query') -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') -@click.option('--top-k', '-k', type=int, default=10, - help='Maximum results to show') -@click.option('--type', '-t', multiple=True, - help='Filter by chunk type (function, class, method)') -@click.option('--lang', multiple=True, - help='Filter by language (python, javascript, etc.)') -@click.option('--show-content', '-c', is_flag=True, - help='Show code content in results') -@click.option('--show-perf', is_flag=True, - help='Show performance metrics') -def search(query: str, path: str, top_k: int, type: tuple, lang: tuple, show_content: bool, show_perf: bool): +@click.argument("query") +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") +@click.option("--top-k", "-k", type=int, default=10, help="Maximum results to show") +@click.option( + "--type", "-t", multiple=True, help="Filter by chunk type (function, class, method)" +) +@click.option("--lang", multiple=True, help="Filter by language (python, javascript, etc.)") +@click.option("--show-content", "-c", is_flag=True, help="Show code content in results") +@click.option("--show-per", is_flag=True, help="Show performance metrics") +def search( + query: str, + path: str, + top_k: int, + type: tuple, + lang: tuple, + show_content: bool, + show_perf: bool, +): """Search codebase using semantic similarity.""" project_path = Path(path).resolve() - + # Check if indexed - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if not rag_dir.exists(): console.print("[red]Error:[/red] Project not indexed. Run 'rag-mini init' first.") sys.exit(1) - + # Get performance monitor monitor = get_monitor() if show_perf else None - + # Check if server is running client = RAGClient() use_server = client.is_running() - + try: if use_server: # Use server for fast queries console.print("[dim]Using RAG server...[/dim]") - + response = client.search(query, top_k=top_k) - - if response.get('success'): + + if response.get("success"): # Convert response to SearchResult objects from .search import SearchResult + results = [] - for r in response['results']: + for r in response["results"]: result = SearchResult( - file_path=r['file_path'], - content=r['content'], - score=r['score'], - start_line=r['start_line'], - end_line=r['end_line'], - chunk_type=r['chunk_type'], - name=r['name'], - language=r['language'] + file_path=r["file_path"], + content=r["content"], + score=r["score"], + start_line=r["start_line"], + end_line=r["end_line"], + chunk_type=r["chunk_type"], + name=r["name"], + language=r["language"], ) results.append(result) - + # Show server stats - search_time = response.get('search_time_ms', 0) - total_queries = response.get('total_queries', 0) - console.print(f"[dim]Search time: {search_time}ms (Query #{total_queries})[/dim]\n") + search_time = response.get("search_time_ms", 0) + total_queries = response.get("total_queries", 0) + console.print( + f"[dim]Search time: {search_time}ms (Query #{total_queries})[/dim]\n" + ) else: console.print(f"[red]Server error:[/red] {response.get('error')}") sys.exit(1) @@ -215,7 +221,7 @@ def search(query: str, path: str, top_k: int, type: tuple, lang: tuple, show_con searcher = CodeSearcher(project_path) else: searcher = CodeSearcher(project_path) - + # Perform search with timing if monitor: with monitor.measure("Execute Vector Search"): @@ -223,7 +229,7 @@ def search(query: str, path: str, top_k: int, type: tuple, lang: tuple, show_con query, top_k=top_k, chunk_types=list(type) if type else None, - languages=list(lang) if lang else None + languages=list(lang) if lang else None, ) else: with console.status(f"[cyan]Searching for: {query}[/cyan]"): @@ -231,9 +237,9 @@ def search(query: str, path: str, top_k: int, type: tuple, lang: tuple, show_con query, top_k=top_k, chunk_types=list(type) if type else None, - languages=list(lang) if lang else None + languages=list(lang) if lang else None, ) - + # Display results if results: if use_server: @@ -243,27 +249,30 @@ def search(query: str, path: str, top_k: int, type: tuple, lang: tuple, show_con display_searcher.display_results(results, show_content=show_content) else: searcher.display_results(results, show_content=show_content) - + # Copy first result to clipboard if available try: import pyperclip + first_result = results[0] location = f"{first_result.file_path}:{first_result.start_line}" pyperclip.copy(location) - console.print(f"\n[dim]First result location copied to clipboard: {location}[/dim]") - except: - pass + console.print( + f"\n[dim]First result location copied to clipboard: {location}[/dim]" + ) + except (ImportError, OSError): + pass # Clipboard not available else: console.print(f"\n[yellow]No results found for: {query}[/yellow]") console.print("\n[dim]Tips:[/dim]") console.print(" โ€ข Try different keywords") console.print(" โ€ข Use natural language queries") - + # Show performance summary if monitor: monitor.print_summary() console.print(" โ€ข Check if files are indexed with 'mini-rag stats'") - + except Exception as e: console.print(f"\n[bold red]Search error:[/bold red] {e}") logger.exception("Search failed") @@ -271,68 +280,69 @@ def search(query: str, path: str, top_k: int, type: tuple, lang: tuple, show_con @cli.command() -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") def stats(path: str): """Show index statistics.""" project_path = Path(path).resolve() - + # Check if indexed - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if not rag_dir.exists(): console.print("[red]Error:[/red] Project not indexed. Run 'rag-mini init' first.") sys.exit(1) - + try: # Get statistics indexer = ProjectIndexer(project_path) index_stats = indexer.get_statistics() - + searcher = CodeSearcher(project_path) search_stats = searcher.get_statistics() - + # Display project info console.print(f"\n[bold cyan]Project:[/bold cyan] {project_path.name}") console.print(f"[dim]Path: {project_path}[/dim]\n") - + # Index statistics table table = Table(title="Index Statistics") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") - - table.add_row("Files Indexed", str(index_stats['file_count'])) - table.add_row("Total Chunks", str(index_stats['chunk_count'])) + + table.add_row("Files Indexed", str(index_stats["file_count"])) + table.add_row("Total Chunks", str(index_stats["chunk_count"])) table.add_row("Index Size", f"{index_stats['index_size_mb']:.2f} MB") - table.add_row("Last Updated", index_stats['indexed_at'] or "Never") - + table.add_row("Last Updated", index_stats["indexed_at"] or "Never") + console.print(table) - + # Language distribution - if 'languages' in search_stats: + if "languages" in search_stats: console.print("\n[bold]Language Distribution:[/bold]") lang_table = Table() lang_table.add_column("Language", style="cyan") lang_table.add_column("Chunks", style="green") - - for lang, count in sorted(search_stats['languages'].items(), - key=lambda x: x[1], reverse=True): + + for lang, count in sorted( + search_stats["languages"].items(), key=lambda x: x[1], reverse=True + ): lang_table.add_row(lang, str(count)) - + console.print(lang_table) - + # Chunk type distribution - if 'chunk_types' in search_stats: + if "chunk_types" in search_stats: console.print("\n[bold]Chunk Types:[/bold]") type_table = Table() type_table.add_column("Type", style="cyan") type_table.add_column("Count", style="green") - - for chunk_type, count in sorted(search_stats['chunk_types'].items(), - key=lambda x: x[1], reverse=True): + + for chunk_type, count in sorted( + search_stats["chunk_types"].items(), key=lambda x: x[1], reverse=True + ): type_table.add_row(chunk_type, str(count)) - + console.print(type_table) - + except Exception as e: console.print(f"\n[bold red]Error:[/bold red] {e}") logger.exception("Failed to get statistics") @@ -340,101 +350,116 @@ def stats(path: str): @cli.command() -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") def debug_schema(path: str): """Debug vector database schema and sample data.""" project_path = Path(path).resolve() - + try: - rag_dir = project_path / '.mini-rag' - + rag_dir = project_path / ".mini-rag" + if not rag_dir.exists(): console.print("[red]No RAG index found. Run 'rag-mini init' first.[/red]") return - + # Connect to database try: import lancedb except ImportError: - console.print("[red]LanceDB not available. Install with: pip install lancedb pyarrow[/red]") + console.print( + "[red]LanceDB not available. Install with: pip install lancedb pyarrow[/red]" + ) return - + db = lancedb.connect(rag_dir) - + if "code_vectors" not in db.table_names(): console.print("[red]No code_vectors table found.[/red]") return - + table = db.open_table("code_vectors") - + # Print schema console.print("\n[bold cyan] Table Schema:[/bold cyan]") console.print(table.schema) - + # Get sample data - import pandas as pd + df = table.to_pandas() - console.print(f"\n[bold cyan] Table Statistics:[/bold cyan]") + console.print("\n[bold cyan] Table Statistics:[/bold cyan]") console.print(f"Total rows: {len(df)}") - + if len(df) > 0: # Check embedding column - console.print(f"\n[bold cyan] Embedding Column Analysis:[/bold cyan]") - first_embedding = df['embedding'].iloc[0] + console.print("\n[bold cyan] Embedding Column Analysis:[/bold cyan]") + first_embedding = df["embedding"].iloc[0] console.print(f"Type: {type(first_embedding)}") - if hasattr(first_embedding, 'shape'): + if hasattr(first_embedding, "shape"): console.print(f"Shape: {first_embedding.shape}") - if hasattr(first_embedding, 'dtype'): + if hasattr(first_embedding, "dtype"): console.print(f"Dtype: {first_embedding.dtype}") - + # Show first few rows - console.print(f"\n[bold cyan] Sample Data (first 3 rows):[/bold cyan]") + console.print("\n[bold cyan] Sample Data (first 3 rows):[/bold cyan]") for i in range(min(3, len(df))): row = df.iloc[i] console.print(f"\n[yellow]Row {i}:[/yellow]") console.print(f" chunk_id: {row['chunk_id']}") console.print(f" file_path: {row['file_path']}") console.print(f" content: {row['content'][:50]}...") - console.print(f" embedding: {type(row['embedding'])} of length {len(row['embedding']) if hasattr(row['embedding'], '__len__') else 'unknown'}") - + embed_len = ( + len(row["embedding"]) + if hasattr(row["embedding"], "__len__") + else "unknown" + ) + console.print(f" embedding: {type(row['embedding'])} of length {embed_len}") + except Exception as e: logger.error(f"Schema debug failed: {e}") console.print(f"[red]Error: {e}[/red]") @cli.command() -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') -@click.option('--delay', '-d', type=float, default=10.0, - help='Update delay in seconds (default: 10s for non-invasive)') -@click.option('--silent', '-s', is_flag=True, default=False, - help='Run silently in background without output') +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") +@click.option( + "--delay", + "-d", + type=float, + default=10.0, + help="Update delay in seconds (default: 10s for non-invasive)", +) +@click.option( + "--silent", + "-s", + is_flag=True, + default=False, + help="Run silently in background without output", +) def watch(path: str, delay: float, silent: bool): """Watch for file changes and update index automatically (non-invasive by default).""" project_path = Path(path).resolve() - + # Check if indexed - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if not rag_dir.exists(): if not silent: console.print("[red]Error:[/red] Project not indexed. Run 'rag-mini init' first.") sys.exit(1) - + try: # Always use non-invasive watcher watcher = NonInvasiveFileWatcher(project_path) - + # Only show startup messages if not silent if not silent: console.print(f"\n[bold green]๐Ÿ•Š๏ธ Non-Invasive Watcher:[/bold green] {project_path}") console.print("[dim]Low CPU/memory usage - won't interfere with development[/dim]") console.print(f"[dim]Update delay: {delay}s[/dim]") console.print("\n[yellow]Press Ctrl+C to stop watching[/yellow]\n") - + # Start watching watcher.start() - + if silent: # Silent mode: just wait for interrupt without any output try: @@ -448,10 +473,10 @@ def watch(path: str, delay: float, silent: bool): while True: try: time.sleep(1) - + # Get current statistics stats = watcher.get_statistics() - + # Only update display if something changed if stats != last_stats: # Clear previous line @@ -459,26 +484,28 @@ def watch(path: str, delay: float, silent: bool): f"\r[green]โœ“[/green] Files updated: {stats.get('files_processed', 0)} | " f"[red]โœ—[/red] Failed: {stats.get('files_dropped', 0)} | " f"[cyan]โง—[/cyan] Queue: {stats['queue_size']}", - end="" + end="", ) last_stats = stats - + except KeyboardInterrupt: break - + # Stop watcher if not silent: console.print("\n\n[yellow]Stopping watcher...[/yellow]") watcher.stop() - + # Show final stats only if not silent if not silent: final_stats = watcher.get_statistics() - console.print(f"\n[bold green]Watch Summary:[/bold green]") + console.print("\n[bold green]Watch Summary:[/bold green]") console.print(f"Files updated: {final_stats.get('files_processed', 0)}") console.print(f"Files failed: {final_stats.get('files_dropped', 0)}") - console.print(f"Total runtime: {final_stats.get('uptime_seconds', 0):.1f} seconds\n") - + console.print( + f"Total runtime: {final_stats.get('uptime_seconds', 0):.1f} seconds\n" + ) + except Exception as e: console.print(f"\n[bold red]Error:[/bold red] {e}") logger.exception("Watch failed") @@ -486,86 +513,81 @@ def watch(path: str, delay: float, silent: bool): @cli.command() -@click.argument('function_name') -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') -@click.option('--top-k', '-k', type=int, default=5, - help='Maximum results') +@click.argument("function_name") +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") +@click.option("--top-k", "-k", type=int, default=5, help="Maximum results") def find_function(function_name: str, path: str, top_k: int): """Find a specific function by name.""" project_path = Path(path).resolve() - + try: searcher = CodeSearcher(project_path) results = searcher.get_function(function_name, top_k=top_k) - + if results: searcher.display_results(results, show_content=True) else: console.print(f"[yellow]No functions found matching: {function_name}[/yellow]") - + except Exception as e: console.print(f"[red]Error:[/red] {e}") sys.exit(1) @cli.command() -@click.argument('class_name') -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') -@click.option('--top-k', '-k', type=int, default=5, - help='Maximum results') +@click.argument("class_name") +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") +@click.option("--top-k", "-k", type=int, default=5, help="Maximum results") def find_class(class_name: str, path: str, top_k: int): """Find a specific class by name.""" project_path = Path(path).resolve() - + try: searcher = CodeSearcher(project_path) results = searcher.get_class(class_name, top_k=top_k) - + if results: searcher.display_results(results, show_content=True) else: console.print(f"[yellow]No classes found matching: {class_name}[/yellow]") - + except Exception as e: console.print(f"[red]Error:[/red] {e}") sys.exit(1) @cli.command() -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") def update(path: str): """Update index for changed files.""" project_path = Path(path).resolve() - + # Check if indexed - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if not rag_dir.exists(): console.print("[red]Error:[/red] Project not indexed. Run 'rag-mini init' first.") sys.exit(1) - + try: indexer = ProjectIndexer(project_path) - + console.print(f"\n[cyan]Checking for changes in {project_path}...[/cyan]\n") - + stats = indexer.index_project(force_reindex=False) - - if stats['files_indexed'] > 0: + + if stats["files_indexed"] > 0: console.print(f"[green][/green] Updated {stats['files_indexed']} files") console.print(f"Created {stats['chunks_created']} new chunks") else: console.print("[green] All files are up to date![/green]") - + except Exception as e: console.print(f"[red]Error:[/red] {e}") sys.exit(1) @cli.command() -@click.option('--show-code', '-c', is_flag=True, help='Show example code') +@click.option("--show-code", "-c", is_flag=True, help="Show example code") def info(show_code: bool): """Show information about Mini RAG.""" # Create info panel @@ -590,13 +612,13 @@ def info(show_code: bool): โ€ข Search: <50ms latency โ€ข Storage: ~200MB for 10k files """ - + panel = Panel(info_text, title="About Mini RAG", border_style="cyan") console.print(panel) - + if show_code: console.print("\n[bold]Example Usage:[/bold]\n") - + code = """# Initialize a project rag-mini init @@ -613,32 +635,30 @@ rag-mini watch # Get statistics rag-mini stats""" - + syntax = Syntax(code, "bash", theme="monokai") console.print(syntax) @cli.command() -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') -@click.option('--port', type=int, default=7777, - help='Server port') +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") +@click.option("--port", type=int, default=7777, help="Server port") def server(path: str, port: int): """Start persistent RAG server (keeps model loaded).""" project_path = Path(path).resolve() - + # Check if indexed - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if not rag_dir.exists(): console.print("[red]Error:[/red] Project not indexed. Run 'rag-mini init' first.") sys.exit(1) - + try: console.print(f"[bold cyan]Starting RAG server for:[/bold cyan] {project_path}") console.print(f"[dim]Port: {port}[/dim]\n") - + start_server(project_path, port) - + except KeyboardInterrupt: console.print("\n[yellow]Server stopped by user[/yellow]") except Exception as e: @@ -648,65 +668,67 @@ def server(path: str, port: int): @cli.command() -@click.option('--path', '-p', type=click.Path(exists=True), default='.', - help='Project path') -@click.option('--port', type=int, default=7777, - help='Server port') -@click.option('--discovery', '-d', is_flag=True, - help='Run codebase discovery analysis') +@click.option("--path", "-p", type=click.Path(exists=True), default=".", help="Project path") +@click.option("--port", type=int, default=7777, help="Server port") +@click.option("--discovery", "-d", is_flag=True, help="Run codebase discovery analysis") def status(path: str, port: int, discovery: bool): """Show comprehensive RAG system status with optional codebase discovery.""" project_path = Path(path).resolve() - + # Print header console.print(f"\n[bold cyan]RAG System Status for:[/bold cyan] {project_path.name}") console.print(f"[dim]Path: {project_path}[/dim]\n") - + # Check folder contents console.print("[bold]๐Ÿ“ Folder Contents:[/bold]") try: all_files = list(project_path.rglob("*")) - source_files = [f for f in all_files if f.is_file() and f.suffix in ['.py', '.js', '.ts', '.go', '.java', '.cpp', '.c', '.h']] - + source_files = [ + f + for f in all_files + if f.is_file() + and f.suffix in [".py", ".js", ".ts", ".go", ".java", ".cpp", ".c", ".h"] + ] + console.print(f" โ€ข Total files: {len([f for f in all_files if f.is_file()])}") console.print(f" โ€ข Source files: {len(source_files)}") console.print(f" โ€ข Directories: {len([f for f in all_files if f.is_dir()])}") except Exception as e: console.print(f" [red]Error reading folder: {e}[/red]") - + # Check index status console.print("\n[bold]๐Ÿ—‚๏ธ Index Status:[/bold]") - rag_dir = project_path / '.mini-rag' + rag_dir = project_path / ".mini-rag" if rag_dir.exists(): try: indexer = ProjectIndexer(project_path) index_stats = indexer.get_statistics() - - console.print(f" โ€ข Status: [green]โœ… Indexed[/green]") + + console.print(" โ€ข Status: [green]โœ… Indexed[/green]") console.print(f" โ€ข Files indexed: {index_stats['file_count']}") console.print(f" โ€ข Total chunks: {index_stats['chunk_count']}") console.print(f" โ€ข Index size: {index_stats['index_size_mb']:.2f} MB") console.print(f" โ€ข Last updated: {index_stats['indexed_at'] or 'Never'}") except Exception as e: - console.print(f" โ€ข Status: [yellow]โš ๏ธ Index exists but has issues[/yellow]") + console.print(" โ€ข Status: [yellow]โš ๏ธ Index exists but has issues[/yellow]") console.print(f" โ€ข Error: {e}") else: console.print(" โ€ข Status: [red]โŒ Not indexed[/red]") console.print(" โ€ข Run 'rag-mini init' to initialize") - + # Check server status console.print("\n[bold]๐Ÿš€ Server Status:[/bold]") client = RAGClient(port) - + if client.is_running(): console.print(f" โ€ข Status: [green]โœ… Running on port {port}[/green]") - + # Try to get server info try: response = client.search("test", top_k=1) # Minimal query to get stats - if response.get('success'): - uptime = response.get('server_uptime', 0) - queries = response.get('total_queries', 0) + if response.get("success"): + uptime = response.get("server_uptime", 0) + queries = response.get("total_queries", 0) console.print(f" โ€ข Uptime: {uptime}s") console.print(f" โ€ข Total queries: {queries}") except Exception as e: @@ -714,47 +736,51 @@ def status(path: str, port: int, discovery: bool): else: console.print(f" โ€ข Status: [red]โŒ Not running on port {port}[/red]") console.print(" โ€ข Run 'rag-mini server' to start the server") - + # Run codebase discovery if requested if discovery and rag_dir.exists(): console.print("\n[bold]๐Ÿง  Codebase Discovery:[/bold]") try: # Import and run intelligent discovery import sys - - # Add tools directory to path + + # Add tools directory to path tools_path = Path(__file__).parent.parent.parent / "tools" if tools_path.exists(): sys.path.insert(0, str(tools_path)) from intelligent_codebase_discovery import IntelligentCodebaseDiscovery - + discovery_system = IntelligentCodebaseDiscovery(project_path) discovery_system.run_lightweight_discovery() else: console.print(" [yellow]Discovery system not found[/yellow]") - + except Exception as e: console.print(f" [red]Discovery failed: {e}[/red]") - + elif discovery and not rag_dir.exists(): console.print("\n[bold]๐Ÿง  Codebase Discovery:[/bold]") console.print(" [yellow]โŒ Cannot run discovery - project not indexed[/yellow]") console.print(" Run 'rag-mini init' first to initialize the system") - + # Show next steps console.print("\n[bold]๐Ÿ“‹ Next Steps:[/bold]") if not rag_dir.exists(): console.print(" 1. Run [cyan]rag-mini init[/cyan] to initialize the RAG system") - console.print(" 2. Use [cyan]rag-mini search \"your query\"[/cyan] to search code") + console.print(' 2. Use [cyan]rag-mini search "your query"[/cyan] to search code') elif not client.is_running(): console.print(" 1. Run [cyan]rag-mini server[/cyan] to start the server") - console.print(" 2. Use [cyan]rag-mini search \"your query\"[/cyan] to search code") + console.print(' 2. Use [cyan]rag-mini search "your query"[/cyan] to search code') else: - console.print(" โ€ข System ready! Use [cyan]rag-mini search \"your query\"[/cyan] to search") - console.print(" โ€ข Add [cyan]--discovery[/cyan] flag to run intelligent codebase analysis") - + console.print( + ' โ€ข System ready! Use [cyan]rag-mini search "your query"[/cyan] to search' + ) + console.print( + " โ€ข Add [cyan]--discovery[/cyan] flag to run intelligent codebase analysis" + ) + console.print() -if __name__ == '__main__': - cli() \ No newline at end of file +if __name__ == "__main__": + cli() diff --git a/mini_rag/config.py b/mini_rag/config.py index 77c4c1e..21d1a49 100644 --- a/mini_rag/config.py +++ b/mini_rag/config.py @@ -3,11 +3,12 @@ Configuration management for FSS-Mini-RAG. Handles loading, saving, and validation of YAML config files. """ -import yaml import logging +from dataclasses import asdict, dataclass from pathlib import Path -from typing import Dict, Any, Optional -from dataclasses import dataclass, asdict +from typing import Any, Dict, Optional + +import yaml logger = logging.getLogger(__name__) @@ -15,6 +16,7 @@ logger = logging.getLogger(__name__) @dataclass class ChunkingConfig: """Configuration for text chunking.""" + max_size: int = 2000 min_size: int = 150 strategy: str = "semantic" # "semantic" or "fixed" @@ -23,6 +25,7 @@ class ChunkingConfig: @dataclass class StreamingConfig: """Configuration for large file streaming.""" + enabled: bool = True threshold_bytes: int = 1048576 # 1MB @@ -30,21 +33,22 @@ class StreamingConfig: @dataclass class FilesConfig: """Configuration for file processing.""" + min_file_size: int = 50 exclude_patterns: list = None include_patterns: list = None - + def __post_init__(self): if self.exclude_patterns is None: self.exclude_patterns = [ "node_modules/**", - ".git/**", + ".git/**", "__pycache__/**", "*.pyc", ".venv/**", "venv/**", "build/**", - "dist/**" + "dist/**", ] if self.include_patterns is None: self.include_patterns = ["**/*"] # Include everything by default @@ -53,6 +57,7 @@ class FilesConfig: @dataclass class EmbeddingConfig: """Configuration for embedding generation.""" + preferred_method: str = "ollama" # "ollama", "ml", "hash", "auto" ollama_model: str = "nomic-embed-text" ollama_host: str = "localhost:11434" @@ -63,52 +68,51 @@ class EmbeddingConfig: @dataclass class SearchConfig: """Configuration for search behavior.""" + default_top_k: int = 10 enable_bm25: bool = True similarity_threshold: float = 0.1 expand_queries: bool = False # Enable automatic query expansion -@dataclass +@dataclass class LLMConfig: """Configuration for LLM synthesis and query expansion.""" + # Core settings synthesis_model: str = "auto" # "auto", "qwen3:1.7b", "qwen2.5:1.5b", etc. expansion_model: str = "auto" # Usually same as synthesis_model - max_expansion_terms: int = 8 # Maximum additional terms to add - enable_synthesis: bool = False # Enable by default when --synthesize used + max_expansion_terms: int = 8 # Maximum additional terms to add + enable_synthesis: bool = False # Enable by default when --synthesize used synthesis_temperature: float = 0.3 enable_thinking: bool = True # Enable thinking mode for Qwen3 models - cpu_optimized: bool = True # Prefer lightweight models - + cpu_optimized: bool = True # Prefer lightweight models + # Context window configuration (critical for RAG performance) - context_window: int = 16384 # Context window size in tokens (16K recommended) - auto_context: bool = True # Auto-adjust context based on model capabilities - + context_window: int = 16384 # Context window size in tokens (16K recommended) + auto_context: bool = True # Auto-adjust context based on model capabilities + # Model preference rankings (configurable) - model_rankings: list = None # Will be set in __post_init__ - + model_rankings: list = None # Will be set in __post_init__ + # Provider-specific settings (for different LLM providers) - provider: str = "ollama" # "ollama", "openai", "anthropic" + provider: str = "ollama" # "ollama", "openai", "anthropic" ollama_host: str = "localhost:11434" # Ollama connection api_key: Optional[str] = None # API key for cloud providers - api_base: Optional[str] = None # Base URL for API (e.g., OpenRouter) - timeout: int = 20 # Request timeout in seconds - + api_base: Optional[str] = None # Base URL for API (e.g., OpenRouter) + timeout: int = 20 # Request timeout in seconds + def __post_init__(self): if self.model_rankings is None: # Default model preference rankings (can be overridden in config file) self.model_rankings = [ # Testing model (prioritized for current testing phase) "qwen3:1.7b", - # Ultra-efficient models (perfect for CPU-only systems) - "qwen3:0.6b", - + "qwen3:0.6b", # Recommended model (excellent quality but larger) "qwen3:4b", - - # Common fallbacks (prioritize Qwen models) + # Common fallbacks (prioritize Qwen models) "qwen2.5:1.5b", "qwen2.5:3b", ] @@ -117,24 +121,26 @@ class LLMConfig: @dataclass class UpdateConfig: """Configuration for auto-update system.""" - auto_check: bool = True # Check for updates automatically + + auto_check: bool = True # Check for updates automatically check_frequency_hours: int = 24 # How often to check (hours) - auto_install: bool = False # Auto-install without asking (not recommended) - backup_before_update: bool = True # Create backup before updating - notify_beta_releases: bool = False # Include beta/pre-releases + auto_install: bool = False # Auto-install without asking (not recommended) + backup_before_update: bool = True # Create backup before updating + notify_beta_releases: bool = False # Include beta/pre-releases @dataclass class RAGConfig: """Main RAG system configuration.""" + chunking: ChunkingConfig = None - streaming: StreamingConfig = None + streaming: StreamingConfig = None files: FilesConfig = None embedding: EmbeddingConfig = None search: SearchConfig = None llm: LLMConfig = None updates: UpdateConfig = None - + def __post_init__(self): if self.chunking is None: self.chunking = ChunkingConfig() @@ -154,12 +160,12 @@ class RAGConfig: class ConfigManager: """Manages configuration loading, saving, and validation.""" - + def __init__(self, project_path: Path): self.project_path = Path(project_path) - self.rag_dir = self.project_path / '.mini-rag' - self.config_path = self.rag_dir / 'config.yaml' - + self.rag_dir = self.project_path / ".mini-rag" + self.config_path = self.rag_dir / "config.yaml" + def load_config(self) -> RAGConfig: """Load configuration from YAML file or create default.""" if not self.config_path.exists(): @@ -167,75 +173,81 @@ class ConfigManager: config = RAGConfig() self.save_config(config) return config - + try: - with open(self.config_path, 'r') as f: + with open(self.config_path, "r") as f: data = yaml.safe_load(f) - + if not data: logger.warning("Empty config file, using defaults") return RAGConfig() - + # Convert nested dicts back to dataclass instances config = RAGConfig() - - if 'chunking' in data: - config.chunking = ChunkingConfig(**data['chunking']) - if 'streaming' in data: - config.streaming = StreamingConfig(**data['streaming']) - if 'files' in data: - config.files = FilesConfig(**data['files']) - if 'embedding' in data: - config.embedding = EmbeddingConfig(**data['embedding']) - if 'search' in data: - config.search = SearchConfig(**data['search']) - if 'llm' in data: - config.llm = LLMConfig(**data['llm']) - + + if "chunking" in data: + config.chunking = ChunkingConfig(**data["chunking"]) + if "streaming" in data: + config.streaming = StreamingConfig(**data["streaming"]) + if "files" in data: + config.files = FilesConfig(**data["files"]) + if "embedding" in data: + config.embedding = EmbeddingConfig(**data["embedding"]) + if "search" in data: + config.search = SearchConfig(**data["search"]) + if "llm" in data: + config.llm = LLMConfig(**data["llm"]) + return config - + except yaml.YAMLError as e: # YAML syntax error - help user fix it instead of silent fallback - error_msg = f"โš ๏ธ Config file has YAML syntax error at line {getattr(e, 'problem_mark', 'unknown')}: {e}" + error_msg = ( + f"โš ๏ธ Config file has YAML syntax error at line " + f"{getattr(e, 'problem_mark', 'unknown')}: {e}" + ) logger.error(error_msg) print(f"\n{error_msg}") print(f"Config file: {self.config_path}") print("๐Ÿ’ก Check YAML syntax (indentation, quotes, colons)") print("๐Ÿ’ก Or delete config file to reset to defaults") return RAGConfig() # Still return defaults but warn user - + except Exception as e: logger.error(f"Failed to load config from {self.config_path}: {e}") logger.info("Using default configuration") return RAGConfig() - + def save_config(self, config: RAGConfig): """Save configuration to YAML file with comments.""" try: self.rag_dir.mkdir(exist_ok=True) - + # Convert to dict for YAML serialization config_dict = asdict(config) - + # Create YAML content with comments yaml_content = self._create_yaml_with_comments(config_dict) - + # Write with basic file locking to prevent corruption - with open(self.config_path, 'w') as f: + with open(self.config_path, "w") as f: try: import fcntl - fcntl.flock(f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) # Non-blocking exclusive lock + + fcntl.flock( + f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB + ) # Non-blocking exclusive lock f.write(yaml_content) fcntl.flock(f.fileno(), fcntl.LOCK_UN) # Unlock except (OSError, ImportError): # Fallback for Windows or if fcntl unavailable f.write(yaml_content) - + logger.info(f"Configuration saved to {self.config_path}") - + except Exception as e: logger.error(f"Failed to save config to {self.config_path}: {e}") - + def _create_yaml_with_comments(self, config_dict: Dict[str, Any]) -> str: """Create YAML content with helpful comments.""" yaml_lines = [ @@ -245,93 +257,97 @@ class ConfigManager: "", "# Text chunking settings", "chunking:", - f" max_size: {config_dict['chunking']['max_size']} # Maximum characters per chunk", - f" min_size: {config_dict['chunking']['min_size']} # Minimum characters per chunk", - f" strategy: {config_dict['chunking']['strategy']} # 'semantic' (language-aware) or 'fixed'", + f" max_size: {config_dict['chunking']['max_size']} # Max chars per chunk", + f" min_size: {config_dict['chunking']['min_size']} # Min chars per chunk", + f" strategy: {config_dict['chunking']['strategy']} # 'semantic' or 'fixed'", "", - "# Large file streaming settings", + "# Large file streaming settings", "streaming:", f" enabled: {str(config_dict['streaming']['enabled']).lower()}", - f" threshold_bytes: {config_dict['streaming']['threshold_bytes']} # Files larger than this use streaming (1MB)", + f" threshold_bytes: {config_dict['streaming']['threshold_bytes']} # Stream files >1MB", "", "# File processing settings", "files:", - f" min_file_size: {config_dict['files']['min_file_size']} # Skip files smaller than this", + f" min_file_size: {config_dict['files']['min_file_size']} # Skip small files", " exclude_patterns:", ] - - for pattern in config_dict['files']['exclude_patterns']: - yaml_lines.append(f" - \"{pattern}\"") - - yaml_lines.extend([ - " include_patterns:", - " - \"**/*\" # Include all files by default", - "", - "# Embedding generation settings", - "embedding:", - f" preferred_method: {config_dict['embedding']['preferred_method']} # 'ollama', 'ml', 'hash', or 'auto'", - f" ollama_model: {config_dict['embedding']['ollama_model']}", - f" ollama_host: {config_dict['embedding']['ollama_host']}", - f" ml_model: {config_dict['embedding']['ml_model']}", - f" batch_size: {config_dict['embedding']['batch_size']} # Embeddings processed per batch", - "", - "# Search behavior settings", - "search:", - f" default_top_k: {config_dict['search']['default_top_k']} # Default number of top results", - f" enable_bm25: {str(config_dict['search']['enable_bm25']).lower()} # Enable keyword matching boost", - f" similarity_threshold: {config_dict['search']['similarity_threshold']} # Minimum similarity score", - f" expand_queries: {str(config_dict['search']['expand_queries']).lower()} # Enable automatic query expansion", - "", - "# LLM synthesis and query expansion settings", - "llm:", - f" ollama_host: {config_dict['llm']['ollama_host']}", - f" synthesis_model: {config_dict['llm']['synthesis_model']} # 'auto', 'qwen3:1.7b', etc.", - f" expansion_model: {config_dict['llm']['expansion_model']} # Usually same as synthesis_model", - f" max_expansion_terms: {config_dict['llm']['max_expansion_terms']} # Maximum terms to add to queries", - f" enable_synthesis: {str(config_dict['llm']['enable_synthesis']).lower()} # Enable synthesis by default", - f" synthesis_temperature: {config_dict['llm']['synthesis_temperature']} # LLM temperature for analysis", - "", - " # Context window configuration (critical for RAG performance)", - " # ๐Ÿ’ก Sizing guide: 2K=1 question, 4K=1-2 questions, 8K=manageable, 16K=most users", - " # 32K=large codebases, 64K+=power users only", - " # โš ๏ธ Larger contexts use exponentially more CPU/memory - only increase if needed", - " # ๐Ÿ”ง Low context limits? Try smaller topk, better search terms, or archive noise", - f" context_window: {config_dict['llm']['context_window']} # Context size in tokens", - f" auto_context: {str(config_dict['llm']['auto_context']).lower()} # Auto-adjust context based on model capabilities", - "", - " model_rankings: # Preferred model order (edit to change priority)", - ]) - + + for pattern in config_dict["files"]["exclude_patterns"]: + yaml_lines.append(f' - "{pattern}"') + + yaml_lines.extend( + [ + " include_patterns:", + ' - "**/*" # Include all files by default', + "", + "# Embedding generation settings", + "embedding:", + f" preferred_method: {config_dict['embedding']['preferred_method']} # Method", + f" ollama_model: {config_dict['embedding']['ollama_model']}", + f" ollama_host: {config_dict['embedding']['ollama_host']}", + f" ml_model: {config_dict['embedding']['ml_model']}", + f" batch_size: {config_dict['embedding']['batch_size']} # Per batch", + "", + "# Search behavior settings", + "search:", + f" default_top_k: {config_dict['search']['default_top_k']} # Top results", + f" enable_bm25: {str(config_dict['search']['enable_bm25']).lower()} # Keyword boost", + f" similarity_threshold: {config_dict['search']['similarity_threshold']} # Min score", + f" expand_queries: {str(config_dict['search']['expand_queries']).lower()} # Auto expand", + "", + "# LLM synthesis and query expansion settings", + "llm:", + f" ollama_host: {config_dict['llm']['ollama_host']}", + f" synthesis_model: {config_dict['llm']['synthesis_model']} # Model name", + f" expansion_model: {config_dict['llm']['expansion_model']} # Model name", + f" max_expansion_terms: {config_dict['llm']['max_expansion_terms']} # Max terms", + f" enable_synthesis: {str(config_dict['llm']['enable_synthesis']).lower()} # Enable synthesis by default", + f" synthesis_temperature: {config_dict['llm']['synthesis_temperature']} # LLM temperature for analysis", + "", + " # Context window configuration (critical for RAG performance)", + " # ๐Ÿ’ก Sizing guide: 2K=1 question, 4K=1-2 questions, 8K=manageable, 16K=most users", + " # 32K=large codebases, 64K+=power users only", + " # โš ๏ธ Larger contexts use exponentially more CPU/memory - only increase if needed", + " # ๐Ÿ”ง Low context limits? Try smaller topk, better search terms, or archive noise", + f" context_window: {config_dict['llm']['context_window']} # Context size in tokens", + f" auto_context: {str(config_dict['llm']['auto_context']).lower()} # Auto-adjust context based on model capabilities", + "", + " model_rankings: # Preferred model order (edit to change priority)", + ] + ) + # Add model rankings list - if 'model_rankings' in config_dict['llm'] and config_dict['llm']['model_rankings']: - for model in config_dict['llm']['model_rankings'][:10]: # Show first 10 - yaml_lines.append(f" - \"{model}\"") - if len(config_dict['llm']['model_rankings']) > 10: + if "model_rankings" in config_dict["llm"] and config_dict["llm"]["model_rankings"]: + for model in config_dict["llm"]["model_rankings"][:10]: # Show first 10 + yaml_lines.append(f' - "{model}"') + if len(config_dict["llm"]["model_rankings"]) > 10: yaml_lines.append(" # ... (edit config to see all options)") - + # Add update settings - yaml_lines.extend([ - "", - "# Auto-update system settings", - "updates:", - f" auto_check: {str(config_dict['updates']['auto_check']).lower()} # Check for updates automatically", - f" check_frequency_hours: {config_dict['updates']['check_frequency_hours']} # Hours between update checks", - f" auto_install: {str(config_dict['updates']['auto_install']).lower()} # Auto-install updates (not recommended)", - f" backup_before_update: {str(config_dict['updates']['backup_before_update']).lower()} # Create backup before updating", - f" notify_beta_releases: {str(config_dict['updates']['notify_beta_releases']).lower()} # Include beta releases in checks", - ]) - - return '\n'.join(yaml_lines) - + yaml_lines.extend( + [ + "", + "# Auto-update system settings", + "updates:", + f" auto_check: {str(config_dict['updates']['auto_check']).lower()} # Check for updates automatically", + f" check_frequency_hours: {config_dict['updates']['check_frequency_hours']} # Hours between update checks", + f" auto_install: {str(config_dict['updates']['auto_install']).lower()} # Auto-install updates (not recommended)", + f" backup_before_update: {str(config_dict['updates']['backup_before_update']).lower()} # Create backup before updating", + f" notify_beta_releases: {str(config_dict['updates']['notify_beta_releases']).lower()} # Include beta releases in checks", + ] + ) + + return "\n".join(yaml_lines) + def update_config(self, **kwargs) -> RAGConfig: """Update specific configuration values.""" config = self.load_config() - + for key, value in kwargs.items(): if hasattr(config, key): setattr(config, key, value) else: logger.warning(f"Unknown config key: {key}") - + self.save_config(config) - return config \ No newline at end of file + return config diff --git a/mini_rag/explorer.py b/mini_rag/explorer.py index be5b3c8..949dda6 100644 --- a/mini_rag/explorer.py +++ b/mini_rag/explorer.py @@ -9,158 +9,169 @@ Perfect for exploring codebases with detailed reasoning and follow-up questions. import json import logging import time -from typing import List, Dict, Any, Optional -from pathlib import Path from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional try: + from .config import RAGConfig from .llm_synthesizer import LLMSynthesizer, SynthesisResult from .search import CodeSearcher - from .config import RAGConfig from .system_context import get_system_context except ImportError: # For direct testing + from config import RAGConfig from llm_synthesizer import LLMSynthesizer, SynthesisResult from search import CodeSearcher - from config import RAGConfig - get_system_context = lambda x=None: "" + + def get_system_context(x=None): + return "" + logger = logging.getLogger(__name__) + @dataclass class ExplorationSession: """Track an exploration session with context history.""" + project_path: Path conversation_history: List[Dict[str, Any]] session_id: str started_at: float - - def add_exchange(self, question: str, search_results: List[Any], response: SynthesisResult): + + def add_exchange( + self, question: str, search_results: List[Any], response: SynthesisResult + ): """Add a question/response exchange to the conversation history.""" - self.conversation_history.append({ - "timestamp": time.time(), - "question": question, - "search_results_count": len(search_results), - "response": { - "summary": response.summary, - "key_points": response.key_points, - "code_examples": response.code_examples, - "suggested_actions": response.suggested_actions, - "confidence": response.confidence + self.conversation_history.append( + { + "timestamp": time.time(), + "question": question, + "search_results_count": len(search_results), + "response": { + "summary": response.summary, + "key_points": response.key_points, + "code_examples": response.code_examples, + "suggested_actions": response.suggested_actions, + "confidence": response.confidence, + }, } - }) + ) + class CodeExplorer: """Interactive code exploration with thinking and context memory.""" - + def __init__(self, project_path: Path, config: RAGConfig = None): self.project_path = project_path self.config = config or RAGConfig() - + # Initialize components with thinking enabled self.searcher = CodeSearcher(project_path) self.synthesizer = LLMSynthesizer( ollama_url=f"http://{self.config.llm.ollama_host}", model=self.config.llm.synthesis_model, enable_thinking=True, # Always enable thinking in explore mode - config=self.config # Pass config for model rankings + config=self.config, # Pass config for model rankings ) - + # Session management self.current_session: Optional[ExplorationSession] = None - + def start_exploration_session(self) -> bool: """Start a new exploration session.""" - + # Simple availability check - don't do complex model restart logic if not self.synthesizer.is_available(): print("โŒ LLM service unavailable. Please check Ollama is running.") return False - + session_id = f"explore_{int(time.time())}" self.current_session = ExplorationSession( project_path=self.project_path, conversation_history=[], session_id=session_id, - started_at=time.time() + started_at=time.time(), ) - + print("๐Ÿง  Exploration Mode Started") print(f"Project: {self.project_path.name}") - + return True - + def explore_question(self, question: str, context_limit: int = 10) -> Optional[str]: """Explore a question with full thinking and context.""" if not self.current_session: return "โŒ No exploration session active. Start one first." - + # Search for relevant information search_start = time.time() results = self.searcher.search( - question, + question, top_k=context_limit, include_context=True, semantic_weight=0.7, - bm25_weight=0.3 + bm25_weight=0.3, ) search_time = time.time() - search_start - + # Build enhanced prompt with conversation context synthesis_prompt = self._build_contextual_prompt(question, results) - + # Get thinking-enabled analysis synthesis_start = time.time() synthesis = self._synthesize_with_context(synthesis_prompt, results) synthesis_time = time.time() - synthesis_start - + # Add to conversation history self.current_session.add_exchange(question, results, synthesis) - + # Streaming already displayed the response # Just return minimal status for caller session_duration = time.time() - self.current_session.started_at exchange_count = len(self.current_session.conversation_history) - + status = f"\n๐Ÿ“Š Session: {session_duration/60:.1f}m | Question #{exchange_count} | Results: {len(results)} | Time: {search_time+synthesis_time:.1f}s" return status - + def _build_contextual_prompt(self, question: str, results: List[Any]) -> str: """Build a prompt that includes conversation context.""" # Get recent conversation context (last 3 exchanges) - context_summary = "" if self.current_session.conversation_history: recent_exchanges = self.current_session.conversation_history[-3:] context_parts = [] - + for i, exchange in enumerate(recent_exchanges, 1): prev_q = exchange["question"] prev_summary = exchange["response"]["summary"] context_parts.append(f"Previous Q{i}: {prev_q}") context_parts.append(f"Previous A{i}: {prev_summary}") - - context_summary = "\n".join(context_parts) - + + # "\n".join(context_parts) # Unused variable removed + # Build search results context results_context = [] for i, result in enumerate(results[:8], 1): - file_path = result.file_path if hasattr(result, 'file_path') else 'unknown' - content = result.content if hasattr(result, 'content') else str(result) - score = result.score if hasattr(result, 'score') else 0.0 - - results_context.append(f""" + # result.file_path if hasattr(result, "file_path") else "unknown" # Unused variable removed + # result.content if hasattr(result, "content") else str(result) # Unused variable removed + # result.score if hasattr(result, "score") else 0.0 # Unused variable removed + + results_context.append( + """ Result {i} (Score: {score:.3f}): File: {file_path} Content: {content[:800]}{'...' if len(content) > 800 else ''} -""") - - results_text = "\n".join(results_context) - +""" + ) + + # "\n".join(results_context) # Unused variable removed + # Get system context for better responses - system_context = get_system_context(self.project_path) - + # get_system_context(self.project_path) # Unused variable removed + # Create comprehensive exploration prompt with thinking - prompt = f""" + prompt = """ The user asked: "{question}" System context: {system_context} @@ -197,7 +208,7 @@ Please provide a helpful, natural explanation that answers their question. Write Structure your response to include: 1. A clear explanation of what you found and how it answers their question -2. The most important insights from the information you discovered +2. The most important insights from the information you discovered 3. Relevant examples or code patterns when helpful 4. Practical next steps they could take @@ -210,37 +221,43 @@ Guidelines: - Use natural language, not structured formats - Break complex topics into understandable pieces """ - + return prompt - + def _synthesize_with_context(self, prompt: str, results: List[Any]) -> SynthesisResult: """Synthesize results with full context and thinking.""" try: # Use streaming with thinking visible (don't collapse) - response = self.synthesizer._call_ollama(prompt, temperature=0.2, disable_thinking=False, use_streaming=True, collapse_thinking=False) - thinking_stream = "" - + response = self.synthesizer._call_ollama( + prompt, + temperature=0.2, + disable_thinking=False, + use_streaming=True, + collapse_thinking=False, + ) + # "" # Unused variable removed + # Streaming already shows thinking and response # No need for additional indicators - + if not response: return SynthesisResult( summary="Analysis unavailable (LLM service error)", key_points=[], code_examples=[], suggested_actions=["Check LLM service status"], - confidence=0.0 + confidence=0.0, ) - + # Use natural language response directly return SynthesisResult( summary=response.strip(), key_points=[], # Not used with natural language responses code_examples=[], # Not used with natural language responses suggested_actions=[], # Not used with natural language responses - confidence=0.85 # High confidence for natural responses + confidence=0.85, # High confidence for natural responses ) - + except Exception as e: logger.error(f"Context synthesis failed: {e}") return SynthesisResult( @@ -248,124 +265,153 @@ Guidelines: key_points=[], code_examples=[], suggested_actions=["Check system status and try again"], - confidence=0.0 + confidence=0.0, ) - - def _format_exploration_response(self, question: str, synthesis: SynthesisResult, - result_count: int, search_time: float, synthesis_time: float) -> str: + + def _format_exploration_response( + self, + question: str, + synthesis: SynthesisResult, + result_count: int, + search_time: float, + synthesis_time: float, + ) -> str: """Format exploration response with context indicators.""" - + output = [] - + # Header with session context session_duration = time.time() - self.current_session.started_at exchange_count = len(self.current_session.conversation_history) - + output.append(f"๐Ÿง  EXPLORATION ANALYSIS (Question #{exchange_count})") - output.append(f"Session: {session_duration/60:.1f}m | Results: {result_count} | " - f"Time: {search_time+synthesis_time:.1f}s") + output.append( + f"Session: {session_duration/60:.1f}m | Results: {result_count} | " + f"Time: {search_time+synthesis_time:.1f}s" + ) output.append("=" * 60) output.append("") - + # Response was already displayed via streaming # Just show completion status output.append("โœ… Analysis complete") output.append("") output.append("") - + # Confidence and context indicator - confidence_emoji = "๐ŸŸข" if synthesis.confidence > 0.7 else "๐ŸŸก" if synthesis.confidence > 0.4 else "๐Ÿ”ด" - context_indicator = f" | Context: {exchange_count-1} previous questions" if exchange_count > 1 else "" - output.append(f"{confidence_emoji} Confidence: {synthesis.confidence:.1%}{context_indicator}") - + confidence_emoji = ( + "๐ŸŸข" + if synthesis.confidence > 0.7 + else "๐ŸŸก" if synthesis.confidence > 0.4 else "๐Ÿ”ด" + ) + context_indicator = ( + f" | Context: {exchange_count-1} previous questions" if exchange_count > 1 else "" + ) + output.append( + f"{confidence_emoji} Confidence: {synthesis.confidence:.1%}{context_indicator}" + ) + return "\n".join(output) - + def get_session_summary(self) -> str: """Get a summary of the current exploration session.""" if not self.current_session: return "No active exploration session." - + duration = time.time() - self.current_session.started_at exchange_count = len(self.current_session.conversation_history) - + summary = [ - f"๐Ÿง  EXPLORATION SESSION SUMMARY", - f"=" * 40, + "๐Ÿง  EXPLORATION SESSION SUMMARY", + "=" * 40, f"Project: {self.project_path.name}", f"Session ID: {self.current_session.session_id}", f"Duration: {duration/60:.1f} minutes", f"Questions explored: {exchange_count}", - f"", + "", ] - + if exchange_count > 0: summary.append("๐Ÿ“‹ Topics explored:") for i, exchange in enumerate(self.current_session.conversation_history, 1): - question = exchange["question"][:50] + "..." if len(exchange["question"]) > 50 else exchange["question"] + question = ( + exchange["question"][:50] + "..." + if len(exchange["question"]) > 50 + else exchange["question"] + ) confidence = exchange["response"]["confidence"] summary.append(f" {i}. {question} (confidence: {confidence:.1%})") - + return "\n".join(summary) - + def end_session(self) -> str: """End the current exploration session.""" if not self.current_session: return "No active session to end." - + summary = self.get_session_summary() self.current_session = None - + return summary + "\n\nโœ… Exploration session ended." - + def _check_model_restart_needed(self) -> bool: """Check if model restart would improve thinking quality.""" try: - # Simple heuristic: if we can detect the model was recently used + # Simple heuristic: if we can detect the model was recently used # with , suggest restart for better thinking quality - + # Test with a simple thinking prompt to see response quality test_response = self.synthesizer._call_ollama( - "Think briefly: what is 2+2?", - temperature=0.1, - disable_thinking=False + "Think briefly: what is 2+2?", temperature=0.1, disable_thinking=False ) - + if test_response: # If response is suspiciously short or shows signs of no-think behavior if len(test_response.strip()) < 10 or "4" == test_response.strip(): return True - + except Exception: pass - + return False - + def _handle_model_restart(self) -> bool: """Handle user confirmation and model restart.""" try: - print("\n๐Ÿค” To ensure best thinking quality, exploration mode works best with a fresh model.") + print( + "\n๐Ÿค” To ensure best thinking quality, exploration mode works best with a fresh model." + ) print(f" Currently running: {self.synthesizer.model}") - print("\n๐Ÿ’ก Stop current model and restart for optimal exploration? (y/N): ", end="", flush=True) - + print( + "\n๐Ÿ’ก Stop current model and restart for optimal exploration? (y/N): ", + end="", + flush=True, + ) + response = input().strip().lower() - - if response in ['y', 'yes']: + + if response in ["y", "yes"]: print("\n๐Ÿ”„ Stopping current model...") - + # Use ollama stop command for clean model restart import subprocess + try: - subprocess.run([ - "ollama", "stop", self.synthesizer.model - ], timeout=10, capture_output=True) - + subprocess.run( + ["ollama", "stop", self.synthesizer.model], + timeout=10, + capture_output=True, + ) + print("โœ… Model stopped successfully.") - print("๐Ÿš€ Exploration mode will restart the model with thinking enabled...") - + print( + "๐Ÿš€ Exploration mode will restart the model with thinking enabled..." + ) + # Reset synthesizer initialization to force fresh start self.synthesizer._initialized = False return True - + except subprocess.TimeoutExpired: print("โš ๏ธ Model stop timed out, continuing anyway...") return False @@ -378,19 +424,18 @@ Guidelines: else: print("๐Ÿ“ Continuing with current model...") return False - + except KeyboardInterrupt: print("\n๐Ÿ“ Continuing with current model...") return False except EOFError: print("\n๐Ÿ“ Continuing with current model...") return False - + def _call_ollama_with_thinking(self, prompt: str, temperature: float = 0.3) -> tuple: """Call Ollama with streaming for fast time-to-first-token.""" import requests - import json - + try: # Use the synthesizer's model and connection model_to_use = self.synthesizer.model @@ -399,14 +444,15 @@ Guidelines: model_to_use = self.synthesizer.available_models[0] else: return None, None - + # Enable thinking by NOT adding final_prompt = prompt - + # Get optimal parameters for this model from .llm_optimization import get_optimal_ollama_parameters + optimal_params = get_optimal_ollama_parameters(model_to_use) - + payload = { "model": model_to_use, "prompt": final_prompt, @@ -418,94 +464,102 @@ Guidelines: "num_ctx": self.synthesizer._get_optimal_context_size(model_to_use), "num_predict": optimal_params.get("num_predict", 2000), "repeat_penalty": optimal_params.get("repeat_penalty", 1.1), - "presence_penalty": optimal_params.get("presence_penalty", 1.0) - } + "presence_penalty": optimal_params.get("presence_penalty", 1.0), + }, } - + response = requests.post( f"{self.synthesizer.ollama_url}/api/generate", json=payload, stream=True, - timeout=65 + timeout=65, ) - + if response.status_code == 200: # Collect streaming response raw_response = "" thinking_displayed = False - + for line in response.iter_lines(): if line: try: - chunk_data = json.loads(line.decode('utf-8')) - chunk_text = chunk_data.get('response', '') - + chunk_data = json.loads(line.decode("utf-8")) + chunk_text = chunk_data.get("response", "") + if chunk_text: raw_response += chunk_text - + # Display thinking stream as it comes in - if not thinking_displayed and '' in raw_response: + if not thinking_displayed and "" in raw_response: # Start displaying thinking self._start_thinking_display() thinking_displayed = True - + if thinking_displayed: self._stream_thinking_chunk(chunk_text) - - if chunk_data.get('done', False): + + if chunk_data.get("done", False): break - + except json.JSONDecodeError: continue - + # Finish thinking display if it was shown if thinking_displayed: self._end_thinking_display() - + # Extract thinking stream and final response thinking_stream, final_response = self._extract_thinking(raw_response) - + return final_response, thinking_stream else: return None, None - + except Exception as e: logger.error(f"Thinking-enabled Ollama call failed: {e}") return None, None - + def _extract_thinking(self, raw_response: str) -> tuple: """Extract thinking content from response.""" thinking_stream = "" final_response = raw_response - + # Look for thinking patterns if "" in raw_response and "" in raw_response: # Extract thinking content between tags start_tag = raw_response.find("") end_tag = raw_response.find("") + len("") - + if start_tag != -1 and end_tag != -1: - thinking_content = raw_response[start_tag + 7:end_tag - 8] # Remove tags + thinking_content = raw_response[start_tag + 7 : end_tag - 8] # Remove tags thinking_stream = thinking_content.strip() - + # Remove thinking from final response final_response = (raw_response[:start_tag] + raw_response[end_tag:]).strip() - + # Alternative patterns for models that use different thinking formats elif "Let me think" in raw_response or "I need to analyze" in raw_response: # Simple heuristic: first paragraph might be thinking - lines = raw_response.split('\n') + lines = raw_response.split("\n") potential_thinking = [] final_lines = [] - - thinking_indicators = ["Let me think", "I need to", "First, I'll", "Looking at", "Analyzing"] + + thinking_indicators = [ + "Let me think", + "I need to", + "First, I'll", + "Looking at", + "Analyzing", + ] in_thinking = False - + for line in lines: if any(indicator in line for indicator in thinking_indicators): in_thinking = True potential_thinking.append(line) - elif in_thinking and (line.startswith('{') or line.startswith('**') or line.startswith('#')): + elif in_thinking and ( + line.startswith("{") or line.startswith("**") or line.startswith("#") + ): # Likely end of thinking, start of structured response in_thinking = False final_lines.append(line) @@ -513,84 +567,87 @@ Guidelines: potential_thinking.append(line) else: final_lines.append(line) - + if potential_thinking: - thinking_stream = '\n'.join(potential_thinking).strip() - final_response = '\n'.join(final_lines).strip() - + thinking_stream = "\n".join(potential_thinking).strip() + final_response = "\n".join(final_lines).strip() + return thinking_stream, final_response - + def _start_thinking_display(self): """Start the thinking stream display.""" print("\n\033[2m\033[3m๐Ÿ’ญ AI Thinking:\033[0m") print("\033[2m\033[3m" + "โ”€" * 40 + "\033[0m") self._thinking_buffer = "" self._in_thinking_tags = False - + def _stream_thinking_chunk(self, chunk: str): """Stream a chunk of thinking as it arrives.""" - import sys - + self._thinking_buffer += chunk - + # Check if we're in thinking tags - if '' in self._thinking_buffer and not self._in_thinking_tags: + if "" in self._thinking_buffer and not self._in_thinking_tags: self._in_thinking_tags = True # Display everything after - start_idx = self._thinking_buffer.find('') + 7 + start_idx = self._thinking_buffer.find("") + 7 thinking_content = self._thinking_buffer[start_idx:] if thinking_content: - print(f"\033[2m\033[3m{thinking_content}\033[0m", end='', flush=True) - elif self._in_thinking_tags and '' not in chunk: + print(f"\033[2m\033[3m{thinking_content}\033[0m", end="", flush=True) + elif self._in_thinking_tags and "" not in chunk: # We're in thinking mode, display the chunk - print(f"\033[2m\033[3m{chunk}\033[0m", end='', flush=True) - elif '' in self._thinking_buffer: + print(f"\033[2m\033[3m{chunk}\033[0m", end="", flush=True) + elif "" in self._thinking_buffer: # End of thinking self._in_thinking_tags = False - + def _end_thinking_display(self): """End the thinking stream display.""" - print(f"\n\033[2m\033[3m" + "โ”€" * 40 + "\033[0m") + print("\n\033[2m\033[3m" + "โ”€" * 40 + "\033[0m") print() - + def _display_thinking_stream(self, thinking_stream: str): """Display thinking stream in light gray and italic (fallback for non-streaming).""" if not thinking_stream: return - + print("\n\033[2m\033[3m๐Ÿ’ญ AI Thinking:\033[0m") print("\033[2m\033[3m" + "โ”€" * 40 + "\033[0m") - + # Split into paragraphs and display with proper formatting - paragraphs = thinking_stream.split('\n\n') + paragraphs = thinking_stream.split("\n\n") for para in paragraphs: if para.strip(): # Wrap long lines nicely - lines = para.strip().split('\n') + lines = para.strip().split("\n") for line in lines: if line.strip(): # Light gray and italic print(f"\033[2m\033[3m{line}\033[0m") print() # Paragraph spacing - + print("\033[2m\033[3m" + "โ”€" * 40 + "\033[0m") print() + # Quick test function + + def test_explorer(): """Test the code explorer.""" explorer = CodeExplorer(Path(".")) - + if not explorer.start_exploration_session(): print("โŒ Could not start exploration session") return - + # Test question response = explorer.explore_question("How does authentication work in this codebase?") if response: print(response) - + print("\n" + explorer.end_session()) + if __name__ == "__main__": - test_explorer() \ No newline at end of file + test_explorer() diff --git a/mini_rag/fast_server.py b/mini_rag/fast_server.py index 940e9df..5d7e62b 100644 --- a/mini_rag/fast_server.py +++ b/mini_rag/fast_server.py @@ -12,40 +12,47 @@ Drop-in replacement for the original server with: """ import json +import logging +import os import socket -import threading -import time import subprocess import sys -import os -import logging +import threading +import time +from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path -from typing import Dict, Any, Optional, Callable -from datetime import datetime -from concurrent.futures import ThreadPoolExecutor, Future -import queue +from typing import Any, Callable, Dict, Optional + +from rich import print as rprint # Rich console for beautiful output from rich.console import Console -from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn, MofNCompleteColumn -from rich.panel import Panel -from rich.table import Table from rich.live import Live -from rich import print as rprint +from rich.panel import Panel +from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + SpinnerColumn, + TextColumn, + TimeRemainingColumn, +) +from rich.table import Table # Fix Windows console first -if sys.platform == 'win32': - os.environ['PYTHONUTF8'] = '1' +if sys.platform == "win32": + os.environ["PYTHONUTF8"] = "1" try: from .windows_console_fix import fix_windows_console + fix_windows_console() - except: + except (ImportError, OSError): pass -from .search import CodeSearcher -from .ollama_embeddings import OllamaEmbedder as CodeEmbedder from .indexer import ProjectIndexer +from .ollama_embeddings import OllamaEmbedder as CodeEmbedder from .performance import PerformanceMonitor +from .search import CodeSearcher logger = logging.getLogger(__name__) console = Console() @@ -53,7 +60,7 @@ console = Console() class ServerStatus: """Real-time server status tracking""" - + def __init__(self): self.phase = "initializing" self.progress = 0.0 @@ -63,7 +70,7 @@ class ServerStatus: self.ready = False self.error = None self.health_checks = {} - + def update(self, phase: str, progress: float = None, message: str = None, **details): """Update server status""" self.phase = phase @@ -72,42 +79,42 @@ class ServerStatus: if message: self.message = message self.details.update(details) - + def set_ready(self): """Mark server as ready""" self.ready = True self.phase = "ready" self.progress = 100.0 self.message = "Server ready and accepting connections" - + def set_error(self, error: str): """Mark server as failed""" self.error = error self.phase = "failed" self.message = f"Server failed: {error}" - + def get_status(self) -> Dict[str, Any]: """Get complete status as dict""" return { - 'phase': self.phase, - 'progress': self.progress, - 'message': self.message, - 'ready': self.ready, - 'error': self.error, - 'uptime': time.time() - self.start_time, - 'health_checks': self.health_checks, - 'details': self.details + "phase": self.phase, + "progress": self.progress, + "message": self.message, + "ready": self.ready, + "error": self.error, + "uptime": time.time() - self.start_time, + "health_checks": self.health_checks, + "details": self.details, } class FastRAGServer: """Ultra-fast RAG server with enhanced feedback and monitoring""" - + def __init__(self, project_path: Path, port: int = 7777, auto_index: bool = True): self.project_path = project_path self.port = port self.auto_index = auto_index - + # Server state self.searcher = None self.embedder = None @@ -115,24 +122,24 @@ class FastRAGServer: self.running = False self.socket = None self.query_count = 0 - + # Status and monitoring self.status = ServerStatus() self.performance = PerformanceMonitor() self.health_check_interval = 30 # seconds self.last_health_check = 0 - + # Threading self.executor = ThreadPoolExecutor(max_workers=3) self.status_callbacks = [] - + # Progress tracking self.indexing_progress = None - + def add_status_callback(self, callback: Callable[[Dict], None]): """Add callback for status updates""" self.status_callbacks.append(callback) - + def _notify_status(self): """Notify all status callbacks""" status = self.status.get_status() @@ -141,112 +148,119 @@ class FastRAGServer: callback(status) except Exception as e: logger.warning(f"Status callback failed: {e}") - + def _kill_existing_server(self) -> bool: """Kill any existing process using our port with better feedback""" try: self.status.update("port_check", 5, "Checking for existing servers...") self._notify_status() - + # Quick port check first test_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) test_sock.settimeout(1.0) # Faster timeout - result = test_sock.connect_ex(('localhost', self.port)) + result = test_sock.connect_ex(("localhost", self.port)) test_sock.close() - + if result != 0: # Port is free return True - + console.print(f"[yellow]โš ๏ธ Port {self.port} is occupied, clearing it...[/yellow]") self.status.update("port_cleanup", 10, f"Clearing port {self.port}...") self._notify_status() - - if sys.platform == 'win32': + + if sys.platform == "win32": # Windows: Enhanced process killing - cmd = ['netstat', '-ano'] + cmd = ["netstat", "-ano"] result = subprocess.run(cmd, capture_output=True, text=True, timeout=5) - - for line in result.stdout.split('\n'): - if f':{self.port}' in line and 'LISTENING' in line: + + for line in result.stdout.split("\n"): + if f":{self.port}" in line and "LISTENING" in line: parts = line.split() if len(parts) >= 5: pid = parts[-1] console.print(f"[dim]Killing process {pid}[/dim]") - subprocess.run(['taskkill', '/PID', pid, '/F'], - capture_output=True, timeout=3) + subprocess.run( + ["taskkill", "/PID", pid, "/F"], + capture_output=True, + timeout=3, + ) time.sleep(0.5) # Reduced wait time break else: # Unix/Linux: Enhanced process killing - result = subprocess.run(['lsof', '-ti', f':{self.port}'], - capture_output=True, text=True, timeout=3) + result = subprocess.run( + ["lso", "-ti", f":{self.port}"], + capture_output=True, + text=True, + timeout=3, + ) if result.stdout.strip(): pids = result.stdout.strip().split() for pid in pids: console.print(f"[dim]Killing process {pid}[/dim]") - subprocess.run(['kill', '-9', pid], capture_output=True) + subprocess.run(["kill", "-9", pid], capture_output=True) time.sleep(0.5) - + # Verify port is free test_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) test_sock.settimeout(1.0) - result = test_sock.connect_ex(('localhost', self.port)) + result = test_sock.connect_ex(("localhost", self.port)) test_sock.close() - + if result == 0: raise RuntimeError(f"Failed to free port {self.port}") - + console.print(f"[green]โœ… Port {self.port} cleared[/green]") return True - + except subprocess.TimeoutExpired: raise RuntimeError("Timeout while clearing port") except Exception as e: raise RuntimeError(f"Failed to clear port {self.port}: {e}") - + def _check_indexing_needed(self) -> bool: """Quick check if indexing is needed""" - rag_dir = self.project_path / '.mini-rag' + rag_dir = self.project_path / ".mini-rag" if not rag_dir.exists(): return True - + # Check if database exists and is not empty - db_path = rag_dir / 'code_vectors.lance' + db_path = rag_dir / "code_vectors.lance" if not db_path.exists(): return True - + # Quick file count check try: import lancedb except ImportError: # If LanceDB not available, assume index is empty and needs creation return True - + try: db = lancedb.connect(rag_dir) - if 'code_vectors' not in db.table_names(): + if "code_vectors" not in db.table_names(): return True - table = db.open_table('code_vectors') + table = db.open_table("code_vectors") count = table.count_rows() return count == 0 - except: + except (OSError, IOError, ValueError, AttributeError): return True - + def _fast_index(self) -> bool: """Fast indexing with enhanced progress reporting""" try: self.status.update("indexing", 20, "Initializing indexer...") self._notify_status() - + # Create indexer with optimized settings self.indexer = ProjectIndexer( self.project_path, embedder=self.embedder, # Reuse loaded embedder - max_workers=min(4, os.cpu_count() or 2) + max_workers=min(4, os.cpu_count() or 2), ) - + console.print("\n[bold cyan]๐Ÿš€ Fast Indexing Starting...[/bold cyan]") - + with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), @@ -256,88 +270,99 @@ class FastRAGServer: console=console, refresh_per_second=10, # More responsive updates ) as progress: - + # Override indexer's progress reporting original_index_project = self.indexer.index_project - + def enhanced_index_project(*args, **kwargs): # Get files to index first files_to_index = self.indexer._get_files_to_index() total_files = len(files_to_index) - + if total_files == 0: self.status.update("indexing", 80, "Index up to date") - return {'files_indexed': 0, 'chunks_created': 0, 'time_taken': 0} - + return { + "files_indexed": 0, + "chunks_created": 0, + "time_taken": 0, + } + task = progress.add_task( - f"[cyan]Indexing {total_files} files...", - total=total_files + f"[cyan]Indexing {total_files} files...", total=total_files ) - + # Track progress by hooking into the processor processed_count = 0 - + def track_progress(): nonlocal processed_count while processed_count < total_files and self.running: time.sleep(0.1) # Fast polling current_progress = (processed_count / total_files) * 60 + 20 - self.status.update("indexing", current_progress, - f"Indexed {processed_count}/{total_files} files") + self.status.update( + "indexing", + current_progress, + f"Indexed {processed_count}/{total_files} files", + ) progress.update(task, completed=processed_count) self._notify_status() - + # Start progress tracking progress_thread = threading.Thread(target=track_progress) progress_thread.daemon = True progress_thread.start() - + # Hook into the processing original_process_file = self.indexer._process_file - + def tracked_process_file(*args, **kwargs): nonlocal processed_count result = original_process_file(*args, **kwargs) processed_count += 1 return result - + self.indexer._process_file = tracked_process_file - + # Run the actual indexing stats = original_index_project(*args, **kwargs) - + progress.update(task, completed=total_files) return stats - + self.indexer.index_project = enhanced_index_project - + # Run indexing stats = self.indexer.index_project(force_reindex=False) - - self.status.update("indexing", 80, - f"Indexed {stats.get('files_indexed', 0)} files, " - f"created {stats.get('chunks_created', 0)} chunks") + + self.status.update( + "indexing", + 80, + f"Indexed {stats.get('files_indexed', 0)} files, " + f"created {stats.get('chunks_created', 0)} chunks", + ) self._notify_status() - - console.print(f"\n[green]โœ… Indexing complete: {stats.get('files_indexed', 0)} files, " - f"{stats.get('chunks_created', 0)} chunks in {stats.get('time_taken', 0):.1f}s[/green]") - + + console.print( + f"\n[green]โœ… Indexing complete: {stats.get('files_indexed', 0)} files, " + f"{stats.get('chunks_created', 0)} chunks in {stats.get('time_taken', 0):.1f}s[/green]" + ) + return True - + except Exception as e: self.status.set_error(f"Indexing failed: {e}") self._notify_status() console.print(f"[red]โŒ Indexing failed: {e}[/red]") return False - + def _initialize_components(self) -> bool: """Fast parallel component initialization""" try: console.print("\n[bold blue]๐Ÿ”ง Initializing RAG Server...[/bold blue]") - + # Check if indexing is needed first needs_indexing = self._check_indexing_needed() - + with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), @@ -345,10 +370,12 @@ class FastRAGServer: TimeRemainingColumn(), console=console, ) as progress: - + # Task 1: Load embedder (this takes the most time) - embedder_task = progress.add_task("[cyan]Loading embedding model...", total=100) - + embedder_task = progress.add_task( + "[cyan]Loading embedding model...", total=100 + ) + def load_embedder(): self.status.update("embedder", 25, "Loading embedding model...") self._notify_status() @@ -357,155 +384,155 @@ class FastRAGServer: progress.update(embedder_task, completed=100) self.status.update("embedder", 50, "Embedding model loaded") self._notify_status() - + # Start embedder loading in background embedder_future = self.executor.submit(load_embedder) - + # Wait for embedder to finish (this is the bottleneck) embedder_future.result(timeout=120) # 2 minute timeout - + # Task 2: Handle indexing if needed if needs_indexing and self.auto_index: if not self._fast_index(): return False - + # Task 3: Initialize searcher (fast with pre-loaded embedder) searcher_task = progress.add_task("[cyan]Connecting to database...", total=100) self.status.update("searcher", 85, "Connecting to database...") self._notify_status() - + self.searcher = CodeSearcher(self.project_path, embedder=self.embedder) progress.update(searcher_task, completed=100) - + self.status.update("searcher", 95, "Database connected") self._notify_status() - + # Final health check self._run_health_checks() - + console.print("[green]โœ… All components initialized successfully[/green]") return True - + except Exception as e: error_msg = f"Component initialization failed: {e}" self.status.set_error(error_msg) self._notify_status() console.print(f"[red]โŒ {error_msg}[/red]") return False - + def _run_health_checks(self): """Comprehensive health checks""" checks = {} - + try: # Check 1: Embedder functionality if self.embedder: test_embedding = self.embedder.embed_code("def test(): pass") - checks['embedder'] = { - 'status': 'healthy', - 'embedding_dim': len(test_embedding), - 'model': getattr(self.embedder, 'model_name', 'unknown') + checks["embedder"] = { + "status": "healthy", + "embedding_dim": len(test_embedding), + "model": getattr(self.embedder, "model_name", "unknown"), } else: - checks['embedder'] = {'status': 'missing'} - + checks["embedder"] = {"status": "missing"} + # Check 2: Database connectivity if self.searcher: stats = self.searcher.get_statistics() - checks['database'] = { - 'status': 'healthy', - 'chunks': stats.get('total_chunks', 0), - 'languages': len(stats.get('languages', {})) + checks["database"] = { + "status": "healthy", + "chunks": stats.get("total_chunks", 0), + "languages": len(stats.get("languages", {})), } else: - checks['database'] = {'status': 'missing'} - + checks["database"] = {"status": "missing"} + # Check 3: Search functionality if self.searcher: test_results = self.searcher.search("test query", top_k=1) - checks['search'] = { - 'status': 'healthy', - 'test_results': len(test_results) + checks["search"] = { + "status": "healthy", + "test_results": len(test_results), } else: - checks['search'] = {'status': 'unavailable'} - + checks["search"] = {"status": "unavailable"} + # Check 4: Port availability try: test_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - test_sock.bind(('localhost', self.port)) + test_sock.bind(("localhost", self.port)) test_sock.close() - checks['port'] = {'status': 'available'} - except: - checks['port'] = {'status': 'occupied'} - + checks["port"] = {"status": "available"} + except (ConnectionError, OSError, TypeError, ValueError, socket.error): + checks["port"] = {"status": "occupied"} + except Exception as e: - checks['health_check_error'] = str(e) - + checks["health_check_error"] = str(e) + self.status.health_checks = checks self.last_health_check = time.time() - + # Display health summary table = Table(title="Health Check Results") table.add_column("Component", style="cyan") table.add_column("Status", style="green") table.add_column("Details", style="dim") - + for component, info in checks.items(): - status = info.get('status', 'unknown') - details = ', '.join([f"{k}={v}" for k, v in info.items() if k != 'status']) - - color = "green" if status in ['healthy', 'available'] else "yellow" + status = info.get("status", "unknown") + details = ", ".join([f"{k}={v}" for k, v in info.items() if k != "status"]) + + color = "green" if status in ["healthy", "available"] else "yellow" table.add_row(component, f"[{color}]{status}[/{color}]", details) - + console.print(table) - + def start(self): """Start the server with enhanced feedback""" try: start_time = time.time() - + # Step 1: Clear existing servers if not self._kill_existing_server(): return False - + # Step 2: Initialize all components if not self._initialize_components(): return False - + # Step 3: Start network server self.status.update("server", 98, "Starting network server...") self._notify_status() - + self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.socket.bind(('localhost', self.port)) + self.socket.bind(("localhost", self.port)) self.socket.listen(10) # Increased backlog - + self.running = True - + # Server is ready! total_time = time.time() - start_time self.status.set_ready() self._notify_status() - + # Display ready status panel = Panel( - f"[bold green]๐ŸŽ‰ RAG Server Ready![/bold green]\n\n" + "[bold green]๐ŸŽ‰ RAG Server Ready![/bold green]\n\n" f"๐ŸŒ Address: localhost:{self.port}\n" f"โšก Startup Time: {total_time:.2f}s\n" f"๐Ÿ“ Project: {self.project_path.name}\n" f"๐Ÿง  Model: {getattr(self.embedder, 'model_name', 'default')}\n" f"๐Ÿ“Š Chunks Indexed: {self.status.health_checks.get('database', {}).get('chunks', 0)}\n\n" - f"[dim]Ready to serve the development environment queries...[/dim]", + "[dim]Ready to serve the development environment queries...[/dim]", title="๐Ÿš€ Server Status", - border_style="green" + border_style="green", ) console.print(panel) - + # Start serving self._serve() - + except KeyboardInterrupt: console.print("\n[yellow]Server interrupted by user[/yellow]") self.stop() @@ -515,29 +542,29 @@ class FastRAGServer: self._notify_status() console.print(f"[red]โŒ {error_msg}[/red]") raise - + def _serve(self): """Main server loop with enhanced monitoring""" console.print("[dim]Waiting for connections... Press Ctrl+C to stop[/dim]\n") - + while self.running: try: client, addr = self.socket.accept() - + # Handle in thread pool for better performance self.executor.submit(self._handle_client, client) - + # Periodic health checks if time.time() - self.last_health_check > self.health_check_interval: self.executor.submit(self._run_health_checks) - + except KeyboardInterrupt: break except Exception as e: if self.running: logger.error(f"Server error: {e}") console.print(f"[red]Server error: {e}[/red]") - + def _handle_client(self, client: socket.socket): """Enhanced client handling with better error reporting""" try: @@ -545,263 +572,255 @@ class FastRAGServer: client.settimeout(30.0) # 30 second timeout data = self._receive_json(client) request = json.loads(data) - + # Handle different request types - if request.get('command') == 'shutdown': + if request.get("command") == "shutdown": console.print("\n[yellow]๐Ÿ›‘ Shutdown requested[/yellow]") - response = {'success': True, 'message': 'Server shutting down'} + response = {"success": True, "message": "Server shutting down"} self._send_json(client, response) self.stop() return - - if request.get('command') == 'status': - response = { - 'success': True, - 'status': self.status.get_status() - } + + if request.get("command") == "status": + response = {"success": True, "status": self.status.get_status()} self._send_json(client, response) return - + # Handle search requests - query = request.get('query', '') - top_k = request.get('top_k', 10) - + query = request.get("query", "") + top_k = request.get("top_k", 10) + if not query: raise ValueError("Empty query") - + self.query_count += 1 - + # Enhanced query logging - console.print(f"[blue]๐Ÿ” Query #{self.query_count}:[/blue] [dim]{query[:50]}{'...' if len(query) > 50 else ''}[/dim]") - + console.print( + f"[blue]๐Ÿ” Query #{self.query_count}:[/blue] [dim]{query[:50]}{'...' if len(query) > 50 else ''}[/dim]" + ) + # Perform search with timing start = time.time() results = self.searcher.search(query, top_k=top_k) search_time = time.time() - start - + # Enhanced response response = { - 'success': True, - 'query': query, - 'count': len(results), - 'search_time_ms': int(search_time * 1000), - 'results': [r.to_dict() for r in results], - 'server_uptime': int(time.time() - self.status.start_time), - 'total_queries': self.query_count, - 'server_status': 'ready' + "success": True, + "query": query, + "count": len(results), + "search_time_ms": int(search_time * 1000), + "results": [r.to_dict() for r in results], + "server_uptime": int(time.time() - self.status.start_time), + "total_queries": self.query_count, + "server_status": "ready", } - + self._send_json(client, response) - + # Enhanced result logging - console.print(f"[green]โœ… {len(results)} results in {search_time*1000:.0f}ms[/green]") - + console.print( + f"[green]โœ… {len(results)} results in {search_time*1000:.0f}ms[/green]" + ) + except Exception as e: error_msg = str(e) logger.error(f"Client handler error: {error_msg}") - + error_response = { - 'success': False, - 'error': error_msg, - 'error_type': type(e).__name__, - 'server_status': self.status.phase + "success": False, + "error": error_msg, + "error_type": type(e).__name__, + "server_status": self.status.phase, } - + try: self._send_json(client, error_response) - except: + except (TypeError, ValueError): pass - + console.print(f"[red]โŒ Query failed: {error_msg}[/red]") finally: try: client.close() - except: + except (ConnectionError, OSError, TypeError, ValueError, socket.error): pass - + def _receive_json(self, sock: socket.socket) -> str: """Receive JSON with length prefix and timeout handling""" try: # Receive length (4 bytes) - length_data = b'' + length_data = b"" while len(length_data) < 4: chunk = sock.recv(4 - len(length_data)) if not chunk: raise ConnectionError("Connection closed while receiving length") length_data += chunk - - length = int.from_bytes(length_data, 'big') + + length = int.from_bytes(length_data, "big") if length > 10_000_000: # 10MB limit raise ValueError(f"Message too large: {length} bytes") - + # Receive data - data = b'' + data = b"" while len(data) < length: chunk = sock.recv(min(65536, length - len(data))) if not chunk: raise ConnectionError("Connection closed while receiving data") data += chunk - - return data.decode('utf-8') + + return data.decode("utf-8") except socket.timeout: raise ConnectionError("Timeout while receiving data") - + def _send_json(self, sock: socket.socket, data: dict): """Send JSON with length prefix""" - json_str = json.dumps(data, ensure_ascii=False, separators=(',', ':')) - json_bytes = json_str.encode('utf-8') - + json_str = json.dumps(data, ensure_ascii=False, separators=(",", ":")) + json_bytes = json_str.encode("utf-8") + # Send length prefix length = len(json_bytes) - sock.send(length.to_bytes(4, 'big')) - + sock.send(length.to_bytes(4, "big")) + # Send data sock.sendall(json_bytes) - + def stop(self): """Graceful server shutdown""" console.print("\n[yellow]๐Ÿ›‘ Shutting down server...[/yellow]") - + self.running = False - + if self.socket: try: self.socket.close() - except: + except (ConnectionError, OSError, TypeError, ValueError, socket.error): pass - + # Shutdown executor self.executor.shutdown(wait=True, timeout=5.0) - + console.print("[green]โœ… Server stopped gracefully[/green]") # Enhanced client with status monitoring + + class FastRAGClient: """Enhanced client with better error handling and status monitoring""" - + def __init__(self, port: int = 7777): self.port = port self.timeout = 30.0 - + def search(self, query: str, top_k: int = 10) -> Dict[str, Any]: """Enhanced search with better error handling""" try: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.settimeout(self.timeout) - sock.connect(('localhost', self.port)) - - request = {'query': query, 'top_k': top_k} + sock.connect(("localhost", self.port)) + + request = {"query": query, "top_k": top_k} self._send_json(sock, request) - + data = self._receive_json(sock) response = json.loads(data) - + sock.close() return response - + except ConnectionRefusedError: return { - 'success': False, - 'error': 'RAG server not running. Start with: python -m mini_rag server', - 'error_type': 'connection_refused' + "success": False, + "error": "RAG server not running. Start with: python -m mini_rag server", + "error_type": "connection_refused", } except socket.timeout: return { - 'success': False, - 'error': f'Request timed out after {self.timeout}s', - 'error_type': 'timeout' + "success": False, + "error": f"Request timed out after {self.timeout}s", + "error_type": "timeout", } except Exception as e: - return { - 'success': False, - 'error': str(e), - 'error_type': type(e).__name__ - } - + return {"success": False, "error": str(e), "error_type": type(e).__name__} + def get_status(self) -> Dict[str, Any]: """Get detailed server status""" try: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.settimeout(5.0) - sock.connect(('localhost', self.port)) - - request = {'command': 'status'} + sock.connect(("localhost", self.port)) + + request = {"command": "status"} self._send_json(sock, request) - + data = self._receive_json(sock) response = json.loads(data) - + sock.close() return response - + except Exception as e: - return { - 'success': False, - 'error': str(e), - 'server_running': False - } - + return {"success": False, "error": str(e), "server_running": False} + def is_running(self) -> bool: """Enhanced server detection""" try: status = self.get_status() - return status.get('success', False) - except: + return status.get("success", False) + except (TypeError, ValueError): return False - + def shutdown(self) -> Dict[str, Any]: """Gracefully shutdown server""" try: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.settimeout(10.0) - sock.connect(('localhost', self.port)) - - request = {'command': 'shutdown'} + sock.connect(("localhost", self.port)) + + request = {"command": "shutdown"} self._send_json(sock, request) - + data = self._receive_json(sock) response = json.loads(data) - + sock.close() return response - + except Exception as e: - return { - 'success': False, - 'error': str(e) - } - + return {"success": False, "error": str(e)} + def _send_json(self, sock: socket.socket, data: dict): """Send JSON with length prefix""" - json_str = json.dumps(data, ensure_ascii=False, separators=(',', ':')) - json_bytes = json_str.encode('utf-8') - + json_str = json.dumps(data, ensure_ascii=False, separators=(",", ":")) + json_bytes = json_str.encode("utf-8") + length = len(json_bytes) - sock.send(length.to_bytes(4, 'big')) + sock.send(length.to_bytes(4, "big")) sock.sendall(json_bytes) - + def _receive_json(self, sock: socket.socket) -> str: """Receive JSON with length prefix""" # Receive length - length_data = b'' + length_data = b"" while len(length_data) < 4: chunk = sock.recv(4 - len(length_data)) if not chunk: raise ConnectionError("Connection closed") length_data += chunk - - length = int.from_bytes(length_data, 'big') - + + length = int.from_bytes(length_data, "big") + # Receive data - data = b'' + data = b"" while len(data) < length: chunk = sock.recv(min(65536, length - len(data))) if not chunk: raise ConnectionError("Connection closed") data += chunk - - return data.decode('utf-8') + + return data.decode("utf-8") def start_fast_server(project_path: Path, port: int = 7777, auto_index: bool = True): @@ -816,4 +835,4 @@ def start_fast_server(project_path: Path, port: int = 7777, auto_index: bool = T # Backwards compatibility RAGServer = FastRAGServer RAGClient = FastRAGClient -start_server = start_fast_server \ No newline at end of file +start_server = start_fast_server diff --git a/mini_rag/indexer.py b/mini_rag/indexer.py index 8cfa580..5922397 100644 --- a/mini_rag/indexer.py +++ b/mini_rag/indexer.py @@ -3,31 +3,39 @@ Parallel indexing engine for efficient codebase processing. Handles file discovery, chunking, embedding, and storage. """ -import os -import json import hashlib +import json import logging -from pathlib import Path -from typing import List, Dict, Any, Optional, Set, Tuple +import os from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional + import numpy as np import pandas as pd -from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn from rich.console import Console +from rich.progress import ( + BarColumn, + Progress, + SpinnerColumn, + TextColumn, + TimeRemainingColumn, +) # Optional LanceDB import try: import lancedb import pyarrow as pa + LANCEDB_AVAILABLE = True except ImportError: lancedb = None pa = None LANCEDB_AVAILABLE = False +from .chunker import CodeChunker from .ollama_embeddings import OllamaEmbedder as CodeEmbedder -from .chunker import CodeChunker, CodeChunk from .path_handler import normalize_path, normalize_relative_path logger = logging.getLogger(__name__) @@ -36,15 +44,17 @@ console = Console() class ProjectIndexer: """Indexes a project directory for semantic search.""" - - def __init__(self, - project_path: Path, - embedder: Optional[CodeEmbedder] = None, - chunker: Optional[CodeChunker] = None, - max_workers: int = 4): + + def __init__( + self, + project_path: Path, + embedder: Optional[CodeEmbedder] = None, + chunker: Optional[CodeChunker] = None, + max_workers: int = 4, + ): """ Initialize the indexer. - + Args: project_path: Path to the project to index embedder: CodeEmbedder instance (creates one if not provided) @@ -52,174 +62,241 @@ class ProjectIndexer: max_workers: Number of parallel workers for indexing """ self.project_path = Path(project_path).resolve() - self.rag_dir = self.project_path / '.mini-rag' - self.manifest_path = self.rag_dir / 'manifest.json' - self.config_path = self.rag_dir / 'config.json' - + self.rag_dir = self.project_path / ".mini-rag" + self.manifest_path = self.rag_dir / "manifest.json" + self.config_path = self.rag_dir / "config.json" + # Create RAG directory if it doesn't exist self.rag_dir.mkdir(exist_ok=True) - + # Initialize components self.embedder = embedder or CodeEmbedder() self.chunker = chunker or CodeChunker() self.max_workers = max_workers - + # Initialize database connection self.db = None self.table = None - + # File patterns to include/exclude self.include_patterns = [ # Code files - '*.py', '*.js', '*.jsx', '*.ts', '*.tsx', - '*.go', '*.java', '*.cpp', '*.c', '*.cs', - '*.rs', '*.rb', '*.php', '*.swift', '*.kt', - '*.scala', '*.r', '*.m', '*.h', '*.hpp', + "*.py", + "*.js", + "*.jsx", + "*.ts", + "*.tsx", + "*.go", + "*.java", + "*.cpp", + "*.c", + "*.cs", + "*.rs", + "*.rb", + "*.php", + "*.swift", + "*.kt", + "*.scala", + "*.r", + "*.m", + "*.h", + "*.hpp", # Documentation files - '*.md', '*.markdown', '*.rst', '*.txt', - '*.adoc', '*.asciidoc', + "*.md", + "*.markdown", + "*.rst", + "*.txt", + "*.adoc", + "*.asciidoc", # Config files - '*.json', '*.yaml', '*.yml', '*.toml', '*.ini', - '*.xml', '*.conf', '*.config', + "*.json", + "*.yaml", + "*.yml", + "*.toml", + "*.ini", + "*.xml", + "*.con", + "*.config", # Other text files - 'README', 'LICENSE', 'CHANGELOG', 'AUTHORS', - 'CONTRIBUTING', 'TODO', 'NOTES' + "README", + "LICENSE", + "CHANGELOG", + "AUTHORS", + "CONTRIBUTING", + "TODO", + "NOTES", ] - + self.exclude_patterns = [ - '__pycache__', '.git', 'node_modules', '.venv', 'venv', - 'env', 'dist', 'build', 'target', '.idea', '.vscode', - '*.pyc', '*.pyo', '*.pyd', '.DS_Store', '*.so', '*.dll', - '*.dylib', '*.exe', '*.bin', '*.log', '*.lock' + "__pycache__", + ".git", + "node_modules", + ".venv", + "venv", + "env", + "dist", + "build", + "target", + ".idea", + ".vscode", + "*.pyc", + "*.pyo", + "*.pyd", + ".DS_Store", + "*.so", + "*.dll", + "*.dylib", + "*.exe", + "*.bin", + "*.log", + "*.lock", ] - + # Load existing manifest if it exists self.manifest = self._load_manifest() - + def _load_manifest(self) -> Dict[str, Any]: """Load existing manifest or create new one.""" if self.manifest_path.exists(): try: - with open(self.manifest_path, 'r') as f: + with open(self.manifest_path, "r") as f: return json.load(f) except Exception as e: logger.warning(f"Failed to load manifest: {e}") - + return { - 'version': '1.0', - 'indexed_at': None, - 'file_count': 0, - 'chunk_count': 0, - 'files': {} + "version": "1.0", + "indexed_at": None, + "file_count": 0, + "chunk_count": 0, + "files": {}, } - + def _save_manifest(self): """Save manifest to disk.""" try: - with open(self.manifest_path, 'w') as f: + with open(self.manifest_path, "w") as f: json.dump(self.manifest, f, indent=2) except Exception as e: logger.error(f"Failed to save manifest: {e}") - + def _load_config(self) -> Dict[str, Any]: """Load or create comprehensive configuration.""" if self.config_path.exists(): try: - with open(self.config_path, 'r') as f: + with open(self.config_path, "r") as f: config = json.load(f) # Apply any loaded settings self._apply_config(config) return config except Exception as e: logger.warning(f"Failed to load config: {e}, using defaults") - + # Default configuration - comprehensive and user-friendly config = { "project": { "name": self.project_path.name, "description": f"RAG index for {self.project_path.name}", - "created_at": datetime.now().isoformat() + "created_at": datetime.now().isoformat(), }, "embedding": { "provider": "ollama", - "model": self.embedder.model_name if hasattr(self.embedder, 'model_name') else 'nomic-embed-text:latest', + "model": ( + self.embedder.model_name + if hasattr(self.embedder, "model_name") + else "nomic-embed-text:latest" + ), "base_url": "http://localhost:11434", "batch_size": 4, - "max_workers": 4 + "max_workers": 4, }, "chunking": { - "max_size": self.chunker.max_chunk_size if hasattr(self.chunker, 'max_chunk_size') else 2500, - "min_size": self.chunker.min_chunk_size if hasattr(self.chunker, 'min_chunk_size') else 100, + "max_size": ( + self.chunker.max_chunk_size + if hasattr(self.chunker, "max_chunk_size") + else 2500 + ), + "min_size": ( + self.chunker.min_chunk_size + if hasattr(self.chunker, "min_chunk_size") + else 100 + ), "overlap": 100, - "strategy": "semantic" - }, - "streaming": { - "enabled": True, - "threshold_mb": 1, - "chunk_size_kb": 64 + "strategy": "semantic", }, + "streaming": {"enabled": True, "threshold_mb": 1, "chunk_size_kb": 64}, "files": { "include_patterns": self.include_patterns, "exclude_patterns": self.exclude_patterns, "max_file_size_mb": 50, - "encoding_fallbacks": ["utf-8", "latin-1", "cp1252", "utf-8-sig"] + "encoding_fallbacks": ["utf-8", "latin-1", "cp1252", "utf-8-sig"], }, "indexing": { "parallel_workers": self.max_workers, "incremental": True, "track_changes": True, - "skip_binary": True + "skip_binary": True, }, "search": { "default_top_k": 10, "similarity_threshold": 0.7, "hybrid_search": True, - "bm25_weight": 0.3 + "bm25_weight": 0.3, }, "storage": { "compress_vectors": False, "index_type": "ivf_pq", - "cleanup_old_chunks": True - } + "cleanup_old_chunks": True, + }, } - + # Save comprehensive config with nice formatting self._save_config(config) return config - + def _apply_config(self, config: Dict[str, Any]): """Apply configuration settings to the indexer.""" try: # Apply embedding settings - if 'embedding' in config: - emb_config = config['embedding'] - if hasattr(self.embedder, 'model_name'): - self.embedder.model_name = emb_config.get('model', self.embedder.model_name) - if hasattr(self.embedder, 'base_url'): - self.embedder.base_url = emb_config.get('base_url', self.embedder.base_url) - + if "embedding" in config: + emb_config = config["embedding"] + if hasattr(self.embedder, "model_name"): + self.embedder.model_name = emb_config.get( + "model", self.embedder.model_name + ) + if hasattr(self.embedder, "base_url"): + self.embedder.base_url = emb_config.get("base_url", self.embedder.base_url) + # Apply chunking settings - if 'chunking' in config: - chunk_config = config['chunking'] - if hasattr(self.chunker, 'max_chunk_size'): - self.chunker.max_chunk_size = chunk_config.get('max_size', self.chunker.max_chunk_size) - if hasattr(self.chunker, 'min_chunk_size'): - self.chunker.min_chunk_size = chunk_config.get('min_size', self.chunker.min_chunk_size) - + if "chunking" in config: + chunk_config = config["chunking"] + if hasattr(self.chunker, "max_chunk_size"): + self.chunker.max_chunk_size = chunk_config.get( + "max_size", self.chunker.max_chunk_size + ) + if hasattr(self.chunker, "min_chunk_size"): + self.chunker.min_chunk_size = chunk_config.get( + "min_size", self.chunker.min_chunk_size + ) + # Apply file patterns - if 'files' in config: - file_config = config['files'] - self.include_patterns = file_config.get('include_patterns', self.include_patterns) - self.exclude_patterns = file_config.get('exclude_patterns', self.exclude_patterns) - + if "files" in config: + file_config = config["files"] + self.include_patterns = file_config.get( + "include_patterns", self.include_patterns + ) + self.exclude_patterns = file_config.get( + "exclude_patterns", self.exclude_patterns + ) + # Apply indexing settings - if 'indexing' in config: - idx_config = config['indexing'] - self.max_workers = idx_config.get('parallel_workers', self.max_workers) - + if "indexing" in config: + idx_config = config["indexing"] + self.max_workers = idx_config.get("parallel_workers", self.max_workers) + except Exception as e: logger.warning(f"Failed to apply some config settings: {e}") - + def _save_config(self, config: Dict[str, Any]): """Save configuration with nice formatting and comments.""" try: @@ -228,17 +305,17 @@ class ProjectIndexer: "_comment": "RAG System Configuration - Edit this file to customize indexing behavior", "_version": "2.0", "_docs": "See README.md for detailed configuration options", - **config + **config, } - - with open(self.config_path, 'w') as f: + + with open(self.config_path, "w") as f: json.dump(config_with_comments, f, indent=2, sort_keys=True) - + logger.info(f"Configuration saved to {self.config_path}") - + except Exception as e: logger.error(f"Failed to save config: {e}") - + def _get_file_hash(self, file_path: Path) -> str: """Calculate SHA256 hash of a file.""" sha256_hash = hashlib.sha256() @@ -250,161 +327,184 @@ class ProjectIndexer: except Exception as e: logger.error(f"Failed to hash {file_path}: {e}") return "" - + def _should_index_file(self, file_path: Path) -> bool: """Check if a file should be indexed based on patterns and content.""" # Check file size (skip files > 1MB) try: if file_path.stat().st_size > 1_000_000: return False - except: + except (OSError, IOError): return False - + # Check exclude patterns first path_str = str(file_path) for pattern in self.exclude_patterns: if pattern in path_str: return False - + # Check include patterns (extension-based) for pattern in self.include_patterns: if file_path.match(pattern): return True - + # NEW: Content-based inclusion for extensionless files if not file_path.suffix: return self._should_index_extensionless_file(file_path) - + return False - + def _should_index_extensionless_file(self, file_path: Path) -> bool: """Check if an extensionless file should be indexed based on content.""" try: # Read first 1KB to check content - with open(file_path, 'rb') as f: + with open(file_path, "rb") as f: first_chunk = f.read(1024) - + # Check if it's a text file (not binary) try: - text_content = first_chunk.decode('utf-8') + text_content = first_chunk.decode("utf-8") except UnicodeDecodeError: return False # Binary file, skip - + # Check for code indicators code_indicators = [ - '#!/usr/bin/env python', '#!/usr/bin/python', '#!.*python', - 'import ', 'from ', 'def ', 'class ', 'if __name__', - 'function ', 'var ', 'const ', 'let ', 'package main', - 'public class', 'private class', 'public static void' + "#!/usr/bin/env python", + "#!/usr/bin/python", + "#!.*python", + "import ", + "from ", + "def ", + "class ", + "if __name__", + "function ", + "var ", + "const ", + "let ", + "package main", + "public class", + "private class", + "public static void", ] - + text_lower = text_content.lower() for indicator in code_indicators: if indicator in text_lower: return True - + # Check for configuration files config_indicators = [ - '#!/bin/bash', '#!/bin/sh', '[', 'version =', 'name =', - 'description =', 'author =', '', '", + " bool: """Smart check if a file needs to be reindexed - optimized for speed.""" file_str = normalize_relative_path(file_path, self.project_path) - + # Not in manifest - needs indexing - if file_str not in self.manifest['files']: + if file_str not in self.manifest["files"]: return True - - file_info = self.manifest['files'][file_str] - + + file_info = self.manifest["files"][file_str] + try: stat = file_path.stat() - + # Quick checks first (no I/O) - check size and modification time - stored_size = file_info.get('size', 0) - stored_mtime = file_info.get('mtime', 0) - + stored_size = file_info.get("size", 0) + stored_mtime = file_info.get("mtime", 0) + current_size = stat.st_size current_mtime = stat.st_mtime - + # If size or mtime changed, definitely needs reindex if current_size != stored_size or current_mtime != stored_mtime: return True - + # Size and mtime same - check hash only if needed (for paranoia) # This catches cases where content changed but mtime didn't (rare but possible) current_hash = self._get_file_hash(file_path) - stored_hash = file_info.get('hash', '') - + stored_hash = file_info.get("hash", "") + return current_hash != stored_hash - + except (OSError, IOError) as e: logger.warning(f"Could not check file stats for {file_path}: {e}") # If we can't check file stats, assume it needs reindex return True - + def _cleanup_removed_files(self): """Remove entries for files that no longer exist from manifest and database.""" - if 'files' not in self.manifest: + if "files" not in self.manifest: return - + removed_files = [] - for file_str in list(self.manifest['files'].keys()): + for file_str in list(self.manifest["files"].keys()): file_path = self.project_path / file_str if not file_path.exists(): removed_files.append(file_str) - + if removed_files: logger.info(f"Cleaning up {len(removed_files)} removed files from index") - + for file_str in removed_files: # Remove from database try: - if hasattr(self, 'table') and self.table: + if hasattr(self, "table") and self.table: self.table.delete(f"file_path = '{file_str}'") logger.debug(f"Removed chunks for deleted file: {file_str}") except Exception as e: logger.warning(f"Could not remove chunks for {file_str}: {e}") - + # Remove from manifest - del self.manifest['files'][file_str] - + del self.manifest["files"][file_str] + # Save updated manifest self._save_manifest() logger.info(f"Cleanup complete - removed {len(removed_files)} files") - + def _get_files_to_index(self) -> List[Path]: """Get all files that need to be indexed.""" files_to_index = [] - + # Walk through project directory for root, dirs, files in os.walk(self.project_path): # Skip excluded directories - dirs[:] = [d for d in dirs if not any(pattern in d for pattern in self.exclude_patterns)] - + dirs[:] = [ + d for d in dirs if not any(pattern in d for pattern in self.exclude_patterns) + ] + root_path = Path(root) for file in files: file_path = root_path / file - + if self._should_index_file(file_path) and self._needs_reindex(file_path): files_to_index.append(file_path) - + return files_to_index - - def _process_file(self, file_path: Path, stream_threshold: int = 1024 * 1024) -> Optional[List[Dict[str, Any]]]: + + def _process_file( + self, file_path: Path, stream_threshold: int = 1024 * 1024 + ) -> Optional[List[Dict[str, Any]]]: """Process a single file: read, chunk, embed. - + Args: file_path: Path to the file to process stream_threshold: Files larger than this (in bytes) use streaming (default: 1MB) @@ -412,30 +512,30 @@ class ProjectIndexer: try: # Check file size for streaming decision file_size = file_path.stat().st_size - + if file_size > stream_threshold: logger.info(f"Streaming large file ({file_size:,} bytes): {file_path}") content = self._read_file_streaming(file_path) else: # Read file content normally for small files - content = file_path.read_text(encoding='utf-8') - + content = file_path.read_text(encoding="utf-8") + # Chunk the file chunks = self.chunker.chunk_file(file_path, content) - + if not chunks: return None - + # Prepare data for embedding chunk_texts = [chunk.content for chunk in chunks] - + # Generate embeddings embeddings = self.embedder.embed_code(chunk_texts) - + # Prepare records for database records = [] expected_dim = self.embedder.get_embedding_dim() - + for i, chunk in enumerate(chunks): # Validate embedding embedding = embeddings[i].astype(np.float32) @@ -444,144 +544,160 @@ class ProjectIndexer: f"Invalid embedding dimension for {file_path} chunk {i}: " f"expected ({expected_dim},), got {embedding.shape}" ) - + record = { - 'file_path': normalize_relative_path(file_path, self.project_path), - 'absolute_path': normalize_path(file_path), - 'chunk_id': f"{file_path.stem}_{i}", - 'content': chunk.content, - 'start_line': int(chunk.start_line), - 'end_line': int(chunk.end_line), - 'chunk_type': chunk.chunk_type, - 'name': chunk.name or f"chunk_{i}", - 'language': chunk.language, - 'embedding': embedding, # Keep as numpy array - 'indexed_at': datetime.now().isoformat(), + "file_path": normalize_relative_path(file_path, self.project_path), + "absolute_path": normalize_path(file_path), + "chunk_id": f"{file_path.stem}_{i}", + "content": chunk.content, + "start_line": int(chunk.start_line), + "end_line": int(chunk.end_line), + "chunk_type": chunk.chunk_type, + "name": chunk.name or f"chunk_{i}", + "language": chunk.language, + "embedding": embedding, # Keep as numpy array + "indexed_at": datetime.now().isoformat(), # Add new metadata fields - 'file_lines': int(chunk.file_lines) if chunk.file_lines else 0, - 'chunk_index': int(chunk.chunk_index) if chunk.chunk_index is not None else i, - 'total_chunks': int(chunk.total_chunks) if chunk.total_chunks else len(chunks), - 'parent_class': chunk.parent_class or '', - 'parent_function': chunk.parent_function or '', - 'prev_chunk_id': chunk.prev_chunk_id or '', - 'next_chunk_id': chunk.next_chunk_id or '', + "file_lines": int(chunk.file_lines) if chunk.file_lines else 0, + "chunk_index": ( + int(chunk.chunk_index) if chunk.chunk_index is not None else i + ), + "total_chunks": ( + int(chunk.total_chunks) if chunk.total_chunks else len(chunks) + ), + "parent_class": chunk.parent_class or "", + "parent_function": chunk.parent_function or "", + "prev_chunk_id": chunk.prev_chunk_id or "", + "next_chunk_id": chunk.next_chunk_id or "", } records.append(record) - + # Update manifest with enhanced tracking file_str = normalize_relative_path(file_path, self.project_path) stat = file_path.stat() - self.manifest['files'][file_str] = { - 'hash': self._get_file_hash(file_path), - 'size': stat.st_size, - 'mtime': stat.st_mtime, - 'chunks': len(chunks), - 'indexed_at': datetime.now().isoformat(), - 'language': chunks[0].language if chunks else 'unknown', - 'encoding': 'utf-8' # Track encoding used + self.manifest["files"][file_str] = { + "hash": self._get_file_hash(file_path), + "size": stat.st_size, + "mtime": stat.st_mtime, + "chunks": len(chunks), + "indexed_at": datetime.now().isoformat(), + "language": chunks[0].language if chunks else "unknown", + "encoding": "utf-8", # Track encoding used } - + return records - + except Exception as e: logger.error(f"Failed to process {file_path}: {e}") return None - + def _read_file_streaming(self, file_path: Path, chunk_size: int = 64 * 1024) -> str: """ Read large files in chunks to avoid loading entirely into memory. - + Args: file_path: Path to the file to read chunk_size: Size of each read chunk in bytes (default: 64KB) - + Returns: Complete file content as string """ content_parts = [] - + try: - with open(file_path, 'r', encoding='utf-8') as f: + with open(file_path, "r", encoding="utf-8") as f: while True: chunk = f.read(chunk_size) if not chunk: break content_parts.append(chunk) - + logger.debug(f"Streamed {len(content_parts)} chunks from {file_path}") - return ''.join(content_parts) - + return "".join(content_parts) + except UnicodeDecodeError: # Try with different encodings for problematic files - for encoding in ['latin-1', 'cp1252', 'utf-8-sig']: + for encoding in ["latin-1", "cp1252", "utf-8-sig"]: try: - with open(file_path, 'r', encoding=encoding) as f: + with open(file_path, "r", encoding=encoding) as f: content_parts = [] while True: chunk = f.read(chunk_size) if not chunk: break content_parts.append(chunk) - - logger.debug(f"Streamed {len(content_parts)} chunks from {file_path} using {encoding}") - return ''.join(content_parts) + + logger.debug( + f"Streamed {len(content_parts)} chunks from {file_path} using {encoding}" + ) + return "".join(content_parts) except UnicodeDecodeError: continue - + # If all encodings fail, return empty string logger.warning(f"Could not decode {file_path} with any encoding") return "" - + def _init_database(self): """Initialize LanceDB connection and table.""" if not LANCEDB_AVAILABLE: - logger.error("LanceDB is not available. Please install LanceDB for full indexing functionality.") + logger.error( + "LanceDB is not available. Please install LanceDB for full indexing functionality." + ) logger.info("For Ollama-only mode, consider using hash-based embeddings instead.") - raise ImportError("LanceDB dependency is required for indexing. Install with: pip install lancedb pyarrow") - + raise ImportError( + "LanceDB dependency is required for indexing. Install with: pip install lancedb pyarrow" + ) + try: self.db = lancedb.connect(self.rag_dir) - + # Define schema with fixed-size vector embedding_dim = self.embedder.get_embedding_dim() - schema = pa.schema([ - pa.field("file_path", pa.string()), - pa.field("absolute_path", pa.string()), - pa.field("chunk_id", pa.string()), - pa.field("content", pa.string()), - pa.field("start_line", pa.int32()), - pa.field("end_line", pa.int32()), - pa.field("chunk_type", pa.string()), - pa.field("name", pa.string()), - pa.field("language", pa.string()), - pa.field("embedding", pa.list_(pa.float32(), embedding_dim)), # Fixed-size list - pa.field("indexed_at", pa.string()), - # New metadata fields - pa.field("file_lines", pa.int32()), - pa.field("chunk_index", pa.int32()), - pa.field("total_chunks", pa.int32()), - pa.field("parent_class", pa.string(), nullable=True), - pa.field("parent_function", pa.string(), nullable=True), - pa.field("prev_chunk_id", pa.string(), nullable=True), - pa.field("next_chunk_id", pa.string(), nullable=True), - ]) - + schema = pa.schema( + [ + pa.field("file_path", pa.string()), + pa.field("absolute_path", pa.string()), + pa.field("chunk_id", pa.string()), + pa.field("content", pa.string()), + pa.field("start_line", pa.int32()), + pa.field("end_line", pa.int32()), + pa.field("chunk_type", pa.string()), + pa.field("name", pa.string()), + pa.field("language", pa.string()), + pa.field( + "embedding", pa.list_(pa.float32(), embedding_dim) + ), # Fixed-size list + pa.field("indexed_at", pa.string()), + # New metadata fields + pa.field("file_lines", pa.int32()), + pa.field("chunk_index", pa.int32()), + pa.field("total_chunks", pa.int32()), + pa.field("parent_class", pa.string(), nullable=True), + pa.field("parent_function", pa.string(), nullable=True), + pa.field("prev_chunk_id", pa.string(), nullable=True), + pa.field("next_chunk_id", pa.string(), nullable=True), + ] + ) + # Create or open table if "code_vectors" in self.db.table_names(): try: # Try to open existing table self.table = self.db.open_table("code_vectors") - + # Check if schema matches by trying to get the schema existing_schema = self.table.schema - + # Check if all required fields exist required_fields = {field.name for field in schema} existing_fields = {field.name for field in existing_schema} - + if not required_fields.issubset(existing_fields): # Schema mismatch - drop and recreate table - logger.warning("Schema mismatch detected. Dropping and recreating table.") + logger.warning( + "Schema mismatch detected. Dropping and recreating table." + ) self.db.drop_table("code_vectors") self.table = self.db.create_table("code_vectors", schema=schema) logger.info("Recreated code_vectors table with updated schema") @@ -596,39 +712,41 @@ class ProjectIndexer: else: # Create empty table with schema self.table = self.db.create_table("code_vectors", schema=schema) - logger.info(f"Created new code_vectors table with embedding dimension {embedding_dim}") - + logger.info( + f"Created new code_vectors table with embedding dimension {embedding_dim}" + ) + except Exception as e: logger.error(f"Failed to initialize database: {e}") raise - + def index_project(self, force_reindex: bool = False) -> Dict[str, Any]: """ Index the entire project. - + Args: force_reindex: If True, reindex all files regardless of changes - + Returns: Dictionary with indexing statistics """ start_time = datetime.now() - + # Initialize database self._init_database() - + # Clean up removed files (essential for portability) if not force_reindex: self._cleanup_removed_files() - + # Clear manifest if force reindex if force_reindex: self.manifest = { - 'version': '1.0', - 'indexed_at': None, - 'file_count': 0, - 'chunk_count': 0, - 'files': {} + "version": "1.0", + "indexed_at": None, + "file_count": 0, + "chunk_count": 0, + "files": {}, } # Clear existing table if "code_vectors" in self.db.table_names(): @@ -636,24 +754,24 @@ class ProjectIndexer: self.table = None # Reinitialize the database to recreate the table self._init_database() - + # Get files to index files_to_index = self._get_files_to_index() - + if not files_to_index: console.print("[green][/green] All files are up to date!") return { - 'files_indexed': 0, - 'chunks_created': 0, - 'time_taken': 0, + "files_indexed": 0, + "chunks_created": 0, + "time_taken": 0, } - + console.print(f"[cyan]Found {len(files_to_index)} files to index[/cyan]") - + # Process files in parallel all_records = [] failed_files = [] - + with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), @@ -662,23 +780,20 @@ class ProjectIndexer: TimeRemainingColumn(), console=console, ) as progress: - - task = progress.add_task( - "[cyan]Indexing files...", - total=len(files_to_index) - ) - + + task = progress.add_task("[cyan]Indexing files...", total=len(files_to_index)) + with ThreadPoolExecutor(max_workers=self.max_workers) as executor: # Submit all files for processing future_to_file = { executor.submit(self._process_file, file_path): file_path for file_path in files_to_index } - + # Process completed files for future in as_completed(future_to_file): file_path = future_to_file[future] - + try: records = future.result() if records: @@ -686,9 +801,9 @@ class ProjectIndexer: except Exception as e: logger.error(f"Failed to process {file_path}: {e}") failed_files.append(file_path) - + progress.advance(task) - + # Batch insert all records if all_records: try: @@ -699,57 +814,59 @@ class ProjectIndexer: df["file_lines"] = df["file_lines"].astype("int32") df["chunk_index"] = df["chunk_index"].astype("int32") df["total_chunks"] = df["total_chunks"].astype("int32") - + # Table should already be created in _init_database if self.table is None: raise RuntimeError("Table not initialized properly") - + self.table.add(df) - + console.print(f"[green][/green] Added {len(all_records)} chunks to database") except Exception as e: logger.error(f"Failed to insert records: {e}") raise - + # Update manifest - self.manifest['indexed_at'] = datetime.now().isoformat() - self.manifest['file_count'] = len(self.manifest['files']) - self.manifest['chunk_count'] = sum( - f['chunks'] for f in self.manifest['files'].values() + self.manifest["indexed_at"] = datetime.now().isoformat() + self.manifest["file_count"] = len(self.manifest["files"]) + self.manifest["chunk_count"] = sum( + f["chunks"] for f in self.manifest["files"].values() ) self._save_manifest() - + # Calculate statistics end_time = datetime.now() time_taken = (end_time - start_time).total_seconds() - + stats = { - 'files_indexed': len(files_to_index) - len(failed_files), - 'files_failed': len(failed_files), - 'chunks_created': len(all_records), - 'time_taken': time_taken, - 'files_per_second': len(files_to_index) / time_taken if time_taken > 0 else 0, + "files_indexed": len(files_to_index) - len(failed_files), + "files_failed": len(failed_files), + "chunks_created": len(all_records), + "time_taken": time_taken, + "files_per_second": (len(files_to_index) / time_taken if time_taken > 0 else 0), } - + # Print summary console.print("\n[bold green]Indexing Complete![/bold green]") console.print(f"Files indexed: {stats['files_indexed']}") console.print(f"Chunks created: {stats['chunks_created']}") console.print(f"Time taken: {stats['time_taken']:.2f} seconds") console.print(f"Speed: {stats['files_per_second']:.1f} files/second") - + if failed_files: - console.print(f"\n[yellow]Warning:[/yellow] {len(failed_files)} files failed to index") - + console.print( + f"\n[yellow]Warning:[/yellow] {len(failed_files)} files failed to index" + ) + return stats - + def update_file(self, file_path: Path) -> bool: """ Update index for a single file with proper vector multiply in/out. - + Args: file_path: Path to the file to update - + Returns: True if successful, False otherwise """ @@ -757,13 +874,13 @@ class ProjectIndexer: # Make sure database is initialized if self.table is None: self._init_database() - + # Get normalized file path for consistent lookup file_str = normalize_relative_path(file_path, self.project_path) - + # Process the file to get new chunks records = self._process_file(file_path) - + if records: # Create DataFrame with proper types df = pd.DataFrame(records) @@ -772,40 +889,44 @@ class ProjectIndexer: df["file_lines"] = df["file_lines"].astype("int32") df["chunk_index"] = df["chunk_index"].astype("int32") df["total_chunks"] = df["total_chunks"].astype("int32") - + # Use vector store's update method (multiply out old, multiply in new) - if hasattr(self, '_vector_store') and self._vector_store: + if hasattr(self, "_vector_store") and self._vector_store: success = self._vector_store.update_file_vectors(file_str, df) else: # Fallback: delete by file path and add new data try: self.table.delete(f"file = '{file_str}'") except Exception as e: - logger.debug(f"Could not delete existing chunks (might not exist): {e}") + logger.debug( + f"Could not delete existing chunks (might not exist): {e}" + ) self.table.add(df) success = True - + if success: # Update manifest with enhanced file tracking file_hash = self._get_file_hash(file_path) stat = file_path.stat() - if 'files' not in self.manifest: - self.manifest['files'] = {} - self.manifest['files'][file_str] = { - 'hash': file_hash, - 'size': stat.st_size, - 'mtime': stat.st_mtime, - 'chunks': len(records), - 'last_updated': datetime.now().isoformat(), - 'language': records[0].get('language', 'unknown') if records else 'unknown', - 'encoding': 'utf-8' + if "files" not in self.manifest: + self.manifest["files"] = {} + self.manifest["files"][file_str] = { + "hash": file_hash, + "size": stat.st_size, + "mtime": stat.st_mtime, + "chunks": len(records), + "last_updated": datetime.now().isoformat(), + "language": ( + records[0].get("language", "unknown") if records else "unknown" + ), + "encoding": "utf-8", } self._save_manifest() logger.debug(f"Successfully updated {len(records)} chunks for {file_str}") return True else: # File exists but has no processable content - remove existing chunks - if hasattr(self, '_vector_store') and self._vector_store: + if hasattr(self, "_vector_store") and self._vector_store: self._vector_store.delete_by_file(file_str) else: try: @@ -814,69 +935,69 @@ class ProjectIndexer: pass logger.debug(f"Removed chunks for empty/unprocessable file: {file_str}") return True - + return False - + except Exception as e: logger.error(f"Failed to update {file_path}: {e}") return False - + def delete_file(self, file_path: Path) -> bool: """ Delete all chunks for a file from the index. - + Args: file_path: Path to the file to delete from index - + Returns: - True if successful, False otherwise + True if successful, False otherwise """ try: if self.table is None: self._init_database() - + file_str = normalize_relative_path(file_path, self.project_path) - + # Delete from vector store - if hasattr(self, '_vector_store') and self._vector_store: + if hasattr(self, "_vector_store") and self._vector_store: success = self._vector_store.delete_by_file(file_str) else: try: self.table.delete(f"file = '{file_str}'") success = True except Exception as e: - logger.error(f"Failed to delete {file_str}: {e}") + logger.error(f"Failed to delete {file_str}: {e}") success = False - + # Update manifest - if success and 'files' in self.manifest and file_str in self.manifest['files']: - del self.manifest['files'][file_str] + if success and "files" in self.manifest and file_str in self.manifest["files"]: + del self.manifest["files"][file_str] self._save_manifest() logger.debug(f"Deleted chunks for file: {file_str}") - + return success - + except Exception as e: logger.error(f"Failed to delete {file_path}: {e}") return False - + def get_statistics(self) -> Dict[str, Any]: """Get indexing statistics.""" stats = { - 'project_path': str(self.project_path), - 'indexed_at': self.manifest.get('indexed_at', 'Never'), - 'file_count': self.manifest.get('file_count', 0), - 'chunk_count': self.manifest.get('chunk_count', 0), - 'index_size_mb': 0, + "project_path": str(self.project_path), + "indexed_at": self.manifest.get("indexed_at", "Never"), + "file_count": self.manifest.get("file_count", 0), + "chunk_count": self.manifest.get("chunk_count", 0), + "index_size_mb": 0, } - + # Calculate index size try: - db_path = self.rag_dir / 'code_vectors.lance' + db_path = self.rag_dir / "code_vectors.lance" if db_path.exists(): - size_bytes = sum(f.stat().st_size for f in db_path.rglob('*') if f.is_file()) - stats['index_size_mb'] = size_bytes / (1024 * 1024) - except: + size_bytes = sum(f.stat().st_size for f in db_path.rglob("*") if f.is_file()) + stats["index_size_mb"] = size_bytes / (1024 * 1024) + except (OSError, IOError, PermissionError): pass - - return stats \ No newline at end of file + + return stats diff --git a/mini_rag/llm_safeguards.py b/mini_rag/llm_safeguards.py index 2fc5238..0b50868 100644 --- a/mini_rag/llm_safeguards.py +++ b/mini_rag/llm_safeguards.py @@ -6,163 +6,173 @@ Provides runaway prevention, context management, and intelligent detection of problematic model behaviors to ensure reliable user experience. """ +import logging import re import time -import logging -from typing import Optional, Dict, List, Tuple from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple logger = logging.getLogger(__name__) + @dataclass class SafeguardConfig: """Configuration for LLM safeguards - gentle and educational.""" - max_output_tokens: int = 4000 # Allow longer responses for learning - max_repetition_ratio: float = 0.7 # Be very permissive - only catch extreme repetition - max_response_time: int = 120 # Allow 2 minutes for complex thinking - min_useful_length: int = 10 # Lower threshold - short answers can be useful - context_window: int = 32000 # Match Qwen3 context length (32K token limit) + + max_output_tokens: int = 4000 # Allow longer responses for learning + max_repetition_ratio: float = 0.7 # Be very permissive - only catch extreme repetition + max_response_time: int = 120 # Allow 2 minutes for complex thinking + min_useful_length: int = 10 # Lower threshold - short answers can be useful + context_window: int = 32000 # Match Qwen3 context length (32K token limit) enable_thinking_detection: bool = True # Detect thinking patterns - + + class ModelRunawayDetector: """Detects and prevents model runaway behaviors.""" - + def __init__(self, config: SafeguardConfig = None): self.config = config or SafeguardConfig() self.response_patterns = self._compile_patterns() - + def _compile_patterns(self) -> Dict[str, re.Pattern]: """Compile regex patterns for runaway detection.""" return { # Excessive repetition patterns - 'word_repetition': re.compile(r'\b(\w+)\b(?:\s+\1\b){3,}', re.IGNORECASE), - 'phrase_repetition': re.compile(r'(.{10,50}?)\1{2,}', re.DOTALL), - + "word_repetition": re.compile(r"\b(\w+)\b(?:\s+\1\b){3,}", re.IGNORECASE), + "phrase_repetition": re.compile(r"(.{10,50}?)\1{2,}", re.DOTALL), # Thinking loop patterns (small models get stuck) - 'thinking_loop': re.compile(r'(let me think|i think|thinking|consider|actually|wait|hmm|well)\s*[.,:]*\s*\1', re.IGNORECASE), - + "thinking_loop": re.compile( + r"(let me think|i think|thinking|consider|actually|wait|hmm|well)\s*[.,:]*\s*\1", + re.IGNORECASE, + ), # Rambling patterns - 'excessive_filler': re.compile(r'\b(um|uh|well|you know|like|basically|actually|so|then|and|but|however)\b(?:\s+[^.!?]*){5,}', re.IGNORECASE), - + "excessive_filler": re.compile( + r"\b(um|uh|well|you know|like|basically|actually|so|then|and|but|however)\b(?:\s+[^.!?]*){5,}", + re.IGNORECASE, + ), # JSON corruption patterns - 'broken_json': re.compile(r'\{[^}]*\{[^}]*\{'), # Nested broken JSON - 'json_repetition': re.compile(r'("[\w_]+"\s*:\s*"[^"]*",?\s*){4,}'), # Repeated JSON fields + "broken_json": re.compile(r"\{[^}]*\{[^}]*\{"), # Nested broken JSON + "json_repetition": re.compile( + r'("[\w_]+"\s*:\s*"[^"]*",?\s*){4,}' + ), # Repeated JSON fields } - - def check_response_quality(self, response: str, query: str, start_time: float) -> Tuple[bool, Optional[str], Optional[str]]: + + def check_response_quality( + self, response: str, query: str, start_time: float + ) -> Tuple[bool, Optional[str], Optional[str]]: """ Check response quality and detect runaway behaviors. - + Returns: (is_valid, issue_type, user_explanation) """ if not response or len(response.strip()) < self.config.min_useful_length: return False, "too_short", self._explain_too_short() - + # Check response time elapsed = time.time() - start_time if elapsed > self.config.max_response_time: return False, "timeout", self._explain_timeout() - + # Check for repetition issues repetition_issue = self._check_repetition(response) if repetition_issue: return False, repetition_issue, self._explain_repetition(repetition_issue) - + # Check for thinking loops if self.config.enable_thinking_detection: thinking_issue = self._check_thinking_loops(response) if thinking_issue: return False, thinking_issue, self._explain_thinking_loop() - + # Check for rambling rambling_issue = self._check_rambling(response) if rambling_issue: return False, rambling_issue, self._explain_rambling() - + # Check JSON corruption (for structured responses) - if '{' in response and '}' in response: + if "{" in response and "}" in response: json_issue = self._check_json_corruption(response) if json_issue: return False, json_issue, self._explain_json_corruption() - + return True, None, None - + def _check_repetition(self, response: str) -> Optional[str]: """Check for excessive repetition.""" # Word repetition - if self.response_patterns['word_repetition'].search(response): + if self.response_patterns["word_repetition"].search(response): return "word_repetition" - - # Phrase repetition - if self.response_patterns['phrase_repetition'].search(response): + + # Phrase repetition + if self.response_patterns["phrase_repetition"].search(response): return "phrase_repetition" - + # Calculate repetition ratio (excluding Qwen3 thinking blocks) analysis_text = response if "" in response and "" in response: # Extract only the actual response (after thinking) for repetition analysis thinking_end = response.find("") if thinking_end != -1: - analysis_text = response[thinking_end + 8:].strip() - + analysis_text = response[thinking_end + 8 :].strip() + # If the actual response (excluding thinking) is short, don't penalize if len(analysis_text.split()) < 20: return None - + words = analysis_text.split() if len(words) > 10: unique_words = set(words) repetition_ratio = 1 - (len(unique_words) / len(words)) if repetition_ratio > self.config.max_repetition_ratio: return "high_repetition_ratio" - + return None - + def _check_thinking_loops(self, response: str) -> Optional[str]: """Check for thinking loops (common in small models).""" - if self.response_patterns['thinking_loop'].search(response): + if self.response_patterns["thinking_loop"].search(response): return "thinking_loop" - + # Check for excessive meta-commentary - thinking_words = ['think', 'considering', 'actually', 'wait', 'hmm', 'let me'] + thinking_words = ["think", "considering", "actually", "wait", "hmm", "let me"] thinking_count = sum(response.lower().count(word) for word in thinking_words) - + if thinking_count > 5 and len(response.split()) < 200: return "excessive_thinking" - + return None - + def _check_rambling(self, response: str) -> Optional[str]: """Check for rambling or excessive filler.""" - if self.response_patterns['excessive_filler'].search(response): + if self.response_patterns["excessive_filler"].search(response): return "excessive_filler" - + # Check for extremely long sentences (sign of rambling) - sentences = re.split(r'[.!?]+', response) + sentences = re.split(r"[.!?]+", response) long_sentences = [s for s in sentences if len(s.split()) > 50] - + if len(long_sentences) > 2: return "excessive_rambling" - + return None - + def _check_json_corruption(self, response: str) -> Optional[str]: """Check for JSON corruption in structured responses.""" - if self.response_patterns['broken_json'].search(response): + if self.response_patterns["broken_json"].search(response): return "broken_json" - - if self.response_patterns['json_repetition'].search(response): + + if self.response_patterns["json_repetition"].search(response): return "json_repetition" - + return None - + def _explain_too_short(self) -> str: return """๐Ÿค” The AI response was too short to be helpful. **Why this happens:** โ€ข The model might be confused by the query -โ€ข Context might be insufficient +โ€ข Context might be insufficient โ€ข Model might be overloaded **What to try:** @@ -180,11 +190,11 @@ class ModelRunawayDetector: **What to try:** โ€ข Try a simpler, more direct question -โ€ข Use synthesis mode for faster responses: `--synthesize` +โ€ข Use synthesis mode for faster responses: `--synthesize` โ€ข Consider using a larger model if available""" def _explain_repetition(self, issue_type: str) -> str: - return f"""๐Ÿ”„ The AI got stuck in repetition loops ({issue_type}). + return """๐Ÿ”„ The AI got stuck in repetition loops ({issue_type}). **Why this happens:** โ€ข Small models sometimes repeat when uncertain @@ -216,7 +226,7 @@ class ModelRunawayDetector: **Why this happens:** โ€ข Small models sometimes lose focus on complex topics -โ€ข Query might be too broad or vague +โ€ข Query might be too broad or vague โ€ข Model trying to cover too much at once **What to try:** @@ -233,7 +243,7 @@ class ModelRunawayDetector: โ€ข Context limits can cause format errors โ€ข Complex analysis might overwhelm formatting -**What to try:** +**What to try:** โ€ข Try the question again (often resolves itself) โ€ข Use simpler questions for better formatting โ€ข Synthesis mode sometimes gives cleaner output @@ -242,90 +252,109 @@ class ModelRunawayDetector: def get_recovery_suggestions(self, issue_type: str, query: str) -> List[str]: """Get specific recovery suggestions based on the issue.""" suggestions = [] - - if issue_type in ['thinking_loop', 'excessive_thinking']: - suggestions.extend([ - f"Try synthesis mode: `rag-mini search . \"{query}\" --synthesize`", - "Ask more direct questions without 'why' or 'how'", - "Break complex questions into smaller parts" - ]) - - elif issue_type in ['word_repetition', 'phrase_repetition', 'high_repetition_ratio']: - suggestions.extend([ - "Try rephrasing your question completely", - "Use more specific technical terms", - f"Try exploration mode: `rag-mini explore .`" - ]) - - elif issue_type == 'timeout': - suggestions.extend([ - "Try a simpler version of your question", - "Use synthesis mode for faster responses", - "Check if Ollama is under heavy load" - ]) - + + if issue_type in ["thinking_loop", "excessive_thinking"]: + suggestions.extend( + [ + f'Try synthesis mode: `rag-mini search . "{query}" --synthesize`', + "Ask more direct questions without 'why' or 'how'", + "Break complex questions into smaller parts", + ] + ) + + elif issue_type in [ + "word_repetition", + "phrase_repetition", + "high_repetition_ratio", + ]: + suggestions.extend( + [ + "Try rephrasing your question completely", + "Use more specific technical terms", + "Try exploration mode: `rag-mini explore .`", + ] + ) + + elif issue_type == "timeout": + suggestions.extend( + [ + "Try a simpler version of your question", + "Use synthesis mode for faster responses", + "Check if Ollama is under heavy load", + ] + ) + # Universal suggestions - suggestions.extend([ - "Consider using a larger model if available (qwen3:1.7b or qwen3:4b)", - "Check model status: `ollama list`" - ]) - + suggestions.extend( + [ + "Consider using a larger model if available (qwen3:1.7b or qwen3:4b)", + "Check model status: `ollama list`", + ] + ) + return suggestions + def get_optimal_ollama_parameters(model_name: str) -> Dict[str, any]: """Get optimal parameters for different Ollama models.""" - + base_params = { - "num_ctx": 32768, # Good context window for most uses - "num_predict": 2000, # Reasonable response length - "temperature": 0.3, # Balanced creativity/consistency + "num_ctx": 32768, # Good context window for most uses + "num_predict": 2000, # Reasonable response length + "temperature": 0.3, # Balanced creativity/consistency } - + # Model-specific optimizations if "qwen3:0.6b" in model_name.lower(): return { **base_params, - "repeat_penalty": 1.15, # Prevent repetition in small model - "presence_penalty": 1.5, # Suppress repetitive outputs - "top_p": 0.8, # Focused sampling - "top_k": 20, # Limit choices - "num_predict": 1500, # Shorter responses for reliability + "repeat_penalty": 1.15, # Prevent repetition in small model + "presence_penalty": 1.5, # Suppress repetitive outputs + "top_p": 0.8, # Focused sampling + "top_k": 20, # Limit choices + "num_predict": 1500, # Shorter responses for reliability } - + elif "qwen3:1.7b" in model_name.lower(): return { **base_params, - "repeat_penalty": 1.1, # Less aggressive for larger model - "presence_penalty": 1.0, # Balanced - "top_p": 0.9, # More creative - "top_k": 40, # More choices + "repeat_penalty": 1.1, # Less aggressive for larger model + "presence_penalty": 1.0, # Balanced + "top_p": 0.9, # More creative + "top_k": 40, # More choices } - + elif any(size in model_name.lower() for size in ["3b", "7b", "8b"]): return { **base_params, - "repeat_penalty": 1.05, # Minimal for larger models - "presence_penalty": 0.5, # Light touch - "top_p": 0.95, # High creativity - "top_k": 50, # Many choices - "num_predict": 3000, # Longer responses OK + "repeat_penalty": 1.05, # Minimal for larger models + "presence_penalty": 0.5, # Light touch + "top_p": 0.95, # High creativity + "top_k": 50, # Many choices + "num_predict": 3000, # Longer responses OK } - + return base_params + # Quick test + + def test_safeguards(): """Test the safeguard system.""" detector = ModelRunawayDetector() - + # Test repetition detection bad_response = "The user authentication system works by checking user credentials. The user authentication system works by checking user credentials. The user authentication system works by checking user credentials." - - is_valid, issue, explanation = detector.check_response_quality(bad_response, "auth", time.time()) - + + is_valid, issue, explanation = detector.check_response_quality( + bad_response, "auth", time.time() + ) + print(f"Repetition test: Valid={is_valid}, Issue={issue}") if explanation: print(explanation) + if __name__ == "__main__": - test_safeguards() \ No newline at end of file + test_safeguards() diff --git a/mini_rag/llm_synthesizer.py b/mini_rag/llm_synthesizer.py index 8653904..5cce0c8 100644 --- a/mini_rag/llm_synthesizer.py +++ b/mini_rag/llm_synthesizer.py @@ -9,78 +9,110 @@ Takes raw search results and generates coherent, contextual summaries. import json import logging import time -from typing import List, Dict, Any, Optional from dataclasses import dataclass -import requests from pathlib import Path +from typing import Any, List, Optional + +import requests try: - from .llm_safeguards import ModelRunawayDetector, SafeguardConfig, get_optimal_ollama_parameters + from .llm_safeguards import ( + ModelRunawayDetector, + SafeguardConfig, + get_optimal_ollama_parameters, + ) from .system_context import get_system_context except ImportError: # Graceful fallback if safeguards not available ModelRunawayDetector = None SafeguardConfig = None - get_optimal_ollama_parameters = lambda x: {} - get_system_context = lambda x=None: "" + + def get_optimal_ollama_parameters(x): + return {} + + def get_system_context(x=None): + return "" + logger = logging.getLogger(__name__) + @dataclass class SynthesisResult: """Result of LLM synthesis.""" + summary: str key_points: List[str] code_examples: List[str] suggested_actions: List[str] confidence: float + class LLMSynthesizer: """Synthesizes RAG search results using Ollama LLMs.""" - - def __init__(self, ollama_url: str = "http://localhost:11434", model: str = None, enable_thinking: bool = False, config=None): - self.ollama_url = ollama_url.rstrip('/') + + def __init__( + self, + ollama_url: str = "http://localhost:11434", + model: str = None, + enable_thinking: bool = False, + config=None, + ): + self.ollama_url = ollama_url.rstrip("/") self.available_models = [] self.model = model self.enable_thinking = enable_thinking # Default False for synthesis mode self._initialized = False self.config = config # For accessing model rankings - + # Initialize safeguards if ModelRunawayDetector: self.safeguard_detector = ModelRunawayDetector(SafeguardConfig()) else: self.safeguard_detector = None - + def _get_available_models(self) -> List[str]: """Get list of available Ollama models.""" try: response = requests.get(f"{self.ollama_url}/api/tags", timeout=5) if response.status_code == 200: data = response.json() - return [model['name'] for model in data.get('models', [])] + return [model["name"] for model in data.get("models", [])] except Exception as e: logger.warning(f"Could not fetch Ollama models: {e}") return [] - + def _select_best_model(self) -> str: """Select the best available model based on configuration rankings.""" if not self.available_models: # Use config fallback if available, otherwise use default - if self.config and hasattr(self.config, 'llm') and hasattr(self.config.llm, 'model_rankings') and self.config.llm.model_rankings: + if ( + self.config + and hasattr(self.config, "llm") + and hasattr(self.config.llm, "model_rankings") + and self.config.llm.model_rankings + ): return self.config.llm.model_rankings[0] # First preferred model return "qwen2.5:1.5b" # System fallback only if no config - + # Get model rankings from config or use defaults - if self.config and hasattr(self.config, 'llm') and hasattr(self.config.llm, 'model_rankings'): + if ( + self.config + and hasattr(self.config, "llm") + and hasattr(self.config.llm, "model_rankings") + ): model_rankings = self.config.llm.model_rankings else: # Fallback rankings if no config model_rankings = [ - "qwen3:1.7b", "qwen3:0.6b", "qwen3:4b", "qwen2.5:3b", - "qwen2.5:1.5b", "qwen2.5-coder:1.5b" + "qwen3:1.7b", + "qwen3:0.6b", + "qwen3:4b", + "qwen2.5:3b", + "qwen2.5:1.5b", + "qwen2.5-coder:1.5b", ] - + # Find first available model from our ranked list (exact matches first) for preferred_model in model_rankings: for available_model in self.available_models: @@ -88,102 +120,113 @@ class LLMSynthesizer: if preferred_model.lower() == available_model.lower(): logger.info(f"Selected exact match model: {available_model}") return available_model - + # Partial match with version handling (e.g., "qwen3:1.7b" matches "qwen3:1.7b-q8_0") - preferred_parts = preferred_model.lower().split(':') - available_parts = available_model.lower().split(':') - + preferred_parts = preferred_model.lower().split(":") + available_parts = available_model.lower().split(":") + if len(preferred_parts) >= 2 and len(available_parts) >= 2: - if (preferred_parts[0] == available_parts[0] and - preferred_parts[1] in available_parts[1]): + if ( + preferred_parts[0] == available_parts[0] + and preferred_parts[1] in available_parts[1] + ): logger.info(f"Selected version match model: {available_model}") return available_model - + # If no preferred models found, use first available fallback = self.available_models[0] logger.warning(f"Using fallback model: {fallback}") return fallback - + def _ensure_initialized(self): """Lazy initialization with LLM warmup.""" if self._initialized: return - + # Load available models self.available_models = self._get_available_models() if not self.model: self.model = self._select_best_model() - + # Skip warmup - models are fast enough and warmup causes delays # Warmup removed to eliminate startup delays and unwanted model calls - + self._initialized = True - + def _get_optimal_context_size(self, model_name: str) -> int: """Get optimal context size based on model capabilities and configuration.""" # Get configured context window - if self.config and hasattr(self.config, 'llm'): + if self.config and hasattr(self.config, "llm"): configured_context = self.config.llm.context_window - auto_context = getattr(self.config.llm, 'auto_context', True) + auto_context = getattr(self.config.llm, "auto_context", True) else: configured_context = 16384 # Default to 16K auto_context = True - + # Model-specific maximum context windows (based on research) model_limits = { # Qwen3 models with native context support - 'qwen3:0.6b': 32768, # 32K native - 'qwen3:1.7b': 32768, # 32K native - 'qwen3:4b': 131072, # 131K with YaRN extension - + "qwen3:0.6b": 32768, # 32K native + "qwen3:1.7b": 32768, # 32K native + "qwen3:4b": 131072, # 131K with YaRN extension # Qwen2.5 models - 'qwen2.5:1.5b': 32768, # 32K native - 'qwen2.5:3b': 32768, # 32K native - 'qwen2.5-coder:1.5b': 32768, # 32K native - + "qwen2.5:1.5b": 32768, # 32K native + "qwen2.5:3b": 32768, # 32K native + "qwen2.5-coder:1.5b": 32768, # 32K native # Fallback for unknown models - 'default': 8192 + "default": 8192, } - + # Find model limit (check for partial matches) - model_limit = model_limits.get('default', 8192) + model_limit = model_limits.get("default", 8192) for model_pattern, limit in model_limits.items(): - if model_pattern != 'default' and model_pattern.lower() in model_name.lower(): + if model_pattern != "default" and model_pattern.lower() in model_name.lower(): model_limit = limit break - + # If auto_context is enabled, respect model limits if auto_context: optimal_context = min(configured_context, model_limit) else: optimal_context = configured_context - + # Ensure minimum usable context for RAG optimal_context = max(optimal_context, 4096) # Minimum 4K for basic RAG - - logger.debug(f"Context for {model_name}: {optimal_context} tokens (configured: {configured_context}, limit: {model_limit})") + + logger.debug( + f"Context for {model_name}: {optimal_context} tokens (configured: {configured_context}, limit: {model_limit})" + ) return optimal_context - + def is_available(self) -> bool: """Check if Ollama is available and has models.""" self._ensure_initialized() return len(self.available_models) > 0 - - def _call_ollama(self, prompt: str, temperature: float = 0.3, disable_thinking: bool = False, use_streaming: bool = True, collapse_thinking: bool = True) -> Optional[str]: + + def _call_ollama( + self, + prompt: str, + temperature: float = 0.3, + disable_thinking: bool = False, + use_streaming: bool = True, + collapse_thinking: bool = True, + ) -> Optional[str]: """Make a call to Ollama API with safeguards.""" start_time = time.time() - + try: # Ensure we're initialized self._ensure_initialized() - + # Use the best available model with retry logic model_to_use = self.model if self.model not in self.available_models: # Refresh model list in case of race condition - logger.warning(f"Configured model {self.model} not in available list, refreshing...") + logger.warning( + f"Configured model {self.model} not in available list, refreshing..." + ) self.available_models = self._get_available_models() - + if self.model in self.available_models: model_to_use = self.model logger.info(f"Model {self.model} found after refresh") @@ -194,19 +237,19 @@ class LLMSynthesizer: else: logger.error("No Ollama models available") return None - + # Handle thinking mode for Qwen3 models final_prompt = prompt use_thinking = self.enable_thinking and not disable_thinking - + # For non-thinking mode, add tag for Qwen3 if not use_thinking and "qwen3" in model_to_use.lower(): if not final_prompt.endswith(" "): final_prompt += " " - - # Get optimal parameters for this model + + # Get optimal parameters for this model optimal_params = get_optimal_ollama_parameters(model_to_use) - + # Qwen3-specific optimal parameters based on research if "qwen3" in model_to_use.lower(): if use_thinking: @@ -216,7 +259,7 @@ class LLMSynthesizer: qwen3_top_k = 20 qwen3_presence = 1.5 else: - # Non-thinking mode: Temperature=0.7, TopP=0.8, TopK=20, PresencePenalty=1.5 + # Non-thinking mode: Temperature=0.7, TopP=0.8, TopK=20, PresencePenalty=1.5 qwen3_temp = 0.7 qwen3_top_p = 0.8 qwen3_top_k = 20 @@ -226,7 +269,7 @@ class LLMSynthesizer: qwen3_top_p = optimal_params.get("top_p", 0.9) qwen3_top_k = optimal_params.get("top_k", 40) qwen3_presence = optimal_params.get("presence_penalty", 1.0) - + payload = { "model": model_to_use, "prompt": final_prompt, @@ -235,74 +278,92 @@ class LLMSynthesizer: "temperature": qwen3_temp, "top_p": qwen3_top_p, "top_k": qwen3_top_k, - "num_ctx": self._get_optimal_context_size(model_to_use), # Dynamic context based on model and config + "num_ctx": self._get_optimal_context_size( + model_to_use + ), # Dynamic context based on model and config "num_predict": optimal_params.get("num_predict", 2000), "repeat_penalty": optimal_params.get("repeat_penalty", 1.1), - "presence_penalty": qwen3_presence - } + "presence_penalty": qwen3_presence, + }, } - + # Handle streaming with thinking display if use_streaming: - return self._handle_streaming_with_thinking_display(payload, model_to_use, use_thinking, start_time, collapse_thinking) - + return self._handle_streaming_with_thinking_display( + payload, model_to_use, use_thinking, start_time, collapse_thinking + ) + response = requests.post( f"{self.ollama_url}/api/generate", json=payload, - timeout=65 # Slightly longer than safeguard timeout + timeout=65, # Slightly longer than safeguard timeout ) - + if response.status_code == 200: result = response.json() - + # All models use standard response format # Qwen3 thinking tokens are embedded in the response content itself as ... - raw_response = result.get('response', '').strip() - + raw_response = result.get("response", "").strip() + # Log thinking content for Qwen3 debugging - if "qwen3" in model_to_use.lower() and use_thinking and "" in raw_response: + if ( + "qwen3" in model_to_use.lower() + and use_thinking + and "" in raw_response + ): thinking_start = raw_response.find("") thinking_end = raw_response.find("") if thinking_start != -1 and thinking_end != -1: - thinking_content = raw_response[thinking_start+7:thinking_end] + thinking_content = raw_response[thinking_start + 7 : thinking_end] logger.info(f"Qwen3 thinking: {thinking_content[:100]}...") - + # Apply safeguards to check response quality if self.safeguard_detector and raw_response: - is_valid, issue_type, explanation = self.safeguard_detector.check_response_quality( - raw_response, prompt[:100], start_time # First 100 chars of prompt for context + is_valid, issue_type, explanation = ( + self.safeguard_detector.check_response_quality( + raw_response, + prompt[:100], + start_time, # First 100 chars of prompt for context + ) ) - + if not is_valid: logger.warning(f"Safeguard triggered: {issue_type}") # Preserve original response but add safeguard warning - return self._create_safeguard_response_with_content(issue_type, explanation, raw_response) - + return self._create_safeguard_response_with_content( + issue_type, explanation, raw_response + ) + # Clean up thinking tags from final response cleaned_response = raw_response - if '' in cleaned_response or '' in cleaned_response: + if "" in cleaned_response or "" in cleaned_response: # Remove thinking content but preserve the rest - cleaned_response = cleaned_response.replace('', '').replace('', '') + cleaned_response = cleaned_response.replace("", "").replace( + "", "" + ) # Clean up extra whitespace that might be left - lines = cleaned_response.split('\n') + lines = cleaned_response.split("\n") cleaned_lines = [] for line in lines: if line.strip(): # Only keep non-empty lines cleaned_lines.append(line) - cleaned_response = '\n'.join(cleaned_lines) - + cleaned_response = "\n".join(cleaned_lines) + return cleaned_response.strip() else: logger.error(f"Ollama API error: {response.status_code}") return None - + except Exception as e: logger.error(f"Ollama call failed: {e}") return None - - def _create_safeguard_response(self, issue_type: str, explanation: str, original_prompt: str) -> str: + + def _create_safeguard_response( + self, issue_type: str, explanation: str, original_prompt: str + ) -> str: """Create a helpful response when safeguards are triggered.""" - return f"""โš ๏ธ Model Response Issue Detected + return """โš ๏ธ Model Response Issue Detected {explanation} @@ -312,25 +373,27 @@ class LLMSynthesizer: **Your options:** 1. **Try again**: Ask the same question (often resolves itself) -2. **Rephrase**: Make your question more specific or break it into parts +2. **Rephrase**: Make your question more specific or break it into parts 3. **Use exploration mode**: `rag-mini explore` for complex questions 4. **Different approach**: Try synthesis mode: `--synthesize` for simpler responses This is normal with smaller AI models and helps ensure you get quality responses.""" - def _create_safeguard_response_with_content(self, issue_type: str, explanation: str, original_response: str) -> str: + def _create_safeguard_response_with_content( + self, issue_type: str, explanation: str, original_response: str + ) -> str: """Create a response that preserves the original content but adds a safeguard warning.""" - + # For Qwen3, extract the actual response (after thinking) actual_response = original_response if "" in original_response and "" in original_response: thinking_end = original_response.find("") if thinking_end != -1: - actual_response = original_response[thinking_end + 8:].strip() - + actual_response = original_response[thinking_end + 8 :].strip() + # If we have useful content, preserve it with a warning if len(actual_response.strip()) > 20: - return f"""โš ๏ธ **Response Quality Warning** ({issue_type}) + return """โš ๏ธ **Response Quality Warning** ({issue_type}) {explanation} @@ -345,7 +408,7 @@ This is normal with smaller AI models and helps ensure you get quality responses ๐Ÿ’ก **Note**: This response may have quality issues. Consider rephrasing your question or trying exploration mode for better results.""" else: # If content is too short or problematic, use the original safeguard response - return f"""โš ๏ธ Model Response Issue Detected + return """โš ๏ธ Model Response Issue Detected {explanation} @@ -353,87 +416,106 @@ This is normal with smaller AI models and helps ensure you get quality responses **Your options:** 1. **Try again**: Ask the same question (often resolves itself) -2. **Rephrase**: Make your question more specific or break it into parts +2. **Rephrase**: Make your question more specific or break it into parts 3. **Use exploration mode**: `rag-mini explore` for complex questions This is normal with smaller AI models and helps ensure you get quality responses.""" - def _handle_streaming_with_thinking_display(self, payload: dict, model_name: str, use_thinking: bool, start_time: float, collapse_thinking: bool = True) -> Optional[str]: + def _handle_streaming_with_thinking_display( + self, + payload: dict, + model_name: str, + use_thinking: bool, + start_time: float, + collapse_thinking: bool = True, + ) -> Optional[str]: """Handle streaming response with real-time thinking token display.""" import json - import sys - + try: response = requests.post( - f"{self.ollama_url}/api/generate", - json=payload, - stream=True, - timeout=65 + f"{self.ollama_url}/api/generate", json=payload, stream=True, timeout=65 ) - + if response.status_code != 200: logger.error(f"Ollama API error: {response.status_code}") return None - + full_response = "" thinking_content = "" is_in_thinking = False is_thinking_complete = False thinking_lines_printed = 0 - + # ANSI escape codes for colors and cursor control - GRAY = '\033[90m' # Dark gray for thinking - LIGHT_GRAY = '\033[37m' # Light gray alternative - RESET = '\033[0m' # Reset color - CLEAR_LINE = '\033[2K' # Clear entire line - CURSOR_UP = '\033[A' # Move cursor up one line - + GRAY = "\033[90m" # Dark gray for thinking + # "\033[37m" # Light gray alternative # Unused variable removed + RESET = "\033[0m" # Reset color + CLEAR_LINE = "\033[2K" # Clear entire line + CURSOR_UP = "\033[A" # Move cursor up one line + print(f"\n๐Ÿ’ญ {GRAY}Thinking...{RESET}", flush=True) - + for line in response.iter_lines(): if line: try: - chunk_data = json.loads(line.decode('utf-8')) - chunk_text = chunk_data.get('response', '') - + chunk_data = json.loads(line.decode("utf-8")) + chunk_text = chunk_data.get("response", "") + if chunk_text: full_response += chunk_text - + # Handle thinking tokens - if use_thinking and '' in chunk_text: + if use_thinking and "" in chunk_text: is_in_thinking = True - chunk_text = chunk_text.replace('', '') - - if is_in_thinking and '' in chunk_text: + chunk_text = chunk_text.replace("", "") + + if is_in_thinking and "" in chunk_text: is_in_thinking = False is_thinking_complete = True - chunk_text = chunk_text.replace('', '') - + chunk_text = chunk_text.replace("", "") + if collapse_thinking: # Clear thinking content and show completion # Move cursor up to clear thinking lines for _ in range(thinking_lines_printed + 1): - print(f"{CURSOR_UP}{CLEAR_LINE}", end='', flush=True) - - print(f"๐Ÿ’ญ {GRAY}Thinking complete โœ“{RESET}", flush=True) + print( + f"{CURSOR_UP}{CLEAR_LINE}", + end="", + flush=True, + ) + + print( + f"๐Ÿ’ญ {GRAY}Thinking complete โœ“{RESET}", + flush=True, + ) thinking_lines_printed = 0 else: # Keep thinking visible, just show completion - print(f"\n๐Ÿ’ญ {GRAY}Thinking complete โœ“{RESET}", flush=True) - + print( + f"\n๐Ÿ’ญ {GRAY}Thinking complete โœ“{RESET}", + flush=True, + ) + print("๐Ÿค– AI Response:", flush=True) continue - + # Display thinking content in gray with better formatting if is_in_thinking and chunk_text.strip(): thinking_content += chunk_text - + # Handle line breaks and word wrapping properly - if ' ' in chunk_text or '\n' in chunk_text or len(thinking_content) > 100: + if ( + " " in chunk_text + or "\n" in chunk_text + or len(thinking_content) > 100 + ): # Split by sentences for better readability - sentences = thinking_content.replace('\n', ' ').split('. ') - - for sentence in sentences[:-1]: # Process complete sentences + sentences = thinking_content.replace("\n", " ").split(". ") + + for sentence in sentences[ + :-1 + ]: # Process complete sentences sentence = sentence.strip() if sentence: # Word wrap long sentences @@ -442,136 +524,172 @@ This is normal with smaller AI models and helps ensure you get quality responses for word in words: if len(line + " " + word) > 70: if line: - print(f"{GRAY} {line.strip()}{RESET}", flush=True) + print( + f"{GRAY} {line.strip()}{RESET}", + flush=True, + ) thinking_lines_printed += 1 line = word else: line += " " + word if line else word - + if line.strip(): - print(f"{GRAY} {line.strip()}.{RESET}", flush=True) + print( + f"{GRAY} {line.strip()}.{RESET}", + flush=True, + ) thinking_lines_printed += 1 - + # Keep the last incomplete sentence for next iteration thinking_content = sentences[-1] if sentences else "" - + # Display regular response content (skip any leftover thinking) - elif not is_in_thinking and is_thinking_complete and chunk_text.strip(): + elif ( + not is_in_thinking + and is_thinking_complete + and chunk_text.strip() + ): # Filter out any remaining thinking tags that might leak through clean_text = chunk_text - if '' in clean_text or '' in clean_text: - clean_text = clean_text.replace('', '').replace('', '') - + if "" in clean_text or "" in clean_text: + clean_text = clean_text.replace("", "").replace( + "", "" + ) + if clean_text: # Remove .strip() here to preserve whitespace # Preserve all formatting including newlines and spaces - print(clean_text, end='', flush=True) - + print(clean_text, end="", flush=True) + # Check if response is done - if chunk_data.get('done', False): + if chunk_data.get("done", False): print() # Final newline break - + except json.JSONDecodeError: continue except Exception as e: logger.error(f"Error processing stream chunk: {e}") continue - + return full_response - + except Exception as e: logger.error(f"Streaming failed: {e}") return None - def _handle_streaming_with_early_stop(self, payload: dict, model_name: str, use_thinking: bool, start_time: float) -> Optional[str]: + def _handle_streaming_with_early_stop( + self, payload: dict, model_name: str, use_thinking: bool, start_time: float + ) -> Optional[str]: """Handle streaming response with intelligent early stopping.""" import json - + try: response = requests.post( - f"{self.ollama_url}/api/generate", - json=payload, - stream=True, - timeout=65 + f"{self.ollama_url}/api/generate", json=payload, stream=True, timeout=65 ) - + if response.status_code != 200: logger.error(f"Ollama API error: {response.status_code}") return None - + full_response = "" word_buffer = [] repetition_window = 30 # Check last 30 words for repetition (more context) - stop_threshold = 0.8 # Stop only if 80% of recent words are repetitive (very permissive) + stop_threshold = ( + 0.8 # Stop only if 80% of recent words are repetitive (very permissive) + ) min_response_length = 100 # Don't early stop until we have at least 100 chars - + for line in response.iter_lines(): if line: try: - chunk_data = json.loads(line.decode('utf-8')) - chunk_text = chunk_data.get('response', '') - + chunk_data = json.loads(line.decode("utf-8")) + chunk_text = chunk_data.get("response", "") + if chunk_text: full_response += chunk_text - + # Add words to buffer for repetition detection new_words = chunk_text.split() word_buffer.extend(new_words) - + # Keep only recent words in buffer if len(word_buffer) > repetition_window: word_buffer = word_buffer[-repetition_window:] - + # Check for repetition patterns after we have enough words AND content - if len(word_buffer) >= repetition_window and len(full_response) >= min_response_length: + if ( + len(word_buffer) >= repetition_window + and len(full_response) >= min_response_length + ): unique_words = set(word_buffer) repetition_ratio = 1 - (len(unique_words) / len(word_buffer)) - + # Early stop only if repetition is EXTREMELY high (80%+) if repetition_ratio > stop_threshold: - logger.info(f"Early stopping due to repetition: {repetition_ratio:.2f}") - + logger.info( + f"Early stopping due to repetition: {repetition_ratio:.2f}" + ) + # Add a gentle completion to the response - if not full_response.strip().endswith(('.', '!', '?')): + if not full_response.strip().endswith((".", "!", "?")): full_response += "..." - + # Send stop signal to model (attempt to gracefully stop) try: - stop_payload = {"model": model_name, "stop": True} - requests.post(f"{self.ollama_url}/api/generate", json=stop_payload, timeout=2) - except: + stop_payload = { + "model": model_name, + "stop": True, + } + requests.post( + f"{self.ollama_url}/api/generate", + json=stop_payload, + timeout=2, + ) + except ( + ConnectionError, + FileNotFoundError, + IOError, + OSError, + TimeoutError, + requests.RequestException, + ): pass # If stop fails, we already have partial response - + break - - if chunk_data.get('done', False): + + if chunk_data.get("done", False): break - + except json.JSONDecodeError: continue - + # Clean up thinking tags from final response cleaned_response = full_response - if '' in cleaned_response or '' in cleaned_response: + if "" in cleaned_response or "" in cleaned_response: # Remove thinking content but preserve the rest - cleaned_response = cleaned_response.replace('', '').replace('', '') + cleaned_response = cleaned_response.replace("", "").replace( + "", "" + ) # Clean up extra whitespace that might be left - lines = cleaned_response.split('\n') + lines = cleaned_response.split("\n") cleaned_lines = [] for line in lines: if line.strip(): # Only keep non-empty lines cleaned_lines.append(line) - cleaned_response = '\n'.join(cleaned_lines) - + cleaned_response = "\n".join(cleaned_lines) + return cleaned_response.strip() - + except Exception as e: logger.error(f"Streaming with early stop failed: {e}") return None - - def synthesize_search_results(self, query: str, results: List[Any], project_path: Path) -> SynthesisResult: + + def synthesize_search_results( + self, query: str, results: List[Any], project_path: Path + ) -> SynthesisResult: """Synthesize search results into a coherent summary.""" - + self._ensure_initialized() if not self.is_available(): return SynthesisResult( @@ -579,29 +697,31 @@ This is normal with smaller AI models and helps ensure you get quality responses key_points=[], code_examples=[], suggested_actions=["Install and run Ollama with a model"], - confidence=0.0 + confidence=0.0, ) - + # Prepare context from search results context_parts = [] for i, result in enumerate(results[:8], 1): # Limit to top 8 results - file_path = result.file_path if hasattr(result, 'file_path') else 'unknown' - content = result.content if hasattr(result, 'content') else str(result) - score = result.score if hasattr(result, 'score') else 0.0 - - context_parts.append(f""" + # result.file_path if hasattr(result, "file_path") else "unknown" # Unused variable removed + # result.content if hasattr(result, "content") else str(result) # Unused variable removed + # result.score if hasattr(result, "score") else 0.0 # Unused variable removed + + context_parts.append( + """ Result {i} (Score: {score:.3f}): File: {file_path} Content: {content[:500]}{'...' if len(content) > 500 else ''} -""") - - context = "\n".join(context_parts) - +""" + ) + + # "\n".join(context_parts) # Unused variable removed + # Get system context for better responses - system_context = get_system_context(project_path) - + # get_system_context(project_path) # Unused variable removed + # Create synthesis prompt with system context - prompt = f"""You are a senior software engineer analyzing code search results. Your task is to synthesize the search results into a helpful, actionable summary. + prompt = """You are a senior software engineer analyzing code search results. Your task is to synthesize the search results into a helpful, actionable summary. SYSTEM CONTEXT: {system_context} SEARCH QUERY: "{query}" @@ -615,7 +735,7 @@ Please provide a synthesis in the following JSON format: "summary": "A 2-3 sentence overview of what the search results show", "key_points": [ "Important finding 1", - "Important finding 2", + "Important finding 2", "Important finding 3" ], "code_examples": [ @@ -639,42 +759,42 @@ Respond with ONLY the JSON, no other text.""" # Get LLM response response = self._call_ollama(prompt, temperature=0.2) - + if not response: return SynthesisResult( summary="LLM synthesis failed (API error)", key_points=[], code_examples=[], suggested_actions=["Check Ollama status and try again"], - confidence=0.0 + confidence=0.0, ) - + # Parse JSON response try: # Extract JSON from response (in case there's extra text) - start_idx = response.find('{') - end_idx = response.rfind('}') + 1 + start_idx = response.find("{") + end_idx = response.rfind("}") + 1 if start_idx >= 0 and end_idx > start_idx: json_str = response[start_idx:end_idx] data = json.loads(json_str) - + return SynthesisResult( - summary=data.get('summary', 'No summary generated'), - key_points=data.get('key_points', []), - code_examples=data.get('code_examples', []), - suggested_actions=data.get('suggested_actions', []), - confidence=float(data.get('confidence', 0.5)) + summary=data.get("summary", "No summary generated"), + key_points=data.get("key_points", []), + code_examples=data.get("code_examples", []), + suggested_actions=data.get("suggested_actions", []), + confidence=float(data.get("confidence", 0.5)), ) else: # Fallback: use the raw response as summary return SynthesisResult( - summary=response[:300] + '...' if len(response) > 300 else response, + summary=response[:300] + "..." if len(response) > 300 else response, key_points=[], code_examples=[], suggested_actions=[], - confidence=0.3 + confidence=0.3, ) - + except Exception as e: logger.error(f"Failed to parse LLM response: {e}") return SynthesisResult( @@ -682,75 +802,89 @@ Respond with ONLY the JSON, no other text.""" key_points=[], code_examples=[], suggested_actions=["Try the search again or check LLM output"], - confidence=0.0 + confidence=0.0, ) - + def format_synthesis_output(self, synthesis: SynthesisResult, query: str) -> str: """Format synthesis result for display.""" - + output = [] output.append("๐Ÿง  LLM SYNTHESIS") output.append("=" * 50) output.append("") - - output.append(f"๐Ÿ“ Summary:") + + output.append("๐Ÿ“ Summary:") output.append(f" {synthesis.summary}") output.append("") - + if synthesis.key_points: output.append("๐Ÿ” Key Findings:") for point in synthesis.key_points: output.append(f" โ€ข {point}") output.append("") - + if synthesis.code_examples: output.append("๐Ÿ’ก Code Patterns:") for example in synthesis.code_examples: output.append(f" {example}") output.append("") - + if synthesis.suggested_actions: output.append("๐ŸŽฏ Suggested Actions:") for action in synthesis.suggested_actions: output.append(f" โ€ข {action}") output.append("") - - confidence_emoji = "๐ŸŸข" if synthesis.confidence > 0.7 else "๐ŸŸก" if synthesis.confidence > 0.4 else "๐Ÿ”ด" + + confidence_emoji = ( + "๐ŸŸข" + if synthesis.confidence > 0.7 + else "๐ŸŸก" if synthesis.confidence > 0.4 else "๐Ÿ”ด" + ) output.append(f"{confidence_emoji} Confidence: {synthesis.confidence:.1%}") output.append("") - + return "\n".join(output) + # Quick test function + + def test_synthesizer(): """Test the synthesizer with sample data.""" from dataclasses import dataclass - - @dataclass + + @dataclass class MockResult: file_path: str content: str score: float - + synthesizer = LLMSynthesizer() - + if not synthesizer.is_available(): print("โŒ Ollama not available for testing") return - + # Mock search results results = [ - MockResult("auth.py", "def authenticate_user(username, password):\n return verify_credentials(username, password)", 0.95), - MockResult("models.py", "class User:\n def login(self):\n return authenticate_user(self.username, self.password)", 0.87) + MockResult( + "auth.py", + "def authenticate_user(username, password):\n return verify_credentials(username, password)", + 0.95, + ), + MockResult( + "models.py", + "class User:\n def login(self):\n return authenticate_user(self.username, self.password)", + 0.87, + ), ] - + synthesis = synthesizer.synthesize_search_results( - "user authentication", - results, - Path("/test/project") + "user authentication", results, Path("/test/project") ) - + print(synthesizer.format_synthesis_output(synthesis, "user authentication")) + if __name__ == "__main__": - test_synthesizer() \ No newline at end of file + test_synthesizer() diff --git a/mini_rag/non_invasive_watcher.py b/mini_rag/non_invasive_watcher.py index 996deff..e0a816b 100644 --- a/mini_rag/non_invasive_watcher.py +++ b/mini_rag/non_invasive_watcher.py @@ -3,16 +3,16 @@ Non-invasive file watcher designed to not interfere with development workflows. Uses minimal resources and gracefully handles high-load scenarios. """ -import os -import time import logging -import threading import queue +import threading +import time +from datetime import datetime from pathlib import Path from typing import Optional, Set -from datetime import datetime + +from watchdog.events import DirModifiedEvent, FileSystemEventHandler from watchdog.observers import Observer -from watchdog.events import FileSystemEventHandler, DirModifiedEvent from .indexer import ProjectIndexer @@ -21,7 +21,7 @@ logger = logging.getLogger(__name__) class NonInvasiveQueue: """Ultra-lightweight queue with aggressive deduplication and backoff.""" - + def __init__(self, delay: float = 5.0, max_queue_size: int = 100): self.queue = queue.Queue(maxsize=max_queue_size) self.pending = set() @@ -29,28 +29,28 @@ class NonInvasiveQueue: self.delay = delay self.last_update = {} self.dropped_count = 0 - + def add(self, file_path: Path) -> bool: """Add file to queue with aggressive filtering.""" with self.lock: file_str = str(file_path) current_time = time.time() - + # Skip if recently processed if file_str in self.last_update: if current_time - self.last_update[file_str] < self.delay: return False - + # Skip if already pending if file_str in self.pending: return False - + # Skip if queue is getting full (backpressure) if self.queue.qsize() > self.queue.maxsize * 0.8: self.dropped_count += 1 logger.debug(f"Dropping update for {file_str} - queue overloaded") return False - + try: self.queue.put_nowait(file_path) self.pending.add(file_str) @@ -59,7 +59,7 @@ class NonInvasiveQueue: except queue.Full: self.dropped_count += 1 return False - + def get(self, timeout: float = 0.1) -> Optional[Path]: """Get next file with very short timeout.""" try: @@ -73,77 +73,87 @@ class NonInvasiveQueue: class MinimalEventHandler(FileSystemEventHandler): """Minimal event handler that only watches for meaningful changes.""" - - def __init__(self, - update_queue: NonInvasiveQueue, - include_patterns: Set[str], - exclude_patterns: Set[str]): + + def __init__( + self, + update_queue: NonInvasiveQueue, + include_patterns: Set[str], + exclude_patterns: Set[str], + ): self.update_queue = update_queue self.include_patterns = include_patterns self.exclude_patterns = exclude_patterns self.last_event_time = {} - + def _should_process(self, file_path: str) -> bool: """Ultra-conservative file filtering.""" path = Path(file_path) - + # Only process files, not directories if not path.is_file(): return False - + # Skip if too large (>1MB) try: if path.stat().st_size > 1024 * 1024: return False except (OSError, PermissionError): return False - + # Skip temporary and system files name = path.name - if (name.startswith('.') or - name.startswith('~') or - name.endswith('.tmp') or - name.endswith('.swp') or - name.endswith('.lock')): + if ( + name.startswith(".") + or name.startswith("~") + or name.endswith(".tmp") + or name.endswith(".swp") + or name.endswith(".lock") + ): return False - + # Check exclude patterns first (faster) path_str = str(path) for pattern in self.exclude_patterns: if pattern in path_str: return False - + # Check include patterns for pattern in self.include_patterns: if path.match(pattern): return True - + return False - + def _rate_limit_event(self, file_path: str) -> bool: """Rate limit events per file.""" current_time = time.time() if file_path in self.last_event_time: - if current_time - self.last_event_time[file_path] < 2.0: # 2 second cooldown per file + if ( + current_time - self.last_event_time[file_path] < 2.0 + ): # 2 second cooldown per file return False - + self.last_event_time[file_path] = current_time return True - + def on_modified(self, event): """Handle file modifications with minimal overhead.""" - if (not event.is_directory and - self._should_process(event.src_path) and - self._rate_limit_event(event.src_path)): + if ( + not event.is_directory + and self._should_process(event.src_path) + and self._rate_limit_event(event.src_path) + ): self.update_queue.add(Path(event.src_path)) - + def on_created(self, event): """Handle file creation.""" - if (not event.is_directory and - self._should_process(event.src_path) and - self._rate_limit_event(event.src_path)): + if ( + not event.is_directory + and self._should_process(event.src_path) + and self._rate_limit_event(event.src_path) + ): self.update_queue.add(Path(event.src_path)) - + def on_deleted(self, event): """Handle file deletion.""" if not event.is_directory and self._rate_limit_event(event.src_path): @@ -157,15 +167,17 @@ class MinimalEventHandler(FileSystemEventHandler): class NonInvasiveFileWatcher: """Non-invasive file watcher that prioritizes system stability.""" - - def __init__(self, - project_path: Path, - indexer: Optional[ProjectIndexer] = None, - cpu_limit: float = 0.1, # Max 10% CPU usage - max_memory_mb: int = 50): # Max 50MB memory + + def __init__( + self, + project_path: Path, + indexer: Optional[ProjectIndexer] = None, + cpu_limit: float = 0.1, # Max 10% CPU usage + max_memory_mb: int = 50, + ): # Max 50MB memory """ Initialize non-invasive watcher. - + Args: project_path: Path to watch indexer: ProjectIndexer instance @@ -176,158 +188,173 @@ class NonInvasiveFileWatcher: self.indexer = indexer or ProjectIndexer(self.project_path) self.cpu_limit = cpu_limit self.max_memory_mb = max_memory_mb - + # Initialize components with conservative settings - self.update_queue = NonInvasiveQueue(delay=10.0, max_queue_size=50) # Very conservative + self.update_queue = NonInvasiveQueue( + delay=10.0, max_queue_size=50 + ) # Very conservative self.observer = Observer() self.worker_thread = None self.running = False - + # Get patterns from indexer self.include_patterns = set(self.indexer.include_patterns) self.exclude_patterns = set(self.indexer.exclude_patterns) - + # Add more aggressive exclusions - self.exclude_patterns.update({ - '__pycache__', '.git', 'node_modules', '.venv', 'venv', - 'dist', 'build', 'target', '.idea', '.vscode', '.pytest_cache', - 'coverage', 'htmlcov', '.coverage', '.mypy_cache', '.tox', - 'logs', 'log', 'tmp', 'temp', '.DS_Store' - }) - + self.exclude_patterns.update( + { + "__pycache__", + ".git", + "node_modules", + ".venv", + "venv", + "dist", + "build", + "target", + ".idea", + ".vscode", + ".pytest_cache", + "coverage", + "htmlcov", + ".coverage", + ".mypy_cache", + ".tox", + "logs", + "log", + "tmp", + "temp", + ".DS_Store", + } + ) + # Stats self.stats = { - 'files_processed': 0, - 'files_dropped': 0, - 'cpu_throttle_count': 0, - 'started_at': None, + "files_processed": 0, + "files_dropped": 0, + "cpu_throttle_count": 0, + "started_at": None, } - + def start(self): """Start non-invasive watching.""" if self.running: return - + logger.info(f"Starting non-invasive file watcher for {self.project_path}") - + # Set up minimal event handler event_handler = MinimalEventHandler( - self.update_queue, - self.include_patterns, - self.exclude_patterns + self.update_queue, self.include_patterns, self.exclude_patterns ) - + # Schedule with recursive watching - self.observer.schedule( - event_handler, - str(self.project_path), - recursive=True - ) - + self.observer.schedule(event_handler, str(self.project_path), recursive=True) + # Start low-priority worker thread self.running = True self.worker_thread = threading.Thread( - target=self._process_updates_gently, - daemon=True, - name="RAG-FileWatcher" + target=self._process_updates_gently, daemon=True, name="RAG-FileWatcher" ) # Set lowest priority self.worker_thread.start() - + # Start observer self.observer.start() - - self.stats['started_at'] = datetime.now() + + self.stats["started_at"] = datetime.now() logger.info("Non-invasive file watcher started") - + def stop(self): """Stop watching gracefully.""" if not self.running: return - + logger.info("Stopping non-invasive file watcher...") - + # Stop observer first self.observer.stop() self.observer.join(timeout=2.0) # Don't wait too long - + # Stop worker thread self.running = False if self.worker_thread and self.worker_thread.is_alive(): self.worker_thread.join(timeout=3.0) # Don't block shutdown - + logger.info("Non-invasive file watcher stopped") - + def _process_updates_gently(self): """Process updates with extreme care not to interfere.""" logger.debug("Non-invasive update processor started") - + process_start_time = time.time() - + while self.running: try: # Yield CPU frequently time.sleep(0.5) # Always sleep between operations - + # Get next file with very short timeout file_path = self.update_queue.get(timeout=0.1) - + if file_path: # Check CPU usage before processing current_time = time.time() elapsed = current_time - process_start_time - + # Simple CPU throttling: if we've been working too much, back off if elapsed > 0: # If we're consuming too much time, throttle aggressively work_ratio = 0.1 # Assume we use 10% of time in this check if work_ratio > self.cpu_limit: - self.stats['cpu_throttle_count'] += 1 + self.stats["cpu_throttle_count"] += 1 time.sleep(2.0) # Back off significantly continue - + # Process single file with error isolation try: if file_path.exists(): success = self.indexer.update_file(file_path) else: success = self.indexer.delete_file(file_path) - + if success: - self.stats['files_processed'] += 1 - + self.stats["files_processed"] += 1 + # Always yield CPU after processing time.sleep(0.1) - + except Exception as e: - logger.debug(f"Non-invasive watcher: failed to process {file_path}: {e}") + logger.debug( + f"Non-invasive watcher: failed to process {file_path}: {e}" + ) # Don't let errors propagate - just continue continue - + # Update dropped count from queue - self.stats['files_dropped'] = self.update_queue.dropped_count - + self.stats["files_dropped"] = self.update_queue.dropped_count + except Exception as e: logger.debug(f"Non-invasive watcher error: {e}") time.sleep(1.0) # Back off on errors - + logger.debug("Non-invasive update processor stopped") - + def get_statistics(self) -> dict: """Get non-invasive watcher statistics.""" stats = self.stats.copy() - stats['queue_size'] = self.update_queue.queue.qsize() - stats['running'] = self.running - - if stats['started_at']: - uptime = datetime.now() - stats['started_at'] - stats['uptime_seconds'] = uptime.total_seconds() - + stats["queue_size"] = self.update_queue.queue.qsize() + stats["running"] = self.running + + if stats["started_at"]: + uptime = datetime.now() - stats["started_at"] + stats["uptime_seconds"] = uptime.total_seconds() + return stats - + def __enter__(self): self.start() return self - + def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() \ No newline at end of file + self.stop() diff --git a/mini_rag/ollama_embeddings.py b/mini_rag/ollama_embeddings.py index c546e74..b58958b 100644 --- a/mini_rag/ollama_embeddings.py +++ b/mini_rag/ollama_embeddings.py @@ -3,15 +3,14 @@ Hybrid code embedding module - Ollama primary with ML fallback. Tries Ollama first, falls back to local ML stack if needed. """ -import requests -import numpy as np -from typing import List, Union, Optional, Dict, Any import logging -from functools import lru_cache import time -import json from concurrent.futures import ThreadPoolExecutor -import threading +from functools import lru_cache +from typing import Any, Dict, List, Optional, Union + +import numpy as np +import requests logger = logging.getLogger(__name__) @@ -19,8 +18,9 @@ logger = logging.getLogger(__name__) FALLBACK_AVAILABLE = False try: import torch - from transformers import AutoTokenizer, AutoModel from sentence_transformers import SentenceTransformer + from transformers import AutoModel, AutoTokenizer + FALLBACK_AVAILABLE = True logger.debug("ML fallback dependencies available") except ImportError: @@ -29,12 +29,16 @@ except ImportError: class OllamaEmbedder: """Hybrid embeddings: Ollama primary with ML fallback.""" - - def __init__(self, model_name: str = "nomic-embed-text:latest", base_url: str = "http://localhost:11434", - enable_fallback: bool = True): + + def __init__( + self, + model_name: str = "nomic-embed-text:latest", + base_url: str = "http://localhost:11434", + enable_fallback: bool = True, + ): """ Initialize the hybrid embedder. - + Args: model_name: Ollama model to use for embeddings base_url: Base URL for Ollama API @@ -44,15 +48,15 @@ class OllamaEmbedder: self.base_url = base_url self.embedding_dim = 768 # Standard for nomic-embed-text self.enable_fallback = enable_fallback and FALLBACK_AVAILABLE - + # State tracking self.ollama_available = False self.fallback_embedder = None self.mode = "unknown" # "ollama", "fallback", or "hash" - + # Try to initialize Ollama first self._initialize_providers() - + def _initialize_providers(self): """Initialize embedding providers in priority order.""" # Try Ollama first @@ -64,13 +68,15 @@ class OllamaEmbedder: except Exception as e: logger.debug(f"Ollama not available: {e}") self.ollama_available = False - + # Try ML fallback if self.enable_fallback: try: self._initialize_fallback_embedder() self.mode = "fallback" - logger.info(f"โœ… ML fallback active: {self.fallback_embedder.model_type if hasattr(self.fallback_embedder, 'model_type') else 'transformer'}") + logger.info( + f"โœ… ML fallback active: {self.fallback_embedder.model_type if hasattr(self.fallback_embedder, 'model_type') else 'transformer'}" + ) except Exception as fallback_error: logger.warning(f"ML fallback failed: {fallback_error}") self.mode = "hash" @@ -78,7 +84,7 @@ class OllamaEmbedder: else: self.mode = "hash" logger.info("โš ๏ธ Using hash-based embeddings (no fallback enabled)") - + def _verify_ollama_connection(self): """Verify Ollama server is running and model is available.""" try: @@ -93,17 +99,17 @@ class OllamaEmbedder: print() raise ConnectionError("Ollama service not running. Start with: ollama serve") except requests.exceptions.Timeout: - print("โฑ๏ธ Ollama Service Timeout") + print("โฑ๏ธ Ollama Service Timeout") print(" Ollama is taking too long to respond") print(" Check if Ollama is overloaded: ollama ps") print(" Restart if needed: killall ollama && ollama serve") print() raise ConnectionError("Ollama service timeout") - + # Check if our model is available - models = response.json().get('models', []) - model_names = [model['name'] for model in models] - + models = response.json().get("models", []) + model_names = [model["name"] for model in models] + if self.model_name not in model_names: print(f"๐Ÿ“ฆ Model '{self.model_name}' Not Found") print(" Embedding models convert text into searchable vectors") @@ -113,19 +119,23 @@ class OllamaEmbedder: print() # Try to pull the model self._pull_model() - + def _initialize_fallback_embedder(self): """Initialize the ML fallback embedder.""" if not FALLBACK_AVAILABLE: raise RuntimeError("ML dependencies not available for fallback") - + # Try lightweight models first for better compatibility fallback_models = [ - ("sentence-transformers/all-MiniLM-L6-v2", 384, self._init_sentence_transformer), + ( + "sentence-transformers/all-MiniLM-L6-v2", + 384, + self._init_sentence_transformer, + ), ("microsoft/codebert-base", 768, self._init_transformer_model), ("microsoft/unixcoder-base", 768, self._init_transformer_model), ] - + for model_name, dim, init_func in fallback_models: try: init_func(model_name) @@ -135,31 +145,33 @@ class OllamaEmbedder: except Exception as e: logger.debug(f"Failed to load {model_name}: {e}") continue - + raise RuntimeError("Could not initialize any fallback embedding model") - + def _init_sentence_transformer(self, model_name: str): """Initialize sentence-transformers model.""" self.fallback_embedder = SentenceTransformer(model_name) - self.fallback_embedder.model_type = 'sentence_transformer' - + self.fallback_embedder.model_type = "sentence_transformer" + def _init_transformer_model(self, model_name: str): """Initialize transformer model.""" - device = 'cuda' if torch.cuda.is_available() else 'cpu' + device = "cuda" if torch.cuda.is_available() else "cpu" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModel.from_pretrained(model_name).to(device) model.eval() - + # Create a simple wrapper + class TransformerWrapper: + def __init__(self, model, tokenizer, device): self.model = model self.tokenizer = tokenizer self.device = device - self.model_type = 'transformer' - + self.model_type = "transformer" + self.fallback_embedder = TransformerWrapper(model, tokenizer, device) - + def _pull_model(self): """Pull the embedding model if not available.""" logger.info(f"Pulling model {self.model_name}...") @@ -167,13 +179,13 @@ class OllamaEmbedder: response = requests.post( f"{self.base_url}/api/pull", json={"name": self.model_name}, - timeout=300 # 5 minutes for model download + timeout=300, # 5 minutes for model download ) response.raise_for_status() logger.info(f"Successfully pulled {self.model_name}") except requests.exceptions.RequestException as e: raise RuntimeError(f"Failed to pull model {self.model_name}: {e}") - + def _get_embedding(self, text: str) -> np.ndarray: """Get embedding using the best available provider.""" if self.mode == "ollama" and self.ollama_available: @@ -183,28 +195,25 @@ class OllamaEmbedder: else: # Hash fallback return self._hash_embedding(text) - + def _get_ollama_embedding(self, text: str) -> np.ndarray: """Get embedding from Ollama API.""" try: response = requests.post( f"{self.base_url}/api/embeddings", - json={ - "model": self.model_name, - "prompt": text - }, - timeout=30 + json={"model": self.model_name, "prompt": text}, + timeout=30, ) response.raise_for_status() - + result = response.json() - embedding = result.get('embedding', []) - + embedding = result.get("embedding", []) + if not embedding: raise ValueError("No embedding returned from Ollama") - + return np.array(embedding, dtype=np.float32) - + except requests.exceptions.RequestException as e: logger.error(f"Ollama API request failed: {e}") # Degrade gracefully - try fallback if available @@ -216,82 +225,88 @@ class OllamaEmbedder: except (ValueError, KeyError) as e: logger.error(f"Invalid response from Ollama: {e}") return self._hash_embedding(text) - + def _get_fallback_embedding(self, text: str) -> np.ndarray: """Get embedding from ML fallback.""" try: - if self.fallback_embedder.model_type == 'sentence_transformer': + if self.fallback_embedder.model_type == "sentence_transformer": embedding = self.fallback_embedder.encode([text], convert_to_numpy=True)[0] return embedding.astype(np.float32) - - elif self.fallback_embedder.model_type == 'transformer': + + elif self.fallback_embedder.model_type == "transformer": # Tokenize and generate embedding inputs = self.fallback_embedder.tokenizer( - text, - padding=True, - truncation=True, + text, + padding=True, + truncation=True, max_length=512, - return_tensors="pt" + return_tensors="pt", ).to(self.fallback_embedder.device) - + with torch.no_grad(): outputs = self.fallback_embedder.model(**inputs) - + # Use pooler output if available, otherwise mean pooling - if hasattr(outputs, 'pooler_output') and outputs.pooler_output is not None: + if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None: embedding = outputs.pooler_output[0] else: # Mean pooling over sequence length - attention_mask = inputs['attention_mask'] + attention_mask = inputs["attention_mask"] token_embeddings = outputs.last_hidden_state[0] - + # Mask and average - input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() + input_mask_expanded = ( + attention_mask.unsqueeze(-1) + .expand(token_embeddings.size()) + .float() + ) sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 0) sum_mask = torch.clamp(input_mask_expanded.sum(0), min=1e-9) embedding = sum_embeddings / sum_mask - + return embedding.cpu().numpy().astype(np.float32) - + else: - raise ValueError(f"Unknown fallback model type: {self.fallback_embedder.model_type}") - + raise ValueError( + f"Unknown fallback model type: {self.fallback_embedder.model_type}" + ) + except Exception as e: logger.error(f"Fallback embedding failed: {e}") return self._hash_embedding(text) - + def _hash_embedding(self, text: str) -> np.ndarray: """Generate deterministic hash-based embedding as fallback.""" import hashlib - + # Create deterministic hash - hash_obj = hashlib.sha256(text.encode('utf-8')) + hash_obj = hashlib.sha256(text.encode("utf-8")) hash_bytes = hash_obj.digest() - + # Convert to numbers and normalize hash_nums = np.frombuffer(hash_bytes, dtype=np.uint8) - + # Expand to target dimension using repetition while len(hash_nums) < self.embedding_dim: hash_nums = np.concatenate([hash_nums, hash_nums]) - + # Take exactly the dimension we need - embedding = hash_nums[:self.embedding_dim].astype(np.float32) - + embedding = hash_nums[: self.embedding_dim].astype(np.float32) + # Normalize to [-1, 1] range embedding = (embedding / 127.5) - 1.0 - + logger.debug(f"Using hash fallback embedding for text: {text[:50]}...") return embedding - + def embed_code(self, code: Union[str, List[str]], language: str = "python") -> np.ndarray: """ Generate embeddings for code snippet(s). - + Args: code: Single code string or list of code strings language: Programming language (used for context) - + Returns: Embedding vector(s) as numpy array """ @@ -300,22 +315,22 @@ class OllamaEmbedder: single_input = True else: single_input = False - + # Preprocess code for better embeddings processed_code = [self._preprocess_code(c, language) for c in code] - + # Generate embeddings embeddings = [] for text in processed_code: embedding = self._get_embedding(text) embeddings.append(embedding) - + embeddings = np.array(embeddings, dtype=np.float32) - + if single_input: return embeddings[0] return embeddings - + def _preprocess_code(self, code: str, language: str = "python") -> str: """ Preprocess code for better embedding quality. @@ -323,25 +338,25 @@ class OllamaEmbedder: """ # Remove leading/trailing whitespace code = code.strip() - + # Normalize whitespace but preserve structure - lines = code.split('\n') + lines = code.split("\n") processed_lines = [] - + for line in lines: # Remove trailing whitespace line = line.rstrip() # Keep non-empty lines if line: processed_lines.append(line) - - cleaned_code = '\n'.join(processed_lines) - + + cleaned_code = "\n".join(processed_lines) + # Add language context for better embeddings if language and cleaned_code: return f"```{language}\n{cleaned_code}\n```" return cleaned_code - + @lru_cache(maxsize=1000) def embed_query(self, query: str) -> np.ndarray: """ @@ -351,149 +366,151 @@ class OllamaEmbedder: # Enhance query for code search enhanced_query = f"Search for code related to: {query}" return self._get_embedding(enhanced_query) - + def batch_embed_files(self, file_contents: List[dict], max_workers: int = 4) -> List[dict]: """ Embed multiple files efficiently using concurrent requests to Ollama. - + Args: file_contents: List of dicts with 'content' and optionally 'language' keys max_workers: Maximum number of concurrent Ollama requests - + Returns: List of dicts with added 'embedding' key (preserves original order) """ if not file_contents: return [] - + # For small batches, use sequential processing to avoid overhead if len(file_contents) <= 2: return self._batch_embed_sequential(file_contents) - + # For very large batches, use chunked processing to prevent memory issues if len(file_contents) > 500: # Process in chunks to manage memory return self._batch_embed_chunked(file_contents, max_workers) - + return self._batch_embed_concurrent(file_contents, max_workers) - + def _batch_embed_sequential(self, file_contents: List[dict]) -> List[dict]: """Sequential processing for small batches.""" results = [] for file_dict in file_contents: - content = file_dict['content'] - language = file_dict.get('language', 'python') + content = file_dict["content"] + language = file_dict.get("language", "python") embedding = self.embed_code(content, language) - + result = file_dict.copy() - result['embedding'] = embedding + result["embedding"] = embedding results.append(result) - + return results - - def _batch_embed_concurrent(self, file_contents: List[dict], max_workers: int) -> List[dict]: + + def _batch_embed_concurrent( + self, file_contents: List[dict], max_workers: int + ) -> List[dict]: """Concurrent processing for larger batches.""" + def embed_single(item_with_index): index, file_dict = item_with_index - content = file_dict['content'] - language = file_dict.get('language', 'python') - + content = file_dict["content"] + language = file_dict.get("language", "python") + try: embedding = self.embed_code(content, language) result = file_dict.copy() - result['embedding'] = embedding + result["embedding"] = embedding return index, result except Exception as e: logger.error(f"Failed to embed content at index {index}: {e}") # Return with hash fallback result = file_dict.copy() - result['embedding'] = self._hash_embedding(content) + result["embedding"] = self._hash_embedding(content) return index, result - + # Create indexed items to preserve order indexed_items = list(enumerate(file_contents)) - + # Process concurrently with ThreadPoolExecutor(max_workers=max_workers) as executor: indexed_results = list(executor.map(embed_single, indexed_items)) - + # Sort by original index and extract results indexed_results.sort(key=lambda x: x[0]) return [result for _, result in indexed_results] - - def _batch_embed_chunked(self, file_contents: List[dict], max_workers: int, chunk_size: int = 200) -> List[dict]: + + def _batch_embed_chunked( + self, file_contents: List[dict], max_workers: int, chunk_size: int = 200 + ) -> List[dict]: """ Process very large batches in smaller chunks to prevent memory issues. This is important for beginners who might try to index huge projects. """ results = [] total_chunks = len(file_contents) - + # Process in chunks for i in range(0, len(file_contents), chunk_size): - chunk = file_contents[i:i + chunk_size] - + chunk = file_contents[i : i + chunk_size] + # Log progress for large operations if total_chunks > chunk_size: chunk_num = i // chunk_size + 1 total_chunk_count = (total_chunks + chunk_size - 1) // chunk_size - logger.info(f"Processing chunk {chunk_num}/{total_chunk_count} ({len(chunk)} files)") - + logger.info( + f"Processing chunk {chunk_num}/{total_chunk_count} ({len(chunk)} files)" + ) + # Process this chunk using concurrent method chunk_results = self._batch_embed_concurrent(chunk, max_workers) results.extend(chunk_results) - + # Brief pause between chunks to prevent overwhelming the system if i + chunk_size < len(file_contents): - import time + time.sleep(0.1) # 100ms pause between chunks - + return results - + def get_embedding_dim(self) -> int: """Return the dimension of embeddings produced by this model.""" return self.embedding_dim - + def get_mode(self) -> str: """Return current embedding mode: 'ollama', 'fallback', or 'hash'.""" return self.mode - + def get_status(self) -> Dict[str, Any]: """Get detailed status of the embedding system.""" return { "mode": self.mode, "ollama_available": self.ollama_available, "fallback_available": FALLBACK_AVAILABLE and self.enable_fallback, - "fallback_model": getattr(self.fallback_embedder, 'model_type', None) if self.fallback_embedder else None, + "fallback_model": ( + getattr(self.fallback_embedder, "model_type", None) + if self.fallback_embedder + else None + ), "embedding_dim": self.embedding_dim, "ollama_model": self.model_name if self.mode == "ollama" else None, - "ollama_url": self.base_url if self.mode == "ollama" else None + "ollama_url": self.base_url if self.mode == "ollama" else None, } - + def get_embedding_info(self) -> Dict[str, str]: """Get human-readable embedding system information for installer.""" status = self.get_status() mode = status.get("mode", "unknown") if mode == "ollama": - return { - "method": f"Ollama ({status['ollama_model']})", - "status": "working" - } + return {"method": f"Ollama ({status['ollama_model']})", "status": "working"} # Treat legacy/alternate naming uniformly if mode in ("fallback", "ml"): return { "method": f"ML Fallback ({status['fallback_model']})", - "status": "working" + "status": "working", } if mode == "hash": - return { - "method": "Hash-based (basic similarity)", - "status": "working" - } - return { - "method": "Unknown", - "status": "error" - } - + return {"method": "Hash-based (basic similarity)", "status": "working"} + return {"method": "Unknown", "status": "error"} + def warmup(self): """Warm up the embedding system with a dummy request.""" dummy_code = "def hello(): pass" @@ -503,14 +520,18 @@ class OllamaEmbedder: # Convenience function for quick embedding -def embed_code(code: Union[str, List[str]], model_name: str = "nomic-embed-text:latest") -> np.ndarray: + + +def embed_code( + code: Union[str, List[str]], model_name: str = "nomic-embed-text:latest" +) -> np.ndarray: """ Quick function to embed code without managing embedder instance. - + Args: code: Code string(s) to embed model_name: Ollama model name to use - + Returns: Embedding vector(s) """ @@ -519,4 +540,4 @@ def embed_code(code: Union[str, List[str]], model_name: str = "nomic-embed-text: # Compatibility alias for drop-in replacement -CodeEmbedder = OllamaEmbedder \ No newline at end of file +CodeEmbedder = OllamaEmbedder diff --git a/mini_rag/path_handler.py b/mini_rag/path_handler.py index f7c8746..8cc40c5 100644 --- a/mini_rag/path_handler.py +++ b/mini_rag/path_handler.py @@ -4,51 +4,50 @@ Handles forward/backward slashes on any file system. Robust cross-platform path handling. """ -import os import sys from pathlib import Path -from typing import Union, List +from typing import List, Union def normalize_path(path: Union[str, Path]) -> str: """ Normalize a path to always use forward slashes. This ensures consistency across platforms in storage. - + Args: path: Path as string or Path object - + Returns: Path string with forward slashes """ # Convert to Path object first path_obj = Path(path) - + # Convert to string and replace backslashes - path_str = str(path_obj).replace('\\', '/') - + path_str = str(path_obj).replace("\\", "/") + # Handle UNC paths on Windows - if sys.platform == 'win32' and path_str.startswith('//'): + if sys.platform == "win32" and path_str.startswith("//"): # Keep UNC paths as they are return path_str - + return path_str def normalize_relative_path(path: Union[str, Path], base: Union[str, Path]) -> str: """ Get a normalized relative path. - + Args: path: Path to make relative base: Base path to be relative to - + Returns: Relative path with forward slashes """ path_obj = Path(path).resolve() base_obj = Path(base).resolve() - + try: rel_path = path_obj.relative_to(base_obj) return normalize_path(rel_path) @@ -61,10 +60,10 @@ def denormalize_path(path_str: str) -> Path: """ Convert a normalized path string back to a Path object. This handles the conversion from storage format to OS format. - + Args: path_str: Normalized path string with forward slashes - + Returns: Path object appropriate for the OS """ @@ -75,10 +74,10 @@ def denormalize_path(path_str: str) -> Path: def join_paths(*parts: Union[str, Path]) -> str: """ Join path parts and return normalized result. - + Args: *parts: Path parts to join - + Returns: Normalized joined path """ @@ -90,46 +89,46 @@ def join_paths(*parts: Union[str, Path]) -> str: def split_path(path: Union[str, Path]) -> List[str]: """ Split a path into its components. - + Args: path: Path to split - + Returns: List of path components """ path_obj = Path(path) parts = [] - + # Handle drive on Windows if path_obj.drive: parts.append(path_obj.drive) - + # Add all other parts parts.extend(path_obj.parts[1:] if path_obj.drive else path_obj.parts) - + return parts def ensure_forward_slashes(path_str: str) -> str: """ Quick function to ensure a path string uses forward slashes. - + Args: path_str: Path string - + Returns: Path with forward slashes """ - return path_str.replace('\\', '/') + return path_str.replace("\\", "/") def ensure_native_slashes(path_str: str) -> str: """ Ensure a path uses the native separator for the OS. - + Args: path_str: Path string - + Returns: Path with native separators """ @@ -137,6 +136,8 @@ def ensure_native_slashes(path_str: str) -> str: # Convenience functions for common operations + + def storage_path(path: Union[str, Path]) -> str: """Convert path to storage format (forward slashes).""" return normalize_path(path) @@ -149,4 +150,4 @@ def display_path(path: Union[str, Path]) -> str: def from_storage_path(path_str: str) -> Path: """Convert from storage format to Path object.""" - return denormalize_path(path_str) \ No newline at end of file + return denormalize_path(path_str) diff --git a/mini_rag/performance.py b/mini_rag/performance.py index 2613ceb..1edf02b 100644 --- a/mini_rag/performance.py +++ b/mini_rag/performance.py @@ -3,85 +3,87 @@ Performance monitoring for RAG system. Track loading times, query times, and resource usage. """ -import time -import psutil -import os -from contextlib import contextmanager -from typing import Dict, Any, Optional import logging +import os +import time +from contextlib import contextmanager +from typing import Any, Dict, Optional + +import psutil logger = logging.getLogger(__name__) class PerformanceMonitor: """Track performance metrics for RAG operations.""" - + def __init__(self): self.metrics = {} self.process = psutil.Process(os.getpid()) - + @contextmanager def measure(self, operation: str): """Context manager to measure operation time and memory.""" # Get initial state start_time = time.time() start_memory = self.process.memory_info().rss / 1024 / 1024 # MB - + try: yield self finally: # Calculate metrics end_time = time.time() end_memory = self.process.memory_info().rss / 1024 / 1024 # MB - + duration = end_time - start_time memory_delta = end_memory - start_memory - + # Store metrics self.metrics[operation] = { - 'duration_seconds': duration, - 'memory_delta_mb': memory_delta, - 'final_memory_mb': end_memory, + "duration_seconds": duration, + "memory_delta_mb": memory_delta, + "final_memory_mb": end_memory, } - + logger.info( f"[PERF] {operation}: {duration:.2f}s, " f"Memory: {end_memory:.1f}MB (+{memory_delta:+.1f}MB)" ) - + def get_summary(self) -> Dict[str, Any]: """Get performance summary.""" - total_time = sum(m['duration_seconds'] for m in self.metrics.values()) - + total_time = sum(m["duration_seconds"] for m in self.metrics.values()) + return { - 'total_time_seconds': total_time, - 'operations': self.metrics, - 'current_memory_mb': self.process.memory_info().rss / 1024 / 1024, + "total_time_seconds": total_time, + "operations": self.metrics, + "current_memory_mb": self.process.memory_info().rss / 1024 / 1024, } - + def print_summary(self): """Print a formatted summary.""" - print("\n" + "="*50) + print("\n" + "=" * 50) print("PERFORMANCE SUMMARY") - print("="*50) - + print("=" * 50) + for op, metrics in self.metrics.items(): print(f"\n{op}:") print(f" Time: {metrics['duration_seconds']:.2f}s") print(f" Memory: +{metrics['memory_delta_mb']:+.1f}MB") - + summary = self.get_summary() print(f"\nTotal Time: {summary['total_time_seconds']:.2f}s") print(f"Current Memory: {summary['current_memory_mb']:.1f}MB") - print("="*50) + print("=" * 50) # Global instance for easy access _monitor = None + def get_monitor() -> PerformanceMonitor: """Get or create global monitor instance.""" global _monitor if _monitor is None: _monitor = PerformanceMonitor() - return _monitor \ No newline at end of file + return _monitor diff --git a/mini_rag/query_expander.py b/mini_rag/query_expander.py index 95c00b0..b50f954 100644 --- a/mini_rag/query_expander.py +++ b/mini_rag/query_expander.py @@ -7,7 +7,7 @@ Automatically expands search queries to find more relevant results. Example: "authentication" becomes "authentication login user verification credentials" -## How It Helps +## How It Helps - 2-3x more relevant search results - Works with any content (code, docs, notes, etc.) - Completely transparent to users @@ -26,22 +26,25 @@ expanded = expander.expand_query("error handling") # Result: "error handling exception try catch fault tolerance" ``` -Perfect for beginners - enable in TUI for exploration, +Perfect for beginners - enable in TUI for exploration, disable in CLI for maximum speed. """ import logging import re import threading -from typing import List, Optional +from typing import Optional + import requests + from .config import RAGConfig logger = logging.getLogger(__name__) + class QueryExpander: """Expands search queries using LLM to improve search recall.""" - + def __init__(self, config: RAGConfig): self.config = config self.ollama_url = f"http://{config.llm.ollama_host}" @@ -49,37 +52,37 @@ class QueryExpander: self.max_terms = config.llm.max_expansion_terms self.enabled = config.search.expand_queries self._initialized = False - + # Cache for expanded queries to avoid repeated API calls self._cache = {} self._cache_lock = threading.RLock() # Thread-safe cache access - + def _ensure_initialized(self): """Lazy initialization with LLM warmup.""" if self._initialized: return - + # Skip warmup - causes startup delays and unwanted model calls # Query expansion works fine on first use without warmup - + self._initialized = True - + def expand_query(self, query: str) -> str: """Expand a search query with related terms.""" if not self.enabled or not query.strip(): return query - + self._ensure_initialized() - + # Check cache first (thread-safe) with self._cache_lock: if query in self._cache: return self._cache[query] - + # Don't expand very short queries or obvious keywords if len(query.split()) <= 1 or len(query) <= 3: return query - + try: expanded = self._llm_expand_query(query) if expanded and expanded != query: @@ -91,23 +94,23 @@ class QueryExpander: self._manage_cache_size() logger.info(f"Expanded query: '{query}' โ†’ '{expanded}'") return expanded - + except Exception as e: logger.warning(f"Query expansion failed: {e}") - + # Return original query if expansion fails return query - + def _llm_expand_query(self, query: str) -> Optional[str]: """Use LLM to expand the query with related terms.""" - + # Use best available model model_to_use = self._select_expansion_model() if not model_to_use: return None - + # Create expansion prompt - prompt = f"""You are a search query expert. Expand the following search query with {self.max_terms} additional related terms that would help find relevant content. + prompt = """You are a search query expert. Expand the following search query with {self.max_terms} additional related terms that would help find relevant content. Original query: "{query}" @@ -134,95 +137,99 @@ Expanded query:""" "options": { "temperature": 0.1, # Very low temperature for consistent expansions "top_p": 0.8, - "max_tokens": 100 # Keep it short - } + "max_tokens": 100, # Keep it short + }, } - + response = requests.post( f"{self.ollama_url}/api/generate", json=payload, - timeout=10 # Quick timeout for low latency + timeout=10, # Quick timeout for low latency ) - + if response.status_code == 200: - result = response.json().get('response', '').strip() - + result = response.json().get("response", "").strip() + # Clean up the response - extract just the expanded query expanded = self._clean_expansion(result, query) return expanded - + except Exception as e: logger.warning(f"LLM expansion failed: {e}") return None - + def _select_expansion_model(self) -> Optional[str]: """Select the best available model for query expansion.""" - + if self.model != "auto": return self.model - + try: # Get available models response = requests.get(f"{self.ollama_url}/api/tags", timeout=5) if response.status_code == 200: data = response.json() - available = [model['name'] for model in data.get('models', [])] - + available = [model["name"] for model in data.get("models", [])] + # Use same model rankings as main synthesizer for consistency expansion_preferences = [ - "qwen3:1.7b", "qwen3:0.6b", "qwen3:4b", "qwen2.5:3b", - "qwen2.5:1.5b", "qwen2.5-coder:1.5b" + "qwen3:1.7b", + "qwen3:0.6b", + "qwen3:4b", + "qwen2.5:3b", + "qwen2.5:1.5b", + "qwen2.5-coder:1.5b", ] - + for preferred in expansion_preferences: for available_model in available: if preferred in available_model: logger.debug(f"Using {available_model} for query expansion") return available_model - + # Fallback to first available model if available: return available[0] - + except Exception as e: logger.warning(f"Could not select expansion model: {e}") - + return None - + def _clean_expansion(self, raw_response: str, original_query: str) -> str: """Clean the LLM response to extract just the expanded query.""" - + # Remove common response artifacts clean_response = raw_response.strip() - + # Remove quotes if the entire response is quoted if clean_response.startswith('"') and clean_response.endswith('"'): clean_response = clean_response[1:-1] - + # Take only the first line if multiline - clean_response = clean_response.split('\n')[0].strip() - + clean_response = clean_response.split("\n")[0].strip() + # Remove excessive punctuation and normalize spaces - clean_response = re.sub(r'[^\w\s-]', ' ', clean_response) - clean_response = re.sub(r'\s+', ' ', clean_response).strip() - + clean_response = re.sub(r"[^\w\s-]", " ", clean_response) + clean_response = re.sub(r"\s+", " ", clean_response).strip() + # Ensure it starts with the original query if not clean_response.lower().startswith(original_query.lower()): clean_response = f"{original_query} {clean_response}" - + # Limit the total length to avoid very long queries words = clean_response.split() if len(words) > len(original_query.split()) + self.max_terms: - words = words[:len(original_query.split()) + self.max_terms] - clean_response = ' '.join(words) - + words = words[: len(original_query.split()) + self.max_terms] + clean_response = " ".join(words) + return clean_response - + def clear_cache(self): """Clear the expansion cache (thread-safe).""" with self._cache_lock: self._cache.clear() - + def _manage_cache_size(self, max_size: int = 1000): """Keep cache from growing too large (prevents memory leaks).""" with self._cache_lock: @@ -232,45 +239,49 @@ Expanded query:""" keep_count = max_size // 2 self._cache = dict(items[-keep_count:]) logger.debug(f"Cache trimmed from {len(items)} to {len(self._cache)} entries") - + def is_available(self) -> bool: """Check if query expansion is available.""" if not self.enabled: return False - + self._ensure_initialized() try: response = requests.get(f"{self.ollama_url}/api/tags", timeout=5) return response.status_code == 200 - except: + except (ConnectionError, TimeoutError, requests.RequestException): return False + # Quick test function + + def test_expansion(): """Test the query expander.""" from .config import RAGConfig - + config = RAGConfig() config.search.expand_queries = True config.llm.max_expansion_terms = 6 - + expander = QueryExpander(config) - + if not expander.is_available(): print("โŒ Ollama not available for testing") return - + test_queries = [ "authentication", - "error handling", + "error handling", "database query", - "user interface" + "user interface", ] - + print("๐Ÿ” Testing Query Expansion:") for query in test_queries: expanded = expander.expand_query(query) print(f" '{query}' โ†’ '{expanded}'") + if __name__ == "__main__": - test_expansion() \ No newline at end of file + test_expansion() diff --git a/mini_rag/search.py b/mini_rag/search.py index 1823fab..8417a27 100644 --- a/mini_rag/search.py +++ b/mini_rag/search.py @@ -4,29 +4,33 @@ Optimized for code search with relevance scoring. """ import logging +from collections import defaultdict +from datetime import datetime from pathlib import Path -from typing import List, Dict, Any, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple + import numpy as np import pandas as pd -from rich.console import Console -from rich.table import Table -from rich.syntax import Syntax from rank_bm25 import BM25Okapi -from collections import defaultdict +from rich.console import Console +from rich.syntax import Syntax +from rich.table import Table # Optional LanceDB import try: import lancedb + LANCEDB_AVAILABLE = True except ImportError: lancedb = None LANCEDB_AVAILABLE = False +from datetime import timedelta + +from .config import ConfigManager from .ollama_embeddings import OllamaEmbedder as CodeEmbedder from .path_handler import display_path from .query_expander import QueryExpander -from .config import ConfigManager -from datetime import datetime, timedelta logger = logging.getLogger(__name__) console = Console() @@ -34,19 +38,21 @@ console = Console() class SearchResult: """Represents a single search result.""" - - def __init__(self, - file_path: str, - content: str, - score: float, - start_line: int, - end_line: int, - chunk_type: str, - name: str, - language: str, - context_before: Optional[str] = None, - context_after: Optional[str] = None, - parent_chunk: Optional['SearchResult'] = None): + + def __init__( + self, + file_path: str, + content: str, + score: float, + start_line: int, + end_line: int, + chunk_type: str, + name: str, + language: str, + context_before: Optional[str] = None, + context_after: Optional[str] = None, + parent_chunk: Optional["SearchResult"] = None, + ): self.file_path = file_path self.content = content self.score = score @@ -58,59 +64,57 @@ class SearchResult: self.context_before = context_before self.context_after = context_after self.parent_chunk = parent_chunk - + def __repr__(self): return f"SearchResult({self.file_path}:{self.start_line}-{self.end_line}, score={self.score:.3f})" - + def to_dict(self) -> Dict[str, Any]: """Convert to dictionary.""" return { - 'file_path': self.file_path, - 'content': self.content, - 'score': self.score, - 'start_line': self.start_line, - 'end_line': self.end_line, - 'chunk_type': self.chunk_type, - 'name': self.name, - 'language': self.language, - 'context_before': self.context_before, - 'context_after': self.context_after, - 'parent_chunk': self.parent_chunk.to_dict() if self.parent_chunk else None, + "file_path": self.file_path, + "content": self.content, + "score": self.score, + "start_line": self.start_line, + "end_line": self.end_line, + "chunk_type": self.chunk_type, + "name": self.name, + "language": self.language, + "context_before": self.context_before, + "context_after": self.context_after, + "parent_chunk": self.parent_chunk.to_dict() if self.parent_chunk else None, } - + def format_for_display(self, max_lines: int = 10) -> str: """Format for display with syntax highlighting.""" lines = self.content.splitlines() if len(lines) > max_lines: # Show first and last few lines half = max_lines // 2 - lines = lines[:half] + ['...'] + lines[-half:] - - return '\n'.join(lines) + lines = lines[:half] + ["..."] + lines[-half:] + + return "\n".join(lines) class CodeSearcher: """Semantic code search using vector similarity.""" - - def __init__(self, - project_path: Path, - embedder: Optional[CodeEmbedder] = None): + + def __init__(self, project_path: Path, embedder: Optional[CodeEmbedder] = None): """ Initialize searcher. - + Args: project_path: Path to the project embedder: CodeEmbedder instance (creates one if not provided) """ self.project_path = Path(project_path).resolve() - self.rag_dir = self.project_path / '.mini-rag' + self.rag_dir = self.project_path / ".mini-rag" self.embedder = embedder or CodeEmbedder() - + # Load configuration and initialize query expander config_manager = ConfigManager(project_path) self.config = config_manager.load_config() self.query_expander = QueryExpander(self.config) - + # Initialize database connection self.db = None self.table = None @@ -119,7 +123,7 @@ class CodeSearcher: self.chunk_ids = [] self._connect() self._build_bm25_index() - + def _connect(self): """Connect to the LanceDB database.""" if not LANCEDB_AVAILABLE: @@ -128,8 +132,10 @@ class CodeSearcher: print(" Install it with: pip install lancedb pyarrow") print(" For basic Ollama functionality, use hash-based search instead") print() - raise ImportError("LanceDB dependency is required for search. Install with: pip install lancedb pyarrow") - + raise ImportError( + "LanceDB dependency is required for search. Install with: pip install lancedb pyarrow" + ) + try: if not self.rag_dir.exists(): print("๐Ÿ—ƒ๏ธ No Search Index Found") @@ -138,149 +144,163 @@ class CodeSearcher: print(" (This analyzes your files and creates semantic search vectors)") print() raise FileNotFoundError(f"No RAG index found at {self.rag_dir}") - + self.db = lancedb.connect(self.rag_dir) - + if "code_vectors" not in self.db.table_names(): - print("๐Ÿ”ง Index Database Corrupted") + print("๐Ÿ”ง Index Database Corrupted") print(" The search index exists but is missing data tables") - print(f" Rebuild index: rm -rf {self.rag_dir} && ./rag-mini index {self.project_path}") + print( + f" Rebuild index: rm -rf {self.rag_dir} && ./rag-mini index {self.project_path}" + ) print(" (This will recreate the search database)") print() raise ValueError("No code_vectors table found. Run indexing first.") - + self.table = self.db.open_table("code_vectors") - + except Exception as e: logger.error(f"Failed to connect to database: {e}") raise - + def _build_bm25_index(self): """Build BM25 index from all chunks in the database.""" if not self.table: return - + try: # Load all chunks into memory for BM25 df = self.table.to_pandas() - + # Prepare texts for BM25 by combining content with metadata self.chunk_texts = [] self.chunk_ids = [] - + for idx, row in df.iterrows(): # Create searchable text combining content, name, and type searchable_text = f"{row['content']} {row['name'] or ''} {row['chunk_type']}" - + # Tokenize for BM25 (simple word splitting) tokens = searchable_text.lower().split() - + self.chunk_texts.append(tokens) self.chunk_ids.append(idx) - + # Build BM25 index self.bm25 = BM25Okapi(self.chunk_texts) logger.info(f"Built BM25 index with {len(self.chunk_texts)} chunks") - + except Exception as e: logger.error(f"Failed to build BM25 index: {e}") self.bm25 = None - - def get_chunk_context(self, chunk_id: str, include_adjacent: bool = True, include_parent: bool = True) -> Dict[str, Any]: + + def get_chunk_context( + self, chunk_id: str, include_adjacent: bool = True, include_parent: bool = True + ) -> Dict[str, Any]: """ Get context for a specific chunk including adjacent and parent chunks. - + Args: chunk_id: The ID of the chunk to get context for include_adjacent: Whether to include previous and next chunks include_parent: Whether to include parent class chunk for methods - + Returns: Dictionary with 'chunk', 'prev', 'next', and 'parent' SearchResults """ if not self.table: raise RuntimeError("Database not connected") - + try: # Get the main chunk by ID df = self.table.to_pandas() - chunk_rows = df[df['chunk_id'] == chunk_id] - + chunk_rows = df[df["chunk_id"] == chunk_id] + if chunk_rows.empty: - return {'chunk': None, 'prev': None, 'next': None, 'parent': None} - + return {"chunk": None, "prev": None, "next": None, "parent": None} + chunk_row = chunk_rows.iloc[0] - context = {'chunk': self._row_to_search_result(chunk_row, score=1.0)} - + context = {"chunk": self._row_to_search_result(chunk_row, score=1.0)} + # Get adjacent chunks if requested if include_adjacent: # Get previous chunk - if pd.notna(chunk_row.get('prev_chunk_id')): - prev_rows = df[df['chunk_id'] == chunk_row['prev_chunk_id']] + if pd.notna(chunk_row.get("prev_chunk_id")): + prev_rows = df[df["chunk_id"] == chunk_row["prev_chunk_id"]] if not prev_rows.empty: - context['prev'] = self._row_to_search_result(prev_rows.iloc[0], score=1.0) + context["prev"] = self._row_to_search_result( + prev_rows.iloc[0], score=1.0 + ) else: - context['prev'] = None + context["prev"] = None else: - context['prev'] = None - + context["prev"] = None + # Get next chunk - if pd.notna(chunk_row.get('next_chunk_id')): - next_rows = df[df['chunk_id'] == chunk_row['next_chunk_id']] + if pd.notna(chunk_row.get("next_chunk_id")): + next_rows = df[df["chunk_id"] == chunk_row["next_chunk_id"]] if not next_rows.empty: - context['next'] = self._row_to_search_result(next_rows.iloc[0], score=1.0) + context["next"] = self._row_to_search_result( + next_rows.iloc[0], score=1.0 + ) else: - context['next'] = None + context["next"] = None else: - context['next'] = None + context["next"] = None else: - context['prev'] = None - context['next'] = None - + context["prev"] = None + context["next"] = None + # Get parent class chunk if requested and applicable - if include_parent and pd.notna(chunk_row.get('parent_class')): + if include_parent and pd.notna(chunk_row.get("parent_class")): # Find the parent class chunk - parent_rows = df[(df['name'] == chunk_row['parent_class']) & - (df['chunk_type'] == 'class') & - (df['file_path'] == chunk_row['file_path'])] + parent_rows = df[ + (df["name"] == chunk_row["parent_class"]) + & (df["chunk_type"] == "class") + & (df["file_path"] == chunk_row["file_path"]) + ] if not parent_rows.empty: - context['parent'] = self._row_to_search_result(parent_rows.iloc[0], score=1.0) + context["parent"] = self._row_to_search_result( + parent_rows.iloc[0], score=1.0 + ) else: - context['parent'] = None + context["parent"] = None else: - context['parent'] = None - + context["parent"] = None + return context - + except Exception as e: logger.error(f"Failed to get chunk context: {e}") - return {'chunk': None, 'prev': None, 'next': None, 'parent': None} - + return {"chunk": None, "prev": None, "next": None, "parent": None} + def _row_to_search_result(self, row: pd.Series, score: float) -> SearchResult: """Convert a DataFrame row to a SearchResult.""" return SearchResult( - file_path=display_path(row['file_path']), - content=row['content'], + file_path=display_path(row["file_path"]), + content=row["content"], score=score, - start_line=row['start_line'], - end_line=row['end_line'], - chunk_type=row['chunk_type'], - name=row['name'], - language=row['language'] + start_line=row["start_line"], + end_line=row["end_line"], + chunk_type=row["chunk_type"], + name=row["name"], + language=row["language"], ) - - def search(self, - query: str, - top_k: int = 10, - chunk_types: Optional[List[str]] = None, - languages: Optional[List[str]] = None, - file_pattern: Optional[str] = None, - semantic_weight: float = 0.7, - bm25_weight: float = 0.3, - include_context: bool = False) -> List[SearchResult]: + + def search( + self, + query: str, + top_k: int = 10, + chunk_types: Optional[List[str]] = None, + languages: Optional[List[str]] = None, + file_pattern: Optional[str] = None, + semantic_weight: float = 0.7, + bm25_weight: float = 0.3, + include_context: bool = False, + ) -> List[SearchResult]: """ Hybrid search for code similar to the query using both semantic and BM25. - + Args: query: Natural language search query top_k: Maximum number of results to return @@ -290,57 +310,56 @@ class CodeSearcher: semantic_weight: Weight for semantic similarity (default 0.7) bm25_weight: Weight for BM25 keyword score (default 0.3) include_context: Whether to include adjacent and parent chunks for each result - + Returns: List of SearchResult objects, sorted by combined relevance """ if not self.table: raise RuntimeError("Database not connected") - + # Expand query for better recall (if enabled) expanded_query = self.query_expander.expand_query(query) - + # Use original query for display but expanded query for search search_query = expanded_query if expanded_query != query else query - + # Embed the expanded query for semantic search query_embedding = self.embedder.embed_query(search_query) - + # Ensure query is a numpy array of float32 if not isinstance(query_embedding, np.ndarray): query_embedding = np.array(query_embedding, dtype=np.float32) else: query_embedding = query_embedding.astype(np.float32) - + # Get more results for hybrid scoring results_df = ( self.table.search(query_embedding) .limit(top_k * 4) # Get extra results for filtering and diversity .to_pandas() ) - + if results_df.empty: return [] - + # Apply filters first if chunk_types: - results_df = results_df[results_df['chunk_type'].isin(chunk_types)] - + results_df = results_df[results_df["chunk_type"].isin(chunk_types)] + if languages: - results_df = results_df[results_df['language'].isin(languages)] - + results_df = results_df[results_df["language"].isin(languages)] + if file_pattern: import fnmatch - mask = results_df['file_path'].apply( - lambda x: fnmatch.fnmatch(x, file_pattern) - ) + + mask = results_df["file_path"].apply(lambda x: fnmatch.fnmatch(x, file_pattern)) results_df = results_df[mask] - + # Calculate BM25 scores if available if self.bm25: # Tokenize expanded query for BM25 query_tokens = search_query.lower().split() - + # Get BM25 scores for all chunks in results bm25_scores = {} for idx, row in results_df.iterrows(): @@ -353,69 +372,79 @@ class CodeSearcher: bm25_scores[idx] = 0.0 else: bm25_scores = {idx: 0.0 for idx in results_df.index} - + # Calculate hybrid scores hybrid_results = [] for idx, row in results_df.iterrows(): # Semantic score (convert distance to similarity) - distance = row['_distance'] + distance = row["_distance"] semantic_score = 1 / (1 + distance) - + # BM25 score bm25_score = bm25_scores.get(idx, 0.0) - + # Combined score - combined_score = (semantic_weight * semantic_score + - bm25_weight * bm25_score) - + combined_score = semantic_weight * semantic_score + bm25_weight * bm25_score + result = SearchResult( - file_path=display_path(row['file_path']), - content=row['content'], + file_path=display_path(row["file_path"]), + content=row["content"], score=combined_score, - start_line=row['start_line'], - end_line=row['end_line'], - chunk_type=row['chunk_type'], - name=row['name'], - language=row['language'] + start_line=row["start_line"], + end_line=row["end_line"], + chunk_type=row["chunk_type"], + name=row["name"], + language=row["language"], ) hybrid_results.append(result) - + # Apply smart re-ranking for better quality (zero overhead) hybrid_results = self._smart_rerank(hybrid_results) - + # Apply diversity constraints diverse_results = self._apply_diversity_constraints(hybrid_results, top_k) - + # Add context if requested if include_context: diverse_results = self._add_context_to_results(diverse_results, results_df) - + return diverse_results - + def _smart_rerank(self, results: List[SearchResult]) -> List[SearchResult]: """ Smart result re-ranking for better quality with zero overhead. - + Boosts scores based on: - File importance (README, main files, configs) - Content freshness (recently modified files) - File type relevance """ now = datetime.now() - + for result in results: # File importance boost (20% boost for important files) file_path_lower = str(result.file_path).lower() important_patterns = [ - 'readme', 'main.', 'index.', '__init__', 'config', - 'setup', 'install', 'getting', 'started', 'docs/', - 'documentation', 'guide', 'tutorial', 'example' + "readme", + "main.", + "index.", + "__init__", + "config", + "setup", + "install", + "getting", + "started", + "docs/", + "documentation", + "guide", + "tutorial", + "example", ] - + if any(pattern in file_path_lower for pattern in important_patterns): result.score *= 1.2 logger.debug(f"Important file boost: {result.file_path}") - + # Recency boost (10% boost for files modified in last week) # Note: This uses file modification time if available in the data try: @@ -423,42 +452,46 @@ class CodeSearcher: file_mtime = Path(result.file_path).stat().st_mtime modified_date = datetime.fromtimestamp(file_mtime) days_old = (now - modified_date).days - + if days_old <= 7: # Modified in last week result.score *= 1.1 - logger.debug(f"Recent file boost: {result.file_path} ({days_old} days old)") + logger.debug( + f"Recent file boost: {result.file_path} ({days_old} days old)" + ) elif days_old <= 30: # Modified in last month result.score *= 1.05 - + except (OSError, ValueError): # File doesn't exist or can't get stats - no boost pass - + # Content type relevance boost - if hasattr(result, 'chunk_type'): - if result.chunk_type in ['function', 'class', 'method']: + if hasattr(result, "chunk_type"): + if result.chunk_type in ["function", "class", "method"]: # Code definitions are usually more valuable result.score *= 1.1 - elif result.chunk_type in ['comment', 'docstring']: + elif result.chunk_type in ["comment", "docstring"]: # Documentation is valuable for understanding result.score *= 1.05 - + # Penalize very short content (likely not useful) if len(result.content.strip()) < 50: result.score *= 0.9 - + # Small boost for content with good structure (has multiple lines) - lines = result.content.strip().split('\n') + lines = result.content.strip().split("\n") if len(lines) >= 3 and any(len(line.strip()) > 10 for line in lines): result.score *= 1.02 - + # Sort by updated scores return sorted(results, key=lambda x: x.score, reverse=True) - - def _apply_diversity_constraints(self, results: List[SearchResult], top_k: int) -> List[SearchResult]: + + def _apply_diversity_constraints( + self, results: List[SearchResult], top_k: int + ) -> List[SearchResult]: """ Apply diversity constraints to search results. - + - Max 2 chunks per file - Prefer different chunk types - Deduplicate overlapping content @@ -467,117 +500,121 @@ class CodeSearcher: file_counts = defaultdict(int) seen_content_hashes = set() chunk_type_counts = defaultdict(int) - + for result in results: # Check file limit if file_counts[result.file_path] >= 2: continue - + # Check for duplicate/overlapping content content_hash = hash(result.content.strip()[:200]) # Hash first 200 chars if content_hash in seen_content_hashes: continue - + # Prefer diverse chunk types - if len(final_results) >= top_k // 2 and chunk_type_counts[result.chunk_type] > top_k // 3: + if ( + len(final_results) >= top_k // 2 + and chunk_type_counts[result.chunk_type] > top_k // 3 + ): # Skip if we have too many of this type already continue - + # Add result final_results.append(result) file_counts[result.file_path] += 1 seen_content_hashes.add(content_hash) chunk_type_counts[result.chunk_type] += 1 - + if len(final_results) >= top_k: break - + return final_results - - def _add_context_to_results(self, results: List[SearchResult], search_df: pd.DataFrame) -> List[SearchResult]: + + def _add_context_to_results( + self, results: List[SearchResult], search_df: pd.DataFrame + ) -> List[SearchResult]: """ Add context (adjacent and parent chunks) to search results. - + Args: results: List of search results to add context to search_df: DataFrame from the initial search (for finding chunk_id) - + Returns: List of SearchResult objects with context added """ # Get full dataframe for context lookups full_df = self.table.to_pandas() - + # Create a mapping from result to chunk_id result_to_chunk_id = {} for result in results: # Find matching row in search_df matching_rows = search_df[ - (search_df['file_path'] == result.file_path) & - (search_df['start_line'] == result.start_line) & - (search_df['end_line'] == result.end_line) + (search_df["file_path"] == result.file_path) + & (search_df["start_line"] == result.start_line) + & (search_df["end_line"] == result.end_line) ] if not matching_rows.empty: - result_to_chunk_id[result] = matching_rows.iloc[0]['chunk_id'] - + result_to_chunk_id[result] = matching_rows.iloc[0]["chunk_id"] + # Add context to each result for result in results: chunk_id = result_to_chunk_id.get(result) if not chunk_id: continue - + # Get the row for this chunk - chunk_rows = full_df[full_df['chunk_id'] == chunk_id] + chunk_rows = full_df[full_df["chunk_id"] == chunk_id] if chunk_rows.empty: continue - + chunk_row = chunk_rows.iloc[0] - + # Add adjacent chunks as context - if pd.notna(chunk_row.get('prev_chunk_id')): - prev_rows = full_df[full_df['chunk_id'] == chunk_row['prev_chunk_id']] + if pd.notna(chunk_row.get("prev_chunk_id")): + prev_rows = full_df[full_df["chunk_id"] == chunk_row["prev_chunk_id"]] if not prev_rows.empty: - result.context_before = prev_rows.iloc[0]['content'] - - if pd.notna(chunk_row.get('next_chunk_id')): - next_rows = full_df[full_df['chunk_id'] == chunk_row['next_chunk_id']] + result.context_before = prev_rows.iloc[0]["content"] + + if pd.notna(chunk_row.get("next_chunk_id")): + next_rows = full_df[full_df["chunk_id"] == chunk_row["next_chunk_id"]] if not next_rows.empty: - result.context_after = next_rows.iloc[0]['content'] - + result.context_after = next_rows.iloc[0]["content"] + # Add parent class chunk if applicable - if pd.notna(chunk_row.get('parent_class')): + if pd.notna(chunk_row.get("parent_class")): parent_rows = full_df[ - (full_df['name'] == chunk_row['parent_class']) & - (full_df['chunk_type'] == 'class') & - (full_df['file_path'] == chunk_row['file_path']) + (full_df["name"] == chunk_row["parent_class"]) + & (full_df["chunk_type"] == "class") + & (full_df["file_path"] == chunk_row["file_path"]) ] if not parent_rows.empty: parent_row = parent_rows.iloc[0] result.parent_chunk = SearchResult( - file_path=display_path(parent_row['file_path']), - content=parent_row['content'], + file_path=display_path(parent_row["file_path"]), + content=parent_row["content"], score=1.0, - start_line=parent_row['start_line'], - end_line=parent_row['end_line'], - chunk_type=parent_row['chunk_type'], - name=parent_row['name'], - language=parent_row['language'] + start_line=parent_row["start_line"], + end_line=parent_row["end_line"], + chunk_type=parent_row["chunk_type"], + name=parent_row["name"], + language=parent_row["language"], ) - + return results - - def search_similar_code(self, - code_snippet: str, - top_k: int = 10, - exclude_self: bool = True) -> List[SearchResult]: + + def search_similar_code( + self, code_snippet: str, top_k: int = 10, exclude_self: bool = True + ) -> List[SearchResult]: """ Find code similar to a given snippet using hybrid search. - + Args: code_snippet: Code to find similar matches for top_k: Maximum number of results exclude_self: Whether to exclude exact matches - + Returns: List of similar code chunks """ @@ -587,9 +624,9 @@ class CodeSearcher: query=code_snippet, top_k=top_k * 2 if exclude_self else top_k, semantic_weight=0.8, # Higher semantic weight for code similarity - bm25_weight=0.2 + bm25_weight=0.2, ) - + if exclude_self: # Filter out exact matches filtered_results = [] @@ -599,114 +636,108 @@ class CodeSearcher: if len(filtered_results) >= top_k: break return filtered_results - + return results[:top_k] - + def get_function(self, function_name: str, top_k: int = 5) -> List[SearchResult]: """ Search for a specific function by name. - + Args: function_name: Name of the function to find top_k: Maximum number of results - + Returns: List of matching functions """ # Create a targeted query query = f"function {function_name} implementation definition" - + # Search with filters - results = self.search( - query, - top_k=top_k * 2, - chunk_types=['function', 'method'] - ) - + results = self.search(query, top_k=top_k * 2, chunk_types=["function", "method"]) + # Further filter by name filtered = [] for result in results: if result.name and function_name.lower() in result.name.lower(): filtered.append(result) - + return filtered[:top_k] - + def get_class(self, class_name: str, top_k: int = 5) -> List[SearchResult]: """ Search for a specific class by name. - + Args: class_name: Name of the class to find top_k: Maximum number of results - + Returns: List of matching classes """ # Create a targeted query query = f"class {class_name} definition implementation" - + # Search with filters - results = self.search( - query, - top_k=top_k * 2, - chunk_types=['class'] - ) - + results = self.search(query, top_k=top_k * 2, chunk_types=["class"]) + # Further filter by name filtered = [] for result in results: if result.name and class_name.lower() in result.name.lower(): filtered.append(result) - + return filtered[:top_k] - + def explain_code(self, query: str, top_k: int = 5) -> List[SearchResult]: """ Find code that helps explain a concept. - + Args: query: Concept to explain (e.g., "how to connect to database") top_k: Maximum number of examples - + Returns: List of relevant code examples """ # Enhance query for explanation enhanced_query = f"example implementation {query}" - + return self.search(enhanced_query, top_k=top_k) - + def find_usage(self, identifier: str, top_k: int = 10) -> List[SearchResult]: """ Find usage examples of an identifier (function, class, variable). - + Args: identifier: The identifier to find usage for top_k: Maximum number of results - + Returns: List of usage examples """ # Search for usage patterns query = f"using {identifier} calling {identifier} import {identifier}" - + results = self.search(query, top_k=top_k * 2) - + # Filter to ensure identifier appears in content filtered = [] for result in results: if identifier in result.content: filtered.append(result) - + return filtered[:top_k] - - def display_results(self, - results: List[SearchResult], - show_content: bool = True, - max_content_lines: int = 10): + + def display_results( + self, + results: List[SearchResult], + show_content: bool = True, + max_content_lines: int = 10, + ): """ Display search results in a formatted table. - + Args: results: List of search results show_content: Whether to show code content @@ -715,7 +746,7 @@ class CodeSearcher: if not results: console.print("[yellow]No results found[/yellow]") return - + # Create table table = Table(title=f"Search Results ({len(results)} matches)") table.add_column("Score", style="cyan", width=6) @@ -723,81 +754,85 @@ class CodeSearcher: table.add_column("Type", style="green", width=10) table.add_column("Name", style="magenta") table.add_column("Lines", style="yellow", width=10) - + for result in results: table.add_row( f"{result.score:.3f}", result.file_path, result.chunk_type, result.name or "-", - f"{result.start_line}-{result.end_line}" + f"{result.start_line}-{result.end_line}", ) - + console.print(table) - + # Show content if requested if show_content and results: console.print("\n[bold]Top Results:[/bold]\n") - + for i, result in enumerate(results[:3], 1): - console.print(f"[bold cyan]#{i}[/bold cyan] {result.file_path}:{result.start_line}") + console.print( + f"[bold cyan]#{i}[/bold cyan] {result.file_path}:{result.start_line}" + ) console.print(f"[dim]Type: {result.chunk_type} | Name: {result.name}[/dim]") - + # Display code with syntax highlighting syntax = Syntax( result.format_for_display(max_content_lines), result.language, theme="monokai", line_numbers=True, - start_line=result.start_line + start_line=result.start_line, ) console.print(syntax) console.print() - + def get_statistics(self) -> Dict[str, Any]: """Get search index statistics.""" if not self.table: - return {'error': 'Database not connected'} - + return {"error": "Database not connected"} + try: # Get table statistics num_rows = len(self.table.to_pandas()) - + # Get unique files df = self.table.to_pandas() - unique_files = df['file_path'].nunique() - + unique_files = df["file_path"].nunique() + # Get chunk type distribution - chunk_types = df['chunk_type'].value_counts().to_dict() - + chunk_types = df["chunk_type"].value_counts().to_dict() + # Get language distribution - languages = df['language'].value_counts().to_dict() - + languages = df["language"].value_counts().to_dict() + return { - 'total_chunks': num_rows, - 'unique_files': unique_files, - 'chunk_types': chunk_types, - 'languages': languages, - 'index_ready': True, + "total_chunks": num_rows, + "unique_files": unique_files, + "chunk_types": chunk_types, + "languages": languages, + "index_ready": True, } - + except Exception as e: logger.error(f"Failed to get statistics: {e}") - return {'error': str(e)} + return {"error": str(e)} # Convenience functions + + def search_code(project_path: Path, query: str, top_k: int = 10) -> List[SearchResult]: """ Quick search function. - + Args: project_path: Path to the project query: Search query top_k: Maximum results - + Returns: List of search results """ searcher = CodeSearcher(project_path) - return searcher.search(query, top_k=top_k) \ No newline at end of file + return searcher.search(query, top_k=top_k) diff --git a/mini_rag/server.py b/mini_rag/server.py index c06aa8d..5367d05 100644 --- a/mini_rag/server.py +++ b/mini_rag/server.py @@ -4,30 +4,30 @@ No more loading/unloading madness! """ import json +import logging +import os import socket +import subprocess +import sys import threading import time -import subprocess from pathlib import Path -from typing import Dict, Any, Optional -import logging -import sys -import os +from typing import Any, Dict, Optional # Fix Windows console -if sys.platform == 'win32': - os.environ['PYTHONUTF8'] = '1' +if sys.platform == "win32": + os.environ["PYTHONUTF8"] = "1" -from .search import CodeSearcher from .ollama_embeddings import OllamaEmbedder as CodeEmbedder from .performance import PerformanceMonitor +from .search import CodeSearcher logger = logging.getLogger(__name__) class RAGServer: """Persistent server that keeps embeddings and DB loaded.""" - + def __init__(self, project_path: Path, port: int = 7777): self.project_path = project_path self.port = port @@ -37,37 +37,36 @@ class RAGServer: self.socket = None self.start_time = None self.query_count = 0 - + def _kill_existing_server(self): """Kill any existing process using our port.""" try: # Check if port is in use test_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - result = test_sock.connect_ex(('localhost', self.port)) + result = test_sock.connect_ex(("localhost", self.port)) test_sock.close() - + if result == 0: # Port is in use print(f"๏ธ Port {self.port} is already in use, attempting to free it...") - - if sys.platform == 'win32': + + if sys.platform == "win32": # Windows: Find and kill process using netstat import subprocess + try: # Get process ID using the port result = subprocess.run( - ['netstat', '-ano'], - capture_output=True, - text=True + ["netstat", "-ano"], capture_output=True, text=True ) - - for line in result.stdout.split('\n'): - if f':{self.port}' in line and 'LISTENING' in line: + + for line in result.stdout.split("\n"): + if f":{self.port}" in line and "LISTENING" in line: parts = line.split() pid = parts[-1] print(f" Found process {pid} using port {self.port}") - + # Kill the process - subprocess.run(['taskkill', '//PID', pid, '//F'], check=False) + subprocess.run(["taskkill", "//PID", pid, "//F"], check=False) print(f" Killed process {pid}") time.sleep(1) # Give it a moment to release the port break @@ -76,15 +75,16 @@ class RAGServer: else: # Unix/Linux: Use lsof and kill import subprocess + try: result = subprocess.run( - ['lsof', '-ti', f':{self.port}'], - capture_output=True, - text=True + ["lso", "-ti", f":{self.port}"], + capture_output=True, + text=True, ) if result.stdout.strip(): pid = result.stdout.strip() - subprocess.run(['kill', '-9', pid], check=False) + subprocess.run(["kill", "-9", pid], check=False) print(f" Killed process {pid}") time.sleep(1) except Exception as e: @@ -92,38 +92,38 @@ class RAGServer: except Exception as e: # Non-critical error, just log it logger.debug(f"Error checking port: {e}") - + def start(self): """Start the RAG server.""" # Kill any existing process on our port first self._kill_existing_server() - + print(f" Starting RAG server on port {self.port}...") - + # Load everything once perf = PerformanceMonitor() - + with perf.measure("Load Embedder"): self.embedder = CodeEmbedder() - + with perf.measure("Connect Database"): self.searcher = CodeSearcher(self.project_path, embedder=self.embedder) - + perf.print_summary() - + # Start server self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.socket.bind(('localhost', self.port)) + self.socket.bind(("localhost", self.port)) self.socket.listen(5) - + self.running = True self.start_time = time.time() - + print(f"\n RAG server ready on localhost:{self.port}") print(" Model loaded, database connected") print(" Waiting for queries...\n") - + # Handle connections while self.running: try: @@ -136,50 +136,50 @@ class RAGServer: except Exception as e: if self.running: logger.error(f"Server error: {e}") - + def _handle_client(self, client: socket.socket): """Handle a client connection.""" try: # Receive query with proper message framing data = self._receive_json(client) request = json.loads(data) - + # Check for shutdown command - if request.get('command') == 'shutdown': + if request.get("command") == "shutdown": print("\n Shutdown requested") - response = {'success': True, 'message': 'Server shutting down'} + response = {"success": True, "message": "Server shutting down"} self._send_json(client, response) self.stop() return - - query = request.get('query', '') - top_k = request.get('top_k', 10) - + + query = request.get("query", "") + top_k = request.get("top_k", 10) + self.query_count += 1 print(f"[Query #{self.query_count}] {query}") - + # Perform search start = time.time() results = self.searcher.search(query, top_k=top_k) search_time = time.time() - start - + # Prepare response response = { - 'success': True, - 'query': query, - 'count': len(results), - 'search_time_ms': int(search_time * 1000), - 'results': [r.to_dict() for r in results], - 'server_uptime': int(time.time() - self.start_time), - 'total_queries': self.query_count, + "success": True, + "query": query, + "count": len(results), + "search_time_ms": int(search_time * 1000), + "results": [r.to_dict() for r in results], + "server_uptime": int(time.time() - self.start_time), + "total_queries": self.query_count, } - + # Send response with proper framing self._send_json(client, response) - + print(f" Found {len(results)} results in {search_time*1000:.0f}ms") - - except ConnectionError as e: + + except ConnectionError: # Normal disconnection - client closed connection # This is expected behavior, don't log as error pass @@ -187,52 +187,49 @@ class RAGServer: # Only log actual errors, not normal disconnections if "Connection closed" not in str(e): logger.error(f"Client handler error: {e}") - error_response = { - 'success': False, - 'error': str(e) - } + error_response = {"success": False, "error": str(e)} try: self._send_json(client, error_response) - except: + except (ConnectionError, OSError, TypeError, ValueError, socket.error): pass finally: client.close() - + def _receive_json(self, sock: socket.socket) -> str: """Receive a complete JSON message with length prefix.""" # First receive the length (4 bytes) - length_data = b'' + length_data = b"" while len(length_data) < 4: chunk = sock.recv(4 - len(length_data)) if not chunk: raise ConnectionError("Connection closed while receiving length") length_data += chunk - - length = int.from_bytes(length_data, 'big') - + + length = int.from_bytes(length_data, "big") + # Now receive the actual data - data = b'' + data = b"" while len(data) < length: chunk = sock.recv(min(65536, length - len(data))) if not chunk: raise ConnectionError("Connection closed while receiving data") data += chunk - - return data.decode('utf-8') - + + return data.decode("utf-8") + def _send_json(self, sock: socket.socket, data: dict): """Send a JSON message with length prefix.""" # Sanitize the data to ensure JSON compatibility - json_str = json.dumps(data, ensure_ascii=False, separators=(',', ':')) - json_bytes = json_str.encode('utf-8') - + json_str = json.dumps(data, ensure_ascii=False, separators=(",", ":")) + json_bytes = json_str.encode("utf-8") + # Send length prefix (4 bytes) length = len(json_bytes) - sock.send(length.to_bytes(4, 'big')) - + sock.send(length.to_bytes(4, "big")) + # Send the data sock.sendall(json_bytes) - + def stop(self): """Stop the server.""" self.running = False @@ -243,101 +240,89 @@ class RAGServer: class RAGClient: """Client to communicate with RAG server.""" - + def __init__(self, port: int = 7777): self.port = port self.use_legacy = False - + def search(self, query: str, top_k: int = 10) -> Dict[str, Any]: """Send search query to server.""" try: # Connect to server sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect(('localhost', self.port)) - + sock.connect(("localhost", self.port)) + # Send request with proper framing - request = { - 'query': query, - 'top_k': top_k - } + request = {"query": query, "top_k": top_k} self._send_json(sock, request) - + # Receive response with proper framing data = self._receive_json(sock) response = json.loads(data) - + sock.close() return response - + except ConnectionRefusedError: return { - 'success': False, - 'error': 'RAG server not running. Start with: rag-mini server' + "success": False, + "error": "RAG server not running. Start with: rag-mini server", } except ConnectionError as e: # Try legacy mode without message framing if not self.use_legacy and "receiving length" in str(e): self.use_legacy = True return self._search_legacy(query, top_k) - return { - 'success': False, - 'error': str(e) - } + return {"success": False, "error": str(e)} except Exception as e: - return { - 'success': False, - 'error': str(e) - } - + return {"success": False, "error": str(e)} + def _receive_json(self, sock: socket.socket) -> str: """Receive a complete JSON message with length prefix.""" # First receive the length (4 bytes) - length_data = b'' + length_data = b"" while len(length_data) < 4: chunk = sock.recv(4 - len(length_data)) if not chunk: raise ConnectionError("Connection closed while receiving length") length_data += chunk - - length = int.from_bytes(length_data, 'big') - + + length = int.from_bytes(length_data, "big") + # Now receive the actual data - data = b'' + data = b"" while len(data) < length: chunk = sock.recv(min(65536, length - len(data))) if not chunk: raise ConnectionError("Connection closed while receiving data") data += chunk - - return data.decode('utf-8') - + + return data.decode("utf-8") + def _send_json(self, sock: socket.socket, data: dict): """Send a JSON message with length prefix.""" - json_str = json.dumps(data, ensure_ascii=False, separators=(',', ':')) - json_bytes = json_str.encode('utf-8') - + json_str = json.dumps(data, ensure_ascii=False, separators=(",", ":")) + json_bytes = json_str.encode("utf-8") + # Send length prefix (4 bytes) length = len(json_bytes) - sock.send(length.to_bytes(4, 'big')) - + sock.send(length.to_bytes(4, "big")) + # Send the data sock.sendall(json_bytes) - + def _search_legacy(self, query: str, top_k: int = 10) -> Dict[str, Any]: """Legacy search without message framing for old servers.""" try: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect(('localhost', self.port)) - + sock.connect(("localhost", self.port)) + # Send request (old way) - request = { - 'query': query, - 'top_k': top_k - } - sock.send(json.dumps(request).encode('utf-8')) - + request = {"query": query, "top_k": top_k} + sock.send(json.dumps(request).encode("utf-8")) + # Receive response (accumulate until we get valid JSON) - data = b'' + data = b"" while True: chunk = sock.recv(65536) if not chunk: @@ -345,32 +330,26 @@ class RAGClient: data += chunk try: # Try to decode as JSON - response = json.loads(data.decode('utf-8')) + response = json.loads(data.decode("utf-8")) sock.close() return response except json.JSONDecodeError: # Keep receiving continue - + sock.close() - return { - 'success': False, - 'error': 'Incomplete response from server' - } + return {"success": False, "error": "Incomplete response from server"} except Exception as e: - return { - 'success': False, - 'error': str(e) - } - + return {"success": False, "error": str(e)} + def is_running(self) -> bool: """Check if server is running.""" try: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - result = sock.connect_ex(('localhost', self.port)) + result = sock.connect_ex(("localhost", self.port)) sock.close() return result == 0 - except: + except (ConnectionError, OSError, TypeError, ValueError, socket.error): return False @@ -389,23 +368,31 @@ def auto_start_if_needed(project_path: Path) -> Optional[subprocess.Popen]: if not client.is_running(): # Start server in background import subprocess - cmd = [sys.executable, "-m", "mini_rag.cli", "server", "--path", str(project_path)] + + cmd = [ + sys.executable, + "-m", + "mini_rag.cli", + "server", + "--path", + str(project_path), + ] process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - creationflags=subprocess.CREATE_NEW_CONSOLE if sys.platform == 'win32' else 0 + creationflags=(subprocess.CREATE_NEW_CONSOLE if sys.platform == "win32" else 0), ) - + # Wait for server to start for _ in range(30): # 30 second timeout time.sleep(1) if client.is_running(): print(" RAG server started automatically") return process - + # Failed to start process.terminate() raise RuntimeError("Failed to start RAG server") - - return None \ No newline at end of file + + return None diff --git a/mini_rag/smart_chunking.py b/mini_rag/smart_chunking.py index 9d2289c..37609ea 100644 --- a/mini_rag/smart_chunking.py +++ b/mini_rag/smart_chunking.py @@ -3,148 +3,140 @@ Smart language-aware chunking strategies for FSS-Mini-RAG. Automatically adapts chunking based on file type and content patterns. """ -from typing import Dict, Any, List from pathlib import Path -import json +from typing import Any, Dict, List + class SmartChunkingStrategy: """Intelligent chunking that adapts to file types and content.""" - + def __init__(self): self.language_configs = { - 'python': { - 'max_size': 3000, # Larger for better function context - 'min_size': 200, - 'strategy': 'function', - 'prefer_semantic': True + "python": { + "max_size": 3000, # Larger for better function context + "min_size": 200, + "strategy": "function", + "prefer_semantic": True, }, - 'javascript': { - 'max_size': 2500, - 'min_size': 150, - 'strategy': 'function', - 'prefer_semantic': True + "javascript": { + "max_size": 2500, + "min_size": 150, + "strategy": "function", + "prefer_semantic": True, }, - 'markdown': { - 'max_size': 2500, - 'min_size': 300, # Larger minimum for complete thoughts - 'strategy': 'header', - 'preserve_structure': True + "markdown": { + "max_size": 2500, + "min_size": 300, # Larger minimum for complete thoughts + "strategy": "header", + "preserve_structure": True, }, - 'json': { - 'max_size': 1000, # Smaller for config files - 'min_size': 50, - 'skip_if_large': True, # Skip huge config JSONs - 'max_file_size': 50000 # 50KB limit + "json": { + "max_size": 1000, # Smaller for config files + "min_size": 50, + "skip_if_large": True, # Skip huge config JSONs + "max_file_size": 50000, # 50KB limit }, - 'yaml': { - 'max_size': 1500, - 'min_size': 100, - 'strategy': 'key_block' - }, - 'text': { - 'max_size': 2000, - 'min_size': 200, - 'strategy': 'paragraph' - }, - 'bash': { - 'max_size': 1500, - 'min_size': 100, - 'strategy': 'function' - } + "yaml": {"max_size": 1500, "min_size": 100, "strategy": "key_block"}, + "text": {"max_size": 2000, "min_size": 200, "strategy": "paragraph"}, + "bash": {"max_size": 1500, "min_size": 100, "strategy": "function"}, } - + # Smart defaults for unknown languages self.default_config = { - 'max_size': 2000, - 'min_size': 150, - 'strategy': 'semantic' + "max_size": 2000, + "min_size": 150, + "strategy": "semantic", } - + def get_config_for_language(self, language: str, file_size: int = 0) -> Dict[str, Any]: """Get optimal chunking config for a specific language.""" config = self.language_configs.get(language, self.default_config).copy() - + # Smart adjustments based on file size if file_size > 0: if file_size < 500: # Very small files - config['max_size'] = max(config['max_size'] // 2, 200) - config['min_size'] = 50 - elif file_size > 20000: # Large files - config['max_size'] = min(config['max_size'] + 1000, 4000) - + config["max_size"] = max(config["max_size"] // 2, 200) + config["min_size"] = 50 + elif file_size > 20000: # Large files + config["max_size"] = min(config["max_size"] + 1000, 4000) + return config - + def should_skip_file(self, language: str, file_size: int) -> bool: """Determine if a file should be skipped entirely.""" lang_config = self.language_configs.get(language, {}) - + # Skip huge JSON config files - if language == 'json' and lang_config.get('skip_if_large'): - max_size = lang_config.get('max_file_size', 50000) + if language == "json" and lang_config.get("skip_if_large"): + max_size = lang_config.get("max_file_size", 50000) if file_size > max_size: return True - + # Skip tiny files that won't provide good context if file_size < 30: return True - + return False - + def get_smart_defaults(self, project_stats: Dict[str, Any]) -> Dict[str, Any]: """Generate smart defaults based on project language distribution.""" - languages = project_stats.get('languages', {}) - total_files = sum(languages.values()) - + languages = project_stats.get("languages", {}) + # sum(languages.values()) # Unused variable removed + # Determine primary language - primary_lang = max(languages.items(), key=lambda x: x[1])[0] if languages else 'python' + primary_lang = max(languages.items(), key=lambda x: x[1])[0] if languages else "python" primary_config = self.language_configs.get(primary_lang, self.default_config) - + # Smart streaming threshold based on large files - large_files = project_stats.get('large_files', 0) + large_files = project_stats.get("large_files", 0) streaming_threshold = 5120 if large_files > 5 else 1048576 # 5KB vs 1MB - + return { "chunking": { - "max_size": primary_config['max_size'], - "min_size": primary_config['min_size'], - "strategy": primary_config.get('strategy', 'semantic'), + "max_size": primary_config["max_size"], + "min_size": primary_config["min_size"], + "strategy": primary_config.get("strategy", "semantic"), "language_specific": { - lang: config for lang, config in self.language_configs.items() + lang: config + for lang, config in self.language_configs.items() if languages.get(lang, 0) > 0 - } + }, }, "streaming": { "enabled": True, "threshold_bytes": streaming_threshold, - "chunk_size_kb": 64 + "chunk_size_kb": 64, }, "files": { "skip_tiny_files": True, "tiny_threshold": 30, - "smart_json_filtering": True - } + "smart_json_filtering": True, + }, } + # Example usage + + def analyze_and_suggest(manifest_data: Dict[str, Any]) -> Dict[str, Any]: """Analyze project and suggest optimal configuration.""" from collections import Counter - - files = manifest_data.get('files', {}) + + files = manifest_data.get("files", {}) languages = Counter() large_files = 0 - + for info in files.values(): - lang = info.get('language', 'unknown') + lang = info.get("language", "unknown") languages[lang] += 1 - if info.get('size', 0) > 10000: + if info.get("size", 0) > 10000: large_files += 1 - + stats = { - 'languages': dict(languages), - 'large_files': large_files, - 'total_files': len(files) + "languages": dict(languages), + "large_files": large_files, + "total_files": len(files), } - + strategy = SmartChunkingStrategy() - return strategy.get_smart_defaults(stats) \ No newline at end of file + return strategy.get_smart_defaults(stats) diff --git a/mini_rag/system_context.py b/mini_rag/system_context.py index 98fe9f9..146c9d8 100644 --- a/mini_rag/system_context.py +++ b/mini_rag/system_context.py @@ -7,22 +7,21 @@ context-aware assistance without compromising privacy. import platform import sys -import os from pathlib import Path from typing import Dict, Optional class SystemContextCollector: """Collects system context information for enhanced LLM grounding.""" - + @staticmethod def get_system_context(project_path: Optional[Path] = None) -> str: """ Get concise system context for LLM grounding. - + Args: project_path: Current project directory - + Returns: Formatted system context string (max 200 chars for privacy) """ @@ -30,14 +29,12 @@ class SystemContextCollector: # Basic system info os_name = platform.system() python_ver = f"{sys.version_info.major}.{sys.version_info.minor}" - + # Simplified OS names - os_short = { - 'Windows': 'Win', - 'Linux': 'Linux', - 'Darwin': 'macOS' - }.get(os_name, os_name) - + os_short = {"Windows": "Win", "Linux": "Linux", "Darwin": "macOS"}.get( + os_name, os_name + ) + # Working directory info if project_path: # Use relative or shortened path for privacy @@ -49,55 +46,55 @@ class SystemContextCollector: path_info = project_path.name else: path_info = Path.cwd().name - + # Trim path if too long for our 200-char limit if len(path_info) > 50: path_info = f".../{path_info[-45:]}" - + # Command style hints cmd_style = "rag.bat" if os_name == "Windows" else "./rag-mini" - + # Format concise context context = f"[{os_short} {python_ver}, {path_info}, use {cmd_style}]" - + # Ensure we stay under 200 chars if len(context) > 200: context = context[:197] + "...]" - + return context - + except Exception: # Fallback to minimal info if anything fails return f"[{platform.system()}, Python {sys.version_info.major}.{sys.version_info.minor}]" - + @staticmethod def get_command_context(os_name: Optional[str] = None) -> Dict[str, str]: """ Get OS-appropriate command examples. - + Returns: Dictionary with command patterns for the current OS """ if os_name is None: os_name = platform.system() - + if os_name == "Windows": return { "launcher": "rag.bat", "index": "rag.bat index C:\\path\\to\\project", - "search": "rag.bat search C:\\path\\to\\project \"query\"", + "search": 'rag.bat search C:\\path\\to\\project "query"', "explore": "rag.bat explore C:\\path\\to\\project", "path_sep": "\\", - "example_path": "C:\\Users\\username\\Documents\\myproject" + "example_path": "C:\\Users\\username\\Documents\\myproject", } else: return { "launcher": "./rag-mini", - "index": "./rag-mini index /path/to/project", - "search": "./rag-mini search /path/to/project \"query\"", + "index": "./rag-mini index /path/to/project", + "search": './rag-mini search /path/to/project "query"', "explore": "./rag-mini explore /path/to/project", "path_sep": "/", - "example_path": "~/Documents/myproject" + "example_path": "~/Documents/myproject", } @@ -112,6 +109,7 @@ def get_command_context() -> Dict[str, str]: # Test function + if __name__ == "__main__": print("System Context Test:") print(f"Context: {get_system_context()}") @@ -120,4 +118,4 @@ if __name__ == "__main__": print("Command Context:") cmds = get_command_context() for key, value in cmds.items(): - print(f" {key}: {value}") \ No newline at end of file + print(f" {key}: {value}") diff --git a/mini_rag/updater.py b/mini_rag/updater.py index f635de3..9e1c53a 100644 --- a/mini_rag/updater.py +++ b/mini_rag/updater.py @@ -6,30 +6,32 @@ Provides seamless GitHub-based updates with user-friendly interface. Checks for new releases, downloads updates, and handles installation safely. """ -import os -import sys import json -import time +import os import shutil -import zipfile -import tempfile import subprocess -from pathlib import Path -from typing import Optional, Dict, Any, Tuple -from datetime import datetime, timedelta +import sys +import tempfile +import zipfile from dataclasses import dataclass +from datetime import datetime, timedelta +from pathlib import Path +from typing import Optional, Tuple try: import requests + REQUESTS_AVAILABLE = True except ImportError: REQUESTS_AVAILABLE = False from .config import ConfigManager + @dataclass class UpdateInfo: """Information about an available update.""" + version: str release_url: str download_url: str @@ -37,42 +39,45 @@ class UpdateInfo: published_at: str is_newer: bool + class UpdateChecker: """ Handles checking for and applying updates from GitHub releases. - + Features: - Checks GitHub API for latest releases - Downloads and applies updates safely with backup - Respects user preferences and rate limiting - Provides graceful fallbacks if network unavailable """ - - def __init__(self, - repo_owner: str = "FSSCoding", - repo_name: str = "Fss-Mini-Rag", - current_version: str = "2.1.0"): + + def __init__( + self, + repo_owner: str = "FSSCoding", + repo_name: str = "Fss-Mini-Rag", + current_version: str = "2.1.0", + ): self.repo_owner = repo_owner self.repo_name = repo_name self.current_version = current_version self.github_api_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}" self.check_frequency_hours = 24 # Check once per day - + # Paths self.app_root = Path(__file__).parent.parent self.cache_file = self.app_root / ".update_cache.json" self.backup_dir = self.app_root / ".backup" - + # User preferences (graceful fallback if config unavailable) try: self.config = ConfigManager(self.app_root) except Exception: self.config = None - + def should_check_for_updates(self) -> bool: """ Determine if we should check for updates now. - + Respects: - User preference to disable updates - Rate limiting (once per day by default) @@ -80,70 +85,74 @@ class UpdateChecker: """ if not REQUESTS_AVAILABLE: return False - + # Check user preference - if hasattr(self.config, 'updates') and not getattr(self.config.updates, 'auto_check', True): + if hasattr(self.config, "updates") and not getattr( + self.config.updates, "auto_check", True + ): return False - + # Check if we've checked recently if self.cache_file.exists(): try: - with open(self.cache_file, 'r') as f: + with open(self.cache_file, "r") as f: cache = json.load(f) - last_check = datetime.fromisoformat(cache.get('last_check', '2020-01-01')) - if datetime.now() - last_check < timedelta(hours=self.check_frequency_hours): + last_check = datetime.fromisoformat(cache.get("last_check", "2020-01-01")) + if datetime.now() - last_check < timedelta( + hours=self.check_frequency_hours + ): return False except (json.JSONDecodeError, ValueError, KeyError): pass # Ignore cache errors, will check anyway - + return True - + def check_for_updates(self) -> Optional[UpdateInfo]: """ Check GitHub API for the latest release. - + Returns: UpdateInfo if an update is available, None otherwise """ if not REQUESTS_AVAILABLE: return None - + try: # Get latest release from GitHub API response = requests.get( f"{self.github_api_url}/releases/latest", timeout=10, - headers={"Accept": "application/vnd.github.v3+json"} + headers={"Accept": "application/vnd.github.v3+json"}, ) - + if response.status_code != 200: return None - + release_data = response.json() - + # Extract version info - latest_version = release_data.get('tag_name', '').lstrip('v') - release_notes = release_data.get('body', 'No release notes available.') - published_at = release_data.get('published_at', '') - release_url = release_data.get('html_url', '') - + latest_version = release_data.get("tag_name", "").lstrip("v") + release_notes = release_data.get("body", "No release notes available.") + published_at = release_data.get("published_at", "") + release_url = release_data.get("html_url", "") + # Find download URL for source code download_url = None - for asset in release_data.get('assets', []): - if asset.get('name', '').endswith('.zip'): - download_url = asset.get('browser_download_url') + for asset in release_data.get("assets", []): + if asset.get("name", "").endswith(".zip"): + download_url = asset.get("browser_download_url") break - + # Fallback to source code zip if not download_url: download_url = f"https://github.com/{self.repo_owner}/{self.repo_name}/archive/refs/tags/v{latest_version}.zip" - + # Check if this is a newer version is_newer = self._is_version_newer(latest_version, self.current_version) - + # Update cache self._update_cache(latest_version, is_newer) - + if is_newer: return UpdateInfo( version=latest_version, @@ -151,92 +160,95 @@ class UpdateChecker: download_url=download_url, release_notes=release_notes, published_at=published_at, - is_newer=True + is_newer=True, ) - - except Exception as e: + + except Exception: # Silently fail for network issues - don't interrupt user experience pass - + return None - + def _is_version_newer(self, latest: str, current: str) -> bool: """ Compare version strings to determine if latest is newer. - + Simple semantic version comparison supporting: - Major.Minor.Patch (e.g., 2.1.0) - Major.Minor (e.g., 2.1) """ + def version_tuple(v): return tuple(map(int, (v.split(".")))) - + try: return version_tuple(latest) > version_tuple(current) except (ValueError, AttributeError): # If version parsing fails, assume it's newer to be safe return latest != current - + def _update_cache(self, latest_version: str, is_newer: bool): """Update the cache file with check results.""" cache_data = { - 'last_check': datetime.now().isoformat(), - 'latest_version': latest_version, - 'is_newer': is_newer + "last_check": datetime.now().isoformat(), + "latest_version": latest_version, + "is_newer": is_newer, } - + try: - with open(self.cache_file, 'w') as f: + with open(self.cache_file, "w") as f: json.dump(cache_data, f, indent=2) except Exception: pass # Ignore cache write errors - - def download_update(self, update_info: UpdateInfo, progress_callback=None) -> Optional[Path]: + + def download_update( + self, update_info: UpdateInfo, progress_callback=None + ) -> Optional[Path]: """ Download the update package to a temporary location. - + Args: update_info: Information about the update to download progress_callback: Optional callback for progress updates - + Returns: Path to downloaded file, or None if download failed """ if not REQUESTS_AVAILABLE: return None - + try: # Create temporary file for download - with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as tmp_file: + with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_file: tmp_path = Path(tmp_file.name) - + # Download with progress tracking response = requests.get(update_info.download_url, stream=True, timeout=30) response.raise_for_status() - - total_size = int(response.headers.get('content-length', 0)) + + total_size = int(response.headers.get("content-length", 0)) downloaded = 0 - - with open(tmp_path, 'wb') as f: + + with open(tmp_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): if chunk: f.write(chunk) downloaded += len(chunk) if progress_callback and total_size > 0: progress_callback(downloaded, total_size) - + return tmp_path - - except Exception as e: + + except Exception: # Clean up on error - if 'tmp_path' in locals() and tmp_path.exists(): + if "tmp_path" in locals() and tmp_path.exists(): tmp_path.unlink() return None - + def create_backup(self) -> bool: """ Create a backup of the current installation. - + Returns: True if backup created successfully """ @@ -244,22 +256,22 @@ class UpdateChecker: # Remove old backup if it exists if self.backup_dir.exists(): shutil.rmtree(self.backup_dir) - + # Create new backup self.backup_dir.mkdir(exist_ok=True) - + # Copy key files and directories important_items = [ - 'mini_rag', - 'rag-mini.py', - 'rag-tui.py', - 'requirements.txt', - 'install_mini_rag.sh', - 'install_windows.bat', - 'README.md', - 'assets' + "mini_rag", + "rag-mini.py", + "rag-tui.py", + "requirements.txt", + "install_mini_rag.sh", + "install_windows.bat", + "README.md", + "assets", ] - + for item in important_items: src = self.app_root / item if src.exists(): @@ -267,20 +279,20 @@ class UpdateChecker: shutil.copytree(src, self.backup_dir / item) else: shutil.copy2(src, self.backup_dir / item) - + return True - - except Exception as e: + + except Exception: return False - + def apply_update(self, update_package_path: Path, update_info: UpdateInfo) -> bool: """ Apply the downloaded update. - + Args: update_package_path: Path to the downloaded update package update_info: Information about the update - + Returns: True if update applied successfully """ @@ -288,133 +300,133 @@ class UpdateChecker: # Extract to temporary directory first with tempfile.TemporaryDirectory() as tmp_dir: tmp_path = Path(tmp_dir) - + # Extract the archive - with zipfile.ZipFile(update_package_path, 'r') as zip_ref: + with zipfile.ZipFile(update_package_path, "r") as zip_ref: zip_ref.extractall(tmp_path) - + # Find the extracted directory (may be nested) extracted_dirs = [d for d in tmp_path.iterdir() if d.is_dir()] if not extracted_dirs: return False - + source_dir = extracted_dirs[0] - + # Copy files to application directory important_items = [ - 'mini_rag', - 'rag-mini.py', - 'rag-tui.py', - 'requirements.txt', - 'install_mini_rag.sh', - 'install_windows.bat', - 'README.md' + "mini_rag", + "rag-mini.py", + "rag-tui.py", + "requirements.txt", + "install_mini_rag.sh", + "install_windows.bat", + "README.md", ] - + for item in important_items: src = source_dir / item dst = self.app_root / item - + if src.exists(): if dst.exists(): if dst.is_dir(): shutil.rmtree(dst) else: dst.unlink() - + if src.is_dir(): shutil.copytree(src, dst) else: shutil.copy2(src, dst) - + # Update version info self._update_version_info(update_info.version) - + return True - - except Exception as e: + + except Exception: return False - + def _update_version_info(self, new_version: str): """Update version information in the application.""" # Update __init__.py version - init_file = self.app_root / 'mini_rag' / '__init__.py' + init_file = self.app_root / "mini_rag" / "__init__.py" if init_file.exists(): try: content = init_file.read_text() updated_content = content.replace( f'__version__ = "{self.current_version}"', - f'__version__ = "{new_version}"' + f'__version__ = "{new_version}"', ) init_file.write_text(updated_content) except Exception: pass - + def rollback_update(self) -> bool: """ Rollback to the backup version if update failed. - + Returns: True if rollback successful """ if not self.backup_dir.exists(): return False - + try: # Restore from backup for item in self.backup_dir.iterdir(): dst = self.app_root / item.name - + if dst.exists(): if dst.is_dir(): shutil.rmtree(dst) else: dst.unlink() - + if item.is_dir(): shutil.copytree(item, dst) else: shutil.copy2(item, dst) - + return True - - except Exception as e: + + except Exception: return False - + def restart_application(self): """Restart the application after update.""" try: # Get the current script path - current_script = sys.argv[0] - + # sys.argv[0] # Unused variable removed + # Restart with the same arguments - if sys.platform.startswith('win'): + if sys.platform.startswith("win"): # Windows subprocess.Popen([sys.executable] + sys.argv) else: # Unix-like systems os.execv(sys.executable, [sys.executable] + sys.argv) - - except Exception as e: + + except Exception: # If restart fails, just exit gracefully - print(f"\nโœ… Update complete! Please restart the application manually.") + print("\nโœ… Update complete! Please restart the application manually.") sys.exit(0) def get_legacy_notification() -> Optional[str]: """ Check if this is a legacy version that needs urgent notification. - + For users who downloaded before the auto-update system. """ try: # Check if this is a very old version by looking for cache file # Old versions won't have update cache, so we can detect them app_root = Path(__file__).parent.parent - cache_file = app_root / ".update_cache.json" - + # app_root / ".update_cache.json" # Unused variable removed + # Also check version in __init__.py to see if it's old - init_file = app_root / 'mini_rag' / '__init__.py' + init_file = app_root / "mini_rag" / "__init__.py" if init_file.exists(): content = init_file.read_text() if '__version__ = "2.0.' in content or '__version__ = "1.' in content: @@ -424,7 +436,7 @@ def get_legacy_notification() -> Optional[str]: Your version of FSS-Mini-RAG is missing critical updates! ๐Ÿ”ง Recent improvements include: -โ€ข Fixed LLM response formatting issues +โ€ข Fixed LLM response formatting issues โ€ข Added context window configuration โ€ข Improved Windows installer reliability โ€ข Added auto-update system (this notification!) @@ -436,26 +448,28 @@ Your version of FSS-Mini-RAG is missing critical updates! """ except Exception: pass - + return None # Global convenience functions _updater_instance = None + def check_for_updates() -> Optional[UpdateInfo]: """Global function to check for updates.""" global _updater_instance if _updater_instance is None: _updater_instance = UpdateChecker() - + if _updater_instance.should_check_for_updates(): return _updater_instance.check_for_updates() return None + def get_updater() -> UpdateChecker: """Get the global updater instance.""" global _updater_instance if _updater_instance is None: _updater_instance = UpdateChecker() - return _updater_instance \ No newline at end of file + return _updater_instance diff --git a/mini_rag/venv_checker.py b/mini_rag/venv_checker.py index 492303d..55e2a2b 100644 --- a/mini_rag/venv_checker.py +++ b/mini_rag/venv_checker.py @@ -4,64 +4,70 @@ Virtual Environment Checker Ensures scripts run in proper Python virtual environment for consistency and safety. """ -import sys import os -import sysconfig +import sys from pathlib import Path + def is_in_virtualenv() -> bool: """Check if we're running in a virtual environment.""" # Check for virtual environment indicators return ( - hasattr(sys, 'real_prefix') or # virtualenv - (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix) or # venv/pyvenv - os.environ.get('VIRTUAL_ENV') is not None # Environment variable + hasattr(sys, "real_prefix") + or (hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix) # virtualenv + or os.environ.get("VIRTUAL_ENV") is not None # venv/pyvenv # Environment variable ) + def get_expected_venv_path() -> Path: """Get the expected virtual environment path for this project.""" # Assume .venv in the same directory as the script script_dir = Path(__file__).parent.parent - return script_dir / '.venv' + return script_dir / ".venv" + def check_correct_venv() -> tuple[bool, str]: """ Check if we're in the correct virtual environment. - + Returns: (is_correct, message) """ if not is_in_virtualenv(): return False, "not in virtual environment" - + expected_venv = get_expected_venv_path() if not expected_venv.exists(): return False, "expected virtual environment not found" - - current_venv = os.environ.get('VIRTUAL_ENV') + + current_venv = os.environ.get("VIRTUAL_ENV") if current_venv: current_venv_path = Path(current_venv).resolve() expected_venv_path = expected_venv.resolve() - + if current_venv_path != expected_venv_path: - return False, f"wrong virtual environment (using {current_venv_path}, expected {expected_venv_path})" - + return ( + False, + f"wrong virtual environment (using {current_venv_path}, expected {expected_venv_path})", + ) + return True, "correct virtual environment" + def show_venv_warning(script_name: str = "script") -> None: """Show virtual environment warning with helpful instructions.""" expected_venv = get_expected_venv_path() - + print("โš ๏ธ VIRTUAL ENVIRONMENT WARNING") print("=" * 50) print() print(f"This {script_name} should be run in a Python virtual environment for:") print(" โ€ข Consistent dependencies") - print(" โ€ข Isolated package versions") + print(" โ€ข Isolated package versions") print(" โ€ข Proper security isolation") print(" โ€ข Reliable functionality") print() - + if expected_venv.exists(): print("โœ… Virtual environment found!") print(f" Location: {expected_venv}") @@ -82,7 +88,7 @@ def show_venv_warning(script_name: str = "script") -> None: print(f" python3 -m venv {expected_venv}") print(f" source {expected_venv}/bin/activate") print(" pip install -r requirements.txt") - + print() print("๐Ÿ’ก Why this matters:") print(" Without a virtual environment, you may experience:") @@ -92,22 +98,23 @@ def show_venv_warning(script_name: str = "script") -> None: print(" โ€ข Potential system-wide package pollution") print() + def check_and_warn_venv(script_name: str = "script", force_exit: bool = False) -> bool: """ Check virtual environment and warn if needed. - + Args: script_name: Name of the script for user-friendly messages force_exit: Whether to exit if not in correct venv - + Returns: True if in correct venv, False otherwise """ is_correct, message = check_correct_venv() - + if not is_correct: show_venv_warning(script_name) - + if force_exit: print(f"โ›” Exiting {script_name} for your safety.") print(" Please activate the virtual environment and try again.") @@ -116,27 +123,32 @@ def check_and_warn_venv(script_name: str = "script", force_exit: bool = False) - print(f"โš ๏ธ Continuing anyway, but {script_name} may not work correctly...") print() return False - + return True + def require_venv(script_name: str = "script") -> None: """Require virtual environment or exit.""" check_and_warn_venv(script_name, force_exit=True) + # Quick test function + + def main(): """Test the virtual environment checker.""" print("๐Ÿงช Virtual Environment Checker Test") print("=" * 40) - + print(f"In virtual environment: {is_in_virtualenv()}") print(f"Expected venv path: {get_expected_venv_path()}") - + is_correct, message = check_correct_venv() print(f"Correct venv: {is_correct} ({message})") - + if not is_correct: show_venv_warning("test script") + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/mini_rag/watcher.py b/mini_rag/watcher.py index 887f54e..3ffecb8 100644 --- a/mini_rag/watcher.py +++ b/mini_rag/watcher.py @@ -4,14 +4,21 @@ Monitors project files and updates the index incrementally. """ import logging -import threading import queue +import threading import time -from pathlib import Path -from typing import Set, Optional, Callable from datetime import datetime +from pathlib import Path +from typing import Callable, Optional, Set + +from watchdog.events import ( + FileCreatedEvent, + FileDeletedEvent, + FileModifiedEvent, + FileMovedEvent, + FileSystemEventHandler, +) from watchdog.observers import Observer -from watchdog.events import FileSystemEventHandler, FileModifiedEvent, FileCreatedEvent, FileDeletedEvent, FileMovedEvent from .indexer import ProjectIndexer @@ -20,11 +27,11 @@ logger = logging.getLogger(__name__) class UpdateQueue: """Thread-safe queue for file updates with deduplication.""" - + def __init__(self, delay: float = 1.0): """ Initialize update queue. - + Args: delay: Delay in seconds before processing updates (for debouncing) """ @@ -33,24 +40,24 @@ class UpdateQueue: self.lock = threading.Lock() self.delay = delay self.last_update = {} # Track last update time per file - + def add(self, file_path: Path): """Add a file to the update queue.""" with self.lock: file_str = str(file_path) current_time = time.time() - + # Check if we should debounce this update if file_str in self.last_update: if current_time - self.last_update[file_str] < self.delay: return # Skip this update - + self.last_update[file_str] = current_time - + if file_str not in self.pending: self.pending.add(file_str) self.queue.put(file_path) - + def get(self, timeout: Optional[float] = None) -> Optional[Path]: """Get next file from queue.""" try: @@ -60,11 +67,11 @@ class UpdateQueue: return file_path except queue.Empty: return None - + def empty(self) -> bool: """Check if queue is empty.""" return self.queue.empty() - + def size(self) -> int: """Get queue size.""" return self.queue.qsize() @@ -72,15 +79,17 @@ class UpdateQueue: class CodeFileEventHandler(FileSystemEventHandler): """Handles file system events for code files.""" - - def __init__(self, - update_queue: UpdateQueue, - include_patterns: Set[str], - exclude_patterns: Set[str], - project_path: Path): + + def __init__( + self, + update_queue: UpdateQueue, + include_patterns: Set[str], + exclude_patterns: Set[str], + project_path: Path, + ): """ Initialize event handler. - + Args: update_queue: Queue for file updates include_patterns: File patterns to include @@ -91,47 +100,47 @@ class CodeFileEventHandler(FileSystemEventHandler): self.include_patterns = include_patterns self.exclude_patterns = exclude_patterns self.project_path = project_path - + def _should_process(self, file_path: str) -> bool: """Check if file should be processed.""" path = Path(file_path) - + # Check if it's a file (not directory) if not path.is_file(): return False - + # Check exclude patterns path_str = str(path) for pattern in self.exclude_patterns: if pattern in path_str: return False - + # Check include patterns for pattern in self.include_patterns: if path.match(pattern): return True - + return False - + def on_modified(self, event: FileModifiedEvent): """Handle file modification.""" if not event.is_directory and self._should_process(event.src_path): logger.debug(f"File modified: {event.src_path}") self.update_queue.add(Path(event.src_path)) - + def on_created(self, event: FileCreatedEvent): """Handle file creation.""" if not event.is_directory and self._should_process(event.src_path): logger.debug(f"File created: {event.src_path}") self.update_queue.add(Path(event.src_path)) - + def on_deleted(self, event: FileDeletedEvent): """Handle file deletion.""" if not event.is_directory and self._should_process(event.src_path): logger.debug(f"File deleted: {event.src_path}") # Add deletion task to queue (we'll handle it differently) self.update_queue.add(Path(event.src_path)) - + def on_moved(self, event: FileMovedEvent): """Handle file move/rename.""" if not event.is_directory: @@ -145,16 +154,18 @@ class CodeFileEventHandler(FileSystemEventHandler): class FileWatcher: """Watches project files and updates index automatically.""" - - def __init__(self, - project_path: Path, - indexer: Optional[ProjectIndexer] = None, - update_delay: float = 1.0, - batch_size: int = 10, - batch_timeout: float = 5.0): + + def __init__( + self, + project_path: Path, + indexer: Optional[ProjectIndexer] = None, + update_delay: float = 1.0, + batch_size: int = 10, + batch_timeout: float = 5.0, + ): """ Initialize file watcher. - + Args: project_path: Path to project to watch indexer: ProjectIndexer instance (creates one if not provided) @@ -167,86 +178,79 @@ class FileWatcher: self.update_delay = update_delay self.batch_size = batch_size self.batch_timeout = batch_timeout - + # Initialize components self.update_queue = UpdateQueue(delay=update_delay) self.observer = Observer() self.worker_thread = None self.running = False - + # Get patterns from indexer self.include_patterns = set(self.indexer.include_patterns) self.exclude_patterns = set(self.indexer.exclude_patterns) - + # Statistics self.stats = { - 'files_updated': 0, - 'files_failed': 0, - 'started_at': None, - 'last_update': None, + "files_updated": 0, + "files_failed": 0, + "started_at": None, + "last_update": None, } - + def start(self): """Start watching for file changes.""" if self.running: logger.warning("Watcher is already running") return - + logger.info(f"Starting file watcher for {self.project_path}") - + # Set up file system observer event_handler = CodeFileEventHandler( self.update_queue, self.include_patterns, self.exclude_patterns, - self.project_path + self.project_path, ) - - self.observer.schedule( - event_handler, - str(self.project_path), - recursive=True - ) - + + self.observer.schedule(event_handler, str(self.project_path), recursive=True) + # Start worker thread self.running = True - self.worker_thread = threading.Thread( - target=self._process_updates, - daemon=True - ) + self.worker_thread = threading.Thread(target=self._process_updates, daemon=True) self.worker_thread.start() - + # Start observer self.observer.start() - - self.stats['started_at'] = datetime.now() + + self.stats["started_at"] = datetime.now() logger.info("File watcher started successfully") - + def stop(self): """Stop watching for file changes.""" if not self.running: return - + logger.info("Stopping file watcher...") - + # Stop observer self.observer.stop() self.observer.join() - + # Stop worker thread self.running = False if self.worker_thread: self.worker_thread.join(timeout=5.0) - + logger.info("File watcher stopped") - + def _process_updates(self): """Worker thread that processes file updates.""" logger.info("Update processor thread started") - + batch = [] batch_start_time = None - + while self.running: try: # Calculate timeout for getting next item @@ -263,46 +267,46 @@ class FileWatcher: else: # Wait for more items or timeout timeout = min(0.1, self.batch_timeout - elapsed) - + # Get next file from queue file_path = self.update_queue.get(timeout=timeout) - + if file_path: # Add to batch if not batch: batch_start_time = time.time() batch.append(file_path) - + # Check if batch is full if len(batch) >= self.batch_size: self._process_batch(batch) batch = [] batch_start_time = None - + except queue.Empty: # Check if we have a pending batch that's timed out if batch and (time.time() - batch_start_time) >= self.batch_timeout: self._process_batch(batch) batch = [] batch_start_time = None - + except Exception as e: logger.error(f"Error in update processor: {e}") time.sleep(1) # Prevent tight loop on error - + # Process any remaining items if batch: self._process_batch(batch) - + logger.info("Update processor thread stopped") - + def _process_batch(self, files: list[Path]): """Process a batch of file updates.""" if not files: return - + logger.info(f"Processing batch of {len(files)} file updates") - + for file_path in files: try: if file_path.exists(): @@ -313,87 +317,91 @@ class FileWatcher: # File doesn't exist - delete from index logger.debug(f"Deleting {file_path} from index - file no longer exists") success = self.indexer.delete_file(file_path) - + if success: - self.stats['files_updated'] += 1 + self.stats["files_updated"] += 1 else: - self.stats['files_failed'] += 1 - - self.stats['last_update'] = datetime.now() - + self.stats["files_failed"] += 1 + + self.stats["last_update"] = datetime.now() + except Exception as e: logger.error(f"Failed to process {file_path}: {e}") - self.stats['files_failed'] += 1 - - logger.info(f"Batch processing complete. Updated: {self.stats['files_updated']}, Failed: {self.stats['files_failed']}") - + self.stats["files_failed"] += 1 + + logger.info( + f"Batch processing complete. Updated: {self.stats['files_updated']}, Failed: {self.stats['files_failed']}" + ) + def get_statistics(self) -> dict: """Get watcher statistics.""" stats = self.stats.copy() - stats['queue_size'] = self.update_queue.size() - stats['is_running'] = self.running - - if stats['started_at']: - uptime = datetime.now() - stats['started_at'] - stats['uptime_seconds'] = uptime.total_seconds() - + stats["queue_size"] = self.update_queue.size() + stats["is_running"] = self.running + + if stats["started_at"]: + uptime = datetime.now() - stats["started_at"] + stats["uptime_seconds"] = uptime.total_seconds() + return stats - + def wait_for_updates(self, timeout: Optional[float] = None) -> bool: """ Wait for pending updates to complete. - + Args: timeout: Maximum time to wait in seconds - + Returns: True if all updates completed, False if timeout """ start_time = time.time() - + while not self.update_queue.empty(): if timeout and (time.time() - start_time) > timeout: return False time.sleep(0.1) - + # Wait a bit more to ensure batch processing completes time.sleep(self.batch_timeout + 0.5) return True - + def __enter__(self): """Context manager entry.""" self.start() return self - + def __exit__(self, exc_type, exc_val, exc_tb): """Context manager exit.""" self.stop() # Convenience function + + def watch_project(project_path: Path, callback: Optional[Callable] = None): """ Watch a project for changes and update index automatically. - + Args: project_path: Path to project callback: Optional callback function called after each update """ watcher = FileWatcher(project_path) - + try: watcher.start() logger.info(f"Watching {project_path} for changes. Press Ctrl+C to stop.") - + while True: time.sleep(1) - + # Call callback if provided if callback: stats = watcher.get_statistics() callback(stats) - + except KeyboardInterrupt: logger.info("Stopping watcher...") finally: - watcher.stop() \ No newline at end of file + watcher.stop() diff --git a/mini_rag/windows_console_fix.py b/mini_rag/windows_console_fix.py index 2e5abfb..9742301 100644 --- a/mini_rag/windows_console_fix.py +++ b/mini_rag/windows_console_fix.py @@ -3,9 +3,9 @@ Windows Console Unicode/Emoji Fix Reliable Windows console Unicode/emoji support for 2025. """ -import sys -import os import io +import os +import sys def fix_windows_console(): @@ -14,28 +14,33 @@ def fix_windows_console(): Call this at the start of any script that needs to output Unicode/emojis. """ # Set environment variable for UTF-8 mode - os.environ['PYTHONUTF8'] = '1' - + os.environ["PYTHONUTF8"] = "1" + # For Python 3.7+ - if hasattr(sys.stdout, 'reconfigure'): - sys.stdout.reconfigure(encoding='utf-8') - sys.stderr.reconfigure(encoding='utf-8') - if hasattr(sys.stdin, 'reconfigure'): - sys.stdin.reconfigure(encoding='utf-8') + if hasattr(sys.stdout, "reconfigure"): + sys.stdout.reconfigure(encoding="utf-8") + sys.stderr.reconfigure(encoding="utf-8") + if hasattr(sys.stdin, "reconfigure"): + sys.stdin.reconfigure(encoding="utf-8") else: # For older Python versions - if sys.platform == 'win32': + if sys.platform == "win32": # Replace streams with UTF-8 versions - sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', line_buffering=True) - sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', line_buffering=True) - + sys.stdout = io.TextIOWrapper( + sys.stdout.buffer, encoding="utf-8", line_buffering=True + ) + sys.stderr = io.TextIOWrapper( + sys.stderr.buffer, encoding="utf-8", line_buffering=True + ) + # Also set the console code page to UTF-8 on Windows - if sys.platform == 'win32': + if sys.platform == "win32": import subprocess + try: # Set console to UTF-8 code page - subprocess.run(['chcp', '65001'], shell=True, capture_output=True) - except: + subprocess.run(["chcp", "65001"], shell=True, capture_output=True) + except (OSError, subprocess.SubprocessError): pass @@ -44,12 +49,14 @@ fix_windows_console() # Test function to verify it works + + def test_emojis(): """Test that emojis work properly.""" print("Testing emoji output:") print(" Check mark") print(" Cross mark") - print(" Rocket") + print(" Rocket") print(" Fire") print(" Computer") print(" Python") @@ -57,7 +64,7 @@ def test_emojis(): print(" Search") print(" Lightning") print(" Sparkles") - + if __name__ == "__main__": - test_emojis() \ No newline at end of file + test_emojis() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..1befaaa --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,34 @@ +[tool.isort] +profile = "black" +line_length = 95 +multi_line_output = 3 +include_trailing_comma = true +force_grid_wrap = 0 +use_parentheses = true +ensure_newline_before_comments = true +src_paths = ["mini_rag", "tests", "examples", "scripts"] +known_first_party = ["mini_rag"] +sections = ["FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER"] +skip = [".venv", ".venv-linting", "__pycache__", ".git"] +skip_glob = ["*.egg-info/*", "build/*", "dist/*"] + +[tool.black] +line-length = 95 +target-version = ['py310'] +include = '\.pyi?$' +extend-exclude = ''' +/( + # directories + \.eggs + | \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | \.venv-linting + | _build + | buck-out + | build + | dist +)/ +''' \ No newline at end of file diff --git a/rag-mini b/rag-mini index 770f267..897c434 100755 --- a/rag-mini +++ b/rag-mini @@ -329,7 +329,7 @@ main() { ;; "index"|"search"|"explore"|"status"|"update"|"check-update") # Direct CLI commands - call Python script - exec "$PYTHON" "$SCRIPT_DIR/rag-mini.py" "$@" + exec "$PYTHON" "$SCRIPT_DIR/bin/rag-mini.py" "$@" ;; *) # Unknown command - show help diff --git a/rag-tui b/rag-tui index 175ac6c..5bb20d4 100755 --- a/rag-tui +++ b/rag-tui @@ -19,4 +19,4 @@ if [ ! -f "$PYTHON" ]; then fi # Launch TUI -exec "$PYTHON" "$SCRIPT_DIR/rag-tui.py" "$@" \ No newline at end of file +exec "$PYTHON" "$SCRIPT_DIR/bin/rag-tui.py" "$@" \ No newline at end of file diff --git a/scripts/setup-github-template.py b/scripts/setup-github-template.py index 7e6a092..d786d92 100755 --- a/scripts/setup-github-template.py +++ b/scripts/setup-github-template.py @@ -6,67 +6,67 @@ Converts a project to use the auto-update template system. This script helps migrate projects from Gitea to GitHub with auto-update capability. """ -import os -import sys +import argparse import json import shutil -import argparse +import sys from pathlib import Path -from typing import Dict, Any, Optional +from typing import Dict, Optional + def setup_project_template( project_path: Path, repo_owner: str, repo_name: str, project_type: str = "python", - include_auto_update: bool = True + include_auto_update: bool = True, ) -> bool: """ Setup a project to use the GitHub auto-update template system. - + Args: project_path: Path to the project directory repo_owner: GitHub username/organization - repo_name: GitHub repository name + repo_name: GitHub repository name project_type: Type of project (python, general) include_auto_update: Whether to include auto-update system - + Returns: True if setup successful """ - + print(f"๐Ÿš€ Setting up GitHub template for: {repo_owner}/{repo_name}") print(f"๐Ÿ“ Project path: {project_path}") print(f"๐Ÿ”ง Project type: {project_type}") print(f"๐Ÿ”„ Auto-update: {'Enabled' if include_auto_update else 'Disabled'}") print() - + try: # Create .github directory structure github_dir = project_path / ".github" workflows_dir = github_dir / "workflows" templates_dir = github_dir / "ISSUE_TEMPLATE" - + # Ensure directories exist workflows_dir.mkdir(parents=True, exist_ok=True) templates_dir.mkdir(parents=True, exist_ok=True) - + # 1. Setup GitHub Actions workflows setup_workflows(workflows_dir, repo_owner, repo_name, project_type) - + # 2. Setup auto-update system if requested if include_auto_update: setup_auto_update_system(project_path, repo_owner, repo_name) - + # 3. Create issue templates setup_issue_templates(templates_dir) - + # 4. Create/update project configuration setup_project_config(project_path, repo_owner, repo_name, include_auto_update) - + # 5. Create README template if needed setup_readme_template(project_path, repo_owner, repo_name) - + print("โœ… GitHub template setup completed successfully!") print() print("๐Ÿ“‹ Next Steps:") @@ -75,20 +75,21 @@ def setup_project_template( print("3. Test auto-update system: ./project check-update") print("4. Enable GitHub Pages for documentation (optional)") print() - + return True - + except Exception as e: print(f"โŒ Setup failed: {e}") return False + def setup_workflows(workflows_dir: Path, repo_owner: str, repo_name: str, project_type: str): """Setup GitHub Actions workflow files.""" - + print("๐Ÿ”ง Setting up GitHub Actions workflows...") - + # Release workflow - release_workflow = f"""name: Auto Release & Update System + release_workflow = """name: Auto Release & Update System on: push: tags: @@ -105,18 +106,18 @@ jobs: runs-on: ubuntu-latest permissions: contents: write - + steps: - name: Checkout code uses: actions/checkout@v4 with: fetch-depth: 0 - + - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.11' - + - name: Extract version id: version run: | @@ -127,18 +128,18 @@ jobs: fi echo "version=$VERSION" >> $GITHUB_OUTPUT echo "clean_version=${{VERSION#v}}" >> $GITHUB_OUTPUT - + - name: Update version in code run: | VERSION="${{{{ steps.version.outputs.clean_version }}}}" # Update version files find . -name "__init__.py" -exec sed -i 's/__version__ = ".*"/__version__ = "'$VERSION'"/' {{}} + - + - name: Generate release notes id: release_notes run: | VERSION="${{{{ steps.version.outputs.version }}}}" - + # Get commits since last tag LAST_TAG=$(git describe --tags --abbrev=0 HEAD~1 2>/dev/null || echo "") if [ -n "$LAST_TAG" ]; then @@ -146,20 +147,20 @@ jobs: else COMMITS=$(git log --oneline --pretty=format:"โ€ข %s" | head -10) fi - + # Create release notes cat > release_notes.md << EOF ## What's New in $VERSION - + ### ๐Ÿš€ Changes $COMMITS - + ### ๐Ÿ“ฅ Installation Download and install the latest version: \`\`\`bash curl -sSL https://github.com/{repo_owner}/{repo_name}/releases/latest/download/install.sh | bash \`\`\` - + ### ๐Ÿ”„ Auto-Update If you have auto-update support: \`\`\`bash @@ -167,7 +168,7 @@ jobs: ./{repo_name} update \`\`\` EOF - + - name: Create GitHub Release uses: softprops/action-gh-release@v2 with: @@ -181,12 +182,12 @@ jobs: *.bat requirements.txt """ - + (workflows_dir / "release.yml").write_text(release_workflow) - + # CI workflow for Python projects if project_type == "python": - ci_workflow = f"""name: CI/CD Pipeline + ci_workflow = """name: CI/CD Pipeline on: push: branches: [ main, develop ] @@ -201,25 +202,25 @@ jobs: matrix: os: [ubuntu-latest, windows-latest, macos-latest] python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] - + steps: - name: Checkout code uses: actions/checkout@v4 - + - name: Set up Python ${{{{ matrix.python-version }}}} uses: actions/setup-python@v5 with: python-version: ${{{{ matrix.python-version }}}} - + - name: Install dependencies run: | python -m pip install --upgrade pip pip install -r requirements.txt - + - name: Run tests run: | python -c "import {repo_name.replace('-', '_')}; print('โœ… Import successful')" - + - name: Test auto-update system run: | python -c " @@ -231,33 +232,38 @@ jobs: " """ (workflows_dir / "ci.yml").write_text(ci_workflow) - + print(" โœ… GitHub Actions workflows created") + def setup_auto_update_system(project_path: Path, repo_owner: str, repo_name: str): """Setup the auto-update system for the project.""" - + print("๐Ÿ”„ Setting up auto-update system...") - + # Copy updater.py from FSS-Mini-RAG as template template_updater = Path(__file__).parent.parent / "mini_rag" / "updater.py" - + if template_updater.exists(): # Create project module directory if needed - module_name = repo_name.replace('-', '_') + module_name = repo_name.replace("-", "_") module_dir = project_path / module_name module_dir.mkdir(exist_ok=True) - + # Copy and customize updater target_updater = module_dir / "updater.py" shutil.copy2(template_updater, target_updater) - + # Customize for this project content = target_updater.read_text() - content = content.replace('repo_owner: str = "FSSCoding"', f'repo_owner: str = "{repo_owner}"') - content = content.replace('repo_name: str = "Fss-Mini-Rag"', f'repo_name: str = "{repo_name}"') + content = content.replace( + 'repo_owner: str = "FSSCoding"', f'repo_owner: str = "{repo_owner}"' + ) + content = content.replace( + 'repo_name: str = "Fss-Mini-Rag"', f'repo_name: str = "{repo_name}"' + ) target_updater.write_text(content) - + # Update __init__.py to include updater init_file = module_dir / "__init__.py" if init_file.exists(): @@ -272,16 +278,17 @@ except ImportError: pass """ init_file.write_text(content) - + print(" โœ… Auto-update system configured") else: print(" โš ๏ธ Template updater not found, you'll need to implement manually") + def setup_issue_templates(templates_dir: Path): """Setup GitHub issue templates.""" - + print("๐Ÿ“ Setting up issue templates...") - + bug_template = """--- name: Bug Report about: Create a report to help us improve @@ -312,7 +319,7 @@ A clear and concise description of what you expected to happen. **Additional context** Add any other context about the problem here. """ - + feature_template = """--- name: Feature Request about: Suggest an idea for this project @@ -334,46 +341,50 @@ A clear and concise description of any alternative solutions you've considered. **Additional context** Add any other context or screenshots about the feature request here. """ - + (templates_dir / "bug_report.md").write_text(bug_template) (templates_dir / "feature_request.md").write_text(feature_template) - + print(" โœ… Issue templates created") -def setup_project_config(project_path: Path, repo_owner: str, repo_name: str, include_auto_update: bool): + +def setup_project_config( + project_path: Path, repo_owner: str, repo_name: str, include_auto_update: bool +): """Setup project configuration file.""" - + print("โš™๏ธ Setting up project configuration...") - + config = { "project": { "name": repo_name, "owner": repo_owner, "github_url": f"https://github.com/{repo_owner}/{repo_name}", - "auto_update_enabled": include_auto_update + "auto_update_enabled": include_auto_update, }, "github": { "template_version": "1.0.0", "last_sync": None, - "workflows_enabled": True - } + "workflows_enabled": True, + }, } - + config_file = project_path / ".github" / "project-config.json" - with open(config_file, 'w') as f: + with open(config_file, "w") as f: json.dump(config, f, indent=2) - + print(" โœ… Project configuration created") + def setup_readme_template(project_path: Path, repo_owner: str, repo_name: str): """Setup README template if one doesn't exist.""" - + readme_file = project_path / "README.md" - + if not readme_file.exists(): print("๐Ÿ“– Creating README template...") - - readme_content = f"""# {repo_name} + + readme_content = """# {repo_name} > A brief description of your project @@ -390,7 +401,7 @@ curl -sSL https://github.com/{repo_owner}/{repo_name}/releases/latest/download/i ## Features - โœจ Feature 1 -- ๐Ÿš€ Feature 2 +- ๐Ÿš€ Feature 2 - ๐Ÿ”ง Feature 3 ## Installation @@ -441,10 +452,11 @@ This project includes automatic update checking: ๐Ÿค– **Auto-Update Enabled**: This project will notify you of new versions automatically! """ - + readme_file.write_text(readme_content) print(" โœ… README template created") + def main(): """Main entry point.""" parser = argparse.ArgumentParser( @@ -454,32 +466,38 @@ def main(): Examples: python setup-github-template.py myproject --owner username --name my-project python setup-github-template.py /path/to/project --owner org --name cool-tool --no-auto-update - """ + """, ) - - parser.add_argument('project_path', type=Path, help='Path to project directory') - parser.add_argument('--owner', required=True, help='GitHub username or organization') - parser.add_argument('--name', required=True, help='GitHub repository name') - parser.add_argument('--type', choices=['python', 'general'], default='python', - help='Project type (default: python)') - parser.add_argument('--no-auto-update', action='store_true', - help='Disable auto-update system') - + + parser.add_argument("project_path", type=Path, help="Path to project directory") + parser.add_argument("--owner", required=True, help="GitHub username or organization") + parser.add_argument("--name", required=True, help="GitHub repository name") + parser.add_argument( + "--type", + choices=["python", "general"], + default="python", + help="Project type (default: python)", + ) + parser.add_argument( + "--no-auto-update", action="store_true", help="Disable auto-update system" + ) + args = parser.parse_args() - + if not args.project_path.exists(): print(f"โŒ Project path does not exist: {args.project_path}") sys.exit(1) - + success = setup_project_template( project_path=args.project_path, repo_owner=args.owner, repo_name=args.name, project_type=args.type, - include_auto_update=not args.no_auto_update + include_auto_update=not args.no_auto_update, ) - + sys.exit(0 if success else 1) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/test-configs.py b/scripts/test-configs.py index 50eb1bb..eddba2c 100755 --- a/scripts/test-configs.py +++ b/scripts/test-configs.py @@ -4,80 +4,87 @@ Test script to validate all config examples are syntactically correct and contain required fields for FSS-Mini-RAG. """ -import yaml import sys from pathlib import Path -from typing import Dict, Any, List +from typing import Any, Dict, List + +import yaml + def validate_config_structure(config: Dict[str, Any], config_name: str) -> List[str]: """Validate that config has required structure.""" errors = [] - + # Required sections - required_sections = ['chunking', 'streaming', 'files', 'embedding', 'search'] + required_sections = ["chunking", "streaming", "files", "embedding", "search"] for section in required_sections: if section not in config: errors.append(f"{config_name}: Missing required section '{section}'") - + # Validate chunking section - if 'chunking' in config: - chunking = config['chunking'] - required_chunking = ['max_size', 'min_size', 'strategy'] + if "chunking" in config: + chunking = config["chunking"] + required_chunking = ["max_size", "min_size", "strategy"] for field in required_chunking: if field not in chunking: errors.append(f"{config_name}: Missing chunking.{field}") - + # Validate types and ranges - if 'max_size' in chunking and not isinstance(chunking['max_size'], int): + if "max_size" in chunking and not isinstance(chunking["max_size"], int): errors.append(f"{config_name}: chunking.max_size must be integer") - if 'min_size' in chunking and not isinstance(chunking['min_size'], int): + if "min_size" in chunking and not isinstance(chunking["min_size"], int): errors.append(f"{config_name}: chunking.min_size must be integer") - if 'strategy' in chunking and chunking['strategy'] not in ['semantic', 'fixed']: + if "strategy" in chunking and chunking["strategy"] not in ["semantic", "fixed"]: errors.append(f"{config_name}: chunking.strategy must be 'semantic' or 'fixed'") - + # Validate embedding section - if 'embedding' in config: - embedding = config['embedding'] - if 'preferred_method' in embedding: - valid_methods = ['ollama', 'ml', 'hash', 'auto'] - if embedding['preferred_method'] not in valid_methods: - errors.append(f"{config_name}: embedding.preferred_method must be one of {valid_methods}") - + if "embedding" in config: + embedding = config["embedding"] + if "preferred_method" in embedding: + valid_methods = ["ollama", "ml", "hash", "auto"] + if embedding["preferred_method"] not in valid_methods: + errors.append( + f"{config_name}: embedding.preferred_method must be one of {valid_methods}" + ) + # Validate LLM section (if present) - if 'llm' in config: - llm = config['llm'] - if 'synthesis_temperature' in llm: - temp = llm['synthesis_temperature'] + if "llm" in config: + llm = config["llm"] + if "synthesis_temperature" in llm: + temp = llm["synthesis_temperature"] if not isinstance(temp, (int, float)) or temp < 0 or temp > 1: - errors.append(f"{config_name}: llm.synthesis_temperature must be number between 0-1") - + errors.append( + f"{config_name}: llm.synthesis_temperature must be number between 0-1" + ) + return errors + def test_config_file(config_path: Path) -> bool: """Test a single config file.""" print(f"Testing {config_path.name}...") - + try: # Test YAML parsing - with open(config_path, 'r') as f: + with open(config_path, "r") as f: config = yaml.safe_load(f) - + if not config: print(f" โŒ {config_path.name}: Empty or invalid YAML") return False - + # Test structure errors = validate_config_structure(config, config_path.name) - + if errors: print(f" โŒ {config_path.name}: Structure errors:") for error in errors: print(f" โ€ข {error}") return False - + print(f" โœ… {config_path.name}: Valid") return True - + except yaml.YAMLError as e: print(f" โŒ {config_path.name}: YAML parsing error: {e}") return False @@ -85,31 +92,32 @@ def test_config_file(config_path: Path) -> bool: print(f" โŒ {config_path.name}: Unexpected error: {e}") return False + def main(): """Test all config examples.""" script_dir = Path(__file__).parent project_root = script_dir.parent - examples_dir = project_root / 'examples' - + examples_dir = project_root / "examples" + if not examples_dir.exists(): print(f"โŒ Examples directory not found: {examples_dir}") sys.exit(1) - + # Find all config files - config_files = list(examples_dir.glob('config*.yaml')) - + config_files = list(examples_dir.glob("config*.yaml")) + if not config_files: print(f"โŒ No config files found in {examples_dir}") sys.exit(1) - + print(f"๐Ÿงช Testing {len(config_files)} config files...\n") - + all_passed = True for config_file in sorted(config_files): passed = test_config_file(config_file) if not passed: all_passed = False - + print(f"\n{'='*50}") if all_passed: print("โœ… All config files are valid!") @@ -120,5 +128,6 @@ def main(): print("โŒ Some config files have issues - please fix before release") sys.exit(1) -if __name__ == '__main__': - main() \ No newline at end of file + +if __name__ == "__main__": + main() diff --git a/tests/01_basic_integration_test.py b/tests/01_basic_integration_test.py index 4fec7a7..e88f5fb 100644 --- a/tests/01_basic_integration_test.py +++ b/tests/01_basic_integration_test.py @@ -14,74 +14,82 @@ import sys import tempfile from pathlib import Path +from mini_rag.chunker import CodeChunker +from mini_rag.indexer import ProjectIndexer +from mini_rag.ollama_embeddings import OllamaEmbedder as CodeEmbedder +from mini_rag.search import CodeSearcher + # Check if virtual environment is activated + + def check_venv(): - if 'VIRTUAL_ENV' not in os.environ: + if "VIRTUAL_ENV" not in os.environ: print("โš ๏ธ WARNING: Virtual environment not detected!") print(" This test requires the virtual environment to be activated.") - print(" Run: source .venv/bin/activate && PYTHONPATH=. python tests/01_basic_integration_test.py") + print( + " Run: source .venv/bin/activate && PYTHONPATH=. python tests/01_basic_integration_test.py" + ) print(" Continuing anyway...\n") + check_venv() # Fix Windows encoding -if sys.platform == 'win32': - os.environ['PYTHONUTF8'] = '1' - sys.stdout.reconfigure(encoding='utf-8') +if sys.platform == "win32": + os.environ["PYTHONUTF8"] = "1" + sys.stdout.reconfigure(encoding="utf-8") -from mini_rag.chunker import CodeChunker -from mini_rag.indexer import ProjectIndexer -from mini_rag.search import CodeSearcher -from mini_rag.ollama_embeddings import OllamaEmbedder as CodeEmbedder def main(): print("=" * 60) print("RAG System Integration Demo") print("=" * 60) - + with tempfile.TemporaryDirectory() as tmpdir: project_path = Path(tmpdir) - + # Create sample project files print("\n1. Creating sample project files...") - + # Main calculator module - (project_path / "calculator.py").write_text('''""" + (project_path / "calculator.py").write_text( + '''""" Advanced calculator module with various mathematical operations. """ import math from typing import List, Union + class BasicCalculator: """Basic calculator with fundamental operations.""" - + def __init__(self): """Initialize calculator with result history.""" self.history = [] self.last_result = 0 - + def add(self, a: float, b: float) -> float: """Add two numbers and store result.""" result = a + b self.history.append(f"{a} + {b} = {result}") self.last_result = result return result - + def subtract(self, a: float, b: float) -> float: """Subtract b from a.""" result = a - b self.history.append(f"{a} - {b} = {result}") self.last_result = result return result - + def multiply(self, a: float, b: float) -> float: """Multiply two numbers.""" result = a * b self.history.append(f"{a} * {b} = {result}") self.last_result = result return result - + def divide(self, a: float, b: float) -> float: """Divide a by b with zero check.""" if b == 0: @@ -91,16 +99,17 @@ class BasicCalculator: self.last_result = result return result + class ScientificCalculator(BasicCalculator): """Scientific calculator extending basic operations.""" - + def power(self, base: float, exponent: float) -> float: """Calculate base raised to exponent.""" result = math.pow(base, exponent) self.history.append(f"{base} ^ {exponent} = {result}") self.last_result = result return result - + def sqrt(self, n: float) -> float: """Calculate square root.""" if n < 0: @@ -109,7 +118,7 @@ class ScientificCalculator(BasicCalculator): self.history.append(f"sqrt({n}) = {result}") self.last_result = result return result - + def logarithm(self, n: float, base: float = 10) -> float: """Calculate logarithm with specified base.""" result = math.log(n, base) @@ -123,6 +132,7 @@ def calculate_mean(numbers: List[float]) -> float: return 0.0 return sum(numbers) / len(numbers) + def calculate_median(numbers: List[float]) -> float: """Calculate median of a list of numbers.""" if not numbers: @@ -133,6 +143,7 @@ def calculate_median(numbers: List[float]) -> float: return (sorted_nums[n//2-1] + sorted_nums[n//2]) / 2 return sorted_nums[n//2] + def calculate_mode(numbers: List[float]) -> float: """Calculate mode (most frequent value).""" if not numbers: @@ -142,79 +153,88 @@ def calculate_mode(numbers: List[float]) -> float: frequency[num] = frequency.get(num, 0) + 1 mode = max(frequency.keys(), key=frequency.get) return mode -''') - +''' + ) + # Test file for the calculator - (project_path / "test_calculator.py").write_text('''""" + (project_path / "test_calculator.py").write_text( + '''""" Unit tests for calculator module. """ import unittest from calculator import BasicCalculator, ScientificCalculator, calculate_mean + class TestBasicCalculator(unittest.TestCase): """Test cases for BasicCalculator.""" - + def setUp(self): """Set up test calculator.""" self.calc = BasicCalculator() - + def test_addition(self): """Test addition operation.""" result = self.calc.add(5, 3) self.assertEqual(result, 8) self.assertEqual(self.calc.last_result, 8) - + def test_division_by_zero(self): """Test division by zero raises error.""" with self.assertRaises(ValueError): self.calc.divide(10, 0) + class TestStatistics(unittest.TestCase): """Test statistical functions.""" - + def test_mean(self): """Test mean calculation.""" numbers = [1, 2, 3, 4, 5] self.assertEqual(calculate_mean(numbers), 3.0) - + def test_empty_list(self): """Test mean of empty list.""" self.assertEqual(calculate_mean([]), 0.0) if __name__ == "__main__": unittest.main() -''') - +''' + ) + print(" Created 2 Python files") - + # 2. Index the project print("\n2. Indexing project with intelligent chunking...") - + # Use realistic chunk size chunker = CodeChunker(min_chunk_size=10, max_chunk_size=100) indexer = ProjectIndexer(project_path, chunker=chunker) stats = indexer.index_project() - + print(f" Indexed {stats['files_indexed']} files") print(f" Created {stats['chunks_created']} chunks") print(f" Time: {stats['time_taken']:.2f} seconds") - + # 3. Demonstrate search capabilities print("\n3. Testing search capabilities...") searcher = CodeSearcher(project_path) - + # Test different search types print("\n a) Semantic search for 'calculate average':") results = searcher.search("calculate average", top_k=3) for i, result in enumerate(results, 1): - print(f" {i}. {result.chunk_type} '{result.name}' in {result.file_path} (score: {result.score:.3f})") - + print( + f" {i}. {result.chunk_type} '{result.name}' in {result.file_path} (score: {result.score:.3f})" + ) + print("\n b) BM25-weighted search for 'divide zero':") results = searcher.search("divide zero", top_k=3, semantic_weight=0.2, bm25_weight=0.8) for i, result in enumerate(results, 1): - print(f" {i}. {result.chunk_type} '{result.name}' in {result.file_path} (score: {result.score:.3f})") - + print( + f" {i}. {result.chunk_type} '{result.name}' in {result.file_path} (score: {result.score:.3f})" + ) + print("\n c) Search with context for 'test addition':") results = searcher.search("test addition", top_k=2, include_context=True) for i, result in enumerate(results, 1): @@ -225,39 +245,39 @@ if __name__ == "__main__": print(f" Has previous context: {len(result.context_before)} chars") if result.context_after: print(f" Has next context: {len(result.context_after)} chars") - + # 4. Test chunk navigation print("\n4. Testing chunk navigation...") - + # Get all chunks to find a method df = searcher.table.to_pandas() - method_chunks = df[df['chunk_type'] == 'method'] - + method_chunks = df[df["chunk_type"] == "method"] + if len(method_chunks) > 0: # Pick a method in the middle mid_idx = len(method_chunks) // 2 - chunk_id = method_chunks.iloc[mid_idx]['chunk_id'] - chunk_name = method_chunks.iloc[mid_idx]['name'] - + chunk_id = method_chunks.iloc[mid_idx]["chunk_id"] + chunk_name = method_chunks.iloc[mid_idx]["name"] + print(f"\n Getting context for method '{chunk_name}':") context = searcher.get_chunk_context(chunk_id) - - if context['chunk']: + + if context["chunk"]: print(f" Current: {context['chunk'].name}") - if context['prev']: + if context["prev"]: print(f" Previous: {context['prev'].name}") - if context['next']: + if context["next"]: print(f" Next: {context['next'].name}") - if context['parent']: + if context["parent"]: print(f" Parent class: {context['parent'].name}") - + # 5. Show statistics print("\n5. Index Statistics:") stats = searcher.get_statistics() print(f" - Total chunks: {stats['total_chunks']}") print(f" - Unique files: {stats['unique_files']}") print(f" - Chunk types: {stats['chunk_types']}") - + print("\n" + "=" * 60) print(" All features working correctly!") print("=" * 60) @@ -268,5 +288,6 @@ if __name__ == "__main__": print("- Context-aware search with adjacent chunks") print("- Chunk navigation following code relationships") + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/tests/02_search_examples.py b/tests/02_search_examples.py index 1708934..afd01f4 100644 --- a/tests/02_search_examples.py +++ b/tests/02_search_examples.py @@ -5,9 +5,10 @@ Simple demo of the hybrid search system showing real results. import sys from pathlib import Path + from rich.console import Console -from rich.syntax import Syntax from rich.panel import Panel +from rich.syntax import Syntax from rich.table import Table from mini_rag.search import CodeSearcher @@ -17,102 +18,110 @@ console = Console() def demo_search(project_path: Path): """Run demo searches showing the hybrid system in action.""" - + console.print("\n[bold cyan]Mini RAG Hybrid Search Demo[/bold cyan]\n") - + # Initialize searcher console.print("Initializing search system...") searcher = CodeSearcher(project_path) - + # Get index stats stats = searcher.get_statistics() - if 'error' not in stats: - console.print(f"\n[green] Index ready:[/green] {stats['total_chunks']} chunks from {stats['unique_files']} files") + if "error" not in stats: + console.print( + f"\n[green] Index ready:[/green] {stats['total_chunks']} chunks from {stats['unique_files']} files" + ) console.print(f"[dim]Languages: {', '.join(stats['languages'].keys())}[/dim]") console.print(f"[dim]Chunk types: {', '.join(stats['chunk_types'].keys())}[/dim]\n") - + # Demo queries demos = [ { - 'title': 'Keyword-Heavy Search', - 'query': 'BM25Okapi rank_bm25 search scoring', - 'description': 'This query has specific technical keywords that BM25 excels at finding', - 'top_k': 5 + "title": "Keyword-Heavy Search", + "query": "BM25Okapi rank_bm25 search scoring", + "description": "This query has specific technical keywords that BM25 excels at finding", + "top_k": 5, }, { - 'title': 'Natural Language Query', - 'query': 'how to build search index from database chunks', - 'description': 'This semantic query benefits from transformer embeddings understanding intent', - 'top_k': 5 + "title": "Natural Language Query", + "query": "how to build search index from database chunks", + "description": "This semantic query benefits from transformer embeddings understanding intent", + "top_k": 5, }, { - 'title': 'Mixed Technical Query', - 'query': 'vector embeddings for semantic code search with transformers', - 'description': 'This hybrid query combines technical terms with conceptual understanding', - 'top_k': 5 + "title": "Mixed Technical Query", + "query": "vector embeddings for semantic code search with transformers", + "description": "This hybrid query combines technical terms with conceptual understanding", + "top_k": 5, }, { - 'title': 'Function Search', - 'query': 'search method implementation with filters', - 'description': 'Looking for specific function implementations', - 'top_k': 5 - } + "title": "Function Search", + "query": "search method implementation with filters", + "description": "Looking for specific function implementations", + "top_k": 5, + }, ] - + for demo in demos: console.rule(f"\n[bold yellow]{demo['title']}[/bold yellow]") console.print(f"[dim]{demo['description']}[/dim]") console.print(f"\n[cyan]Query:[/cyan] '{demo['query']}'") - + # Run search with hybrid mode results = searcher.search( - query=demo['query'], - top_k=demo['top_k'], + query=demo["query"], + top_k=demo["top_k"], semantic_weight=0.7, - bm25_weight=0.3 + bm25_weight=0.3, ) - + if not results: console.print("[red]No results found![/red]") continue - + console.print(f"\n[green]Found {len(results)} results:[/green]\n") - + # Show each result for i, result in enumerate(results, 1): # Create result panel header = f"#{i} {result.file_path}:{result.start_line}-{result.end_line}" - + # Get code preview lines = result.content.splitlines() if len(lines) > 10: - preview_lines = lines[:8] + ['...'] + lines[-2:] + preview_lines = lines[:8] + ["..."] + lines[-2:] else: preview_lines = lines - - preview = '\n'.join(preview_lines) - + + preview = "\n".join(preview_lines) + # Create info table info = Table.grid(padding=0) info.add_column(style="cyan", width=12) info.add_column(style="white") - + info.add_row("Score:", f"{result.score:.3f}") info.add_row("Type:", result.chunk_type) info.add_row("Name:", result.name or "N/A") info.add_row("Language:", result.language) - + # Display result - console.print(Panel( - f"{info}\n\n[dim]{preview}[/dim]", - title=header, - title_align="left", - border_style="blue" - )) - + console.print( + Panel( + f"{info}\n\n[dim]{preview}[/dim]", + title=header, + title_align="left", + border_style="blue", + ) + ) + # Show scoring breakdown for top result if results: - console.print("\n[dim]Top result hybrid score: {:.3f} (70% semantic + 30% BM25)[/dim]".format(results[0].score)) + console.print( + "\n[dim]Top result hybrid score: {:.3f} (70% semantic + 30% BM25)[/dim]".format( + results[0].score + ) + ) def main(): @@ -122,14 +131,14 @@ def main(): else: # Use the RAG system itself as the demo project project_path = Path(__file__).parent - - if not (project_path / '.mini-rag').exists(): + + if not (project_path / ".mini-rag").exists(): console.print("[red]Error: No RAG index found. Run 'rag-mini index' first.[/red]") console.print(f"[dim]Looked in: {project_path / '.mini-rag'}[/dim]") return - + demo_search(project_path) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/tests/03_system_validation.py b/tests/03_system_validation.py index ea47134..83e64a5 100644 --- a/tests/03_system_validation.py +++ b/tests/03_system_validation.py @@ -2,53 +2,55 @@ Integration test to verify all three agents' work integrates properly. """ -import sys import os +import sys import tempfile from pathlib import Path # Fix Windows encoding -if sys.platform == 'win32': - os.environ['PYTHONUTF8'] = '1' - sys.stdout.reconfigure(encoding='utf-8') +if sys.platform == "win32": + os.environ["PYTHONUTF8"] = "1" + sys.stdout.reconfigure(encoding="utf-8") from mini_rag.chunker import CodeChunker +from mini_rag.config import RAGConfig from mini_rag.indexer import ProjectIndexer -from mini_rag.search import CodeSearcher from mini_rag.ollama_embeddings import OllamaEmbedder as CodeEmbedder from mini_rag.query_expander import QueryExpander -from mini_rag.config import RAGConfig +from mini_rag.search import CodeSearcher + def test_chunker(): """Test that chunker creates chunks with all required metadata.""" print("1. Testing Chunker...") - + # Create test Python file with more substantial content test_code = '''"""Test module for integration testing the chunker.""" import os import sys + class TestClass: """A test class with multiple methods.""" - + def __init__(self): """Initialize the test class.""" self.value = 42 self.name = "test" - + def method_one(self): """First method with some logic.""" result = self.value * 2 return result - + def method_two(self, x): """Second method that takes a parameter.""" if x > 0: return self.value + x else: return self.value - x - + def method_three(self): """Third method for testing.""" data = [] @@ -56,13 +58,14 @@ class TestClass: data.append(i * self.value) return data + class AnotherClass: """Another test class.""" - + def __init__(self, name): """Initialize with name.""" self.name = name - + def process(self): """Process something.""" return f"Processing {self.name}" @@ -72,22 +75,25 @@ def standalone_function(arg1, arg2): result = arg1 + arg2 return result * 2 + def another_function(): """Another standalone function.""" data = {"key": "value", "number": 123} return data ''' - + chunker = CodeChunker(min_chunk_size=1) # Use small chunk size for testing chunks = chunker.chunk_file(Path("test.py"), test_code) - + print(f" Created {len(chunks)} chunks") - + # Debug: Show what chunks were created print(" Chunks created:") for chunk in chunks: - print(f" - Type: {chunk.chunk_type}, Name: {chunk.name}, Lines: {chunk.start_line}-{chunk.end_line}") - + print( + f" - Type: {chunk.chunk_type}, Name: {chunk.name}, Lines: {chunk.start_line}-{chunk.end_line}" + ) + # Check metadata issues = [] for i, chunk in enumerate(chunks): @@ -97,68 +103,82 @@ def another_function(): issues.append(f"Chunk {i} missing total_chunks") if chunk.file_lines is None: issues.append(f"Chunk {i} missing file_lines") - + # Check links (except first/last) if i > 0 and chunk.prev_chunk_id is None: issues.append(f"Chunk {i} missing prev_chunk_id") if i < len(chunks) - 1 and chunk.next_chunk_id is None: issues.append(f"Chunk {i} missing next_chunk_id") - + # Check parent_class for methods - if chunk.chunk_type == 'method' and chunk.parent_class is None: + if chunk.chunk_type == "method" and chunk.parent_class is None: issues.append(f"Method chunk {chunk.name} missing parent_class") - - print(f" - Chunk {i}: {chunk.chunk_type} '{chunk.name}' " - f"[{chunk.chunk_index}/{chunk.total_chunks}] " - f"prev={chunk.prev_chunk_id} next={chunk.next_chunk_id}") - + + print( + f" - Chunk {i}: {chunk.chunk_type} '{chunk.name}' " + f"[{chunk.chunk_index}/{chunk.total_chunks}] " + f"prev={chunk.prev_chunk_id} next={chunk.next_chunk_id}" + ) + if issues: print(" Issues found:") for issue in issues: print(f" - {issue}") else: print(" All metadata present") - + return len(issues) == 0 + def test_indexer_storage(): """Test that indexer stores the new metadata.""" print("\n2. Testing Indexer Storage...") - + with tempfile.TemporaryDirectory() as tmpdir: project_path = Path(tmpdir) - + # Create test file test_file = project_path / "test.py" - test_file.write_text(''' + test_file.write_text( + """ + + class MyClass: + def my_method(self): return 42 -''') - +""" + ) + # Index the project with small chunk size for testing from mini_rag.chunker import CodeChunker + chunker = CodeChunker(min_chunk_size=1) indexer = ProjectIndexer(project_path, chunker=chunker) stats = indexer.index_project() - + print(f" Indexed {stats['chunks_created']} chunks") - + # Check what was stored if indexer.table: df = indexer.table.to_pandas() columns = df.columns.tolist() - - required_fields = ['chunk_id', 'prev_chunk_id', 'next_chunk_id', 'parent_class'] + + required_fields = [ + "chunk_id", + "prev_chunk_id", + "next_chunk_id", + "parent_class", + ] missing_fields = [f for f in required_fields if f not in columns] - + if missing_fields: print(f" Missing fields in database: {missing_fields}") print(f" Current fields: {columns}") return False else: print(" All required fields in database schema") - + # Check if data is actually stored sample = df.iloc[0] if len(df) > 0 else None if sample is not None: @@ -166,38 +186,41 @@ class MyClass: print(f" Sample prev_chunk_id: {sample.get('prev_chunk_id', 'MISSING')}") print(f" Sample next_chunk_id: {sample.get('next_chunk_id', 'MISSING')}") print(f" Sample parent_class: {sample.get('parent_class', 'MISSING')}") - + return len(missing_fields) == 0 + def test_search_integration(): """Test that search uses the new metadata.""" print("\n3. Testing Search Integration...") - + with tempfile.TemporaryDirectory() as tmpdir: project_path = Path(tmpdir) - + # Create test files with proper content that will create multiple chunks - (project_path / "math_utils.py").write_text('''"""Math utilities module.""" + (project_path / "math_utils.py").write_text( + '''"""Math utilities module.""" import math + class Calculator: """A simple calculator class.""" - + def __init__(self): """Initialize calculator.""" self.result = 0 - + def add(self, a, b): """Add two numbers.""" self.result = a + b return self.result - + def multiply(self, a, b): """Multiply two numbers.""" self.result = a * b return self.result - + def divide(self, a, b): """Divide two numbers.""" if b == 0: @@ -205,14 +228,15 @@ class Calculator: self.result = a / b return self.result + class AdvancedCalculator(Calculator): """Advanced calculator with more operations.""" - + def power(self, a, b): """Raise a to power b.""" self.result = a ** b return self.result - + def sqrt(self, a): """Calculate square root.""" self.result = math.sqrt(a) @@ -224,6 +248,7 @@ def compute_average(numbers): return 0 return sum(numbers) / len(numbers) + def compute_median(numbers): """Compute median of a list.""" if not numbers: @@ -233,20 +258,22 @@ def compute_median(numbers): if n % 2 == 0: return (sorted_nums[n//2-1] + sorted_nums[n//2]) / 2 return sorted_nums[n//2] -''') - +''' + ) + # Index with small chunk size for testing chunker = CodeChunker(min_chunk_size=1) indexer = ProjectIndexer(project_path, chunker=chunker) indexer.index_project() - + # Search searcher = CodeSearcher(project_path) - + # Test BM25 integration - results = searcher.search("multiply numbers", top_k=5, - semantic_weight=0.3, bm25_weight=0.7) - + results = searcher.search( + "multiply numbers", top_k=5, semantic_weight=0.3, bm25_weight=0.7 + ) + if results: print(f" BM25 + semantic search returned {len(results)} results") for r in results[:2]: @@ -254,45 +281,50 @@ def compute_median(numbers): else: print(" No search results returned") return False - + # Test context retrieval print("\n Testing context retrieval...") if searcher.table: df = searcher.table.to_pandas() print(f" Total chunks in DB: {len(df)}") - - # Find a method chunk to test parent context - method_chunks = df[df['chunk_type'] == 'method'] + + # Find a method/function chunk to test parent context + method_chunks = df[df["chunk_type"].isin(["method", "function"])] if len(method_chunks) > 0: - method_chunk_id = method_chunks.iloc[0]['chunk_id'] + method_chunk_id = method_chunks.iloc[0]["chunk_id"] context = searcher.get_chunk_context(method_chunk_id) - - if context['chunk']: + + if context["chunk"]: print(f" Got main chunk: {context['chunk'].name}") - if context['prev']: + if context["prev"]: print(f" Got previous chunk: {context['prev'].name}") else: - print(f" - No previous chunk (might be first)") - if context['next']: + print(" - No previous chunk (might be first)") + if context["next"]: print(f" Got next chunk: {context['next'].name}") else: - print(f" - No next chunk (might be last)") - if context['parent']: + print(" - No next chunk (might be last)") + if context["parent"]: print(f" Got parent chunk: {context['parent'].name}") else: - print(f" - No parent chunk") - + print(" - No parent chunk") + # Test include_context in search results_with_context = searcher.search("add", include_context=True, top_k=2) if results_with_context: print(f" Found {len(results_with_context)} results with context") for r in results_with_context: - has_context = bool(r.context_before or r.context_after or r.parent_chunk) - print(f" - {r.name}: context_before={bool(r.context_before)}, " - f"context_after={bool(r.context_after)}, parent={bool(r.parent_chunk)}") - + # Check if result has context (unused variable removed) + print( + f" - {r.name}: context_before={bool(r.context_before)}, " + f"context_after={bool(r.context_after)}, parent={bool(r.parent_chunk)}" + ) + # Check if at least one result has some context - if any(r.context_before or r.context_after or r.parent_chunk for r in results_with_context): + if any( + r.context_before or r.context_after or r.parent_chunk + for r in results_with_context + ): print(" Search with context working") return True else: @@ -304,112 +336,117 @@ def compute_median(numbers): else: print(" No method chunks found in database") return False - + return True + def test_server(): """Test that server still works.""" print("\n4. Testing Server...") - + # Just check if we can import and create server instance try: from mini_rag.server import RAGServer - server = RAGServer(Path("."), port=7778) + + # RAGServer(Path("."), port=7778) # Unused variable removed print(" Server can be instantiated") return True except Exception as e: print(f" Server error: {e}") return False + def test_new_features(): """Test new features: query expansion and smart ranking.""" print("\n5. Testing New Features (Query Expansion & Smart Ranking)...") - + try: # Test configuration loading config = RAGConfig() - print(f" โœ… Configuration loaded successfully") + print(" โœ… Configuration loaded successfully") print(f" Query expansion enabled: {config.search.expand_queries}") print(f" Max expansion terms: {config.llm.max_expansion_terms}") - + # Test query expander (will use mock if Ollama unavailable) expander = QueryExpander(config) test_query = "authentication" - + if expander.is_available(): expanded = expander.expand_query(test_query) print(f" โœ… Query expansion working: '{test_query}' โ†’ '{expanded}'") else: - print(f" โš ๏ธ Query expansion offline (Ollama not available)") + print(" โš ๏ธ Query expansion offline (Ollama not available)") # Test that it still returns original query expanded = expander.expand_query(test_query) if expanded == test_query: - print(f" โœ… Graceful degradation working: returns original query") + print(" โœ… Graceful degradation working: returns original query") else: - print(f" โŒ Error: should return original query when offline") + print(" โŒ Error: should return original query when offline") return False - + # Test smart ranking (this always works as it's zero-overhead) print(" ๐Ÿงฎ Testing smart ranking...") - + # Create a simple test to verify the method exists and can be called with tempfile.TemporaryDirectory() as temp_dir: temp_path = Path(temp_dir) - + # Create a simple test project test_file = temp_path / "README.md" test_file.write_text("# Test Project\nThis is a test README file.") - + try: searcher = CodeSearcher(temp_path) # Test that the _smart_rerank method exists - if hasattr(searcher, '_smart_rerank'): + if hasattr(searcher, "_smart_rerank"): print(" โœ… Smart ranking method available") return True else: print(" โŒ Smart ranking method not found") return False - + except Exception as e: print(f" โŒ Smart ranking test failed: {e}") return False - + except Exception as e: print(f" โŒ New features test failed: {e}") return False + def main(): """Run all integration tests.""" print("=" * 50) print("RAG System Integration Check") print("=" * 50) - + results = { "Chunker": test_chunker(), - "Indexer": test_indexer_storage(), + "Indexer": test_indexer_storage(), "Search": test_search_integration(), "Server": test_server(), - "New Features": test_new_features() + "New Features": test_new_features(), } - + print("\n" + "=" * 50) print("SUMMARY:") print("=" * 50) - + all_passed = True for component, passed in results.items(): status = " PASS" if passed else " FAIL" print(f"{component}: {status}") if not passed: all_passed = False - + if all_passed: print("\n All integration tests passed!") else: print("\n๏ธ Some tests failed - fixes needed!") - + return all_passed + if __name__ == "__main__": success = main() - sys.exit(0 if success else 1) \ No newline at end of file + sys.exit(0 if success else 1) diff --git a/tests/show_index_contents.py b/tests/show_index_contents.py index 177af07..75a4e74 100644 --- a/tests/show_index_contents.py +++ b/tests/show_index_contents.py @@ -3,19 +3,19 @@ Show what files are actually indexed in the RAG system. """ -import sys import os +import sys +from collections import Counter from pathlib import Path -if sys.platform == 'win32': - os.environ['PYTHONUTF8'] = '1' - sys.stdout.reconfigure(encoding='utf-8') +from mini_rag.vector_store import VectorStore + +if sys.platform == "win32": + os.environ["PYTHONUTF8"] = "1" + sys.stdout.reconfigure(encoding="utf-8") sys.path.insert(0, str(Path(__file__).parent)) -from mini_rag.vector_store import VectorStore -from collections import Counter - project_path = Path.cwd() store = VectorStore(project_path) store._connect() @@ -32,16 +32,16 @@ for row in store.table.to_pandas().itertuples(): unique_files = sorted(set(files)) -print(f"\n Indexed Files Summary") +print("\n Indexed Files Summary") print(f"Total files: {len(unique_files)}") print(f"Total chunks: {len(files)}") print(f"\nChunk types: {dict(chunk_types)}") -print(f"\n Files with most chunks:") +print("\n Files with most chunks:") for file, count in chunks_by_file.most_common(10): print(f" {count:3d} chunks: {file}") -print(f"\n Text-to-speech files:") -tts_files = [f for f in unique_files if 'text-to-speech' in f or 'speak' in f.lower()] +print("\n Text-to-speech files:") +tts_files = [f for f in unique_files if "text-to-speech" in f or "speak" in f.lower()] for f in tts_files: - print(f" - {f} ({chunks_by_file[f]} chunks)") \ No newline at end of file + print(f" - {f} ({chunks_by_file[f]} chunks)") diff --git a/tests/test_context_retrieval.py b/tests/test_context_retrieval.py index 5c1a6cd..588a983 100644 --- a/tests/test_context_retrieval.py +++ b/tests/test_context_retrieval.py @@ -12,30 +12,37 @@ Or run directly with venv: import os from pathlib import Path -from mini_rag.search import CodeSearcher + from mini_rag.ollama_embeddings import OllamaEmbedder as CodeEmbedder +from mini_rag.search import CodeSearcher # Check if virtual environment is activated + + def check_venv(): - if 'VIRTUAL_ENV' not in os.environ: + if "VIRTUAL_ENV" not in os.environ: print("โš ๏ธ WARNING: Virtual environment not detected!") print(" This test requires the virtual environment to be activated.") - print(" Run: source .venv/bin/activate && PYTHONPATH=. python tests/test_context_retrieval.py") + print( + " Run: source .venv/bin/activate && PYTHONPATH=. python tests/test_context_retrieval.py" + ) print(" Continuing anyway...\n") + check_venv() + def test_context_retrieval(): """Test the new context retrieval functionality.""" - + # Initialize searcher project_path = Path(__file__).parent try: embedder = CodeEmbedder() searcher = CodeSearcher(project_path, embedder) - + print("Testing search with context...") - + # Test 1: Search without context print("\n1. Search WITHOUT context:") results = searcher.search("chunk metadata", top_k=3, include_context=False) @@ -45,7 +52,7 @@ def test_context_retrieval(): print(f" Has context_before: {result.context_before is not None}") print(f" Has context_after: {result.context_after is not None}") print(f" Has parent_chunk: {result.parent_chunk is not None}") - + # Test 2: Search with context print("\n2. Search WITH context:") results = searcher.search("chunk metadata", top_k=3, include_context=True) @@ -55,39 +62,51 @@ def test_context_retrieval(): print(f" Has context_before: {result.context_before is not None}") print(f" Has context_after: {result.context_after is not None}") print(f" Has parent_chunk: {result.parent_chunk is not None}") - + if result.context_before: print(f" Context before preview: {result.context_before[:50]}...") if result.context_after: print(f" Context after preview: {result.context_after[:50]}...") if result.parent_chunk: - print(f" Parent chunk: {result.parent_chunk.name} ({result.parent_chunk.chunk_type})") - + print( + f" Parent chunk: {result.parent_chunk.name} ({result.parent_chunk.chunk_type})" + ) + # Test 3: get_chunk_context method print("\n3. Testing get_chunk_context method:") # Get a sample chunk_id from the first result df = searcher.table.to_pandas() if not df.empty: - sample_chunk_id = df.iloc[0]['chunk_id'] + sample_chunk_id = df.iloc[0]["chunk_id"] print(f" Getting context for chunk_id: {sample_chunk_id}") - + context = searcher.get_chunk_context(sample_chunk_id) - - if context['chunk']: - print(f" Main chunk: {context['chunk'].file_path}:{context['chunk'].start_line}") - if context['prev']: - print(f" Previous chunk: lines {context['prev'].start_line}-{context['prev'].end_line}") - if context['next']: - print(f" Next chunk: lines {context['next'].start_line}-{context['next'].end_line}") - if context['parent']: - print(f" Parent chunk: {context['parent'].name} ({context['parent'].chunk_type})") - + + if context["chunk"]: + print( + f" Main chunk: {context['chunk'].file_path}:{context['chunk'].start_line}" + ) + if context["prev"]: + print( + f" Previous chunk: lines {context['prev'].start_line}-{context['prev'].end_line}" + ) + if context["next"]: + print( + f" Next chunk: lines {context['next'].start_line}-{context['next'].end_line}" + ) + if context["parent"]: + print( + f" Parent chunk: {context['parent'].name} ({context['parent'].chunk_type})" + ) + print("\nAll tests completed successfully!") - + except Exception as e: print(f"Error during testing: {e}") import traceback + traceback.print_exc() + if __name__ == "__main__": - test_context_retrieval() \ No newline at end of file + test_context_retrieval() diff --git a/test_fixes.py b/tests/test_fixes.py similarity index 80% rename from test_fixes.py rename to tests/test_fixes.py index cdcbc3f..9d9877a 100644 --- a/test_fixes.py +++ b/tests/test_fixes.py @@ -10,55 +10,61 @@ Or run directly with venv: source .venv/bin/activate && python test_fixes.py """ -import sys import os +import sys import tempfile from pathlib import Path # Check if virtual environment is activated + + def check_venv(): - if 'VIRTUAL_ENV' not in os.environ: + if "VIRTUAL_ENV" not in os.environ: print("โš ๏ธ WARNING: Virtual environment not detected!") print(" This test requires the virtual environment to be activated.") print(" Run: source .venv/bin/activate && python test_fixes.py") print(" Continuing anyway...\n") + check_venv() # Add current directory to Python path -sys.path.insert(0, '.') +sys.path.insert(0, ".") + def test_config_model_rankings(): """Test that model rankings are properly configured.""" print("=" * 60) print("TESTING CONFIG AND MODEL RANKINGS") print("=" * 60) - + try: # Test config loading without heavy dependencies from mini_rag.config import ConfigManager, LLMConfig - + # Create a temporary directory for testing with tempfile.TemporaryDirectory() as tmpdir: config_manager = ConfigManager(tmpdir) config = config_manager.load_config() - + print("โœ“ Config loads successfully") - + # Check LLM config and model rankings - if hasattr(config, 'llm'): + if hasattr(config, "llm"): llm_config = config.llm print(f"โœ“ LLM config found: {type(llm_config)}") - - if hasattr(llm_config, 'model_rankings'): + + if hasattr(llm_config, "model_rankings"): rankings = llm_config.model_rankings print(f"โœ“ Model rankings: {rankings}") - + if rankings and rankings[0] == "qwen3:1.7b": print("โœ“ qwen3:1.7b is FIRST priority - CORRECT!") return True else: - print(f"โœ— WRONG: First model is {rankings[0] if rankings else 'None'}, should be qwen3:1.7b") + print( + f"โœ— WRONG: First model is {rankings[0] if rankings else 'None'}, should be qwen3:1.7b" + ) return False else: print("โœ— Model rankings not found in LLM config") @@ -66,7 +72,7 @@ def test_config_model_rankings(): else: print("โœ— LLM config not found") return False - + except ImportError as e: print(f"โœ— Import error: {e}") return False @@ -74,17 +80,18 @@ def test_config_model_rankings(): print(f"โœ— Error: {e}") return False + def test_context_length_fix(): """Test that context length is correctly set to 32K.""" print("\n" + "=" * 60) print("TESTING CONTEXT LENGTH FIXES") print("=" * 60) - + try: # Read the synthesizer file and check for 32000 - with open('mini_rag/llm_synthesizer.py', 'r') as f: + with open("mini_rag/llm_synthesizer.py", "r") as f: synthesizer_content = f.read() - + if '"num_ctx": 32000' in synthesizer_content: print("โœ“ LLM Synthesizer: num_ctx is correctly set to 32000") elif '"num_ctx": 80000' in synthesizer_content: @@ -92,133 +99,139 @@ def test_context_length_fix(): return False else: print("? LLM Synthesizer: num_ctx setting not found clearly") - + # Read the safeguards file and check for 32000 - with open('mini_rag/llm_safeguards.py', 'r') as f: + with open("mini_rag/llm_safeguards.py", "r") as f: safeguards_content = f.read() - - if 'context_window: int = 32000' in safeguards_content: + + if "context_window: int = 32000" in safeguards_content: print("โœ“ Safeguards: context_window is correctly set to 32000") return True - elif 'context_window: int = 80000' in safeguards_content: + elif "context_window: int = 80000" in safeguards_content: print("โœ— Safeguards: context_window is still 80000 - NEEDS FIX") return False else: print("? Safeguards: context_window setting not found clearly") return False - + except Exception as e: print(f"โœ— Error checking context length: {e}") return False + def test_safeguard_preservation(): """Test that safeguards preserve content instead of dropping it.""" print("\n" + "=" * 60) print("TESTING SAFEGUARD CONTENT PRESERVATION") print("=" * 60) - + try: # Read the synthesizer file and check for the preservation method - with open('mini_rag/llm_synthesizer.py', 'r') as f: + with open("mini_rag/llm_synthesizer.py", "r") as f: synthesizer_content = f.read() - - if '_create_safeguard_response_with_content' in synthesizer_content: + + if "_create_safeguard_response_with_content" in synthesizer_content: print("โœ“ Safeguard content preservation method exists") else: print("โœ— Safeguard content preservation method missing") return False - + # Check for the specific preservation logic - if 'AI Response (use with caution):' in synthesizer_content: + if "AI Response (use with caution):" in synthesizer_content: print("โœ“ Content preservation warning format found") else: print("โœ— Content preservation warning format missing") return False - + # Check that it's being called instead of dropping content - if 'return self._create_safeguard_response_with_content(issue_type, explanation, raw_response)' in synthesizer_content: + if ( + "return self._create_safeguard_response_with_content(issue_type, explanation, raw_response)" + in synthesizer_content + ): print("โœ“ Preservation method is called when safeguards trigger") return True else: print("โœ— Preservation method not called properly") return False - + except Exception as e: print(f"โœ— Error checking safeguard preservation: {e}") return False + def test_import_fixes(): """Test that import statements are fixed from claude_rag to mini_rag.""" print("\n" + "=" * 60) print("TESTING IMPORT STATEMENT FIXES") print("=" * 60) - + test_files = [ - 'tests/test_rag_integration.py', - 'tests/01_basic_integration_test.py', - 'tests/test_hybrid_search.py', - 'tests/test_context_retrieval.py' + "tests/test_rag_integration.py", + "tests/01_basic_integration_test.py", + "tests/test_hybrid_search.py", + "tests/test_context_retrieval.py", ] - + all_good = True - + for test_file in test_files: if Path(test_file).exists(): try: - with open(test_file, 'r') as f: + with open(test_file, "r") as f: content = f.read() - - if 'claude_rag' in content: + + if "claude_rag" in content: print(f"โœ— {test_file}: Still contains 'claude_rag' imports") all_good = False - elif 'mini_rag' in content: + elif "mini_rag" in content: print(f"โœ“ {test_file}: Uses correct 'mini_rag' imports") else: print(f"? {test_file}: No rag imports found") - + except Exception as e: print(f"โœ— Error reading {test_file}: {e}") all_good = False else: print(f"? {test_file}: File not found") - + return all_good + def main(): """Run all tests.""" print("FSS-Mini-RAG Fix Verification Tests") print("Testing all the critical fixes...") - + tests = [ ("Model Rankings", test_config_model_rankings), - ("Context Length", test_context_length_fix), + ("Context Length", test_context_length_fix), ("Safeguard Preservation", test_safeguard_preservation), - ("Import Fixes", test_import_fixes) + ("Import Fixes", test_import_fixes), ] - + results = {} - + for test_name, test_func in tests: try: results[test_name] = test_func() except Exception as e: print(f"โœ— {test_name} test crashed: {e}") results[test_name] = False - + # Summary print("\n" + "=" * 60) print("TEST SUMMARY") print("=" * 60) - + passed = sum(1 for result in results.values() if result) total = len(results) - + for test_name, result in results.items(): status = "โœ“ PASS" if result else "โœ— FAIL" print(f"{status} {test_name}") - + print(f"\nOverall: {passed}/{total} tests passed") - + if passed == total: print("๐ŸŽ‰ ALL TESTS PASSED - System should be working properly!") return 0 @@ -226,5 +239,6 @@ def main(): print("โŒ SOME TESTS FAILED - System needs more fixes!") return 1 + if __name__ == "__main__": - sys.exit(main()) \ No newline at end of file + sys.exit(main()) diff --git a/tests/test_hybrid_search.py b/tests/test_hybrid_search.py index 4b728e4..02cdd65 100644 --- a/tests/test_hybrid_search.py +++ b/tests/test_hybrid_search.py @@ -12,46 +12,49 @@ Or run directly with venv: """ import time -import json from pathlib import Path -from typing import List, Dict, Any -from rich.console import Console -from rich.table import Table -from rich.panel import Panel -from rich.columns import Columns -from rich.syntax import Syntax -from rich.progress import track +from typing import Any, Dict + +from rich.console import Console +from rich.progress import track +from rich.table import Table -from mini_rag.search import CodeSearcher, SearchResult from mini_rag.ollama_embeddings import OllamaEmbedder as CodeEmbedder +from mini_rag.search import CodeSearcher console = Console() class SearchTester: """Test harness for hybrid search evaluation.""" - + def __init__(self, project_path: Path): self.project_path = project_path console.print(f"\n[cyan]Initializing search system for: {project_path}[/cyan]") - + # Initialize searcher start = time.time() self.searcher = CodeSearcher(project_path) init_time = time.time() - start - + console.print(f"[green] Initialized in {init_time:.2f}s[/green]") - + # Get statistics stats = self.searcher.get_statistics() - if 'error' not in stats: - console.print(f"[dim]Index contains {stats['total_chunks']} chunks from {stats['unique_files']} files[/dim]\n") - - def run_query(self, query: str, top_k: int = 10, - semantic_only: bool = False, - bm25_only: bool = False) -> Dict[str, Any]: + if "error" not in stats: + console.print( + f"[dim]Index contains {stats['total_chunks']} chunks from {stats['unique_files']} files[/dim]\n" + ) + + def run_query( + self, + query: str, + top_k: int = 10, + semantic_only: bool = False, + bm25_only: bool = False, + ) -> Dict[str, Any]: """Run a single query and return metrics.""" - + # Set weights based on mode if semantic_only: semantic_weight, bm25_weight = 1.0, 0.0 @@ -62,150 +65,156 @@ class SearchTester: else: semantic_weight, bm25_weight = 0.7, 0.3 mode = "Hybrid (70/30)" - + # Run search start = time.time() results = self.searcher.search( query=query, top_k=top_k, semantic_weight=semantic_weight, - bm25_weight=bm25_weight + bm25_weight=bm25_weight, ) search_time = time.time() - start - + return { - 'query': query, - 'mode': mode, - 'results': results, - 'search_time_ms': search_time * 1000, - 'num_results': len(results), - 'top_score': results[0].score if results else 0, - 'avg_score': sum(r.score for r in results) / len(results) if results else 0, + "query": query, + "mode": mode, + "results": results, + "search_time_ms": search_time * 1000, + "num_results": len(results), + "top_score": results[0].score if results else 0, + "avg_score": sum(r.score for r in results) / len(results) if results else 0, } - + def compare_search_modes(self, query: str, top_k: int = 5): """Compare results across different search modes.""" console.print(f"\n[bold cyan]Query:[/bold cyan] '{query}'") console.print(f"[dim]Top {top_k} results per mode[/dim]\n") - + # Run searches in all modes modes = [ - ('hybrid', False, False), - ('semantic', True, False), - ('bm25', False, True) + ("hybrid", False, False), + ("semantic", True, False), + ("bm25", False, True), ] - + all_results = {} for mode_name, semantic_only, bm25_only in modes: result = self.run_query(query, top_k, semantic_only, bm25_only) all_results[mode_name] = result - + # Create comparison table table = Table(title="Search Mode Comparison") table.add_column("Metric", style="cyan", width=20) table.add_column("Hybrid (70/30)", style="green") table.add_column("Semantic Only", style="blue") table.add_column("BM25 Only", style="magenta") - + # Add metrics table.add_row( "Search Time (ms)", f"{all_results['hybrid']['search_time_ms']:.1f}", f"{all_results['semantic']['search_time_ms']:.1f}", - f"{all_results['bm25']['search_time_ms']:.1f}" + f"{all_results['bm25']['search_time_ms']:.1f}", ) - + table.add_row( "Results Found", - str(all_results['hybrid']['num_results']), - str(all_results['semantic']['num_results']), - str(all_results['bm25']['num_results']) + str(all_results["hybrid"]["num_results"]), + str(all_results["semantic"]["num_results"]), + str(all_results["bm25"]["num_results"]), ) - + table.add_row( "Top Score", f"{all_results['hybrid']['top_score']:.3f}", f"{all_results['semantic']['top_score']:.3f}", - f"{all_results['bm25']['top_score']:.3f}" + f"{all_results['bm25']['top_score']:.3f}", ) - + table.add_row( "Avg Score", f"{all_results['hybrid']['avg_score']:.3f}", f"{all_results['semantic']['avg_score']:.3f}", - f"{all_results['bm25']['avg_score']:.3f}" + f"{all_results['bm25']['avg_score']:.3f}", ) - + console.print(table) - + # Show top results from each mode console.print("\n[bold]Top Results by Mode:[/bold]") - + for mode_name, result_data in all_results.items(): console.print(f"\n[bold cyan]{result_data['mode']}:[/bold cyan]") - for i, result in enumerate(result_data['results'][:3], 1): - console.print(f"\n{i}. [green]{result.file_path}[/green]:{result.start_line}-{result.end_line}") - console.print(f" [dim]Type: {result.chunk_type} | Name: {result.name} | Score: {result.score:.3f}[/dim]") - + for i, result in enumerate(result_data["results"][:3], 1): + console.print( + f"\n{i}. [green]{result.file_path}[/green]:{result.start_line}-{result.end_line}" + ) + console.print( + f" [dim]Type: {result.chunk_type} | Name: {result.name} | Score: {result.score:.3f}[/dim]" + ) + # Show snippet lines = result.content.splitlines()[:5] for line in lines: - console.print(f" [dim]{line[:80]}{'...' if len(line) > 80 else ''}[/dim]") - + console.print( + f" [dim]{line[:80]}{'...' if len(line) > 80 else ''}[/dim]" + ) + def test_query_types(self): """Test different types of queries to show system capabilities.""" test_queries = [ # Keyword-heavy queries (should benefit from BM25) { - 'query': 'class CodeSearcher search method', - 'description': 'Specific class and method names', - 'expected': 'Should find exact matches with BM25 boost' + "query": "class CodeSearcher search method", + "description": "Specific class and method names", + "expected": "Should find exact matches with BM25 boost", }, { - 'query': 'import pandas numpy torch', - 'description': 'Multiple import keywords', - 'expected': 'BM25 should excel at finding import statements' + "query": "import pandas numpy torch", + "description": "Multiple import keywords", + "expected": "BM25 should excel at finding import statements", }, - # Semantic queries (should benefit from embeddings) { - 'query': 'find similar code chunks using vector similarity', - 'description': 'Natural language description', - 'expected': 'Semantic search should understand intent' + "query": "find similar code chunks using vector similarity", + "description": "Natural language description", + "expected": "Semantic search should understand intent", }, { - 'query': 'how to initialize database connection', - 'description': 'How-to question', - 'expected': 'Semantic search should find relevant implementations' + "query": "how to initialize database connection", + "description": "How-to question", + "expected": "Semantic search should find relevant implementations", }, - # Mixed queries (benefit from hybrid) { - 'query': 'BM25 scoring implementation for search ranking', - 'description': 'Technical terms + intent', - 'expected': 'Hybrid should balance keyword and semantic matching' + "query": "BM25 scoring implementation for search ranking", + "description": "Technical terms + intent", + "expected": "Hybrid should balance keyword and semantic matching", }, { - 'query': 'embedding vectors for code search with transformers', - 'description': 'Domain-specific terminology', - 'expected': 'Hybrid should leverage both approaches' - } + "query": "embedding vectors for code search with transformers", + "description": "Domain-specific terminology", + "expected": "Hybrid should leverage both approaches", + }, ] - + console.print("\n[bold yellow]Query Type Analysis[/bold yellow]") - console.print("[dim]Testing different query patterns to demonstrate hybrid search benefits[/dim]\n") - + console.print( + "[dim]Testing different query patterns to demonstrate hybrid search benefits[/dim]\n" + ) + for test_case in test_queries: console.rule(f"\n[cyan]{test_case['description']}[/cyan]") console.print(f"[dim]{test_case['expected']}[/dim]") - self.compare_search_modes(test_case['query'], top_k=3) + self.compare_search_modes(test_case["query"], top_k=3) time.sleep(0.5) # Brief pause between tests - + def benchmark_performance(self, num_queries: int = 50): """Run performance benchmarks.""" console.print("\n[bold yellow]Performance Benchmark[/bold yellow]") console.print(f"[dim]Running {num_queries} queries to measure performance[/dim]\n") - + # Sample queries for benchmarking benchmark_queries = [ "search function implementation", @@ -217,28 +226,28 @@ class SearchTester: "test cases unit testing", "configuration settings", "logging and debugging", - "performance optimization" + "performance optimization", ] * (num_queries // 10 + 1) - + benchmark_queries = benchmark_queries[:num_queries] - + # Benchmark each mode modes = [ - ('Hybrid (70/30)', 0.7, 0.3), - ('Semantic Only', 1.0, 0.0), - ('BM25 Only', 0.0, 1.0) + ("Hybrid (70/30)", 0.7, 0.3), + ("Semantic Only", 1.0, 0.0), + ("BM25 Only", 0.0, 1.0), ] - + results_table = Table(title="Performance Benchmark Results") results_table.add_column("Mode", style="cyan") results_table.add_column("Avg Time (ms)", style="green") results_table.add_column("Min Time (ms)", style="blue") results_table.add_column("Max Time (ms)", style="red") results_table.add_column("Total Time (s)", style="magenta") - + for mode_name, sem_weight, bm25_weight in modes: times = [] - + console.print(f"[cyan]Testing {mode_name}...[/cyan]") for query in track(benchmark_queries, description=f"Running {mode_name}"): start = time.time() @@ -246,69 +255,75 @@ class SearchTester: query=query, limit=10, semantic_weight=sem_weight, - bm25_weight=bm25_weight + bm25_weight=bm25_weight, ) elapsed = (time.time() - start) * 1000 times.append(elapsed) - + # Calculate statistics avg_time = sum(times) / len(times) min_time = min(times) max_time = max(times) total_time = sum(times) / 1000 - + results_table.add_row( mode_name, f"{avg_time:.2f}", f"{min_time:.2f}", f"{max_time:.2f}", - f"{total_time:.2f}" + f"{total_time:.2f}", ) - + console.print("\n") console.print(results_table) - + def test_diversity_constraints(self): """Test the diversity constraints in search results.""" console.print("\n[bold yellow]Diversity Constraints Test[/bold yellow]") console.print("[dim]Verifying max 2 chunks per file and chunk type diversity[/dim]\n") - + # Query that might return many results from same files query = "function implementation code search" results = self.searcher.search(query, top_k=20) - + # Analyze diversity file_counts = {} chunk_types = {} - + for result in results: file_counts[result.file_path] = file_counts.get(result.file_path, 0) + 1 chunk_types[result.chunk_type] = chunk_types.get(result.chunk_type, 0) + 1 - + # Create diversity report table = Table(title="Result Diversity Analysis") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") - + table.add_row("Total Results", str(len(results))) table.add_row("Unique Files", str(len(file_counts))) - table.add_row("Max Chunks per File", str(max(file_counts.values()) if file_counts else 0)) + table.add_row( + "Max Chunks per File", str(max(file_counts.values()) if file_counts else 0) + ) table.add_row("Unique Chunk Types", str(len(chunk_types))) - + console.print(table) - + # Show file distribution if len(file_counts) > 0: console.print("\n[bold]File Distribution:[/bold]") - for file_path, count in sorted(file_counts.items(), key=lambda x: x[1], reverse=True)[:5]: + for file_path, count in sorted( + file_counts.items(), key=lambda x: x[1], reverse=True + )[:5]: console.print(f" {count}x {file_path}") - + # Show chunk type distribution if len(chunk_types) > 0: console.print("\n[bold]Chunk Type Distribution:[/bold]") - for chunk_type, count in sorted(chunk_types.items(), key=lambda x: x[1], reverse=True): + for chunk_type, count in sorted( + chunk_types.items(), key=lambda x: x[1], reverse=True + ): console.print(f" {chunk_type}: {count} chunks") - + # Verify constraints console.print("\n[bold]Constraint Verification:[/bold]") max_per_file = max(file_counts.values()) if file_counts else 0 @@ -321,45 +336,45 @@ class SearchTester: def main(): """Run comprehensive hybrid search tests.""" import sys - + if len(sys.argv) > 1: project_path = Path(sys.argv[1]) else: project_path = Path.cwd() - - if not (project_path / '.mini-rag').exists(): + + if not (project_path / ".mini-rag").exists(): console.print("[red]Error: No RAG index found. Run 'rag-mini index' first.[/red]") return - + # Create tester tester = SearchTester(project_path) - + # Run all tests - console.print("\n" + "="*80) + console.print("\n" + "=" * 80) console.print("[bold green]Mini RAG Hybrid Search Test Suite[/bold green]") - console.print("="*80) - + console.print("=" * 80) + # Test 1: Query type analysis tester.test_query_types() - + # Test 2: Performance benchmark - console.print("\n" + "-"*80) + console.print("\n" + "-" * 80) tester.benchmark_performance(num_queries=30) - + # Test 3: Diversity constraints - console.print("\n" + "-"*80) + console.print("\n" + "-" * 80) tester.test_diversity_constraints() - + # Summary - console.print("\n" + "="*80) + console.print("\n" + "=" * 80) console.print("[bold green]Test Suite Complete![/bold green]") console.print("\n[dim]The hybrid search combines:") console.print(" โ€ข Semantic understanding from transformer embeddings") console.print(" โ€ข Keyword relevance from BM25 scoring") console.print(" โ€ข Result diversity through intelligent filtering") console.print(" โ€ข Performance optimization through concurrent processing[/dim]") - console.print("="*80 + "\n") + console.print("=" * 80 + "\n") if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/tests/test_min_chunk_size.py b/tests/test_min_chunk_size.py index 64d08ef..a2f2a32 100644 --- a/tests/test_min_chunk_size.py +++ b/tests/test_min_chunk_size.py @@ -1,13 +1,16 @@ """Test with smaller min_chunk_size.""" -from mini_rag.chunker import CodeChunker from pathlib import Path +from mini_rag.chunker import CodeChunker + test_code = '''"""Test module.""" import os + class MyClass: + def method(self): return 42 @@ -24,4 +27,4 @@ for i, chunk in enumerate(chunks): print(f"\nChunk {i}: {chunk.chunk_type} '{chunk.name}'") print(f"Lines {chunk.start_line}-{chunk.end_line}") print(f"Size: {len(chunk.content.splitlines())} lines") - print("-" * 40) \ No newline at end of file + print("-" * 40) diff --git a/tests/test_mode_separation.py b/tests/test_mode_separation.py index a5c90dc..7ee1114 100644 --- a/tests/test_mode_separation.py +++ b/tests/test_mode_separation.py @@ -7,7 +7,6 @@ between thinking and no-thinking modes. """ import sys -import os import tempfile import unittest from pathlib import Path @@ -16,51 +15,54 @@ from pathlib import Path sys.path.insert(0, str(Path(__file__).parent.parent)) try: - from mini_rag.llm_synthesizer import LLMSynthesizer - from mini_rag.explorer import CodeExplorer from mini_rag.config import RAGConfig + from mini_rag.explorer import CodeExplorer from mini_rag.indexer import ProjectIndexer + from mini_rag.llm_synthesizer import LLMSynthesizer from mini_rag.search import CodeSearcher except ImportError as e: print(f"โŒ Could not import RAG components: {e}") print(" This test requires the full RAG system to be installed") sys.exit(1) + class TestModeSeparation(unittest.TestCase): """Test the clean separation between synthesis and exploration modes.""" - + def setUp(self): """Set up test environment.""" self.temp_dir = tempfile.mkdtemp() self.project_path = Path(self.temp_dir) - + # Create a simple test project test_file = self.project_path / "test_module.py" - test_file.write_text('''"""Test module for mode separation testing.""" + test_file.write_text( + '''"""Test module for mode separation testing.""" def authenticate_user(username: str, password: str) -> bool: """Authenticate a user with username and password.""" # Simple authentication logic if not username or not password: return False - + # Check against database (simplified) valid_users = {"admin": "secret", "user": "password"} return valid_users.get(username) == password + class UserManager: """Manages user operations.""" - + def __init__(self): self.users = {} - + def create_user(self, username: str) -> bool: """Create a new user.""" if username in self.users: return False self.users[username] = {"created": True} return True - + def get_user_info(self, username: str) -> dict: """Get user information.""" return self.users.get(username, {}) @@ -71,196 +73,216 @@ def process_login_request(username: str, password: str) -> dict: return {"success": True, "message": "Login successful"} else: return {"success": False, "message": "Invalid credentials"} -''') - +''' + ) + # Index the project for testing try: indexer = ProjectIndexer(self.project_path) indexer.index_project() except Exception as e: self.skipTest(f"Could not index test project: {e}") - + def tearDown(self): """Clean up test environment.""" import shutil + shutil.rmtree(self.temp_dir, ignore_errors=True) - + def test_01_synthesis_mode_defaults(self): """Test that synthesis mode has correct defaults.""" synthesizer = LLMSynthesizer() - + # Should default to no thinking - self.assertFalse(synthesizer.enable_thinking, - "Synthesis mode should default to no thinking") - + self.assertFalse( + synthesizer.enable_thinking, "Synthesis mode should default to no thinking" + ) + print("โœ… Synthesis mode defaults to no thinking") - + def test_02_exploration_mode_defaults(self): """Test that exploration mode enables thinking.""" config = RAGConfig() explorer = CodeExplorer(self.project_path, config) - + # Should enable thinking in exploration mode - self.assertTrue(explorer.synthesizer.enable_thinking, - "Exploration mode should enable thinking") - + self.assertTrue( + explorer.synthesizer.enable_thinking, + "Exploration mode should enable thinking", + ) + print("โœ… Exploration mode enables thinking by default") - + def test_03_no_runtime_thinking_toggle(self): """Test that thinking mode cannot be toggled at runtime.""" synthesizer = LLMSynthesizer(enable_thinking=False) - + # Should not have public methods to toggle thinking - thinking_methods = [method for method in dir(synthesizer) - if 'thinking' in method.lower() and not method.startswith('_')] - + thinking_methods = [ + method + for method in dir(synthesizer) + if "thinking" in method.lower() and not method.startswith("_") + ] + # The only thinking-related attribute should be the readonly enable_thinking - self.assertEqual(len(thinking_methods), 0, - "Should not have public thinking toggle methods") - + self.assertEqual( + len(thinking_methods), 0, "Should not have public thinking toggle methods" + ) + print("โœ… No runtime thinking toggle methods available") - + def test_04_mode_contamination_prevention(self): """Test that modes don't contaminate each other.""" if not self._ollama_available(): self.skipTest("Ollama not available for contamination testing") - + # Create synthesis mode synthesizer synthesis_synthesizer = LLMSynthesizer(enable_thinking=False) - - # Create exploration mode synthesizer + + # Create exploration mode synthesizer exploration_synthesizer = LLMSynthesizer(enable_thinking=True) - + # Both should maintain their thinking settings - self.assertFalse(synthesis_synthesizer.enable_thinking, - "Synthesis synthesizer should remain no-thinking") - self.assertTrue(exploration_synthesizer.enable_thinking, - "Exploration synthesizer should remain thinking-enabled") - + self.assertFalse( + synthesis_synthesizer.enable_thinking, + "Synthesis synthesizer should remain no-thinking", + ) + self.assertTrue( + exploration_synthesizer.enable_thinking, + "Exploration synthesizer should remain thinking-enabled", + ) + print("โœ… Mode contamination prevented") - + def test_05_exploration_session_management(self): """Test exploration session management.""" config = RAGConfig() explorer = CodeExplorer(self.project_path, config) - + # Should start with no active session - self.assertIsNone(explorer.current_session, - "Should start with no active session") - + self.assertIsNone(explorer.current_session, "Should start with no active session") + # Should be able to create session summary even without session summary = explorer.get_session_summary() - self.assertIn("No active", summary, - "Should handle no active session gracefully") - + self.assertIn("No active", summary, "Should handle no active session gracefully") + print("โœ… Session management working correctly") - + def test_06_context_memory_structure(self): """Test that exploration mode has context memory structure.""" config = RAGConfig() explorer = CodeExplorer(self.project_path, config) - + # Should have context tracking attributes - self.assertTrue(hasattr(explorer, 'current_session'), - "Explorer should have session tracking") - + self.assertTrue( + hasattr(explorer, "current_session"), + "Explorer should have session tracking", + ) + print("โœ… Context memory structure present") - + def test_07_synthesis_mode_no_thinking_prompts(self): """Test that synthesis mode properly handles no-thinking.""" if not self._ollama_available(): self.skipTest("Ollama not available for prompt testing") - + synthesizer = LLMSynthesizer(enable_thinking=False) - + # Test the _call_ollama method handling - if hasattr(synthesizer, '_call_ollama'): + if hasattr(synthesizer, "_call_ollama"): # Should append when thinking disabled # This is a white-box test of the implementation try: # Mock test - just verify the method exists and can be called - result = synthesizer._call_ollama("test", temperature=0.1, disable_thinking=True) + # Test call (result unused) + synthesizer._call_ollama("test", temperature=0.1, disable_thinking=True) # Don't assert on result since Ollama might not be available print("โœ… No-thinking prompt handling available") except Exception as e: print(f"โš ๏ธ Prompt handling test skipped: {e}") else: self.fail("Synthesizer should have _call_ollama method") - + def test_08_mode_specific_initialization(self): """Test that modes initialize correctly with lazy loading.""" # Synthesis mode synthesis_synthesizer = LLMSynthesizer(enable_thinking=False) - self.assertFalse(synthesis_synthesizer._initialized, - "Should start uninitialized for lazy loading") - - # Exploration mode + self.assertFalse( + synthesis_synthesizer._initialized, + "Should start uninitialized for lazy loading", + ) + + # Exploration mode config = RAGConfig() explorer = CodeExplorer(self.project_path, config) - self.assertFalse(explorer.synthesizer._initialized, - "Should start uninitialized for lazy loading") - + self.assertFalse( + explorer.synthesizer._initialized, + "Should start uninitialized for lazy loading", + ) + print("โœ… Lazy initialization working correctly") - + def test_09_search_vs_exploration_integration(self): """Test integration differences between search and exploration.""" # Regular search (synthesis mode) searcher = CodeSearcher(self.project_path) search_results = searcher.search("authentication", top_k=3) - - self.assertGreater(len(search_results), 0, - "Search should return results") - + + self.assertGreater(len(search_results), 0, "Search should return results") + # Exploration mode setup config = RAGConfig() explorer = CodeExplorer(self.project_path, config) - + # Both should work with same project but different approaches - self.assertTrue(hasattr(explorer, 'synthesizer'), - "Explorer should have thinking-enabled synthesizer") - + self.assertTrue( + hasattr(explorer, "synthesizer"), + "Explorer should have thinking-enabled synthesizer", + ) + print("โœ… Search and exploration integration working") - + def test_10_mode_guidance_detection(self): """Test that the system can detect when to recommend different modes.""" # Words that should trigger exploration mode recommendation - exploration_triggers = ['why', 'how', 'explain', 'debug'] - + exploration_triggers = ["why", "how", "explain", "debug"] + for trigger in exploration_triggers: query = f"{trigger} does authentication work" # This would typically be tested in the main CLI # Here we just verify the trigger detection logic exists has_trigger = any(word in query.lower() for word in exploration_triggers) - self.assertTrue(has_trigger, - f"Should detect '{trigger}' as exploration trigger") - + self.assertTrue(has_trigger, f"Should detect '{trigger}' as exploration trigger") + print("โœ… Mode guidance detection working") - + def _ollama_available(self) -> bool: """Check if Ollama is available for testing.""" try: import requests + response = requests.get("http://localhost:11434/api/tags", timeout=5) return response.status_code == 200 except Exception: return False + def main(): """Run mode separation tests.""" print("๐Ÿงช Testing Mode Separation") print("=" * 40) - + # Check if we're in the right environment if not Path("mini_rag").exists(): print("โŒ Tests must be run from the FSS-Mini-RAG root directory") sys.exit(1) - + # Run tests loader = unittest.TestLoader() suite = loader.loadTestsFromTestCase(TestModeSeparation) runner = unittest.TextTestRunner(verbosity=2) result = runner.run(suite) - + # Summary print("\n" + "=" * 40) if result.wasSuccessful(): @@ -269,9 +291,10 @@ def main(): else: print("โŒ Some tests failed") print(f" Failed: {len(result.failures)}, Errors: {len(result.errors)}") - + return result.wasSuccessful() + if __name__ == "__main__": success = main() - sys.exit(0 if success else 1) \ No newline at end of file + sys.exit(0 if success else 1) diff --git a/tests/test_ollama_integration.py b/tests/test_ollama_integration.py index 65673bf..b7227f1 100755 --- a/tests/test_ollama_integration.py +++ b/tests/test_ollama_integration.py @@ -8,72 +8,71 @@ what's working and what needs attention. Run with: python3 tests/test_ollama_integration.py """ -import unittest -import requests -import json import sys +import unittest from pathlib import Path -from unittest.mock import patch, MagicMock +from unittest.mock import MagicMock, patch + +import requests + +from mini_rag.config import RAGConfig +from mini_rag.llm_synthesizer import LLMSynthesizer +from mini_rag.query_expander import QueryExpander # Add project to path sys.path.insert(0, str(Path(__file__).parent.parent)) -from mini_rag.query_expander import QueryExpander -from mini_rag.llm_synthesizer import LLMSynthesizer -from mini_rag.config import RAGConfig - class TestOllamaIntegration(unittest.TestCase): """ Tests to help beginners troubleshoot their Ollama setup. - + Each test explains what it's checking and gives clear feedback about what's working or needs to be fixed. """ - + def setUp(self): """Set up test configuration.""" self.config = RAGConfig() print(f"\n๐Ÿงช Testing with Ollama host: {self.config.llm.ollama_host}") - + def test_01_ollama_server_running(self): """ โœ… Check if Ollama server is running and responding. - + This test verifies that: - Ollama is installed and running - The API endpoint is accessible - Basic connectivity works """ print("\n๐Ÿ“ก Testing Ollama server connectivity...") - + try: response = requests.get( - f"http://{self.config.llm.ollama_host}/api/tags", - timeout=5 + f"http://{self.config.llm.ollama_host}/api/tags", timeout=5 ) - + if response.status_code == 200: data = response.json() - models = data.get('models', []) - print(f" โœ… Ollama server is running!") + models = data.get("models", []) + print(" โœ… Ollama server is running!") print(f" ๐Ÿ“ฆ Found {len(models)} models available") - + if models: print(" ๐ŸŽฏ Available models:") for model in models[:5]: # Show first 5 - name = model.get('name', 'unknown') - size = model.get('size', 0) + name = model.get("name", "unknown") + size = model.get("size", 0) print(f" โ€ข {name} ({size//1000000:.0f}MB)") if len(models) > 5: print(f" ... and {len(models)-5} more") else: print(" โš ๏ธ No models found. Install with: ollama pull qwen3:4b") - + self.assertTrue(True) else: self.fail(f"Ollama server responded with status {response.status_code}") - + except requests.exceptions.ConnectionError: self.fail( "โŒ Cannot connect to Ollama server.\n" @@ -84,35 +83,32 @@ class TestOllamaIntegration(unittest.TestCase): ) except Exception as e: self.fail(f"โŒ Unexpected error: {e}") - + def test_02_embedding_model_available(self): """ โœ… Check if embedding model is available. - + This test verifies that: - The embedding model (nomic-embed-text) is installed - Embedding API calls work correctly - Model responds with valid embeddings """ print("\n๐Ÿง  Testing embedding model availability...") - + try: # Test embedding generation response = requests.post( f"http://{self.config.llm.ollama_host}/api/embeddings", - json={ - "model": "nomic-embed-text", - "prompt": "test embedding" - }, - timeout=10 + json={"model": "nomic-embed-text", "prompt": "test embedding"}, + timeout=10, ) - + if response.status_code == 200: data = response.json() - embedding = data.get('embedding', []) - + embedding = data.get("embedding", []) + if embedding and len(embedding) > 0: - print(f" โœ… Embedding model working!") + print(" โœ… Embedding model working!") print(f" ๐Ÿ“Š Generated {len(embedding)}-dimensional vectors") self.assertTrue(len(embedding) > 100) # Should be substantial vectors else: @@ -126,285 +122,283 @@ class TestOllamaIntegration(unittest.TestCase): ) else: self.fail(f"Embedding API error: {response.status_code}") - + except Exception as e: self.fail(f"โŒ Embedding test failed: {e}") - + def test_03_llm_model_available(self): """ โœ… Check if LLM models are available for synthesis/expansion. - + This test verifies that: - At least one LLM model is available - The model can generate text responses - Response quality is reasonable """ print("\n๐Ÿค– Testing LLM model availability...") - + synthesizer = LLMSynthesizer(config=self.config) - + if not synthesizer.is_available(): self.fail( "โŒ No LLM models available.\n" " ๐Ÿ’ก Install a model like: ollama pull qwen3:4b" ) - + print(f" โœ… Found {len(synthesizer.available_models)} LLM models") print(f" ๐ŸŽฏ Will use: {synthesizer.model}") - + # Test basic text generation try: response = synthesizer._call_ollama( - "Complete this: The capital of France is", - temperature=0.1 + "Complete this: The capital of France is", temperature=0.1 ) - + if response and len(response.strip()) > 0: - print(f" โœ… Model generating responses!") + print(" โœ… Model generating responses!") print(f" ๐Ÿ’ฌ Sample response: '{response[:50]}...'") - + # Basic quality check if "paris" in response.lower(): print(" ๐ŸŽฏ Response quality looks good!") else: print(" โš ๏ธ Response quality might be low") - + self.assertTrue(len(response) > 5) else: self.fail("Model produced empty response") - + except Exception as e: self.fail(f"โŒ LLM generation test failed: {e}") - + def test_04_query_expansion_working(self): """ โœ… Check if query expansion is working correctly. - + This test verifies that: - QueryExpander can connect to Ollama - Expansion produces reasonable results - Caching is working """ print("\n๐Ÿ” Testing query expansion...") - + # Enable expansion for testing self.config.search.expand_queries = True expander = QueryExpander(self.config) - + if not expander.is_available(): self.skipTest("โญ๏ธ Skipping - Ollama not available (tested above)") - + # Test expansion test_query = "authentication" expanded = expander.expand_query(test_query) - + print(f" ๐Ÿ“ Original: '{test_query}'") print(f" โžก๏ธ Expanded: '{expanded}'") - + # Quality checks if expanded == test_query: print(" โš ๏ธ No expansion occurred (might be normal for simple queries)") else: # Should contain original query self.assertIn(test_query.lower(), expanded.lower()) - + # Should be longer self.assertGreater(len(expanded.split()), len(test_query.split())) - + # Test caching cached = expander.expand_query(test_query) self.assertEqual(expanded, cached) print(" โœ… Expansion and caching working!") - + def test_05_synthesis_mode_no_thinking(self): """ โœ… Test synthesis mode operates without thinking. - + Verifies that LLMSynthesizer in synthesis mode: - Defaults to no thinking - Handles tokens properly - Works independently of exploration mode """ print("\n๐Ÿš€ Testing synthesis mode (no thinking)...") - + # Create synthesis mode synthesizer (default behavior) synthesizer = LLMSynthesizer() - + # Should default to no thinking - self.assertFalse(synthesizer.enable_thinking, - "Synthesis mode should default to no thinking") + self.assertFalse( + synthesizer.enable_thinking, "Synthesis mode should default to no thinking" + ) print(" โœ… Defaults to no thinking") - + if synthesizer.is_available(): print(" ๐Ÿ“ Testing with live Ollama...") - + # Create mock search results from dataclasses import dataclass - + @dataclass class MockResult: file_path: str content: str score: float - - results = [ - MockResult("auth.py", "def authenticate(user): return True", 0.95) - ] - - # Test synthesis + + results = [MockResult("auth.py", "def authenticate(user): return True", 0.95)] + + # Test synthesis synthesis = synthesizer.synthesize_search_results( "user authentication", results, Path(".") ) - + # Should get reasonable synthesis self.assertIsNotNone(synthesis) self.assertGreater(len(synthesis.summary), 10) print(" โœ… Synthesis mode working without thinking") else: print(" โญ๏ธ Live test skipped - Ollama not available") - + def test_06_exploration_mode_thinking(self): """ โœ… Test exploration mode enables thinking. - + Verifies that CodeExplorer: - Enables thinking by default - Has session management - Works independently of synthesis mode """ print("\n๐Ÿง  Testing exploration mode (with thinking)...") - + try: from mini_rag.explorer import CodeExplorer except ImportError: self.skipTest("โญ๏ธ CodeExplorer not available") - + # Create exploration mode explorer = CodeExplorer(Path("."), self.config) - + # Should enable thinking - self.assertTrue(explorer.synthesizer.enable_thinking, - "Exploration mode should enable thinking") + self.assertTrue( + explorer.synthesizer.enable_thinking, + "Exploration mode should enable thinking", + ) print(" โœ… Enables thinking by default") - + # Should have session management - self.assertIsNone(explorer.current_session, - "Should start with no active session") + self.assertIsNone(explorer.current_session, "Should start with no active session") print(" โœ… Session management available") - + # Should handle session summary gracefully summary = explorer.get_session_summary() self.assertIn("No active", summary) print(" โœ… Graceful session handling") - + def test_07_mode_separation(self): """ โœ… Test that synthesis and exploration modes don't interfere. - + Verifies clean separation: - Different thinking settings - Independent operation - No cross-contamination """ print("\n๐Ÿ”„ Testing mode separation...") - + # Create both modes synthesizer = LLMSynthesizer(enable_thinking=False) - + try: from mini_rag.explorer import CodeExplorer + explorer = CodeExplorer(Path("."), self.config) except ImportError: self.skipTest("โญ๏ธ CodeExplorer not available") - + # Should have different thinking settings - self.assertFalse(synthesizer.enable_thinking, - "Synthesis should not use thinking") - self.assertTrue(explorer.synthesizer.enable_thinking, - "Exploration should use thinking") - + self.assertFalse(synthesizer.enable_thinking, "Synthesis should not use thinking") + self.assertTrue( + explorer.synthesizer.enable_thinking, "Exploration should use thinking" + ) + # Both should be uninitialized (lazy loading) - self.assertFalse(synthesizer._initialized, - "Should use lazy loading") - self.assertFalse(explorer.synthesizer._initialized, - "Should use lazy loading") - + self.assertFalse(synthesizer._initialized, "Should use lazy loading") + self.assertFalse(explorer.synthesizer._initialized, "Should use lazy loading") + print(" โœ… Clean mode separation confirmed") - + def test_08_with_mocked_ollama(self): """ โœ… Test components work with mocked Ollama (for offline testing). - + This test verifies that: - System gracefully handles Ollama being unavailable - Fallback behaviors work correctly - Error messages are helpful """ print("\n๐ŸŽญ Testing with mocked Ollama responses...") - + # Mock successful embedding response mock_embedding_response = MagicMock() mock_embedding_response.status_code = 200 mock_embedding_response.json.return_value = { - 'embedding': [0.1] * 768 # Standard embedding size + "embedding": [0.1] * 768 # Standard embedding size } - + # Mock LLM response mock_llm_response = MagicMock() mock_llm_response.status_code = 200 mock_llm_response.json.return_value = { - 'response': 'authentication login user verification credentials' + "response": "authentication login user verification credentials" } - - with patch('requests.post', side_effect=[mock_embedding_response, mock_llm_response]): + + with patch("requests.post", side_effect=[mock_embedding_response, mock_llm_response]): # Test query expansion with mocked response expander = QueryExpander(self.config) expander.enabled = True - + expanded = expander._llm_expand_query("authentication") if expanded: print(f" โœ… Mocked expansion: '{expanded}'") self.assertIn("authentication", expanded) else: print(" โš ๏ธ Expansion returned None (might be expected)") - + # Test graceful degradation when Ollama unavailable - with patch('requests.get', side_effect=requests.exceptions.ConnectionError()): + with patch("requests.get", side_effect=requests.exceptions.ConnectionError()): expander_offline = QueryExpander(self.config) - + # Should handle unavailable server gracefully self.assertFalse(expander_offline.is_available()) - + # Should return original query when offline result = expander_offline.expand_query("test query") self.assertEqual(result, "test query") print(" โœ… Graceful offline behavior working!") - + def test_06_configuration_validation(self): """ โœ… Check if configuration is valid and complete. - + This test verifies that: - All required config sections exist - Values are reasonable - Host/port settings are valid """ print("\nโš™๏ธ Testing configuration validation...") - + # Check LLM config self.assertIsNotNone(self.config.llm) self.assertTrue(self.config.llm.ollama_host) self.assertTrue(isinstance(self.config.llm.max_expansion_terms, int)) self.assertGreater(self.config.llm.max_expansion_terms, 0) - - print(f" โœ… LLM config valid") + + print(" โœ… LLM config valid") print(f" Host: {self.config.llm.ollama_host}") print(f" Max expansion terms: {self.config.llm.max_expansion_terms}") - - # Check search config + + # Check search config self.assertIsNotNone(self.config.search) self.assertGreater(self.config.search.default_top_k, 0) - print(f" โœ… Search config valid") + print(" โœ… Search config valid") print(f" Default top-k: {self.config.search.default_top_k}") print(f" Query expansion: {self.config.search.expand_queries}") @@ -418,10 +412,10 @@ def run_troubleshooting(): print("These tests help you troubleshoot your Ollama setup.") print("Each test explains what it's checking and how to fix issues.") print() - + # Run tests with detailed output unittest.main(verbosity=2, exit=False) - + print("\n" + "=" * 50) print("๐Ÿ’ก Common Solutions:") print(" โ€ข Install Ollama: https://ollama.ai/download") @@ -432,5 +426,5 @@ def run_troubleshooting(): print("๐Ÿ“š For more help, see docs/QUERY_EXPANSION.md") -if __name__ == '__main__': - run_troubleshooting() \ No newline at end of file +if __name__ == "__main__": + run_troubleshooting() diff --git a/tests/test_rag_integration.py b/tests/test_rag_integration.py index 00313e8..c0ac74f 100644 --- a/tests/test_rag_integration.py +++ b/tests/test_rag_integration.py @@ -10,21 +10,26 @@ Or run directly with venv: source .venv/bin/activate && PYTHONPATH=. python tests/test_rag_integration.py """ -import tempfile -import shutil import os +import tempfile from pathlib import Path + from mini_rag.indexer import ProjectIndexer from mini_rag.search import CodeSearcher # Check if virtual environment is activated + + def check_venv(): - if 'VIRTUAL_ENV' not in os.environ: + if "VIRTUAL_ENV" not in os.environ: print("โš ๏ธ WARNING: Virtual environment not detected!") print(" This test requires the virtual environment to be activated.") - print(" Run: source .venv/bin/activate && PYTHONPATH=. python tests/test_rag_integration.py") + print( + " Run: source .venv/bin/activate && PYTHONPATH=. python tests/test_rag_integration.py" + ) print(" Continuing anyway...\n") + check_venv() # Sample Python file with proper structure @@ -35,15 +40,16 @@ This module demonstrates various Python constructs. import os import sys -from typing import List, Dict, Optional +from typing import List, Optional from dataclasses import dataclass # Module-level constants DEFAULT_TIMEOUT = 30 MAX_RETRIES = 3 - @dataclass + + class Config: """Configuration dataclass.""" timeout: int = DEFAULT_TIMEOUT @@ -53,73 +59,71 @@ class Config: class DataProcessor: """ Main data processor class. - + This class handles the processing of various data types and provides a unified interface for data operations. """ - + def __init__(self, config: Config): """ Initialize the processor with configuration. - + Args: config: Configuration object """ self.config = config self._cache = {} self._initialized = False - + def process(self, data: List[Dict]) -> List[Dict]: """ Process a list of data items. - + Args: data: List of dictionaries to process - + Returns: Processed data list """ if not self._initialized: self._initialize() - + results = [] for item in data: processed = self._process_item(item) results.append(processed) - + return results - + def _initialize(self): """Initialize internal state.""" self._cache.clear() self._initialized = True - + def _process_item(self, item: Dict) -> Dict: """Process a single item.""" # Implementation details return {**item, 'processed': True} - def main(): """Main entry point.""" config = Config() processor = DataProcessor(config) - + test_data = [ {'id': 1, 'value': 'test1'}, {'id': 2, 'value': 'test2'}, ] - + results = processor.process(test_data) print(f"Processed {len(results)} items") - if __name__ == "__main__": main() ''' # Sample markdown file -sample_markdown = '''# RAG System Documentation +sample_markdown = """# RAG System Documentation ## Overview @@ -175,103 +179,103 @@ Main class for indexing projects. ### CodeSearcher Provides semantic search capabilities. -''' +""" def test_integration(): """Test the complete RAG system with smart chunking.""" - + # Create temporary project directory with tempfile.TemporaryDirectory() as tmpdir: project_path = Path(tmpdir) - + # Create test files (project_path / "processor.py").write_text(sample_code) (project_path / "README.md").write_text(sample_markdown) - + print("=" * 60) print("TESTING RAG SYSTEM INTEGRATION") print("=" * 60) - + # Index the project print("\n1. Indexing project...") indexer = ProjectIndexer(project_path) stats = indexer.index_project() - + print(f" - Files indexed: {stats['files_indexed']}") print(f" - Total chunks: {stats['chunks_created']}") print(f" - Indexing time: {stats['time_taken']:.2f}s") - + # Verify chunks were created properly print("\n2. Verifying chunk metadata...") - + # Initialize searcher searcher = CodeSearcher(project_path) - + # Search for specific content print("\n3. Testing search functionality...") - + # Test 1: Search for class with docstring results = searcher.search("data processor class unified interface", top_k=3) - print(f"\n Test 1 - Class search:") + print("\n Test 1 - Class search:") for i, result in enumerate(results[:1]): print(f" - Match {i+1}: {result.file_path}") print(f" Chunk type: {result.chunk_type}") print(f" Score: {result.score:.3f}") - if 'This class handles' in result.content: + if "This class handles" in result.content: print(" [OK] Docstring included with class") else: print(" [FAIL] Docstring not found") - + # Test 2: Search for method with docstring results = searcher.search("process list of data items", top_k=3) - print(f"\n Test 2 - Method search:") + print("\n Test 2 - Method search:") for i, result in enumerate(results[:1]): print(f" - Match {i+1}: {result.file_path}") print(f" Chunk type: {result.chunk_type}") print(f" Parent class: {getattr(result, 'parent_class', 'N/A')}") - if 'Args:' in result.content and 'Returns:' in result.content: + if "Args:" in result.content and "Returns:" in result.content: print(" [OK] Docstring included with method") else: print(" [FAIL] Method docstring not complete") - + # Test 3: Search markdown content results = searcher.search("smart chunking capabilities markdown", top_k=3) - print(f"\n Test 3 - Markdown search:") + print("\n Test 3 - Markdown search:") for i, result in enumerate(results[:1]): print(f" - Match {i+1}: {result.file_path}") print(f" Chunk type: {result.chunk_type}") print(f" Lines: {result.start_line}-{result.end_line}") - + # Test 4: Verify chunk navigation - print(f"\n Test 4 - Chunk navigation:") + print("\n Test 4 - Chunk navigation:") all_results = searcher.search("", top_k=100) # Get all chunks - py_chunks = [r for r in all_results if r.file_path.endswith('.py')] - + py_chunks = [r for r in all_results if r.file_path.endswith(".py")] + if py_chunks: first_chunk = py_chunks[0] print(f" - First chunk: index={getattr(first_chunk, 'chunk_index', 'N/A')}") print(f" Next chunk ID: {getattr(first_chunk, 'next_chunk_id', 'N/A')}") - + # Verify chain valid_chain = True for i in range(len(py_chunks) - 1): curr = py_chunks[i] - next_chunk = py_chunks[i + 1] + # py_chunks[i + 1] # Unused variable removed expected_next = f"processor_{i+1}" - if getattr(curr, 'next_chunk_id', None) != expected_next: + if getattr(curr, "next_chunk_id", None) != expected_next: valid_chain = False break - + if valid_chain: print(" [OK] Chunk navigation chain is valid") else: print(" [FAIL] Chunk navigation chain broken") - + print("\n" + "=" * 60) print("INTEGRATION TEST COMPLETED") print("=" * 60) if __name__ == "__main__": - test_integration() \ No newline at end of file + test_integration() diff --git a/tests/test_smart_ranking.py b/tests/test_smart_ranking.py index 4100f1c..4c8be61 100755 --- a/tests/test_smart_ranking.py +++ b/tests/test_smart_ranking.py @@ -8,26 +8,26 @@ and producing better quality results. Run with: python3 tests/test_smart_ranking.py """ -import unittest import sys -from pathlib import Path +import unittest from datetime import datetime, timedelta -from unittest.mock import patch, MagicMock +from pathlib import Path +from unittest.mock import MagicMock, patch + +from mini_rag.search import CodeSearcher, SearchResult # Add project to path sys.path.insert(0, str(Path(__file__).parent.parent)) -from mini_rag.search import SearchResult, CodeSearcher - class TestSmartRanking(unittest.TestCase): """ Test smart result re-ranking for better search quality. - + These tests verify that important files, recent files, and well-structured content get appropriate boosts. """ - + def setUp(self): """Set up test results for ranking.""" # Create mock search results with different characteristics @@ -40,27 +40,31 @@ class TestSmartRanking(unittest.TestCase): end_line=2, chunk_type="text", name="temp", - language="text" + language="text", ), SearchResult( - file_path=Path("README.md"), - content="This is a comprehensive README file\nwith detailed installation instructions\nand usage examples for beginners.", + file_path=Path("README.md"), + content=( + "This is a comprehensive README file\n" + "with detailed installation instructions\n" + "and usage examples for beginners." + ), score=0.7, # Lower initial score start_line=1, end_line=5, chunk_type="markdown", name="Installation Guide", - language="markdown" + language="markdown", ), SearchResult( file_path=Path("src/main.py"), - content="def main():\n \"\"\"Main application entry point.\"\"\"\n app = create_app()\n return app.run()", + content='def main():\n """Main application entry point."""\n app = create_app()\n return app.run()', score=0.75, start_line=10, end_line=15, chunk_type="function", name="main", - language="python" + language="python", ), SearchResult( file_path=Path("temp/cache_123.log"), @@ -68,123 +72,123 @@ class TestSmartRanking(unittest.TestCase): score=0.85, start_line=1, end_line=1, - chunk_type="text", + chunk_type="text", name="log", - language="text" - ) + language="text", + ), ] - + def test_01_important_file_boost(self): """ โœ… Test that important files get ranking boosts. - + README files, main files, config files, etc. should be ranked higher than random temporary files. """ print("\n๐Ÿ“ˆ Testing important file boost...") - + # Create a minimal CodeSearcher to test ranking searcher = MagicMock() searcher._smart_rerank = CodeSearcher._smart_rerank.__get__(searcher) - + # Test re-ranking ranked = searcher._smart_rerank(self.mock_results.copy()) - + # Find README and temp file results - readme_result = next((r for r in ranked if 'README' in str(r.file_path)), None) - temp_result = next((r for r in ranked if 'temp' in str(r.file_path)), None) - + readme_result = next((r for r in ranked if "README" in str(r.file_path)), None) + temp_result = next((r for r in ranked if "temp" in str(r.file_path)), None) + self.assertIsNotNone(readme_result) self.assertIsNotNone(temp_result) - + # README should be boosted (original 0.7 * 1.2 = 0.84) self.assertGreater(readme_result.score, 0.8) - + # README should now rank higher than the temp file readme_index = ranked.index(readme_result) temp_index = ranked.index(temp_result) self.assertLess(readme_index, temp_index) - + print(f" โœ… README boosted from 0.7 to {readme_result.score:.3f}") print(f" ๐Ÿ“Š README now ranks #{readme_index + 1}, temp file ranks #{temp_index + 1}") - + def test_02_content_quality_boost(self): """ โœ… Test that well-structured content gets boosts. - + Content with multiple lines and good structure should rank higher than very short snippets. """ print("\n๐Ÿ“ Testing content quality boost...") - + searcher = MagicMock() searcher._smart_rerank = CodeSearcher._smart_rerank.__get__(searcher) - + ranked = searcher._smart_rerank(self.mock_results.copy()) - + # Find short and long content results short_result = next((r for r in ranked if len(r.content.strip()) < 20), None) - structured_result = next((r for r in ranked if 'README' in str(r.file_path)), None) - + structured_result = next((r for r in ranked if "README" in str(r.file_path)), None) + if short_result: # Short content should be penalized (score * 0.9) print(f" ๐Ÿ“‰ Short content penalized: {short_result.score:.3f}") # Original was likely reduced - + if structured_result: # Well-structured content gets small boost (score * 1.02) - lines = structured_result.content.strip().split('\n') + lines = structured_result.content.strip().split("\n") if len(lines) >= 3: print(f" ๐Ÿ“ˆ Structured content boosted: {structured_result.score:.3f}") print(f" ({len(lines)} lines of content)") - + self.assertTrue(True) # Test passes if no exceptions - + def test_03_chunk_type_relevance(self): """ โœ… Test that relevant chunk types get appropriate boosts. - + Functions, classes, and documentation should be ranked higher than random text snippets. """ print("\n๐Ÿท๏ธ Testing chunk type relevance...") - + searcher = MagicMock() searcher._smart_rerank = CodeSearcher._smart_rerank.__get__(searcher) - + ranked = searcher._smart_rerank(self.mock_results.copy()) - + # Find function result - function_result = next((r for r in ranked if r.chunk_type == 'function'), None) - + function_result = next((r for r in ranked if r.chunk_type == "function"), None) + if function_result: # Function should get boost (original score * 1.1) print(f" โœ… Function chunk boosted: {function_result.score:.3f}") print(f" Function: {function_result.name}") - + # Should rank well compared to original score original_score = 0.75 self.assertGreater(function_result.score, original_score) - + self.assertTrue(True) - - @patch('pathlib.Path.stat') + + @patch("pathlib.Path.stat") def test_04_recency_boost(self, mock_stat): """ โœ… Test that recently modified files get ranking boosts. - + Files modified in the last week should rank higher than very old files. """ print("\nโฐ Testing recency boost...") - + # Mock file stats for different modification times now = datetime.now() - + def mock_stat_side_effect(file_path): mock_stat_obj = MagicMock() - - if 'README' in str(file_path): + + if "README" in str(file_path): # Recent file (2 days ago) recent_time = (now - timedelta(days=2)).timestamp() mock_stat_obj.st_mtime = recent_time @@ -192,98 +196,102 @@ class TestSmartRanking(unittest.TestCase): # Old file (2 months ago) old_time = (now - timedelta(days=60)).timestamp() mock_stat_obj.st_mtime = old_time - + return mock_stat_obj - + # Apply mock to Path.stat for each result mock_stat.side_effect = lambda: mock_stat_side_effect("dummy") - + # Patch the Path constructor to return mocked paths - with patch.object(Path, 'stat', side_effect=mock_stat_side_effect): + with patch.object(Path, "stat", side_effect=mock_stat_side_effect): searcher = MagicMock() searcher._smart_rerank = CodeSearcher._smart_rerank.__get__(searcher) - + ranked = searcher._smart_rerank(self.mock_results.copy()) - - readme_result = next((r for r in ranked if 'README' in str(r.file_path)), None) - + + readme_result = next((r for r in ranked if "README" in str(r.file_path)), None) + if readme_result: # Recent file should get boost # Original 0.7 * 1.2 (important) * 1.1 (recent) * 1.02 (structured) โ‰ˆ 0.88 print(f" โœ… Recent file boosted: {readme_result.score:.3f}") self.assertGreater(readme_result.score, 0.8) - + print(" ๐Ÿ“… Recency boost system working!") - + def test_05_overall_ranking_quality(self): """ โœ… Test that overall ranking produces sensible results. - + After all boosts and penalties, the ranking should make sense: - Important, recent, well-structured files should rank highest - Short, temporary, old files should rank lowest """ print("\n๐Ÿ† Testing overall ranking quality...") - + searcher = MagicMock() searcher._smart_rerank = CodeSearcher._smart_rerank.__get__(searcher) - + # Test with original unsorted results unsorted = self.mock_results.copy() ranked = searcher._smart_rerank(unsorted) - + print(" ๐Ÿ“Š Final ranking:") for i, result in enumerate(ranked, 1): file_name = Path(result.file_path).name print(f" {i}. {file_name} (score: {result.score:.3f})") - + # Quality checks: # 1. Results should be sorted by score (descending) scores = [r.score for r in ranked] self.assertEqual(scores, sorted(scores, reverse=True)) - + # 2. README should rank higher than temp files - readme_pos = next((i for i, r in enumerate(ranked) if 'README' in str(r.file_path)), None) - temp_pos = next((i for i, r in enumerate(ranked) if 'temp' in str(r.file_path)), None) - + readme_pos = next( + (i for i, r in enumerate(ranked) if "README" in str(r.file_path)), None + ) + temp_pos = next((i for i, r in enumerate(ranked) if "temp" in str(r.file_path)), None) + if readme_pos is not None and temp_pos is not None: self.assertLess(readme_pos, temp_pos) print(f" โœ… README ranks #{readme_pos + 1}, temp file ranks #{temp_pos + 1}") - + # 3. Function/code should rank well - function_pos = next((i for i, r in enumerate(ranked) if r.chunk_type == 'function'), None) + function_pos = next( + (i for i, r in enumerate(ranked) if r.chunk_type == "function"), None + ) if function_pos is not None: self.assertLess(function_pos, len(ranked) // 2) # Should be in top half print(f" โœ… Function code ranks #{function_pos + 1}") - + print(" ๐ŸŽฏ Ranking quality looks good!") - + def test_06_zero_overhead_verification(self): """ โœ… Verify that smart ranking adds zero overhead. - + The ranking should only use existing data and lightweight operations. No additional API calls or expensive operations. """ print("\nโšก Testing zero overhead...") - + searcher = MagicMock() searcher._smart_rerank = CodeSearcher._smart_rerank.__get__(searcher) - + import time - + # Time the ranking operation start_time = time.time() - ranked = searcher._smart_rerank(self.mock_results.copy()) + # searcher._smart_rerank(self.mock_results.copy()) # Unused variable removed end_time = time.time() - + ranking_time = (end_time - start_time) * 1000 # Convert to milliseconds - + print(f" โฑ๏ธ Ranking took {ranking_time:.2f}ms for {len(self.mock_results)} results") - + # Should be very fast (< 10ms for small result sets) self.assertLess(ranking_time, 50) # Very generous upper bound - + # Verify no external calls were made (check that we only use existing data) # This is implicitly tested by the fact that we're using mock objects print(" โœ… Zero overhead verified - only uses existing result data!") @@ -297,18 +305,18 @@ def run_ranking_tests(): print("=" * 40) print("Testing the zero-overhead ranking improvements.") print() - + unittest.main(verbosity=2, exit=False) - + print("\n" + "=" * 40) print("๐Ÿ’ก Smart Ranking Features:") print(" โ€ข Important files (README, main, config) get 20% boost") - print(" โ€ข Recent files (< 1 week) get 10% boost") + print(" โ€ข Recent files (< 1 week) get 10% boost") print(" โ€ข Functions/classes get 10% boost") print(" โ€ข Well-structured content gets 2% boost") print(" โ€ข Very short content gets 10% penalty") print(" โ€ข All boosts are cumulative for maximum quality") -if __name__ == '__main__': - run_ranking_tests() \ No newline at end of file +if __name__ == "__main__": + run_ranking_tests() diff --git a/tests/troubleshoot.py b/tests/troubleshoot.py index 3e6255e..0fa60c8 100755 --- a/tests/troubleshoot.py +++ b/tests/troubleshoot.py @@ -8,21 +8,22 @@ and helps identify what's working and what needs attention. Run with: python3 tests/troubleshoot.py """ -import sys import subprocess +import sys from pathlib import Path # Add project to path sys.path.insert(0, str(Path(__file__).parent.parent)) + def main(): """Run comprehensive troubleshooting checks.""" - + print("๐Ÿ”ง FSS-Mini-RAG Troubleshooting Tool") print("=" * 50) print("This tool checks your setup and helps fix common issues.") print() - + # Menu of available tests print("Available tests:") print(" 1. Full Ollama Integration Test") @@ -30,21 +31,21 @@ def main(): print(" 3. Basic System Validation") print(" 4. All Tests (recommended)") print() - + choice = input("Select test (1-4, or Enter for all): ").strip() - + if choice == "1" or choice == "" or choice == "4": print("\n" + "๐Ÿค– OLLAMA INTEGRATION TESTS".center(50, "=")) run_test("test_ollama_integration.py") - + if choice == "2" or choice == "" or choice == "4": print("\n" + "๐Ÿงฎ SMART RANKING TESTS".center(50, "=")) run_test("test_smart_ranking.py") - + if choice == "3" or choice == "" or choice == "4": print("\n" + "๐Ÿ” SYSTEM VALIDATION TESTS".center(50, "=")) run_test("03_system_validation.py") - + print("\n" + "โœ… TROUBLESHOOTING COMPLETE".center(50, "=")) print("๐Ÿ’ก If you're still having issues:") print(" โ€ข Check docs/QUERY_EXPANSION.md for setup help") @@ -52,35 +53,37 @@ def main(): print(" โ€ข Start Ollama server: ollama serve") print(" โ€ข Install models: ollama pull qwen3:4b") + def run_test(test_file): """Run a specific test file.""" test_path = Path(__file__).parent / test_file - + if not test_path.exists(): print(f"โŒ Test file not found: {test_file}") return - + try: # Run the test - result = subprocess.run([ - sys.executable, str(test_path) - ], capture_output=True, text=True, timeout=60) - + result = subprocess.run( + [sys.executable, str(test_path)], capture_output=True, text=True, timeout=60 + ) + # Show output if result.stdout: print(result.stdout) if result.stderr: print("STDERR:", result.stderr) - + if result.returncode == 0: print(f"โœ… {test_file} completed successfully!") else: print(f"โš ๏ธ {test_file} had some issues (return code: {result.returncode})") - + except subprocess.TimeoutExpired: print(f"โฐ {test_file} timed out after 60 seconds") except Exception as e: print(f"โŒ Error running {test_file}: {e}") + if __name__ == "__main__": - main() \ No newline at end of file + main()