Refactor vectorization process by removing openai_vectorize calls and updating vectorizer initialization

controllers/utils.py

@@ -22,7 +22,7 @@ import PyPDF2
 from transformers import pipeline
 
 from controllers.summarizer import summarize
-from controllers.vectorizer import vectorize, openai_vectorize
+from controllers.vectorizer import vectorize
 
 load_dotenv()
 
@@ -703,6 +703,6 @@ def crawl_by_url(url, article):
     article['referenceid'] = None
     update_content(article)
     vectorize(article)
-    openai_vectorize(article)
+    # openai_vectorize(article)
 
 data = download_files_from_s3('data')

controllers/vectorizer.py

@@ -1,7 +1,7 @@
 """Module to upsert data into AstraDB"""
 import os
 import logging
-import uuid
+# import uuid
 import time
 
 import tiktoken
@@ -24,13 +24,13 @@ embedding = AzureOpenAIEmbeddings(
     api_version="2024-07-01-preview",
     azure_endpoint="https://openai-oe.openai.azure.com/")
 
-vstore = AstraDBVectorStore(embedding=embedding,
-                            namespace="default_keyspace",
-                            collection_name="FinFast_China",
-                            token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
-                            api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"])
+# vstore = AstraDBVectorStore(embedding=embedding,
+#                             namespace="default_keyspace",
+#                             collection_name="FinFast_China",
+#                             token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
+#                             api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"])
 
-openai_vstore = AstraDBVectorStore(
+vstore = AstraDBVectorStore(
     collection_vector_service_options=CollectionVectorServiceOptions(
         provider="azureOpenAI",
         model_name="text-embedding-3-small",
@@ -43,7 +43,7 @@ openai_vstore = AstraDBVectorStore(
         },
     ),
     namespace="default_keyspace",
-    collection_name="text_embedding_3_small",
+    collection_name="article",
     token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
     api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"])
 
@@ -74,7 +74,7 @@ def add_documents_with_retry(chunks, ids, max_retries=3):
     """
     for attempt in range(max_retries):
         try:
-            openai_vstore.add_documents(chunks, ids=ids)
+            vstore.add_documents(chunks, ids=ids)
         except (ConnectionError, TimeoutError) as e:
             logging.info("Attempt %d failed: %s", attempt + 1, e)
             if attempt < max_retries - 1:
@@ -82,8 +82,9 @@ def add_documents_with_retry(chunks, ids, max_retries=3):
             else:
                 logging.error("Max retries reached. Operation failed.")
                 logging.error(ids)
+                print(ids)
 
-def openai_vectorize(article):
+def vectorize(article):
     """
     Process the given article.
 
@@ -118,37 +119,37 @@ def openai_vectorize(article):
     except (ConnectionError, TimeoutError, ValueError) as e:
         logging.error("Failed to add documents: %s", e)
 
-def vectorize(article):
-    """
-    Process the given article.
-
-    Parameters:
-    article (DataFrame): The article to be processed.
-
-    Returns:
-    None
-    """
-    article['id'] = str(article['id'])
-    if isinstance(article, dict):
-        article = [article]  # Convert single dictionary to list of dictionaries
-    df = pd.DataFrame(article)
-    df = df[['id','site','title','titleCN','category','author','content',
-             'publishDate','link']]
-    df['publishDate'] = pd.to_datetime(df['publishDate'])
-    loader = DataFrameLoader(df, page_content_column="content")
-    documents = loader.load()
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=800,
-        chunk_overlap=20,
-        length_function=len,
-        is_separator_regex=False,
-    )
-
-    chunks = text_splitter.split_documents(documents)
-    ids = []
-    for chunk in chunks:
-        _id = f"{chunk.metadata['id']}-{str(uuid.uuid5(uuid.NAMESPACE_OID,chunk.page_content))}"
-        ids.append(_id)
-    inserted_ids = vstore.add_documents(chunks, ids=ids)
-    print(inserted_ids)
-    logging.info(inserted_ids)
+# def vectorize(article):
+#     """
+#     Process the given article.
+
+#     Parameters:
+#     article (DataFrame): The article to be processed.
+
+#     Returns:
+#     None
+#     """
+#     article['id'] = str(article['id'])
+#     if isinstance(article, dict):
+#         article = [article]  # Convert single dictionary to list of dictionaries
+#     df = pd.DataFrame(article)
+#     df = df[['id','site','title','titleCN','category','author','content',
+#              'publishDate','link']]
+#     df['publishDate'] = pd.to_datetime(df['publishDate'])
+#     loader = DataFrameLoader(df, page_content_column="content")
+#     documents = loader.load()
+#     text_splitter = RecursiveCharacterTextSplitter(
+#         chunk_size=800,
+#         chunk_overlap=20,
+#         length_function=len,
+#         is_separator_regex=False,
+#     )
+
+#     chunks = text_splitter.split_documents(documents)
+#     ids = []
+#     for chunk in chunks:
+#         _id = f"{chunk.metadata['id']}-{str(uuid.uuid5(uuid.NAMESPACE_OID,chunk.page_content))}"
+#         ids.append(_id)
+#     inserted_ids = vstore.add_documents(chunks, ids=ids)
+#     print(inserted_ids)
+#     logging.info(inserted_ids)

source/eastmoney.py

@@ -22,7 +22,7 @@ from controllers.utils import (
     translate,
     update_content
 )
-from controllers.vectorizer import openai_vectorize, vectorize
+from controllers.vectorizer import vectorize
 
 with open('xpath.json', 'r', encoding='UTF-8') as f:
     xpath_dict = json.load(f)
@@ -98,7 +98,7 @@ def _crawl(url, article, retries=3):
         article['referenceid'] = reference_id
     update_content(article)
     vectorize(article)
-    openai_vectorize(article)
+    # openai_vectorize(article)
 
 @task(name = "Data Collection - eastmoney", log_prints = True)
 def crawl(delta):