RAG Introduction with Ollama OpenAI API¶

import os
import requests
import numpy as np
import pandas as pd
from textwrap import wrap
from math import sqrt
from openai import OpenAI

# === Configuration ===
OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://ollama:11434")

# === Model Configuration ===
OLLAMA_EMBEDDING_MODEL = "llama3.2:latest"
OLLAMA_LLM_MODEL = "llama3.2:latest"

# Initialize OpenAI client pointing to Ollama
client = OpenAI(
    base_url=f"{OLLAMA_HOST}/v1",
    api_key="ollama"  # Required by library but ignored by Ollama
)

print(f"Ollama host: {OLLAMA_HOST}")
print(f"Embedding model: {OLLAMA_EMBEDDING_MODEL}")
print(f"LLM model: {OLLAMA_LLM_MODEL}")

# Sample document: COVID-19 Omicron variant information
# (Embedded for self-contained demo - based on scientific literature)

SAMPLE_TEXT = """
The Omicron variant of SARS-CoV-2, first identified in South Africa in November 2021, 
rapidly spread across the globe and became the dominant variant in many countries by early 2022. 
This variant exhibited significant mutations in the spike protein, raising concerns about 
vaccine efficacy and therapeutic interventions.

In France, the emergence of Omicron led to a rapid replacement of the Delta variant during 
the winter of 2021-2022. Epidemiological surveillance showed that Omicron cases doubled 
approximately every two to three days during its initial spread, significantly faster than 
previous variants.

The Omicron variant is characterized by approximately 30 mutations in the spike protein alone, 
including mutations at positions K417N, N440K, G446S, S477N, T478K, E484A, Q493R, G496S, 
Q498R, N501Y, and Y505H. Many of these mutations are located in the receptor-binding domain 
(RBD), which is crucial for viral entry into host cells.

Studies in France demonstrated that while Omicron showed increased transmissibility compared 
to Delta, it was associated with reduced severity of disease. Hospitalization rates and 
intensive care unit admissions were lower per infection compared to the Delta wave, though 
the sheer number of cases still strained healthcare systems.

The immune evasion properties of Omicron were substantial. Research showed reduced neutralization 
by antibodies elicited by previous infection with earlier variants or by primary vaccination 
series. However, booster doses significantly improved protection against severe disease.

Mathematical modeling of the Omicron invasion in France utilized multi-variant epidemiological 
models to understand the dynamics of variant replacement. These models incorporated factors 
such as cross-immunity between variants, vaccine coverage, and waning immunity over time.

The basic reproduction number (R0) of Omicron was estimated to be significantly higher than 
Delta, with estimates ranging from 8 to 15 depending on the population and setting. This 
high transmissibility was a key factor in its rapid global spread.

French public health authorities responded to the Omicron wave with enhanced testing capacity, 
acceleration of booster vaccination campaigns, and implementation of sanitary passes requiring 
up-to-date vaccination status for access to certain venues and activities.

Subsequent sub-lineages of Omicron, including BA.2, BA.4, BA.5, and later BQ and XBB variants, 
continued to evolve with additional mutations conferring further immune evasion properties. 
This ongoing evolution necessitated updates to vaccine formulations and continued surveillance.

The experience with Omicron in France and globally highlighted the importance of genomic 
surveillance, rapid response capabilities, and adaptable public health strategies in managing 
emerging variants of concern during a pandemic.
"""

# Chunk the text into smaller pieces for embedding
wrapped_text = wrap(SAMPLE_TEXT.strip(), 1000)
print(f"Document chunked into {len(wrapped_text)} pieces")

len(wrapped_text)

# Test embedding a single chunk
response = client.embeddings.create(
    model=OLLAMA_EMBEDDING_MODEL,
    input=wrapped_text[0]
)

embedding = response.data[0].embedding
print(f"Text chunk (first 100 chars): {wrapped_text[0][:100]}...")
print(f"Embedding dimensions: {len(embedding)}")
print(f"First 5 values: {embedding[:5]}")

# Generate embeddings for all chunks
embeddings = []

for i, text in enumerate(wrapped_text):
    response = client.embeddings.create(
        model=OLLAMA_EMBEDDING_MODEL,
        input=text
    )
    embedding = np.array(response.data[0].embedding)
    embeddings.append(embedding)

print(f"Created {len(embeddings)} embeddings")

len(embeddings[0])

len(embeddings)

import pandas as pd

vector_data_base = pd.DataFrame({ 'text': wrapped_text,
                                  'embeddings': embeddings })

vector_data_base.head()

vector_data_base['embeddings'][0]

def cosine_distance(a,b):
    return(float(1.-((np.dot(a,b))/(sqrt(np.dot(a,a))*sqrt(np.dot(b,b))))))

def best_answers(n,query,database):
    distances = []
    for i in range(len(vector_data_base)):
        distances.append(cosine_distance(database['embeddings'][i],query))
    local_db = database.copy()
    local_db['distances'] = distances
    local_db = local_db.nsmallest(n,'distances')
    return(list(local_db['text']))

def encode_into_single_embedding(intext):
    """Encode text into a single embedding vector using Ollama."""
    embedded_chunks = []
    wrapped = wrap(intext, 1000)
    
    for text in wrapped:
        response = client.embeddings.create(
            model=OLLAMA_EMBEDDING_MODEL,
            input=text
        )
        embedding = np.array(response.data[0].embedding)
        embedded_chunks.append(embedding)
    
    # If only one chunk, return it directly
    if len(embedded_chunks) == 1:
        return embedded_chunks[0]
    
    # Combine multiple embeddings by averaging
    combined = np.mean(embedded_chunks, axis=0)
    
    # Normalize the combined embedding
    norm = np.linalg.norm(combined)
    if norm > 0:
        combined = combined / norm
    
    return combined

# Test query encoding
encoded = encode_into_single_embedding("What do you know about the Omicron variant?")
print(f"Query embedding shape: {len(encoded)} dimensions")

# Find the most relevant chunks
relevant_chunks = best_answers(3, encoded, vector_data_base)
print("Top 3 relevant chunks:")
for i, chunk in enumerate(relevant_chunks):
    print(f"\n--- Chunk {i+1} ---")
    print(chunk[:200] + "...")

# System prompt for RAG
SYSTEM_PROMPT = """You are a helpful AI assistant. Answer questions based on the provided context.
If the answer is not in the context, say so clearly. Be concise but thorough."""

# Conversation history for multi-turn chat
conversation_history = []

def build_messages(context_chunks, user_query, history=None):
    """Build messages list for chat completion with RAG context."""
    context = "\n\n---\n\n".join(context_chunks)
    
    system_content = f"""{SYSTEM_PROMPT}

## Retrieved Context:
{context}
"""
    
    messages = [{"role": "system", "content": system_content}]
    
    # Add conversation history if provided
    if history:
        messages.extend(history)
    
    # Add current user query
    messages.append({"role": "user", "content": user_query})
    
    return messages

def get_llm_response(messages, max_tokens=500):
    """Get response from LLM via Ollama's OpenAI-compatible API."""
    response = client.chat.completions.create(
        model=OLLAMA_LLM_MODEL,
        messages=messages,
        max_tokens=max_tokens,
        temperature=0.7
    )
    return response.choices[0].message.content

def chat(user_query, n_docs=5):
    """RAG chat function with conversation memory."""
    global conversation_history
    
    # Encode query and retrieve relevant chunks
    query_embedding = encode_into_single_embedding(user_query)
    relevant_chunks = best_answers(n_docs, query_embedding, vector_data_base)
    
    # Build messages with context and history
    messages = build_messages(relevant_chunks, user_query, conversation_history)
    
    # Get response from LLM
    response = get_llm_response(messages)
    
    # Update conversation history (without the system message)
    conversation_history.append({"role": "user", "content": user_query})
    conversation_history.append({"role": "assistant", "content": response})
    
    print(response)
    return response

# First question about the Omicron variant
chat("What do you know about the Omicron variant in France?")

# Follow-up question (uses conversation history)
chat("What mutations does it have?")

# View conversation history
print("=== Conversation History ===")
for msg in conversation_history:
    role = msg["role"].upper()
    content = msg["content"][:200] + "..." if len(msg["content"]) > 200 else msg["content"]
    print(f"\n[{role}]: {content}")

## 10. Utility Functions

def reset_conversation():
    """Reset conversation history to start fresh."""
    global conversation_history
    conversation_history = []
    print("✓ Conversation history cleared")

# Uncomment to reset:
# reset_conversation()

# Try another question!
# chat("How did France respond to the Omicron wave?")