Spaces:
Runtime error
Runtime error
File size: 1,169 Bytes
b5fafa1 8dfc9f8 0159aaf b5fafa1 8dfc9f8 b5fafa1 8dfc9f8 e040f4f 8dfc9f8 b5fafa1 48a564a b5fafa1 e040f4f 48a564a 4fbfc2b b5fafa1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import pytest
from deepengineer.common_path import DATA_DIR
from deepengineer.deepsearch.main_pdf_agent import (
FindInMarkdownTool,
GetMarkdownTool,
GetPagesContentTool,
GetTableOfContentsTool,
create_agent,
)
from mistralai import OCRResponse
def load_mock_ocr_response() -> OCRResponse:
with open(DATA_DIR / "report_thermal_neutron.json") as f:
return OCRResponse.model_validate_json(f.read())
def test_pdf_agent():
ocr_response = load_mock_ocr_response()
pdf_agent = create_agent(ocr_response)
assert pdf_agent is not None
assert pdf_agent.name == "markdown_agent"
assert pdf_agent.tools is not None
assert len(pdf_agent.tools) == 4 + 1 # +1 for the final answer
GetTableOfContentsTool(ocr_response).forward()
GetMarkdownTool(ocr_response).forward()
GetPagesContentTool(ocr_response).forward([1, 2, 3])
FindInMarkdownTool(ocr_response).forward(["thermal neutron", "neutron"])
@pytest.mark.expensive
def test_run_pdf_agent():
ocr_response = load_mock_ocr_response()
pdf_agent = create_agent(ocr_response)
assert pdf_agent.run("Give me a summary of the document.") is not None
|