13
13
14
14
from langchain_vectorize .retrievers import VectorizeRetriever
15
15
16
+ logger = logging .getLogger (__name__ )
17
+
16
18
17
19
@pytest .fixture (scope = "session" )
18
20
def api_token () -> str :
@@ -35,7 +37,7 @@ def org_id() -> str:
35
37
@pytest .fixture (scope = "session" )
36
38
def environment () -> Literal ["prod" , "dev" , "local" , "staging" ]:
37
39
env = os .getenv ("VECTORIZE_ENV" , "prod" )
38
- if env not in [ "prod" , "dev" , "local" , "staging" ] :
40
+ if env not in { "prod" , "dev" , "local" , "staging" } :
39
41
msg = "Invalid VECTORIZE_ENV environment variable."
40
42
raise ValueError (msg )
41
43
return env
@@ -76,7 +78,7 @@ def pipeline_id(api_client: v.ApiClient, org_id: str) -> Iterator[str]:
76
78
),
77
79
)
78
80
source_connector_id = response .connector .id
79
- logging .info ("Created source connector %s" , source_connector_id )
81
+ logger .info ("Created source connector %s" , source_connector_id )
80
82
81
83
uploads_api = v .UploadsApi (api_client )
82
84
upload_response = uploads_api .start_file_upload_to_connector (
@@ -107,23 +109,23 @@ def pipeline_id(api_client: v.ApiClient, org_id: str) -> Iterator[str]:
107
109
msg = "Upload failed:"
108
110
raise ValueError (msg )
109
111
else :
110
- logging .info ("Upload successful" )
112
+ logger .info ("Upload successful" )
111
113
112
114
ai_platforms = v .AIPlatformConnectorsApi (api_client ).get_ai_platform_connectors (
113
115
org_id
114
116
)
115
117
builtin_ai_platform = next (
116
118
c .id for c in ai_platforms .ai_platform_connectors if c .type == "VECTORIZE"
117
119
)
118
- logging .info ("Using AI platform %s" , builtin_ai_platform )
120
+ logger .info ("Using AI platform %s" , builtin_ai_platform )
119
121
120
122
vector_databases = v .DestinationConnectorsApi (
121
123
api_client
122
124
).get_destination_connectors (org_id )
123
125
builtin_vector_db = next (
124
126
c .id for c in vector_databases .destination_connectors if c .type == "VECTORIZE"
125
127
)
126
- logging .info ("Using destination connector %s" , builtin_vector_db )
128
+ logger .info ("Using destination connector %s" , builtin_vector_db )
127
129
128
130
pipeline_response = pipelines .create_pipeline (
129
131
org_id ,
@@ -150,14 +152,14 @@ def pipeline_id(api_client: v.ApiClient, org_id: str) -> Iterator[str]:
150
152
),
151
153
)
152
154
pipeline_id = pipeline_response .data .id
153
- logging .info ("Created pipeline %s" , pipeline_id )
155
+ logger .info ("Created pipeline %s" , pipeline_id )
154
156
155
157
yield pipeline_id
156
158
157
159
try :
158
160
pipelines .delete_pipeline (org_id , pipeline_id )
159
161
except Exception :
160
- logging .exception ("Failed to delete pipeline %s" , pipeline_id )
162
+ logger .exception ("Failed to delete pipeline %s" , pipeline_id )
161
163
162
164
163
165
def test_retrieve_init_args (
0 commit comments