feat: add test for crawl_database

src/deepengineer/deepsearch/analyse_markdown_agent.py

@@ -1,3 +1,7 @@
+"""
+Simple agent to analyse a markdown, just to test some ideas.
+"""
+
 from smolagents import CodeAgent, tool, Tool, LiteLLMModel
 from deepengineer.webcrawler.pdf_utils import get_markdown_by_page_numbers, get_table_of_contents_per_page_markdown, find_in_markdown, convert_ocr_response_to_markdown
 from mistralai import OCRResponse

src/deepengineer/webcrawler/pdf_utils.py

@@ -7,7 +7,7 @@ from mistralai import Mistral
 import os
 from litellm import completion
 
-from mistralai.models import OCRResponse, OCRPageObject
+from mistralai.models import OCRResponse, OCRPageObject, OCRUsageInfo
 import yaml
 from tenacity import retry, stop_after_attempt, wait_fixed, RetryError
 from litellm.exceptions import BadRequestError
@@ -106,9 +106,10 @@ def get_table_of_contents_per_page_markdown(markdown: OCRResponse) -> str:
     return table_of_contents
 
 def convert_raw_markdown_to_ocr_response(raw_markdown: str) -> OCRResponse:
-    # split by big title starting with # and then a space
-    pages = raw_markdown.split("\n# ")
-    return OCRResponse(pages=[OCRPageObject(markdown="# " + page, page_number=i) for i, page in enumerate(pages)])
+    pages = raw_markdown.split("# ")
+    usage_info_empty = OCRUsageInfo(pages_processed=0)
+    return OCRResponse(pages=[OCRPageObject(index=i, markdown="# " + page, images=[], dimensions=None) for i, page in enumerate(pages)], usage_info=usage_info_empty, model="",)
+
 
 
 

tests/webcrawler/test_crawl_database.py

@@ -16,4 +16,9 @@ def test_crawl_database_arxiv_link():
     assert db.get_markdown_of_url("https://arxiv.org/abs/2105.00643").pages[0].markdown is not None
     assert len(db.get_markdown_of_url("https://arxiv.org/abs/2105.00643").pages) == 20
     
-    
+def test_crawl_database_wikipedia_url():
+    db = DataBase()
+    db.crawl_url("https://en.wikipedia.org/wiki/Deep_learning")
+    assert db.get_markdown_of_url("https://en.wikipedia.org/wiki/Deep_learning") is not None
+    assert db.get_markdown_of_url("https://en.wikipedia.org/wiki/Deep_learning").pages[0].markdown is not None
+    assert len(db.get_markdown_of_url("https://en.wikipedia.org/wiki/Deep_learning").pages) >= 40