@@ -130,6 +130,44 @@ async def test_classify_file_ids(
130
130
assert item .result .type == expected_type
131
131
132
132
133
+ @pytest .mark .asyncio
134
+ async def test_classify_file_ids_from_api_key (
135
+ e2e_test_settings : EndToEndTestSettings ,
136
+ file_client : FileClient ,
137
+ simple_pdf_file_path : str ,
138
+ research_paper_path : str ,
139
+ classification_rules : list [ClassifierRule ],
140
+ ):
141
+ """Test classifying files by their IDs"""
142
+ # Upload test files first to get their IDs
143
+ pdf_file = await file_client .upload_file (simple_pdf_file_path )
144
+ research_paper_file = await file_client .upload_file (research_paper_path )
145
+
146
+ classify_client = ClassifyClient .from_api_key (
147
+ api_key = e2e_test_settings .LLAMA_CLOUD_API_KEY .get_secret_value (),
148
+ base_url = e2e_test_settings .LLAMA_CLOUD_BASE_URL ,
149
+ project_id = pdf_file .project_id ,
150
+ organization_id = e2e_test_settings .LLAMA_CLOUD_ORGANIZATION_ID ,
151
+ )
152
+
153
+ # Classify the uploaded files
154
+ results = await classify_client .aclassify_file_ids (
155
+ rules = classification_rules , file_ids = [pdf_file .id , research_paper_file .id ]
156
+ )
157
+
158
+ assert isinstance (results , ClassifyJobResults )
159
+ assert len (results .items ) == 2
160
+
161
+ file_id_to_expected_type = {
162
+ pdf_file .id : "number" ,
163
+ research_paper_file .id : "research_paper" ,
164
+ }
165
+ # Verify each file got classified
166
+ for item in results .items :
167
+ expected_type = file_id_to_expected_type [item .file_id ]
168
+ assert item .result .type == expected_type
169
+
170
+
133
171
@parameterize_sync_and_async
134
172
@pytest .mark .asyncio
135
173
async def test_classify_file_path (
0 commit comments