diff --git a/.gitignore b/.gitignore
index 4f9ddb1..df62ec6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -137,3 +137,8 @@ src/tmp.py
 
 .vercel/
 api/dataset.pkl
+temp/
+
+api/dataset_big.pkl
+
+api/dataset_300.pkl
diff --git a/api/chat.py b/api/chat.py
index 2254223..fb5378e 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -10,12 +10,15 @@ import tiktoken
 # OpenAI models
 EMBEDDING_MODEL = "text-embedding-ada-002"
 COMPLETIONS_MODEL = "gpt-3.5-turbo"
+# COMPLETIONS_MODEL = "gpt-4"
 MODERATION_ENDPOINT = "https://api.openai.com/v1/moderations"
 
 # OpenAI parameters
 LEN_EMBEDDINGS = 1536
-MAX_TOKEN_LEN_PROMPT = 4095 # This may be 8191, unsure.
-TRUNCATE_CONTEXT = 2000
+MAX_TOKEN_LEN_PROMPT = 8191 if COMPLETIONS_MODEL == 'gpt-4' else 4095
+TRUNCATE_CONTEXT_LEN = 2300 if COMPLETIONS_MODEL == 'gpt-4' else 1500
+TRUNCATE_HISTORY_LEN = 500
+MAX_RESPONSE_LEN = 900
 
 # --------------------------------- prompt code --------------------------------
 
@@ -24,7 +27,8 @@ def limit_tokens(text: str, max_tokens: int, encoding_name: str = "cl100k_base")
     tokens = encoding.encode(text)[:max_tokens]
     return encoding.decode(tokens)
 
-def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Block]) -> List[Dict[str, str]]:
+def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Block], encoding_name: str = "cl100k_base"):
+
     # History takes the format: history=[
     #     {"role": "system", "content": "You are a helpful assistant."},
     #     {"role": "user", "content": "Who won the world series in 2020?"},
@@ -33,51 +37,68 @@ def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Bl
     #     {"role": "assistant", "content": "Los Angeles, California."}
     # ]
 
-    # Initialize prompt with system description
-    prompt = [{"role": "system", "content": "You are a helpful assistant knowledgeable about AI Alignment and Safety."}]
+    # Encoder to count tokens
+    enc = tiktoken.get_encoding(encoding_name)
+    total_tokens = 0
 
-    # Add previous dialogue
-    prompt.extend(history) 
+    prompt = []
 
-    instruction_prompt = \
+    system_prompt = "You are a helpful assistant knowledgeable about AI Alignment and Saftey."
+    total_tokens += len(enc.encode(system_prompt))
+
+    # Get past user queries
+    past_user_queries = "\nQ: ".join([message["content"] for message in history if message["role"] == "user"][-5:])
+    past_user_queries = f"My previous queries in our conversation have been:\n" + past_user_queries
+
+    # Instruction prompt
+    instruction_context_query_prompt = \
         "Please give a clear and coherent answer to my question (written after \"Q:\") " \
         "using the following sources. Each source is labeled with a letter. Feel free to " \
         "use the sources in any order, and try to use multiple sources in your answer."
 
-    prompt.append({"role": "user", "content": instruction_prompt})
-    
-    # Add context from top-k blocks
+    # Context from top-k blocks
     context_prompt = ""
+
     for i, block in enumerate(context):
         context_prompt += f"[{chr(ord('a') + i)}] {block.title} - {block.author} - {block.date}\n\n{block.text}\n\n\n"
-
     context_prompt = context_prompt[:-2] # trim last two newlines
+    context_prompt = limit_tokens(context_prompt, TRUNCATE_CONTEXT_LEN)  # truncate the context_prompt to max TRUNCATE_CONTEXT tokens
+    context_prompt += "\n" if (context_prompt[-1] != "\n") else ""
 
-    context_prompt = limit_tokens(context_prompt, TRUNCATE_CONTEXT) # truncate to about 2k tokens
+    # Question prompt
+    question_prompt = f"In your answer, please cite any claims you make back to each source " \
+                    f"using the format: [a], [b], etc. If you use multiple sources to make a claim " \
+                    f"cite all of them. For example: \"AGI is concerning [c, d, e].\"\n\nQ: " + query
 
-    prompt.append({"role": "user", "content": f"{context_prompt}"})
-    
-    # Add user query
-    question_prompt = "In your answer, please cite any claims you make back to each source " \
-                      "using the format: [a], [b], etc. If you use multiple sources to make a claim " \
-                      "cite all of them. For example: \"AGI is concerning [c, d, e].\""
+    instruction_context_query_prompt = f"{instruction_context_query_prompt}\n\n{context_prompt}\n\n{question_prompt}"
 
-    question_prompt += "\n\n\nQ: " + query
-
-    prompt.append({"role": "user", "content": question_prompt})
-    
-    return prompt
+    total_tokens += len(enc.encode(past_user_queries))
+    total_tokens += len(enc.encode(history[-2]["content"])) if (len(history) >= 2) else 0
+    total_tokens += len(enc.encode(history[-1]["content"])) if (len(history) >= 1) else 0
+    total_tokens += len(enc.encode(instruction_context_query_prompt))
 
+    # If the prompt is too long, truncate the last answer
+    if total_tokens > MAX_TOKEN_LEN_PROMPT - TRUNCATE_HISTORY_LEN:
+        tokens_left = MAX_TOKEN_LEN_PROMPT - total_tokens
+        print(f"WARNING: Prompt is too long! Prompt length: {total_tokens} tokens")
+        last_assistant_reply_trunctated = limit_tokens(prompt[-1]["content"], tokens_left)
+        prompt[-1]["content"] = f"{last_assistant_reply_trunctated}"
 
+    prompt.append({"role": "system", "content": system_prompt})
+    prompt.append({"role": "user", "content": past_user_queries})
+    prompt.extend(history[-2:])
+    prompt.append({"role": "user", "content": instruction_context_query_prompt})
 
+    return prompt, MAX_TOKEN_LEN_PROMPT - (total_tokens + 50) # add 50 tokens for safety
 
 # ------------------------------------------------------------------------------
-
-def normal_completion(prompt: List[Dict[str, str]]) -> str:
+        
+def normal_completion(prompt: List[Dict[str, str]], max_tokens_completion: int) -> str:
     try:
         return openai.ChatCompletion.create(
             model=COMPLETIONS_MODEL,
-            messages=prompt
+            messages=prompt,
+            max_tokens=max_tokens_completion
         )["choices"][0]["message"]["content"]
     except Exception as e:
         print(e)
@@ -90,11 +111,14 @@ def talk_to_robot(dataset_dict, query: str, history: List[Dict[str, str]] = [],
     top_k_blocks: List[Block] = get_top_k_blocks(dataset_dict, query, k)
     
     # 2. Generate a prompt for the ChatCompletions API
-    prompt: List[Dict[str, str]] = construct_prompt(query, history, top_k_blocks)
+    prompt, max_tokens_completion = construct_prompt(query, history, top_k_blocks)
+
+    # print(" ------------------------------ prompt: -----------------------------")
+    # for message in prompt:
+    #     print(f"{message['role']}: {message['content']}\n\n")
 
     # if we were to error out, return something like this
     # return (False, "Example error message", None)
     
     # 3. Answer the user query
-    return (True, normal_completion(prompt), top_k_blocks)
-    
+    return (True, normal_completion(prompt, max_tokens_completion), top_k_blocks)
diff --git a/api/get_blocks.py b/api/get_blocks.py
index 6f30cb1..d4af74e 100644
--- a/api/get_blocks.py
+++ b/api/get_blocks.py
@@ -47,13 +47,25 @@ def get_embedding(text: str) -> np.ndarray:
 # Get the k blocks most semantically similar to the query.
 def get_top_k_blocks(data, user_query: str, k: int = 10) -> List[Block]:
 
+    # print time
+    t = time.time()
+
     # Get the embedding for the query.
     query_embedding = get_embedding(user_query)
-    
+
+    t1 = time.time()
+    print("Time to get embedding: ", t1 - t)
+
     similarity_scores = np.dot(data["embeddings"], query_embedding) # big fat calculation
+
+    t2 = time.time()
+    print("Time to get similarity scores: ", t2 - t1)
     
     top_k_block_indices = list(reversed(np.argpartition(similarity_scores, -k)[-k:])) # Get the top k indices of the blocks
 
+    t3 = time.time()
+    print("Time to get top k indices: ", t3 - t2)
+
     top_k_metadata_indexes = [data["embeddings_metadata_index"][i] for i in top_k_block_indices]
     top_k_texts = [strip_block(data["embedding_strings"][i]) for i in top_k_block_indices]
     top_k_metadata = [data["metadata"][i] for i in top_k_metadata_indexes]
@@ -76,6 +88,8 @@ def get_top_k_blocks(data, user_query: str, k: int = 10) -> List[Block]:
     for key, group in itertools.groupby(blocks_plus_old_index, key=key):
         group = list(group)
         if len(group) == 0: continue
+        
+        group = group[:3] # limit to a max of 3 blocks from any one source
 
         text = "\n\n\n.....\n\n\n".join([block[0].text for block in group])
 
diff --git a/web/src/header.tsx b/web/src/header.tsx
index 059a96e..46320a7 100644
--- a/web/src/header.tsx
+++ b/web/src/header.tsx
@@ -16,7 +16,7 @@ const Header: React.FC<{page: "index" | "semantic"}> = ({page}) => {
         
     return (<>
                 <div className="flex my-4">
-                    <h1 className="flex-1 my-0">Alignment Search</h1>
+                    <h1 className="flex-1 my-0">AlignmentSearch</h1>
                     {sidebar}
                 </div>
                 <p>
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index 5a76ec2..dee86c5 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -87,7 +87,7 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
 
     // system reply
     return (
-        <div className="my-3">
+        <div className="mt-3 mb-8">
             {   // split into paragraphs
                 entry.display_content.split("\n").map(paragraph => ( <p> {
                     paragraph.split(in_text_citation_regex).map((text, i) => {
diff --git a/web/src/styles/globals.css b/web/src/styles/globals.css
index 2e88fa2..f304db2 100644
--- a/web/src/styles/globals.css
+++ b/web/src/styles/globals.css
@@ -15,6 +15,8 @@ main {
   max-width: 800px;
   margin: 0 auto;
   padding: 0 2rem;
+  margin-top: 4rem;
+  margin-bottom: 4rem;
 }
 
 a {