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

@@ -20,7 +20,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()
@@ -47,26 +47,30 @@ 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 ComparatorTableToCSV(BaseModel):
+    file_name: str
 
 class UserInDB:
     def __init__(self, username: str, hashed_password: str):
@@ -246,6 +250,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())
@@ -338,14 +353,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.
@@ -438,17 +464,19 @@ async def retrieve_comparator(
             user: Union[None, UserInDB] = Depends(self.get_auth_dependency())
         ):
             try:
-                rows = component["component"].comparator_db.retrieve_all()
+                rows = component["component"].comparator_db.retrieve_all_question_answers()
                 data = []
                 for row in rows:
-                    _, model_name, qid, rank, answer, _ = row
+                    a_id, model_name, qid, question, answer, rank, _ = row
 
                     data.append(
                         {
+                            "id": a_id,
                             "model": model_name,
                             "qid": qid,
-                            "rank": rank,
+                            "question": question,
                             "answer": answer,
+                            "rank": rank,
                         }
                     )
                 return {"data": data, "log": "Table retrieved", "status": "200"}
@@ -466,6 +494,19 @@ async def close_comparator(
             except Exception as ex:
                 return {"log": f"Table close failed: {ex}", "status": "500"}
 
+        @app.post("/chat/comparator/db/save_to_csv")
+        async def save_comparator_table_to_csv(
+            request_body: ComparatorTableToCSV,
+            user: Union[None, UserInDB] = Depends(self.get_auth_dependency()),
+        ):
+            try:
+                print("Saving Comparator to CSV", request_body.file_name)
+                component["component"].comparator_db.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_qa_retrieval_route(self, app: FastAPI, component: Dict[str, Any]):
         """
         Create QA retrieval routes for the application.

pykoi/chat/db/comparator_database.py

@@ -1,4 +1,5 @@
 """Comparator Database"""
+import csv
 import datetime
 import os
 
@@ -8,6 +9,7 @@
 
 
 from pykoi.chat.db.abs_database import AbsDatabase
+from pykoi.chat.db.constants import COMPARATOR_CSV_HEADER
 
 
 class ComparatorQuestionDatabase(AbsDatabase):
@@ -199,6 +201,25 @@ def retrieve_all(self) -> List[Tuple]:
             rows = cursor.fetchall()
         return rows
 
+    def retrieve_all_question_answers(self):
+        """
+        Retrieves all question-answer pairs from the database.
+
+        Returns:
+            rows: rows of data of the question-answer pairs.
+        """
+        query = """
+        SELECT comparator.id, comparator.model, comparator.qid, comparator_question.question, comparator.answer, comparator.rank, comparator.timestamp
+        FROM comparator
+        JOIN comparator_question
+        ON comparator.qid = comparator_question.id;
+        """
+        with self._lock:
+            cursor = self.get_cursor()
+            cursor.execute(query)
+            rows = cursor.fetchall()
+        return rows
+
     def print_table(self, rows: List[Tuple]) -> None:
         """
         Prints the comparator table.
@@ -217,6 +238,29 @@ def print_table(self, rows: List[Tuple]) -> None:
                 f"Timestamp: {row[5]}"
             )
 
+
+    def save_to_csv(self, csv_file_name="comparator_table"):
+        """
+        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 "comparator_table".
+
+        The CSV file will have the following columns: TODO. 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 method. It then writes this data to the CSV file.
+        """
+
+        my_sql_data = self.retrieve_all_question_answers()
+
+        with open(csv_file_name + ".csv", "w", newline="") as file:
+            writer = csv.writer(file)
+            writer.writerow(COMPARATOR_CSV_HEADER)
+            writer.writerows(my_sql_data)
+
     def retrieve_all_question_answers_as_pandas(self) -> pd.DataFrame:
         """
         Retrieves all data by joining the comparator and comparator_question tables as a pandas dataframe.
@@ -248,3 +292,4 @@ def retrieve_all_question_answers_as_pandas(self) -> pd.DataFrame:
             columns=["ID", "Model", "QID", "Question", "Rank", "Answer", "Timestamp"],
         )
         return df
+

pykoi/chat/db/constants.py

@@ -26,5 +26,45 @@
     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 = "||"
+
+# Comparator table
+COMPARATOR_CSV_HEADER_ID = "ID"
+COMPARATOR_CSV_HEADER_MODEL = "Model"
+COMPARATOR_CSV_HEADER_QID = "Question ID"
+COMPARATOR_CSV_HEADER_QUESTION = "Question"
+COMPARATOR_CSV_HEADER_ANSWER = "Answer"
+COMPARATOR_CSV_HEADER_RANK = "Rank"
+COMPARATOR_CSV_HEADER_TIMESTAMP = "Timestamp"
+COMPARATOR_CSV_HEADER = (
+    COMPARATOR_CSV_HEADER_ID,
+    COMPARATOR_CSV_HEADER_MODEL,
+    COMPARATOR_CSV_HEADER_QID,
+    COMPARATOR_CSV_HEADER_QUESTION,
+    COMPARATOR_CSV_HEADER_ANSWER,
+    COMPARATOR_CSV_HEADER_RANK,
+    COMPARATOR_CSV_HEADER_TIMESTAMP,
+)

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)