This is an automated email from the ASF dual-hosted git repository.
rahulvats pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 430f0e08d30 Fixing LLM providers tests (#59031)
430f0e08d30 is described below
commit 430f0e08d30425f950f62ebf16a8482da150ce2d
Author: Atul Singh <[email protected]>
AuthorDate: Thu Dec 4 17:34:57 2025 +0530
Fixing LLM providers tests (#59031)
Fixing LLM providers tests
---
.../system/pinecone/example_create_pod_index.py | 13 ++++-----
.../system/pinecone/example_pinecone_cohere.py | 32 ++++++++++++----------
.../system/pinecone/example_pinecone_openai.py | 10 +++++--
3 files changed, 31 insertions(+), 24 deletions(-)
diff --git
a/providers/pinecone/tests/system/pinecone/example_create_pod_index.py
b/providers/pinecone/tests/system/pinecone/example_create_pod_index.py
index 3583abb0a87..daebe225ab0 100644
--- a/providers/pinecone/tests/system/pinecone/example_create_pod_index.py
+++ b/providers/pinecone/tests/system/pinecone/example_create_pod_index.py
@@ -26,27 +26,26 @@ try:
except ImportError:
# Airflow 2 path
from airflow.decorators import task, teardown # type:
ignore[attr-defined,no-redef]
-from airflow.providers.pinecone.operators.pinecone import
CreatePodIndexOperator
+from airflow.providers.pinecone.operators.pinecone import
CreateServerlessIndexOperator
index_name = os.getenv("INDEX_NAME", "test")
with DAG(
"example_pinecone_create_pod_index",
- schedule="@once",
+ schedule=None,
start_date=datetime(2024, 1, 1),
catchup=False,
) as dag:
# [START howto_operator_create_pod_index]
# reference:
https://docs.pinecone.io/reference/api/control-plane/create_index
- create_index = CreatePodIndexOperator(
+ create_index = CreateServerlessIndexOperator(
task_id="pinecone_create_pod_index",
index_name=index_name,
dimension=3,
- replicas=1,
- shards=1,
- pods=1,
- pod_type="p1.x1",
+ cloud="aws",
+ region="us-west-2",
+ metric="cosine",
)
# [END howto_operator_create_pod_index]
diff --git
a/providers/pinecone/tests/system/pinecone/example_pinecone_cohere.py
b/providers/pinecone/tests/system/pinecone/example_pinecone_cohere.py
index 68ec1c14e5a..29d763ef46f 100644
--- a/providers/pinecone/tests/system/pinecone/example_pinecone_cohere.py
+++ b/providers/pinecone/tests/system/pinecone/example_pinecone_cohere.py
@@ -22,12 +22,16 @@ from datetime import datetime
from airflow import DAG
try:
- from airflow.sdk import setup, task, teardown
+ from airflow.sdk import task, teardown
except ImportError:
# Airflow 2 path
- from airflow.decorators import setup, task, teardown # type:
ignore[attr-defined,no-redef]
+ from airflow.decorators import task, teardown # type:
ignore[attr-defined,no-redef]
from airflow.providers.cohere.operators.embedding import
CohereEmbeddingOperator
-from airflow.providers.pinecone.operators.pinecone import
PineconeIngestOperator
+from airflow.providers.pinecone.hooks.pinecone import PineconeHook
+from airflow.providers.pinecone.operators.pinecone import (
+ CreateServerlessIndexOperator,
+ PineconeIngestOperator,
+)
index_name = os.getenv("INDEX_NAME", "example-pinecone-index")
namespace = os.getenv("NAMESPACE", "example-pinecone-index")
@@ -41,15 +45,14 @@ with DAG(
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()
- pod_spec = hook.get_pod_spec_obj()
- hook.create_index(index_name=index_name, dimension=1024, spec=pod_spec)
+ create_index = CreateServerlessIndexOperator(
+ task_id="pinecone_create_serverless_index",
+ index_name=index_name,
+ dimension=1024,
+ cloud="aws",
+ region="us-west-2",
+ metric="cosine",
+ )
embed_task = CohereEmbeddingOperator(
task_id="embed_task",
@@ -74,12 +77,11 @@ with DAG(
@teardown
@task
def delete_index():
- from airflow.providers.pinecone.hooks.pinecone import PineconeHook
-
hook = PineconeHook()
hook.delete_index(index_name=index_name)
- create_index() >> embed_task >> transformed_output >> perform_ingestion >>
delete_index()
+ create_index >> embed_task >> transformed_output >> perform_ingestion >>
delete_index()
+
from tests_common.test_utils.system_tests import get_test_run # noqa: E402
diff --git
a/providers/pinecone/tests/system/pinecone/example_pinecone_openai.py
b/providers/pinecone/tests/system/pinecone/example_pinecone_openai.py
index 593a7490f96..26458049a9c 100644
--- a/providers/pinecone/tests/system/pinecone/example_pinecone_openai.py
+++ b/providers/pinecone/tests/system/pinecone/example_pinecone_openai.py
@@ -27,7 +27,10 @@ except ImportError:
# Airflow 2 path
from airflow.decorators import task, teardown # type:
ignore[attr-defined,no-redef]
from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator
-from airflow.providers.pinecone.operators.pinecone import
CreatePodIndexOperator, PineconeIngestOperator
+from airflow.providers.pinecone.operators.pinecone import (
+ CreateServerlessIndexOperator,
+ PineconeIngestOperator,
+)
index_name = os.getenv("INDEX_NAME", "example-pinecone-index")
namespace = os.getenv("NAMESPACE", "example-pinecone-index")
@@ -79,10 +82,13 @@ with DAG(
start_date=datetime(2023, 1, 1),
catchup=False,
) as dag:
- create_index = CreatePodIndexOperator(
+ create_index = CreateServerlessIndexOperator(
task_id="create_index",
index_name=index_name,
dimension=1536,
+ cloud="aws",
+ region="us-west-2",
+ metric="cosine",
)
embed_task = OpenAIEmbeddingOperator(
