app.py

# Basic flask stuff for building http APIs and rendering html templates
from flask import Flask, render_template, redirect, url_for, request, session, jsonify

# Bootstrap integration with flask so we can make pretty pages
from flask_bootstrap import Bootstrap

# Flask forms integrations which save insane amounts of time
from flask_wtf import FlaskForm
from wtforms import StringField, SubmitField, PasswordField, SelectField, TextAreaField, IntegerField, FloatField, DateTimeField
from wtforms.validators import DataRequired

# Basic python stuff
import os
import json
import functools
import requests
import datetime

# Mongo stuff
import pymongo
from bson import ObjectId

# Some nice formatting for code
import misaka

# Import OpenAI, Azure and Mistral libraries
from openai import OpenAI
from openai import AzureOpenAI
from mistralai import Mistral

# Nice way to load environment variables for deployments
from dotenv import load_dotenv
load_dotenv()

# Create the Flask app object
app = Flask(__name__)

# Session key
app.config['SECRET_KEY'] = os.environ["SECRET_KEY"]
app.config['SESSION_COOKIE_NAME'] = 'extbrain_admin'

# API Key for app to serve API requests to clients
API_KEY = os.environ["API_KEY"]

# Determine which LLM/embedding service we will be using
# this could be llamacpp, mistral or openai
service_type = os.environ["SERVICE"]

# optionally connect the oai/mistral clients if they're configured
if "MISTRAL_API_KEY" in os.environ:
    mistral_client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])    
    model_name = os.environ["MODEL_NAME"]
    DEFAULT_SCORE_CUT = 0.9

if "OPENAI_API_KEY" in os.environ:
    oai_client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
    model_name = os.environ["MODEL_NAME"]
    embed_model_name = "text-embedding-3-small"
    DEFAULT_SCORE_CUT = 0.7

if "AZURE_API_KEY" in os.environ:
    oai_client = AzureOpenAI(api_key=os.environ["AZURE_API_KEY"], api_version="2023-12-01-preview", azure_endpoint=os.environ["AZURE_ENDPOINT"])
    model_name = os.environ["MODEL_NAME"]
    embed_model_name = os.environ["AZURE_EMBED_MODEL"]
    DEFAULT_SCORE_CUT = 0.85

# User Auth
users_string = os.environ["USERS"]
users = json.loads(users_string)

# Some handy defaults that control the app
DEFAULT_SYSTEM_MESSAGE = "You are a friendly chatbot. You help the user answer questions, solve problems and make plans.  You think deeply about the question and provide a detailed, accurate response."
DEFAULT_TEMPERATURE = 0.7
DEFAULT_LIMIT = 3
DEFAULT_CANDIDATES = 100
DEFAULT_FACTS_PER_PAGE = 25
DEFAULT_TEXT_SCORE_CUT = 3.0 # This is a hack, reranking is better!

## Controls the fact synthesis 
FACT_SYSTEM_MESSAGE = "You are sumbot, you take articles, blog posts and social media posts and you summarize the content from these into facts.  Each fact should be a single line and you think very carefully about what facts are important to the source material.  You capture enough facts in each article that it could be reproduced later using just the facts you provide.  You are expert level at this task and never miss a fact."
FACT_PROMPT = "Summaraize all the facts in this {context} into bullet points, one per line: {paste_data}"
FACT_TEMPERATURE = 0.1

# Load up the local model and embedder config if using the "local" service
# Local in this context means llama.cpp running in server mode with an InstructML format model
# and InstructorVec service running
if service_type == "local":
    # Load the llm model config
    with open("model.json", 'r',  encoding='utf-8') as file:
        local_model = json.load(file)

    # Load the embedder config
    with open("embedder.json", 'r',  encoding='utf-8') as file:
        local_embedder = json.load(file)

    DEFAULT_SCORE_CUT = 0.81

# Connect to mongo using environment variables
client = pymongo.MongoClient(os.environ["MONGO_CON"])
db = client[os.environ["MONGO_DB"]]

# Set the dimensions for the embedding models based on what service we've configured
if service_type == "local":
    vector_dimensions = 1024  # BreadVec 1024d
if service_type == "mistral":
    vector_dimensions = 1024 # mistral-embed 1024d
if service_type == "openai" or service_type == "azure":
    vector_dimensions = 1536  # text-embedding-3-small or text-ada-002 1536d

# If this is the first time we've run extBrain, we will test to see if the collections exist
# and if they do not, we create them and the lexical and vector search indexes
try:
    # Collection does not exist, create it
    db.create_collection('facts')
    db.create_collection('chunks')

    # Create the text search index on facts collection
    command = {
    "createSearchIndexes": "facts",
    "indexes": [
        {
            "name": 'default',
            "definition": {
                "mappings": {
                    "dynamic": False,
                    "fields": {
                    "context": {
                        "type": "string"
                    },
                    "fact": {
                        "type": "string"
                    }
                    }
                },
                "analyzer": "lucene.english"
            },
        }
    ]}
    db.command(command)

    # Create the vector search and text index on chunks collection
    command = {
    "createSearchIndexes": "chunks",
    "indexes": [
        {
            "name": 'default',
            "definition": {
                "analyzer": "lucene.english",
                "mappings": {
                    "dynamic": False,
                    "fields": {
                        "fact_chunk": {
                            "type": "string"
                        },
                        "chunk_embedding": [
                            {
                                "type": "knnVector",
                                "dimensions": vector_dimensions,
                                "similarity": "cosine"
                            }
                        ]    
                    } 
                }
            },
        }
    ]}
    db.command(command)
except:
    # We've already configured the collection/search indexes before just keep going
    pass

# Make it pretty because I can't :(
Bootstrap(app)

# A form for asking your external brain questions
class QuestionForm(FlaskForm):
    question = StringField('Question 💬', validators=[DataRequired()])
    submit = SubmitField('Submit')
    rag = SelectField('Generation Type', choices=[("augmented", "Augmented (Facts)"), ("freeform", "Freeform")])
    hybrid = SelectField('Hybrid Search', choices=[("hybrid", "Vector + Lexical"), ("vector", "Vector Only")])
    temperature = FloatField('LLM Temperature', default=DEFAULT_TEMPERATURE, validators=[DataRequired()])
    candidates = IntegerField('Vector Candidates', default=DEFAULT_CANDIDATES, validators=[DataRequired()])
    limit = IntegerField('Chunks', default=DEFAULT_LIMIT, validators=[DataRequired()])
    score_cut = FloatField('Score Cut Off', default=DEFAULT_SCORE_CUT, validators=[DataRequired()])
    
# A form for testings semantic search of the chunks
class VectorSearchForm(FlaskForm):
    question = StringField('Question 💬', validators=[DataRequired()])
    k = IntegerField("K Value", validators=[DataRequired()])
    score_cut = FloatField("Score Cut Off", validators=[DataRequired()])
    submit = SubmitField('Search')
  
# A form for pasting in your text for summarization
class PasteForm(FlaskForm):
    context = StringField('Context (article name, url, todo, idea)', validators=[DataRequired()])
    paste_data = TextAreaField('Paste in Text ✂️', validators=[DataRequired()])
    submit = SubmitField('Summarize')

# A form for reviewing and saving the summarized facts
class FactSearchForm(FlaskForm):
    query = StringField('Search Facts', validators=[DataRequired()])
    submit = SubmitField('Search')

# A form to edit the facts
class FactEditForm(FlaskForm):
    context = StringField('Context (article name, url, todo, idea)', validators=[DataRequired()])
    fact = StringField('Fact', validators=[DataRequired()])
    save = SubmitField('Save')

# A form for reviewing and saving the summarized facts
class SaveFactsForm(FlaskForm):
    context = StringField('Context (article name, url, todo, idea)', validators=[DataRequired()])
    fact_data = TextAreaField('Summarized Facts (one per line)', validators=[DataRequired()])
    submit = SubmitField('Save Facts')

# A form for configuring the chunk regen process
class ChunksForm(FlaskForm):
    fact_limit = IntegerField('Number of facts per chunk', validators=[DataRequired()])
    which_facts = SelectField('Which facts', choices=[("new", "New/Changed Facts"), ("all", "All Facts")])
    submit = SubmitField('Generate Chunks')

# Amazing, I hate writing this stuff
class LoginForm(FlaskForm):
    username = StringField('Username', validators=[DataRequired()])
    password = PasswordField('Password', validators=[DataRequired()])
    submit = SubmitField('Login')

# Function to call the local text embedder (768d)
def embed_local(text):
    response = requests.get(local_embedder["embedding_endpoint"], params={"text":text, "instruction": "Represent this text for retrieval:" }, headers={"accept": "application/json"})
    vector_embedding = response.json()
    return vector_embedding

# Call OpenAI's new embedder (1536d)
def embed_oai(text):
    text = text.replace("\n", " ")
    return oai_client.embeddings.create(input = [text], model=embed_model_name).data[0].embedding

# Call mistral's embedder (1024d)
def embed_mistral(text):
    text = text.replace("\n", " ")
    return mistral_client.embeddings(model="mistral-embed", input=[text]).data[0].embedding

# Use whichever embedder makes sense for the configured service
def embed(text):
    if service_type == "local":
        return embed_local(text)
    if service_type == "openai" or service_type == "azure":
        return embed_oai(text)
    if service_type == "mistral":
        return embed_mistral(text)

# Query mistral models
def llm_mistral(prompt, system_message, temperature):
    messages = [{"role": "system", "content": system_message},{"role": "user", "content": prompt}]
    response = mistral_client.chat.complete(model=model_name, temperature=temperature, messages=messages)
    return response.choices[0].message.content

# Query OpenAI models
def llm_oai(prompt, system_message, temperature):
    messages = [{"role": "system", "content": system_message},{"role": "user", "content": prompt}]
    response = oai_client.chat.completions.create(model=model_name, temperature=temperature, messages=messages)
    return response.choices[0].message.content

# Call llm using the llm configuration
def llm_local(prompt, system_message, temperature):
    client = OpenAI(api_key=local_model["api_key"], base_url=local_model["base_url"])
    messages=[{"role": "system", "content": system_message},{"role": "user", "content": prompt}]
    response = client.chat.completions.create(model=local_model["model"], temperature=temperature, messages=messages)
    return response.choices[0].message.content

# Determine which LLM we should call depending on what's configured
def llm(user_prompt, system_message, temperature=DEFAULT_TEMPERATURE):
    if service_type == "local":
        return llm_local(user_prompt, system_message, temperature)
    if service_type == "openai" or service_type == "azure":
        return llm_oai(user_prompt, system_message, temperature)
    if service_type == "mistral":
        return llm_mistral(user_prompt, system_message, temperature)


# Chat with model with or without augmentation
def chat(prompt, system_message, augmented=True, temperature=DEFAULT_TEMPERATURE, candidates=DEFAULT_CANDIDATES, limit=DEFAULT_LIMIT, score_cut=DEFAULT_SCORE_CUT, hybrid=True, text_score_cut=DEFAULT_TEXT_SCORE_CUT):
    # If we're doing RAG, vector search, assemble chunks and query with them
    fact_chunks = []
    chunk_string = ""
    if augmented:
        chunks = search_chunks(prompt, candidates, limit, score_cut)
        # Vector search chunks
        for chunk in chunks:
            fact_chunks.append(chunk["fact_chunk"])
            chunk_string += chunk["fact_chunk"]
        # Text search chunks if hybrid is enabled
        if hybrid:
            lexical_chunks = search_chunks_text(prompt, limit, text_score_cut)
            for chunk in lexical_chunks:
                fact_chunks.append(chunk["fact_chunk"])
                chunk_string += chunk["fact_chunk"]
        # The most important guardrail for any LLM app:  Don't answer the question if there's no chunks!
        if chunk_string != "":
            llm_prompt = F"Facts:\n{chunk_string}\nAnswer this question using only the relevant facts above: {prompt}"
        else:
            store_dontknow(prompt)
            return {"chunks": [], "completion": "I couldn't find any stored facts to answer that question. Try another question."}
    # If we're not, just query the model directly, without augmentation
    else:
        llm_prompt = prompt

    completion = llm(llm_prompt, system_message, temperature)
    return {"chunks": fact_chunks, "completion": completion}

# Store a fact array in mongo
def store_facts(facts, user, context):
    col = db["facts"]
    dt = datetime.datetime.now()
    for fact in facts:
        col.insert_one({"user": user, "date": dt, "context": context, "fact": fact })

# When we don't know the answer, log the question to the "dontknow" collection
# Brain admins can review this later to see what new knowledge needs to be added
def store_dontknow(prompt):
    col = db["dontknow"]
    dt = datetime.datetime.now()
    col.insert_one({"date": dt, "prompt": prompt })

# Get chunks based on semantic search
def search_chunks(prompt, candidates, limit, score_cut):
    # Get the embedding for the prompt first
    vector = embed(prompt)

    # Build the Atlas vector search aggregation
    vector_search_agg = [
        {
            "$vectorSearch": { 
                "index": "default",
                "path": "chunk_embedding",
                "queryVector": vector,
                "numCandidates": candidates, 
                "limit": limit
            }
        },
        {
            "$project": {
                "fact_chunk": 1,
                "score": {"$meta": "vectorSearchScore"}
            }
        },
        {
            "$match": {
                "score": { "$gte": score_cut }
            }
        }
    ]

    # Connect to chunks, run query, return results
    col = db["chunks"]
    chunk_records = col.aggregate(vector_search_agg)
    # Return this as a list instead of a cursor
    return chunk_records

# Get chunks based on text search 
def search_chunks_text(prompt, limit, text_score_cut):
    # Build the Atlas vector search aggregation
    text_search_agg = [
        {   
            "$search": {
                "text": {
                    "path": "fact_chunk",
                    "query": prompt
                }
            }   
        },
        {
            "$limit": limit
        },
        {
            "$project": {
                "fact_chunk": 1,
                "score": {"$meta": "searchScore"}
            }
        },
        {
            "$match": {
                "score": { "$gte": text_score_cut }
            }
        }
    ]

    # Connect to chunks, run query, return results
    col = db["chunks"]
    chunk_records = col.aggregate(text_search_agg)
    # Return this as a list instead of a cursor
    return chunk_records

# Get all facts
def get_facts(skip,limit):
    col = db["facts"]
    fact_records = col.find().skip(skip).limit(limit).sort([("date", -1)])
    return fact_records

# Searches facts using text search
def search_facts(query):
    
    # Build the Atlas vector search aggregation
    search_agg = [
        {
            "$search": {
                "text": {
                    "path": ["context", "fact"],
                    "query": query
                }
            }
        },
        {
            "$project": {
                "_id": 1,
                "user": 1,
                "date": 1,
                "context": 1,
                "fact": 1,
                "score": {"$meta": "searchScore"}
            }
        },
        {
            "$limit": DEFAULT_FACTS_PER_PAGE
        }
    ]

    # Connect to chunks, run query, return results
    col = db["facts"]
    fact_records = col.aggregate(search_agg)
    # Return this as a list instead of a cursor
    return fact_records

# Return the count of facts in the system
def count_facts():
    col = db["facts"]
    return  col.count_documents({})

# Clean up a list of facts
def clean_facts(facts):            
    # Parse and save the facts here!
    facts = facts.replace("- ", "") # remove bullet points
    facts = facts.replace("* ", "") # also bullet points
    facts = facts.replace("\r", "") # No carraige return
    facts = facts.replace("\t", "") # No tabs

    # Split the facts by line
    fact_list = facts.split("\n")

    # Remove any empty facts from the list
    facts_clean = list(filter(None, fact_list))

    return(facts_clean)

# Save a chunk to the chunks collection and it's embedding
def save_chunk(col, chunk, context):
    col.insert_one({"fact_chunk": chunk, "chunk_embedding": embed(chunk), "context": context})

# Generate chunk collection using facts, fixed fact limit but chunks can only contain a single context
def chunk_by_context(chunk_limit, force=False):
    # Use the facts collection
    facts_col = db["facts"]

    # Connect to chunks
    chunks_col = db["chunks"]
    
    # Regenerate all chunks or just the changed facts
    if force:
        # Nuke all the existing facts
        chunks_col.delete_many({})
        facts = facts_col.find()
    else:
        facts = facts_col.find({"chunked": {'$ne': True}})

    # Build up a dict with the key being context, user and date
    # Store all the facts under each one.   
    result_dict = {}
    for doc in facts:
        context = doc["context"]
        user = doc["user"]
        date = doc["date"].strftime('%Y-%m-%d')
        fact = doc["fact"]
        
        # Create a unique key using the three fields
        key = (context, user, date)
        
        # Add the facts to the key
        if key not in result_dict:
            result_dict[key] = []
        result_dict[key].append(fact)

    for key, value in result_dict.items():
        context, user, date = key
        facts = value

        # Iterate through the facts and group them
        fact_header = f"{context} from {user} on {date}"
        fact_count = 0
        fact_string = ""
        for fact in facts:
            fact_count += 1
            fact_string += f" - {fact}\n"
            if fact_count == chunk_limit:
                chunk_string = f"{fact_header}\n{fact_string}"
                save_chunk(chunks_col, chunk_string, context)
                fact_string = ""
                fact_count = 0
        # Catch the last one
        if fact_count > 0:
            chunk_string = f"{fact_header}\n{fact_string}"
            save_chunk(chunks_col, chunk_string, context)

    # Update the facts collection to indicate that all outstanding facts have been chunked
    # This allows us to selectively chunk new and updated data
    filter_chunked = { "$or": [ { "chunked": False }, { "chunked": { "$exists": False } } ] }
    update_chunked = { "$set": { "chunked": True } }
    facts_col.update_many(filter_chunked, update_chunked)

# Define a decorator to check if the user is authenticated
# No idea how this works...  Magic.
def login_required(view):
    @functools.wraps(view)
    def wrapped_view(**kwargs):
        if users != None:
            if session.get("user") is None:
                return redirect(url_for('login'))
        return view(**kwargs)        
    return wrapped_view

# The default question view
@app.route('/', methods=['GET', 'POST'])
@login_required
def index():

    # Question form for the external brain
    form = QuestionForm()

    # If user is prompting send it
    if form.validate_on_submit():

         # Get the form variables
        form_result = request.form.to_dict(flat=True)
        q = form_result["question"]

        if form_result["hybrid"] == "hybrid":
            hybrid = True
        else:
            hybrid = False

        if form_result["rag"] == "augmented":
            llm_result = chat(q, DEFAULT_SYSTEM_MESSAGE, True, float(form_result["temperature"]), int(form_result["candidates"]), int(form_result["limit"]), float(form_result["score_cut"]), hybrid)
        else:
            llm_result = chat(q, DEFAULT_SYSTEM_MESSAGE, False, float(form_result["temperature"]), int(form_result["candidates"]), int(form_result["limit"]), float(form_result["score_cut"]), hybrid)
                                    
        # Format with Misaka
        formatted_result = misaka.html(llm_result["completion"])
        return render_template('index.html', llm_result=formatted_result, form=form)
    
    # Spit out the template
    return render_template('index.html', llm_result="", form=form)

# The paste in data, get facts, store facts
@app.route('/paste', methods=['GET', 'POST'])
@login_required
def pastetext():
    # Question form for the external brain
    form = PasteForm()

    # Send the text to the llm or the facts to mongo
    if form.is_submitted():
        # Get the form variables
        form_result = request.form.to_dict(flat=True)
        
        if form_result["submit"] == "Save Facts":
            # Parse and save the facts here!
            facts = clean_facts(form_result["fact_data"])
            store_facts(facts, session["user"], form_result["context"])
            return redirect(url_for('index'))
        else:
            paste_data = form_result["paste_data"]
            context =  form_result["context"]
            llm_prompt = FACT_PROMPT.format(context=context, paste_data=paste_data)
            llm_result = llm(llm_prompt, FACT_SYSTEM_MESSAGE).lstrip().rstrip()

             # Change the form to the fact review form to save facts
            form = SaveFactsForm(fact_data=llm_result, context=form_result["context"])
            return render_template('factreview.html', form=form)
    
    # Spit out the template
    return render_template('pastedata.html', form=form)

# The paste in data, get facts, store facts
@app.route('/manual', methods=['GET', 'POST'])
@login_required
def manual():
    form = SaveFactsForm()

    if form.is_submitted():
        # Get the form variables
        form_result = request.form.to_dict(flat=True)
        
        # Parse and save the facts here!
        facts = clean_facts(form_result["fact_data"])
        store_facts(facts, session["user"], form_result["context"])
        return redirect(url_for('index'))
        
    return render_template('factreview.html', form=form)

# Regenerate the chunks!
@app.route('/facts', methods=['GET', 'POST'])
@login_required
def facts():
    form = FactSearchForm()
    fact_count = count_facts()

    # For paginating the display
    page = request.args.get('page')
    if not page:
        page = 0
    else:
        page = int(page)

    # Make sure the page doesn't go too low or too high
    if page < 0:
        page = 0
    if page >= fact_count / DEFAULT_FACTS_PER_PAGE:
        page = int(fact_count / DEFAULT_FACTS_PER_PAGE)

    if form.is_submitted():
        form_result = request.form.to_dict(flat=True)
        facts = search_facts(form_result["query"])
    else:
        skip = page * DEFAULT_FACTS_PER_PAGE
        limit = DEFAULT_FACTS_PER_PAGE
        facts = get_facts(skip, limit)
    return render_template('facts.html', page=page, form=form, facts=facts, facts_count=fact_count)

# Regenerate the chunks!
@app.route('/chunks', methods=['GET', 'POST'])
@login_required
def regenchunks():
    # Question form for the external brain
    form = ChunksForm(fact_limit=DEFAULT_LIMIT)
    fact_count = count_facts()

    # Regen the chunks based on settings
    if form.is_submitted():
        # Get the form variables
        form_result = request.form.to_dict(flat=True)
        fact_limit = int(form_result["fact_limit"])
        which_facts = form_result["which_facts"]
        
        # Regenerate just the changed facts or all of them?
        if which_facts == "all":
            chunk_by_context(fact_limit, True)
        else:
            chunk_by_context(fact_limit)

        return redirect(url_for('index'))

    # Spit out the template
    return render_template('chunks.html', fact_count=fact_count, form=form)

# Search chunks
@app.route('/search', methods=['GET', 'POST'])
@login_required
def search():
    chunks = []
    form = VectorSearchForm(k=DEFAULT_CANDIDATES, score_cut=DEFAULT_SCORE_CUT)
    # Regen the chunks based on settings
    if form.is_submitted():
        form_result = request.form.to_dict(flat=True)
        # Search chunks using vector search
        prompt = form_result["question"]
        candidates = int(form_result["k"])
        score_cut = float(form_result["score_cut"])
        chunks = search_chunks(prompt, candidates, 10, score_cut)

    return render_template('search.html', chunks=chunks, form=form)


# This fact is wrong and will be corrected harshly
@app.route('/edit/<id>', methods=['GET', 'POST'])
@login_required
def fact_edit(id):
    
    # Pull up that fact
    facts_col = db["facts"]
    chunks_col = db["chunks"]
    fact_record = facts_col.find_one({'_id': ObjectId(id)})

    # Fact edit form
    form = FactEditForm()

    if form.is_submitted():
        form_result = request.form.to_dict(flat=True)
        form_result.pop('csrf_token')
        form_result.pop('save')
        # Indicate we need to re-chunk all the facts in this category because we edited them
        facts_col.update_many({'context': form_result["context"]}, {'$set': {'chunked': False}})
        # Remove the chunks with this context, they're no longer balanced
        chunks_col.delete_many({'context': form_result["context"]})
        facts_col.update_one({'_id': ObjectId(id)}, {'$set': form_result})
        return redirect(url_for('index'))
    else:        
        # Populate the edit form.
        form.context.data = fact_record["context"]
        form.fact.data = fact_record["fact"]
        return render_template('edit.html', form=form)

# This fact is bad, and it should feel bad
@app.route('/delete/<id>')
@login_required
def fact_delete(id):
    facts_col = db["facts"]
    chunks_col = db["chunks"]
    
    # Find the fact record, and set all facts with the same context to not chunked
    # and delete the existing chunks in the same category
    fact_record = facts_col.find_one({'_id': ObjectId(id)})
    facts_col.update_many({'context': fact_record["context"]}, {'$set': {'chunked': False}})
    chunks_col.delete_many({'context': fact_record["context"]})

    facts_col.delete_one({'_id': ObjectId(id)})
    return redirect(url_for('facts'))


# API route for chatting with bot
@app.route('/api/chat', methods=['GET'])
def api_query():
    # Get the API key from the URL and bail out if it's wrong
    api_key = request.args.get('api_key')
    if api_key != API_KEY:
        return jsonify({"error": "Invalid API Key"})
    
    # Get the prompt, validate it's something
    prompt = request.args.get('prompt')
    if not prompt:
        return jsonify({"error": "you must provide a prompt"})
    
    # Optional system message tunable
    system_message = request.args.get('system_message')
    if not system_message:
        system_message = DEFAULT_SYSTEM_MESSAGE
    
    # Augmented or raw model?
    augmented = request.args.get('augmented')
    if not augmented:
        augmented = True
    else:
        if augmented == "False" or "false":
            augmented = False
        else:
            augmented = True

    # Spicy or boring?
    temperature = request.args.get('temperature')
    if not temperature:
        temperature = DEFAULT_TEMPERATURE

    # Working hard or hardly working?
    candidates = request.args.get('candidates')
    if not candidates:
        candidates = DEFAULT_CANDIDATES

    # Lots of facts or just a few?
    limit = request.args.get('limit')
    if not limit:
        limit = DEFAULT_LIMIT

    # Highly discerning or sloppy?
    score_cut = request.args.get('score_cut')
    if not score_cut:
        score_cut = DEFAULT_SCORE_CUT

    # Ask for a good result
    response = chat(prompt, system_message, augmented, float(temperature), int(candidates), int(limit), float(score_cut))
    return jsonify(response)

# Login/logout routes that rely on the user being stored in session
@app.route('/login', methods=['GET', 'POST'])
def login():
    form = LoginForm()
    if form.validate_on_submit():
        if form.username.data in users:
            if form.password.data == users[form.username.data]:
                session["user"] = form.username.data
                return redirect(url_for('index'))
    return render_template('login.html', form=form)

# We finally have a link for this now!
@app.route('/logout')
def logout():
    session["user"] = None
    return redirect(url_for('login'))