Update insert_into_milvus_db.py

insert_into_milvus_db.py

@@ -54,73 +54,82 @@ def get_secrets():
 
 
 def create_schema(uri: str, token: str, collection_name: str):
-    logger.warning('Entering create_schema()')
-    connections.connect("default", uri=uri, token=token)
-
-    if not utility.has_collection(collection_name):
-        dim = 768 # embeddings dim
-        
-        article_url = FieldSchema(name="article_url", dtype=DataType.VARCHAR, max_length=10000, 
-                                   is_primary=True, description="url of the article")
-        article_title = FieldSchema(name="article_title", dtype=DataType.VARCHAR, max_length=5000, 
-                                   is_primary=False, description="headline of the article") 
-        
-        article_src = FieldSchema(name="article_src", dtype=DataType.VARCHAR, max_length=1000, 
-                                   is_primary=False, description="src of the article")
-        
-        article_date = FieldSchema(name="article_date", dtype=DataType.VARCHAR, max_length=1000, 
-                                   is_primary=False, description="date of the article")
-        
-        article_age = FieldSchema(name="article_age", dtype=DataType.INT64,
-                                   is_primary=False, description="age of the article")
+    try:
+        logger.warning('Entering create_schema()')
+        connections.connect("default", uri=uri, token=token)
     
-        article_embed = FieldSchema(name="article_embed", dtype=DataType.FLOAT_VECTOR, dim=dim) # description embeddings
+        if not utility.has_collection(collection_name):
+            dim = 768 # embeddings dim
+            
+            article_url = FieldSchema(name="article_url", dtype=DataType.VARCHAR, max_length=10000, 
+                                       is_primary=True, description="url of the article")
+            article_title = FieldSchema(name="article_title", dtype=DataType.VARCHAR, max_length=5000, 
+                                       is_primary=False, description="headline of the article") 
             
-        schema = CollectionSchema(fields=[article_url, article_title, article_src, 
-                                          article_date, article_age, article_embed], 
-                                  auto_id=False, description="collection of news articles")
-        logger.warning("Creating the collection")
-        collection = Collection(name=collection_name, schema=schema)
-        # logger.warning(f"Schema: {schema}")
-        logger.warning("Successfully created collection")
-    else:
-        collection = Collection(name=collection_name)
-        logger.warning("Using existing collection")
-    logger.warning('Exiting create_schema()')
-    return collection
+            article_src = FieldSchema(name="article_src", dtype=DataType.VARCHAR, max_length=1000, 
+                                       is_primary=False, description="src of the article")
+            
+            article_date = FieldSchema(name="article_date", dtype=DataType.VARCHAR, max_length=1000, 
+                                       is_primary=False, description="date of the article")
+            
+            article_age = FieldSchema(name="article_age", dtype=DataType.INT64,
+                                       is_primary=False, description="age of the article")
+        
+            article_embed = FieldSchema(name="article_embed", dtype=DataType.FLOAT_VECTOR, dim=dim) # description embeddings
+                
+            schema = CollectionSchema(fields=[article_url, article_title, article_src, 
+                                              article_date, article_age, article_embed], 
+                                      auto_id=False, description="collection of news articles")
+            logger.warning("Creating the collection")
+            collection = Collection(name=collection_name, schema=schema)
+            # logger.warning(f"Schema: {schema}")
+            logger.warning("Successfully created collection")
+        else:
+            collection = Collection(name=collection_name)
+            logger.warning("Using existing collection")
+        logger.warning('Exiting create_schema()')
+        return collection
+    except:
+        raise
     
 
 def prepare_docs(vectorizer):
-    logger.warning('Entering prepare_docs()')
-    logger.warning('Retrieving latest news')
-    news_df = get_news()
-    if news_df is None:
-        raise Exception("ERROR: No latest news in retrieved")
-    logger.warning('Successfully retrieved latest news')
-    article_url = news_df['url'].tolist()
-    article_title = news_df['title'].tolist()
-    article_src = news_df['src'].tolist()
-    article_date = news_df['parsed_date'].tolist()
-    article_age = news_df['news_age'].tolist()
-    article_embed = vectorizer.vectorize_(article_title)
-    logger.warning('Exiting prepare_docs()')
-    return [article_url, article_title, article_src, 
-            article_date, article_age, article_embed]
+    try:
+        logger.warning('Entering prepare_docs()')
+        logger.warning('Retrieving latest news')
+        news_df = get_news()
+        if news_df is None:
+            raise Exception("ERROR: No latest news in retrieved")
+        logger.warning('Successfully retrieved latest news')
+        article_url = news_df['url'].tolist()
+        article_title = news_df['title'].tolist()
+        article_src = news_df['src'].tolist()
+        article_date = news_df['parsed_date'].tolist()
+        article_age = news_df['news_age'].tolist()
+        article_embed = vectorizer.vectorize_(article_title)
+        logger.warning('Exiting prepare_docs()')
+        return [article_url, article_title, article_src, 
+                article_date, article_age, article_embed]
+    except:
+        raise
     
     
 def upsert_db(vectorizer, collection):
-    logger.warning('Entering upsert_db()')
-    collection_is_empty = 0
-    if collection.is_empty:
-        collection_is_empty = 1
-        
-    docs_to_upsert = prepare_docs(vectorizer)
-    ins_resp = collection.upsert(docs_to_upsert)
-    if ins_resp.err_count != 0:
-        raise Exception(f'Milvus Insertion not successful. {ins_resp.err_count} errors reported.')
-    if collection_is_empty:
-        index_params = {"index_type": "AUTOINDEX", "metric_type": "IP", "params": {}} 
-        collection.create_index(field_name='article_embed', index_params=index_params)
-        collection.load()
-    logger.warning('Upsert successful')
-    logger.warning('Exiting upsert_db()')
+    try:
+        logger.warning('Entering upsert_db()')
+        collection_is_empty = 0
+        if collection.is_empty:
+            collection_is_empty = 1
+            
+        docs_to_upsert = prepare_docs(vectorizer)
+        ins_resp = collection.upsert(docs_to_upsert)
+        if ins_resp.err_count != 0:
+            raise Exception(f'Milvus Insertion not successful. {ins_resp.err_count} errors reported.')
+        if collection_is_empty:
+            index_params = {"index_type": "AUTOINDEX", "metric_type": "IP", "params": {}} 
+            collection.create_index(field_name='article_embed', index_params=index_params)
+            collection.load()
+        logger.warning('Upsert successful')
+        logger.warning('Exiting upsert_db()')
+    except:
+        raise