Update streamlit_app.py

streamlit_app.py

@@ -20,8 +20,13 @@ import chromadb
 def list_of_collections():
     client = chromadb.Client(CHROMA_SETTINGS)
     return (client.list_collections())
-    
-def main():
+
+def get_collection_names():
+
+    collections = list_of_collections()
+    return [collection.name for collection in collections]
+
+if __name__ == '__main__':
     st.title("PrivateGPT App: Document Embedding and Retrieval")
     
     # Document upload section
@@ -39,53 +44,3 @@ def main():
         query = st.text_input("Query")
         if st.button("Retrieve"):
             retrieve_documents(query, selected_collection)
-
-# def embed_documents(files:List[st.runtime.uploaded_file_manager.UploadedFile], collection_name:str):
-#     endpoint = f"{API_BASE_URL}/embed"
-#     files_data = [("files", file) for file in files]
-#     data = {"collection_name": collection_name}
-
-#     response = requests.post(endpoint, files=files_data, data=data)
-#     if response.status_code == 200:
-#         st.success("Documents embedded successfully!")
-#     else:
-#         st.error("Document embedding failed.")
-#         st.write(response.text)
-
-
-def get_collection_names():
-
-    collections = list_of_collections()
-    return [collection.name for collection in collections]
-
-
-
-# def retrieve_documents(query: str, collection_name: str):
-#     endpoint = f"{API_BASE_URL}/retrieve"
-#     data = {"query": query, "collection_name": collection_name}
-
-#     # Modify socket options for the HTTPConnection class
-#     HTTPConnection.default_socket_options = (
-#         HTTPConnection.default_socket_options + [
-#             (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
-#             (socket.SOL_TCP, socket.TCP_KEEPIDLE, 45),
-#             (socket.SOL_TCP, socket.TCP_KEEPINTVL, 10),
-#             (socket.SOL_TCP, socket.TCP_KEEPCNT, 6)
-#         ]
-#     )
-    
-#     response = requests.post(endpoint, params=data)
-#     if response.status_code == 200:
-#         result = response.json()
-#         st.subheader("Results")
-#         st.text(result["results"])
-        
-#         st.subheader("Documents")
-#         for doc in result["docs"]:
-#             st.text(doc)
-#     else:
-#         st.error("Failed to retrieve documents.")
-#         st.write(response.text)
-
-if __name__ == '__main__':
-    main()