Merge pull request #93 from StampyAI/new-pinecone

Use new pinecone embeddings

api/src/stampy_chat/chat.py

@@ -71,7 +71,7 @@ def construct_prompt(query: str, mode: str, history: List[Dict[str, str]], conte
 
     # Context from top-k blocks
     for i, block in enumerate(context):
-        block_str = f"[{chr(ord('a') + i)}] {block.title} - {block.author} - {block.date}\n{block.text}\n\n"
+        block_str = f"[{chr(ord('a') + i)}] {block.title} - {','.join(block.authors)} - {block.date}\n{block.text}\n\n"
         block_tc = len(ENCODER.encode(block_str))
 
         if token_count + block_tc > int(NUM_TOKENS * CONTEXT_FRACTION):
@@ -148,7 +148,7 @@ def talk_to_robot_internal(index, query: str, mode: str, history: List[Dict[str,
         yield {"state": "loading", "phase": "semantic"}
         top_k_blocks = get_top_k_blocks(index, query, k)
 
-        yield {"state": "loading", "phase": "semantic", 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in top_k_blocks]}
+        yield {"state": "loading", "phase": "semantic", 'citations': [{'title': block.title, 'author': block.authors, 'date': block.date, 'url': block.url} for block in top_k_blocks]}
 
         # 2. Generate a prompt
         yield {"state": "loading", "phase": "prompt"}

api/src/stampy_chat/env.py

@@ -1,7 +1,6 @@
 import os
 import openai
 import pinecone
-from discord_webhook import DiscordWebhook
 
 if os.path.exists('.env'):
     from dotenv import load_dotenv
@@ -9,33 +8,38 @@ if os.path.exists('.env'):
 else:
     print("'api/.env' not found. Rename the 'api/.env.example' file and fill in values.")
 
-
-OPENAI_API_KEY   = os.environ.get('OPENAI_API_KEY')
-PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
-LOGGING_URL      = os.environ.get('LOGGING_URL')
-PINECONE_INDEX   = None
-
 FLASK_PORT = int(os.environ.get('FLASK_PORT', '3001'))
 
+LOG_LEVEL           = os.environ.get("LOG_LEVEL", "WARNING").upper()
+DISCORD_LOG_LEVEL   = os.environ.get("DISCORD_LOG_LEVEL", "WARNING").upper()
+DISCORD_LOGGING_URL = os.environ.get('LOGGING_URL')
+
+### OpenAI ###
+OPENAI_API_KEY   = os.environ.get('OPENAI_API_KEY')
 openai.api_key = OPENAI_API_KEY # non-optional
 
+### Pinecone ###
+PINECONE_API_KEY     = os.environ.get('PINECONE_API_KEY')
+PINECONE_ENVIRONMENT = os.environ.get('PINECONE_ENVIRONMENT', "us-east1-gcp")
+PINECONE_INDEX_NAME  = os.environ.get("PINECONE_INDEX_NAME", "alignment-search")
+PINECONE_INDEX       = None
+PINECONE_NAMESPACE   = os.environ.get("PINECONE_NAMESPACE", "alignment-search")  # "normal" or "finetuned" for the new index, "alignment-search" for the old one
 # Only init pinecone if we have an env value for it.
-if PINECONE_API_KEY is not None and PINECONE_API_KEY != "":
-
+if PINECONE_API_KEY:
     pinecone.init(
         api_key = PINECONE_API_KEY,
-        environment = "us-east1-gcp",
+        environment = PINECONE_ENVIRONMENT,
     )
 
-    PINECONE_INDEX = pinecone.Index(index_name="alignment-search")
+    PINECONE_INDEX = pinecone.Index(index_name=PINECONE_INDEX_NAME)
 
 # log something only if the logging url is set
 def log(*args, end="\n"):
     message = " ".join([str(arg) for arg in args]) + end
     # print(message)
-    if LOGGING_URL is not None and LOGGING_URL != "":
+    if DISCORD_LOGGING_URL is not None and DISCORD_LOGGING_URL != "":
         while len(message) > 2000 - 8:
             m_section, message = message[:2000 - 8], message[2000 - 8:]
             m_section = "```\n" + m_section + "\n```"
-            DiscordWebhook(url=LOGGING_URL, content=m_section).execute()
-        DiscordWebhook(url=LOGGING_URL, content="```\n" + message + "\n```").execute()
+            DiscordWebhook(url=DISCORD_LOGGING_URL, content=m_section).execute()
+        DiscordWebhook(url=DISCORD_LOGGING_URL, content="```\n" + message + "\n```").execute()

api/src/stampy_chat/get_blocks.py

@@ -7,6 +7,7 @@ import openai
 import regex as re
 import requests
 import time
+from stampy_chat.env import PINECONE_NAMESPACE
 
 # ---------------------------------- constants ---------------------------------
 
@@ -16,13 +17,14 @@ EMBEDDING_MODEL = "text-embedding-ada-002"
 
 @dataclasses.dataclass
 class Block:
+    id: str
     title: str
-    author: str
+    authors: List[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.
@@ -79,7 +81,7 @@ def get_top_k_blocks(index, user_query: str, k: int) -> List[Block]:
     print(f'Time to get embedding: {t1-t:.2f}s')
 
     query_response = index.query(
-        namespace="alignment-search",  # ugly, sorry
+        namespace=PINECONE_NAMESPACE,
         top_k=k,
         include_values=False,
         include_metadata=True,
@@ -88,28 +90,38 @@ def get_top_k_blocks(index, user_query: str, k: int) -> List[Block]:
     blocks = []
     for match in query_response['matches']:
 
-        date = match['metadata']['date']
+        date = match['metadata']['date_published']
 
-        if type(date) == datetime.date: date = date.strftime("%Y-%m-%d") # iso8601
+        if isinstance(date, datetime.date):
+            date = date.isoformat()
+        elif isinstance(date, datetime.datetime):
+            date = date.date().isoformat()
+        elif isinstance(date, float):
+            date = datetime.datetime.fromtimestamp(date).date().isoformat()
+
+        authors = match['metadata'].get('authors')
+        if not authors and match['metadata'].get('author'):
+            authors = [match['metadata'].get('author')]
 
         blocks.append(Block(
+            id = match['id'],
             title = match['metadata']['title'],
-            author = match['metadata']['author'],
+            authors = authors,
             date = date,
             url = match['metadata']['url'],
-            tags = match['metadata']['tags'],
+            tags = match['metadata'].get('tags'),
             text = strip_block(match['metadata']['text'])
         ))
 
     t2 = time.time()
 
     print(f'Time to get top-k blocks: {t2-t1:.2f}s')
-    
+
     # for all blocks that are "the same" (same title, author, date, url, tags),
     # combine their text with "....." 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 "")
+    key = lambda bi: (bi[0].id, bi[0].title or "", bi[0].authors 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)
@@ -119,14 +131,14 @@ def get_top_k_blocks(index, user_query: str, k: int) -> 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".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.append((Block(*key, text), min_index))
 
     unified_blocks.sort(key=lambda bi: bi[1])
     return [block for block, _ in unified_blocks]

web/src/pages/index.tsx

@@ -13,7 +13,7 @@ import { API_URL } from "../settings"
 
 type Citation = {
   title: string;
-  author: string;
+  authors: string[];
   date: string;
   url: string;
 }
@@ -68,8 +68,8 @@ const ShowCitation: React.FC<{citation: Citation, i: number}> = ({citation, i})
 
   var c_str = citation.title;
 
-  if (citation.author && citation.author !== "")
-    c_str += " - " + citation.author;
+  if (citation.authors && citation.authors.length > 0)
+    c_str += " - " + citation.authors.join(', ');
   if (citation.date && citation.date !== "")
     c_str += " - " + citation.date;
 
@@ -528,7 +528,7 @@ const Home: NextPage = () => {
         </div>
 
 
-        <h2 className="bg-red-100 text-red-800"><b>WARNING</b>: This is a very <b>early prototype</b> using data through June 2022. <Link href="http://bit.ly/stampy-chat-issues" target="_blank">Feedback</Link> welcomed.</h2>
+        <h2 className="bg-red-100 text-red-800"><b>WARNING</b>: This is a very <b>early prototype</b>. <Link href="http://bit.ly/stampy-chat-issues" target="_blank">Feedback</Link> welcomed.</h2>
 
 
         <ul>

web/src/pages/semantic.tsx

@@ -65,7 +65,7 @@ const Semantic: NextPage = () => {
 
 type SemanticEntry = {
   title: string;
-  author: string;
+  authors: string[];
   date: string;
   url: string;
   tags: string;
@@ -77,10 +77,10 @@ const ShowSemanticEntry: React.FC<{entry: SemanticEntry}> = ({entry}) => {
   return (
     <div className="my-3">
 
-      {/* horizontally split first row, title on left, author on right */}
+      {/* horizontally split first row, title on left, authors on right */}
       <div className="flex">
         <h3 className="text-xl flex-1">{entry.title}</h3>
-        <p className="flex-1 text-right my-0">{entry.author} - {entry.date}</p>
+        <p className="flex-1 text-right my-0">{entry.authors.join(', ')} - {entry.date}</p>
       </div>
       { entry.text.split("\n").map((paragraph, i) => {
         const p = paragraph.trim();