mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-17 11:29:56 +00:00
576589bdc5
* feat: add websearch in a sandbox
* fix: celery callbacks
(cherry picked from commit fe7d0ce9d5)
* fix: mypy
118 lines
3.2 KiB
Python
118 lines
3.2 KiB
Python
import asyncio
|
|
from typing import Any
|
|
|
|
from private_gpt.components.web.scraper.pool import ScrapeSessionPool
|
|
|
|
|
|
class _FakeSession:
|
|
pass
|
|
|
|
|
|
class _FakeSandboxProvider:
|
|
def __init__(self) -> None:
|
|
self.created: list[_FakeSession] = []
|
|
self.killed: list[_FakeSession] = []
|
|
|
|
async def create_session(self, **_: Any) -> _FakeSession:
|
|
session = _FakeSession()
|
|
self.created.append(session)
|
|
return session
|
|
|
|
async def kill_session(self, session: _FakeSession) -> None:
|
|
self.killed.append(session)
|
|
|
|
|
|
def _pool(
|
|
provider: _FakeSandboxProvider,
|
|
*,
|
|
pool_size: int = 2,
|
|
max_requests_per_session: int = 1,
|
|
idle_timeout_seconds: int = 300,
|
|
) -> ScrapeSessionPool:
|
|
return ScrapeSessionPool(
|
|
provider, # type: ignore[arg-type]
|
|
pool_size=pool_size,
|
|
max_requests_per_session=max_requests_per_session,
|
|
idle_timeout_seconds=idle_timeout_seconds,
|
|
)
|
|
|
|
|
|
async def test_single_request_sessions_are_killed_on_release() -> None:
|
|
provider = _FakeSandboxProvider()
|
|
pool = _pool(provider, max_requests_per_session=1)
|
|
|
|
session = await pool.acquire()
|
|
await pool.release(session)
|
|
|
|
assert provider.killed == [session]
|
|
|
|
other = await pool.acquire()
|
|
assert other is not session, "exhausted sessions must not be reused"
|
|
await pool.release(other)
|
|
|
|
|
|
async def test_sessions_are_reused_until_request_budget_is_spent() -> None:
|
|
provider = _FakeSandboxProvider()
|
|
pool = _pool(provider, max_requests_per_session=3)
|
|
|
|
first = await pool.acquire()
|
|
await pool.release(first)
|
|
assert provider.killed == []
|
|
|
|
assert await pool.acquire() is first
|
|
await pool.release(first)
|
|
assert await pool.acquire() is first
|
|
await pool.release(first)
|
|
|
|
assert provider.killed == [first], "third release exhausts the budget"
|
|
await pool.close()
|
|
|
|
|
|
async def test_batch_release_counts_all_requests() -> None:
|
|
provider = _FakeSandboxProvider()
|
|
pool = _pool(provider, max_requests_per_session=3)
|
|
|
|
session = await pool.acquire()
|
|
await pool.release(session, requests=3)
|
|
|
|
assert provider.killed == [session], "a batch of 3 exhausts the budget at once"
|
|
await pool.close()
|
|
|
|
|
|
async def test_broken_sessions_are_killed_even_with_budget_left() -> None:
|
|
provider = _FakeSandboxProvider()
|
|
pool = _pool(provider, max_requests_per_session=10)
|
|
|
|
session = await pool.acquire()
|
|
await pool.release(session, broken=True)
|
|
|
|
assert provider.killed == [session]
|
|
await pool.close()
|
|
|
|
|
|
async def test_pool_size_caps_concurrent_sessions() -> None:
|
|
provider = _FakeSandboxProvider()
|
|
pool = _pool(provider, pool_size=1)
|
|
|
|
first = await pool.acquire()
|
|
|
|
waiter = asyncio.create_task(pool.acquire())
|
|
await asyncio.sleep(0.01)
|
|
assert not waiter.done(), "second acquire must wait for the free slot"
|
|
|
|
await pool.release(first)
|
|
second = await asyncio.wait_for(waiter, timeout=1)
|
|
await pool.release(second)
|
|
|
|
|
|
async def test_close_kills_idle_sessions() -> None:
|
|
provider = _FakeSandboxProvider()
|
|
pool = _pool(provider, max_requests_per_session=5)
|
|
|
|
session = await pool.acquire()
|
|
await pool.release(session)
|
|
assert provider.killed == []
|
|
|
|
await pool.close()
|
|
assert provider.killed == [session]
|