tools
/
drakkenheim-rag
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit

This commit is contained in:
Thom Werring 2025-09-21 21:55:57 +02:00
commit 712f188f73
No known key found for this signature in database
17 changed files with 1631 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

285
.gitignore vendored Normal file
View File

@ -0,0 +1,285 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be added to the global gitignore or merged into this project gitignore. For a PyCharm
# project, it is recommended to include the following files:
.idea/
*.iws
*.iml
*.ipr
# Visual Studio Code
.vscode/
*.code-workspace
# Local History for Visual Studio Code
.history/
# Built Visual Studio Code Extensions
*.vsix
# macOS
.DS_Store
.AppleDouble
.LSOverride
# Icon must end with two \r
Icon
# Thumbnails
._*
# Files that might appear in the root of a volume
.DocumentRevisions-V100
.fseventsd
.Spotlight-V100
.TemporaryItems
.Trashes
.VolumeIcon.icns
.com.apple.timemachine.donotpresent
# Directories potentially created on remote AFP share
.AppleDB
.AppleDesktop
Network Trash Folder
Temporary Items
.apdisk
# Windows
Thumbs.db
Thumbs.db:encryptable
ehthumbs.db
ehthumbs_vista.db
# Dump file
*.stackdump
# Folder config file
[Dd]esktop.ini
# Recycle Bin used on file shares
$RECYCLE.BIN/
# Windows Installer files
*.cab
*.msi
*.msix
*.msm
*.msp
# Windows shortcuts
*.lnk
# Linux
*~
# temporary files which can be created if a process still has a handle open of a deleted file
.fuse_hidden*
# KDE directory preferences
.directory
# Linux trash folder which might appear on any partition or disk
.Trash-*
# .nfs files are created when an open file is removed but is still being accessed
.nfs*
# Streamlit
.streamlit/
# ChromaDB data
demo-rag-chroma/
*.sqlite3
# Backup files
backups/
*.backup
*.bak
# Logs
*.log
logs/
# Temporary files
*.tmp
*.temp
temp/
# Docker
.dockerignore
# Environment files
.env.local
.env.development
.env.test
.env.production
# IDE specific files
*.swp
*.swo
*~
# OS generated files
.DS_Store?
.Spotlight-V100
.Trashes
ehthumbs.db
Thumbs.db
# Project specific
input/
processed_files.json

289
DEPLOYMENT.md Normal file
View File

@ -0,0 +1,289 @@
# Docker Deployment Guide
This guide explains how to deploy the Drakkenheim RAG application using Docker Compose with persistent storage.
## Prerequisites
- Docker and Docker Compose installed on your server
- OpenAI API key
## Project Structure
The project follows Python best practices with a `src/` layout:
```
drakkenheim/
├── src/
│ └── drakkenheim/ # Main package
│ ├── __init__.py # Package initialization
│ ├── auth.py # Authentication & session management
│ ├── llm.py # LLM, embeddings & vector operations
│ ├── document.py # Document processing
│ ├── storage.py # File persistence
│ └── ui.py # User interface components
├── app.py # Main application entry point
├── generate_password.py # Password generation utility
├── docker-compose.yml # Docker orchestration
├── Dockerfile # Container definition
└── requirements/ # Dependencies
```
**Benefits of the `src/` layout:**
- **Clean separation** between source code and project files
- **Prevents import issues** during development and testing
- **Industry standard** for Python projects
- **Better packaging** and distribution support
- **Easier testing** and CI/CD integration
## Quick Start
1. **Clone and navigate to the project:**
```bash
git clone <your-repo-url>
cd drakkenheim
```
2. **Set up environment variables:**
```bash
# Copy the example environment file
cp env.example .env
# Generate secure password hash
python3 generate_password.py
# Edit .env and add your OpenAI API key and generated auth config
nano .env
```
3. **Build and start the application:**
```bash
docker compose up -d
```
4. **Access the application:**
- Open your browser to `http://your-server-ip:8501`
- The application will be available and ready to use
## Persistent Storage
The application uses Docker volumes to persist data between restarts:
- **ChromaDB Database**: `chroma_data` volume
- **Processed Files Metadata**: `processed_files` volume
- **Uploaded Files**: `uploaded_files` volume
## Management Commands
### Start the application:
```bash
docker compose up -d
```
### Stop the application:
```bash
docker compose down
```
### View logs:
```bash
docker compose logs -f rag-app
```
### Restart the application:
```bash
docker compose restart rag-app
```
### Update the application:
```bash
# Pull latest changes
git pull
# Rebuild and restart
docker compose up -d --build
```
### Backup data:
```bash
# Create backup directory
mkdir -p backups
# Backup ChromaDB data
docker run --rm -v drakkenheim_chroma_data:/data -v $(pwd)/backups:/backup alpine tar czf /backup/chroma_data_$(date +%Y%m%d_%H%M%S).tar.gz -C /data .
# Backup processed files
docker run --rm -v drakkenheim_processed_files:/data -v $(pwd)/backups:/backup alpine tar czf /backup/processed_files_$(date +%Y%m%d_%H%M%S).tar.gz -C /data .
```
### Restore data:
```bash
# Restore ChromaDB data
docker run --rm -v drakkenheim_chroma_data:/data -v $(pwd)/backups:/backup alpine tar xzf /backup/chroma_data_YYYYMMDD_HHMMSS.tar.gz -C /data
# Restore processed files
docker run --rm -v drakkenheim_processed_files:/data -v $(pwd)/backups:/backup alpine tar xzf /backup/processed_files_YYYYMMDD_HHMMSS.tar.gz -C /data
```
## Configuration
### Environment Variables
Edit `.env` file to customize:
```bash
# Required
OPENAI_API_KEY=your_api_key_here
# Authentication (REQUIRED for production)
APP_PASSWORD_HASH=your_hashed_password_here
PASSWORD_SALT=your_random_salt_here
SESSION_TIMEOUT=3600
# Optional overrides
OPENAI_EMBEDDING_MODEL=text-embedding-3-small
OPENAI_LLM_MODEL=gpt-4o-mini
STREAMLIT_SERVER_PORT=8501
```
### Authentication Setup
**Generate secure credentials:**
```bash
# Run the password generator
python3 generate_password.py
# Follow the prompts to set your password
# Copy the generated values to your .env file
```
**Security Features:**
- **Password hashing**: SHA-256 with salt
- **Session timeout**: Automatic logout after inactivity
- **Secure storage**: Passwords never stored in plain text
- **Development mode**: No auth required if no password is set
### Port Configuration
To change the port, modify `docker-compose.yml`:
```yaml
ports:
- "8080:8501" # Change 8080 to your desired port
```
## Monitoring
### Health Check
The application includes a health check that monitors:
- Streamlit server availability
- Automatic restart on failure
### Logs
```bash
# View all logs
docker-compose logs
# Follow logs in real-time
docker-compose logs -f
# View specific service logs
docker-compose logs rag-app
```
## Troubleshooting
### Common Issues
1. **Port already in use:**
```bash
# Check what's using the port
sudo netstat -tulpn | grep :8501
# Kill the process or change the port in docker-compose.yml
```
2. **Permission issues:**
```bash
# Fix volume permissions
sudo chown -R 1000:1000 /var/lib/docker/volumes/drakkenheim_*/
```
3. **Out of disk space:**
```bash
# Clean up unused Docker resources
docker system prune -a
# Check volume sizes
docker system df -v
```
4. **API key issues:**
- Verify your `.env` file has the correct API key
- Check the logs for authentication errors
- Ensure the API key has sufficient credits
### Reset Application
```bash
# Stop and remove containers
docker compose down
# Remove volumes (WARNING: This deletes all data)
docker volume rm drakkenheim_chroma_data drakkenheim_processed_files drakkenheim_uploaded_files
# Start fresh
docker compose up -d
```
## Security Considerations
1. **API Key Security:**
- Never commit `.env` files to version control
- Use Docker secrets for production deployments
- Rotate API keys regularly
2. **Network Security:**
- Use a reverse proxy (nginx) for production
- Enable HTTPS with SSL certificates
- Restrict access with firewall rules
3. **Data Security:**
- Regular backups of persistent volumes
- Encrypt sensitive data at rest
- Monitor access logs
## Production Deployment
For production deployment, consider:
1. **Reverse Proxy Setup:**
```nginx
server {
listen 80;
server_name your-domain.com;
location / {
proxy_pass http://localhost:8501;
proxy_set_header Host $host;
proxy_set_header X-Real-IP $remote_addr;
}
}
```
2. **SSL Certificate:**
```bash
# Using Let's Encrypt
sudo certbot --nginx -d your-domain.com
```
3. **Monitoring:**
- Set up log aggregation
- Monitor resource usage
- Set up alerts for failures
## Support
If you encounter issues:
1. Check the logs: `docker compose logs -f`
2. Verify your environment variables
3. Ensure Docker has sufficient resources
4. Check the troubleshooting section above

38
Dockerfile Normal file
View File

@ -0,0 +1,38 @@
# Use Python 3.11 slim image
FROM python:3.11-slim
# Set working directory
WORKDIR /app
# Install system dependencies
RUN apt-get update && apt-get install -y \
build-essential \
curl \
git \
&& rm -rf /var/lib/apt/lists/*
# Copy requirements first for better caching
COPY requirements/requirements.txt .
# Install Python dependencies
RUN pip install --no-cache-dir -r requirements.txt
# Copy application code
COPY . .
# Create directories for persistent storage
RUN mkdir -p /app/data/chroma-db
RUN mkdir -p /app/data/processed-files
RUN mkdir -p /app/data/uploaded-files
# Add src directory to Python path
ENV PYTHONPATH=/app/src:$PYTHONPATH
# Expose Streamlit port
EXPOSE 8501
# Health check
HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
# Run the application
CMD ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]

72
Makefile Normal file
View File

@ -0,0 +1,72 @@
SHELL :=/bin/bash
.PHONY: clean check setup docker-build docker-up docker-down docker-logs docker-restart docker-backup
.DEFAULT_GOAL=help
VENV_DIR = .venv
PYTHON_VERSION = python3
check: # Ruff check
@ruff check .
@echo "✅ Check complete!"
fix: # Fix auto-fixable linting issues
@ruff check app.py --fix
clean: # Clean temporary files
@rm -rf __pycache__ .pytest_cache
@find . -name '*.pyc' -exec rm -r {} +
@find . -name '__pycache__' -exec rm -r {} +
@rm -rf build dist
@find . -name '*.egg-info' -type d -exec rm -r {} +
run: # Run the application
@PYTHONPATH=src streamlit run app.py
setup: # Initial project setup
@echo "Creating virtual env at: $(VENV_DIR)"
@$(PYTHON_VERSION) -m venv $(VENV_DIR)
@echo "Installing dependencies..."
@source $(VENV_DIR)/bin/activate && pip install -r requirements/requirements-dev.txt && pip install -r requirements/requirements.txt
@echo -e "\n✅ Done.\n🎉 Run the following commands to get started:\n\n ➡️ source $(VENV_DIR)/bin/activate\n ➡️ export OPENAI_API_KEY=your_api_key_here\n ➡️ make run\n"
docker-build: # Build Docker image
@echo "Building Docker image..."
@docker compose build
@echo "✅ Docker image built!"
docker-up: # Start the application with Docker Compose
@echo "Starting RAG application with Docker Compose..."
@docker compose up -d
@echo "✅ Application started! Access at http://localhost:8501"
docker-down: # Stop the application
@echo "Stopping RAG application..."
@docker compose down
@echo "✅ Application stopped!"
docker-logs: # View application logs
@docker compose logs -f rag-app
docker-restart: # Restart the application
@echo "Restarting RAG application..."
@docker compose restart rag-app
@echo "✅ Application restarted!"
docker-backup: # Backup persistent data
@echo "Creating backup directory..."
@mkdir -p backups
@echo "Backing up ChromaDB data..."
@docker run --rm -v drakkenheim_chroma_data:/data -v $(PWD)/backups:/backup alpine tar czf /backup/chroma_data_$$(date +%Y%m%d_%H%M%S).tar.gz -C /data .
@echo "Backing up processed files..."
@docker run --rm -v drakkenheim_processed_files:/data -v $(PWD)/backups:/backup alpine tar czf /backup/processed_files_$$(date +%Y%m%d_%H%M%S).tar.gz -C /data .
@echo "✅ Backup completed! Check the backups/ directory"
docker-clean: # Clean up Docker resources
@echo "Cleaning up Docker resources..."
@docker compose down
@docker system prune -f
@echo "✅ Docker cleanup completed!"
help: # Show this help
@egrep -h '\s#\s' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?# "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}'

60
app.py Normal file
View File

@ -0,0 +1,60 @@
"""
Drakkenheim RAG Application - Main Entry Point
A Streamlit-based RAG (Retrieval-Augmented Generation) application that allows users to
upload documents and ask questions about their content using OpenAI's GPT models.
This is the main application file that orchestrates all the modules.
"""
import os
import warnings
import streamlit as st
# Import our custom modules
from src.drakkenheim import check_authentication, login_form, load_processed_files, render_sidebar, render_chat_interface
# Set environment variables to suppress warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["CHROMA_TELEMETRY"] = "false"
os.environ["TOKENIZERS_VERBOSITY"] = "error"
# Suppress specific warnings
warnings.filterwarnings("ignore", category=UserWarning, module="sentence_transformers")
warnings.filterwarnings("ignore", category=UserWarning, module="transformers")
warnings.filterwarnings("ignore", category=FutureWarning)
def main():
"""Main application function."""
# Set page configuration
st.set_page_config(
page_title="RAG Chat Assistant",
page_icon="🤖",
layout="wide"
)
# Check authentication first
if not check_authentication():
login_form()
return
# Initialize session state
if "messages" not in st.session_state:
st.session_state.messages = []
if "documents_loaded" not in st.session_state:
st.session_state.documents_loaded = False
if "processed_files" not in st.session_state:
st.session_state.processed_files = []
if "files_loaded" not in st.session_state:
st.session_state.files_loaded = False
load_processed_files()
# Render the application
render_sidebar()
render_chat_interface()
if __name__ == "__main__":
main()

88
deploy.sh Executable file
View File

@ -0,0 +1,88 @@
#!/bin/bash
# Drakkenheim RAG Application Deployment Script
# This script helps you deploy the RAG application using Docker Compose
set -e
echo "🚀 Drakkenheim RAG Application Deployment"
echo "========================================"
# Check if Docker is installed
if ! command -v docker &> /dev/null; then
echo "❌ Docker is not installed. Please install Docker first."
exit 1
fi
# Check if Docker Compose is available
if ! docker compose version &> /dev/null; then
echo "❌ Docker Compose is not available. Please ensure Docker is up to date."
exit 1
fi
# Check if .env file exists
if [ ! -f .env ]; then
echo "📝 Creating .env file from template..."
cp env.example .env
echo ""
echo "🔐 Setting up authentication..."
echo "Generating secure password hash..."
python3 generate_password.py
echo ""
echo "⚠️ Please edit .env file and add your OpenAI API key:"
echo " nano .env"
echo ""
read -p "Press Enter after you've added your API key..."
fi
# Check if OPENAI_API_KEY is set
if ! grep -q "OPENAI_API_KEY=sk-" .env; then
echo "❌ OPENAI_API_KEY not found in .env file."
echo "Please edit .env file and add your OpenAI API key:"
echo " OPENAI_API_KEY=sk-your-key-here"
exit 1
fi
# Check if authentication is configured
if ! grep -q "APP_PASSWORD_HASH=" .env || grep -q "APP_PASSWORD_HASH=your_hashed_password_here" .env; then
echo "⚠️ Authentication not configured. Generating secure password..."
python3 generate_password.py
echo ""
echo "Please copy the generated values to your .env file and restart deployment."
exit 1
fi
echo "✅ Environment configuration looks good!"
# Build the Docker image
echo "🔨 Building Docker image..."
docker compose build
# Start the application
echo "🚀 Starting RAG application..."
docker compose up -d
# Wait for the application to start
echo "⏳ Waiting for application to start..."
sleep 10
# Check if the application is running
if docker compose ps | grep -q "Up"; then
echo "✅ Application is running!"
echo ""
echo "🌐 Access your RAG application at:"
echo " http://localhost:8501"
echo ""
echo "📊 To view logs:"
echo " make docker-logs"
echo ""
echo "🛑 To stop the application:"
echo " make docker-down"
echo ""
echo "📋 To backup your data:"
echo " make docker-backup"
else
echo "❌ Application failed to start. Check logs with:"
echo " make docker-logs"
exit 1
fi

42
docker-compose.yml Normal file
View File

@ -0,0 +1,42 @@
version: '3.8'
services:
rag-app:
build: .
container_name: drakkenheim-rag
ports:
- "8501:8501"
environment:
- OPENAI_API_KEY=${OPENAI_API_KEY}
- APP_PASSWORD_HASH=${APP_PASSWORD_HASH}
- PASSWORD_SALT=${PASSWORD_SALT}
- SESSION_TIMEOUT=${SESSION_TIMEOUT:-3600}
- TOKENIZERS_PARALLELISM=false
- CHROMA_TELEMETRY=false
- TOKENIZERS_VERBOSITY=error
- CHROMA_DB_PATH=/app/data/chroma-db
- PROCESSED_FILES_PATH=/app/data/processed-files
volumes:
# Persistent storage for ChromaDB database
- chroma_data:/app/data/chroma-db
# Persistent storage for processed files metadata
- processed_files:/app/data/processed-files
# Persistent storage for uploaded files (optional)
- uploaded_files:/app/data/uploaded-files
# Mount the input directory for initial documents
- ./input:/app/input:ro
restart: unless-stopped
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:8501/_stcore/health"]
interval: 30s
timeout: 10s
retries: 3
start_period: 40s
volumes:
chroma_data:
driver: local
processed_files:
driver: local
uploaded_files:
driver: local

18
env.example Normal file
View File

@ -0,0 +1,18 @@
# OpenAI API Configuration
OPENAI_API_KEY=your_openai_api_key_here
# Authentication Configuration (REQUIRED for production)
APP_PASSWORD_HASH=your_hashed_password_here
PASSWORD_SALT=your_random_salt_here
SESSION_TIMEOUT=3600
# Optional: Override default models
# OPENAI_EMBEDDING_MODEL=text-embedding-3-small
# OPENAI_LLM_MODEL=gpt-4o-mini
# Optional: Streamlit configuration
# STREAMLIT_SERVER_PORT=8501
# STREAMLIT_SERVER_ADDRESS=0.0.0.0
# Optional: ChromaDB configuration
# CHROMA_DB_PATH=/app/data/chroma-db

56
generate_password.py Executable file
View File

@ -0,0 +1,56 @@
#!/usr/bin/env python3
"""
Password Generator for Drakkenheim RAG Application
This script helps you generate secure password hashes and salts for authentication.
"""
import hashlib
import secrets
import sys
import getpass
def generate_salt(length=32):
"""Generate a random salt."""
return secrets.token_hex(length)
def hash_password(password, salt):
"""Hash a password with salt using SHA-256."""
return hashlib.sha256((password + salt).encode()).hexdigest()
def main():
print("🔐 Drakkenheim RAG - Password Generator")
print("=" * 40)
# Get password from user
password = getpass.getpass("Enter password: ")
if not password:
print("❌ Password cannot be empty!")
sys.exit(1)
# Confirm password
confirm_password = getpass.getpass("Confirm password: ")
if password != confirm_password:
print("❌ Passwords don't match!")
sys.exit(1)
# Generate salt and hash
salt = generate_salt()
password_hash = hash_password(password, salt)
print("\n✅ Password configuration generated!")
print("\n📋 Add these to your .env file:")
print("-" * 40)
print(f"APP_PASSWORD_HASH={password_hash}")
print(f"PASSWORD_SALT={salt}")
print("SESSION_TIMEOUT=3600")
print("-" * 40)
print("\n🔒 Security notes:")
print("• Keep your .env file secure and never commit it to version control")
print("• The salt should be unique for each deployment")
print("• SESSION_TIMEOUT is in seconds (3600 = 1 hour)")
print("• Users will be logged out automatically after the timeout")
if __name__ == "__main__":
main()

1
requirements/requirements-dev.txt Normal file
View File

@ -0,0 +1 @@
ruff==0.7.4

6
requirements/requirements.txt Normal file
View File

@ -0,0 +1,6 @@
openai==1.3.0 # OpenAI API (compatible with ChromaDB 0.4.22)
chromadb==0.4.22 # Vector Database (compatible with OpenAI 1.12.0)
sentence-transformers==3.3.1 # CrossEncoder Re-ranking
streamlit==1.40.1 # Application UI
PyMuPDF==1.24.14 # PDF Document loader
langchain-community==0.3.7 # Utils for text splitting

35
src/drakkenheim/__init__.py Normal file
View File

@ -0,0 +1,35 @@
"""
Drakkenheim RAG Application
A Streamlit-based RAG (Retrieval-Augmented Generation) application that allows users to
upload documents and ask questions about their content using OpenAI's GPT models.
This package contains all the core modules for the application.
"""
__version__ = "1.0.0"
__author__ = "Drakkenheim Team"
__description__ = "RAG Chat Assistant with Document Processing"
# Import main components for easy access
from .auth import check_authentication, login_form, logout
from .llm import get_vector_collection, add_to_vector_collection, query_collection, call_llm
from .document import process_document, normalize_uploaded_file_name
from .storage import save_processed_files, load_processed_files
from .ui import render_sidebar, render_chat_interface
__all__ = [
"check_authentication",
"login_form",
"logout",
"get_vector_collection",
"add_to_vector_collection",
"query_collection",
"call_llm",
"process_document",
"normalize_uploaded_file_name",
"save_processed_files",
"load_processed_files",
"render_sidebar",
"render_chat_interface",
]

93
src/drakkenheim/auth.py Normal file
View File

@ -0,0 +1,93 @@
"""
Authentication module for Drakkenheim RAG Application.
Handles user authentication, password verification, and session management.
"""
import os
import hashlib
from datetime import datetime
import streamlit as st
# Authentication configuration
AUTH_CONFIG = {
"password_hash": os.getenv("APP_PASSWORD_HASH", ""),
"session_timeout": int(os.getenv("SESSION_TIMEOUT", "3600")), # 1 hour default
}
def hash_password(password: str) -> str:
"""Hash a password using SHA-256 with salt."""
salt = os.getenv("PASSWORD_SALT", "default_salt_change_in_production")
return hashlib.sha256((password + salt).encode()).hexdigest()
def verify_password(password: str) -> bool:
"""Verify if the provided password is correct."""
if not AUTH_CONFIG["password_hash"]:
# If no password is set, allow access (for development)
return True
password_hash = hash_password(password)
return password_hash == AUTH_CONFIG["password_hash"]
def check_authentication():
"""Check if user is authenticated and session is valid."""
if "authenticated" not in st.session_state:
st.session_state.authenticated = False
st.session_state.login_time = None
if st.session_state.authenticated:
# Check session timeout
if st.session_state.login_time:
current_time = datetime.now().timestamp()
if current_time - st.session_state.login_time > AUTH_CONFIG["session_timeout"]:
st.session_state.authenticated = False
st.session_state.login_time = None
st.rerun()
return st.session_state.authenticated
def login_form():
"""Display login form."""
st.title("🔐 Drakkenheim RAG - Login")
st.markdown("Please enter the password to access the application.")
with st.form("login_form"):
password = st.text_input("Password", type="password", placeholder="Enter password")
submit_button = st.form_submit_button("Login", type="primary")
if submit_button:
if verify_password(password):
st.session_state.authenticated = True
st.session_state.login_time = datetime.now().timestamp()
st.success("✅ Login successful!")
st.rerun()
else:
st.error("❌ Invalid password. Please try again.")
# Show help for setting up password
with st.expander("🔧 Setup Instructions"):
st.markdown("""
**To set up authentication:**
1. Set the `APP_PASSWORD_HASH` environment variable with a hashed password
2. Set the `PASSWORD_SALT` environment variable for additional security
3. Restart the application
**Generate password hash:**
```bash
python3 -c "import hashlib; print(hashlib.sha256(('your_password' + 'your_salt').encode()).hexdigest())"
```
**Example .env configuration:**
```
APP_PASSWORD_HASH=your_hashed_password_here
PASSWORD_SALT=your_random_salt_here
SESSION_TIMEOUT=3600
```
""")
def logout():
"""Logout the current user."""
st.session_state.authenticated = False
st.session_state.login_time = None
st.rerun()

82
src/drakkenheim/document.py Normal file
View File

@ -0,0 +1,82 @@
"""
Document Processing module for Drakkenheim RAG Application.
Handles document loading, processing, and text splitting for various file formats.
"""
import os
import tempfile
import json
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from streamlit.runtime.uploaded_file_manager import UploadedFile
def process_document(uploaded_file: UploadedFile) -> list[Document]:
"""Process uploaded document and split into chunks for vector storage.
Handles multiple file formats (PDF, TXT, JSON, MD) and converts them into
Document objects that can be processed by the vector database.
Args:
uploaded_file: Streamlit UploadedFile object containing the document
Returns:
list[Document]: List of Document objects ready for vector storage
Raises:
ValueError: If file type is not supported or file is invalid
"""
file_extension = uploaded_file.name.split('.')[-1].lower()
if file_extension == 'pdf':
# Save uploaded file temporarily
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
tmp_file.write(uploaded_file.getvalue())
tmp_file_path = tmp_file.name
# Load PDF using PyMuPDF
loader = PyMuPDFLoader(tmp_file_path)
docs = loader.load()
# Clean up temporary file
os.unlink(tmp_file_path)
elif file_extension in ['txt', 'text']:
content = uploaded_file.read().decode('utf-8')
docs = [Document(page_content=content, metadata={"source": uploaded_file.name})]
elif file_extension == 'json':
content = uploaded_file.read().decode('utf-8')
try:
json_data = json.loads(content)
formatted_content = json.dumps(json_data, indent=2, ensure_ascii=False)
docs = [Document(page_content=formatted_content, metadata={"source": uploaded_file.name, "type": "json"})]
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON file: {str(e)}")
elif file_extension in ['md', 'markdown']:
content = uploaded_file.read().decode('utf-8')
docs = [Document(page_content=content, metadata={"source": uploaded_file.name, "type": "markdown"})]
else:
raise ValueError(f"Unsupported file type: {file_extension}. Supported types: pdf, txt, json, md")
# Split documents into chunks
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=200, # Smaller chunks to avoid token limits
chunk_overlap=50,
separators=["\n\n", "\n", ".", "?", "!", " ", ""],
)
return text_splitter.split_documents(docs)
def normalize_uploaded_file_name(uploaded_file: UploadedFile) -> str:
"""Normalize uploaded file name for consistent storage.
Args:
uploaded_file: Streamlit UploadedFile object
Returns:
str: Normalized file name safe for use as identifier
"""
return uploaded_file.name.replace(" ", "_").replace("/", "_").replace("\\", "_")

197
src/drakkenheim/llm.py Normal file
View File

@ -0,0 +1,197 @@
"""
LLM and Vector Operations module for Drakkenheim RAG Application.
Handles OpenAI API interactions, embeddings, vector operations, and document re-ranking.
"""
import os
import openai
import chromadb
from sentence_transformers import CrossEncoder
# Custom OpenAI embedding function to avoid ChromaDB compatibility issues
class CustomOpenAIEmbeddingFunction:
def __init__(self, api_key: str, model_name: str = "text-embedding-3-small"):
self.api_key = api_key
self.model_name = model_name
self.client = openai.OpenAI(api_key=api_key)
def __call__(self, input):
"""Generate embeddings for input texts."""
try:
# Truncate very long texts to avoid token limits
max_chars = 8000 # Conservative limit for text-embedding-3-small
truncated_input = []
for text in input:
if len(text) > max_chars:
truncated_input.append(text[:max_chars])
print(f"Warning: Truncated text from {len(text)} to {max_chars} characters")
else:
truncated_input.append(text)
response = self.client.embeddings.create(
model=self.model_name,
input=truncated_input
)
return [embedding.embedding for embedding in response.data]
except Exception as e:
print(f"Error generating embeddings: {e}")
return [[0.0] * 1536 for _ in input] # Return zero vectors as fallback
def get_vector_collection():
"""Gets or creates a ChromaDB collection for vector storage.
Creates a custom OpenAI embedding function using the text-embedding-3-small model and initializes
a persistent ChromaDB client. Returns a collection that can be used to store and
query document embeddings.
Returns:
chromadb.Collection: A ChromaDB collection configured with the custom OpenAI embedding
function and cosine similarity space.
"""
openai_ef = CustomOpenAIEmbeddingFunction(
api_key=os.getenv("OPENAI_API_KEY"),
model_name="text-embedding-3-small",
)
# Use Docker volume path if available, fallback to local path
chroma_path = os.getenv("CHROMA_DB_PATH", "./demo-rag-chroma")
chroma_client = chromadb.PersistentClient(path=chroma_path)
return chroma_client.get_or_create_collection(
name="rag_app",
embedding_function=openai_ef,
metadata={"hnsw:space": "cosine"},
)
def add_to_vector_collection(all_splits, file_name: str):
"""Adds document splits to a vector collection for semantic search.
Processes document splits in batches to avoid token limits and adds them to the
ChromaDB collection with proper metadata and unique IDs.
Args:
all_splits: List of Document objects to add to the collection
file_name: Name of the file being processed (used for ID generation)
"""
collection = get_vector_collection()
batch_size = 50 # Process documents in smaller batches
for batch_start in range(0, len(all_splits), batch_size):
batch_end = min(batch_start + batch_size, len(all_splits))
batch_splits = all_splits[batch_start:batch_end]
documents, metadatas, ids = [], [], []
for idx, split in enumerate(batch_splits):
documents.append(split.page_content)
metadatas.append(split.metadata)
ids.append(f"{file_name}_{batch_start + idx}")
collection.upsert(
documents=documents,
metadatas=metadatas,
ids=ids,
)
def query_collection(prompt: str, n_results: int = 10):
"""Queries the vector collection for relevant documents.
Searches the ChromaDB collection for documents similar to the given prompt
and returns the most relevant results.
Args:
prompt: The search query string
n_results: Number of results to return (default: 10)
Returns:
dict: Query results containing documents, metadatas, distances, and ids
"""
collection = get_vector_collection()
results = collection.query(query_texts=[prompt], n_results=n_results)
return results
def call_llm(context: str, prompt: str):
"""Calls the language model with context and prompt to generate a response.
Uses OpenAI GPT-4o-mini API to stream responses from a language model by providing context and a
question prompt. The model uses a system prompt to format and ground its responses appropriately.
Args:
context: String containing the relevant context for answering the question
prompt: String containing the user's question
Yields:
String chunks of the generated response as they become available from the model
Raises:
OpenAIError: If there are issues communicating with the OpenAI API
"""
client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
response = client.chat.completions.create(
model="gpt-4o-mini",
stream=True,
messages=[
{
"role": "system",
"content": f"""You are an AI assistant tasked with providing detailed answers based solely on the given context. Your goal is to analyze the information provided and formulate a comprehensive, well-structured response to the question.
context will be passed as "Context:"
user question will be passed as "Question:"
Instructions:
1. Read the context carefully and identify the most relevant information
2. Use only the information provided in the context to answer the question
3. If the context doesn't contain enough information to answer the question, say so
4. Structure your response in a clear and organized manner
5. Cite specific parts of the context when relevant
6. Be concise but comprehensive
Context: {context}""",
},
{"role": "user", "content": prompt},
],
)
for chunk in response:
if chunk.choices[0].delta.content is not None:
yield chunk.choices[0].delta.content
def re_rank_cross_encoders(documents: list[str], prompt: str) -> tuple[str, list[int]]:
"""Re-ranks documents using a cross-encoder model for more accurate relevance scoring.
Uses the MS MARCO MiniLM cross-encoder model to re-rank the input documents based on
their relevance to the given prompt. This helps improve the quality of retrieved
documents for RAG applications.
Args:
documents: List of document texts to re-rank
prompt: The query prompt to rank documents against
Returns:
tuple: A tuple containing:
- The most relevant document text
- List of document indices in order of relevance
Raises:
RuntimeError: If cross-encoder model fails to load or rank documents
"""
if not documents:
return "", []
try:
encoder_model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
ranks = encoder_model.rank(prompt, documents, top_k=3)
relevant_text_ids = []
relevant_text = ""
for rank in ranks:
relevant_text_ids.append(rank["corpus_id"])
if not relevant_text:
relevant_text = documents[rank["corpus_id"]]
return relevant_text, relevant_text_ids
except Exception as e:
print(f"Error in cross-encoder ranking: {e}")
# Fallback to returning first document
return documents[0] if documents else "", [0]

59
src/drakkenheim/storage.py Normal file
View File

@ -0,0 +1,59 @@
"""
Storage module for Drakkenheim RAG Application.
Handles persistent storage of processed files metadata and session state.
"""
import os
import json
from datetime import datetime
import streamlit as st
def save_processed_files():
"""Save processed files metadata to persistent storage."""
try:
files_data = {
"processed_files": st.session_state.processed_files,
"documents_loaded": st.session_state.documents_loaded,
"timestamp": datetime.now().isoformat()
}
# Create metadata directory if it doesn't exist
metadata_path = os.getenv("PROCESSED_FILES_PATH", "./demo-rag-chroma/metadata")
os.makedirs(metadata_path, exist_ok=True)
with open(os.path.join(metadata_path, "processed_files.json"), "w") as f:
json.dump(files_data, f, indent=2)
except Exception as e:
print(f"Error saving processed files: {e}")
def load_processed_files():
"""Load processed files metadata from persistent storage."""
try:
metadata_path = os.getenv("PROCESSED_FILES_PATH", "./demo-rag-chroma/metadata")
metadata_file = os.path.join(metadata_path, "processed_files.json")
if os.path.exists(metadata_file):
with open(metadata_file, "r") as f:
files_data = json.load(f)
st.session_state.processed_files = files_data.get("processed_files", [])
st.session_state.documents_loaded = files_data.get("documents_loaded", False)
st.session_state.files_loaded = True
# Check if the vector collection actually exists
try:
from .llm import get_vector_collection
collection = get_vector_collection()
if collection.count() == 0:
# Collection is empty, reset the state
st.session_state.documents_loaded = False
st.session_state.processed_files = []
except Exception:
# Collection doesn't exist or is corrupted, reset the state
st.session_state.documents_loaded = False
st.session_state.processed_files = []
else:
st.session_state.files_loaded = True
except Exception as e:
print(f"Error loading processed files: {e}")
st.session_state.files_loaded = True

210
src/drakkenheim/ui.py Normal file
View File

@ -0,0 +1,210 @@
"""
UI Components module for Drakkenheim RAG Application.
Handles the main user interface components and layout.
"""
import streamlit as st
from datetime import datetime
from .auth import logout
from .llm import query_collection, call_llm, re_rank_cross_encoders, add_to_vector_collection
from .document import process_document, normalize_uploaded_file_name
from .storage import save_processed_files
def render_sidebar():
"""Render the sidebar with document upload and file management."""
with st.sidebar:
# Logout button at the top
col1, col2 = st.columns([3, 1])
with col1:
st.title("📚 Document Upload")
with col2:
if st.button("🚪 Logout", help="Logout from the application"):
logout()
st.divider()
uploaded_files = st.file_uploader(
"Upload your documents",
type=["pdf", "txt", "json", "md"],
accept_multiple_files=True,
help="Supported formats: PDF, Text, JSON, Markdown. You can select multiple files at once!"
)
if st.button("📥 Process Documents", type="primary"):
if uploaded_files:
try:
# Create progress elements
progress_bar = st.progress(0)
status_text = st.empty()
file_progress_text = st.empty()
total_files = len(uploaded_files)
processed_count = 0
total_chunks_processed = 0
# Estimate total chunks for better progress tracking
total_chunks_estimated = 0
for uploaded_file in uploaded_files:
try:
# Quick estimation of chunks (rough estimate)
content_length = len(uploaded_file.getvalue())
estimated_chunks = max(1, content_length // 200) # Rough estimate based on 200 char chunks
total_chunks_estimated += estimated_chunks
except Exception:
total_chunks_estimated += 10 # Fallback estimate
for i, uploaded_file in enumerate(uploaded_files):
# File-level progress
file_progress_text.text(f"📁 File {i+1}/{total_files}: {uploaded_file.name}")
try:
# Process document
all_splits = process_document(uploaded_file)
actual_chunks = len(all_splits)
# Add to vector collection with chunk-level progress
add_to_vector_collection(all_splits, normalize_uploaded_file_name(uploaded_file))
st.session_state.documents_loaded = True
# Add file to processed files list
current_time = datetime.now().strftime("%H:%M:%S")
file_info = {
"name": uploaded_file.name,
"size": f"{len(uploaded_file.getvalue()) / 1024:.1f} KB",
"chunks": actual_chunks,
"timestamp": current_time
}
st.session_state.processed_files.append(file_info)
save_processed_files()
# Update status with chunk info
status_text.text(f"✅ {uploaded_file.name} - {actual_chunks} chunks processed")
except Exception as e:
st.error(f"❌ Error processing {uploaded_file.name}: {str(e)}")
# Final status
progress_bar.progress(1.0)
if processed_count == total_files:
status_text.text(f"🎉 All {processed_count} files processed! ({total_chunks_processed} total chunks)")
file_progress_text.text("✅ Processing complete!")
st.success(f"🎉 Successfully processed {processed_count} file(s) with {total_chunks_processed} total chunks!")
else:
status_text.text(f"⚠️ Processed {processed_count}/{total_files} files ({total_chunks_processed} chunks)")
file_progress_text.text("⚠️ Partial completion")
st.warning(f"Processed {processed_count} out of {total_files} files successfully")
# Clear progress after 3 seconds
import time
time.sleep(3)
progress_bar.empty()
status_text.empty()
file_progress_text.empty()
except Exception as e:
st.error(f"Error processing files: {str(e)}")
else:
st.warning("Please select at least one file to process.")
# Clear buttons
col1, col2 = st.columns(2)
with col1:
if st.button("🗑️ Clear Chat"):
st.session_state.messages = []
st.rerun()
with col2:
if st.button("🗑️ Clear All Files"):
st.session_state.processed_files = []
st.session_state.documents_loaded = False
st.session_state.messages = []
save_processed_files()
st.rerun()
# Processed Files Overview
st.divider()
st.subheader("📋 Processed Files")
if st.session_state.processed_files:
# Show total files and chunks
total_chunks = sum(file_info['chunks'] for file_info in st.session_state.processed_files)
st.metric("Total Files", len(st.session_state.processed_files))
st.metric("Total Chunks", total_chunks)
# List individual files
for i, file_info in enumerate(st.session_state.processed_files):
with st.expander(f"📄 {file_info['name']}"):
col1, col2 = st.columns(2)
with col1:
st.write(f"**Size:** {file_info['size']}")
st.write(f"**Chunks:** {file_info['chunks']}")
with col2:
st.write(f"**Processed:** {file_info['timestamp']}")
# Remove file button
if st.button("🗑️ Remove", key=f"remove_{i}"):
st.session_state.processed_files.pop(i)
if not st.session_state.processed_files:
st.session_state.documents_loaded = False
save_processed_files()
st.rerun()
else:
st.info("No files processed yet")
# Document status
st.divider()
if st.session_state.documents_loaded:
st.success(f"📄 {len(st.session_state.processed_files)} file(s) loaded and ready for questions!")
else:
st.info("📄 Upload a document to start chatting")
def render_chat_interface():
"""Render the main chat interface."""
# Display chat messages
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# Chat input
if prompt := st.chat_input("Ask a question about your documents..."):
# Add user message to chat history
st.session_state.messages.append({"role": "user", "content": prompt})
# Display user message
with st.chat_message("user"):
st.markdown(prompt)
# Generate and display assistant response
with st.chat_message("assistant"):
try:
# Query the vector collection
results = query_collection(prompt)
relevant_text = results.get("documents", [[]])[0]
relevant_text_ids = list(range(len(relevant_text)))
# Re-rank results for better relevance
if relevant_text:
relevant_text, relevant_text_ids = re_rank_cross_encoders(relevant_text, prompt)
# Prepare context for LLM
context = "\n\n".join(relevant_text) if relevant_text else "No relevant context found."
# Generate response with streaming
message_placeholder = st.empty()
full_response = ""
for chunk in call_llm(context, prompt):
full_response += chunk
message_placeholder.markdown(full_response + "▌")
# Remove the cursor and display final response
message_placeholder.markdown(full_response)
# Add assistant response to chat history
st.session_state.messages.append({"role": "assistant", "content": full_response})
except Exception as e:
error_message = f"Sorry, I encountered an error: {str(e)}"
st.error(error_message)
st.session_state.messages.append({"role": "assistant", "content": error_message})