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"Query: {query}\nContext: {context}",))
    
model_response = cur.fetchone()[0]
print(model_response['response'])
    
cur.close()
conn.close()