Cuongtruong's request (#1)

- SQLite3 fix db_instance on websocket creation (1ebe5e1fc9e08fa37c143808aaf4dc4ebd182440)


Co-authored-by: Trương Tấn Cường <[email protected]>

.gitignore

@@ -0,0 +1,5 @@
+arxivdb/
+models/
+apikey.txt
+db.sqlite3
+hotfix.ipynb

chat/consumers.py

@@ -1,18 +1,21 @@
 import json
 from . import model_manage as md
+from chat.arxiv_bot.arxiv_bot_utils import ArxivSQL
 from channels.generic.websocket import WebsocketConsumer
 
 
 class ChatConsumer(WebsocketConsumer):
     def connect(self):
         self.accept()
+        self.db_instance = ArxivSQL()
+        
     def disconnect(self, close_code):
         pass
     def receive(self, text_data):
         text_data_json = json.loads(text_data)
         message = text_data_json["message"]
         print(message)
-        record, messagee = md.full_chain_single_question(message)
+        record, messagee = md.full_chain_single_question(message, self.db_instance)
         print("First answer: ",record)
         self.send(text_data=json.dumps({"message": messagee}))
         

chat/model_manage.py

@@ -104,12 +104,12 @@ def response(args):
             if type(new_records) == str:
                 return "Error occured, information not found", new_records
             utils.db.add(new_records)
-            utils.sqldb.add(new_records)
+            db_instance.add(new_records)
             results = utils.db.query_relevant(keywords=keywords, query_texts=query_texts)
             ids = results['metadatas'][0]
             print("Re-queried on chromadb, results: ",ids)
         paper_id = [id['paper_id'] for id in ids]
-        paper_info = utils.sqldb.query_id(paper_id)
+        paper_info = db_instance.query_id(paper_id)
         print(paper_info)
         records = [] # get title (2), author (3), link (6)
         result_string = ""
@@ -125,7 +125,7 @@ def response(args):
     if "title" in keys:
         title = args['title']
         authors = utils.authors_str_to_list(args['author'])
-        paper_info = utils.sqldb.query(title = title,author = authors)
+        paper_info = db_instance.query(title = title,author = authors)
         # if query not found then go crawl brh
         # print(paper_info)
 
@@ -136,8 +136,8 @@ def response(args):
                 # print(new_records)
                 return "Error occured, information not found", "Information not found"
             utils.db.add(new_records)
-            utils.sqldb.add(new_records)
-            paper_info = utils.sqldb.query(title = title,author = authors)
+            db_instance.add(new_records)
+            paper_info = db_instance.query(title = title,author = authors)
             print("Re-queried on chromadb, results: ",paper_info)
         # -------------------------------------
         records = [] # get title (2), author (3), link (6)
@@ -150,13 +150,13 @@ def response(args):
             return "Information not found", "Information not found"
         return result_string, records
         # invoke llm and return result
-def full_chain_single_question(input_prompt):
+def full_chain_single_question(input_prompt, db_instance):
     try:
         first_prompt = extract_keyword_prompt(input_prompt)
         temp_answer = model.generate_content(first_prompt).text
 
         args = json.loads(utils.trimming(temp_answer))
-        contexts, results = response(args)
+        contexts, results = response(args, db_instance)
         if not results:
             # print(contexts)
             return "Random question, direct return", contexts