# Install Ollama
curl -fsSL https://ollama.ai/install.sh | sh

# Pull the LLM for chat
ollama pull llama3.1:8b

# Pull a local embedding model (no cloud required)
ollama pull nomic-embed-text

# Optional: pull a more capable model for complex reasoning
ollama pull qwen2.5:32b

# Verify models are available
ollama list

# Create a virtual environment
python -m venv .venv
source .venv/bin/activate  # Linux/macOS
# .venv\Scripts\activate   # Windows

# Install dependencies
pip install \
  langchain \
  langchain-ollama \
  langchain-chroma \
  chromadb \
  pypdf \
  python-dotenv \
  rich  # For prettier terminal output

local-ai/
├── .env                    # Configuration (no API keys needed!)
├── main.py                 # Chatbot entrypoint
├── ingest.py               # Document ingestion pipeline
├── rag_chain.py            # RAG chain construction
├── documents/              # Source documents for RAG
│   ├── manual.pdf
│   └── notes.md
└── vectorstore/            # ChromaDB storage (auto-created)

# ingest.py
import os
from pathlib import Path
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma

VECTORSTORE_PATH = "./vectorstore"
DOCUMENTS_PATH = "./documents"
EMBED_MODEL = "nomic-embed-text"

def load_documents(docs_path: str) -> list:
    """Load all PDF and text documents from a directory."""
    documents = []
    path = Path(docs_path)

    for file in path.rglob("*.pdf"):
        loader = PyPDFLoader(str(file))
        documents.extend(loader.load())
        print(f"Loaded PDF: {file.name} ({len(documents)} total pages)")

    for file in path.rglob("*.md"):
        loader = TextLoader(str(file), encoding="utf-8")
        documents.extend(loader.load())
        print(f"Loaded Markdown: {file.name}")

    for file in path.rglob("*.txt"):
        loader = TextLoader(str(file), encoding="utf-8")
        documents.extend(loader.load())
        print(f"Loaded Text: {file.name}")

    return documents

def ingest_documents(docs_path: str = DOCUMENTS_PATH) -> Chroma:
    """Ingest documents into ChromaDB with local embeddings."""
    print(f"Loading documents from {docs_path}...")
    raw_docs = load_documents(docs_path)

    if not raw_docs:
        raise ValueError(f"No documents found in {docs_path}")

    print(f"Splitting {len(raw_docs)} documents into chunks...")
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,       # Characters per chunk
        chunk_overlap=200,     # Overlap to preserve context at boundaries
        length_function=len,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    chunks = splitter.split_documents(raw_docs)
    print(f"Created {len(chunks)} chunks")

    print(f"Generating embeddings with {EMBED_MODEL} (local)...")
    embeddings = OllamaEmbeddings(
        model=EMBED_MODEL,
        base_url="http://localhost:11434",
    )

    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=VECTORSTORE_PATH,
        collection_name="local_rag",
    )
    print(f"Vector store created with {vectorstore._collection.count()} vectors")
    return vectorstore

if __name__ == "__main__":
    ingest_documents()
    print("Ingestion complete. Run main.py to start the chatbot.")

# rag_chain.py
from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_chroma import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

VECTORSTORE_PATH = "./vectorstore"
EMBED_MODEL = "nomic-embed-text"
CHAT_MODEL = "llama3.1:8b"

RAG_PROMPT_TEMPLATE = """You are an assistant answering questions based on provided context.

Use only the information in the context below to answer. If the answer is not in the context,
say "I don't have enough information in the provided documents to answer this question."

Context:
{context}

Question: {question}

Answer:"""

def format_docs(docs) -> str:
    """Format retrieved documents into a single context string."""
    return "\n\n---\n\n".join(
        f"Source: {doc.metadata.get('source', 'Unknown')}\n{doc.page_content}"
        for doc in docs
    )

def build_rag_chain():
    """Build the local RAG chain with Ollama + ChromaDB."""
    embeddings = OllamaEmbeddings(
        model=EMBED_MODEL,
        base_url="http://localhost:11434",
    )

    vectorstore = Chroma(
        persist_directory=VECTORSTORE_PATH,
        embedding_function=embeddings,
        collection_name="local_rag",
    )

    retriever = vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 4},  # Retrieve top 4 most relevant chunks
    )

    llm = ChatOllama(
        model=CHAT_MODEL,
        base_url="http://localhost:11434",
        temperature=0.3,
        num_ctx=8192,
    )

    prompt = ChatPromptTemplate.from_template(RAG_PROMPT_TEMPLATE)

    rag_chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    return rag_chain, retriever

def stream_rag_response(chain, question: str) -> str:
    """Stream a RAG response to the terminal."""
    full_response = ""
    for chunk in chain.stream(question):
        print(chunk, end="", flush=True)
        full_response += chunk
    print()
    return full_response

# main.py
import sys
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rag_chain import build_rag_chain, stream_rag_response

console = Console()

def main():
    console.print(Panel(
        "[bold green]Local RAG Chatbot[/bold green]\n"
        "Powered by Ollama + ChromaDB + LangChain\n"
        "No cloud. No API keys. Your data stays local.",
        title="AI Assistant",
        border_style="green"
    ))

    console.print("[yellow]Loading local models and vector store...[/yellow]")

    try:
        chain, retriever = build_rag_chain()
    except Exception as e:
        console.print(f"[red]Failed to initialize: {e}[/red]")
        console.print("[yellow]Make sure Ollama is running and documents are ingested.[/yellow]")
        console.print("[yellow]Run: python ingest.py[/yellow]")
        sys.exit(1)

    console.print("[green]Ready. Type your question or 'quit' to exit.[/green]\n")

    while True:
        try:
            question = input("You: ").strip()
        except (KeyboardInterrupt, EOFError):
            console.print("\n[yellow]Goodbye.[/yellow]")
            break

        if not question:
            continue
        if question.lower() in ("quit", "exit", "q"):
            console.print("[yellow]Goodbye.[/yellow]")
            break

        # Show retrieved context sources for transparency
        docs = retriever.invoke(question)
        if docs:
            sources = list({doc.metadata.get("source", "unknown") for doc in docs})
            console.print(f"[dim]Sources: {', '.join(sources)}[/dim]")

        console.print("[bold]Assistant:[/bold] ", end="")
        stream_rag_response(chain, question)
        print()

if __name__ == "__main__":
    main()

# 1. Start Ollama (if not auto-started)
ollama serve

# 2. Create a test document
mkdir -p documents
cat > documents/notes.md << 'EOF'
# AI Learning Notes

## Transformer Architecture
Transformers use self-attention mechanisms to process sequences in parallel.
The attention mechanism allows each token to attend to all other tokens.
Key components: multi-head attention, feed-forward layers, layer normalization.

## RAG Systems
RAG (Retrieval-Augmented Generation) combines retrieval and generation.
Documents are chunked, embedded, and stored in a vector database.
At query time, relevant chunks are retrieved and passed to the LLM as context.
EOF

# 3. Ingest documents
python ingest.py

# 4. Start chatbot
python main.py

# test_local_setup.py — Verify your setup is working
import requests
import ollama

def test_ollama_connection() -> bool:
    """Test Ollama service is running."""
    try:
        resp = requests.get("http://localhost:11434/api/tags", timeout=5)
        models = [m["name"] for m in resp.json().get("models", [])]
        print(f"Ollama running. Models: {models}")
        return True
    except requests.exceptions.ConnectionError:
        print("Ollama not running. Start with: ollama serve")
        return False

def test_inference(model: str = "llama3.1:8b") -> bool:
    """Test basic inference is working."""
    try:
        response = ollama.chat(
            model=model,
            messages=[{"role": "user", "content": "Reply with just the word: working"}]
        )
        output = response["message"]["content"].strip().lower()
        print(f"Inference test: {output}")
        return "working" in output
    except Exception as e:
        print(f"Inference failed: {e}")
        return False

def test_embeddings(model: str = "nomic-embed-text") -> bool:
    """Test embedding model is working."""
    try:
        response = ollama.embeddings(model=model, prompt="test embedding")
        dim = len(response["embedding"])
        print(f"Embeddings test: {dim}-dimensional vector")
        return dim > 0
    except Exception as e:
        print(f"Embeddings failed: {e}")
        return False

if __name__ == "__main__":
    all_pass = all([
        test_ollama_connection(),
        test_inference(),
        test_embeddings(),
    ])
    print(f"\nSetup status: {'All tests passed' if all_pass else 'Some tests failed'}")