deepnote update

app.py

@@ -7,7 +7,7 @@ import gradio as gr
 app = FastAPI()
 
 
-class Request(BaseModel):
+class AskRequest(BaseModel):
     question: str
     sheet_url: str
     page_content_column: str
@@ -15,31 +15,43 @@ class Request(BaseModel):
 
 
 @app.post("/api/v1/ask")
-async def ask_api(request: Request):
+async def ask_api(request: AskRequest):
     return ask(
         request.sheet_url, request.page_content_column, request.k, request.question
     )
 
 
+@app.delete("/api/v1/")
+async def delete_vectordb_api():
+    return delete_vectordb()
+
+
 def ask(sheet_url: str, page_content_column: str, k: int, question: str):
     vectordb = faq.load_vectordb(sheet_url, page_content_column)
     result = faq.similarity_search(vectordb, question, k=k)
     return result
 
 
-iface = gr.Interface(
-    fn=ask,
-    inputs=[
-        gr.Textbox(label="Google Sheet URL"),
-        gr.Textbox(label="Question Column"),
-        gr.Slider(1, 5, step=1, label="K"),
-        gr.Textbox(label="Question"),
-    ],
-    outputs=[gr.JSON(label="Answer")],
-    allow_flagging="never",
-)
+def delete_vectordb():
+    faq.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")
+    question = gr.Textbox(label="Question")
+    ask_button = gr.Button("Ask")
+    answer_output = gr.JSON(label="Answer")
+    delete_button = gr.Button("Delete Vector DB")
+    ask_button.click(
+        ask,
+        inputs=[sheet_url, page_content_column, k, question],
+        outputs=answer_output,
+    )
+    delete_button.click(delete_vectordb)
 
-app = gr.mount_gradio_app(app, iface, path="/")
+app = gr.mount_gradio_app(app, block, path="/")
 
 
 if __name__ == "__main__":

faq.py

@@ -7,12 +7,14 @@ from langchain.docstore.document import Document
 from langchain.embeddings.base import Embeddings
 from langchain.vectorstores.base import VectorStore
 import os
+import shutil
 
 SHEET_URL_X = "https://docs.google.com/spreadsheets/d/"
 SHEET_URL_Y = "/edit#gid="
 SHEET_URL_Y_EXPORT = "/export?gid="
 CACHE_FOLDER = ".embedding-model"
 VECTORDB_FOLDER = ".vectordb"
+EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
 
 
 def faq_id(sheet_url: str) -> str:
@@ -69,17 +71,40 @@ def similarity_search(
     return vectordb.similarity_search_with_relevance_scores(query=query, k=k)
 
 
-def load_vectordb_id(faq_id: str, page_content_column: str) -> VectorStore:
-    embedding_function = define_embedding_function("sentence-transformers/all-mpnet-base-v2")
+def load_vectordb_id(
+    faq_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)
     except Exception as e:
-        df = read_df(xlsx_url(faq_id))
-        documents = create_documents(df, page_content_column)
-        vectordb = get_vectordb(faq_id=faq_id, embedding_function=embedding_function, documents=documents)
+        vectordb = create_vectordb_id(faq_id, page_content_column, embedding_function)
+
+    return vectordb
+
+
+def create_vectordb_id(
+    faq_id: str,
+    page_content_column: str,
+    embedding_function: HuggingFaceEmbeddings = None,
+) -> VectorStore:
+    if embedding_function is None:
+        embedding_function = define_embedding_function(EMBEDDING_MODEL)
+
+    df = read_df(xlsx_url(faq_id))
+    documents = create_documents(df, page_content_column)
+    vectordb = get_vectordb(
+        faq_id=faq_id, embedding_function=embedding_function, documents=documents
+    )
     return vectordb
 
 
 def load_vectordb(sheet_url: str, page_content_column: str) -> VectorStore:
     return load_vectordb_id(faq_id(sheet_url), page_content_column)
+
+
+def delete_vectordb():
+    shutil.rmtree(VECTORDB_FOLDER, ignore_errors=True)