1313
1414from langchain_vectorize .retrievers import VectorizeRetriever
1515
16+ logger = logging .getLogger (__name__ )
17+
1618
1719@pytest .fixture (scope = "session" )
1820def api_token () -> str :
@@ -35,7 +37,7 @@ def org_id() -> str:
3537@pytest .fixture (scope = "session" )
3638def environment () -> Literal ["prod" , "dev" , "local" , "staging" ]:
3739 env = os .getenv ("VECTORIZE_ENV" , "prod" )
38- if env not in [ "prod" , "dev" , "local" , "staging" ] :
40+ if env not in { "prod" , "dev" , "local" , "staging" } :
3941 msg = "Invalid VECTORIZE_ENV environment variable."
4042 raise ValueError (msg )
4143 return env
@@ -76,7 +78,7 @@ def pipeline_id(api_client: v.ApiClient, org_id: str) -> Iterator[str]:
7678 ),
7779 )
7880 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 )
8082
8183 uploads_api = v .UploadsApi (api_client )
8284 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]:
107109 msg = "Upload failed:"
108110 raise ValueError (msg )
109111 else :
110- logging .info ("Upload successful" )
112+ logger .info ("Upload successful" )
111113
112114 ai_platforms = v .AIPlatformConnectorsApi (api_client ).get_ai_platform_connectors (
113115 org_id
114116 )
115117 builtin_ai_platform = next (
116118 c .id for c in ai_platforms .ai_platform_connectors if c .type == "VECTORIZE"
117119 )
118- logging .info ("Using AI platform %s" , builtin_ai_platform )
120+ logger .info ("Using AI platform %s" , builtin_ai_platform )
119121
120122 vector_databases = v .DestinationConnectorsApi (
121123 api_client
122124 ).get_destination_connectors (org_id )
123125 builtin_vector_db = next (
124126 c .id for c in vector_databases .destination_connectors if c .type == "VECTORIZE"
125127 )
126- logging .info ("Using destination connector %s" , builtin_vector_db )
128+ logger .info ("Using destination connector %s" , builtin_vector_db )
127129
128130 pipeline_response = pipelines .create_pipeline (
129131 org_id ,
@@ -150,14 +152,14 @@ def pipeline_id(api_client: v.ApiClient, org_id: str) -> Iterator[str]:
150152 ),
151153 )
152154 pipeline_id = pipeline_response .data .id
153- logging .info ("Created pipeline %s" , pipeline_id )
155+ logger .info ("Created pipeline %s" , pipeline_id )
154156
155157 yield pipeline_id
156158
157159 try :
158160 pipelines .delete_pipeline (org_id , pipeline_id )
159161 except Exception :
160- logging .exception ("Failed to delete pipeline %s" , pipeline_id )
162+ logger .exception ("Failed to delete pipeline %s" , pipeline_id )
161163
162164
163165def test_retrieve_init_args (
0 commit comments