first commit

This commit is contained in:
Bob Roberts 2025-10-22 13:29:28 -05:00
commit 8e5a46aa59
24 changed files with 3254 additions and 0 deletions

22
.env.example Normal file
View 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
View 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
View 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
View 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
View 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
View 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
View file

View file

0
app/core/__init__.py Normal file
View file

158
app/core/chroma_manager.py Normal file
View 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
View 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
View 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
View 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
View 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
}

View 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
View 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
View 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
View file

68
docker-compose.yml Normal file
View 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
View file

216
mcp_server/server.py Normal file
View 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()

View file

29
requirements.txt Normal file
View 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
View 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