deepnote update

app.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 import faq as faq
+import util as util
 import uvicorn
 import gradio as gr
 
@@ -21,6 +22,21 @@ async def ask_api(request: AskRequest):
     )
 
 
+@app.post("/api/v2/ask")
+async def ask_api(request: AskRequest):
+    faq_id = faq.faq_id(request.sheet_url)
+    xlsx_url = faq.xlsx_url(faq_id)
+    df = faq.read_df(xlsx_url)
+    df_update = util.split_page_breaks(df, request.page_content_column)
+    documents = faq.create_documents(df_update, request.page_content_column)
+    embedding_function = faq.define_embedding_function("sentence-transformers/all-mpnet-base-v2")
+    vectordb = faq.get_vectordb(faq_id=faq_id, embedding_function=embedding_function, documents=documents, vectordb_type=faq.VECTORDB_TYPE.Chroma)
+    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, "ID")
+    return util.serialize_dataframe_as_json(df_filter)
+
+
 @app.delete("/api/v1/")
 async def delete_vectordb_api():
     return delete_vectordb()

util.py

@@ -0,0 +1,49 @@
+import pandas as pd
+
+def split_page_breaks(df, column_name):
+    split_values = df[column_name].str.split("\n")
+
+    new_df = pd.DataFrame({column_name: split_values.explode()})
+    new_df.reset_index(drop=True, inplace=True)
+
+    column_order = df.columns
+
+    new_df = new_df.reindex(column_order, axis=1)
+
+    other_columns = column_order.drop(column_name)
+    for column in other_columns:
+        new_df[column] = (
+            df[column].repeat(split_values.str.len()).reset_index(drop=True)
+        )
+
+    return new_df
+
+
+def transform_documents_to_dataframe(documents):
+    metadata_keys = set()
+    for doc, _ in documents:
+        metadata_keys.update(doc.metadata.keys())
+
+    metadata_values = {key: [] for key in metadata_keys}
+    for doc, _ in documents:
+        for key, value in doc.metadata.items():
+            metadata_values[key].append(value)
+
+    metadata_values["Score"] = [score for _, score in documents]
+
+    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)
+
+    return df
+
+
+def serialize_dataframe_as_json(df):
+    json_array = df.to_dict(orient='records')
+
+    return json_array