diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index 16a1f973cea..b02a2bd5977 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -24,7 +24,7 @@ def create_thread( metadata: Optional[dict] = None, tool_resources: Optional[OpenAICreateThreadParamsToolResources] = None, client: Optional[OpenAI] = None, - **kwargs + **kwargs, ) -> Thread: """ - get the llm provider @@ -115,4 +115,85 @@ def create_thread( ### MESSAGES ### + +def add_message( + custom_llm_provider: Literal["openai"], + thread_id: str, + role: Literal["user", "assistant"], + content: str, + attachments: Optional[List[Attachment]] = None, + metadata: Optional[dict] = None, + client: Optional[OpenAI] = None, + **kwargs, +) -> OpenAIMessage: + ### COMMON OBJECTS ### + message_data = MessageData( + role=role, content=content, attachments=attachments, metadata=metadata + ) + optional_params = GenericLiteLLMParams(**kwargs) + + ### TIMEOUT LOGIC ### + timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 + # set timeout for 10 minutes by default + + if ( + timeout is not None + and isinstance(timeout, httpx.Timeout) + and supports_httpx_timeout(custom_llm_provider) == False + ): + read_timeout = timeout.read or 600 + timeout = read_timeout # default 10 min timeout + elif timeout is not None and not isinstance(timeout, httpx.Timeout): + timeout = float(timeout) # type: ignore + elif timeout is None: + timeout = 600.0 + + response: Optional[OpenAIMessage] = None + if custom_llm_provider == "openai": + api_base = ( + optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or os.getenv("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + organization = ( + optional_params.organization + or litellm.organization + or os.getenv("OPENAI_ORGANIZATION", None) + or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 + ) + # set API KEY + api_key = ( + optional_params.api_key + or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there + or litellm.openai_key + or os.getenv("OPENAI_API_KEY") + ) + response = openai_assistants_api.add_message( + thread_id=thread_id, + message_data=message_data, + api_base=api_base, + api_key=api_key, + timeout=timeout, + max_retries=optional_params.max_retries, + organization=organization, + client=client, + ) + else: + raise litellm.exceptions.BadRequestError( + message="LiteLLM doesn't support {} for 'create_thread'. Only 'openai' is supported.".format( + custom_llm_provider + ), + model="n/a", + llm_provider=custom_llm_provider, + response=httpx.Response( + status_code=400, + content="Unsupported provider", + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + ), + ) + + return response + + ### RUNS ### diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index 9cc6d86bb68..01d3bd2f2e7 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -1268,6 +1268,8 @@ class OpenAIAssistantsAPI(BaseLLM): for k, v in received_args.items(): if k == "self" or k == "client": pass + elif k == "api_base" and v is not None: + data["base_url"] = v elif v is not None: data[k] = v openai_client = OpenAI(**data) # type: ignore @@ -1306,10 +1308,10 @@ class OpenAIAssistantsAPI(BaseLLM): self, thread_id: str, message_data: MessageData, - api_key: str, + api_key: Optional[str], api_base: Optional[str], timeout: Union[float, httpx.Timeout], - max_retries: int, + max_retries: Optional[int], organization: Optional[str], client: Optional[OpenAI] = None, ) -> OpenAIMessage: diff --git a/litellm/tests/test_assistants.py b/litellm/tests/test_assistants.py index 58c8c4c1f84..a3acdae187c 100644 --- a/litellm/tests/test_assistants.py +++ b/litellm/tests/test_assistants.py @@ -26,11 +26,11 @@ V0 Scope: """ -def test_create_thread_litellm(): +def test_create_thread_litellm() -> Thread: message: MessageData = {"role": "user", "content": "Hey, how's it going?"} # type: ignore new_thread = create_thread( custom_llm_provider="openai", - messages=[message], + messages=[message], # type: ignore ) assert isinstance( @@ -39,7 +39,19 @@ def test_create_thread_litellm(): return new_thread -test_create_thread_litellm() +def test_add_message_litellm(): + message: MessageData = {"role": "user", "content": "Hey, how's it going?"} # type: ignore + new_thread = test_create_thread_litellm() + + # add message to thread + message: MessageData = {"role": "user", "content": "Hey, how's it going?"} # type: ignore + added_message = litellm.add_message( + thread_id=new_thread.id, custom_llm_provider="openai", **message + ) + + print(f"added message: {added_message}") + + assert isinstance(added_message, Message) def test_create_thread_openai_direct() -> Thread: @@ -68,7 +80,7 @@ def test_create_thread_openai_direct() -> Thread: def test_add_message_openai_direct(): openai_api = OpenAIAssistantsAPI() # create thread - new_thread = test_create_thread() + new_thread = test_create_thread_openai_direct() # add message to thread message: MessageData = {"role": "user", "content": "Hey, how's it going?"} # type: ignore added_message = openai_api.add_message(