Cleanup

alembic/versions/2024_04_12_1357-4ae5c037da0f_add_page_data_table.py

@@ -22,7 +22,7 @@ def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
     op.create_table('page_data',
     sa.Column('id', sa.Integer(), nullable=False),
-    sa.Column('page_id', sa.String(), nullable=True),
+    sa.Column('page_id', sa.BigInteger(), nullable=False),
     sa.Column('space_key', sa.String(), nullable=True),
     sa.Column('title', sa.String(), nullable=True),
     sa.Column('author', sa.String(), nullable=True),

api/endpoint.py

@@ -32,11 +32,11 @@ class FeedbackEvent(BaseModel):
 
 
 class EmbedRequest(BaseModel):
-    page_id: str
+    page_id: int
 
 
 class InteractionEmbedRequest(BaseModel):
-    interaction_id: str
+    interaction_id: int
 
 
 @processor.post("/api/v1/questions")
@@ -64,6 +64,7 @@ def create_embeds(EmbedRequest: EmbedRequest):
     """
     Endpoint to initiate the embedding generation and storage process in the background.
     :param EmbedRequest:
+    :return:
     """
     page_id = EmbedRequest.page_id
     thread = threading.Thread(target=vector.pages.generate_one_embedding_to_database, args=(page_id,))

api/request.py

@@ -2,10 +2,9 @@
 import logging
 
 
-def post_request(url, payload, headers=None, data_type=None):
+def post_request(url, payload, headers=None):
     """
     Post a request to a given URL.
-    :param data_type: Type of payload data.
     :param url: The URL to post to.
     :param payload: The payload to send.
     :param headers: The headers to send.
@@ -16,8 +15,8 @@ def post_request(url, payload, headers=None, data_type=None):
     try:
         response = requests.post(url, json=payload, headers=headers)
         response.raise_for_status()
-        logging.info(f"INFO: {data_type} request submitted to {url}")
+        logging.info(f"INFO: request submitted to {url}")
     except requests.exceptions.HTTPError as e:
-        logging.error(f"ERROR: An HTTP error with {data_type} request occurred: {e}")
+        logging.error(f"ERROR: An HTTP error occurred: {e}")
     except Exception as e:
-        logging.error(f"ERROR: An error with {data_type} request occurred: {e}")
+        logging.error(f"ERROR: An error occurred: {e}")

confluence/importer.py

@@ -2,7 +2,7 @@
 import logging
 
 from database.page_manager import PageManager
-from database.space_manager import SpaceManager
+from database.space_manager import upsert_space_info
 import vector.pages
 
 from .client import ConfluenceClient
@@ -32,7 +32,7 @@ def import_space(space_key, space_name, session):
 
     vector.pages.generate_missing_embeddings_to_database(session)
 
-    SpaceManager().upsert_space_info(
+    upsert_space_info(
         session,
         space_key=space_key,
         space_name=space_name,

database/space_manager.py

@@ -2,42 +2,18 @@
 from models.space_info import SpaceInfo
 
 
-class SpaceManager:
-    def __init__(self):
-        pass
+def upsert_space_info(session, space_key, space_name, last_import_date):
+    """Create or update space information based on the existence of the space key."""
+    existing_space = session.query(SpaceInfo).filter_by(space_key=space_key).first()
+    last_import_date_formatted = datetime.strptime(last_import_date, '%Y-%m-%d %H:%M:%S')
 
-    def add_space_info(self, space_key, space_name, last_import_date, session):
-        """Add a new space to the database."""
+    if existing_space:
+        existing_space.last_import_date = last_import_date_formatted
+    else:
         new_space = SpaceInfo(
             space_key=space_key,
             space_name=space_name,
-            last_import_date=datetime.strptime(last_import_date, '%Y-%m-%d %H:%M:%S')
+            last_import_date=last_import_date_formatted
         )
         session.add(new_space)
-        session.commit()
-
-    def update_space_info(self, space_key, last_import_date, session):
-        """Update the last import date of an existing space."""
-        space = session.query(SpaceInfo).filter_by(space_key=space_key).first()
-
-        if space:
-            space.last_import_date = datetime.strptime(last_import_date, '%Y-%m-%d %H:%M:%S')
-            session.commit()
-        else:
-            print(f"Space with key {space_key} not found.")
-
-    def upsert_space_info(self, session, space_key, space_name, last_import_date):
-        """Create or update space information based on the existence of the space key."""
-        existing_space = session.query(SpaceInfo).filter_by(space_key=space_key).first()
-        last_import_date_formatted = datetime.strptime(last_import_date, '%Y-%m-%d %H:%M:%S')
-
-        if existing_space:
-            existing_space.last_import_date = last_import_date_formatted
-        else:
-            new_space = SpaceInfo(
-                space_key=space_key,
-                space_name=space_name,
-                last_import_date=last_import_date_formatted
-            )
-            session.add(new_space)
-        session.commit()
+    session.commit()

interactions/identify_knowledge_gap.py

@@ -130,7 +130,7 @@ def query_assistant_with_context(context, formatted_interactions, thread_id=None
     return assistant_response, thread_id
 
 
-def process_and_store_questions(assistant_response_json, db_session):
+def process_and_store_questions(assistant_response_json, session):
     """
     Processes the JSON response from the assistant, extracts questions, stores them in the database,
     and collects the QuizQuestionDTO objects.
@@ -147,7 +147,7 @@ def process_and_store_questions(assistant_response_json, db_session):
         logging.error(f"Error decoding JSON: {e}")
         return []
 
-    quiz_question_manager = QuizQuestionManager(db_session)
+    quiz_question_manager = QuizQuestionManager(session)
 
     quiz_question_dtos = []
     for item in questions_data:

main.py

@@ -10,9 +10,9 @@
 from database.database import get_db_session
 
 
-def answer_question_with_assistant(question, db_session):
+def answer_question_with_assistant(question, session):
     page_ids = vector.pages.retrieve_relevant_ids(question, count=question_context_pages_count)
-    response, thread_id = query_assistant_with_context(question, page_ids, db_session)
+    response, thread_id = query_assistant_with_context(question, page_ids, session)
     return response, thread_id
 
 

models/page_data.py

@@ -1,5 +1,6 @@
 from database.database import Base
-from sqlalchemy import Column, Integer, String, Text, DateTime
+from sqlalchemy import Column, Integer, String, Text, DateTime, BigInteger
+import json
 
 
 class PageData(Base):
@@ -9,14 +10,14 @@ class PageData(Base):
     __tablename__ = 'page_data'
 
     id = Column(Integer, primary_key=True)
-    page_id = Column(String)
+    page_id = Column(BigInteger)
     space_key = Column(String)
     title = Column(String)
     author = Column(String)
     createdDate = Column(DateTime)
     lastUpdated = Column(DateTime)
     content = Column(Text)
-    comments = Column(Text)
+    comments = Column(Text, default=json.dumps([]))
     last_embedded = Column(DateTime)
-    embed = Column(Text)
+    embed = Column(Text, default=json.dumps([]))
 

models/qa_interaction.py

@@ -1,5 +1,6 @@
 from database.database import Base
 from sqlalchemy import Column, Integer, String, Text, DateTime
+import json
 
 
 class QAInteraction(Base):
@@ -17,6 +18,6 @@ class QAInteraction(Base):
     slack_user_id = Column(String)
     question_timestamp = Column(DateTime)
     answer_timestamp = Column(DateTime)
-    comments = Column(Text)
+    comments = Column(Text, default=json.dumps([]))
     last_embedded = Column(DateTime)
-    embed = Column(Text)
+    embed = Column(Text, default=json.dumps([]))

open_ai/assistants/query_assistant_from_documents.py

@@ -42,7 +42,7 @@ def format_pages_as_context(pages: List[PageData], max_length=30000, truncation_
     return context
 
 
-def query_assistant_with_context(question, page_ids, db_session, thread_id=None):
+def query_assistant_with_context(question, page_ids, session, thread_id=None):
     """
     Queries the assistant with a specific question, after setting up the necessary context by adding relevant files.
 
@@ -71,7 +71,7 @@ def query_assistant_with_context(question, page_ids, db_session, thread_id=None)
     print(f"IDs of pages to load in context : {page_ids}\n")
 
     # Format the context
-    pages = PageManager().find_pages(page_ids, db_session)
+    pages = PageManager().find_pages(page_ids, session)
     context = format_pages_as_context(pages)
     print(f"\n\nContext formatted: {context}\n")
 

slack/bot.py

@@ -44,7 +44,7 @@ def start(self):
 def load_slack_bot():
     logging.basicConfig(level=logging.INFO)
     bot_user_id = get_bot_user_id(slack_bot_user_oauth_token)
-    event_handlers = [ChannelMessageHandler(), ]
+    event_handlers = [ChannelMessageHandler()]
     bot = SlackBot(slack_bot_user_oauth_token, slack_app_level_token, bot_user_id, event_handlers)
     bot.start()
 

slack/channel_message_handler.py

@@ -115,7 +115,7 @@ def handle(self, client: SocketModeClient, req: SocketModeRequest, web_client: W
                 "channel": channel,
                 "user": user_id
             }
-            post_request(questions_endpoint, payload, data_type='Question')
+            post_request(questions_endpoint, payload)
 
         # Identify and handle feedback
         elif thread_ts in self.questions:  # Message is a reply to a question
@@ -129,7 +129,7 @@ def handle(self, client: SocketModeClient, req: SocketModeRequest, web_client: W
                 "user": user_id,
                 "parent_question": parent_question
             }
-            post_request(feedback_endpoint, payload, data_type='Feedback')
+            post_request(feedback_endpoint, payload)
 
         else:
             logging.info(f"Skipping message with ID {ts} from user {user_id}. Reason: {skip_reason}")

slack/event_consumer.py

@@ -45,32 +45,34 @@ def get_user_name_from_id(slack_web_client, user_id):
 
 
 class EventConsumer:
-    def __init__(self, session):
-        self.session = session
+    def __init__(self):
         self.web_client = WebClient(token=slack_bot_user_oauth_token)
-        self.interaction_manager = QAInteractionManager(session)
         logging.log(logging.DEBUG, f"Slack Event Consumer initiated successfully")
 
-    def add_question_and_response_to_database(self, question_event, response_text, assistant_thread_id):
-        self.interaction_manager.add_question_and_answer(question=question_event.text,
-                                                         answer=response_text,
-                                                         thread_id=question_event.ts,
-                                                         assistant_thread_id=assistant_thread_id,
-                                                         channel_id=question_event.channel,
-                                                         question_ts=datetime.fromtimestamp(float(question_event.ts)),
-                                                         answer_ts=datetime.now(),
-                                                         slack_user_id=question_event.user)
-
-        print(f"\n\nQuestion and answer stored in the database: question: {question_event.dict()},\nAnswer: {response_text},\nAssistant_id {assistant_thread_id}\n\n")
-
-    def process_question(self, question_event: QuestionEvent):  # TODO: Refactor this method
+    def add_question_and_response_to_database(self, session, question_event, response_text, assistant_thread_id):
+        QAInteractionManager(session).add_question_and_answer(question=question_event.text,
+                                                              answer=response_text,
+                                                              thread_id=question_event.ts,
+                                                              assistant_thread_id=assistant_thread_id,
+                                                              channel_id=question_event.channel,
+                                                              question_ts=datetime.fromtimestamp(
+                                                                  float(question_event.ts)),
+                                                              answer_ts=datetime.now(),
+                                                              slack_user_id=question_event.user)
+
+        print(
+            f"\n\nQuestion and answer stored in the database: question: {question_event.dict()},\nAnswer: {response_text},\nAssistant_id {assistant_thread_id}\n\n")
+
+    def process_question(self, session, question_event: QuestionEvent):  # TODO: Refactor this method
         channel_id = question_event.channel
         message_ts = question_event.ts
+
         try:
-            context_page_ids = vector.pages.retrieve_relevant_ids(question_event.text, count=question_context_pages_count)
+            context_page_ids = vector.pages.retrieve_relevant_ids(question_event.text,
+                                                                  count=question_context_pages_count)
             response_text, assistant_thread_id = query_assistant_with_context(question_event.text,
                                                                               context_page_ids,
-                                                                              self.session,
+                                                                              session,
                                                                               None)
         except Exception as e:
             print(f"Error processing question: {e}")
@@ -79,33 +81,36 @@ def process_question(self, question_event: QuestionEvent):  # TODO: Refactor thi
         if response_text:
             print(f"Response from assistant: {response_text}\n")
             try:
-                self.add_question_and_response_to_database(question_event,
+                self.add_question_and_response_to_database(session,
+                                                           question_event,
                                                            response_text,
                                                            assistant_thread_id=assistant_thread_id)
                 try:
-                    ScoreManager(self.session).add_or_update_score(slack_user_id=question_event.user, category='seeker')
+                    ScoreManager(session).add_or_update_score(slack_user_id=question_event.user, category='seeker')
                     print(f"Score updated for user {question_event.user}")
                 except Exception as e:
                     print(f"Error updating score for user {question_event.user}: {e}")
                 self.web_client.chat_postMessage(channel=channel_id, text=response_text, thread_ts=message_ts)
                 print(f"\nResponse posted to Slack thread: {message_ts}\n")
             except Exception as e:
-                print(f"Error registering message as processed, adding to db and responding to the question on slack: {e}")
+                print(
+                    f"Error registering message as processed, adding to db and responding to the question on slack: {e}")
 
     def generate_extended_context_query(self, existing_interaction, feedback_text):
         extended_context_query = ""
         if existing_interaction:
             extended_context_query = f"Follow up: {feedback_text}, Initial question: {existing_interaction.question_text}, Initial answer: {existing_interaction.answer_text}"
         return extended_context_query
 
-    def process_feedback(self, feedback_event: FeedbackEvent):  # TODO: Refactor this method
+    def process_feedback(self, session, feedback_event: FeedbackEvent):  # TODO: Refactor this method
         channel_id = feedback_event.channel
         message_ts = feedback_event.ts
         thread_ts = feedback_event.thread_ts
         response_text = None
+        interaction_manager = QAInteractionManager(session)
 
         try:
-            existing_interaction = self.interaction_manager.get_interaction_by_thread_id(thread_ts)
+            existing_interaction = interaction_manager.get_interaction_by_thread_id(thread_ts)
             assistant_thread_id = existing_interaction.assistant_thread_id if existing_interaction else None
             print(f"\n\nExisting interaction found: {existing_interaction}\n\n")
         except Exception as e:
@@ -117,32 +122,32 @@ def process_feedback(self, feedback_event: FeedbackEvent):  # TODO: Refactor thi
             print(f"\n\nExtended context: {extended_context_query}\n\n")
             page_ids = vector.pages.retrieve_relevant_ids(extended_context_query, count=question_context_pages_count)
             try:
-                response_text, assistant_thread_id = query_assistant_with_context(feedback_event.text, page_ids, assistant_thread_id)
+                response_text, assistant_thread_id = query_assistant_with_context(feedback_event.text, page_ids,
+                                                                                  assistant_thread_id)
             except Exception as e:
                 print(f"Error processing feedback: {e}")
                 response_text = None
 
         if response_text:
             print(f"Response from assistant: {response_text}\n")
             timestamp_str = datetime.now().isoformat()
-            comment = {"text": feedback_event.text, "user": feedback_event.user, "timestamp": timestamp_str, "assistant response": response_text}
-            self.interaction_manager.add_comment_to_interaction(thread_id=thread_ts, comment=comment)
+            comment = {"text": feedback_event.text, "user": feedback_event.user, "timestamp": timestamp_str,
+                       "assistant response": response_text}
+            interaction_manager.add_comment_to_interaction(thread_id=thread_ts, comment=comment)
             print(f"Feedback appended to the interaction in the database: {feedback_event.dict()}\n")
             self.web_client.chat_postMessage(channel=channel_id, text=response_text, thread_ts=thread_ts)
             print(f"Feedback response posted to Slack thread: {message_ts}\n")
         else:
             print(f"No response generated for feedback: {feedback_event.dict()}\n")
 
 
-# TODO: Move session call to method level
 def process_question(question_event: QuestionEvent):
     """Directly processes a question event without using the queue."""
     with get_db_session() as session:
-        EventConsumer(session).process_question(question_event)
+        EventConsumer().process_question(session, question_event)
 
 
-# TODO: Move session call to method level
 def process_feedback(feedback_event: FeedbackEvent):
     """Directly processes a feedback event without using the queue."""
     with get_db_session() as session:
-        EventConsumer(session).process_feedback(feedback_event)
+        EventConsumer().process_feedback(session, feedback_event)

slack/message_manager.py

@@ -7,7 +7,7 @@
 from slack_sdk.errors import SlackApiError
 
 
-def post_questions_to_slack(db_session, channel_id, quiz_question_dtos, user_ids):
+def post_questions_to_slack(session, channel_id, quiz_question_dtos, user_ids):
     """
     Posts a list of QuizQuestionDTO objects to a specified Slack channel, tags all users in the first reply,
     invites them to contribute to the information gathering related to their questions, asks them to tag domain
@@ -23,8 +23,8 @@ def post_questions_to_slack(db_session, channel_id, quiz_question_dtos, user_ids
     """
 
     client = WebClient(token=slack_bot_user_oauth_token)
-    quiz_question_manager = QuizQuestionManager(db_session)
-    score_manager = ScoreManager(db_session)
+    quiz_question_manager = QuizQuestionManager(session)
+    score_manager = ScoreManager(session)
 
     for quiz_question_dto in quiz_question_dtos:
         try: