47
47
48
48
from langchain_vectorize .retrievers import VectorizeRetriever
49
49
50
+ logger = logging .getLogger (__name__ )
51
+
50
52
51
53
@pytest .fixture (scope = "session" )
52
54
def api_token () -> str :
@@ -69,7 +71,7 @@ def org_id() -> str:
69
71
@pytest .fixture (scope = "session" )
70
72
def environment () -> Literal ["prod" , "dev" , "local" , "staging" ]:
71
73
env = os .getenv ("VECTORIZE_ENV" , "prod" )
72
- if env not in [ "prod" , "dev" , "local" , "staging" ] :
74
+ if env not in { "prod" , "dev" , "local" , "staging" } :
73
75
msg = "Invalid VECTORIZE_ENV environment variable."
74
76
raise ValueError (msg )
75
77
return env # type: ignore[return-value]
@@ -108,7 +110,7 @@ def pipeline_id(api_client: ApiClient, org_id: str) -> Iterator[str]:
108
110
CreateSourceConnectorRequest (FileUpload (name = "from api" , type = "FILE_UPLOAD" )),
109
111
)
110
112
source_connector_id = response .connector .id
111
- logging .info ("Created source connector %s" , source_connector_id )
113
+ logger .info ("Created source connector %s" , source_connector_id )
112
114
113
115
uploads_api = UploadsApi (api_client )
114
116
upload_response = uploads_api .start_file_upload_to_connector (
@@ -139,23 +141,23 @@ def pipeline_id(api_client: ApiClient, org_id: str) -> Iterator[str]:
139
141
msg = "Upload failed:"
140
142
raise ValueError (msg )
141
143
else :
142
- logging .info ("Upload successful" )
144
+ logger .info ("Upload successful" )
143
145
144
146
ai_platforms = AIPlatformConnectorsApi (api_client ).get_ai_platform_connectors (
145
147
org_id
146
148
)
147
149
builtin_ai_platform = next (
148
150
c .id for c in ai_platforms .ai_platform_connectors if c .type == "VECTORIZE"
149
151
)
150
- logging .info ("Using AI platform %s" , builtin_ai_platform )
152
+ logger .info ("Using AI platform %s" , builtin_ai_platform )
151
153
152
154
vector_databases = DestinationConnectorsApi (api_client ).get_destination_connectors (
153
155
org_id
154
156
)
155
157
builtin_vector_db = next (
156
158
c .id for c in vector_databases .destination_connectors if c .type == "VECTORIZE"
157
159
)
158
- logging .info ("Using destination connector %s" , builtin_vector_db )
160
+ logger .info ("Using destination connector %s" , builtin_vector_db )
159
161
160
162
pipeline_response = pipelines .create_pipeline (
161
163
org_id ,
@@ -182,14 +184,14 @@ def pipeline_id(api_client: ApiClient, org_id: str) -> Iterator[str]:
182
184
),
183
185
)
184
186
pipeline_id = pipeline_response .data .id
185
- logging .info ("Created pipeline %s" , pipeline_id )
187
+ logger .info ("Created pipeline %s" , pipeline_id )
186
188
187
189
yield pipeline_id
188
190
189
191
try :
190
192
pipelines .delete_pipeline (org_id , pipeline_id )
191
193
except Exception :
192
- logging .exception ("Failed to delete pipeline %s" , pipeline_id )
194
+ logger .exception ("Failed to delete pipeline %s" , pipeline_id )
193
195
194
196
195
197
def test_retrieve_init_args (
0 commit comments