deepnote update

app.py

@@ -5,6 +5,7 @@ import util as util
 import uvicorn
 import gradio as gr
 from typing import List, Optional
+from fastapi.responses import JSONResponse
 
 app = FastAPI()
 
@@ -13,15 +14,9 @@ class AskRequest(BaseModel):
     question: str
     sheet_url: str
     page_content_column: str
-    k: int
-
-
-class AskRequestEx(BaseModel):
-    question: str
-    sheet_url: str
-    page_content_column: str
-    k: int
-    id_column: str
+    k: int = 20
+    reload_collection: Optional[bool] = None
+    id_column: Optional[str] = None
     synonyms: Optional[List[List[str]]] = None
 
 
@@ -33,15 +28,17 @@ async def ask_api(request: AskRequest):
 
 
 @app.post("/api/v2/ask")
-async def ask_api(request: AskRequestEx):
-    util.SPLIT_PAGE_BREAKS = True
+async def ask_api(request: AskRequest):
+    if request.id_column is not None:
+        util.SPLIT_PAGE_BREAKS = True
     if request.synonyms is not None:
         util.SYNONYMS = request.synonyms
     vectordb = faq.load_vectordb(request.sheet_url, request.page_content_column)
     documents = faq.similarity_search(vectordb, request.question, k=request.k)
     df_doc = util.transform_documents_to_dataframe(documents)
-    df_filter = util.remove_duplicates_by_column(df_doc, request.id_column)
-    return util.dataframe_to_dict(df_filter)
+    if request.id_column is not None:
+        df_doc = util.remove_duplicates_by_column(df_doc, request.id_column)
+    return JSONResponse(util.dataframe_to_dict(df_doc))
 
 
 @app.delete("/api/v1/")
@@ -52,8 +49,9 @@ async def delete_vectordb_api():
 def ask(sheet_url: str, page_content_column: str, k: int, question: str):
     util.SPLIT_PAGE_BREAKS = False
     vectordb = faq.load_vectordb(sheet_url, page_content_column)
-    result = faq.similarity_search(vectordb, question, k=k)
-    return result
+    documents = faq.similarity_search(vectordb, question, k=k)
+    df_doc = util.transform_documents_to_dataframe(documents)
+    return util.dataframe_to_dict(df_doc)
 
 
 def delete_vectordb():
@@ -63,7 +61,7 @@ def delete_vectordb():
 with gr.Blocks() as block:
     sheet_url = gr.Textbox(label="Google Sheet URL")
     page_content_column = gr.Textbox(label="Question Column")
-    k = gr.Slider(2, 5, step=1, label="K")
+    k = gr.Slider(1, 30, step=1, label="K")
     question = gr.Textbox(label="Question")
     ask_button = gr.Button("Ask")
     answer_output = gr.JSON(label="Answer")

faq.py

@@ -32,7 +32,7 @@ def define_embedding_function(model_name: str) -> HuggingFaceEmbeddings:
 
 
 def get_vectordb(
-    faq_id: str,
+    collection_id: str,
     embedding_function: Embeddings,
     documents: List[Document] = None,
     vectordb_type: str = VECTORDB_TYPE,
@@ -44,31 +44,32 @@ def get_vectordb(
             vectordb = AwaDB(
                 embedding=embedding_function, log_and_data_dir=VECTORDB_FOLDER
             )
-            if not vectordb.load_local(table_name=faq_id):
-                raise Exception("faq_id may not exists")
+            if not vectordb.load_local(table_name=collection_id):
+                raise Exception("collection_id may not exists")
         else:
             vectordb = AwaDB.from_documents(
                 documents=documents,
                 embedding=embedding_function,
-                table_name=faq_id,
+                table_name=collection_id,
                 log_and_data_dir=VECTORDB_FOLDER,
             )
     if vectordb_type is VECTORDB_TYPES.Chroma:
         if documents is None:
             vectordb = Chroma(
-                collection_name=faq_id,
+                collection_name=collection_id,
                 embedding_function=embedding_function,
                 persist_directory=VECTORDB_FOLDER,
             )
             if not vectordb.get()["ids"]:
-                raise Exception("faq_id may not exists")
+                raise Exception("collection_id may not exists")
         else:
             vectordb = Chroma.from_documents(
                 documents=documents,
                 embedding=embedding_function,
-                collection_name=faq_id,
+                collection_name=collection_id,
                 persist_directory=VECTORDB_FOLDER,
             )
+            vectordb.persist()
     return vectordb
 
 
@@ -80,33 +81,33 @@ def similarity_search(
 
 
 def load_vectordb_id(
-    faq_id: str,
+    collection_id: str,
     page_content_column: str,
     embedding_function_name: str = EMBEDDING_MODEL,
 ) -> VectorStore:
     embedding_function = define_embedding_function(embedding_function_name)
     vectordb = None
     try:
-        vectordb = get_vectordb(faq_id=faq_id, embedding_function=embedding_function)
+        vectordb = get_vectordb(collection_id=collection_id, embedding_function=embedding_function)
     except Exception as e:
         print(e)
-        vectordb = create_vectordb_id(faq_id, page_content_column, embedding_function)
+        vectordb = create_vectordb_id(collection_id, page_content_column, embedding_function)
 
     return vectordb
 
 
 def create_vectordb_id(
-    faq_id: str,
+    collection_id: str,
     page_content_column: str,
     embedding_function: HuggingFaceEmbeddings = None,
 ) -> VectorStore:
     if embedding_function is None:
         embedding_function = define_embedding_function(EMBEDDING_MODEL)
 
-    df = util.read_df(util.xlsx_url(faq_id), page_content_column)
+    df = util.read_df(util.xlsx_url(collection_id), page_content_column)
     documents = create_documents(df, page_content_column)
     vectordb = get_vectordb(
-        faq_id=faq_id, embedding_function=embedding_function, documents=documents
+        collection_id=collection_id, embedding_function=embedding_function, documents=documents
     )
     return vectordb
 
@@ -115,5 +116,10 @@ def load_vectordb(sheet_url: str, page_content_column: str) -> VectorStore:
     return load_vectordb_id(util.get_id(sheet_url), page_content_column)
 
 
-def delete_vectordb():
+def delete_vectordb() -> None:
     shutil.rmtree(VECTORDB_FOLDER, ignore_errors=True)
+
+
+def delete_vectordb_current_collection(vectordb: VectorStore) -> None:
+    vectordb.delete_collection()
+    vectordb.persist()

util.py

@@ -1,4 +1,5 @@
 import pandas as pd
+from langchain.docstore.document import Document
 
 SHEET_URL_X = "https://docs.google.com/spreadsheets/d/"
 SHEET_URL_Y = "/edit#gid="
@@ -27,7 +28,7 @@ def read_df(xlsx_url: str, page_content_column: str) -> pd.DataFrame:
     return df
 
 
-def split_page_breaks(df, column_name):
+def split_page_breaks(df: pd.DataFrame, column_name: str) -> pd.DataFrame:
     split_values = df[column_name].str.split("\n")
 
     new_df = pd.DataFrame({column_name: split_values.explode()})
@@ -46,37 +47,35 @@ def split_page_breaks(df, column_name):
     return new_df
 
 
-def transform_documents_to_dataframe(documents):
-    metadata_keys = set()
-    for doc, _ in documents:
-        metadata_keys.update(doc.metadata.keys())
+def transform_documents_to_dataframe(documents: Document) -> pd.DataFrame:
+    keys = []
+    values = {"document_score": [], "page_content": []}
 
-    metadata_values = {key: [] for key in metadata_keys}
-    for doc, _ in documents:
+    for doc, score in documents:
         for key, value in doc.metadata.items():
-            metadata_values[key].append(value)
+            if key not in keys:
+                keys.append(key)
+                values[key] = []
+            values[key].append(value)
+        values["document_score"].append(score)
+        values["page_content"].append(doc.page_content)
 
-    metadata_values["Score"] = [score for _, score in documents]
+    return pd.DataFrame(values)
 
-    df = pd.DataFrame(metadata_values)
 
-    return df
-
-
-def remove_duplicates_by_column(df, column):
-    df.drop_duplicates(subset=column, inplace=True)
-    df.reset_index(drop=True, inplace=True)
+def remove_duplicates_by_column(df: pd.DataFrame, column_name: str) -> pd.DataFrame:
+    df.drop_duplicates(subset=column_name, inplace=True, ignore_index=True)
 
     return df
 
 
-def dataframe_to_dict(df):
+def dataframe_to_dict(df: pd.DataFrame) -> dict:
     df_records = df.to_dict(orient="records")
 
     return df_records
 
 
-def duplicate_rows_with_synonyms(df, column, synonyms):
+def duplicate_rows_with_synonyms(df: pd.DataFrame, column: str, synonyms: list[list[str]]) -> pd.DataFrame:
     new_rows = []
     for index, row in df.iterrows():
         new_rows.append(row)