Fix the pinecone system test and add default value for metric

apache · May 2, 2024 · d175973 · d175973
1 parent 1d504ca
commit d175973
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 27 deletions.
diff --git a/airflow/providers/pinecone/operators/pinecone.py b/airflow/providers/pinecone/operators/pinecone.py
@@ -119,7 +119,7 @@ def __init__(
         pod_type: str | None = None,
         metadata_config: dict | None = None,
         source_collection: str | None = None,
-        metric: str | None = None,
+        metric: str | None = "cosine",
         timeout: int | None = None,
         **kwargs: Any,
     ):

diff --git a/tests/system/providers/pinecone/example_pinecone_cohere.py b/tests/system/providers/pinecone/example_pinecone_cohere.py
@@ -17,13 +17,12 @@
 from __future__ import annotations
 
 import os
-import time
 from datetime import datetime
 
 from airflow import DAG
-from airflow.decorators import setup, task, teardown
+from airflow.decorators import task, teardown
 from airflow.providers.cohere.operators.embedding import CohereEmbeddingOperator
-from airflow.providers.pinecone.operators.pinecone import PineconeIngestOperator
+from airflow.providers.pinecone.operators.pinecone import CreatePodIndexOperator, PineconeIngestOperator
 
 index_name = os.getenv("INDEX_NAME", "example-pinecone-index")
 namespace = os.getenv("NAMESPACE", "example-pinecone-index")
@@ -37,15 +36,15 @@
     start_date=datetime(2023, 1, 1),
     catchup=False,
 ) as dag:
-
-    @setup
-    @task
-    def create_index():
-        from airflow.providers.pinecone.hooks.pinecone import PineconeHook
-
-        hook = PineconeHook()
-        hook.create_index(index_name=index_name, dimension=768)
-        time.sleep(60)
+    create_index = CreatePodIndexOperator(
+        task_id="create_index",
+        index_name=index_name,
+        dimension=768,
+        replicas=1,
+        shards=1,
+        pods=1,
+        pod_type="p1.x1",
+    )
 
     embed_task = CohereEmbeddingOperator(
         task_id="embed_task",
@@ -70,7 +69,7 @@ def delete_index():
         hook = PineconeHook()
         hook.delete_index(index_name=index_name)
 
-    create_index() >> embed_task >> perform_ingestion >> delete_index()
+    create_index >> embed_task >> perform_ingestion >> delete_index()
 
 from tests.system.utils import get_test_run  # noqa: E402
 

diff --git a/tests/system/providers/pinecone/example_pinecone_openai.py b/tests/system/providers/pinecone/example_pinecone_openai.py
@@ -17,13 +17,12 @@
 from __future__ import annotations
 
 import os
-import time
 from datetime import datetime
 
 from airflow import DAG
-from airflow.decorators import setup, task, teardown
+from airflow.decorators import task, teardown
 from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator
-from airflow.providers.pinecone.operators.pinecone import PineconeIngestOperator
+from airflow.providers.pinecone.operators.pinecone import CreatePodIndexOperator, PineconeIngestOperator
 
 index_name = os.getenv("INDEX_NAME", "example-pinecone-index")
 namespace = os.getenv("NAMESPACE", "example-pinecone-index")
@@ -75,15 +74,15 @@
     start_date=datetime(2023, 1, 1),
     catchup=False,
 ) as dag:
-
-    @setup
-    @task
-    def create_index():
-        from airflow.providers.pinecone.hooks.pinecone import PineconeHook
-
-        hook = PineconeHook()
-        hook.create_index(index_name=index_name, dimension=1536)
-        time.sleep(60)
+    create_index = CreatePodIndexOperator(
+        task_id="create_index",
+        index_name=index_name,
+        dimension=1536,
+        replicas=1,
+        shards=1,
+        pods=1,
+        pod_type="p1.x1",
+    )
 
     embed_task = OpenAIEmbeddingOperator(
         task_id="embed_task",
@@ -110,7 +109,7 @@ def delete_index():
         hook = PineconeHook()
         hook.delete_index(index_name=index_name)
 
-    create_index() >> embed_task >> perform_ingestion >> delete_index()
+    create_index >> embed_task >> perform_ingestion >> delete_index()
 
 from tests.system.utils import get_test_run  # noqa: E402