Add AUTO routing for priority=cost

Author: dilithjay

lexoid/api.py

@@ -37,6 +37,7 @@
     download_file,
     get_file_type,
     get_webpage_soup,
+    has_image_in_pdf,
     is_supported_file_type,
     is_supported_url_file_type,
     recursive_read_html,
@@ -62,6 +63,23 @@ def wrapper(*args, **kwargs):
                 if args[1] == ParserType.AUTO:
                     router_priority = kwargs.get("router_priority", "speed")
                     autoselect_llm = kwargs.get("autoselect_llm", False)
+                    if router_priority == "cost" and has_image_in_pdf(kwargs["path"]):
+                        # Handling this outside of router to allow for multiple func calls
+                        kwargs["parser_type"] = ParserType.STATIC_PARSE
+                        kwargs["framework"] = "paddleocr"
+                        result = func(**kwargs)
+                        character_threshold = kwargs.get("character_threshold", 100)
+                        len_result = len(result["raw"].strip())
+                        if len_result < character_threshold:
+                            logger.debug(
+                                f"Low character count detected ({len_result} < {character_threshold}), returning result"
+                            )
+                            return result
+                        logger.debug(
+                            f"Character count above threshold ({len_result} >= {character_threshold}), switching to LLM_PARSE"
+                        )
+                        kwargs["parser_type"] = ParserType.LLM_PARSE
+                        return func(**kwargs)
                     routed_parser_type, model = router(
                         kwargs["path"], router_priority, autoselect_llm=autoselect_llm
                     )
@@ -173,15 +191,15 @@ def parse_chunk_list(
     """
     combined_segments = []
     raw_texts = []
+    parsers_used = []
     token_usage = {"input": 0, "output": 0, "llm_page_count": 0}
     for file_path in file_paths:
         result = parse_chunk(file_path, parser_type, **kwargs)
         combined_segments.extend(result["segments"])
         raw_texts.append(result["raw"])
-        if (
-            result.get("parser_used") == ParserType.LLM_PARSE
-            and "token_usage" in result
-        ):
+        parser_used = result.get("parser_used")
+        parsers_used.append(parser_used.value if parser_used else "UNKNOWN")
+        if parser_used == ParserType.LLM_PARSE and "token_usage" in result:
             token_usage["input"] += result["token_usage"]["input"]
             token_usage["output"] += result["token_usage"]["output"]
             token_usage["llm_page_count"] += len(result["segments"])
@@ -195,6 +213,7 @@ def parse_chunk_list(
         "parent_title": kwargs.get("parent_title", ""),
         "recursive_docs": [],
         "token_usage": token_usage,
+        "parsers_used": parsers_used,
     }
 
 
@@ -328,6 +347,11 @@ def parse(
                     ),
                     "total": sum(r["token_usage"]["total"] for r in chunk_results),
                 },
+                "parsers_used": [
+                    parser
+                    for r in chunk_results
+                    for parser in r.get("parsers_used", [])
+                ],
             }
 
         if "api_cost_mapping" in kwargs and "token_usage" in result:

tests/test_parser.py

@@ -432,6 +432,31 @@ async def test_parse_with_schema():
     assert all(key in result for key in sample_schema.keys())
 
 
+@pytest.mark.parametrize("character_threshold", [160, 100])
+@pytest.mark.asyncio
+async def test_cost_priority_routing(character_threshold):
+    sample = "examples/inputs/NHL_agreement_2022_pg1.pdf"
+    parser_type = "AUTO"
+    router_priority = "cost"
+    api_cost_path = os.path.join(os.path.dirname(__file__), "api_cost_mapping.json")
+    config = {
+        "parser_type": parser_type,
+        "router_priority": router_priority,
+        "page_nums": 1,
+        "pages_per_split": 1,
+        "character_threshold": character_threshold,
+        "api_cost_mapping": api_cost_path,
+    }
+    result = parse(sample, **config)
+    logger.debug(f"Token cost: {result['token_cost']}")
+    if character_threshold > 150:
+        assert "STATIC_PARSE" in result["parsers_used"]
+        assert result["token_cost"]["total"] == 0
+    else:
+        assert "LLM_PARSE" in result["parsers_used"]
+        assert result["token_cost"]["total"] > 0
+
+
 @pytest.mark.asyncio
 async def test_audio_parse():
     path = "examples/inputs/harvard.wav"