File size: 2,968 Bytes
b5fafa1
b1060b0
 
 
b5fafa1
 
b1060b0
84c66cd
b5fafa1
b1060b0
 
 
 
f1368c4
 
b1060b0
 
 
f1368c4
 
b1060b0
b5fafa1
b1060b0
 
 
 
 
 
 
 
 
 
 
 
 
f0e5174
 
b1060b0
 
 
 
 
b5fafa1
b1060b0
f1368c4
 
b1060b0
 
 
84c66cd
f1368c4
b1060b0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f0e5174
b5fafa1
 
f0e5174
 
 
 
 
 
 
 
b5fafa1
f0e5174
 
 
 
 
 
b5fafa1
 
 
 
f0e5174
 
 
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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import numpy as np
import pytest
from deepengineer.webcrawler.async_search import (
    SearchResponse,
    arxiv_search_async,
    get_linkup_balance,
    get_tavily_usage,
    linkup_search_async,
    tavily_search_async,
)


@pytest.mark.expensive
@pytest.mark.asyncio
async def test_tavily_search_async():
    usage_before = get_tavily_usage()
    print(usage_before)

    response = await tavily_search_async(
        search_query="Would it be possible to make a thermal reactor with graphite and lead?",
    )

    print(response.answer)
    assert response is not None
    assert isinstance(response, SearchResponse)
    assert (
        response.query
        == "Would it be possible to make a thermal reactor with graphite and lead?"
    )
    assert response.answer is not None
    assert response.search_results is not None
    assert len(response.search_results) == 10
    assert response.search_results[0].title is not None
    assert response.search_results[0].url is not None
    assert response.search_results[0].content is not None
    # raw content is often not available for tavily
    # assert any(result.raw_content is not None for result in response.search_results)

    usage_after = get_tavily_usage()
    print(usage_after)
    assert usage_after == usage_before + 1


@pytest.mark.expensive
@pytest.mark.asyncio
async def test_linkup_search_async():
    balance_before = get_linkup_balance()
    print(balance_before)

    response = await linkup_search_async(
        search_query="Would it be possible to make a thermal reactor with graphite and lead?",
    )
    print(response.answer)
    assert response is not None
    assert isinstance(response, SearchResponse)
    assert (
        response.query
        == "Would it be possible to make a thermal reactor with graphite and lead?"
    )
    assert len(response.search_results) >= 10
    assert response.search_results[0].title is not None
    assert response.search_results[0].url is not None
    assert response.search_results[0].content is not None
    assert response.search_results[0].raw_content is None

    balance_after = get_linkup_balance()
    print(balance_after)
    assert np.isclose(balance_after, balance_before - 0.005)


@pytest.mark.expensive
@pytest.mark.asyncio
async def test_arxiv_search_async():
    balance_before = get_linkup_balance()

    response = await arxiv_search_async(
        search_query="Would it be possible to make a thermal reactor with graphite and lead?",
    )

    assert response is not None
    assert isinstance(response, SearchResponse)
    assert response.query is not None
    assert response.answer is not None
    assert response.search_results is not None
    assert len(response.search_results) >= 10
    assert any(
        result.url.startswith("https://arxiv.org/abs/")
        for result in response.search_results
    )

    balance_after = get_linkup_balance()
    assert np.isclose(balance_after, balance_before - 0.005)