first commit
This commit is contained in:
commit
8e5a46aa59
24 changed files with 3254 additions and 0 deletions
22
.env.example
Normal file
22
.env.example
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
# Ollama Configuration
|
||||
OLLAMA_HOST=http://ollama:11434
|
||||
OLLAMA_MODEL=llama3.1
|
||||
OLLAMA_EMBED_MODEL=nomic-embed-text
|
||||
|
||||
# ChromaDB Configuration
|
||||
CHROMA_HOST=chromadb
|
||||
CHROMA_PORT=8000
|
||||
|
||||
# MCP Server Configuration
|
||||
MCP_SERVER_HOST=0.0.0.0
|
||||
MCP_SERVER_PORT=8051
|
||||
|
||||
# RAG Configuration
|
||||
CHUNK_SIZE=1024
|
||||
CHUNK_OVERLAP=200
|
||||
TOP_K_RETRIEVAL=5
|
||||
RERANK_TOP_N=3
|
||||
|
||||
# Storage Configuration
|
||||
UPLOAD_DIR=/workspace/data/uploads
|
||||
|
||||
63
.gitignore
vendored
Normal file
63
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# Virtual Environment
|
||||
venv/
|
||||
env/
|
||||
ENV/
|
||||
.venv
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
.DS_Store
|
||||
|
||||
# Environment Variables
|
||||
.env
|
||||
|
||||
# Data Directories
|
||||
data/chroma/
|
||||
data/uploads/
|
||||
*.db
|
||||
*.sqlite
|
||||
|
||||
# Docker
|
||||
.dockerignore
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
logs/
|
||||
|
||||
# Streamlit
|
||||
.streamlit/secrets.toml
|
||||
|
||||
# Temporary Files
|
||||
tmp/
|
||||
temp/
|
||||
*.tmp
|
||||
|
||||
# OS
|
||||
Thumbs.db
|
||||
|
||||
17
.streamlit/config.toml
Normal file
17
.streamlit/config.toml
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
[theme]
|
||||
# Archon-inspired dark theme with neon accents
|
||||
base = "dark"
|
||||
primaryColor = "#E91E63" # Neon pink
|
||||
backgroundColor = "#0E0E0E" # Pure black
|
||||
secondaryBackgroundColor = "#1A1A1A" # Dark gray
|
||||
textColor = "#FFFFFF" # White
|
||||
font = "monospace"
|
||||
|
||||
[server]
|
||||
headless = true
|
||||
port = 8501
|
||||
enableCORS = false
|
||||
|
||||
[browser]
|
||||
gatherUsageStats = false
|
||||
|
||||
41
Dockerfile
Normal file
41
Dockerfile
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
FROM python:3.11-slim
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /workspace
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
curl \
|
||||
build-essential \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy requirements first for layer caching
|
||||
COPY requirements.txt .
|
||||
|
||||
# Install Python dependencies
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy application code
|
||||
COPY ./app /workspace/app
|
||||
COPY ./mcp_server /workspace/mcp_server
|
||||
COPY ./.streamlit /workspace/.streamlit
|
||||
|
||||
# Create data directories
|
||||
RUN mkdir -p /workspace/data/uploads
|
||||
|
||||
# Copy startup script
|
||||
COPY start.sh /workspace/start.sh
|
||||
RUN chmod +x /workspace/start.sh
|
||||
|
||||
# Expose ports
|
||||
EXPOSE 8501 8051
|
||||
|
||||
# Set environment variables
|
||||
ENV STREAMLIT_SERVER_PORT=8501
|
||||
ENV STREAMLIT_SERVER_ADDRESS=0.0.0.0
|
||||
ENV MCP_SERVER_PORT=8051
|
||||
ENV PYTHONPATH=/workspace:$PYTHONPATH
|
||||
|
||||
# Run startup script
|
||||
CMD ["/workspace/start.sh"]
|
||||
|
||||
429
README.md
Normal file
429
README.md
Normal file
|
|
@ -0,0 +1,429 @@
|
|||
# 🔥 Code-Forge
|
||||
|
||||
**Localized Multi-Knowledge-Base RAG System with MCP Integration**
|
||||
|
||||
A self-hosted, privacy-first RAG (Retrieval-Augmented Generation) application designed for development teams. Features a badass Archon-inspired UI and exposes knowledge bases via Model Context Protocol (MCP) for AI agent integration.
|
||||
|
||||
---
|
||||
|
||||
## ✨ Features
|
||||
|
||||
### 🎯 Core Capabilities
|
||||
- **Multiple Knowledge Bases**: Create isolated KBs for different projects
|
||||
- **Advanced RAG Strategies**:
|
||||
- 🔍 Vector Search (semantic similarity)
|
||||
- 🔀 Hybrid Search (vector + BM25 keyword)
|
||||
- 🎯 Reranking (cross-encoder scoring)
|
||||
- 🤖 Agentic RAG (multi-step reasoning)
|
||||
- **MCP Integration**: Expose KBs to AI agents like Claude Desktop
|
||||
- **Dual Interface**: Streamlit UI for humans + MCP server for agents
|
||||
|
||||
### 🎨 UI/UX
|
||||
- **Archon-Inspired Design**: Dark theme with neon accents
|
||||
- **Color Scheme**:
|
||||
- Primary: `hsl(271, 91%, 65%)` (Purple)
|
||||
- Accent: `hsl(330, 90%, 65%)` (Pink)
|
||||
- Success: `hsl(160, 84%, 39%)` (Green)
|
||||
- Background: Pure black `#0E0E0E`
|
||||
- **Real-time Health Monitoring**: Ollama & ChromaDB status
|
||||
- **Document Management**: Upload files or import entire directories
|
||||
|
||||
### 📁 Supported File Types
|
||||
- **Documents**: `.md`, `.txt`, `.pdf`
|
||||
- **Code**: `.py`, `.js`, `.jsx`, `.ts`, `.tsx`, `.go`, `.rs`, `.java`, `.cpp`, `.c`, `.h`
|
||||
- **Config**: `.json`, `.yaml`, `.yml`
|
||||
|
||||
---
|
||||
|
||||
## 🏗️ Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────┐
|
||||
│ Code-Forge │
|
||||
├─────────────────────────────────────────────────────────┤
|
||||
│ │
|
||||
│ ┌──────────────┐ ┌──────────────┐ │
|
||||
│ │ Streamlit │ │ MCP Server │ │
|
||||
│ │ UI (8501) │ │ (8051) │ │
|
||||
│ └──────┬───────┘ └──────┬───────┘ │
|
||||
│ │ │ │
|
||||
│ └─────────────┬───────────────┘ │
|
||||
│ │ │
|
||||
│ ┌────────▼────────┐ │
|
||||
│ │ RAG Engine │ │
|
||||
│ │ (LlamaIndex) │ │
|
||||
│ └────────┬────────┘ │
|
||||
│ │ │
|
||||
│ ┌─────────────┴─────────────┐ │
|
||||
│ │ │ │
|
||||
│ ┌────▼─────┐ ┌──────▼──────┐ │
|
||||
│ │ ChromaDB │ │ Ollama │ │
|
||||
│ │ (8000) │ │ (11434) │ │
|
||||
│ └──────────┘ └─────────────┘ │
|
||||
│ │
|
||||
└─────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Services
|
||||
1. **Ollama** (port 11434): Local LLM and embeddings
|
||||
2. **ChromaDB** (port 8000): Vector database
|
||||
3. **App Container**:
|
||||
- Streamlit UI (port 8501)
|
||||
- MCP Server (port 8051)
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Quick Start
|
||||
|
||||
### Prerequisites
|
||||
- **Docker** and **Docker Compose** installed
|
||||
- **Ollama** installed locally (or use Docker Ollama)
|
||||
- At least 8GB RAM recommended
|
||||
|
||||
### 1. Clone and Navigate
|
||||
```bash
|
||||
git clone <your-repo-url>
|
||||
cd code-forge
|
||||
```
|
||||
|
||||
### 2. Pull Required Ollama Models
|
||||
```bash
|
||||
# Pull LLM model
|
||||
ollama pull llama3.1
|
||||
|
||||
# Pull embedding model
|
||||
ollama pull nomic-embed-text
|
||||
```
|
||||
|
||||
### 3. One-Command Deployment
|
||||
```bash
|
||||
docker compose up --build
|
||||
```
|
||||
|
||||
This will:
|
||||
- Start Ollama service
|
||||
- Start ChromaDB service
|
||||
- Build and start the Code-Forge application
|
||||
- Wait for all services to become healthy
|
||||
|
||||
### 4. Access the Application
|
||||
|
||||
**Streamlit UI**: http://localhost:8501
|
||||
|
||||
**MCP Server**: http://localhost:8051/mcp
|
||||
|
||||
---
|
||||
|
||||
## 📚 Usage Guide
|
||||
|
||||
### Creating a Knowledge Base
|
||||
|
||||
1. Open the Streamlit UI at http://localhost:8501
|
||||
2. In the sidebar, expand **"➕ Create New KB"**
|
||||
3. Enter a name (e.g., `my-project`)
|
||||
4. Click **"Create"**
|
||||
|
||||
### Uploading Documents
|
||||
|
||||
**Option 1: File Upload**
|
||||
1. Select your KB from the sidebar
|
||||
2. Go to the **"📚 Knowledge Base"** tab
|
||||
3. Click **"Choose files"** and select documents
|
||||
4. Click **"Upload Files"**
|
||||
|
||||
**Option 2: Directory Import**
|
||||
1. Select your KB from the sidebar
|
||||
2. Go to the **"📚 Knowledge Base"** tab
|
||||
3. Enter the full path to your project directory
|
||||
4. Click **"Import Directory"**
|
||||
|
||||
### Chatting with Your KB
|
||||
|
||||
1. Select your KB from the sidebar
|
||||
2. Go to the **"💬 Chat"** tab
|
||||
3. (Optional) Enable advanced RAG strategies:
|
||||
- **🔀 Hybrid Search**: Better for keyword-heavy queries
|
||||
- **🎯 Reranking**: Improves relevance of top results
|
||||
- **🤖 Agentic RAG**: Best for complex, multi-part questions
|
||||
4. Type your question and press Enter
|
||||
5. View sources by expanding **"📚 View Sources"**
|
||||
|
||||
### Sharing with Your Team
|
||||
|
||||
Want to share Code-Forge with coworkers? See **[SHARING.md](SHARING.md)** for:
|
||||
- ✅ How to share the repository
|
||||
- ✅ What your coworker needs to install
|
||||
- ✅ Quick start guide for new users
|
||||
- ✅ Exporting/importing knowledge bases
|
||||
- ✅ Network sharing options
|
||||
- ✅ Troubleshooting guide
|
||||
|
||||
**TL;DR:** Share the repo → They run `docker compose up -d` → Pull models → Done! 🎉
|
||||
|
||||
---
|
||||
|
||||
## 🔗 MCP Integration
|
||||
|
||||
### Connecting Claude Desktop
|
||||
|
||||
1. Ensure the MCP server is running (it starts automatically with `docker compose up`)
|
||||
2. Get the MCP endpoint URL from the Streamlit sidebar:
|
||||
```
|
||||
http://localhost:8051/mcp
|
||||
```
|
||||
3. Add to Claude Desktop:
|
||||
```bash
|
||||
claude mcp add --transport http code-forge http://localhost:8051/mcp
|
||||
```
|
||||
|
||||
### Available MCP Tools
|
||||
|
||||
#### `search_knowledge_base`
|
||||
Search a KB using RAG with optional advanced features.
|
||||
|
||||
**Parameters:**
|
||||
- `kb_name` (string): Name of the knowledge base
|
||||
- `query` (string): Search query
|
||||
- `use_hybrid` (boolean): Enable hybrid search
|
||||
- `use_rerank` (boolean): Enable reranking
|
||||
- `use_agentic` (boolean): Enable agentic RAG
|
||||
|
||||
**Example:**
|
||||
```json
|
||||
{
|
||||
"kb_name": "my-project",
|
||||
"query": "How does authentication work?",
|
||||
"use_hybrid": true,
|
||||
"use_rerank": true
|
||||
}
|
||||
```
|
||||
|
||||
#### `list_knowledge_bases`
|
||||
List all available knowledge bases.
|
||||
|
||||
**Returns:** Array of KB names
|
||||
|
||||
#### `get_kb_stats`
|
||||
Get statistics for a knowledge base.
|
||||
|
||||
**Parameters:**
|
||||
- `kb_name` (string): Name of the knowledge base
|
||||
|
||||
**Returns:**
|
||||
```json
|
||||
{
|
||||
"total_documents": 42,
|
||||
"total_chunks": 387,
|
||||
"last_updated": "2025-10-22T10:30:00"
|
||||
}
|
||||
```
|
||||
|
||||
#### `list_documents`
|
||||
List all documents in a knowledge base.
|
||||
|
||||
**Parameters:**
|
||||
- `kb_name` (string): Name of the knowledge base
|
||||
|
||||
**Returns:** Array of document metadata objects
|
||||
|
||||
---
|
||||
|
||||
## ⚙️ Configuration
|
||||
|
||||
### Environment Variables
|
||||
|
||||
Create a `.env` file to customize settings:
|
||||
|
||||
```bash
|
||||
# Ollama Configuration
|
||||
OLLAMA_HOST=http://ollama:11434
|
||||
OLLAMA_MODEL=llama3.1
|
||||
OLLAMA_EMBED_MODEL=nomic-embed-text
|
||||
|
||||
# ChromaDB Configuration
|
||||
CHROMA_HOST=chromadb
|
||||
CHROMA_PORT=8000
|
||||
|
||||
# MCP Server Configuration
|
||||
MCP_SERVER_HOST=0.0.0.0
|
||||
MCP_SERVER_PORT=8051
|
||||
|
||||
# RAG Configuration
|
||||
CHUNK_SIZE=1024
|
||||
CHUNK_OVERLAP=200
|
||||
TOP_K_RETRIEVAL=5
|
||||
RERANK_TOP_N=3
|
||||
|
||||
# Storage
|
||||
UPLOAD_DIR=/workspace/data/uploads
|
||||
```
|
||||
|
||||
### Changing Models
|
||||
|
||||
To use different Ollama models:
|
||||
|
||||
1. Pull the model:
|
||||
```bash
|
||||
ollama pull mistral
|
||||
```
|
||||
|
||||
2. Update `.env`:
|
||||
```bash
|
||||
OLLAMA_MODEL=mistral
|
||||
```
|
||||
|
||||
3. Restart the application:
|
||||
```bash
|
||||
docker compose restart app
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🛠️ Development
|
||||
|
||||
### Project Structure
|
||||
|
||||
```
|
||||
code-forge/
|
||||
├── app/
|
||||
│ ├── core/
|
||||
│ │ ├── config.py # Configuration management
|
||||
│ │ ├── health.py # Service health checks
|
||||
│ │ ├── chroma_manager.py # ChromaDB operations
|
||||
│ │ ├── ingestion.py # Document processing
|
||||
│ │ ├── kb_metadata.py # KB metadata management
|
||||
│ │ └── rag_engine.py # RAG implementation
|
||||
│ └── main.py # Streamlit application
|
||||
├── mcp_server/
|
||||
│ └── server.py # MCP server (FastMCP + FastAPI)
|
||||
├── .streamlit/
|
||||
│ └── config.toml # Streamlit theme
|
||||
├── data/
|
||||
│ ├── chroma/ # ChromaDB persistent storage
|
||||
│ └── uploads/ # Uploaded documents
|
||||
├── docker-compose.yml # Service orchestration
|
||||
├── Dockerfile # Application container
|
||||
├── requirements.txt # Python dependencies
|
||||
├── start.sh # Startup script
|
||||
└── README.md # This file
|
||||
```
|
||||
|
||||
### Running Locally (Without Docker)
|
||||
|
||||
1. Install dependencies:
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
2. Start Ollama and ChromaDB separately
|
||||
|
||||
3. Start the MCP server:
|
||||
```bash
|
||||
python mcp_server/server.py
|
||||
```
|
||||
|
||||
4. Start Streamlit:
|
||||
```bash
|
||||
streamlit run app/main.py
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🐛 Troubleshooting
|
||||
|
||||
### Services Not Starting
|
||||
|
||||
**Check service health:**
|
||||
```bash
|
||||
docker compose ps
|
||||
```
|
||||
|
||||
**View logs:**
|
||||
```bash
|
||||
docker compose logs ollama
|
||||
docker compose logs chromadb
|
||||
docker compose logs app
|
||||
```
|
||||
|
||||
### Ollama Models Not Found
|
||||
|
||||
**Pull models manually:**
|
||||
```bash
|
||||
docker compose exec ollama ollama pull llama3.1
|
||||
docker compose exec ollama ollama pull nomic-embed-text
|
||||
```
|
||||
|
||||
### ChromaDB Connection Issues
|
||||
|
||||
**Restart ChromaDB:**
|
||||
```bash
|
||||
docker compose restart chromadb
|
||||
```
|
||||
|
||||
**Check ChromaDB health:**
|
||||
```bash
|
||||
curl http://localhost:8000/api/v1/heartbeat
|
||||
```
|
||||
|
||||
### MCP Connection Fails
|
||||
|
||||
**Verify MCP server is running:**
|
||||
```bash
|
||||
curl http://localhost:8051/health
|
||||
```
|
||||
|
||||
**Check MCP logs:**
|
||||
```bash
|
||||
docker compose logs app | grep MCP
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📊 Performance Tips
|
||||
|
||||
1. **Chunk Size**: Larger chunks (1024-2048 tokens) for technical docs, smaller (512-1024) for code
|
||||
2. **Hybrid Search**: Best for queries with specific keywords or technical terms
|
||||
3. **Reranking**: Adds latency but significantly improves relevance
|
||||
4. **Agentic RAG**: Use for complex questions requiring multi-step reasoning (slower)
|
||||
|
||||
---
|
||||
|
||||
## 🔒 Privacy & Security
|
||||
|
||||
- **100% Local**: All data stays on your machine
|
||||
- **No External APIs**: Ollama runs models locally
|
||||
- **No Telemetry**: ChromaDB telemetry disabled
|
||||
- **Isolated KBs**: Each project has its own knowledge base
|
||||
|
||||
---
|
||||
|
||||
## 📝 License
|
||||
|
||||
MIT License - See LICENSE file for details
|
||||
|
||||
---
|
||||
|
||||
## 🙏 Acknowledgments
|
||||
|
||||
- **LlamaIndex**: RAG framework
|
||||
- **Ollama**: Local LLM runtime
|
||||
- **ChromaDB**: Vector database
|
||||
- **Streamlit**: UI framework
|
||||
- **FastMCP**: MCP server library
|
||||
- **Archon**: Design inspiration
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Roadmap
|
||||
|
||||
- [ ] Multi-modal support (images, diagrams)
|
||||
- [ ] Graph RAG for relationship extraction
|
||||
- [ ] Query caching for faster responses
|
||||
- [ ] Batch document processing
|
||||
- [ ] Export/import knowledge bases
|
||||
- [ ] Advanced analytics dashboard
|
||||
|
||||
---
|
||||
|
||||
**Built with 🔥 by The Augster**
|
||||
|
||||
302
SHARING.md
Normal file
302
SHARING.md
Normal file
|
|
@ -0,0 +1,302 @@
|
|||
# 🚀 Sharing Code-Forge with Your Team
|
||||
|
||||
This guide explains how to share Code-Forge with coworkers so they can run it on their own machines.
|
||||
|
||||
## 📦 What Your Coworker Needs
|
||||
|
||||
### Prerequisites
|
||||
1. **Docker Desktop** (includes Docker Compose)
|
||||
- Mac: https://docs.docker.com/desktop/install/mac-install/
|
||||
- Windows: https://docs.docker.com/desktop/install/windows-install/
|
||||
- Linux: https://docs.docker.com/desktop/install/linux-install/
|
||||
|
||||
2. **Git** (to clone the repository)
|
||||
- Mac: Pre-installed or via Homebrew: `brew install git`
|
||||
- Windows: https://git-scm.com/download/win
|
||||
- Linux: `sudo apt install git` or `sudo yum install git`
|
||||
|
||||
That's it! No Python, Node.js, or other dependencies needed - Docker handles everything.
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Quick Start for Your Coworker
|
||||
|
||||
### Option 1: Share via Git Repository (Recommended)
|
||||
|
||||
**If you have a private Git repo:**
|
||||
|
||||
1. **Push your code to the repo:**
|
||||
```bash
|
||||
git add .
|
||||
git commit -m "Add Code-Forge RAG system"
|
||||
git push origin main
|
||||
```
|
||||
|
||||
2. **Share the repo URL with your coworker**
|
||||
|
||||
3. **Your coworker clones and runs:**
|
||||
```bash
|
||||
# Clone the repository
|
||||
git clone <your-repo-url>
|
||||
cd code-forge
|
||||
|
||||
# Start the system
|
||||
docker compose up -d
|
||||
|
||||
# Pull Ollama models (one-time setup)
|
||||
docker compose exec ollama ollama pull nomic-embed-text
|
||||
docker compose exec ollama ollama pull llama3.1
|
||||
```
|
||||
|
||||
4. **Access the UI:**
|
||||
- Open browser: http://localhost:8501
|
||||
- Done! 🎉
|
||||
|
||||
---
|
||||
|
||||
### Option 2: Share as ZIP File
|
||||
|
||||
**If you don't have a Git repo:**
|
||||
|
||||
1. **Create a ZIP of your project:**
|
||||
```bash
|
||||
cd /Users/iamanmp/Projects
|
||||
zip -r code-forge.zip code-forge \
|
||||
-x "code-forge/data/*" \
|
||||
-x "code-forge/.git/*" \
|
||||
-x "code-forge/__pycache__/*" \
|
||||
-x "code-forge/**/__pycache__/*"
|
||||
```
|
||||
|
||||
2. **Share `code-forge.zip` via:**
|
||||
- Email (if < 25MB)
|
||||
- Google Drive / Dropbox / OneDrive
|
||||
- Slack / Teams file sharing
|
||||
|
||||
3. **Your coworker extracts and runs:**
|
||||
```bash
|
||||
# Extract the ZIP
|
||||
unzip code-forge.zip
|
||||
cd code-forge
|
||||
|
||||
# Start the system
|
||||
docker compose up -d
|
||||
|
||||
# Pull Ollama models (one-time setup)
|
||||
docker compose exec ollama ollama pull nomic-embed-text
|
||||
docker compose exec ollama ollama pull llama3.1
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🔧 What Gets Shared vs. What Doesn't
|
||||
|
||||
### ✅ Shared (in the repository/ZIP)
|
||||
- **All source code** (`app/`, `mcp_server/`, etc.)
|
||||
- **Configuration files** (`docker-compose.yml`, `Dockerfile`, etc.)
|
||||
- **Requirements** (`requirements.txt`)
|
||||
- **Styling** (`.streamlit/config.toml`)
|
||||
|
||||
### ❌ NOT Shared (stays local)
|
||||
- **Ollama models** (~5GB) - Each person downloads their own
|
||||
- **ChromaDB data** (knowledge bases) - Each person creates their own
|
||||
- **Uploaded files** - Each person uploads their own documents
|
||||
- **Docker volumes** - Created fresh on each machine
|
||||
|
||||
---
|
||||
|
||||
## 📊 Data Sharing Options
|
||||
|
||||
### If You Want to Share Knowledge Bases Too:
|
||||
|
||||
**Option A: Export/Import Knowledge Bases**
|
||||
|
||||
Currently, knowledge bases are stored in Docker volumes. To share them:
|
||||
|
||||
1. **Export ChromaDB data:**
|
||||
```bash
|
||||
# Create a backup directory
|
||||
mkdir -p backups
|
||||
|
||||
# Copy ChromaDB data from Docker volume
|
||||
docker compose cp chromadb:/chroma/chroma ./backups/chroma_data
|
||||
```
|
||||
|
||||
2. **Share the backup:**
|
||||
- Add `backups/` to your ZIP or Git repo
|
||||
- Your coworker copies it to their ChromaDB volume:
|
||||
```bash
|
||||
docker compose up -d chromadb
|
||||
docker compose cp ./backups/chroma_data chromadb:/chroma/chroma
|
||||
docker compose restart chromadb app
|
||||
```
|
||||
|
||||
**Option B: Share Original Documents**
|
||||
|
||||
Simpler approach:
|
||||
1. Create a `shared_docs/` folder in the repo
|
||||
2. Add your markdown/PDF files there
|
||||
3. Your coworker uploads them via the UI or directory import
|
||||
|
||||
---
|
||||
|
||||
## 🌐 Network Sharing (Same Office/VPN)
|
||||
|
||||
If you want your coworker to access **your running instance** instead of running their own:
|
||||
|
||||
### 1. Find Your Local IP Address
|
||||
|
||||
**Mac/Linux:**
|
||||
```bash
|
||||
ifconfig | grep "inet " | grep -v 127.0.0.1
|
||||
```
|
||||
|
||||
**Windows:**
|
||||
```cmd
|
||||
ipconfig
|
||||
```
|
||||
|
||||
Look for something like `192.168.1.100`
|
||||
|
||||
### 2. Update docker-compose.yml
|
||||
|
||||
Change the ports to bind to all interfaces:
|
||||
|
||||
```yaml
|
||||
services:
|
||||
app:
|
||||
ports:
|
||||
- "0.0.0.0:8501:8501" # Streamlit UI
|
||||
- "0.0.0.0:8051:8051" # MCP Server
|
||||
```
|
||||
|
||||
### 3. Restart the containers
|
||||
|
||||
```bash
|
||||
docker compose down
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
### 4. Share the URL
|
||||
|
||||
Your coworker can access:
|
||||
- **UI**: `http://<your-ip>:8501`
|
||||
- **MCP**: `http://<your-ip>:8051/mcp`
|
||||
|
||||
**⚠️ Security Note:** This exposes your instance to your local network. Only do this on trusted networks (office LAN, VPN). For internet sharing, use a reverse proxy with authentication.
|
||||
|
||||
---
|
||||
|
||||
## 🔐 Security Considerations
|
||||
|
||||
### For Local Sharing (Same Machine)
|
||||
- ✅ Safe - Everything runs locally
|
||||
- ✅ No authentication needed
|
||||
- ✅ Data stays on the machine
|
||||
|
||||
### For Network Sharing (Same Office/VPN)
|
||||
- ⚠️ Use only on trusted networks
|
||||
- ⚠️ No built-in authentication
|
||||
- ⚠️ Anyone on the network can access
|
||||
|
||||
### For Internet Sharing
|
||||
- ❌ **NOT RECOMMENDED** without additional security
|
||||
- If needed, use:
|
||||
- Reverse proxy (nginx, Caddy)
|
||||
- Authentication (OAuth, basic auth)
|
||||
- HTTPS/SSL certificates
|
||||
- Firewall rules
|
||||
|
||||
---
|
||||
|
||||
## 🐛 Troubleshooting for Your Coworker
|
||||
|
||||
### "Docker command not found"
|
||||
- Install Docker Desktop: https://www.docker.com/products/docker-desktop/
|
||||
|
||||
### "Port 8501 already in use"
|
||||
- Another app is using that port
|
||||
- Change ports in `docker-compose.yml`:
|
||||
```yaml
|
||||
ports:
|
||||
- "8502:8501" # Use 8502 instead
|
||||
```
|
||||
|
||||
### "Ollama models not found"
|
||||
- They forgot to pull the models:
|
||||
```bash
|
||||
docker compose exec ollama ollama pull nomic-embed-text
|
||||
docker compose exec ollama ollama pull llama3.1
|
||||
```
|
||||
|
||||
### "ChromaDB connection error"
|
||||
- Wait 10 seconds for services to start
|
||||
- Check status: `docker compose ps`
|
||||
- Restart: `docker compose restart`
|
||||
|
||||
### "Out of disk space"
|
||||
- Ollama models are ~5GB
|
||||
- Docker images are ~2GB
|
||||
- Free up space or use external drive
|
||||
|
||||
---
|
||||
|
||||
## 📝 Quick Reference Commands
|
||||
|
||||
```bash
|
||||
# Start everything
|
||||
docker compose up -d
|
||||
|
||||
# Stop everything
|
||||
docker compose down
|
||||
|
||||
# View logs
|
||||
docker compose logs -f app
|
||||
|
||||
# Restart a service
|
||||
docker compose restart app
|
||||
|
||||
# Check status
|
||||
docker compose ps
|
||||
|
||||
# Pull Ollama models
|
||||
docker compose exec ollama ollama pull nomic-embed-text
|
||||
docker compose exec ollama ollama pull llama3.1
|
||||
|
||||
# Access Ollama shell
|
||||
docker compose exec ollama bash
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🎓 What Your Coworker Should Know
|
||||
|
||||
1. **First startup takes 5-10 minutes** (downloading models)
|
||||
2. **Subsequent startups take ~10 seconds**
|
||||
3. **Each person has their own knowledge bases** (unless you share ChromaDB data)
|
||||
4. **Models run locally** - no API keys or internet needed after setup
|
||||
5. **Everything is containerized** - won't mess up their system
|
||||
|
||||
---
|
||||
|
||||
## 📞 Support
|
||||
|
||||
If your coworker has issues:
|
||||
1. Check the troubleshooting section above
|
||||
2. Run `docker compose logs -f` to see errors
|
||||
3. Verify Docker Desktop is running
|
||||
4. Try `docker compose down && docker compose up -d` (restart)
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Next Steps
|
||||
|
||||
After your coworker gets it running:
|
||||
1. Create a knowledge base
|
||||
2. Upload some documents
|
||||
3. Try the chat interface
|
||||
4. Experiment with RAG strategies (Hybrid, Reranking, Agentic)
|
||||
5. Connect Claude Desktop via MCP (optional)
|
||||
|
||||
Enjoy! 🎉
|
||||
|
||||
0
app/__init__.py
Normal file
0
app/__init__.py
Normal file
0
app/components/__init__.py
Normal file
0
app/components/__init__.py
Normal file
0
app/core/__init__.py
Normal file
0
app/core/__init__.py
Normal file
158
app/core/chroma_manager.py
Normal file
158
app/core/chroma_manager.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
"""
|
||||
ChromaDB client manager for knowledge base operations.
|
||||
Implements singleton pattern for connection reuse.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
from app.core.config import Config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChromaManager:
|
||||
"""Singleton manager for ChromaDB operations."""
|
||||
|
||||
_instance = None
|
||||
_client = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if self._client is None:
|
||||
self._initialize_client()
|
||||
|
||||
def _initialize_client(self):
|
||||
"""Initialize ChromaDB HTTP client."""
|
||||
try:
|
||||
self._client = chromadb.HttpClient(
|
||||
host=Config.CHROMA_HOST,
|
||||
port=Config.CHROMA_PORT
|
||||
)
|
||||
logger.info(f"ChromaDB client initialized: {Config.get_chroma_url()}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize ChromaDB client: {e}")
|
||||
raise
|
||||
|
||||
def get_client(self) -> chromadb.HttpClient:
|
||||
"""Get the ChromaDB client instance."""
|
||||
if self._client is None:
|
||||
self._initialize_client()
|
||||
return self._client
|
||||
|
||||
def list_collections(self) -> List[str]:
|
||||
"""
|
||||
List all collection names.
|
||||
|
||||
Returns:
|
||||
List of collection names
|
||||
"""
|
||||
try:
|
||||
collections = self._client.list_collections()
|
||||
return [col.name for col in collections]
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list collections: {e}")
|
||||
return []
|
||||
|
||||
def create_collection(self, name: str) -> bool:
|
||||
"""
|
||||
Create a new collection (knowledge base).
|
||||
|
||||
Args:
|
||||
name: Collection name
|
||||
|
||||
Returns:
|
||||
bool: True if successful
|
||||
"""
|
||||
try:
|
||||
self._client.get_or_create_collection(name=name)
|
||||
logger.info(f"Collection created: {name}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create collection {name}: {e}")
|
||||
return False
|
||||
|
||||
def delete_collection(self, name: str) -> bool:
|
||||
"""
|
||||
Delete a collection.
|
||||
|
||||
Args:
|
||||
name: Collection name
|
||||
|
||||
Returns:
|
||||
bool: True if successful
|
||||
"""
|
||||
try:
|
||||
self._client.delete_collection(name=name)
|
||||
logger.info(f"Collection deleted: {name}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete collection {name}: {e}")
|
||||
return False
|
||||
|
||||
def get_collection(self, name: str) -> Optional[chromadb.Collection]:
|
||||
"""
|
||||
Get a collection by name.
|
||||
|
||||
Args:
|
||||
name: Collection name
|
||||
|
||||
Returns:
|
||||
Collection object or None if not found
|
||||
"""
|
||||
try:
|
||||
return self._client.get_collection(name=name)
|
||||
except Exception as e:
|
||||
logger.warning(f"Collection {name} not found: {e}")
|
||||
return None
|
||||
|
||||
def get_collection_count(self, name: str) -> int:
|
||||
"""
|
||||
Get the number of documents in a collection.
|
||||
|
||||
Args:
|
||||
name: Collection name
|
||||
|
||||
Returns:
|
||||
int: Document count (0 if collection not found)
|
||||
"""
|
||||
try:
|
||||
collection = self.get_collection(name)
|
||||
if collection:
|
||||
return collection.count()
|
||||
return 0
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get count for {name}: {e}")
|
||||
return 0
|
||||
|
||||
def collection_exists(self, name: str) -> bool:
|
||||
"""
|
||||
Check if a collection exists.
|
||||
|
||||
Args:
|
||||
name: Collection name
|
||||
|
||||
Returns:
|
||||
bool: True if collection exists
|
||||
"""
|
||||
return name in self.list_collections()
|
||||
|
||||
|
||||
# Global singleton instance
|
||||
_chroma_manager = None
|
||||
|
||||
|
||||
def get_chroma_manager() -> ChromaManager:
|
||||
"""Get the global ChromaManager instance."""
|
||||
global _chroma_manager
|
||||
if _chroma_manager is None:
|
||||
_chroma_manager = ChromaManager()
|
||||
return _chroma_manager
|
||||
|
||||
75
app/core/config.py
Normal file
75
app/core/config.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
"""
|
||||
Configuration management for Code-Forge.
|
||||
Centralizes all application settings with environment variable support.
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
|
||||
class Config:
|
||||
"""Central configuration for Code-Forge application."""
|
||||
|
||||
# Ollama Configuration
|
||||
OLLAMA_HOST: str = os.getenv("OLLAMA_HOST", "http://ollama:11434")
|
||||
OLLAMA_MODEL: str = os.getenv("OLLAMA_MODEL", "llama3.2:3b")
|
||||
OLLAMA_EMBED_MODEL: str = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
|
||||
|
||||
# ChromaDB Configuration
|
||||
CHROMA_HOST: str = os.getenv("CHROMA_HOST", "chromadb")
|
||||
CHROMA_PORT: int = int(os.getenv("CHROMA_PORT", "8000"))
|
||||
|
||||
# MCP Server Configuration
|
||||
MCP_SERVER_HOST: str = os.getenv("MCP_SERVER_HOST", "0.0.0.0")
|
||||
MCP_SERVER_PORT: int = int(os.getenv("MCP_SERVER_PORT", "8051"))
|
||||
|
||||
# Supported File Types
|
||||
SUPPORTED_FILE_TYPES: List[str] = [
|
||||
".md", ".txt", ".pdf",
|
||||
".py", ".js", ".jsx", ".ts", ".tsx",
|
||||
".json", ".yaml", ".yml",
|
||||
".go", ".rs", ".java", ".cpp", ".c", ".h"
|
||||
]
|
||||
|
||||
# RAG Configuration
|
||||
CHUNK_SIZE: int = 1024
|
||||
CHUNK_OVERLAP: int = 200
|
||||
TOP_K_RETRIEVAL: int = 5
|
||||
RERANK_TOP_N: int = 3
|
||||
|
||||
# Storage Configuration
|
||||
UPLOAD_DIR: str = os.getenv("UPLOAD_DIR", "/workspace/data/uploads")
|
||||
|
||||
@classmethod
|
||||
def get_ollama_url(cls) -> str:
|
||||
"""Get the full Ollama API URL."""
|
||||
return cls.OLLAMA_HOST
|
||||
|
||||
@classmethod
|
||||
def get_chroma_url(cls) -> str:
|
||||
"""Get the full ChromaDB URL."""
|
||||
return f"http://{cls.CHROMA_HOST}:{cls.CHROMA_PORT}"
|
||||
|
||||
@classmethod
|
||||
def get_mcp_url(cls) -> str:
|
||||
"""Get the full MCP server URL."""
|
||||
return f"http://{cls.MCP_SERVER_HOST}:{cls.MCP_SERVER_PORT}"
|
||||
|
||||
@classmethod
|
||||
def is_supported_file(cls, filename: str) -> bool:
|
||||
"""Check if a file type is supported."""
|
||||
return any(filename.lower().endswith(ext) for ext in cls.SUPPORTED_FILE_TYPES)
|
||||
|
||||
@classmethod
|
||||
def ensure_upload_dir(cls) -> Path:
|
||||
"""Ensure upload directory exists and return Path object."""
|
||||
upload_path = Path(cls.UPLOAD_DIR)
|
||||
upload_path.mkdir(parents=True, exist_ok=True)
|
||||
return upload_path
|
||||
|
||||
135
app/core/health.py
Normal file
135
app/core/health.py
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
"""
|
||||
Health check utilities for monitoring service connectivity.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import Dict
|
||||
|
||||
import requests
|
||||
|
||||
from app.core.config import Config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def check_ollama_health() -> Dict[str, any]:
|
||||
"""
|
||||
Check Ollama service health and available models.
|
||||
|
||||
Returns:
|
||||
dict: Status information with 'status', 'models', and optional 'error'
|
||||
"""
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{Config.OLLAMA_HOST}/api/tags",
|
||||
timeout=5
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
models = [model.get("name", "unknown") for model in data.get("models", [])]
|
||||
|
||||
return {
|
||||
"status": "healthy",
|
||||
"models": models,
|
||||
"url": Config.OLLAMA_HOST
|
||||
}
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
logger.warning(f"Ollama connection failed: {e}")
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"error": "Connection refused - is Ollama running?",
|
||||
"url": Config.OLLAMA_HOST
|
||||
}
|
||||
except requests.exceptions.Timeout as e:
|
||||
logger.warning(f"Ollama timeout: {e}")
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"error": "Request timed out",
|
||||
"url": Config.OLLAMA_HOST
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Ollama health check failed: {e}")
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"error": str(e),
|
||||
"url": Config.OLLAMA_HOST
|
||||
}
|
||||
|
||||
|
||||
def check_chroma_health() -> Dict[str, any]:
|
||||
"""
|
||||
Check ChromaDB service health.
|
||||
|
||||
Returns:
|
||||
dict: Status information with 'status' and optional 'error'
|
||||
"""
|
||||
try:
|
||||
# Use v2 API endpoint for health check
|
||||
response = requests.get(
|
||||
f"{Config.get_chroma_url()}/api/v2/auth/identity",
|
||||
timeout=5
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
return {
|
||||
"status": "healthy",
|
||||
"url": Config.get_chroma_url()
|
||||
}
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
logger.warning(f"ChromaDB connection failed: {e}")
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"error": "Connection refused - is ChromaDB running?",
|
||||
"url": Config.get_chroma_url()
|
||||
}
|
||||
except requests.exceptions.Timeout as e:
|
||||
logger.warning(f"ChromaDB timeout: {e}")
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"error": "Request timed out",
|
||||
"url": Config.get_chroma_url()
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"ChromaDB health check failed: {e}")
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"error": str(e),
|
||||
"url": Config.get_chroma_url()
|
||||
}
|
||||
|
||||
|
||||
def wait_for_services(max_retries: int = 30, delay: int = 2) -> bool:
|
||||
"""
|
||||
Wait for both Ollama and ChromaDB services to become healthy.
|
||||
|
||||
Args:
|
||||
max_retries: Maximum number of retry attempts
|
||||
delay: Seconds to wait between retries
|
||||
|
||||
Returns:
|
||||
bool: True if both services are healthy, False if max retries exceeded
|
||||
"""
|
||||
logger.info("Waiting for services to become healthy...")
|
||||
|
||||
for attempt in range(1, max_retries + 1):
|
||||
ollama_status = check_ollama_health()
|
||||
chroma_status = check_chroma_health()
|
||||
|
||||
if ollama_status["status"] == "healthy" and chroma_status["status"] == "healthy":
|
||||
logger.info(f"All services healthy after {attempt} attempts")
|
||||
return True
|
||||
|
||||
logger.info(
|
||||
f"Attempt {attempt}/{max_retries}: "
|
||||
f"Ollama={ollama_status['status']}, "
|
||||
f"ChromaDB={chroma_status['status']}"
|
||||
)
|
||||
|
||||
if attempt < max_retries:
|
||||
time.sleep(delay)
|
||||
|
||||
logger.error(f"Services did not become healthy after {max_retries} attempts")
|
||||
return False
|
||||
|
||||
224
app/core/ingestion.py
Normal file
224
app/core/ingestion.py
Normal file
|
|
@ -0,0 +1,224 @@
|
|||
"""
|
||||
Document ingestion pipeline for Code-Forge.
|
||||
Handles file upload, text extraction, chunking, embedding, and storage.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List
|
||||
|
||||
import fitz # PyMuPDF
|
||||
from llama_index.core import Document, Settings
|
||||
from llama_index.core.node_parser import SentenceSplitter
|
||||
from llama_index.embeddings.ollama import OllamaEmbedding
|
||||
|
||||
from app.core.chroma_manager import get_chroma_manager
|
||||
from app.core.config import Config
|
||||
from app.core.markdown_preprocessor import preprocess_markdown
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def extract_text_from_file(file_path: str) -> str:
|
||||
"""
|
||||
Extract text content from a file.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
str: Extracted text content
|
||||
"""
|
||||
file_path = Path(file_path)
|
||||
extension = file_path.suffix.lower()
|
||||
|
||||
try:
|
||||
if extension == ".pdf":
|
||||
# Extract text from PDF using PyMuPDF
|
||||
text = ""
|
||||
with fitz.open(file_path) as doc:
|
||||
for page in doc:
|
||||
text += page.get_text()
|
||||
return text
|
||||
else:
|
||||
# Read text/code files with UTF-8 encoding
|
||||
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
|
||||
return f.read()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to extract text from {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
|
||||
def chunk_document(text: str, metadata: Dict) -> List[Document]:
|
||||
"""
|
||||
Chunk document text into smaller pieces with overlap.
|
||||
|
||||
Args:
|
||||
text: Document text
|
||||
metadata: Document metadata
|
||||
|
||||
Returns:
|
||||
List of Document objects
|
||||
"""
|
||||
splitter = SentenceSplitter(
|
||||
chunk_size=Config.CHUNK_SIZE,
|
||||
chunk_overlap=Config.CHUNK_OVERLAP
|
||||
)
|
||||
|
||||
# Create a single document
|
||||
doc = Document(text=text, metadata=metadata)
|
||||
|
||||
# Split into chunks
|
||||
nodes = splitter.get_nodes_from_documents([doc])
|
||||
|
||||
# Convert nodes back to documents
|
||||
chunks = []
|
||||
for i, node in enumerate(nodes):
|
||||
chunk_metadata = metadata.copy()
|
||||
chunk_metadata["chunk_index"] = i
|
||||
chunk_metadata["chunk_id"] = node.node_id
|
||||
chunks.append(Document(text=node.text, metadata=chunk_metadata))
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
def ingest_file(
|
||||
file_path: str,
|
||||
collection_name: str,
|
||||
filename: str
|
||||
) -> Dict[str, any]:
|
||||
"""
|
||||
Ingest a single file into a knowledge base.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
collection_name: Target collection name
|
||||
filename: Original filename
|
||||
|
||||
Returns:
|
||||
dict: Ingestion result with status and details
|
||||
"""
|
||||
try:
|
||||
# Extract text
|
||||
text = extract_text_from_file(file_path)
|
||||
if not text.strip():
|
||||
return {
|
||||
"status": "error",
|
||||
"filename": filename,
|
||||
"error": "No text content extracted"
|
||||
}
|
||||
|
||||
# Create base metadata
|
||||
file_ext = Path(filename).suffix.lower()
|
||||
metadata = {
|
||||
"filename": filename,
|
||||
"file_type": file_ext,
|
||||
"upload_date": datetime.now().isoformat(),
|
||||
"source_path": str(file_path)
|
||||
}
|
||||
|
||||
# Preprocess markdown files
|
||||
if file_ext in ['.md', '.markdown']:
|
||||
try:
|
||||
processed_text, enriched_metadata = preprocess_markdown(
|
||||
text, filename, str(file_path)
|
||||
)
|
||||
text = processed_text
|
||||
metadata.update(enriched_metadata)
|
||||
logger.info(f"Preprocessed markdown: {filename}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Markdown preprocessing failed for {filename}: {e}")
|
||||
# Continue with original text if preprocessing fails
|
||||
|
||||
# Chunk document
|
||||
chunks = chunk_document(text, metadata)
|
||||
|
||||
# Initialize embedding model
|
||||
embed_model = OllamaEmbedding(
|
||||
model_name=Config.OLLAMA_EMBED_MODEL,
|
||||
base_url=Config.OLLAMA_HOST
|
||||
)
|
||||
|
||||
# Get ChromaDB collection
|
||||
chroma_manager = get_chroma_manager()
|
||||
collection = chroma_manager.get_collection(collection_name)
|
||||
|
||||
if not collection:
|
||||
return {
|
||||
"status": "error",
|
||||
"filename": filename,
|
||||
"error": f"Collection {collection_name} not found"
|
||||
}
|
||||
|
||||
# Generate embeddings and add to collection
|
||||
for chunk in chunks:
|
||||
# Generate embedding
|
||||
embedding = embed_model.get_text_embedding(chunk.text)
|
||||
|
||||
# Create unique ID
|
||||
chunk_id = f"{filename}_{chunk.metadata['chunk_index']}_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Add to ChromaDB
|
||||
collection.add(
|
||||
documents=[chunk.text],
|
||||
embeddings=[embedding],
|
||||
metadatas=[chunk.metadata],
|
||||
ids=[chunk_id]
|
||||
)
|
||||
|
||||
logger.info(f"Ingested {filename}: {len(chunks)} chunks")
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"filename": filename,
|
||||
"chunks": len(chunks),
|
||||
"file_type": file_ext
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to ingest {filename}: {e}")
|
||||
return {
|
||||
"status": "error",
|
||||
"filename": filename,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
|
||||
def ingest_directory(
|
||||
dir_path: str,
|
||||
collection_name: str
|
||||
) -> List[Dict[str, any]]:
|
||||
"""
|
||||
Recursively ingest all supported files from a directory.
|
||||
|
||||
Args:
|
||||
dir_path: Path to directory
|
||||
collection_name: Target collection name
|
||||
|
||||
Returns:
|
||||
List of ingestion results
|
||||
"""
|
||||
results = []
|
||||
dir_path = Path(dir_path)
|
||||
|
||||
if not dir_path.exists() or not dir_path.is_dir():
|
||||
logger.error(f"Directory not found: {dir_path}")
|
||||
return [{
|
||||
"status": "error",
|
||||
"error": f"Directory not found: {dir_path}"
|
||||
}]
|
||||
|
||||
# Recursively find all supported files
|
||||
for file_path in dir_path.rglob("*"):
|
||||
if file_path.is_file() and Config.is_supported_file(file_path.name):
|
||||
result = ingest_file(
|
||||
str(file_path),
|
||||
collection_name,
|
||||
file_path.name
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
175
app/core/kb_metadata.py
Normal file
175
app/core/kb_metadata.py
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
"""
|
||||
Knowledge base metadata management for UI display.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from datetime import datetime
|
||||
from typing import Dict, List
|
||||
|
||||
from app.core.chroma_manager import get_chroma_manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def generate_tags(file_type: str) -> List[str]:
|
||||
"""
|
||||
Generate display tags based on file type.
|
||||
|
||||
Args:
|
||||
file_type: File extension (e.g., '.py', '.md')
|
||||
|
||||
Returns:
|
||||
List of tag strings
|
||||
"""
|
||||
tag_map = {
|
||||
".py": ["Code", "Python"],
|
||||
".js": ["Code", "JavaScript"],
|
||||
".jsx": ["Code", "React"],
|
||||
".ts": ["Code", "TypeScript"],
|
||||
".tsx": ["Code", "React", "TypeScript"],
|
||||
".go": ["Code", "Go"],
|
||||
".rs": ["Code", "Rust"],
|
||||
".java": ["Code", "Java"],
|
||||
".cpp": ["Code", "C++"],
|
||||
".c": ["Code", "C"],
|
||||
".h": ["Code", "Header"],
|
||||
".md": ["Document", "Markdown"],
|
||||
".txt": ["Document", "Text"],
|
||||
".pdf": ["Document", "PDF"],
|
||||
".json": ["Config", "JSON"],
|
||||
".yaml": ["Config", "YAML"],
|
||||
".yml": ["Config", "YAML"],
|
||||
}
|
||||
|
||||
return tag_map.get(file_type.lower(), ["Document", "Unknown"])
|
||||
|
||||
|
||||
def format_timestamp(iso_timestamp: str) -> str:
|
||||
"""
|
||||
Format ISO timestamp for display.
|
||||
|
||||
Args:
|
||||
iso_timestamp: ISO format timestamp string
|
||||
|
||||
Returns:
|
||||
Formatted string (e.g., "Updated: 10/22/2025")
|
||||
"""
|
||||
try:
|
||||
dt = datetime.fromisoformat(iso_timestamp)
|
||||
return f"Updated: {dt.strftime('%m/%d/%Y')}"
|
||||
except Exception:
|
||||
return "Updated: Unknown"
|
||||
|
||||
|
||||
def get_documents_metadata(collection_name: str) -> List[Dict[str, any]]:
|
||||
"""
|
||||
Get metadata for all documents in a collection.
|
||||
Groups chunks by filename and aggregates metadata.
|
||||
|
||||
Args:
|
||||
collection_name: Name of the collection
|
||||
|
||||
Returns:
|
||||
List of document metadata dicts
|
||||
"""
|
||||
try:
|
||||
chroma_manager = get_chroma_manager()
|
||||
collection = chroma_manager.get_collection(collection_name)
|
||||
|
||||
if not collection:
|
||||
logger.warning(f"Collection {collection_name} not found")
|
||||
return []
|
||||
|
||||
# Get all documents with metadata
|
||||
results = collection.get(include=["metadatas"])
|
||||
|
||||
if not results or not results.get("metadatas"):
|
||||
return []
|
||||
|
||||
# Group chunks by filename
|
||||
docs_by_filename = defaultdict(list)
|
||||
for metadata in results["metadatas"]:
|
||||
filename = metadata.get("filename", "unknown")
|
||||
docs_by_filename[filename].append(metadata)
|
||||
|
||||
# Aggregate metadata for each document
|
||||
documents = []
|
||||
for filename, chunks in docs_by_filename.items():
|
||||
# Get metadata from first chunk
|
||||
first_chunk = chunks[0]
|
||||
file_type = first_chunk.get("file_type", "")
|
||||
upload_date = first_chunk.get("upload_date", "")
|
||||
|
||||
documents.append({
|
||||
"filename": filename,
|
||||
"file_type": file_type,
|
||||
"tags": generate_tags(file_type),
|
||||
"upload_date": upload_date,
|
||||
"formatted_date": format_timestamp(upload_date),
|
||||
"chunk_count": len(chunks)
|
||||
})
|
||||
|
||||
# Sort by upload date (newest first)
|
||||
documents.sort(
|
||||
key=lambda x: x.get("upload_date", ""),
|
||||
reverse=True
|
||||
)
|
||||
|
||||
return documents
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get documents metadata: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def get_collection_stats(collection_name: str) -> Dict[str, any]:
|
||||
"""
|
||||
Get statistics for a collection.
|
||||
|
||||
Args:
|
||||
collection_name: Name of the collection
|
||||
|
||||
Returns:
|
||||
dict: Statistics including total docs, chunks, last updated
|
||||
"""
|
||||
try:
|
||||
chroma_manager = get_chroma_manager()
|
||||
collection = chroma_manager.get_collection(collection_name)
|
||||
|
||||
if not collection:
|
||||
return {
|
||||
"total_documents": 0,
|
||||
"total_chunks": 0,
|
||||
"last_updated": None
|
||||
}
|
||||
|
||||
# Get all metadata
|
||||
results = collection.get(include=["metadatas"])
|
||||
metadatas = results.get("metadatas", [])
|
||||
|
||||
# Count unique documents
|
||||
unique_filenames = set(
|
||||
m.get("filename", "") for m in metadatas
|
||||
)
|
||||
|
||||
# Get latest upload date
|
||||
upload_dates = [
|
||||
m.get("upload_date", "") for m in metadatas if m.get("upload_date")
|
||||
]
|
||||
last_updated = max(upload_dates) if upload_dates else None
|
||||
|
||||
return {
|
||||
"total_documents": len(unique_filenames),
|
||||
"total_chunks": len(metadatas),
|
||||
"last_updated": last_updated
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get collection stats: {e}")
|
||||
return {
|
||||
"total_documents": 0,
|
||||
"total_chunks": 0,
|
||||
"last_updated": None
|
||||
}
|
||||
|
||||
286
app/core/markdown_preprocessor.py
Normal file
286
app/core/markdown_preprocessor.py
Normal file
|
|
@ -0,0 +1,286 @@
|
|||
"""
|
||||
Markdown preprocessing for Code-Forge.
|
||||
Standardizes markdown files for better chunking and retrieval.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def extract_frontmatter(text: str) -> Tuple[Optional[Dict], str]:
|
||||
"""
|
||||
Extract YAML/TOML frontmatter from markdown.
|
||||
|
||||
Args:
|
||||
text: Markdown content
|
||||
|
||||
Returns:
|
||||
Tuple of (frontmatter_dict, content_without_frontmatter)
|
||||
"""
|
||||
frontmatter = {}
|
||||
content = text
|
||||
|
||||
# Check for YAML frontmatter (--- ... ---)
|
||||
yaml_pattern = r'^---\s*\n(.*?)\n---\s*\n'
|
||||
yaml_match = re.match(yaml_pattern, text, re.DOTALL)
|
||||
|
||||
if yaml_match:
|
||||
frontmatter_text = yaml_match.group(1)
|
||||
content = text[yaml_match.end():]
|
||||
|
||||
# Parse simple YAML (key: value pairs)
|
||||
for line in frontmatter_text.split('\n'):
|
||||
line = line.strip()
|
||||
if ':' in line:
|
||||
key, value = line.split(':', 1)
|
||||
key = key.strip()
|
||||
value = value.strip().strip('"').strip("'")
|
||||
frontmatter[key] = value
|
||||
|
||||
# Check for TOML frontmatter (+++ ... +++)
|
||||
toml_pattern = r'^\+\+\+\s*\n(.*?)\n\+\+\+\s*\n'
|
||||
toml_match = re.match(toml_pattern, text, re.DOTALL)
|
||||
|
||||
if toml_match:
|
||||
frontmatter_text = toml_match.group(1)
|
||||
content = text[toml_match.end():]
|
||||
|
||||
# Parse simple TOML (key = value pairs)
|
||||
for line in frontmatter_text.split('\n'):
|
||||
line = line.strip()
|
||||
if '=' in line:
|
||||
key, value = line.split('=', 1)
|
||||
key = key.strip()
|
||||
value = value.strip().strip('"').strip("'")
|
||||
frontmatter[key] = value
|
||||
|
||||
return frontmatter if frontmatter else None, content
|
||||
|
||||
|
||||
def extract_title(text: str) -> Optional[str]:
|
||||
"""
|
||||
Extract title from markdown (first H1 or frontmatter).
|
||||
|
||||
Args:
|
||||
text: Markdown content
|
||||
|
||||
Returns:
|
||||
Title string or None
|
||||
"""
|
||||
# Try to find first H1 header
|
||||
h1_pattern = r'^#\s+(.+)$'
|
||||
match = re.search(h1_pattern, text, re.MULTILINE)
|
||||
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def normalize_headers(text: str) -> str:
|
||||
"""
|
||||
Normalize markdown headers to ATX style (# ## ###).
|
||||
|
||||
Args:
|
||||
text: Markdown content
|
||||
|
||||
Returns:
|
||||
Normalized markdown
|
||||
"""
|
||||
lines = text.split('\n')
|
||||
normalized = []
|
||||
|
||||
i = 0
|
||||
while i < len(lines):
|
||||
line = lines[i]
|
||||
|
||||
# Check if next line is a setext-style header underline
|
||||
if i + 1 < len(lines):
|
||||
next_line = lines[i + 1].strip()
|
||||
|
||||
# H1: underlined with ===
|
||||
if next_line and all(c == '=' for c in next_line):
|
||||
normalized.append(f"# {line}")
|
||||
i += 2 # Skip the underline
|
||||
continue
|
||||
|
||||
# H2: underlined with ---
|
||||
if next_line and all(c == '-' for c in next_line) and len(next_line) >= 3:
|
||||
normalized.append(f"## {line}")
|
||||
i += 2 # Skip the underline
|
||||
continue
|
||||
|
||||
# Keep line as-is
|
||||
normalized.append(line)
|
||||
i += 1
|
||||
|
||||
return '\n'.join(normalized)
|
||||
|
||||
|
||||
def clean_whitespace(text: str) -> str:
|
||||
"""
|
||||
Clean up excessive whitespace while preserving code blocks.
|
||||
|
||||
Args:
|
||||
text: Markdown content
|
||||
|
||||
Returns:
|
||||
Cleaned markdown
|
||||
"""
|
||||
# Preserve code blocks
|
||||
code_blocks = []
|
||||
code_pattern = r'```[\s\S]*?```'
|
||||
|
||||
def save_code_block(match):
|
||||
code_blocks.append(match.group(0))
|
||||
return f"__CODE_BLOCK_{len(code_blocks) - 1}__"
|
||||
|
||||
text = re.sub(code_pattern, save_code_block, text)
|
||||
|
||||
# Remove trailing whitespace from lines
|
||||
lines = [line.rstrip() for line in text.split('\n')]
|
||||
|
||||
# Collapse multiple blank lines to max 2
|
||||
cleaned = []
|
||||
blank_count = 0
|
||||
|
||||
for line in lines:
|
||||
if not line.strip():
|
||||
blank_count += 1
|
||||
if blank_count <= 2:
|
||||
cleaned.append(line)
|
||||
else:
|
||||
blank_count = 0
|
||||
cleaned.append(line)
|
||||
|
||||
text = '\n'.join(cleaned)
|
||||
|
||||
# Restore code blocks
|
||||
for i, code_block in enumerate(code_blocks):
|
||||
text = text.replace(f"__CODE_BLOCK_{i}__", code_block)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def normalize_code_fences(text: str) -> str:
|
||||
"""
|
||||
Normalize code fence markers to triple backticks.
|
||||
|
||||
Args:
|
||||
text: Markdown content
|
||||
|
||||
Returns:
|
||||
Normalized markdown
|
||||
"""
|
||||
# Replace ~~~ with ```
|
||||
text = re.sub(r'^~~~', '```', text, flags=re.MULTILINE)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def slugify(text: str) -> str:
|
||||
"""
|
||||
Convert text to URL-friendly slug.
|
||||
|
||||
Args:
|
||||
text: Text to slugify
|
||||
|
||||
Returns:
|
||||
Slugified string
|
||||
"""
|
||||
# Convert to lowercase
|
||||
text = text.lower()
|
||||
|
||||
# Replace spaces and underscores with hyphens
|
||||
text = re.sub(r'[\s_]+', '-', text)
|
||||
|
||||
# Remove non-alphanumeric characters (except hyphens)
|
||||
text = re.sub(r'[^a-z0-9-]', '', text)
|
||||
|
||||
# Remove multiple consecutive hyphens
|
||||
text = re.sub(r'-+', '-', text)
|
||||
|
||||
# Remove leading/trailing hyphens
|
||||
text = text.strip('-')
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def preprocess_markdown(
|
||||
text: str,
|
||||
filename: str,
|
||||
file_path: str
|
||||
) -> Tuple[str, Dict]:
|
||||
"""
|
||||
Preprocess markdown file for better chunking and retrieval.
|
||||
|
||||
Args:
|
||||
text: Raw markdown content
|
||||
filename: Original filename
|
||||
file_path: File path
|
||||
|
||||
Returns:
|
||||
Tuple of (processed_text, enriched_metadata)
|
||||
"""
|
||||
logger.info(f"Preprocessing markdown: {filename}")
|
||||
|
||||
# Extract frontmatter
|
||||
frontmatter, content = extract_frontmatter(text)
|
||||
|
||||
# Extract title
|
||||
title = None
|
||||
if frontmatter and 'title' in frontmatter:
|
||||
title = frontmatter['title']
|
||||
else:
|
||||
title = extract_title(content)
|
||||
|
||||
# If no title found, use filename
|
||||
if not title:
|
||||
title = Path(filename).stem.replace('-', ' ').replace('_', ' ').title()
|
||||
|
||||
# Normalize content
|
||||
content = normalize_headers(content)
|
||||
content = normalize_code_fences(content)
|
||||
content = clean_whitespace(content)
|
||||
|
||||
# Build enriched metadata
|
||||
metadata = {
|
||||
"title": title,
|
||||
"original_filename": filename,
|
||||
"source_path": file_path,
|
||||
"preprocessed": True,
|
||||
"preprocessed_date": datetime.now().isoformat()
|
||||
}
|
||||
|
||||
# Add frontmatter fields to metadata
|
||||
if frontmatter:
|
||||
for key, value in frontmatter.items():
|
||||
if key not in metadata: # Don't override existing keys
|
||||
metadata[f"fm_{key}"] = value # Prefix with 'fm_' to indicate frontmatter
|
||||
|
||||
# Extract tags if present
|
||||
if frontmatter and 'tags' in frontmatter:
|
||||
tags = frontmatter['tags']
|
||||
if isinstance(tags, str):
|
||||
# Parse comma-separated or space-separated tags
|
||||
tags = re.split(r'[,\s]+', tags)
|
||||
tags = [t.strip() for t in tags if t.strip()]
|
||||
metadata['tags'] = ','.join(tags) if tags else ''
|
||||
|
||||
# Extract author if present
|
||||
if frontmatter and 'author' in frontmatter:
|
||||
metadata['author'] = frontmatter['author']
|
||||
|
||||
# Extract date if present
|
||||
if frontmatter and 'date' in frontmatter:
|
||||
metadata['document_date'] = frontmatter['date']
|
||||
|
||||
logger.info(f"Preprocessed {filename}: title='{title}', metadata_fields={len(metadata)}")
|
||||
|
||||
return content, metadata
|
||||
|
||||
314
app/core/rag_engine.py
Normal file
314
app/core/rag_engine.py
Normal file
|
|
@ -0,0 +1,314 @@
|
|||
"""
|
||||
RAG Engine with advanced retrieval strategies.
|
||||
Supports vector search, hybrid search, reranking, and agentic modes.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from llama_index.core import Settings, StorageContext, VectorStoreIndex, get_response_synthesizer
|
||||
from llama_index.core.query_engine import SubQuestionQueryEngine
|
||||
from llama_index.core.retrievers import VectorIndexRetriever
|
||||
from llama_index.core.tools import QueryEngineTool
|
||||
from llama_index.embeddings.ollama import OllamaEmbedding
|
||||
from llama_index.llms.ollama import Ollama
|
||||
from llama_index.core.postprocessor import SentenceTransformerRerank
|
||||
from llama_index.vector_stores.chroma import ChromaVectorStore
|
||||
|
||||
# BM25 retriever - needs to be installed separately
|
||||
try:
|
||||
from llama_index.retrievers.bm25 import BM25Retriever
|
||||
HAS_BM25 = True
|
||||
except ImportError:
|
||||
HAS_BM25 = False
|
||||
|
||||
from app.core.chroma_manager import get_chroma_manager
|
||||
from app.core.config import Config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RAGEngine:
|
||||
"""Advanced RAG engine with multiple retrieval strategies."""
|
||||
|
||||
def __init__(self, collection_name: str):
|
||||
"""
|
||||
Initialize RAG engine for a specific knowledge base.
|
||||
|
||||
Args:
|
||||
collection_name: Name of the ChromaDB collection
|
||||
"""
|
||||
self.collection_name = collection_name
|
||||
|
||||
# Initialize Ollama LLM with optimized context window
|
||||
# Context needs to fit: retrieved chunks (5 × 1024 = 5120) + question + answer
|
||||
# Using 4096 is a good balance between memory usage and functionality
|
||||
Settings.llm = Ollama(
|
||||
model=Config.OLLAMA_MODEL,
|
||||
base_url=Config.OLLAMA_HOST,
|
||||
request_timeout=120.0,
|
||||
context_window=4096, # Balanced: fits most RAG queries without excessive memory
|
||||
num_ctx=4096 # Ollama-specific parameter for context size
|
||||
)
|
||||
|
||||
# Initialize Ollama embeddings
|
||||
Settings.embed_model = OllamaEmbedding(
|
||||
model_name=Config.OLLAMA_EMBED_MODEL,
|
||||
base_url=Config.OLLAMA_HOST
|
||||
)
|
||||
|
||||
# Get ChromaDB collection
|
||||
chroma_manager = get_chroma_manager()
|
||||
collection = chroma_manager.get_collection(collection_name)
|
||||
|
||||
if not collection:
|
||||
raise ValueError(f"Collection {collection_name} not found")
|
||||
|
||||
# Create vector store and index
|
||||
vector_store = ChromaVectorStore(chroma_collection=collection)
|
||||
storage_context = StorageContext.from_defaults(vector_store=vector_store)
|
||||
self.index = VectorStoreIndex.from_vector_store(
|
||||
vector_store,
|
||||
storage_context=storage_context
|
||||
)
|
||||
|
||||
# Create base query engine
|
||||
self.query_engine = self.index.as_query_engine(
|
||||
similarity_top_k=Config.TOP_K_RETRIEVAL
|
||||
)
|
||||
|
||||
# Initialize retrievers
|
||||
self.vector_retriever = VectorIndexRetriever(
|
||||
index=self.index,
|
||||
similarity_top_k=Config.TOP_K_RETRIEVAL
|
||||
)
|
||||
|
||||
# Initialize BM25 retriever (for hybrid search)
|
||||
if HAS_BM25:
|
||||
try:
|
||||
all_nodes = self.index.docstore.docs.values()
|
||||
self.bm25_retriever = BM25Retriever.from_defaults(
|
||||
nodes=list(all_nodes),
|
||||
similarity_top_k=Config.TOP_K_RETRIEVAL
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize BM25 retriever: {e}")
|
||||
self.bm25_retriever = None
|
||||
else:
|
||||
logger.warning("BM25 retriever not available - llama-index-retrievers-bm25 not installed")
|
||||
self.bm25_retriever = None
|
||||
|
||||
# Initialize reranker
|
||||
try:
|
||||
self.reranker = SentenceTransformerRerank(
|
||||
model="cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
top_n=Config.RERANK_TOP_N
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize reranker: {e}")
|
||||
self.reranker = None
|
||||
|
||||
# Initialize agentic query engine
|
||||
try:
|
||||
query_tool = QueryEngineTool.from_defaults(
|
||||
query_engine=self.query_engine,
|
||||
name="knowledge_base",
|
||||
description=f"Searches the {collection_name} knowledge base for relevant information"
|
||||
)
|
||||
self.sub_question_engine = SubQuestionQueryEngine.from_defaults(
|
||||
query_engine_tools=[query_tool]
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize agentic engine: {e}")
|
||||
self.sub_question_engine = None
|
||||
|
||||
def _reciprocal_rank_fusion(
|
||||
self,
|
||||
results_list: List[List],
|
||||
k: int = 60
|
||||
) -> List:
|
||||
"""
|
||||
Fuse multiple retrieval results using reciprocal rank fusion.
|
||||
|
||||
Args:
|
||||
results_list: List of retrieval result lists
|
||||
k: RRF parameter (default 60)
|
||||
|
||||
Returns:
|
||||
Fused and sorted list of nodes
|
||||
"""
|
||||
# Aggregate scores by node ID
|
||||
node_scores = {}
|
||||
|
||||
for results in results_list:
|
||||
for rank, node in enumerate(results, start=1):
|
||||
node_id = node.node_id
|
||||
score = 1.0 / (k + rank)
|
||||
|
||||
if node_id in node_scores:
|
||||
node_scores[node_id]["score"] += score
|
||||
else:
|
||||
node_scores[node_id] = {
|
||||
"node": node,
|
||||
"score": score
|
||||
}
|
||||
|
||||
# Sort by score
|
||||
sorted_nodes = sorted(
|
||||
node_scores.values(),
|
||||
key=lambda x: x["score"],
|
||||
reverse=True
|
||||
)
|
||||
|
||||
return [item["node"] for item in sorted_nodes[:Config.TOP_K_RETRIEVAL]]
|
||||
|
||||
def _query_hybrid(self, question: str) -> Dict[str, any]:
|
||||
"""
|
||||
Execute hybrid search (vector + BM25).
|
||||
|
||||
Args:
|
||||
question: Query string
|
||||
|
||||
Returns:
|
||||
dict: Answer and sources
|
||||
"""
|
||||
if not self.bm25_retriever:
|
||||
logger.warning("BM25 retriever not available, falling back to vector search")
|
||||
return self._query_base(question)
|
||||
|
||||
# Retrieve from both methods
|
||||
vector_results = self.vector_retriever.retrieve(question)
|
||||
bm25_results = self.bm25_retriever.retrieve(question)
|
||||
|
||||
# Fuse results
|
||||
fused_nodes = self._reciprocal_rank_fusion([vector_results, bm25_results])
|
||||
|
||||
# Synthesize answer
|
||||
synthesizer = get_response_synthesizer()
|
||||
response = synthesizer.synthesize(question, nodes=fused_nodes)
|
||||
|
||||
return self._format_response(response, fused_nodes)
|
||||
|
||||
def _apply_reranking(self, nodes: List, query: str) -> List:
|
||||
"""
|
||||
Apply reranking to retrieved nodes.
|
||||
|
||||
Args:
|
||||
nodes: Retrieved nodes
|
||||
query: Query string
|
||||
|
||||
Returns:
|
||||
Reranked nodes
|
||||
"""
|
||||
if not self.reranker:
|
||||
logger.warning("Reranker not available")
|
||||
return nodes
|
||||
|
||||
reranked = self.reranker.postprocess_nodes(nodes, query_str=query)
|
||||
return reranked
|
||||
|
||||
def _query_agentic(self, question: str) -> Dict[str, any]:
|
||||
"""
|
||||
Execute agentic RAG with sub-question decomposition.
|
||||
|
||||
Args:
|
||||
question: Query string
|
||||
|
||||
Returns:
|
||||
dict: Answer, sub-questions, and sources
|
||||
"""
|
||||
if not self.sub_question_engine:
|
||||
logger.warning("Agentic engine not available, falling back to base query")
|
||||
return self._query_base(question)
|
||||
|
||||
response = self.sub_question_engine.query(question)
|
||||
|
||||
# Extract sub-questions if available
|
||||
sub_questions = []
|
||||
if hasattr(response, "metadata") and response.metadata:
|
||||
sub_qa = response.metadata.get("sub_qa", [])
|
||||
sub_questions = [
|
||||
{"question": item[0], "answer": item[1]}
|
||||
for item in sub_qa
|
||||
]
|
||||
|
||||
result = self._format_response(response, response.source_nodes)
|
||||
result["sub_questions"] = sub_questions
|
||||
|
||||
return result
|
||||
|
||||
def _query_base(self, question: str) -> Dict[str, any]:
|
||||
"""
|
||||
Execute base vector search query.
|
||||
|
||||
Args:
|
||||
question: Query string
|
||||
|
||||
Returns:
|
||||
dict: Answer and sources
|
||||
"""
|
||||
response = self.query_engine.query(question)
|
||||
return self._format_response(response, response.source_nodes)
|
||||
|
||||
def _format_response(self, response, source_nodes) -> Dict[str, any]:
|
||||
"""Format query response for return."""
|
||||
sources = []
|
||||
for node in source_nodes:
|
||||
sources.append({
|
||||
"text": node.text,
|
||||
"metadata": node.metadata if hasattr(node, "metadata") else {},
|
||||
"score": node.score if hasattr(node, "score") else None
|
||||
})
|
||||
|
||||
return {
|
||||
"answer": str(response),
|
||||
"sources": sources
|
||||
}
|
||||
|
||||
def query(
|
||||
self,
|
||||
question: str,
|
||||
use_hybrid: bool = False,
|
||||
use_rerank: bool = False,
|
||||
use_agentic: bool = False
|
||||
) -> Dict[str, any]:
|
||||
"""
|
||||
Execute query with specified strategies.
|
||||
|
||||
Args:
|
||||
question: Query string
|
||||
use_hybrid: Enable hybrid search
|
||||
use_rerank: Enable reranking
|
||||
use_agentic: Enable agentic RAG
|
||||
|
||||
Returns:
|
||||
dict: Answer and sources
|
||||
"""
|
||||
try:
|
||||
# Route to appropriate query method
|
||||
if use_agentic:
|
||||
return self._query_agentic(question)
|
||||
elif use_hybrid:
|
||||
result = self._query_hybrid(question)
|
||||
if use_rerank and result.get("sources"):
|
||||
# Note: Reranking after hybrid fusion
|
||||
pass # Already applied in hybrid
|
||||
return result
|
||||
else:
|
||||
result = self._query_base(question)
|
||||
if use_rerank and result.get("sources"):
|
||||
# Apply reranking to base results
|
||||
nodes = self.vector_retriever.retrieve(question)
|
||||
reranked_nodes = self._apply_reranking(nodes, question)
|
||||
synthesizer = get_response_synthesizer()
|
||||
response = synthesizer.synthesize(question, nodes=reranked_nodes)
|
||||
return self._format_response(response, reranked_nodes)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Query failed: {e}")
|
||||
return {
|
||||
"answer": f"Error executing query: {str(e)}",
|
||||
"sources": []
|
||||
}
|
||||
|
||||
682
app/main.py
Normal file
682
app/main.py
Normal file
|
|
@ -0,0 +1,682 @@
|
|||
"""
|
||||
Code-Forge: Localized Multi-KB RAG System
|
||||
Main Streamlit application with badass Archon-inspired UI.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import streamlit as st
|
||||
|
||||
from app.core.chroma_manager import get_chroma_manager
|
||||
from app.core.config import Config
|
||||
from app.core.health import check_chroma_health, check_ollama_health
|
||||
from app.core.ingestion import ingest_directory, ingest_file
|
||||
from app.core.kb_metadata import get_documents_metadata
|
||||
from app.core.rag_engine import RAGEngine
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Page configuration
|
||||
st.set_page_config(
|
||||
page_title="Code-Forge 🔥",
|
||||
page_icon="🔥",
|
||||
layout="wide",
|
||||
initial_sidebar_state="expanded"
|
||||
)
|
||||
|
||||
# Initialize session state
|
||||
if "selected_kb" not in st.session_state:
|
||||
st.session_state.selected_kb = None
|
||||
if "chat_history" not in st.session_state:
|
||||
st.session_state.chat_history = []
|
||||
if "rag_settings" not in st.session_state:
|
||||
st.session_state.rag_settings = {
|
||||
"hybrid": False,
|
||||
"rerank": False,
|
||||
"agentic": False
|
||||
}
|
||||
|
||||
# Custom CSS with Archon aesthetic
|
||||
st.markdown("""
|
||||
<style>
|
||||
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;700&display=swap');
|
||||
@import url('https://fonts.googleapis.com/css2?family=Material+Icons&display=swap');
|
||||
|
||||
* {
|
||||
font-family: 'JetBrains Mono', monospace !important;
|
||||
}
|
||||
|
||||
/* Hide Material Icons text in expanders and buttons */
|
||||
[data-testid="stIconMaterial"] {
|
||||
font-size: 0 !important;
|
||||
width: 20px !important;
|
||||
height: 20px !important;
|
||||
display: inline-block !important;
|
||||
position: relative !important;
|
||||
}
|
||||
|
||||
/* Replace expander arrow with CSS */
|
||||
[data-testid="stIconMaterial"]::before {
|
||||
content: "▶" !important;
|
||||
font-size: 14px !important;
|
||||
color: hsl(271, 91%, 65%) !important;
|
||||
position: absolute !important;
|
||||
top: 50% !important;
|
||||
left: 50% !important;
|
||||
transform: translate(-50%, -50%) !important;
|
||||
font-family: 'JetBrains Mono', monospace !important;
|
||||
}
|
||||
|
||||
/* Rotate arrow when expander is open */
|
||||
details[open] [data-testid="stIconMaterial"]::before {
|
||||
transform: translate(-50%, -50%) rotate(90deg) !important;
|
||||
}
|
||||
|
||||
/* Hide sidebar collapse button if it exists */
|
||||
[data-testid="collapsedControl"] {
|
||||
display: none !important;
|
||||
}
|
||||
|
||||
/* Root app container */
|
||||
.stApp {
|
||||
background-color: #0E0E0E;
|
||||
}
|
||||
|
||||
/* Fix main content area - prevent cutoff */
|
||||
.main .block-container {
|
||||
padding-top: 2rem !important;
|
||||
padding-bottom: 2rem !important;
|
||||
padding-left: 5rem !important;
|
||||
padding-right: 5rem !important;
|
||||
max-width: none !important;
|
||||
}
|
||||
|
||||
.main {
|
||||
background-color: #0E0E0E;
|
||||
}
|
||||
|
||||
/* Ensure all content has proper spacing */
|
||||
.element-container {
|
||||
margin-left: 0 !important;
|
||||
padding-left: 0 !important;
|
||||
}
|
||||
|
||||
/* Fix sidebar width and padding - make it wider and disable collapse */
|
||||
section[data-testid="stSidebar"] {
|
||||
width: 450px !important;
|
||||
min-width: 450px !important;
|
||||
max-width: 450px !important;
|
||||
background-color: #0E0E0E;
|
||||
border-right: 2px solid hsl(271, 91%, 65%);
|
||||
overflow-x: hidden !important;
|
||||
}
|
||||
|
||||
[data-testid="stSidebar"] > div:first-child {
|
||||
padding: 2rem 1.5rem !important;
|
||||
width: 100% !important;
|
||||
overflow-x: hidden !important;
|
||||
}
|
||||
|
||||
[data-testid="stSidebar"] .element-container {
|
||||
padding-left: 0 !important;
|
||||
margin-left: 0 !important;
|
||||
width: 100% !important;
|
||||
max-width: 100% !important;
|
||||
overflow-x: hidden !important;
|
||||
}
|
||||
|
||||
/* Completely hide sidebar collapse button */
|
||||
[data-testid="collapsedControl"],
|
||||
button[kind="header"] {
|
||||
display: none !important;
|
||||
visibility: hidden !important;
|
||||
}
|
||||
|
||||
/* Make sidebar text wrap properly with better font sizes */
|
||||
[data-testid="stSidebar"] .stMarkdown {
|
||||
width: 100% !important;
|
||||
max-width: 100% !important;
|
||||
overflow-wrap: break-word !important;
|
||||
word-wrap: break-word !important;
|
||||
}
|
||||
|
||||
[data-testid="stSidebar"] code {
|
||||
font-size: 11px !important;
|
||||
word-break: break-all !important;
|
||||
white-space: pre-wrap !important;
|
||||
line-height: 1.4 !important;
|
||||
}
|
||||
|
||||
[data-testid="stSidebar"] pre {
|
||||
white-space: pre-wrap !important;
|
||||
word-wrap: break-word !important;
|
||||
overflow-x: hidden !important;
|
||||
max-width: 100% !important;
|
||||
font-size: 11px !important;
|
||||
line-height: 1.4 !important;
|
||||
}
|
||||
|
||||
[data-testid="stSidebar"] p {
|
||||
overflow-wrap: break-word !important;
|
||||
word-wrap: break-word !important;
|
||||
font-size: 13px !important;
|
||||
}
|
||||
|
||||
/* Larger text for headings in sidebar */
|
||||
[data-testid="stSidebar"] h3 {
|
||||
font-size: 16px !important;
|
||||
}
|
||||
|
||||
[data-testid="stSidebar"] h4 {
|
||||
font-size: 14px !important;
|
||||
}
|
||||
|
||||
/* Fix expander headers in sidebar */
|
||||
[data-testid="stSidebar"] .streamlit-expanderHeader {
|
||||
font-size: 13px !important;
|
||||
overflow-wrap: break-word !important;
|
||||
word-wrap: break-word !important;
|
||||
}
|
||||
|
||||
[data-testid="stSidebar"] details summary {
|
||||
overflow-wrap: break-word !important;
|
||||
word-wrap: break-word !important;
|
||||
}
|
||||
|
||||
/* Fix any negative margins */
|
||||
div[data-testid="column"] {
|
||||
padding-left: 0 !important;
|
||||
}
|
||||
|
||||
/* Status indicators */
|
||||
.status-healthy {
|
||||
color: hsl(160, 84%, 39%);
|
||||
font-weight: bold;
|
||||
text-shadow: 0 0 10px hsl(160, 84%, 39%);
|
||||
}
|
||||
|
||||
.status-unhealthy {
|
||||
color: #FF0000;
|
||||
font-weight: bold;
|
||||
text-shadow: 0 0 10px #FF0000;
|
||||
}
|
||||
|
||||
/* Document cards */
|
||||
.doc-card {
|
||||
background: linear-gradient(135deg, #1A1A1A 0%, #0E0E0E 100%);
|
||||
border: 2px solid hsl(330, 90%, 65%);
|
||||
border-radius: 12px;
|
||||
padding: 20px;
|
||||
margin: 10px 0;
|
||||
box-shadow: 0 8px 16px rgba(233, 30, 99, 0.3);
|
||||
transition: all 0.3s ease;
|
||||
}
|
||||
|
||||
.doc-card:hover {
|
||||
transform: translateY(-4px);
|
||||
box-shadow: 0 12px 24px rgba(233, 30, 99, 0.5);
|
||||
border-color: hsl(271, 91%, 65%);
|
||||
}
|
||||
|
||||
/* Tags */
|
||||
.tag {
|
||||
display: inline-block;
|
||||
background: linear-gradient(135deg, hsl(271, 91%, 65%) 0%, hsl(330, 90%, 65%) 100%);
|
||||
color: white;
|
||||
padding: 4px 10px;
|
||||
border-radius: 6px;
|
||||
font-size: 11px;
|
||||
font-weight: bold;
|
||||
margin: 2px;
|
||||
box-shadow: 0 2px 4px rgba(156, 39, 176, 0.3);
|
||||
}
|
||||
|
||||
/* Buttons */
|
||||
.stButton > button {
|
||||
background: linear-gradient(135deg, #E91E63 0%, hsl(271, 91%, 65%) 100%);
|
||||
color: white;
|
||||
border: none;
|
||||
border-radius: 8px;
|
||||
font-weight: bold;
|
||||
padding: 12px 24px;
|
||||
transition: all 0.3s ease;
|
||||
box-shadow: 0 4px 8px rgba(233, 30, 99, 0.3);
|
||||
width: 100% !important;
|
||||
max-width: 100% !important;
|
||||
}
|
||||
|
||||
.stButton > button:hover {
|
||||
box-shadow: 0 0 30px rgba(233, 30, 99, 0.8);
|
||||
transform: scale(1.05);
|
||||
}
|
||||
|
||||
/* Fix sidebar buttons to not overflow */
|
||||
[data-testid="stSidebar"] .stButton > button {
|
||||
font-size: 12px !important;
|
||||
padding: 8px 12px !important;
|
||||
white-space: normal !important;
|
||||
word-wrap: break-word !important;
|
||||
}
|
||||
|
||||
/* Chat messages */
|
||||
.user-message {
|
||||
background: linear-gradient(135deg, #E91E63 0%, hsl(271, 91%, 65%) 100%);
|
||||
color: white;
|
||||
padding: 16px;
|
||||
border-radius: 12px;
|
||||
margin: 12px 0;
|
||||
box-shadow: 0 4px 8px rgba(233, 30, 99, 0.3);
|
||||
}
|
||||
|
||||
.assistant-message {
|
||||
background: #1A1A1A;
|
||||
color: white;
|
||||
padding: 16px;
|
||||
border-radius: 12px;
|
||||
margin: 12px 0;
|
||||
border-left: 4px solid hsl(160, 84%, 39%);
|
||||
box-shadow: 0 4px 8px rgba(0, 0, 0, 0.5);
|
||||
}
|
||||
|
||||
/* Header */
|
||||
.main-header {
|
||||
text-align: center;
|
||||
padding: 20px;
|
||||
background: linear-gradient(135deg, #0E0E0E 0%, #1A1A1A 100%);
|
||||
border-bottom: 2px solid hsl(271, 91%, 65%);
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
|
||||
.main-header h1 {
|
||||
color: hsl(271, 91%, 65%);
|
||||
text-shadow: 0 0 20px hsl(271, 91%, 65%);
|
||||
}
|
||||
|
||||
.main-header p {
|
||||
color: hsl(160, 84%, 39%);
|
||||
}
|
||||
</style>
|
||||
""", unsafe_allow_html=True)
|
||||
|
||||
# Header
|
||||
st.markdown("""
|
||||
<div class="main-header">
|
||||
<h1>🔥 Code-Forge</h1>
|
||||
<p>Localized Multi-KB RAG System with MCP Integration</p>
|
||||
</div>
|
||||
""", unsafe_allow_html=True)
|
||||
|
||||
# Sidebar
|
||||
with st.sidebar:
|
||||
st.markdown("## ⚙️ Settings")
|
||||
|
||||
# Connection Status
|
||||
st.markdown("### 🔌 Connection Status")
|
||||
|
||||
ollama_status = check_ollama_health()
|
||||
chroma_status = check_chroma_health()
|
||||
|
||||
if ollama_status["status"] == "healthy":
|
||||
st.markdown('<p class="status-healthy">🟢 Ollama Connected</p>', unsafe_allow_html=True)
|
||||
if ollama_status.get("models"):
|
||||
st.caption(f"Models: {', '.join(ollama_status['models'][:3])}")
|
||||
else:
|
||||
st.markdown('<p class="status-unhealthy">🔴 Ollama Disconnected</p>', unsafe_allow_html=True)
|
||||
st.caption(ollama_status.get("error", "Unknown error"))
|
||||
|
||||
if chroma_status["status"] == "healthy":
|
||||
st.markdown('<p class="status-healthy">🟢 ChromaDB Connected</p>', unsafe_allow_html=True)
|
||||
else:
|
||||
st.markdown('<p class="status-unhealthy">🔴 ChromaDB Disconnected</p>', unsafe_allow_html=True)
|
||||
st.caption(chroma_status.get("error", "Unknown error"))
|
||||
|
||||
st.divider()
|
||||
|
||||
# MCP Server Info
|
||||
st.markdown("### 🔗 MCP Server")
|
||||
mcp_url = f"http://localhost:{Config.MCP_SERVER_PORT}/mcp"
|
||||
st.code(mcp_url, language="text")
|
||||
|
||||
st.caption("**Connection Command:**")
|
||||
st.code(f"claude mcp add --transport http code-forge {mcp_url}", language="bash")
|
||||
|
||||
st.divider()
|
||||
|
||||
# Knowledge Base Selector
|
||||
st.markdown("### 📚 Knowledge Base")
|
||||
|
||||
chroma_manager = get_chroma_manager()
|
||||
collections = chroma_manager.list_collections()
|
||||
|
||||
if collections:
|
||||
selected = st.selectbox(
|
||||
"Select KB",
|
||||
collections,
|
||||
key="kb_selector"
|
||||
)
|
||||
st.session_state.selected_kb = selected
|
||||
else:
|
||||
st.warning("No knowledge bases found. Create one below!")
|
||||
st.session_state.selected_kb = None
|
||||
|
||||
# Create new KB
|
||||
with st.expander("➕ Create New KB"):
|
||||
new_kb_name = st.text_input("KB Name", key="new_kb_name")
|
||||
if st.button("Create"):
|
||||
if new_kb_name:
|
||||
if chroma_manager.create_collection(new_kb_name):
|
||||
st.success(f"Created KB: {new_kb_name}")
|
||||
st.rerun()
|
||||
else:
|
||||
st.error("Failed to create KB")
|
||||
else:
|
||||
st.error("Please enter a KB name")
|
||||
|
||||
st.divider()
|
||||
|
||||
# RAG Strategy Toggles
|
||||
st.markdown("### 🎯 RAG Strategies")
|
||||
|
||||
st.session_state.rag_settings["hybrid"] = st.checkbox(
|
||||
"🔀 Hybrid Search",
|
||||
value=st.session_state.rag_settings["hybrid"],
|
||||
help="**Hybrid Search**: Combines semantic vector search with BM25 keyword search for better recall. Uses reciprocal rank fusion to merge results. Best for queries with specific technical terms."
|
||||
)
|
||||
|
||||
st.session_state.rag_settings["rerank"] = st.checkbox(
|
||||
"🎯 Reranking",
|
||||
value=st.session_state.rag_settings["rerank"],
|
||||
help="**Reranking**: Re-scores retrieved chunks using a cross-encoder model (ms-marco-MiniLM) for higher precision. Keeps top 3 most relevant results. Improves answer quality but adds latency."
|
||||
)
|
||||
|
||||
st.session_state.rag_settings["agentic"] = st.checkbox(
|
||||
"🤖 Agentic RAG",
|
||||
value=st.session_state.rag_settings["agentic"],
|
||||
help="**Agentic RAG**: Breaks complex questions into sub-questions and answers them independently before synthesizing. Best for multi-part or analytical queries. Slower but more thorough."
|
||||
)
|
||||
|
||||
# Main content tabs
|
||||
tab1, tab2 = st.tabs(["📚 Knowledge Base", "💬 Chat"])
|
||||
|
||||
# Tab 1: Knowledge Base Management
|
||||
with tab1:
|
||||
st.markdown("## 📚 Knowledge Base Management")
|
||||
|
||||
if not st.session_state.selected_kb:
|
||||
st.info("👈 Please select or create a knowledge base from the sidebar")
|
||||
else:
|
||||
st.markdown(f"### Current KB: **{st.session_state.selected_kb}**")
|
||||
|
||||
st.markdown("---")
|
||||
|
||||
# Document Upload
|
||||
col1, col2 = st.columns(2, gap="large")
|
||||
|
||||
with col1:
|
||||
st.markdown("#### 📤 Upload Documents")
|
||||
st.caption("Choose files to upload to the knowledge base")
|
||||
|
||||
# Initialize upload counter in session state
|
||||
if "upload_counter" not in st.session_state:
|
||||
st.session_state.upload_counter = 0
|
||||
|
||||
# Use dynamic key to reset file uploader after upload
|
||||
uploaded_files = st.file_uploader(
|
||||
"Choose files",
|
||||
accept_multiple_files=True,
|
||||
type=[ext.replace(".", "") for ext in Config.SUPPORTED_FILE_TYPES],
|
||||
label_visibility="collapsed",
|
||||
key=f"file_uploader_{st.session_state.upload_counter}"
|
||||
)
|
||||
|
||||
if uploaded_files and st.button("Upload Files", key="upload_files_btn", use_container_width=True):
|
||||
try:
|
||||
Config.ensure_upload_dir()
|
||||
progress_bar = st.progress(0)
|
||||
status_container = st.container()
|
||||
|
||||
for i, uploaded_file in enumerate(uploaded_files):
|
||||
with status_container:
|
||||
st.write(f"Processing: {uploaded_file.name}...")
|
||||
|
||||
# Save file
|
||||
file_path = Path(Config.UPLOAD_DIR) / uploaded_file.name
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(uploaded_file.getbuffer())
|
||||
|
||||
# Ingest file
|
||||
result = ingest_file(
|
||||
str(file_path),
|
||||
st.session_state.selected_kb,
|
||||
uploaded_file.name
|
||||
)
|
||||
|
||||
with status_container:
|
||||
if result["status"] == "success":
|
||||
st.success(f"✅ {uploaded_file.name}: {result['chunks']} chunks")
|
||||
else:
|
||||
st.error(f"❌ {uploaded_file.name}: {result.get('error', 'Unknown error')}")
|
||||
|
||||
progress_bar.progress((i + 1) / len(uploaded_files))
|
||||
|
||||
with status_container:
|
||||
st.success("🎉 Upload complete! Refreshing document list...")
|
||||
|
||||
# Increment counter to reset file uploader on next render
|
||||
st.session_state.upload_counter += 1
|
||||
|
||||
# Wait a moment before rerun so user can see the success message
|
||||
import time
|
||||
time.sleep(1.5)
|
||||
st.rerun()
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"❌ Upload failed: {str(e)}")
|
||||
import traceback
|
||||
st.code(traceback.format_exc())
|
||||
|
||||
with col2:
|
||||
st.markdown("#### 📁 Import Directory")
|
||||
st.caption("Import all supported files from a directory")
|
||||
|
||||
st.info("💡 **Tip:** Mount your directory as a volume in docker-compose.yml first. Paths are relative to `/workspace` in the container.")
|
||||
|
||||
dir_path = st.text_input(
|
||||
"Directory Path",
|
||||
placeholder="/workspace/mounted_docs",
|
||||
help="Path must exist inside the Docker container",
|
||||
label_visibility="collapsed"
|
||||
)
|
||||
|
||||
if st.button("Import Directory", key="import_dir_btn", use_container_width=True):
|
||||
if not dir_path:
|
||||
st.error("Please enter a directory path")
|
||||
else:
|
||||
with st.spinner("Importing files..."):
|
||||
results = ingest_directory(dir_path, st.session_state.selected_kb)
|
||||
|
||||
if not results:
|
||||
st.warning("No supported files found in directory")
|
||||
elif results[0].get("error") and "not found" in results[0].get("error", ""):
|
||||
st.error(f"Directory not found: {dir_path}")
|
||||
else:
|
||||
success_count = sum(1 for r in results if r.get("status") == "success")
|
||||
error_count = len(results) - success_count
|
||||
|
||||
if success_count > 0:
|
||||
st.success(f"✅ Imported {success_count} file(s)")
|
||||
if error_count > 0:
|
||||
st.warning(f"⚠️ {error_count} file(s) failed")
|
||||
|
||||
# Show first 5 results
|
||||
for result in results[:5]:
|
||||
if result.get("status") == "success":
|
||||
st.caption(f"✅ {result.get('filename', 'Unknown')}")
|
||||
else:
|
||||
st.caption(f"❌ {result.get('filename', 'Unknown')}: {result.get('error', 'Unknown error')}")
|
||||
|
||||
if len(results) > 5:
|
||||
st.caption(f"... and {len(results) - 5} more")
|
||||
|
||||
st.rerun()
|
||||
|
||||
st.divider()
|
||||
|
||||
# Document Grid
|
||||
st.markdown("#### 📄 Documents")
|
||||
|
||||
docs = get_documents_metadata(st.session_state.selected_kb)
|
||||
|
||||
if not docs:
|
||||
st.info("No documents in this knowledge base. Upload files to get started!")
|
||||
else:
|
||||
# Display in 3-column grid
|
||||
cols = st.columns(3)
|
||||
|
||||
for i, doc in enumerate(docs):
|
||||
with cols[i % 3]:
|
||||
# File icon based on type
|
||||
icon_map = {
|
||||
".py": "🐍", ".js": "📜", ".jsx": "⚛️",
|
||||
".ts": "📘", ".tsx": "⚛️", ".md": "📄",
|
||||
".pdf": "📕", ".json": "📋", ".yaml": "⚙️"
|
||||
}
|
||||
icon = icon_map.get(doc["file_type"], "📄")
|
||||
|
||||
# Create card
|
||||
tags_html = "".join([f'<span class="tag">{tag}</span>' for tag in doc["tags"]])
|
||||
|
||||
st.markdown(f"""
|
||||
<div class="doc-card">
|
||||
<div style="font-size: 24px;">{icon}</div>
|
||||
<div style="font-weight: bold; margin: 8px 0;">{doc["filename"]}</div>
|
||||
<div>{tags_html}</div>
|
||||
<div style="font-size: 11px; color: #888; margin-top: 8px;">
|
||||
{doc["formatted_date"]}<br>
|
||||
{doc["chunk_count"]} chunks
|
||||
</div>
|
||||
</div>
|
||||
""", unsafe_allow_html=True)
|
||||
|
||||
# Tab 2: Chat Interface
|
||||
with tab2:
|
||||
st.markdown("## 💬 Chat with Your Knowledge Base")
|
||||
|
||||
if not st.session_state.selected_kb:
|
||||
st.info("👈 Please select a knowledge base from the sidebar")
|
||||
else:
|
||||
# Display active strategies banner
|
||||
active_strategies = []
|
||||
if st.session_state.rag_settings["hybrid"]:
|
||||
active_strategies.append("🔀 Hybrid Search")
|
||||
if st.session_state.rag_settings["rerank"]:
|
||||
active_strategies.append("🎯 Reranking")
|
||||
if st.session_state.rag_settings["agentic"]:
|
||||
active_strategies.append("🤖 Agentic RAG")
|
||||
|
||||
if active_strategies:
|
||||
strategies_text = " + ".join(active_strategies)
|
||||
st.markdown(f"""
|
||||
<div style="background: linear-gradient(90deg, hsl(271 91% 65% / 0.15), hsl(160 84% 39% / 0.15));
|
||||
border-left: 3px solid hsl(271 91% 65%);
|
||||
padding: 12px 16px;
|
||||
border-radius: 4px;
|
||||
margin-bottom: 16px;
|
||||
font-family: 'JetBrains Mono', monospace;">
|
||||
<strong>🎯 Active Retrieval Strategies:</strong> {strategies_text}
|
||||
</div>
|
||||
""", unsafe_allow_html=True)
|
||||
else:
|
||||
st.markdown(f"""
|
||||
<div style="background: hsl(0 0% 10%);
|
||||
border-left: 3px solid hsl(0 0% 30%);
|
||||
padding: 12px 16px;
|
||||
border-radius: 4px;
|
||||
margin-bottom: 16px;
|
||||
font-family: 'JetBrains Mono', monospace;
|
||||
color: hsl(0 0% 60%);">
|
||||
<strong>📊 Mode:</strong> Standard Vector Search
|
||||
</div>
|
||||
""", unsafe_allow_html=True)
|
||||
|
||||
# Display chat history
|
||||
for message in st.session_state.chat_history:
|
||||
with st.chat_message(message["role"]):
|
||||
st.markdown(message["content"])
|
||||
|
||||
if message["role"] == "assistant" and message.get("sources"):
|
||||
with st.expander("📚 View Sources"):
|
||||
for i, source in enumerate(message["sources"][:3], 1):
|
||||
st.markdown(f"**Source {i}:**")
|
||||
st.code(source["text"][:200] + "..." if len(source["text"]) > 200 else source["text"])
|
||||
if source.get("metadata"):
|
||||
st.caption(f"From: {source['metadata'].get('filename', 'Unknown')}")
|
||||
if source.get("score"):
|
||||
st.caption(f"Score: {source['score']:.3f}")
|
||||
st.divider()
|
||||
|
||||
# Chat input
|
||||
if prompt := st.chat_input("Ask a question about your knowledge base..."):
|
||||
# Add user message
|
||||
st.session_state.chat_history.append({
|
||||
"role": "user",
|
||||
"content": prompt
|
||||
})
|
||||
|
||||
# Display user message
|
||||
with st.chat_message("user"):
|
||||
st.markdown(prompt)
|
||||
|
||||
# Generate response
|
||||
with st.chat_message("assistant"):
|
||||
with st.spinner("Thinking..."):
|
||||
try:
|
||||
rag_engine = RAGEngine(st.session_state.selected_kb)
|
||||
result = rag_engine.query(
|
||||
prompt,
|
||||
use_hybrid=st.session_state.rag_settings["hybrid"],
|
||||
use_rerank=st.session_state.rag_settings["rerank"],
|
||||
use_agentic=st.session_state.rag_settings["agentic"]
|
||||
)
|
||||
|
||||
answer = result.get("answer", "No answer generated")
|
||||
sources = result.get("sources", [])
|
||||
|
||||
st.markdown(answer)
|
||||
|
||||
# Add to history
|
||||
st.session_state.chat_history.append({
|
||||
"role": "assistant",
|
||||
"content": answer,
|
||||
"sources": sources
|
||||
})
|
||||
|
||||
# Display sources
|
||||
if sources:
|
||||
with st.expander("📚 View Sources"):
|
||||
for i, source in enumerate(sources[:3], 1):
|
||||
st.markdown(f"**Source {i}:**")
|
||||
st.code(source["text"][:200] + "..." if len(source["text"]) > 200 else source["text"])
|
||||
if source.get("metadata"):
|
||||
st.caption(f"From: {source['metadata'].get('filename', 'Unknown')}")
|
||||
if source.get("score"):
|
||||
st.caption(f"Score: {source['score']:.3f}")
|
||||
st.divider()
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Error: {str(e)}"
|
||||
st.error(error_msg)
|
||||
st.session_state.chat_history.append({
|
||||
"role": "assistant",
|
||||
"content": error_msg
|
||||
})
|
||||
|
||||
# Clear chat button
|
||||
if st.session_state.chat_history and st.button("🗑️ Clear Chat"):
|
||||
st.session_state.chat_history = []
|
||||
st.rerun()
|
||||
|
||||
0
app/utils/__init__.py
Normal file
0
app/utils/__init__.py
Normal file
68
docker-compose.yml
Normal file
68
docker-compose.yml
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
version: '3.8'
|
||||
|
||||
services:
|
||||
ollama:
|
||||
image: ollama/ollama:latest
|
||||
container_name: code-forge-ollama
|
||||
ports:
|
||||
- "11434:11434"
|
||||
volumes:
|
||||
- ollama_models:/root/.ollama
|
||||
environment:
|
||||
- OLLAMA_NUM_PARALLEL=1
|
||||
- OLLAMA_MAX_LOADED_MODELS=1
|
||||
- OLLAMA_KEEP_ALIVE=5m
|
||||
networks:
|
||||
- code-forge-network
|
||||
restart: unless-stopped
|
||||
|
||||
chromadb:
|
||||
image: chromadb/chroma:latest
|
||||
container_name: code-forge-chromadb
|
||||
ports:
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
- chroma_data:/chroma/chroma
|
||||
environment:
|
||||
- IS_PERSISTENT=TRUE
|
||||
- ANONYMIZED_TELEMETRY=FALSE
|
||||
networks:
|
||||
- code-forge-network
|
||||
restart: unless-stopped
|
||||
|
||||
app:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: code-forge-app
|
||||
ports:
|
||||
- "8501:8501" # Streamlit UI
|
||||
- "8051:8051" # MCP Server
|
||||
volumes:
|
||||
- ./app:/workspace/app
|
||||
- ./mcp_server:/workspace/mcp_server
|
||||
- uploads:/workspace/data/uploads
|
||||
# Uncomment to mount a local directory for bulk import:
|
||||
# - /path/to/your/docs:/workspace/mounted_docs:ro
|
||||
environment:
|
||||
- OLLAMA_HOST=http://ollama:11434
|
||||
- CHROMA_HOST=chromadb
|
||||
- CHROMA_PORT=8000
|
||||
- MCP_SERVER_HOST=0.0.0.0
|
||||
- MCP_SERVER_PORT=8051
|
||||
networks:
|
||||
- code-forge-network
|
||||
depends_on:
|
||||
- ollama
|
||||
- chromadb
|
||||
restart: unless-stopped
|
||||
|
||||
networks:
|
||||
code-forge-network:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
ollama_models:
|
||||
chroma_data:
|
||||
uploads:
|
||||
|
||||
0
mcp_server/__init__.py
Normal file
0
mcp_server/__init__.py
Normal file
216
mcp_server/server.py
Normal file
216
mcp_server/server.py
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
"""
|
||||
MCP Server for Code-Forge.
|
||||
Exposes RAG functionality via Model Context Protocol with HTTP transport.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
import uvicorn
|
||||
|
||||
from app.core.chroma_manager import get_chroma_manager
|
||||
from app.core.config import Config
|
||||
from app.core.kb_metadata import get_collection_stats, get_documents_metadata
|
||||
from app.core.rag_engine import RAGEngine
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize FastAPI and FastMCP
|
||||
app = FastAPI(title="Code-Forge MCP Server")
|
||||
mcp = FastMCP("code-forge")
|
||||
|
||||
# Add CORS middleware for Claude Desktop
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def search_knowledge_base(
|
||||
kb_name: str,
|
||||
query: str,
|
||||
use_hybrid: bool = False,
|
||||
use_rerank: bool = False,
|
||||
use_agentic: bool = False
|
||||
) -> dict:
|
||||
"""
|
||||
Search a knowledge base using RAG with optional advanced features.
|
||||
|
||||
Args:
|
||||
kb_name: Name of the knowledge base to search
|
||||
query: The search query
|
||||
use_hybrid: Enable hybrid search (vector + BM25 keyword search)
|
||||
use_rerank: Enable reranking of results for better relevance
|
||||
use_agentic: Enable agentic RAG for complex multi-step queries
|
||||
|
||||
Returns:
|
||||
dict with 'answer' and 'sources' keys
|
||||
"""
|
||||
try:
|
||||
# Validate KB exists
|
||||
chroma_manager = get_chroma_manager()
|
||||
if not chroma_manager.collection_exists(kb_name):
|
||||
return {
|
||||
"error": f"Knowledge base '{kb_name}' not found",
|
||||
"answer": "",
|
||||
"sources": []
|
||||
}
|
||||
|
||||
# Check if KB is empty
|
||||
count = chroma_manager.get_collection_count(kb_name)
|
||||
if count == 0:
|
||||
return {
|
||||
"error": f"Knowledge base '{kb_name}' is empty. Please upload documents first.",
|
||||
"answer": "",
|
||||
"sources": []
|
||||
}
|
||||
|
||||
# Initialize RAG engine
|
||||
rag_engine = RAGEngine(kb_name)
|
||||
|
||||
# Execute query
|
||||
result = rag_engine.query(
|
||||
question=query,
|
||||
use_hybrid=use_hybrid,
|
||||
use_rerank=use_rerank,
|
||||
use_agentic=use_agentic
|
||||
)
|
||||
|
||||
logger.info(f"Query executed on {kb_name}: {query[:50]}...")
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Search failed: {e}")
|
||||
return {
|
||||
"error": str(e),
|
||||
"answer": "",
|
||||
"sources": []
|
||||
}
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def list_knowledge_bases() -> List[str]:
|
||||
"""
|
||||
List all available knowledge bases.
|
||||
|
||||
Returns:
|
||||
List of knowledge base names
|
||||
"""
|
||||
try:
|
||||
chroma_manager = get_chroma_manager()
|
||||
kbs = chroma_manager.list_collections()
|
||||
logger.info(f"Listed {len(kbs)} knowledge bases")
|
||||
return kbs
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list knowledge bases: {e}")
|
||||
return []
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_kb_stats(kb_name: str) -> dict:
|
||||
"""
|
||||
Get statistics for a knowledge base.
|
||||
|
||||
Args:
|
||||
kb_name: Name of the knowledge base
|
||||
|
||||
Returns:
|
||||
dict with total_documents, total_chunks, and last_updated
|
||||
"""
|
||||
try:
|
||||
stats = get_collection_stats(kb_name)
|
||||
logger.info(f"Retrieved stats for {kb_name}")
|
||||
return stats
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get stats for {kb_name}: {e}")
|
||||
return {
|
||||
"error": str(e),
|
||||
"total_documents": 0,
|
||||
"total_chunks": 0,
|
||||
"last_updated": None
|
||||
}
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def list_documents(kb_name: str) -> List[dict]:
|
||||
"""
|
||||
List all documents in a knowledge base.
|
||||
|
||||
Args:
|
||||
kb_name: Name of the knowledge base
|
||||
|
||||
Returns:
|
||||
List of document metadata dicts
|
||||
"""
|
||||
try:
|
||||
docs = get_documents_metadata(kb_name)
|
||||
logger.info(f"Listed {len(docs)} documents in {kb_name}")
|
||||
return docs
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list documents in {kb_name}: {e}")
|
||||
return []
|
||||
|
||||
|
||||
# Health check endpoint
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
"""Health check endpoint."""
|
||||
return {"status": "healthy", "service": "code-forge-mcp"}
|
||||
|
||||
|
||||
# MCP endpoint
|
||||
@app.post("/mcp")
|
||||
async def mcp_endpoint(request: dict):
|
||||
"""
|
||||
MCP protocol endpoint.
|
||||
Handles JSON-RPC style requests from MCP clients.
|
||||
"""
|
||||
try:
|
||||
# FastMCP handles the request routing
|
||||
return await mcp.handle_request(request)
|
||||
except Exception as e:
|
||||
logger.error(f"MCP request failed: {e}")
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32603,
|
||||
"message": str(e)
|
||||
},
|
||||
"id": request.get("id")
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""Start the MCP server."""
|
||||
logger.info(f"Starting Code-Forge MCP Server on {Config.MCP_SERVER_HOST}:{Config.MCP_SERVER_PORT}")
|
||||
logger.info(f"MCP endpoint: http://{Config.MCP_SERVER_HOST}:{Config.MCP_SERVER_PORT}/mcp")
|
||||
|
||||
uvicorn.run(
|
||||
app,
|
||||
host=Config.MCP_SERVER_HOST,
|
||||
port=Config.MCP_SERVER_PORT,
|
||||
log_level="info"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
0
mcp_server/tools/__init__.py
Normal file
0
mcp_server/tools/__init__.py
Normal file
29
requirements.txt
Normal file
29
requirements.txt
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
# Core Framework
|
||||
streamlit>=1.30.0
|
||||
|
||||
# LlamaIndex and RAG
|
||||
llama-index>=0.10.0
|
||||
llama-index-llms-ollama>=0.1.0
|
||||
llama-index-embeddings-ollama>=0.1.0
|
||||
llama-index-vector-stores-chroma>=0.1.0
|
||||
|
||||
# Vector Database
|
||||
chromadb>=0.4.22
|
||||
|
||||
# Document Processing
|
||||
pymupdf>=1.23.0
|
||||
|
||||
# Reranking and Hybrid Search
|
||||
sentence-transformers>=2.2.0
|
||||
rank-bm25>=0.2.2
|
||||
|
||||
# MCP Server
|
||||
mcp>=1.0.0
|
||||
fastapi>=0.104.0
|
||||
uvicorn>=0.24.0
|
||||
pydantic>=2.0.0
|
||||
|
||||
# Utilities
|
||||
python-dotenv>=1.0.0
|
||||
requests>=2.31.0
|
||||
|
||||
18
start.sh
Normal file
18
start.sh
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
#!/bin/bash
|
||||
|
||||
# Set Python path to include workspace
|
||||
export PYTHONPATH=/workspace:$PYTHONPATH
|
||||
|
||||
# Start MCP server in background
|
||||
echo "Starting MCP server on port 8051..."
|
||||
cd /workspace && python mcp_server/server.py &
|
||||
MCP_PID=$!
|
||||
|
||||
# Start Streamlit UI
|
||||
echo "Starting Streamlit UI on port 8501..."
|
||||
cd /workspace && streamlit run app/main.py --server.headless true --server.port 8501 --server.address 0.0.0.0 &
|
||||
STREAMLIT_PID=$!
|
||||
|
||||
# Wait for both processes
|
||||
wait $MCP_PID $STREAMLIT_PID
|
||||
|
||||
Loading…
Add table
Reference in a new issue