wehub-resource-sync
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
import asyncio
|
|
import unittest
|
|
|
|
from gpt_researcher.llm_provider.generic.base import GenericLLMProvider
|
|
|
|
|
|
class _Chunk:
|
|
def __init__(self, content, usage_metadata=None, response_metadata=None):
|
|
self.content = content
|
|
self.usage_metadata = usage_metadata
|
|
self.response_metadata = response_metadata or {}
|
|
|
|
|
|
class _StreamingLLM:
|
|
async def astream(self, messages, **kwargs):
|
|
yield _Chunk("Hello")
|
|
yield _Chunk(
|
|
"",
|
|
usage_metadata={"input_tokens": 321, "output_tokens": 123},
|
|
response_metadata={"usage": {"input_tokens": 321, "output_tokens": 123}},
|
|
)
|
|
|
|
|
|
class TestLLMUsageTracking(unittest.TestCase):
|
|
def test_stream_response_captures_usage_from_empty_final_chunk(self):
|
|
provider = GenericLLMProvider(_StreamingLLM(), verbose=False)
|
|
|
|
response = asyncio.run(provider.stream_response([{"role": "user", "content": "hi"}]))
|
|
|
|
self.assertEqual(response, "Hello")
|
|
self.assertEqual(
|
|
provider.last_usage_metadata,
|
|
{"input_tokens": 321, "output_tokens": 123},
|
|
)
|
|
self.assertEqual(
|
|
provider.last_response_metadata.get("usage"),
|
|
{"input_tokens": 321, "output_tokens": 123},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|