Custom LLM with a RAG.

## Create a network so that all systems can talk to each other
docker network create rag-net

## Start Ollama - this manages and runs your models
docker run -d --network rag-net -v ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama

## Download model llama 3.2 using ollama (around 2GB)
docker exec -it ollama ollama pull llama3.2

## Download the Nomic model (can translate your content to vectors)
docker exec -it ollama ollama pull nomic-embed-text

## List models in Ollama
docker exec -it ollama ollama list

NAME                       ID              SIZE      MODIFIED       
nomic-embed-text:latest    0a109f422b47    274 MB    38 seconds ago    
llama3.2:latest            a80c4f17acd5    2.0 GB    3 minutes ago   

## Check containers that are running
docker ps -a

70dcc4e8535c   ollama/ollama   "/bin/ollama serve"      9 minutes ago   Up 9 minutes              0.0.0.0:11434->11434/tcp   ollama

## Install the vector database PgAI (postgres + timescale extension)
docker run -d --network rag-net -p 5432:5432 --name timescaledb -e POSTGRES_PASSWORD=password timescale/timescaledb-ha:pg16

## Run psql (to manage postgresql) inside the container
docker exec -it timescaledb psql -d postgres 

## Install the pgai extension in the postgres database (happens to be the default database)
CREATE EXTENSION IF NOT EXISTS ai CASCADE;

NOTICE:  installing required extension "vector"
NOTICE:  installing required extension "plpython3u"
CREATE EXTENSION

## Then we can verify that the extension got installed
postgres=# \dx
                                                    List of installed extensions
        Name         | Version |   Schema   |                                      Description                                      
---------------------+---------+------------+---------------------------------------------------------------------------------------
ai                  | 0.3.0   | public     | helper functions for ai workflows
plpgsql             | 1.0     | pg_catalog | PL/pgSQL procedural language
plpython3u          | 1.0     | pg_catalog | PL/Python3U untrusted procedural language
timescaledb         | 2.17.0  | public     | Enables scalable inserts and complex queries for time-series data (Community Edition)
timescaledb_toolkit | 1.18.0  | public     | Library of analytical hyperfunctions, time-series pipelining, and other SQL utilities
vector              | 0.7.4   | public     | vector data type and ivfflat and hnsw access methods
(6 rows)

# Let's create e virtual environment to encapsulate all libraries from the global installation
python3 -m venv llm-pipeline
source llm-pipeline/bin/activate

# To connect with our database
pip install psycopg2
# To parse our hugo markdown files
pip install markdown python-frontmatter
# Our IDE
pip install jupyterlab

# This starts the IDE and you can access it in your browser
jupyter lab

import sys
import psycopg2
import os
import frontmatter

def parse_markdown_files(directory):
    markdown_data = []

    # Use os.walk to traverse the directory tree
    for root, _, files in os.walk(directory):
        for filename in files:
            if filename.endswith('.md'):
                # Construct the full file path
                filepath = os.path.join(root, filename)
                
                # Open and read the markdown file
                with open(filepath, 'r', encoding='utf-8') as file:
                    # Parse front matter and content using frontmatter library
                    post = frontmatter.load(file)
                    
                    # Extract title from front matter
                    title = post.get('title', 'No Title')
                    
                    # Extract content (the markdown content itself)
                    content = post.content
                    
                    # Append to markdown_data list as a dictionary
                    markdown_data.append({
                        "title": title,
                        "content": content
                    })

    return markdown_data

directory_path = '/Users/I/workspace/raphaelbauer.com/content'

markdown_data = parse_markdown_files(directory_path)

def connect_db():
    return psycopg2.connect( # use the credentials of your postgresql database 
        host = 'localhost',
        database = 'postgres',
        user = 'postgres',
        password = 'password',
        port = '5432'
    )

conn = connect_db()
cur = conn.cursor()
cur.execute("""
        CREATE TABLE IF NOT EXISTS documents (
            id SERIAL PRIMARY KEY,
            title TEXT,
            content TEXT,
            embedding VECTOR(768)
        );
    """)
conn.commit()
cur.close()
conn.close()

conn = connect_db()
cur = conn.cursor()

# use the port at which your ollama service is running.
for doc in markdown_data:
    cur.execute("""
        INSERT INTO documents (title, content, embedding)
        VALUES (
            %(title)s,
            %(content)s,
            ollama_embed('nomic-embed-text', concat(%(title)s, ' - ', %(content)s), _host=>'http://ollama:11434')
        )
    """, doc)

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

conn = connect_db()
cur = conn.cursor()
    
cur.execute("""
    SELECT title, content, vector_dims(embedding) 
    FROM documents LIMIT 10;
""")

rows = cur.fetchall()
for row in rows:
    print(f"Title: {row[0]}, Content: {row[1]}, Embedding Dimensions: {row[2]}")

cur.close()
conn.close()

query = "Can you describe how modern QA should look like?"

conn = connect_db()
cur = conn.cursor()
    
# Embed the query using the ollama_embed function
cur.execute("""
    SELECT ollama_embed('nomic-embed-text', %s, _host=>'http://ollama:11434');
""", (query,))
query_embedding = cur.fetchone()[0]

# Retrieve relevant documents based on cosine distance
cur.execute("""
    SELECT title, content, 1 - (embedding <=> %s) AS similarity
    FROM documents
    ORDER BY similarity DESC
    LIMIT 1;
""", (query_embedding,))

rows = cur.fetchall()
    
# Prepare the context for generating the response
context = "\n\n".join([f"Title: {row[0]}\nContent: {row[1]}" for row in rows])
print(context)

cur.close()
conn.close()

conn = connect_db()
cur = conn.cursor()

# Generate the response using the ollama_generate function
cur.execute("""
    SELECT ollama_generate('llama3.2', %s, _host=>'http://ollama:11434');
""", (f"""
DOCUMENT:
{context}

QUESTION:
{query}

INSTRUCTIONS:
Answer the users QUESTION using the DOCUMENT text above.
Keep your answer ground in the facts of the DOCUMENT.
If the DOCUMENT doesn’t contain the facts to answer the QUESTION then please say so.
""",))
    
model_response = cur.fetchone()[0]
print(model_response['response'])
    
cur.close()
conn.close()