commit 8e5a46aa599befb865625cd51262eff5006ccdc9 Author: Bob Roberts Date: Wed Oct 22 13:29:28 2025 -0500 first commit diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..eea3180 --- /dev/null +++ b/.env.example @@ -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 + diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..64b836b --- /dev/null +++ b/.gitignore @@ -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 + diff --git a/.streamlit/config.toml b/.streamlit/config.toml new file mode 100644 index 0000000..9ebbc12 --- /dev/null +++ b/.streamlit/config.toml @@ -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 + diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..83332b6 --- /dev/null +++ b/Dockerfile @@ -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"] + diff --git a/README.md b/README.md new file mode 100644 index 0000000..081cbd4 --- /dev/null +++ b/README.md @@ -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 +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** + diff --git a/SHARING.md b/SHARING.md new file mode 100644 index 0000000..9f995be --- /dev/null +++ b/SHARING.md @@ -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 + 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://:8501` +- **MCP**: `http://: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! 🎉 + diff --git a/app/__init__.py b/app/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/components/__init__.py b/app/components/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/core/__init__.py b/app/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/core/chroma_manager.py b/app/core/chroma_manager.py new file mode 100644 index 0000000..7320893 --- /dev/null +++ b/app/core/chroma_manager.py @@ -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 + diff --git a/app/core/config.py b/app/core/config.py new file mode 100644 index 0000000..dbfa285 --- /dev/null +++ b/app/core/config.py @@ -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 + diff --git a/app/core/health.py b/app/core/health.py new file mode 100644 index 0000000..9603612 --- /dev/null +++ b/app/core/health.py @@ -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 + diff --git a/app/core/ingestion.py b/app/core/ingestion.py new file mode 100644 index 0000000..fcbfb0d --- /dev/null +++ b/app/core/ingestion.py @@ -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 + diff --git a/app/core/kb_metadata.py b/app/core/kb_metadata.py new file mode 100644 index 0000000..4424153 --- /dev/null +++ b/app/core/kb_metadata.py @@ -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 + } + diff --git a/app/core/markdown_preprocessor.py b/app/core/markdown_preprocessor.py new file mode 100644 index 0000000..3563b85 --- /dev/null +++ b/app/core/markdown_preprocessor.py @@ -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 + diff --git a/app/core/rag_engine.py b/app/core/rag_engine.py new file mode 100644 index 0000000..7382366 --- /dev/null +++ b/app/core/rag_engine.py @@ -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": [] + } + diff --git a/app/main.py b/app/main.py new file mode 100644 index 0000000..1661801 --- /dev/null +++ b/app/main.py @@ -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(""" + +""", unsafe_allow_html=True) + +# Header +st.markdown(""" +
+

🔥 Code-Forge

+

Localized Multi-KB RAG System with MCP Integration

+
+""", 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('

🟢 Ollama Connected

', unsafe_allow_html=True) + if ollama_status.get("models"): + st.caption(f"Models: {', '.join(ollama_status['models'][:3])}") + else: + st.markdown('

🔴 Ollama Disconnected

', unsafe_allow_html=True) + st.caption(ollama_status.get("error", "Unknown error")) + + if chroma_status["status"] == "healthy": + st.markdown('

🟢 ChromaDB Connected

', unsafe_allow_html=True) + else: + st.markdown('

🔴 ChromaDB Disconnected

', 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'{tag}' for tag in doc["tags"]]) + + st.markdown(f""" +
+
{icon}
+
{doc["filename"]}
+
{tags_html}
+
+ {doc["formatted_date"]}
+ {doc["chunk_count"]} chunks +
+
+ """, 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""" +
+ 🎯 Active Retrieval Strategies: {strategies_text} +
+ """, unsafe_allow_html=True) + else: + st.markdown(f""" +
+ 📊 Mode: Standard Vector Search +
+ """, 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() + diff --git a/app/utils/__init__.py b/app/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..16811da --- /dev/null +++ b/docker-compose.yml @@ -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: + diff --git a/mcp_server/__init__.py b/mcp_server/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mcp_server/server.py b/mcp_server/server.py new file mode 100644 index 0000000..5260183 --- /dev/null +++ b/mcp_server/server.py @@ -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() + diff --git a/mcp_server/tools/__init__.py b/mcp_server/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..1140518 --- /dev/null +++ b/requirements.txt @@ -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 + diff --git a/start.sh b/start.sh new file mode 100644 index 0000000..1561fcf --- /dev/null +++ b/start.sh @@ -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 +