diff --git a/api/src/stampy_chat/chat.py b/api/src/stampy_chat/chat.py index b837d1b..024c0f3 100644 --- a/api/src/stampy_chat/chat.py +++ b/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"} diff --git a/api/src/stampy_chat/env.py b/api/src/stampy_chat/env.py index b4c66b1..9ac6f98 100644 --- a/api/src/stampy_chat/env.py +++ b/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() diff --git a/api/src/stampy_chat/get_blocks.py b/api/src/stampy_chat/get_blocks.py index 131f140..3122db6 100644 --- a/api/src/stampy_chat/get_blocks.py +++ b/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] diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx index 79097b3..df734e8 100644 --- a/web/src/pages/index.tsx +++ b/web/src/pages/index.tsx @@ -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; diff --git a/web/src/pages/semantic.tsx b/web/src/pages/semantic.tsx index 58ff80c..23c6f1c 100644 --- a/web/src/pages/semantic.tsx +++ b/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 (
{entry.author} - {entry.date}
+{entry.authors.join(', ')} - {entry.date}