test-ollama.py

Go to the documentation of this file.

import asyncio
import ollama
from typing import *
 
T = TypeVar('T')
 
#https://stackoverflow.com/questions/50241696/how-to-iterate-over-an-asynchronous-iterator-with-a-timeout
 
# async generator, needs python 3.6
async def timeout(it: AsyncIterator[T], timeo: float, sentinel: T) -> AsyncGenerator[T, None]:
    try:
        nxt = asyncio.ensure_future(it.__anext__())
        while True:
            try:
                yield await asyncio.wait_for(asyncio.shield(nxt), timeo)
                nxt = asyncio.ensure_future(it.__anext__())
            except asyncio.TimeoutError:
                yield sentinel
    except StopAsyncIteration:
        pass
    finally:
        nxt.cancel()  # in case we're getting cancelled our self
 
async def something():
    yield 1
    await asyncio.sleep(1.1)
    yield 2
    await asyncio.sleep(2.1)
    yield 3
 
 
async def test():
    expect = [1, None, 2, None, None, 3]
    async for item in timeout(something(), 1, None):
        print("Check", item)
        assert item == expect.pop(0)
 
async def main():
  client = ollama.AsyncClient()
  messages = []
 
  while True:
    if content_in := input('>>> '):
      messages.append({'role': 'user', 'content': content_in})
      content_out = ''
      message = {'role': 'assistant', 'content': ''}
      respaiter = await client.chat(model='tinydolphin', messages=messages, stream=True)
      print("type is {}".format(type(respaiter)))
      zob = timeout(respaiter, 0.01, None)
      print("zob type is {}".format(type(zob)))
      
      keepgoing = True
      while keepgoing:
        response = await anext(zob)
        if response is not None:
          if response['done']:
            messages.append(message)
            keepgoing = False
          content = response['message']['content']
          print(content, end='', flush=True)
          message['content'] += content
        else:
          print("[timout]")
      print()
 
 
#asyncio.get_event_loop().run_until_complete(test())
 
try:
  asyncio.run(main())
except (KeyboardInterrupt, EOFError):
  ...