Simplify extraction: just call catllm.extract() instead of reimplementing logic

app.py

@@ -849,172 +849,37 @@ with col_input:
                     st.error(f"{provider} API key not configured")
                 else:
                     model_source = get_model_source(model)
-                    items_list = input_data if isinstance(input_data, list) else [input_data]
 
-                    # For PDF text mode, extract text and process like text corpus
-                    if input_type_selected == "pdf" and mode == "text":
-                        status_text = st.empty()
-                        status_text.text("Extracting text from PDFs...")
+                    with st.spinner("Extracting categories..."):
+                        start_time = time.time()
 
-                        # Extract text from all PDF pages
-                        extracted_texts = extract_text_from_pdfs(items_list)
+                        extract_kwargs = {
+                            'input_data': input_data,
+                            'api_key': actual_api_key,
+                            'input_type': input_type_selected,
+                            'description': description,
+                            'user_model': model,
+                            'model_source': model_source,
+                            'max_categories': int(max_categories)
+                        }
+                        if mode:
+                            extract_kwargs['mode'] = mode
 
-                        if not extracted_texts:
-                            st.error("No text could be extracted from the PDFs")
-                        else:
-                            # Now process as text corpus with random chunks
-                            import random
-                            random.shuffle(extracted_texts)  # Randomize order
-
-                            num_items = len(extracted_texts)
-                            divisions = max(1, num_items // 15)
-                            divisions = min(divisions, 5)
-                            chunk_size_calc = num_items // max(1, divisions)
-                            categories_per_chunk = min(10, chunk_size_calc - 1) if chunk_size_calc > 1 else 5
-
-                            # Progress UI
-                            progress_bar = st.progress(0)
-                            start_time = time.time()
-
-                            all_categories = []
-                            chunk_size = len(extracted_texts) // divisions if divisions > 0 else len(extracted_texts)
-
-                            for div_idx in range(divisions):
-                                progress = div_idx / divisions
-                                progress_bar.progress(progress)
-
-                                elapsed = time.time() - start_time
-                                if div_idx > 0:
-                                    avg_time = elapsed / div_idx
-                                    eta_seconds = avg_time * (divisions - div_idx)
-                                    eta_str = f" | ETA: {eta_seconds:.0f}s" if eta_seconds < 60 else f" | ETA: {eta_seconds/60:.1f}m"
-                                else:
-                                    eta_str = ""
-
-                                status_text.text(f"Analyzing text chunk {div_idx + 1} of {divisions} ({progress*100:.0f}%){eta_str}")
-
-                                start_idx = div_idx * chunk_size
-                                end_idx = start_idx + chunk_size if div_idx < divisions - 1 else len(extracted_texts)
-                                chunk_data = extracted_texts[start_idx:end_idx]
-
-                                if not chunk_data:
-                                    continue
-
-                                try:
-                                    # Process as text type since we extracted the text
-                                    extract_kwargs = {
-                                        'input_data': chunk_data,
-                                        'api_key': actual_api_key,
-                                        'input_type': 'text',  # Treat as text now
-                                        'description': description,
-                                        'user_model': model,
-                                        'model_source': model_source,
-                                        'divisions': 1,
-                                        'categories_per_chunk': categories_per_chunk,
-                                        'max_categories': int(max_categories)
-                                    }
-
-                                    chunk_result = catllm.extract(**extract_kwargs)
-                                    chunk_cats = chunk_result.get('top_categories', [])
-                                    all_categories.extend(chunk_cats)
-                                except Exception as e:
-                                    st.warning(f"Error on chunk {div_idx + 1}: {str(e)}")
-                                    continue
-
-                            progress_bar.progress(1.0)
-                            processing_time = time.time() - start_time
-                            status_text.text(f"Extraction complete in {processing_time:.1f}s")
+                        try:
+                            extract_result = catllm.extract(**extract_kwargs)
+                            categories = extract_result.get('top_categories', [])
 
-                            if all_categories:
-                                from collections import Counter
-                                cat_counts = Counter(all_categories)
-                                top_cats = [cat for cat, _ in cat_counts.most_common(int(max_categories))]
+                            processing_time = time.time() - start_time
 
-                                st.session_state.extracted_categories = top_cats
+                            if categories:
+                                st.session_state.extracted_categories = categories
                                 st.session_state.task_mode = "manual"
-                                st.success(f"Extracted {len(top_cats)} categories from {len(extracted_texts)} pages in {processing_time:.1f}s")
+                                st.success(f"Extracted {len(categories)} categories in {processing_time:.1f}s")
                                 st.rerun()
                             else:
                                 st.error("No categories were extracted")
-                    else:
-                        # Original logic for images, PDFs in image/both mode, and text
-                        num_items = len(items_list)
-
-                        if input_type_selected == "image":
-                            divisions = min(3, max(1, num_items // 5))
-                            categories_per_chunk = 12
-                        else:
-                            divisions = max(1, num_items // 15)
-                            divisions = min(divisions, 5)
-                            chunk_size_calc = num_items // max(1, divisions)
-                            categories_per_chunk = min(10, chunk_size_calc - 1) if chunk_size_calc > 1 else 5
-
-                        # Progress UI
-                        progress_bar = st.progress(0)
-                        status_text = st.empty()
-                        start_time = time.time()
-
-                        all_categories = []
-                        chunk_size = len(items_list) // divisions if divisions > 0 else len(items_list)
-
-                        for div_idx in range(divisions):
-                            progress = div_idx / divisions
-                            progress_bar.progress(progress)
-
-                            elapsed = time.time() - start_time
-                            if div_idx > 0:
-                                avg_time = elapsed / div_idx
-                                eta_seconds = avg_time * (divisions - div_idx)
-                                eta_str = f" | ETA: {eta_seconds:.0f}s" if eta_seconds < 60 else f" | ETA: {eta_seconds/60:.1f}m"
-                            else:
-                                eta_str = ""
-
-                            status_text.text(f"Analyzing chunk {div_idx + 1} of {divisions} ({progress*100:.0f}%){eta_str}")
-
-                            start_idx = div_idx * chunk_size
-                            end_idx = start_idx + chunk_size if div_idx < divisions - 1 else len(items_list)
-                            chunk_data = items_list[start_idx:end_idx]
-
-                            if not chunk_data:
-                                continue
-
-                            try:
-                                extract_kwargs = {
-                                    'input_data': chunk_data,
-                                    'api_key': actual_api_key,
-                                    'input_type': input_type_selected,
-                                    'description': description,
-                                    'user_model': model,
-                                    'model_source': model_source,
-                                    'divisions': 1,
-                                    'categories_per_chunk': categories_per_chunk,
-                                    'max_categories': int(max_categories)
-                                }
-                                if mode:
-                                    extract_kwargs['mode'] = mode
-
-                                chunk_result = catllm.extract(**extract_kwargs)
-                                chunk_cats = chunk_result.get('top_categories', [])
-                                all_categories.extend(chunk_cats)
-                            except Exception as e:
-                                st.warning(f"Error on chunk {div_idx + 1}: {str(e)}")
-                                continue
-
-                        progress_bar.progress(1.0)
-                        processing_time = time.time() - start_time
-                        status_text.text(f"Extraction complete in {processing_time:.1f}s")
-
-                        if all_categories:
-                            from collections import Counter
-                            cat_counts = Counter(all_categories)
-                            top_cats = [cat for cat, _ in cat_counts.most_common(int(max_categories))]
-
-                            st.session_state.extracted_categories = top_cats
-                            st.session_state.task_mode = "manual"
-                            st.success(f"Extracted {len(top_cats)} categories in {processing_time:.1f}s")
-                            st.rerun()
-                        else:
-                            st.error("No categories were extracted")
+                        except Exception as e:
+                            st.error(f"Error: {str(e)}")
 
     # Category inputs (shown for manual mode or after extraction)
     if st.session_state.task_mode == "manual":