Merge pull request #104 from StampyAI/tidy-up

Tidy up

api/src/stampy_chat/chat.py

@@ -1,12 +1,14 @@
-# ------------------------------- env, constants -------------------------------
-from dataclasses import asdict
-from typing import List, Dict, Callable
-import openai
-import re
-from sqlalchemy.orm import PropComparator
-import tiktoken
 import time
+import json
+import re
+import time
+from dataclasses import asdict
+from typing import List, Dict
 
+import openai
+import tiktoken
+
+from stampy_chat.env import COMPLETIONS_MODEL
 from stampy_chat.followups import multisearch_authored
 from stampy_chat.get_blocks import get_top_k_blocks, Block
 from stampy_chat import logging
@@ -15,11 +17,6 @@ from stampy_chat import logging
 logger = logging.getLogger(__name__)
 
 
-# OpenAI models
-EMBEDDING_MODEL = "text-embedding-ada-002"
-COMPLETIONS_MODEL = "gpt-3.5-turbo"
-# COMPLETIONS_MODEL = "gpt-4"
-
 STANDARD_K = 20 if COMPLETIONS_MODEL == 'gpt-4' else 10
 
 # parameters
@@ -139,9 +136,6 @@ def construct_prompt(query: str, mode: str, history: Prompt, context: List[Block
     return prompt
 
 # ------------------------------- completion code -------------------------------
-import time
-import json
-
 
 def check_openai_moderation(prompt: Prompt, query: str):
     prompt_string = '\n\n'.join([message["content"] for message in prompt])

api/src/stampy_chat/env.py

@@ -18,6 +18,10 @@ DISCORD_LOGGING_URL = os.environ.get('LOGGING_URL')
 OPENAI_API_KEY   = os.environ.get('OPENAI_API_KEY')
 openai.api_key = OPENAI_API_KEY # non-optional
 
+### Models ###
+EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-ada-002")
+COMPLETIONS_MODEL = os.environ.get("COMPLETIONS_MODEL", "gpt-3.5-turbo")
+
 ### Pinecone ###
 PINECONE_API_KEY     = os.environ.get('PINECONE_API_KEY')
 PINECONE_ENVIRONMENT = os.environ.get('PINECONE_ENVIRONMENT', "us-east1-gcp")
@@ -41,3 +45,6 @@ host = os.environ.get("CHAT_DB_HOST", "127.0.0.1")
 port = os.environ.get("CHAT_DB_PORT", "3306")
 db_name = os.environ.get("CHAT_DB_NAME", "stampy_chat")
 DB_CONNECTION_URI = f"mysql+mysqlconnector://{user}:{password}@{host}:{port}/{db_name}"
+
+### Local testing helpers ###
+REMOTE_CHAT_INSTANCE = os.environ.get("REMOTE_CHAT_INSTANCE", "https://chat.stampy.ai:8443")

api/src/stampy_chat/get_blocks.py

@@ -7,16 +7,12 @@ import regex as re
 import requests
 import time
 from typing import List, Tuple
-from stampy_chat.env import PINECONE_NAMESPACE
+from stampy_chat.env import PINECONE_NAMESPACE, REMOTE_CHAT_INSTANCE, EMBEDDING_MODEL
 from stampy_chat import logging
 
 
 logger = logging.getLogger(__name__)
 
-# ---------------------------------- constants ---------------------------------
-
-EMBEDDING_MODEL = "text-embedding-ada-002"
-
 # ------------------------------------ types -----------------------------------
 
 @dataclasses.dataclass
@@ -66,7 +62,7 @@ def get_top_k_blocks(index, user_query: str, k: int) -> List[Block]:
 
         logger.info('Pinecone index not found, performing semantic search on chat.stampy.ai endpoint.')
         response = requests.post(
-            "https://chat.stampy.ai:8443/semantic",
+            REMOTE_CHAT_INSTANCE,
             json = {
                 "query": user_query,
                 "k": k
@@ -93,8 +89,9 @@ def get_top_k_blocks(index, user_query: str, k: int) -> List[Block]:
     )
     blocks = []
     for match in query_response['matches']:
+        metadata = match['metadata']
 
-        date = match['metadata']['date_published']
+        date = metadata.get('date_published') or metadata.get('date')
 
         if isinstance(date, datetime.date):
             date = date.isoformat()
@@ -103,18 +100,18 @@ def get_top_k_blocks(index, user_query: str, k: int) -> List[Block]:
         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')]
+        authors = metadata.get('authors')
+        if not authors and metadata.get('author'):
+            authors = [metadata.get('author')]
 
         blocks.append(Block(
-            id = match['metadata']['hash_id'],
-            title = match['metadata']['title'],
+            id = metadata.get('hash_id'),
+            title = metadata['title'],
             authors = authors,
             date = date,
-            url = match['metadata']['url'],
-            tags = match['metadata'].get('tags'),
-            text = strip_block(match['metadata']['text'])
+            url = metadata['url'],
+            tags = metadata.get('tags'),
+            text = strip_block(metadata['text'])
         ))
 
     t2 = time.time()

api/src/stampy_chat/logging.py

@@ -6,6 +6,8 @@ from stampy_chat.db.session import ItemAdder
 from stampy_chat.db.models import Interaction
 
 
+MAX_MESSAGE_LEN = 2000 - 8
+
 class DiscordHandler(StreamHandler):
     def emit(self, record):
         # Ignore messages that come from non chat modules
@@ -22,8 +24,8 @@ class DiscordHandler(StreamHandler):
         if not DISCORD_LOGGING_URL:
             return
 
-        while len(message) > 2000 - 8:
-            m_section, message = message[:2000 - 8], message[2000 - 8:]
+        while len(message) > MAX_MESSAGE_LEN:
+            m_section, message = message[:MAX_MESSAGE_LEN], message[MAX_MESSAGE_LEN:]
             m_section = "```\n" + m_section + "\n```"
             DiscordWebhook(url=DISCORD_LOGGING_URL, content=m_section).execute()
         DiscordWebhook(url=DISCORD_LOGGING_URL, content="```\n" + message + "\n```").execute()