version: "3.8"

services:
  # Ollama — Self-hosted LLM runtime
  ollama:
    image: ollama/ollama:latest
    ports:
      - "127.0.0.1:11434:11434"
    volumes:
      - ollama_data:/root/.ollama
    restart: unless-stopped

  # ChromaDB — Persistent vector database
  chromadb:
    image: chromadb/chromadb:latest
    ports:
      - "127.0.0.1:8000:8000"
    volumes:
      - chroma_data:/chroma/chroma
    environment:
      CHROMA_AUTH_TOKEN: ${CHROMA_AUTH_TOKEN}
      CHROMA_SERVER_AUTH_CREDENTIALS: ${CHROMA_CREDENTIALS}
    restart: unless-stopped

  # RAG API — FastAPI service connecting everything
  rag-api:
    build: ./rag-api
    ports:
      - "127.0.0.1:8080:8080"
    environment:
      OLLAMA_BASE_URL: http://ollama:11434
      CHROMADB_URL: http://chromadb:8000
      EMBEDDING_MODEL: nomic-embed-text
      LLM_MODEL: llama3.1:8b
      COLLECTION_NAME: company-kb
    depends_on:
      - ollama
      - chromadb
    restart: unless-stopped

  # Cloudflare Tunnel — Secure access without open ports
  cloudflared:
    image: cloudflare/cloudflared:latest
    command: tunnel --no-autoupdate run --token ${CLOUDFLARE_TUNNEL_TOKEN}
    restart: unless-stopped
    depends_on:
      - rag-api

volumes:
  ollama_data:
  chroma_data:

FROM python:3.11-slim

WORKDIR /app

RUN pip install fastapi uvicorn pydantic ollama chromadb langchain-community langchain-text-splitters pypdf

COPY requirements.txt .
RUN pip install -r requirements.txt

COPY . .

EXPOSE 8080

CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8080"]

import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import ollama
import chromadb

app = FastAPI(title="Self-Hosted RAG API")

OLLAMA_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
CHROMADB_URL = os.getenv("CHROMADB_URL", "http://localhost:8000")
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
LLM_MODEL = os.getenv("LLM_MODEL", "llama3.1:8b")
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "company-kb")

# Initialize clients
ollama_client = ollama.Client(base_url=OLLAMA_URL)
chroma_client = chromadb.HttpClient(host=CHROMADB_URL.split(":")[1].strip("/"), port=int(CHROMADB_URL.split(":")[2]))
collection = chroma_client.get_or_create_collection(COLLECTION_NAME)


class QueryRequest(BaseModel):
    question: str
    top_k: int = 3


class IngestRequest(BaseModel):
    text: str
    metadata: dict = {}


@app.post("/ingest")
def ingest_document(req: IngestRequest):
    """Ingest a document chunk into the vector store."""
    try:
        # Generate embedding
        embedding_response = ollama_client.embeddings(model=EMBEDDING_MODEL, prompt=req.text)
        embedding = embedding_response["embedding"]

        # Store in ChromaDB
        doc_id = f"doc_{len(collection.get()['ids']) + 1}"
        collection.upsert(
            ids=[doc_id],
            embeddings=[embedding],
            documents=[req.text],
            metadatas=[req.metadata]
        )
        return {"status": "ok", "document_id": doc_id}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/query")
def query_rag(req: QueryRequest):
    """Query the RAG system: retrieve relevant docs + generate answer."""
    try:
        # Get embedding for the question
        question_embedding = ollama_client.embeddings(model=EMBEDDING_MODEL, prompt=req.question)["embedding"]

        # Retrieve relevant documents
        results = collection.query(query_embeddings=[question_embedding], n_results=req.top_k)

        if not results["documents"][0]:
            return {"answer": "No relevant documents found in the knowledge base.", "sources": []}

        # Build context from retrieved documents
        context = "\n\n".join(results["documents"][0])

        # Generate answer using LLM
        prompt = f"""You are a helpful AI assistant. Answer the user's question based only on the provided context.

Context:
{context}

Question: {req.question}

Answer:"""

        response = ollama_client.generate(
            model=LLM_MODEL,
            prompt=prompt,
            stream=False
        )

        return {
            "answer": response["response"],
            "sources": results["documents"][0],
            "distances": results["distances"][0]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
def health():
    return {"status": "healthy", "ollama_model": LLM_MODEL, "embedding_model": EMBEDDING_MODEL}

# Pull the LLM and embedding model
docker exec ollama ollama pull llama3.1:8b
docker exec ollama ollama pull nomic-embed-text

# Verify models are loaded
docker exec ollama ollama list

CHROMA_AUTH_TOKEN=your-chroma-auth-token
CHROMA_SERVER_AUTH_CREDENTIALS=admin:secret
CLOUDFLARE_TUNNEL_TOKEN=your-tunnel-token

# Ingest a text snippet
curl -X POST http://localhost:8080/ingest \
  -H "Content-Type: application/json" \
  -d '{
    "text": "Our company policy states that employees get 20 days of paid time off per year...",
    "metadata": {"source": "employee-handbook.pdf", "page": 12}
  }'

# Query the knowledge base
curl -X POST http://localhost:8080/query \
  -H "Content-Type: application/json" \
  -d '{
    "question": "How many days of PTO do employees get?",
    "top_k": 3
  }'

import os
from pathlib import Path
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import requests

KB_URL = "http://localhost:8080/ingest"

def ingest_directory(directory: str):
    """Ingest all PDFs and text files from a directory."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        separators=["\n\n", "\n", ". ", " "]
    )

    for filepath in Path(directory).rglob("*"):
        if filepath.suffix in (".pdf", ".txt", ".md"):
            print(f"Ingesting: {filepath}")

            if filepath.suffix == ".pdf":
                loader = PyPDFLoader(str(filepath))
            else:
                loader = TextLoader(str(filepath))

            documents = loader.load()
            chunks = splitter.split_documents(documents)

            for i, chunk in enumerate(chunks):
                response = requests.post(KB_URL, json={
                    "text": chunk.page_content,
                    "metadata": {
                        "source": str(filepath),
                        "chunk": i
                    }
                })
                if response.status_code != 200:
                    print(f"  Failed chunk {i}: {response.text}")

    print("Ingestion complete.")

# Usage: ingest_directory("./documents/")

sudo fallocate -l 4G /swapfile
sudo chmod 600 /swapfile
sudo mkswap /swapfile
sudo swapon /swapfile
echo '/swapfile none swap sw 0 0' | sudo tee -a /etc/fstab

# Watch Ollama memory usage
docker stats ollama chromadb rag-api

# Check inference speed
time curl -s http://localhost:11434/api/generate -d '{"model":"llama3.1:8b","prompt":"Hello","stream":false}'

# 1. Bind all services to localhost only (already in docker-compose)
# 2. Use Cloudflare Tunnel for encrypted access
# 3. Enable ChromaDB auth
# 4. Rate-limit the API

# Add to rag-api Dockerfile:
# pip install slowapi

# Add to main.py:
from slowapi import Limiter
from slowapi.util import get_remote_address
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter

@app.post("/query")
@limiter.limit("10/minute")
def query_rag(req: QueryRequest):
    ...