Backend tweak (#45)

* feat: added routes to delete knowledge bases and pipelines
* feat: custom prompts function
fix: unnested eval data
* feat: can handle custom prompt or not

cleanup.py

@@ -3,7 +3,7 @@
 import pymongo
 from dotenv import load_dotenv
 
-load_dotenv(override=True, dotenv_path=".env.testing")
+load_dotenv(override=True)
 
 
 MONGO_URI = os.environ["MONGO_URI"]

knowledge_base/mongo_helper.py

@@ -13,6 +13,7 @@
 
 MONGO_URI = os.environ["MONGO_URI"]
 CONFIG_DB = os.environ["CONFIG_DB"]
+CONFIG_PIPELINE_COL = os.environ["CONFIG_PIPELINE_COL"]
 CONFIG_KB_COL = os.environ["CONFIG_KB_COL"]
 
 def get(db_name, db_collection, query=None, projection=None):
@@ -103,5 +104,24 @@ def get_kb_id(kb_name):
         return str(result["_id"])
     else:
         return None
+
+def delete_knowledge_base(id):
+    mongo = pymongo.MongoClient(MONGO_URI)
+    kb_result = mongo[CONFIG_DB][CONFIG_KB_COL].delete_one({"id": id})
+    pipeline_result = remove_kb_from_pipeline(id)
+    mongo.close()
+
+    return kb_result
+
+def remove_kb_from_pipeline(kb_id):
+    mongo = pymongo.MongoClient(MONGO_URI)
+    result = mongo[CONFIG_DB][CONFIG_PIPELINE_COL].update_many(
+        {},
+        { "$pull": { "knowledge_bases": kb_id } }
+    )
+    log.info(f"remove_kb_from_pipeline: {result}")
+    mongo.close()
+    return result
+
 
 # get_kb_id("Sentence Split")

knowledge_base/routes.py

@@ -55,6 +55,15 @@ async def upload_file(id, file):
         except Exception as e:
             return {"message": f"Error: {e}"}
 
+def delete(id):
+    result = mongo.delete_knowledge_base(id)
+    if result.deleted_count == 1:
+        return {"message": f"{id} deleted"}
+    else:
+        return {"message": f"{id} does not exist"}
+
+
+
 # config helpers to convert strings to numbers
 # probaly should be moved to help james too!
 def str_to_nums(config_dict):

pipeline/mongo_util.py

@@ -40,6 +40,13 @@ def pipeline_name_taken(name):
     )
     return result
 
+def delete_pipeline(id):
+    mongo = pymongo.MongoClient(MONGO_URI)
+    result = mongo[CONFIG_DB][CONFIG_PIPELINE_COL].delete_one(
+        { "id": id }
+    )
+    mongo.close()
+    return result
 
 def insert_pipeline(doc):
     mongo = pymongo.MongoClient(MONGO_URI)

pipeline/pipeline_class.py

@@ -2,7 +2,7 @@
 
 import os
 
-from llama_index.core import get_response_synthesizer, QueryBundle
+from llama_index.core import get_response_synthesizer, QueryBundle, PromptTemplate
 from llama_index.core.postprocessor import SimilarityPostprocessor
 from llama_index.core.postprocessor import LongContextReorder
 from llama_index.postprocessor.colbert_rerank import ColbertRerank
@@ -36,10 +36,16 @@ def query(self, user_query):
 
             # synthesizer, nodes = self._pipeline(user_query).values()
             log.info('nodes:', nodes)
+            custom_prompt = self._config['prompt']
+            if custom_prompt:
+                log.info("using custom prompts")
+                new_template = PromptTemplate(self._config['prompt'])
+                synthesizer.update_prompts(
+                    {"text_qa_template": new_template}
+                )
+            else:
+                log.info("no custom prompts")
 
-            # Note: if we pass in a `simple template` kwarg here,
-            # I believe we can incorporate the custom prompt from 
-            # the pipeline config
             return synthesizer.synthesize(user_query, nodes=nodes)
         except Exception as err:
             log.error('pipeline.py query: ******** ERROR *********', type(err), err)

routers/chatbots.py

@@ -61,3 +61,11 @@ async def post_chatbots(request: Request):
         print("new_pipeline: ", new_pipeline)
 
         return new_pipeline
+
+@router.delete('/{id}/delete')
+async def delete_chatbot(id: str):
+    result = mutil.delete_pipeline(id)
+    if result.deleted_count == 1:
+        return {"message": f"{id} deleted"}
+    else:
+        return {"message": f"{id} does not exist"}

server.py

@@ -49,6 +49,10 @@ async def create_knowledge_base(request: Request):
 async def get_knowledge_base(id: str):
     return kb.get_one(id)
 
+@app.delete("/api/knowledge-bases/{id}/delete")
+async def delete_knowledge_base(id: str):
+    return kb.delete(id)
+
 # add a file to a knowledge base
 @app.post('/api/knowledge-bases/{id}/upload')
 async def upload_file(id: str, file: UploadFile=File(...)):
@@ -72,7 +76,7 @@ async def post_query(body: pq.QueryBody):
 @app.get('/api/history')
 async def get_evals():
     data = evals.get_chat_history()
-    return {"table_data": data}
+    return data
 
 if __name__ == "__main__":
     import uvicorn