mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-30 04:15:45 +08:00
102 lines
5.0 KiB
Python
102 lines
5.0 KiB
Python
"""Real HTTPX requests against a FAKE loopback provider (not provider E2E)."""
|
|
import asyncio
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from aiohttp import web
|
|
from langbot_plugin.api.entities.builtin.rag import (
|
|
RetrievalContext, IngestionContext, FileObject, FileMetadata, DocumentStatus,
|
|
)
|
|
|
|
from test_shared_state import load_plugin, bind, StorageFixture, CONNECTORS, configuration
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('name', CONNECTORS)
|
|
async def test_fake_http_ingest_retrieve_restart_delete_contract(name):
|
|
pc, ec = load_plugin(name)
|
|
received = []
|
|
async def fake_service(request):
|
|
body = await request.read()
|
|
auth = request.headers.get('Authorization')
|
|
received.append((request.method, request.path, auth, body))
|
|
await asyncio.sleep(.002)
|
|
if request.path.endswith('/retrieve'):
|
|
return web.json_response({'records': [{'segment': {'id': 'seg', 'content': auth}, 'score': .8}]})
|
|
if request.path.endswith('/searchTest'):
|
|
return web.json_response({'data': [{'id': 'seg', 'q': auth, 'a': 'answer', 'score': .8}]})
|
|
if request.path.endswith('/retrieval'):
|
|
return web.json_response({'code': 0, 'data': {'chunks': [{'id': 'seg', 'content': auth, 'similarity': .8}]}})
|
|
if request.path.endswith('/create-by-file'):
|
|
assert b'fixture file' in body
|
|
return web.json_response({'document': {'id': 'remote-doc'}})
|
|
if request.path.endswith('/localFile'):
|
|
assert b'fixture file' in body
|
|
return web.json_response({'code': 200, 'data': {'collectionId': 'remote-doc', 'results': {'insertLen': 2}}})
|
|
if request.method == 'POST' and request.path.endswith('/documents'):
|
|
assert b'fixture file' in body
|
|
return web.json_response({'code': 0, 'data': [{'id': 'remote-doc'}]})
|
|
if request.method == 'GET':
|
|
return web.json_response({'code': 0, 'data': []})
|
|
return web.json_response({'code': 200 if name == 'FastGPTConnector' else 0, 'data': {}})
|
|
app = web.Application()
|
|
app.router.add_route('*', '/{path:.*}', fake_service)
|
|
runner = web.AppRunner(app)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, '127.0.0.1', 0)
|
|
await site.start()
|
|
port = site._server.sockets[0].getsockname()[1]
|
|
store = StorageFixture()
|
|
plugin = bind(pc, store)
|
|
async def read_file(path):
|
|
assert path == 'fixture-file'
|
|
return b'fixture file'
|
|
plugin.get_knowledge_file_stream = read_file
|
|
engine = ec()
|
|
engine.plugin = plugin
|
|
cfg = {**configuration('A'), 'api_base_url': f'http://127.0.0.1:{port}', 'auto_graphrag': True, 'auto_raptor': True}
|
|
try:
|
|
ctx = IngestionContext(file_object=FileObject(metadata=FileMetadata(filename='doc.txt', file_size=12, mime_type='text/plain', document_id='local-doc', knowledge_base_id='kb'), storage_path='fixture-file'), knowledge_base_id='kb', creation_settings=cfg)
|
|
result = await engine.ingest(ctx)
|
|
assert result.status == DocumentStatus.PROCESSING
|
|
assert result.document_id == 'remote-doc'
|
|
if name == 'RAGFlowConnector':
|
|
assert any(x[1].endswith('/chunks') for x in received)
|
|
assert any(x[1].endswith('/run_graphrag') for x in received)
|
|
assert any(x[1].endswith('/run_raptor') for x in received)
|
|
async def retrieve(i):
|
|
config = {**cfg, 'api_key': f'T{i}', 'dify_apikey': f'T{i}'}
|
|
return await engine.retrieve(RetrievalContext(query=f'query{i}', knowledge_base_id='kb', creation_settings=config, retrieval_settings={'top_k': 2}))
|
|
outputs = await asyncio.gather(*(retrieve(i) for i in range(8)))
|
|
for i, output in enumerate(outputs):
|
|
assert output.total_found == 1
|
|
assert f'Bearer T{i}' in output.results[0].content[0].text
|
|
assert output.results[0].score == .8
|
|
restarted = ec()
|
|
restarted.plugin = bind(pc, store)
|
|
assert await restarted.delete_document('kb', 'remote-doc')
|
|
assert received[-1][2] == 'Bearer A'
|
|
await restarted.on_knowledge_base_delete('kb')
|
|
assert not await restarted.delete_document('kb', 'remote-doc')
|
|
finally:
|
|
await runner.cleanup()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('name', CONNECTORS)
|
|
async def test_http_has_total_deadline_not_only_inactivity_timeout(name, monkeypatch):
|
|
pc, ec = load_plugin(name)
|
|
from components import shared_state
|
|
assert hasattr(shared_state, 'HTTP_TOTAL_TIMEOUT'), 'Require a total operation deadline'
|
|
monkeypatch.setattr(shared_state, 'HTTP_TOTAL_TIMEOUT', .02)
|
|
import httpx
|
|
async def stalls(request):
|
|
await asyncio.sleep(10)
|
|
client = httpx.AsyncClient
|
|
monkeypatch.setattr(httpx, 'AsyncClient', lambda **kw: client(transport=httpx.MockTransport(stalls), **kw))
|
|
engine = ec()
|
|
engine.plugin = bind(pc, StorageFixture())
|
|
output = await asyncio.wait_for(engine.retrieve(RetrievalContext(query='q', knowledge_base_id='kb', creation_settings=configuration('A'))), .5)
|
|
assert output.total_found == 0
|