Add CSV Download frontend and backend

example/chatbot/demo_launch_app_cpu_openai.py

@@ -6,24 +6,29 @@
     You can follow [the installation guide](https://github.com/CambioML/pykoi/tree/install#option-1-rag-cpu)
     to set up the environment.
 - Run the demo:
-    1. Enter your OpenAI API key in the `api_key` below.
+    1. Enter your OpenAI API key a .env file in the `~/pykoi` directory with the name OPEN_API_KEY, e.g.
+        ```
+        OPENAI_API_KEY=your_api_key
+        ```
     2. On terminal and `~/pykoi` directory, run
         ```
         python -m example.chatbot.demo_launch_app_cpu_openai
         ```
 """
+import os
+
+from dotenv import load_dotenv
 
 from pykoi import Application
 from pykoi.chat import ModelFactory
 from pykoi.chat import QuestionAnswerDatabase
 from pykoi.component import Chatbot, Dashboard
 
-
 ##########################################################
 # Creating an OpenAI model (requires an OpenAI API key) #
 ##########################################################
-# enter openai api key here
-api_key = "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"
+load_dotenv()
+api_key = os.getenv("OPENAI_API_KEY")
 
 # Creating an OpenAI model
 model = ModelFactory.create_model(

example/comparator/demo_model_comparator_cpu_openai.py

@@ -2,26 +2,34 @@
 Demo for the chatbot application using multiple OpenAI models.
 
 - Prerequisites:
-    To run this jupyter notebook, you need a `pykoi` environment with the `rag` option. 
-    You can follow [the installation guide](https://github.com/CambioML/pykoi/tree/install#option-1-rag-cpu) 
-    to set up the environment. 
+    To run this jupyter notebook, you need a `pykoi` environment with the `rag` option.
+    You can follow [the installation guide](https://github.com/CambioML/pykoi/tree/install#option-1-rag-cpu)
+    to set up the environment.
 - Run the demo:
-    1. Enter your OpenAI API key in the `api_key` below.
-    2. On terminal and `~/pykoi` directory, run 
+    1. Enter your OpenAI API key a .env file in the `~/pykoi` directory with the name OPEN_API_KEY, e.g.
+        ```
+        OPENAI_API_KEY=your_api_key
+        ```
+    2. On terminal and `~/pykoi` directory, run
         ```
         python -m example.comparator.demo_model_comparator_cpu_openai
         ```
 """
 
+import os
+
+from dotenv import load_dotenv
+
 from pykoi import Application
 from pykoi.chat import ModelFactory
 from pykoi.component import Compare
 
+
 ##########################################################
 # Creating an OpenAI model (requires an OpenAI API key) #
 ##########################################################
-# enter openai api key here
-api_key = ""
+load_dotenv()
+api_key = os.getenv("OPENAI_API_KEY")
 
 # Creating an OpenAI model
 openai_model_1 = ModelFactory.create_model(

pykoi/application.py

@@ -17,7 +17,7 @@
 from pykoi.interactives.chatbot import Chatbot
 from pykoi.telemetry.telemetry import Telemetry
 from pykoi.telemetry.events import AppStartEvent, AppStopEvent
-from pykoi.chat.db.constants import QA_LIST_SEPARATOR
+from pykoi.chat.db.constants import RAG_LIST_SEPARATOR
 
 
 oauth_scheme = HTTPBasic()
@@ -44,26 +44,27 @@ class RankingTableUpdate(BaseModel):
     up_ranking_answer: str
     low_ranking_answer: str
 
-
 class InferenceRankingTable(BaseModel):
     n: Optional[int] = 2
 
-
 class ModelAnswer(BaseModel):
     model: str
     qid: int
     rank: int
     answer: str
 
-
 class ComparatorInsertRequest(BaseModel):
     data: List[ModelAnswer]
 
-
 class RetrievalNewMessage(BaseModel):
     prompt: str
     file_names: List[str]
 
+class QATableToCSV(BaseModel):
+    file_name: str
+
+class RAGTableToCSV(BaseModel):
+    file_name: str
 
 class UserInDB:
     def __init__(self, username: str, hashed_password: str):
@@ -243,6 +244,17 @@ async def update_qa_table_response(
             except Exception as ex:
                 return {"log": f"Table update failed: {ex}", "status": "500"}
 
+        @app.post("/chat/qa_table/save_to_csv")
+        async def save_qa_table_to_csv(
+            request_body: QATableToCSV,
+            user: Union[None, UserInDB] = Depends(self.get_auth_dependency()),
+        ):
+            try:
+                component["component"].database.save_to_csv(request_body.file_name)
+                return {"log": f"Saved to {request_body.file_name}.csv", "status": "200"}
+            except Exception as ex:
+                return {"log": f"Save to CSV failed: {ex}", "status": "500"}
+
         @app.get("/chat/qa_table/close")
         async def close_qa_table(
             user: Union[None, UserInDB] = Depends(self.get_auth_dependency())
@@ -335,14 +347,25 @@ async def retrieve_rag_table(
                 modified_rows = []
                 for row in rows:
                     row_list = list(row)  # Convert the tuple to a list
-                    row_list[5] = row_list[5].split(QA_LIST_SEPARATOR)
-                    row_list[6] = row_list[6].split(QA_LIST_SEPARATOR)
-                    row_list[7] = row_list[7].split(QA_LIST_SEPARATOR)
+                    row_list[5] = row_list[5].split(RAG_LIST_SEPARATOR)
+                    row_list[6] = row_list[6].split(RAG_LIST_SEPARATOR)
+                    row_list[7] = row_list[7].split(RAG_LIST_SEPARATOR)
                     modified_rows.append(row_list)  # Append the modified list to the new list
                 return {"rows": modified_rows, "log": "RAG Table retrieved", "status": "200"}
             except Exception as ex:
                 return {"log": f"Table retrieval failed: {ex}", "status": "500"}
 
+        @app.post("/chat/rag_table/save_to_csv")
+        async def save_rag_table_to_csv(
+            request_body: RAGTableToCSV,
+            user: Union[None, UserInDB] = Depends(self.get_auth_dependency()),
+        ):
+            try:
+                component["component"].database.save_to_csv(request_body.file_name)
+                return {"log": f"Saved to {request_body.file_name}.csv", "status": "200"}
+            except Exception as ex:
+                return {"log": f"Save to CSV failed: {ex}", "status": "500"}
+
     def create_feedback_route(self, app: FastAPI, component: Dict[str, Any]):
         """
         Create feedback routes for the application.

pykoi/chat/db/constants.py

@@ -26,5 +26,27 @@
     RANKING_CSV_HEADER_LOW_RANKING_ANSWER,
 )
 
+# RAG table
+RAG_CSV_HEADER_ID = "ID"
+RAG_CSV_HEADER_QUESTION = "Question"
+RAG_CSV_HEADER_ANSWER = "Answer"
+RAG_CSV_HEADER_EDITED = "Edited Answer"
+RAG_CSV_HEADER_VOTE_STATUS = "Vote Status"
+RAG_CSV_HEADER_RAG_SOURCES = "RAG Sources"
+RAG_CSV_HEADER_SOURCE = "Source"
+RAG_CSV_HEADER_SOURCE_CONTENT = "Source Content"
+RAG_CSV_HEADER_TIMESTAMP = "Timestamp"
+RAG_CSV_HEADER = (
+    RAG_CSV_HEADER_ID,
+    RAG_CSV_HEADER_QUESTION,
+    RAG_CSV_HEADER_ANSWER,
+    RAG_CSV_HEADER_EDITED,
+    RAG_CSV_HEADER_VOTE_STATUS,
+    RAG_CSV_HEADER_RAG_SOURCES,
+    RAG_CSV_HEADER_SOURCE,
+    RAG_CSV_HEADER_SOURCE_CONTENT,
+    RAG_CSV_HEADER_TIMESTAMP,
+)
+
 # list separator
-QA_LIST_SEPARATOR = "||"
+RAG_LIST_SEPARATOR = "||"

pykoi/chat/db/qa_database.py

@@ -180,13 +180,13 @@ def print_table(self, rows):
                 f"Answer: {row[2]}, Vote Status: {row[3]}, Timestamp: {row[4]}"
             )
 
-    def save_to_csv(self, csv_file_name="question_answer_votes.csv"):
+    def save_to_csv(self, csv_file_name="question_answer_votes"):
         """
         This method saves the contents of the question_answer table into a CSV file.
 
         Args:
             csv_file_name (str, optional): The name of the CSV file to which the data will be written.
-            Defaults to "question_answer_votes.csv".
+            Defaults to "question_answer_votes".
 
         The CSV file will have the following columns: ID, Question, Answer, Vote Status. Each row in the
         CSV file corresponds to a row in the question_answer table.
@@ -196,7 +196,7 @@ def save_to_csv(self, csv_file_name="question_answer_votes.csv"):
         """
         my_sql_data = self.retrieve_all_question_answers()
 
-        with open(csv_file_name, "w", newline="") as file:
+        with open(csv_file_name + ".csv", "w", newline="") as file:
             writer = csv.writer(file)
             writer.writerow(QA_CSV_HEADER)
             writer.writerows(my_sql_data)

pykoi/chat/db/rag_database.py

@@ -7,7 +7,7 @@
 
 import pandas as pd
 
-from pykoi.chat.db.constants import QA_CSV_HEADER, QA_LIST_SEPARATOR
+from pykoi.chat.db.constants import RAG_CSV_HEADER, RAG_LIST_SEPARATOR
 
 
 class RAGDatabase:
@@ -93,9 +93,9 @@ def insert_question_answer(self, question: str, answer: str, rag_sources: list,
         INSERT INTO rag_question_answer (question, answer, edited_answer, rag_sources, source, source_content, vote_status, timestamp)
         VALUES (?, ?, '', ?, ?, ?, 'n/a', ?);
         """
-        rag_sources = QA_LIST_SEPARATOR.join(rag_sources)
-        source = QA_LIST_SEPARATOR.join(source)
-        source_content = QA_LIST_SEPARATOR.join(source_content)
+        rag_sources = RAG_LIST_SEPARATOR.join(rag_sources)
+        source = RAG_LIST_SEPARATOR.join(source)
+        source_content = RAG_LIST_SEPARATOR.join(source_content)
         print("rag insert question answer", rag_sources)
 
         with self._lock:
@@ -194,7 +194,7 @@ def retrieve_all_question_answers_as_pandas(self):
         """
         rows = self.retrieve_all_question_answers()
         rows_to_pd = pd.DataFrame(rows)
-        rows_to_pd.columns = QA_CSV_HEADER
+        rows_to_pd.columns = RAG_CSV_HEADER
         return rows_to_pd
 
     def close_connection(self):
@@ -222,23 +222,23 @@ def print_table(self, rows):
                 f"Answer: {row[2]}, Vote Status: {row[3]}, Timestamp: {row[4]}"
             )
 
-    def save_to_csv(self, csv_file_name="question_answer_votes.csv"):
+    def save_to_csv(self, csv_file_name="rag_table"):
         """
-        This method saves the contents of the question_answer table into a CSV file.
+        This method saves the contents of the RAG table into a CSV file.
 
         Args:
             csv_file_name (str, optional): The name of the CSV file to which the data will be written.
-            Defaults to "question_answer_votes.csv".
+            Defaults to "rag_table".
 
-        The CSV file will have the following columns: ID, Question, Answer, Vote Status. Each row in the
+        The CSV file will have the following columns: ID, Question, Answer, Edited Answer, Vote Status, RAG Sources, Source, Source Content, and Timestamp. Each row in the
         CSV file corresponds to a row in the question_answer table.
 
         This method first retrieves all question-answer pairs from the database by calling the
         retrieve_all_question_answers method. It then writes this data to the CSV file.
         """
         my_sql_data = self.retrieve_all_question_answers()
 
-        with open(csv_file_name, "w", newline="") as file:
+        with open(csv_file_name + ".csv", "w", newline="") as file:
             writer = csv.writer(file)
-            writer.writerow(QA_CSV_HEADER)
+            writer.writerow(RAG_CSV_HEADER)
             writer.writerows(my_sql_data)