add scripts and other minor changes to make it easier to use custom documents with Chat REPL (#7)

backend/app/chat/engine.py

@@ -275,6 +275,7 @@ async def get_chat_engine(
     api_query_engine_tools = [
         get_api_query_engine_tool(doc, service_context)
         for doc in conversation.documents
+        if DocumentMetadataKeysEnum.SEC_DOCUMENT in doc.metadata_map
     ]
 
     quantitative_question_engine = SubQuestionQueryEngine.from_defaults(

backend/scripts/chat_llama.py

@@ -45,6 +45,14 @@ class DocumentPickerCmd(cmd.Cmd):
         except ValueError:
             print("Invalid index. Please enter a number.")
 
+    def do_select_id(self, document_id):
+        "Select a document by it's ID"
+        if not document_id:
+            print("Please enter a valid document ID")
+        else:
+            self.selected_documents.append({"id": document_id})
+            print(f"Selected document ID {document_id}")
+
     def do_finish(self, args):
         "Finish the document selection process: FINISH"
         if len(self.selected_documents) > 0:
@@ -141,14 +149,7 @@ class ConversationCmd(cmd.Cmd):
 
     def do_quit(self, args):
         "Quits the program."
-        try:
-            if self.conversation_id is not None:
-                print("Quitting but first deleting convo.")
-                self.do_delete("")
-            else:
-                print("Quitting.")
-        except:
-            print("failed to delete")
+        print("Quitting.")
         raise SystemExit
 
 

backend/scripts/upsert_document.py

+from fire import Fire
+from app.schema import Document
+from app.db.session import SessionLocal
+from app.api import crud
+import asyncio
+
+async def upsert_single_document(doc_url: str):
+    """
+    Upserts a single SEC document into the database using its URL.
+    """
+    if not doc_url or not doc_url.startswith('http'):
+        print("DOC_URL must be an http(s) based url value")
+        return
+    metadata_map = {}
+    doc = Document(url=doc_url, metadata_map=metadata_map)
+
+    async with SessionLocal() as db:
+        document = await crud.upsert_document_by_url(db, doc)
+        print(f"Upserted document. Database ID:\n{document.id}")
+
+
+def main_upsert_single_document(doc_url: str):
+    """
+    Script to upsert a single document by URL. metada_map parameter will be empty dict ({})
+    This script is useful when trying to use your own PDF files.
+    """
+    asyncio.run(upsert_single_document(doc_url))
+
+if __name__ == "__main__":
+    Fire(main_upsert_single_document)