semantic search working in flask

api/get_blocks.py

@@ -0,0 +1,106 @@
+from typing import List, Tuple
+import dataclasses
+import itertools
+import pickle
+import numpy as np
+import openai
+import regex as re
+import time
+
+# ---------------------------------- constants ---------------------------------
+
+EMBEDDING_MODEL = "text-embedding-ada-002"
+COMPLETIONS_MODEL = "gpt-3.5-turbo"
+
+import pathlib
+project_path = pathlib.Path(__file__).parent
+PATH_TO_DATASET_DICT = project_path / "dataset_dict.pkl"
+
+with open(PATH_TO_DATASET_DICT, 'rb') as f:
+    data = pickle.load(f)
+
+# ------------------------------------ types -----------------------------------
+
+@dataclasses.dataclass
+class Block:
+    title: str
+    author: str
+    date: str
+    url: str
+    tags: str
+    text: str
+    
+# ------------------------------------------------------------------------------
+
+# Get the embedding for a given text. The function will retry with exponential backoff if the API rate limit is reached, up to 4 times.
+def get_embedding(text: str) -> np.ndarray:
+
+    max_retries = 4
+    max_wait_time = 10
+    attempt = 0
+
+    while True:
+        try:
+            result = openai.Embedding.create(model=EMBEDDING_MODEL, input=text)
+            return result["data"][0]["embedding"]
+
+        except openai.error.RateLimitError as e:
+
+            attempt += 1
+
+            if attempt > max_retries: raise e
+
+            time.sleep(min(max_wait_time, 2 ** attempt))
+
+# Get the k blocks most semantically similar to the query.
+def get_top_k_blocks(user_query: str, k: int = 10) -> List[Block]:
+
+    # Get the embedding for the query.
+    query_embedding = get_embedding(user_query)
+    
+    similarity_scores = np.dot(data["embeddings"], query_embedding) # big fat calculation
+    
+    top_k_block_indices = list(reversed(np.argpartition(similarity_scores, -k)[-k:])) # Get the top k indices of the blocks
+
+    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]
+
+    # Combine the top k texts and metadata into a list of Block objects
+    top_k_metadata_and_text = [list(top_k_metadata[i]) + [top_k_texts[i]] for i in range(len(top_k_metadata))]
+    blocks = [Block(*block) for block in top_k_metadata_and_text]
+
+    # for all blocks that are "the same" (same title, author, date, url, tags),
+    # combine their text with "\n\n.....\n\n" in between. Return them in order such
+    # that the combined block has the minimum index of the blocks combined.
+
+    key = lambda bi: (bi[0].title or "", bi[0].author or "", bi[0].date or "", bi[0].url or "", bi[0].tags or "")
+
+    blocks_plus_old_index = [(block, i) for i, block in enumerate(blocks)]
+    blocks_plus_old_index.sort(key=key)
+
+    unified_blocks: List[Tuple[Block, int]] = []
+
+    for key, group in itertools.groupby(blocks_plus_old_index, key=key):
+        group = list(group)
+        if len(group) == 0: continue
+
+        text = "\n\n\n.....\n\n\n".join([block[0].text for block in group])
+
+        min_index = min([block[1] for block in group])
+
+        unified_blocks.append((Block(key[0], key[1], key[2], key[3], key[4], text), min_index))
+
+    unified_blocks.sort(key=lambda bi: bi[1])
+    return [block for block, _ in unified_blocks]
+
+
+# we add the title and authors inside the contents of the block, so that
+# searches for the title or author will be more likely to pull it up. This
+# strips it back out.
+def strip_block(text: str) -> str:
+    r = re.match(r"^\"(.*)\"\s*-\s*Title:.*$", text, re.DOTALL)
+    if not r:
+        print("Warning: couldn't strip block")
+        print(text)
+    return r.group(1) if r else text

api/main.py

@@ -1,13 +1,31 @@
-from flask import Flask, jsonify
+from flask import Flask, jsonify, request
+from flask_cors import CORS, cross_origin
+from get_blocks import get_top_k_blocks
+import dataclasses
 import os
+import openai
 
 app = Flask(__name__)
+cors = CORS(app)
+app.config['CORS_HEADERS'] = 'Content-Type'
+
+# -------------------------------- general setup -------------------------------
+
+OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
+openai.api_key = OPENAI_API_KEY
+
+# ------------------------------- semantic search ------------------------------
+
+@app.route('/semantic', methods=['POST'])
+@cross_origin()
+def semantic():
+    query = request.json['query']
+    return jsonify([dataclasses.asdict(block) for block in get_top_k_blocks(query)])
 
 
-@app.route('/')
-def index():
-    return jsonify({"general kenobi": "hello there"})
 
 
-if __name__ == '__main__':
-    app.run(debug=True, port=5000)
+
+
+
+if __name__ == '__main__': app.run(debug=True, port=3000)

api/requirements.txt

@@ -1,6 +1,7 @@
 # ---- <flask stuff> ----
 
 Flask==1.1.2
+
 click==7.1.2
 gunicorn==20.0.4
 itsdangerous==1.1.0
@@ -8,5 +9,11 @@ Jinja2==2.11.3
 MarkupSafe==1.1.1
 Werkzeug==1.0.1
 
+flask-cors
+
 # ---- </flask stuff> ----
 
+openai==0.27.2
+numpy==1.24.2
+tenacity==8.2.2
+tiktoken

web/src/pages/index.tsx

@@ -1,3 +1,5 @@
+const API_URL = process.env.NEXT_PUBLIC_API_URL || "http://127.0.0.1:3000";
+
 import Head from "next/head";
 import React from "react";
 import { type NextPage } from "next";
@@ -25,8 +27,6 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
     );
 };
 
-const API_URL = process.env.NEXT_PUBLIC_API_URL || "http://127.0.0.1:5000/";
-
 const Home: NextPage = () => {
 
     const [ entries, setEntries ] = useState<Entry[]>([]);

web/src/pages/semantic.tsx

@@ -1,3 +1,5 @@
+const API_URL = process.env.NEXT_PUBLIC_API_URL || "http://127.0.0.1:3000";
+
 import { type NextPage } from "next";
 import React from "react";
 import Head from "next/head";
@@ -18,9 +20,12 @@ const Semantic: NextPage = () => {
         setLoading(true);
         setQuery("");
 
-        const res = await fetch("/api/semantic_search", {
+        const res = await fetch(API_URL + "/semantic", {
             method: "POST",
-            headers: { "Content-Type": "application/json", },
+            headers: { "Content-Type": "application/json",
+                // allow cross-origin requests
+                "Access-Control-Allow-Origin": "*",
+            },
             body: JSON.stringify({query: query}),
         })
 
@@ -47,7 +52,7 @@ const Semantic: NextPage = () => {
                 <SearchBox search={semantic_search} />
                 <ul>
                     {results.map((entry, i) => (
-                        <li key={i}>
+                        <li key={"entry" + i}>
                             <ShowSemanticEntry entry={entry} />
                         </li>
                     ))}
@@ -84,8 +89,8 @@ const ShowSemanticEntry: React.FC<{entry: SemanticEntry}> = ({entry}) => {
             { entry.text.split("\n").map((paragraph, i) => {
                 const p = paragraph.trim();
                 if (p === "") return <></>;
-                if (p === ".....") return <hr key={i} />;
-                return <p className="text-sm" key={i}> {paragraph} </p>
+                if (p === ".....") return <hr key={"b" + i} />;
+                return <p className="text-sm" key={"p" + i}> {paragraph} </p>
               })
             }