diff --git a/api/get_blocks.py b/api/get_blocks.py new file mode 100644 index 0000000..54a78bb --- /dev/null +++ b/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 diff --git a/api/main.py b/api/main.py index 8640559..b8505dc 100644 --- a/api/main.py +++ b/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) diff --git a/api/requirements.txt b/api/requirements.txt index c14420e..a0f4e56 100644 --- a/api/requirements.txt +++ b/api/requirements.txt @@ -1,6 +1,7 @@ # ---- ---- 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 + # ---- ---- +openai==0.27.2 +numpy==1.24.2 +tenacity==8.2.2 +tiktoken diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx index 5eaac96..6d29b02 100644 --- a/web/src/pages/index.tsx +++ b/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([]); diff --git a/web/src/pages/semantic.tsx b/web/src/pages/semantic.tsx index 6449497..60c5dda 100644 --- a/web/src/pages/semantic.tsx +++ b/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 = () => {