Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Thread Deleted Callback needed to delete corresponding OpenAI Assistant API thread #1451

Open
gloveboxes opened this issue Oct 21, 2024 · 1 comment
Labels
enhancement New feature or request needs-triage

Comments

@gloveboxes
Copy link

Is your feature request related to a problem? Please describe.
Currently when a user deletes a Chainlit thread it leaves the OpenAI Assistant API Thread orphaned. I need to be able to delete an OpenAI Assistant API Thread when a user deletes a thread from the conversation history.

Describe the solution you'd like
A thread deleted callback is required that includes context, allowing me to identify and delete the corresponding thread from the OpenAI Assistant API.

Describe alternatives you've considered
I explored using Actions but give a very counter intuitive/confusing user experience.

Additional context
n/a

@dosubot dosubot bot added the enhancement New feature or request label Oct 21, 2024
@hadarsharon
Copy link

hadarsharon commented Oct 31, 2024

Hi @gloveboxes - while this is not an "official" solution by a long shot, I can suggest my own solution for the exact same problem. Perhaps it would be of help to you (and others). It is not quite straightforward, but with enough effort it works well.

  • First off, a bit of preface - while you could leverage Chainlit's Data persistence option (e.g. with SQLAlchemy) and add your own table for mapping a Chainlit thread to an OpenAI thread, I had opted against that for two major reasons:
  1. I didn't want to break compatibility or modify the existing relations between Chainlit's data tables
  2. I liked the separation of concerns between Chainlit's native data and an external data source (i.e. OpenAI), so it made sense to split this mapping on my end as well. Again this is just a personal preference.
  • You do need to enable data persistence, so as to override the delete_thread() method of the data layer, at least based on my own implementation.

With that in mind, having a Postgres (SQLAlchemy) data layer in place, I also spun up a Redis instance as part of my Chainlit app deployment (in my case it's part of the docker-compose stack in which Chainlit runs as its own service, and Chainlit's database is in a Postgres container, but again, you could manage anything however you'd like).

In this Redis instance I store a simple mapping between a Chainlit thread id and an OpenAI thread id. Whenever the Chainlit thread gets deleted, I look for a match between the deleted thread id and the now-to-be-orphaned OpenAI thread id, and if there's one, I delete the OpenAI thread as well. Finally, I then delete it from Redis too.

Some code for reference:

class RedisDataLayer:
    """
    Additional Data Layer to manage OpenAI and Chainlit thread mappings
    """

    def __init__(self, conn_info: str, **kwargs):
        try:
            self.client = redis.from_url(conn_info, decode_responses=True, **kwargs)
            self.logger.info("RedisDataLayer::initialized")
        except Exception as e:
            self.logger.critical(f"RedisDataLayer::initialization error: {e}")
            raise

    @cached_property
    def logger(self):
        # set some logger here...

    async def get_thread(self, chainlit_uuid: str) -> Optional[str]:
        key = f"chainlit_thread:{chainlit_uuid}"
        self.logger.info(f"RedisDataLayer, get_thread, chainlit_uuid = {chainlit_uuid}")
        data = await self.client.hgetall(key)
        if data:
            return data["openai_thread"]
        else:
            self.logger.warning(
                f"RedisDataLayer, get_thread, no mapping found for Chainlit Thread with ID: {chainlit_uuid}"
            )
            return None

    async def set_thread(self, chainlit_uuid: str, openai_uuid: str):
        key = f"chainlit_thread:{chainlit_uuid}"
        self.logger.info(f"RedisDataLayer, set_thread, chainlit_uuid = {chainlit_uuid}, openai_uuid = {openai_uuid}")
        value = {"openai_thread": openai_uuid, "timestamp": datetime.now(UTC).isoformat()}
        await self.client.hset(name=key, mapping=value)

    async def delete_thread(self, chainlit_uuid: str):
        openai_uuid = await self.get_thread(chainlit_uuid=chainlit_uuid)
        self.logger.info(f"RedisDataLayer, delete_thread, chainlit_uuid = {chainlit_uuid}, openai_uuid = {openai_uuid}")

        try:
            openai_client = AsyncOpenAI(...)  # set this however you'd like
        except Exception as e:
            self.logger.error(f"RedisDataLayer::OpenAI client initialization error: {e}")
            raise
        else:
            try:
                self.logger.info(f"RedisDataLayer::Deleting thread from OpenAI: {openai_uuid}")
                await openai_client.beta.threads.delete(thread_id=openai_uuid)
            except Exception as e:
                self.logger.warning(f"RedisDataLayer::OpenAI thread not deleted (perhaps it no longer exists?): {e}")
            finally:
                # Assuming OpenAI client didn't fail, and we called deletion, remove from Redis
                self.logger.info(
                    f"RedisDataLayer::Deleting Chainlit thread from Redis: {chainlit_uuid} (matched with OpenAI thread {openai_uuid})"
                )
                await self.client.delete(f"chainlit_thread:{chainlit_uuid}")

Then, in my PostgresDataLayer (which is basically almost a dummy class that inherits SQLAlchemyDataLayer and enriches the delete_thread method for deleting the OpenAI thread:

class PostgresDataLayer(SQLAlchemyDataLayer):
    def __init__(
        self,
        conn_info: str,
        ssl_require: bool = False,
        storage_provider: Optional[BaseStorageClient] = None,
        user_thread_limit: Optional[int] = 100,
        show_logger: Optional[bool] = True,
    ):
        try:
            super().__init__(
                conninfo=conn_info,
                ssl_require=ssl_require,
                storage_provider=storage_provider,
                user_thread_limit=user_thread_limit,
                show_logger=show_logger,
            )
            self.logger.info("PostgresDataLayer::initialized")
            self.redis = RedisDataLayer(conn_info=ChainlitSettings.redis_conn_info)
        except Exception as e:
            self.logger.critical(f"PostgresDataLayer::initialization error: {e}")

    @cached_property
    def logger(self):
        # set some logger here...

    async def delete_thread(self, thread_id: str):
        await super().delete_thread(thread_id=thread_id)
        await self.redis.delete_thread(chainlit_uuid=thread_id)

That should allow you to set, get and delete. Deletion is already invoked implicitly when you delete a Chainlit thread - all that remains to do is to actually set this mapping in Redis from somewhere.

In your Chainlit app code, you can do something like:

redis_data_layer = RedisDataLayer(conn_info=...)  # set it however is comfortable for you based on the above implementation

async def manage_openai_thread(chainlit_thread_id: str) -> str:
    openai_thread_id = await redis_data_layer.get_thread(chainlit_uuid=chainlit_thread_id)
    if openai_thread_id:
        logger.info(f"Chainlit::Found OpenAI thread id {openai_thread_id} for Chainlit thread {chainlit_thread_id}")
    else:
        logger.info(f"Chainlit::Creating new OpenAI thread for Chainlit thread {chainlit_thread_id}")
        client = AsyncOpenAI(...)  # initialize however you want
        openai_thread_id = await client.beta.threads.create()
    await redis_data_layer.set_thread(chainlit_uuid=chainlit_thread_id, openai_uuid=openai_thread_id)
    return openai_thread_id

def get_current_chainlit_thread_id() -> str:
    return cl.context.session.thread_id

The final step is to call this method from cl.on_message callback, to implement this routine setting/creating feature behind the scenes:

@cl.on_message
async def on_message(message: cl.Message) -> None:
    chainlit_thread_id = get_current_chainlit_thread_id()

    await manage_openai_thread(chainlit_thread_id)

    # ... do what you will from here

Again, this is not straightforward and if there is a better (or more native) way of doing it in Chainlit, let me know. I hope this helps.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request needs-triage
Projects
None yet
Development

No branches or pull requests

2 participants