thinkcol/restaurants

1import asyncio2import os3import re4from inspect import cleandoc5 6import chainlit as cl7import openai8from chainlit import LLMSettings9from chainlit.config import config10 11 12def replace_newlines(match: re.Match) -> str:13    newlines = match.group(0)14    count = len(newlines)15    if count <= 1:16        return " "17    return newlines[1:]18 19 20# TODO each chain should be able to make a child chain?21# root = Chain()22# first = root.child("something")23# first.llm('foo')24class Chain:25    def __init__(self, message_id: str | None, llm_settings: LLMSettings | None = None):26        self.llm_settings = llm_settings27        self.root_id = message_id28 29    def make_message(self, name, final, **kwargs) -> cl.Message:30        if not name:31            name = config.ui.name if final else "Child Chain"32        return cl.Message(33            author=name,34            parent_id=None if final else self.root_id,35            **kwargs,36        )37 38    async def text(self, text, final=False, name=None) -> cl.Message:39        message = self.make_message(content=text, final=final, name=name)40        await message.send()41        return message42 43    async def text_stream(self, text: str, delay=.1, name=None, final=False) -> cl.Message:44        message = self.make_message(content='', final=final, name=name)45        tokens = text.split(" ")46        first = True47        for token in tokens:48            if not first:49                token = " " + token50            await message.stream_token(token)51            await asyncio.sleep(delay)52            first = False53        await message.send()54        return message55 56    async def llm(self, template, *args, name=None, final=False, **kwargs) -> cl.Message:57        template = cleandoc(template)58        template = re.sub('\n+', replace_newlines, template)  # remove a newline59 60        variables = re.findall(r'\{(.*?)}', template)61        if len(args) > 1:62            raise RuntimeError("If there is more than one argument, use kwargs")63        if len(args) > 0 and len(kwargs) > 0:64            raise RuntimeError("Cannot combine args and kwargs")65        if len(args) > 0:66            if len(variables) > 1:67                raise RuntimeError("This chain expects more than one argument. Use kwargs instead.")68            variable_dict = {variables[0]: args[0]}69        else:70            variable_dict = kwargs71 72        prompt = template.format(**variable_dict)73        message = self.make_message(content='', name=name, prompt=prompt, llm_settings=self.llm_settings, final=final)74 75        async for response in await openai.ChatCompletion.acreate(76                **self.llm_settings.to_settings_dict(), api_key=os.environ.get('OPENAI_API_KEY'), stream=True,77                messages=[{'role': 'user', 'content': prompt}]78        ):79            token = response.choices[0]["delta"].get("content", "")80            await message.stream_token(token)81 82        await message.send()83        return message84