Added support for datasets.

Browse Source
This commit is contained in:
retoor 2024-12-09 18:49:50 +01:00
parent 81cbe351d6
commit e3c41aedcb
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/yura/client.py
View File

@ -72,18 +72,21 @@ class AsyncClient:
self.queue_out = asyncio.Queue() self.queue_out = asyncio.Queue()
self.communication_task = None self.communication_task = None
self.session_id = None self.session_id = None
self.ws = None
async def __aenter__(self): async def __aenter__(self):
self.ws = await self.client.ws
return self return self
async def __aexit__(self, *args, **kwargs): async def __aexit__(self, *args, **kwargs):
pass await self.client.close()
self.ws = None
async def create(self, name, extends, system): async def create(self, name, extends, system):
return await self.client.create(name=name, extends=extends, system=system) return await self.client.create(name=name, extends=extends, system=system)
async def chat(self, token, message): async def chat(self, token, message, datasets=None):
yield await self.client.chat(uid=token, message=message) yield await self.client.chat(uid=token, message=message, datasets=datasets or [])
async for msg in self.client: async for msg in self.client:
yield msg yield msg
if msg.get("done"): if msg.get("done"):
@ -94,7 +97,7 @@ class AsyncClient:
async def close(self): async def close(self):
await self.client.close() await self.client.close()
self.client = None self.ws = None
async def cli_client(url="ws://127.0.0.1:8470"): async def cli_client(url="ws://127.0.0.1:8470"):