# Update system
apt update && apt upgrade -y

# Install base tools
apt install -y python3-pip git curl wget tmux htop

# Install Docker (recommended for simplified deployment)
curl -fsSL https://get.docker.com | sh
usermod -aG docker $USER

# requirements.txt
langchain==0.3.1
langchain-community==0.3.1
langchain-text-splitters==0.3.0
chromadb==0.5.0
ollama==0.3.0
unstructured==0.15.0
python-docx==1.1.0
openpyxl==3.1.0

# document_loader.py
from langchain_community.document_loaders import (
    PyPDFLoader,
    TextLoader,
    UnstructuredMarkdownLoader,
    Docx2txtLoader,
    DirectoryLoader
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os

def load_documents(directory: str):
    """Load all supported documents from directory"""
    
    # Load by file type
    pdf_loader = DirectoryLoader(directory, glob="**/*.pdf", loader_cls=PyPDFLoader)
    text_loader = DirectoryLoader(directory, glob="**/*.txt", loader_cls=TextLoader)
    md_loader = DirectoryLoader(directory, glob="**/*.md", loader_cls=UnstructuredMarkdownLoader)
    docx_loader = DirectoryLoader(directory, glob="**/*.docx", loader_cls=Docx2txtLoader)
    
    # Merge all documents
    documents = []
    documents.extend(pdf_loader.load())
    documents.extend(text_loader.load())
    documents.extend(md_loader.load())
    documents.extend(docx_loader.load())
    
    print(f"Loaded {len(documents)} documents")
    return documents

def chunk_documents(documents, chunk_size=500, chunk_overlap=50):
    """Split documents into appropriately sized text chunks"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ".", " ", ""]
    )
    chunks = text_splitter.split_documents(documents)
    print(f"Split into {len(chunks)} chunks")
    return chunks

# embedding_store.py
import chromadb
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma

def create_vector_store(chunks, persist_directory="./chroma_db"):
    """Create vector database and store document chunks"""
    
    # Use local Ollama embeddings
    embeddings = OllamaEmbeddings(
        model="nomic-embed-text",
        base_url="http://localhost:11434"
    )
    
    # Create/load ChromaDB
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=persist_directory
    )
    
    # Persist to disk
    vectorstore.persist()
    print(f"Vector store saved to {persist_directory}")
    
    return vectorstore

def load_vector_store(persist_directory="./chroma_db"):
    """Load existing vector store from disk"""
    embeddings = OllamaEmbeddings(
        model="nomic-embed-text",
        base_url="http://localhost:11434"
    )
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embeddings
    )
    return vectorstore

# rag_qa.py
from langchain.chains import RetrievalQA
from langchain_community.llms import Ollama
from langchain.prompts import PromptTemplate

# Load vector store
vectorstore = load_vector_store()

# Create retriever
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5}
)

# Local LLM
llm = Ollama(
    model="qwen2.5:7b",
    base_url="http://localhost:11434",
    temperature=0.1
)

# Custom prompt template
qa_prompt = PromptTemplate(
    input_variables=["context", "question"],
    template="""Answer the following question based on the context below. If the context
does not contain relevant information, honestly say "I cannot find the answer
in the available materials."

Context:
{context}

Question: {question}

Answer:"""
)

# Create RAG chain
rag_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": qa_prompt}
)

# Execute Q&A
def ask_question(question: str):
    result = rag_chain.invoke({"query": question})
    return {
        "answer": result["result"],
        "sources": [doc.page_content[:200] for doc in result["source_documents"][:3]]
    }

# Hybrid search: Vector + Keyword
retriever = vectorstore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"k": 10, "score_threshold": 0.7}
)

# Query Expansion
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor

compressor = LLMChainExtractor.from_llm(llm)
compression_retriever = ContextualCompressionRetriever(
    base_retriever=retriever,
    document_compressor=compressor
)

# Reranking for better quality
# Use Cross-Encoder to rerank retrieved results

# main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import uvicorn

app = FastAPI(title="RAG Knowledge Base API", version="1.0.0")

class QuestionRequest(BaseModel):
    question: str
    k: int = 5
    include_sources: bool = True

class QuestionResponse(BaseModel):
    answer: str
    sources: Optional[List[str]] = None
    confidence: float

@app.post("/ask", response_model=QuestionResponse)
async def ask_question(req: QuestionRequest):
    try:
        result = ask_question(req.question)
        return QuestionResponse(
            answer=result["answer"],
            sources=result["sources"] if req.include_sources else None
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/stats")
async def get_stats():
    """Get knowledge base statistics"""
    vectorstore = load_vector_store()
    collection = vectorstore._collection
    count = collection.count()
    return {"document_count": count, "collection": "knowledge_base"}

@app.post("/ingest")
async def ingest_documents(directory: str):
    """Reload documents and update knowledge base"""
    documents = load_documents(directory)
    chunks = chunk_documents(documents)
    vectorstore = create_vector_store(chunks)
    return {"status": "success", "documents": len(documents), "chunks": len(chunks)}

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000)

# Start Ollama (if not already running)
ollama serve &

# Pull required models
ollama pull qwen2.5:7b
ollama pull nomic-embed-text

# Start API service
uvicorn main:app --host 0.0.0.0 --port 8000 --reload

# Or use Docker Compose for one-click deployment
docker-compose up -d

# monitoring.py
import time
from datetime import datetime

class RAGMetrics:
    def __init__(self):
        self.query_count = 0
        self.total_latency = 0.0
        self.error_count = 0
    
    def record_query(self, latency: float, success: bool = True):
        self.query_count += 1
        self.total_latency += latency
        if not success:
            self.error_count += 1
    
    def get_stats(self):
        avg_latency = self.total_latency / self.query_count if self.query_count > 0 else 0
        return {
            "total_queries": self.query_count,
            "avg_latency_ms": round(avg_latency * 1000, 2),
            "error_rate": round(self.error_count / self.query_count * 100, 2) if self.query_count > 0 else 0,
            "uptime": datetime.now().isoformat()
        }

metrics = RAGMetrics()

# docker-compose.yml
version: '3.8'

services:
  rag-api:
    build: .
    ports:
      - "8000:8000"
    volumes:
      - ./docs:/app/docs
      - ./chroma_db:/app/chroma_db
    environment:
      - OLLAMA_HOST=http://ollama:11434
    depends_on:
      - ollama

  ollama:
    image: ollama/ollama:latest
    ports:
      - "11434:11434"
    volumes:
      - ./ollama_models:/root/.ollama
    deploy:
      resources:
        limits:
          memory: 8G

  chromadb:
    image: chromadb/chroma:latest
    ports:
      - "8001:8000"
    volumes:
      - ./chroma_db:/chroma/chroma