46 lines
1.8 KiB
Python
46 lines
1.8 KiB
Python
import asyncio
|
|
import json
|
|
import sys
|
|
from unittest import mock
|
|
|
|
sys.path.insert(0, r'C:\kefu\wechat_rpa')
|
|
import httpx
|
|
import model_gateway as gateway
|
|
from model_usage import UsageRecorder
|
|
from test_model_gateway import _outlet, _openai_body
|
|
|
|
|
|
async def run():
|
|
database = mock.Mock()
|
|
saved = []
|
|
database.record_model_usage_batch.side_effect = lambda records: saved.extend(dict(item) for item in records)
|
|
recorder = UsageRecorder(database, tenant_id='isolated-audit')
|
|
first, second = _outlet(), _outlet()
|
|
first.config['id'], second.config['id'] = 'malformed', 'slow-valid'
|
|
|
|
async def handler(request):
|
|
text = json.loads(request.content)['messages'][0]['content']
|
|
if text == 'first':
|
|
return httpx.Response(200, json={
|
|
'choices': [{'message': {'content': '', 'tool_calls': 1}}],
|
|
'usage': {'total_tokens': 10}})
|
|
await asyncio.sleep(.1)
|
|
return httpx.Response(200, json={**_openai_body('answer'), 'usage': {'total_tokens': 20}})
|
|
|
|
error = None
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
try:
|
|
async with recorder:
|
|
await asyncio.gather(
|
|
gateway.call_outlet(client, first, [{'role': 'user', 'content': 'first'}], usage_recorder=recorder),
|
|
gateway.call_outlet(client, second, [{'role': 'user', 'content': 'second'}], usage_recorder=recorder))
|
|
except Exception as exc:
|
|
error = type(exc).__name__
|
|
await asyncio.sleep(.2)
|
|
result = {'raised': error, 'saved': [{'provider_id': e['provider_id'], 'total_tokens': e['total_tokens']} for e in saved],
|
|
'unflushed': [{'provider_id': e['provider_id'], 'total_tokens': e['total_tokens']} for e in recorder.events]}
|
|
print(json.dumps(result))
|
|
|
|
|
|
asyncio.run(run())
|