Yura update.

This commit is contained in:
retoor 2024-12-09 20:16:18 +01:00
parent e3c41aedcb
commit 1f4dc697dc

View File

@ -67,36 +67,43 @@ class AsyncClient:
def __init__(self, url="ws://127.0.0.1:8470"): def __init__(self, url="ws://127.0.0.1:8470"):
self.url = url self.url = url
self.client = AsyncRPCClient(self.url)
self.queue_in = asyncio.Queue() self.queue_in = asyncio.Queue()
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 self.ws = None
async def __aenter__(self): @property
def _connection(self):
if not self.client:
self.client = AsyncRPCClient(self.url)
self.ws = await self.client.ws self.ws = await self.client.ws
return self.client
async def __aenter__(self):
conn = self._connection
return self return self
async def __aexit__(self, *args, **kwargs): async def __aexit__(self, *args, **kwargs):
await self.client.close() await self.close(0
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._connection.create(name=name, extends=extends, system=system)
async def chat(self, token, message, datasets=None): async def chat(self, token, message, datasets=None):
yield await self.client.chat(uid=token, message=message, datasets=datasets or []) yield await self._connection.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"):
break break
async def connect(self, name): async def connect(self, name):
return await self.client.connect(name) return await self._connection.connect(name)
async def close(self): async def close(self):
await self.client.close() await self.client.close()
self.client = None
self.ws = None self.ws = None