Simple AI Workflow with Local RAG using Ollama and Postgres

%%capture
!pip install ollama psycopg2-binary python-dotenv

# Ollama: Python client for running local LLMs and embeddings
import ollama

# psycopg2: PostgreSQL adapter for Python
import psycopg2

# Standard library
import os
import hashlib  # For generating content hashes to prevent duplicate document insertions

# Database connection parameters
DB_HOST = "localhost"
DB_PORT = "5432"
DB_NAME = "vectors"
DB_USER = "postgres"
DB_PASSWORD = "postgres"

# Models to use
CHAT_MODEL = "llama3.2"
EMBEDDING_MODEL = "mxbai-embed-large"
EMBEDDING_DIM = 1024  # mxbai-embed-large produces 1024-dimensional vectors

# Test Ollama connection
try:
    response = ollama.list()
    print(f"Ollama is running. Available models: {[m.model for m in response.models]}")
except Exception as e:
    print(f"Error connecting to Ollama: {e}")
    print("Make sure Ollama is running (try 'ollama serve' in terminal)")

# Test Postgres connection
try:
    conn = psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD
    )
    conn.close()
    print(f"Postgres connection successful to {DB_NAME}")
except Exception as e:
    print(f"Error connecting to Postgres: {e}")
    print("Make sure the Docker container is running")

def get_db_connection():
    """Create and return a database connection."""
    return psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD
    )

# Enable pgvector extension and create table
conn = get_db_connection()
cur = conn.cursor() # Create a cursor

# Enable the vector extension
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")

# Create the documents table
cur.execute(f"""
    CREATE TABLE IF NOT EXISTS documents (
        id BIGSERIAL PRIMARY KEY,
        content TEXT NOT NULL,
        embedding vector({EMBEDDING_DIM}),
        content_hash TEXT UNIQUE
    );
""")

conn.commit()
cur.close()
conn.close()

print("Database setup complete: pgvector enabled and documents table created")

# Load the FastHTML documentation
with open("FastHTML.txt", "r", encoding="utf-8") as f:
    document = f.read()

print(f"Loaded document with {len(document)} characters")
document[:200]

def chunk_text(text: str, chunk_size: int = 1000, overlap: int = 200) -> list[str]:
    """
    Split text into overlapping chunks.
    
    Args:
        text: The text to split
        chunk_size: Target size for each chunk (in characters)
        overlap: Number of characters to overlap between chunks
        
    Returns:
        List of text chunks
    """
    chunks = []
    start = 0
    
    while start < len(text):
        # Get the chunk
        end = start + chunk_size
        chunk = text[start:end]
        
        # Try to break at a natural boundary (newline or period)
        if end < len(text):
            # Look for a good break point near the end
            last_newline = chunk.rfind('\n')
            last_period = chunk.rfind('. ')
            
            # Use the later of the two break points, if found in the last 20% of the chunk
            break_point = max(last_newline, last_period)
            if break_point > chunk_size * 0.8:
                chunk = text[start:start + break_point + 1]
                end = start + break_point + 1
        
        chunks.append(chunk.strip())
        
        # Move start position, accounting for overlap
        start = end - overlap
    
    return chunks

# Split the document into chunks
chunks = chunk_text(document, chunk_size=1000, overlap=200)

print(f"Created {len(chunks)} chunks")
print(f"\nExample chunk (chunk 0):\n{chunks[0][:500]}...")
print(f"\nChunk sizes: min={min(len(c) for c in chunks)}, max={max(len(c) for c in chunks)}, avg={sum(len(c) for c in chunks)//len(chunks)}")

def get_embedding(text: str) -> list[float]:
    """
    Get the embedding vector for a piece of text using Ollama.
    
    Args:
        text: The text to embed
        
    Returns:
        A list of floats representing the embedding vector
    """
    response = ollama.embed(
        model=EMBEDDING_MODEL,
        input=text
    )
    return response.embeddings[0]

response = ollama.embed(
    model=EMBEDDING_MODEL,
    input="What is FastHTML?"
)

print(f"Response fields: {response.model_dump().keys()}")
print(f"Number of embeddings: {len(response.embeddings)}")
print(f"Embedding dimension: {len(response.embeddings[0])}")

# Test the embedding function with a sample text
sample_embedding = get_embedding("What is FastHTML?")
print(f"Embedding dimension: {len(sample_embedding)}")
print(f"First 10 values: {sample_embedding[:10]}")

# Generate embeddings for all chunks
# This may take a few minutes depending on the number of chunks
print(f"Generating embeddings for {len(chunks)} chunks...")

embeddings = []
for i, chunk in enumerate(chunks):
    embedding = get_embedding(chunk)
    embeddings.append(embedding)
    if (i + 1) % 50 == 0:
        print(f"  Processed {i + 1}/{len(chunks)} chunks")

print(f"Done! Generated {len(embeddings)} embeddings")

def get_content_hash(text: str) -> str:
    """Generate a hash of the content to use as a unique identifier."""
    return hashlib.md5(text.encode()).hexdigest() # creates a hash as a 32-character hex string

# Insert all chunks and embeddings into Postgres using upsert
# This prevents duplicates if you re-run the notebook
print(f"Upserting {len(chunks)} documents into Postgres...")

conn = get_db_connection()
cur = conn.cursor()

for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
    content_hash = get_content_hash(chunk)
    
    cur.execute("""
        INSERT INTO documents (content, embedding, content_hash)
        VALUES (%s, %s, %s)
        ON CONFLICT (content_hash) 
        DO UPDATE SET content = EXCLUDED.content, embedding = EXCLUDED.embedding;
    """, (chunk, embedding, content_hash))
    
    if (i + 1) % 50 == 0:
        print(f"  Processed {i + 1}/{len(chunks)} documents")
        conn.commit()

conn.commit()
cur.close()
conn.close()

print(f"Done! Upserted {len(chunks)} documents into Postgres")

# Sanity check: verify documents are in the database
conn = get_db_connection()
cur = conn.cursor()

cur.execute("SELECT COUNT(*) FROM documents;")
count = cur.fetchone()[0]

cur.close()
conn.close()

print(f"Documents table contains {count} entries")

def search_documents(query: str, match_count: int = 5) -> list[dict]:
    """
    Search for documents similar to the query.
    
    Args:
        query: The search query
        match_count: Number of results to return
        
    Returns:
        List of matching documents with their similarity scores
    """
    # Get embedding for the query
    query_embedding = get_embedding(query)
    
    # Search for similar documents
    conn = get_db_connection()
    cur = conn.cursor()

    """
    <=> is cosine distance operator (0 = identical, 2 = opposite)
    %s::vector casts the Python list parameter to a Postgres vector type
    1 - (embedding <=> %s::vector) as similarity - Converts distance to similarity (1 = identical, -1 = opposite)
    """
    
    cur.execute("""
        SELECT id, content, 1 - (embedding <=> %s::vector) as similarity
        FROM documents
        ORDER BY embedding <=> %s::vector
        LIMIT %s;
    """, (query_embedding, query_embedding, match_count))
    
    results = []
    for row in cur.fetchall():
        results.append({
            "id": row[0],
            "content": row[1],
            "similarity": row[2]
        })
    
    cur.close()
    conn.close()
    
    return results

# Test the search function
results = search_documents("How do I create a route in FastHTML?", match_count=3)

print(f"Found {len(results)} matching documents:\n")
for i, doc in enumerate(results):
    print(f"--- Result {i+1} (similarity: {doc['similarity']:.4f}) ---")
    print(doc['content'][:300] + "...")
    print()

def ask_fasthtml_tutor(question: str, num_docs: int = 5) -> str:
    """
    Ask the FastHTML tutor a question.
    
    Args:
        question: The user's question about FastHTML
        num_docs: Number of documents to retrieve for context
        
    Returns:
        The tutor's response
    """
    # Step 1: Retrieve relevant documents
    docs = search_documents(question, match_count=num_docs)
    
    # Step 2: Build context from retrieved documents
    context = "\n\n---\n\n".join([doc["content"] for doc in docs])
    
    # Step 3: Create the prompt with context
    system_prompt = """You are a helpful FastHTML tutor. Answer questions about FastHTML based on the provided documentation context. 

If the context doesn't contain enough information to answer the question, say so honestly. 
Always be accurate and cite specific examples from the documentation when possible."""

    user_prompt = f"""Context from FastHTML documentation:

{context}

---

Question: {question}

Please provide a clear, helpful answer based on the documentation above."""

    # Step 4: Generate response using Ollama
    response = ollama.chat(
        model=CHAT_MODEL,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
    )
    
    return response['message']['content']

# Test the complete RAG pipeline
question = "How do I create a route in FastHTML?"
answer = ask_fasthtml_tutor(question)

print(f"Question: {question}\n")
print(f"Answer:\n{answer}")

# Try your own question!
my_question = "What are FastTags in FastHTML?"
answer = ask_fasthtml_tutor(my_question)

print(f"Question: {my_question}\n")
print(f"Answer:\n{answer}")

# Verify connection status and demonstrate proper cleanup
# In our implementation, connections are opened and closed within each function,
# but here's how you would explicitly close a connection if needed:

def close_connection_safely(conn):
    """Safely close a database connection."""
    if conn is not None and not conn.closed:
        conn.close()
        print("Connection closed successfully")
    else:
        print("No open connection to close")

# Example: Create a connection, use it, and close it properly
test_conn = get_db_connection()
cur = test_conn.cursor()
cur.execute("SELECT COUNT(*) FROM documents;")
count = cur.fetchone()[0]
print(f"Documents in database: {count}")
cur.close()
close_connection_safely(test_conn)

# Drop the documents table to start fresh
# Uncomment and run this cell only when you want to reset the database

# conn = get_db_connection()
# cur = conn.cursor()
# cur.execute("DROP TABLE IF EXISTS documents;")
# conn.commit()
# cur.close()
# conn.close()
# print("Documents table dropped. You can now re-run the notebook with a different embedding model.")