mirror of
https://github.com/wassname/stampy-chat.git
synced 2026-09-11 12:50:34 +08:00
Use new pinecone embeddings
This commit is contained in:
@@ -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"}
|
||||
|
||||
+18
-14
@@ -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()
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user