Pgvector port

.env.example

@@ -17,3 +17,10 @@ LOG_LEVEL=DEBUG
 STANDALONE=true
 SLACK_USER_ID=U01JZQZQZQZ
 USE_FALLBACK=false
+
+# POSTGRES
+POSTGRES_HOST=localhost
+POSTGRES_PORT=5432
+POSTGRES_DATABASE=db-name
+POSTGRES_USER=postgres
+POSTGRES_PASSWORD=super-secret

README.md

@@ -33,6 +33,7 @@ https://haly.ai
 1. You must have a Slack organization where either you or an administrator can approve a new application.
 2. The ability to git clone a repo and run commands in either a Windows, Mac, or Linux terminal.
 3. Install [python](https://www.python.org/downloads/) and [pip](https://pip.pypa.io/en/stable/installation/).
+4. Finally, Install [PostgreSQL](https://www.postgresql.org/download/) and Postgres vector extention using [pgvector](https://github.com/pgvector/pgvector#installation). Additional [Install Methods](https://github.com/pgvector/pgvector#installation) for pgvector.
 
 ### Create your Slack bot:
 1. Go to https://api.slack.com/apps and hit the "Create New App" green button. Select "From an app manifest" option.
@@ -132,6 +133,14 @@ API_BASE_URL=not-needed-for-standalone
 LOG_LEVEL=DEBUG
 STANDALONE=true
 SLACK_USER_ID=U01JZQZQZQZ # Put a your workspace admin user ID if you know it
+
+# POSTGRES
+POSTGRES_HOST=localhost
+POSTGRES_PORT=5432
+POSTGRES_DATABASE=db-name
+POSTGRES_USER=postgres
+POSTGRES_PASSWORD=super-secret
+
 ```
 - Update SLACK_BOT_TOKEN (OAuth token), SLACK_SIGNING_SECRET, OPENAI_API_KEY ([Click here to learn how to get an API key from OpenAI](https://www.maisieai.com/help/how-to-get-an-openai-api-key-for-chatgpt)), and SLACK_USER_ID ([Click here how to get your Slack user ID](https://www.workast.com/help/article/how-to-find-a-slack-user-id/))
 - Have venv installed `python3 -m pip install virtualenv`

requirements.txt

@@ -3,4 +3,6 @@ python-dotenv==1.0.0
 slack-bolt==1.18.0
 openai==0.28
 gunicorn==20.1.0
-replicate==0.18.1
+replicate==0.18.1
+psycopg[binary]==3.1.14
+psycopg2-binary==2.9.9

src/semantic_search/semantic_search/config.py

@@ -54,3 +54,18 @@ def get_api_shared_secret() -> str:
 
 def is_standalone() -> bool:
     return os.environ.get('STANDALONE') == 'true'
+
+def get_postgres_host()-> str :
+    return os.environ.get('POSTGRES_HOST')
+
+def get_postgres_port()-> str :
+    return os.environ.get('POSTGRES_PORT')
+
+def get_postgres_database()-> str :
+    return os.environ.get('POSTGRES_DATABASE')
+
+def get_postgres_user()-> str :
+    return os.environ.get('POSTGRES_USER')
+
+def get_postgres_password()-> str :
+    return os.environ.get('POSTGRES_PASSWORD')

src/semantic_search/semantic_search/external_services/slack_api.py

@@ -114,7 +114,7 @@ def slack_names_map(team_id):
 
 
 def load_previous_messages(team_id: str, channel_id: str, last_message_id: str, number: int):
-    (messages, ) = load_previous_messages_with_pointer(team_id, channel_id, last_message_id, number)
+    (messages, _, _) = load_previous_messages_with_pointer(team_id, channel_id, last_message_id, number)
     return messages[-number:]
 
 

src/semantic_search/semantic_search/external_services/vector_databases/pinecone.py

@@ -0,0 +1,42 @@
+import pinecone
+from .vector_database import VectorDatabase
+from ...config import get_pinecone_key, get_pinecone_environment, get_pinecone_index_name
+
+class Pinecone(VectorDatabase): 
+    def __init__(self):
+        pinecone.init(api_key=get_pinecone_key(), environment=get_pinecone_environment())
+
+    # overriding abstract method 
+    def insert(self, embeddings, chunk, namespace):
+        items = []
+
+        for i in range(len(chunk)):
+            items.append({
+                'id': chunk[i].id,
+                'values': embeddings[i],
+                'metadata': chunk[i].to_metadata()
+            })
+
+        self.get_pinecone_index().upsert(
+            vectors=items,
+            namespace=namespace
+        )
+
+    # overriding abstract method 
+    def delete(self, ids, namespace):
+        self.get_pinecone_index().delete(ids=ids, namespace=namespace)
+
+    # overriding abstract method 
+    def select(self, query_vector, namespace):
+        query_results = self.get_pinecone_index().query(
+            queries=[query_vector],
+            top_k=50,
+            namespace=namespace,
+            include_values=False,
+            includeMetadata=True
+        )
+        return query_results['results'][0]['matches']
+
+    def get_pinecone_index() -> 'pinecone.Index':
+        return pinecone.Index(get_pinecone_index_name())
+

src/semantic_search/semantic_search/external_services/vector_databases/postgres.py

@@ -0,0 +1,50 @@
+from pgvector.psycopg2 import register_vector
+import psycopg2
+from .vector_database import VectorDatabase
+import json
+import numpy as np
+
+from ...config import get_postgres_host, get_postgres_port, get_postgres_database, get_postgres_user, get_postgres_password
+
+class Postgres(VectorDatabase): 
+    def __init__(self):
+        try:
+            self.conn = psycopg2.connect(
+                host=get_postgres_host(),
+                database=get_postgres_database(),
+                user=get_postgres_user(),
+                password=get_postgres_password(),
+                port=get_postgres_port())
+
+            self.cur = self.conn.cursor()
+
+            self.cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
+            register_vector(self.cur)
+
+            self.cur.execute('CREATE TABLE IF NOT EXISTS embedding (id bigserial PRIMARY KEY, namespace text, chunk_id text, metadata text, values vector)')
+            self.conn.commit()
+
+        except (Exception, psycopg2.DatabaseError) as error:
+            print(error)
+
+    # overriding abstract method 
+    def insert(self, embeddings, chunk, namespace):
+        for i in range(len(chunk)):
+            metadata = json.dumps(chunk[i].to_metadata())
+            self.cur.execute('INSERT INTO embedding (namespace, chunk_id, metadata, values) VALUES (%s, %s, %s, %s)', (namespace, chunk[i].id, metadata, embeddings[i]))
+
+        self.conn.commit()
+
+    # overriding abstract method 
+    def delete(self, ids, namespace):
+        self.cur.execute('DELETE FROM embedding WHERE chunk_id IN (%s) AND namespace=%s',(ids, namespace))
+        self.conn.commit()
+
+    # overriding abstract method 
+    def select(self, query_vector):
+        self.cur.execute('SELECT * FROM embedding ORDER BY values <-> %s LIMIT 50', (np.array(query_vector),))
+        embeddings = self.cur.fetchall()
+        output = [dict(id=chunk_id, metadata=json.loads(metadata)) for id, namespace, chunk_id, metadata, values in embeddings]
+        return output
+
+

src/semantic_search/semantic_search/external_services/vector_databases/vector_database.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod 
+
+class VectorDatabase(ABC): 
+    @abstractmethod
+    def insert(self, embeddings, chunk, namespace):
+        pass
+
+    @abstractmethod
+    def delete(self, ids, namespace):
+        pass
+
+    @abstractmethod
+    def select(self, query_vector):
+        pass

src/semantic_search/semantic_search/external_services/vector_databases/vector_instance.py

@@ -0,0 +1,17 @@
+from ...config import get_postgres_host, get_pinecone_environment
+from .postgres import Postgres
+from .pinecone import Pinecone
+
+postgres_instance = None
+pinecone_instance = None
+
+def get_db_instance():
+    global postgres_instance, pinecone_instance
+    if get_postgres_host():
+        if postgres_instance is None:
+            postgres_instance = Postgres()
+        return postgres_instance
+    else:
+        if pinecone_instance is None:
+            pinecone_instance = Pinecone()
+        return pinecone_instance

src/semantic_search/semantic_search/handle_indexation_tasks.py

@@ -3,7 +3,6 @@
 import os
 from flask import Flask, request, jsonify
 from .config import get_google_tasks_service_account
-from .external_services.pinecone import get_pinecone_index
 from .google_tasks import queue_task
 from .load_messages import index_messages
 from .external_services.slack_api import load_previous_messages_with_pointer
@@ -43,7 +42,7 @@ def handle_task():
         [messages, next_last_message, start_from] = load_previous_messages_with_pointer(namespace, channel_id, last_message_id, BULK_SIZE)
         logging.info(f"Task: {task_id}, Iteration Number: {iteration_number}")
         logging.info(f"Task: {task_id}, Number of Actual Messages: {len(messages)}")
-        index_messages(channel_id, messages, start_from, get_pinecone_index(), namespace)
+        index_messages(channel_id, messages, start_from, namespace)
         if next_last_message is not None:
             queue_task({
                 'task_id': task_id,

src/semantic_search/semantic_search/load_messages.py

@@ -1,15 +1,13 @@
 import logging
 import copy
 from typing import List, Dict
-
 from .config import CONTEXT_LENGTH
-from .external_services.pinecone import get_pinecone_index
 from .external_services.openai import create_embeddings, summarize_thread_with_chat_gpt_3_5
 import datetime
 from .external_services.slack_api import fetch_thread_messages, fetch_channel_messages, is_thread, \
     is_actual_message, \
     slack_names_map, filter_messages, load_previous_messages, load_subsequent_messages
-
+from .external_services.vector_databases.vector_instance import get_db_instance
 
 class Embedding:
     def __init__(self, channel_id, id, text, ts, thread_ts=None, author_id=None):
@@ -118,12 +116,12 @@ def attach_header(embeddings: List[Embedding], header: Embedding) -> List[Embedd
     return part_with_header + [embedding.add_header(header) for embedding in part_without_header]
 
 
-def index_messages(channel_id, messages, start_from, pinecone_index, pinecone_namespace):
+def index_messages(channel_id, messages, start_from, namespace):
     total_messages = len(messages)
 
     logging.info("Replacing User IDs with User Names in the messages")
     embeddings = generate_embeddings(channel_id, messages)
-    embeddings = replace_ids_with_names(embeddings, team_id=pinecone_namespace)
+    embeddings = replace_ids_with_names(embeddings, team_id=namespace)
     embeddings = enrich_with_datetime(embeddings)
     embeddings_without_context = embeddings
     embeddings = enrich_with_adjacent_messages(embeddings)
@@ -135,9 +133,9 @@ def index_messages(channel_id, messages, start_from, pinecone_index, pinecone_na
         if is_thread(message):
             logging.info(
                 f"{counter + 1}/{total_messages} Appending thread messages for {message['ts']} : {message['thread_ts']}")
-            thread_messages = filter_messages(fetch_thread_messages(pinecone_namespace, channel_id, message["thread_ts"]))
+            thread_messages = filter_messages(fetch_thread_messages(namespace, channel_id, message["thread_ts"]))
             thread_embeddings = generate_embeddings(channel_id, thread_messages)
-            thread_embeddings = replace_ids_with_names(thread_embeddings, team_id=pinecone_namespace)
+            thread_embeddings = replace_ids_with_names(thread_embeddings, team_id=namespace)
             thread_embeddings = enrich_with_datetime(thread_embeddings)
             thread_header = thread_embeddings[0]
             raw_messages_for_summary = list(map(lambda e: e.text, thread_embeddings))
@@ -164,54 +162,42 @@ def index_messages(channel_id, messages, start_from, pinecone_index, pinecone_na
     messages_for_embedding = list(filter(lambda emb_t: len(emb_t.text) != 0, messages_for_embedding))
     logging.info(f"Removed empty messages, {str(len(messages_for_embedding))} messages left")
 
-    insert_pinecone_embeddings(messages_for_embedding, pinecone_index, pinecone_namespace)
+    insert_db_embeddings(messages_for_embedding, namespace)
 
 
-def index_whole_channel(pinecone_namespace, channel_id):
+def index_whole_channel(namespace, channel_id):
     logging.info(f"Fetching all messages from {channel_id} channel")
-    messages = list(reversed(fetch_channel_messages(pinecone_namespace, channel_id)))
+    messages = list(reversed(fetch_channel_messages(namespace, channel_id)))
     logging.info(f"Loaded {str(len(messages))} messages")
 
     messages = filter_messages(messages)
     total_messages = len(messages)
     logging.info(f"Filtering out service messages, left {str(total_messages)} messages")
 
-    index_messages(channel_id, messages, 0, get_pinecone_index(), pinecone_namespace)
+    index_messages(channel_id, messages, 0, namespace)
 
 
-def insert_pinecone_embeddings(messages_for_embedding: List[Embedding], pinecone_index, pinecone_namespace):
+def insert_db_embeddings(messages_for_embedding: List[Embedding], namespace):
     logging.info("Starting embeddings creation for the generated messages")
     chunk_size = 30  # for OpenAI
     embedding_chunks = [messages_for_embedding[i:i + chunk_size] for i in
                         range(0, len(messages_for_embedding), chunk_size)]
     counter = 0
     for chunk in embedding_chunks:
-        logging.info(f"Inserting a chunk of Pinecone embeddings: [{counter} - {counter + len(chunk) - 1}]")
+        logging.info(f"Inserting a chunk of Postgre embeddings: [{counter} - {counter + len(chunk) - 1}]")
         counter += len(chunk)
         try:
             embeddings = create_embeddings([embedding_message.text for embedding_message in chunk])
-            items = []
-
-            for i in range(len(chunk)):
-                items.append({
-                    'id': chunk[i].id,
-                    'values': embeddings[i],
-                    'metadata': chunk[i].to_metadata()
-                })
-
-            pinecone_index.upsert(
-                vectors=items,
-                namespace=pinecone_namespace
-            )
+            get_db_instance().insert(embeddings, chunk, namespace)
         except:
             logging.exception("Couldn't insert embeddings")
 
 
-def delete_pinecone_embedding(embeddings: list[Embedding], pinecone_index, pinecone_namespace):
+def delete_db_embedding(embeddings: List[Embedding], namespace):
     ids = list(map(lambda emb: emb.id, embeddings))
     logging.info(f"Deleting embeddings for {str(ids)}")
-    pinecone_index.delete(ids=ids, namespace=pinecone_namespace)
-
+    ids_to_delete_str = ", ".join(map(str, ids))
+    get_db_instance().delete(ids_to_delete_str, namespace)
 
 def handle_message_update_and_reindex(body):
     event = body['event']
@@ -224,15 +210,15 @@ def handle_message_update_and_reindex(body):
         if not is_actual_message(message):
             return
         embedding = slack_message_to_embedding(channel_id, message)
-        delete_pinecone_embedding([embedding], get_pinecone_index(), team_id)
+        delete_db_embedding([embedding], team_id)
         if message.get('thread_ts') is not None:
             # just reindex the whole thread
-            index_messages(channel_id, load_previous_messages(team_id, channel_id, message.get('thread_ts'), 1), 0, get_pinecone_index(), team_id)
+            index_messages(channel_id, load_previous_messages(team_id, channel_id, message.get('thread_ts'), 1), 0, team_id)
             return
         message_ts = message['ts']
         messages_for_reindex = load_previous_messages(team_id, channel_id, message_ts, CONTEXT_LENGTH - 1) + load_subsequent_messages(team_id, channel_id, message_ts, CONTEXT_LENGTH - 1)
         # reindex surrounding messages
-        index_messages(channel_id, messages_for_reindex, CONTEXT_LENGTH - 1, get_pinecone_index(), team_id)
+        index_messages(channel_id, messages_for_reindex, CONTEXT_LENGTH - 1, team_id)
         return
     if 'subtype' in event and event['subtype'] == 'message_changed':
         # processing a message update
@@ -242,12 +228,12 @@ def handle_message_update_and_reindex(body):
             return
         if message.get('thread_ts') is not None:
             # just reindex the whole thread
-            index_messages(channel_id, load_previous_messages(team_id, channel_id, message.get('thread_ts'), 1), 0, get_pinecone_index(), team_id)
+            index_messages(channel_id, load_previous_messages(team_id, channel_id, message.get('thread_ts'), 1), 0, team_id)
             return
         message_ts = message['ts']
         messages_for_reindex = load_previous_messages(team_id, channel_id, message_ts, CONTEXT_LENGTH) + load_subsequent_messages(team_id, channel_id, message_ts, CONTEXT_LENGTH)[1:]
         # reindex surrounding messages
-        index_messages(channel_id, messages_for_reindex, CONTEXT_LENGTH - 1, get_pinecone_index(), team_id)
+        index_messages(channel_id, messages_for_reindex, CONTEXT_LENGTH - 1, team_id)
         return
     if 'subtype' not in event:
         message = event
@@ -260,9 +246,8 @@ def handle_message_update_and_reindex(body):
     if not is_actual_message(message):
         return
     embeddings = generate_embedding_for_message(team_id, channel_id, message_id, thread_ts)
-    insert_pinecone_embeddings(
+    insert_db_embeddings(
         embeddings,
-        get_pinecone_index(),
         team_id
     )