Refactor vectorizer to create collection in AstraDB and improve logging

controllers/vectorizer.py

@@ -9,36 +9,72 @@ from langchain_astradb import AstraDBVectorStore
 from langchain_openai import AzureOpenAIEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import DataFrameLoader
+from astrapy import DataAPIClient
 from astrapy.info import CollectionVectorServiceOptions
+from astrapy.exceptions import CollectionAlreadyExistsException
 
 logging.basicConfig(
     format='%(asctime)s - %(levelname)s - %(funcName)s - %(message)s',
     datefmt="%Y-%m-%d %H:%M:%S",
     level=logging.ERROR)
 
+# from astrapy import AstraClient
 ASTRA_DB_APPLICATION_TOKEN = os.environ['ASTRA_DB_APPLICATION_TOKEN']
 ASTRA_DB_API_ENDPOINT = os.environ['ASTRA_DB_API_ENDPOINT']
+COLLECTION_NAME = "article"
+VECTOR_OPTIONS = CollectionVectorServiceOptions(
+    provider="azureOpenAI",
+    model_name="text-embedding-3-small",
+    authentication={"providerKey": "AZURE_OPENAI_API_KEY"},
+    parameters={
+        "resourceName": "openai-oe",
+        "deploymentId": "text-embedding-3-small",
+    },
+)
+
+client = DataAPIClient(token=ASTRA_DB_APPLICATION_TOKEN)
+database = client.get_database(ASTRA_DB_API_ENDPOINT)
 
 embedding = AzureOpenAIEmbeddings(
     api_version="2024-07-01-preview",
     azure_endpoint="https://openai-oe.openai.azure.com/")
 
+try:
+    # Try to create the collection
+    database.create_collection(
+        COLLECTION_NAME,
+        dimension=1536,  # Default dimension for text-embedding-3-small
+        metric="cosine",
+        service=VECTOR_OPTIONS
+    )
+    logging.info("Created new collection '%s'", COLLECTION_NAME)
+except CollectionAlreadyExistsException:
+    logging.info("Collection '%s' already exists. Verifying settings...", COLLECTION_NAME)
+    collection = database.get_collection(COLLECTION_NAME)
+
 vstore = AstraDBVectorStore(
-    collection_vector_service_options=CollectionVectorServiceOptions(
-        provider="azureOpenAI",
-        model_name="text-embedding-3-small",
-        authentication={
-            "providerKey": "AZURE_OPENAI_API_KEY",
-        },
-        parameters={
-            "resourceName": "openai-oe",
-            "deploymentId": "text-embedding-3-small",
-        },
-    ),
+    collection_name=COLLECTION_NAME,
     namespace="default_keyspace",
-    collection_name="article",
-    token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
-    api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"])
+    embedding=embedding,
+    token=ASTRA_DB_APPLICATION_TOKEN,
+    api_endpoint=ASTRA_DB_API_ENDPOINT)
+
+# vstore = AstraDBVectorStore(
+#     collection_vector_service_options=CollectionVectorServiceOptions(
+#         provider="azureOpenAI",
+#         model_name="text-embedding-3-small",
+#         authentication={
+#             "providerKey": "AZURE_OPENAI_API_KEY",
+#         },
+#         parameters={
+#             "resourceName": "openai-oe",
+#             "deploymentId": "text-embedding-3-small",
+#         },
+#     ),
+#     namespace="default_keyspace",
+#     collection_name="article",
+#     token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
+#     api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"])
 
 def token_length(text):
     """