From 5587dbbd325638e56183544fe35eb4ae354e1d79 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 10 Jul 2024 11:14:40 -0700 Subject: [PATCH] add async assistants delete support --- litellm/assistants/main.py | 56 ++++++++++++++++++++++++++++++++++++-- litellm/llms/openai.py | 46 ++++++++++++++++++++++++------- 2 files changed, 89 insertions(+), 13 deletions(-) diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index 25d502c8b55..ba169f5e203 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -4,11 +4,12 @@ import asyncio import contextvars import os from functools import partial -from typing import Any, Dict, Iterable, List, Literal, Optional, Union +from typing import Any, Coroutine, Dict, Iterable, List, Literal, Optional, Union import httpx from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI from openai.types.beta.assistant import Assistant +from openai.types.beta.assistant_deleted import AssistantDeleted import litellm from litellm import client @@ -339,6 +340,44 @@ def create_assistants( return response +async def adelete_assistant( + custom_llm_provider: Literal["openai", "azure"], + client: Optional[AsyncOpenAI] = None, + **kwargs, +) -> AssistantDeleted: + loop = asyncio.get_event_loop() + ### PASS ARGS TO GET ASSISTANTS ### + kwargs["async_delete_assistants"] = True + try: + kwargs["client"] = client + # Use a partial function to pass your keyword arguments + func = partial(delete_assistant, custom_llm_provider, **kwargs) + + # Add the context to the function + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + + _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore + model="", custom_llm_provider=custom_llm_provider + ) # type: ignore + + # Await normally + init_response = await loop.run_in_executor(None, func_with_context) + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + return response # type: ignore + except Exception as e: + raise exception_type( + model="", + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs={}, + extra_kwargs=kwargs, + ) + + def delete_assistant( custom_llm_provider: Literal["openai", "azure"], assistant_id: str, @@ -347,11 +386,21 @@ def delete_assistant( api_base: Optional[str] = None, api_version: Optional[str] = None, **kwargs, -) -> Dict[str, Any]: +) -> AssistantDeleted: optional_params = GenericLiteLLMParams( api_key=api_key, api_base=api_base, api_version=api_version, **kwargs ) + async_delete_assistants: Optional[bool] = kwargs.pop( + "async_delete_assistants", None + ) + if async_delete_assistants is not None and not isinstance( + async_delete_assistants, bool + ): + raise ValueError( + "Invalid value passed in for async_delete_assistants. Only bool or None allowed" + ) + ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -368,7 +417,7 @@ def delete_assistant( elif timeout is None: timeout = 600.0 - response: Optional[Dict[str, Any]] = None + response: Optional[AssistantDeleted] = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base @@ -398,6 +447,7 @@ def delete_assistant( organization=organization, assistant_id=assistant_id, client=client, + async_delete_assistants=async_delete_assistants, ) else: raise litellm.exceptions.BadRequestError( diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index f5c7aef48d6..e6056691e73 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -17,6 +17,7 @@ from typing import ( import httpx import openai from openai import AsyncOpenAI, OpenAI +from openai.types.beta.assistant_deleted import AssistantDeleted from pydantic import BaseModel from typing_extensions import overload, override @@ -2440,6 +2441,30 @@ class OpenAIAssistantsAPI(BaseLLM): response = openai_client.beta.assistants.create(**create_assistant_data) return response + # Delete Assistant + async def async_delete_assistant( + self, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + organization: Optional[str], + client: Optional[AsyncOpenAI], + assistant_id: str, + ) -> AssistantDeleted: + openai_client = self.async_get_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + client=client, + ) + + response = await openai_client.beta.assistants.delete(assistant_id=assistant_id) + + return response + def delete_assistant( self, api_key: Optional[str], @@ -2449,17 +2474,18 @@ class OpenAIAssistantsAPI(BaseLLM): organization: Optional[str], assistant_id: str, client=None, - async_create_assistants=None, + async_delete_assistants=None, ): - # if async_create_assistants is not None and async_create_assistants == True: - # return self.async_create_assistants( - # api_key=api_key, - # api_base=api_base, - # timeout=timeout, - # max_retries=max_retries, - # organization=organization, - # client=client, - # ) + if async_delete_assistants is not None and async_delete_assistants == True: + return self.async_delete_assistant( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + client=client, + assistant_id=assistant_id, + ) openai_client = self.get_openai_client( api_key=api_key, api_base=api_base,