Updates for new openai API

Author: ajhai

llmstack/play/actors/agent.py

@@ -1,13 +1,14 @@
-import importlib
 import logging
 import time
 import uuid
-import orjson as json
 from typing import Any
-from jinja2 import Template
 
+import orjson as json
 from asgiref.sync import async_to_sync
+from jinja2 import Template
+from openai import OpenAI
 from pydantic import BaseModel
+
 from llmstack.play.actor import Actor, BookKeepingData
 from llmstack.play.actors.output import OutputResponse
 from llmstack.play.output_stream import Message, MessageType
@@ -54,6 +55,10 @@ def __init__(self, output_stream, processor_configs, dependencies=[], all_depend
         self._input = kwargs.get('input')
         self._config = kwargs.get('config', {})
 
+        self._openai_client = OpenAI(
+            api_key=self._env['openai_api_key']
+        )
+
         self._agent_messages = [{
             'role': 'system',
             'content': self._config.get('system_message', 'You are a helpful assistant that uses provided tools to perform actions.')
@@ -94,8 +99,6 @@ def _on_error(self, message) -> None:
         )
 
     def on_receive(self, message: Message) -> Any:
-        import openai
-        importlib.reload(openai)
         max_steps = self._config.get('max_steps', 10) + 2
 
         if len(self._agent_messages) > max_steps:
@@ -123,14 +126,12 @@ def on_receive(self, message: Message) -> Any:
 
             model = self._config.get('model', 'gpt-3.5-turbo')
 
-            openai.api_key = self._env['openai_api_key']
-
             # Make one call to the model
             full_content = ''
             function_name = ''
             function_args = ''
             finish_reason = None
-            result = openai.ChatCompletion.create(
+            result = self._openai_client.chat.completions.create(
                 model=model,
                 messages=self._agent_messages,
                 stream=True,
@@ -139,30 +140,31 @@ def on_receive(self, message: Message) -> Any:
             agent_message_id = str(uuid.uuid4())
 
             for data in result:
-                if data.get('object') and data.get('object') == 'chat.completion.chunk' and data.get('choices') and len(data.get('choices')) > 0:
-                    finish_reason = data['choices'][0]['finish_reason']
-                    delta = data['choices'][0]['delta']
-                    function_call = delta.get('function_call')
-                    content = delta.get('content')
-
-                    if function_call and function_call.get('name'):
-                        function_name += function_call['name']
+                logger.info(data)
+                if data.object == 'chat.completion.chunk' and len(data.choices) > 0 and data.choices[0].delta:
+                    finish_reason = data.choices[0].finish_reason
+                    delta = data.choices[0].delta
+                    function_call = delta.function_call
+                    content = delta.content
+
+                    if function_call and function_call.name:
+                        function_name += function_call.name
                         async_to_sync(self._output_stream.write)(
                             AgentOutput(
                                 content=FunctionCall(
-                                    name=function_call['name'],
+                                    name=function_call.name,
                                 ),
                                 id=agent_message_id,
                                 from_id='agent',
                                 type='step',
                             )
                         )
-                    elif function_call and function_call.get('arguments'):
-                        function_args += function_call['arguments']
+                    elif function_call and function_call.arguments:
+                        function_args += function_call.arguments
                         async_to_sync(self._output_stream.write)(
                             AgentOutput(
                                 content=FunctionCall(
-                                    arguments=function_call['arguments'],
+                                    arguments=function_call.arguments,
                                 ),
                                 id=agent_message_id,
                                 from_id='agent',

llmstack/processors/providers/openai/chat_completions_vision.py

@@ -1,18 +1,24 @@
 import importlib
-import openai
-from enum import Enum
 import logging
+from enum import Enum
 from typing import Annotated, List, Literal, Optional, Union
 
+import openai
 from asgiref.sync import async_to_sync
+from openai import OpenAI
 from pydantic import BaseModel, Field, confloat, conint
 
-from llmstack.common.blocks.llm.openai import OpenAIChatCompletionsAPIProcessorConfiguration
-from llmstack.processors.providers.api_processor_interface import ApiProcessorInterface, ApiProcessorSchema
-from llmstack.processors.providers.api_processor_interface import ApiProcessorInterface, ApiProcessorSchema
+from llmstack.common.blocks.llm.openai import (
+    OpenAIChatCompletionsAPIProcessorConfiguration,
+)
+from llmstack.processors.providers.api_processor_interface import (
+    ApiProcessorInterface,
+    ApiProcessorSchema,
+)
 
 logger = logging.getLogger(__name__)
 
+
 class Role(str, Enum):
     SYSTEM = 'system'
     USER = 'user'
@@ -21,42 +27,53 @@ class Role(str, Enum):
     def __str__(self):
         return self.value
 
+
 class ChatCompletionsVisionModel(str, Enum):
     GPT_4_Vision = 'gpt-4-vision-preview'
 
     def __str__(self):
         return self.value
-    
+
+
 class TextMessage(BaseModel):
     type: Literal["text"]
-    
+
     text: str = Field(
         default='', description='The message text.')
-    
+
+
 class UrlImageMessage(BaseModel):
     type: Literal["image_url"]
-    
+
     image_url: str = Field(
         default='', description='The image data URI.')
 
-Message = Annotated[Union[TextMessage, UrlImageMessage], Field(discriminator='type')]
+
+Message = Annotated[Union[TextMessage, UrlImageMessage],
+                    Field(discriminator='type')]
+
+
 class ChatMessage(ApiProcessorSchema):
     role: Optional[Role] = Field(
         default=Role.USER, description="The role of the message sender. Can be 'user' or 'assistant' or 'system'.",
     )
-    content: List[Union[TextMessage, UrlImageMessage]] = Field(default=[], description='The message text.')
-    
+    content: List[Union[TextMessage, UrlImageMessage]] = Field(
+        default=[], description='The message text.')
+
+
 class ChatCompletionsVisionInput(ApiProcessorSchema):
     system_message: Optional[str] = Field(
         default='', description='A message from the system, which will be prepended to the chat history.', widget='textarea',
     )
     messages:  List[Message] = Field(
         default=[], description='A list of messages, each with a role and message text.'
     )
-    
+
+
 class ChatCompletionsVisionOutput(ApiProcessorSchema):
     result: str = Field(default='', description='The model-generated message.')
-    
+
+
 class ChatCompletionsVisionConfiguration(OpenAIChatCompletionsAPIProcessorConfiguration, ApiProcessorSchema):
     model: ChatCompletionsVisionModel = Field(
         default=ChatCompletionsVisionModel.GPT_4_Vision,
@@ -83,6 +100,7 @@ class ChatCompletionsVisionConfiguration(OpenAIChatCompletionsAPIProcessorConfig
         default=False, description="Automatically prune chat history. This is only applicable if 'retain_history' is set to 'true'.",
     )
 
+
 class ChatCompletionsVision(ApiProcessorInterface[ChatCompletionsVisionInput, ChatCompletionsVisionOutput, ChatCompletionsVisionConfiguration]):
     """
     OpenAI Chat Completions with vision API
@@ -114,32 +132,32 @@ def session_data_to_persist(self) -> dict:
     def process(self) -> dict:
         importlib.reload(openai)
         output_stream = self._output_stream
-                
+
         chat_history = self._chat_history if self._config.retain_history else []
         messages = []
-        messages.append({'role': 'system', 'content': self._input.system_message})
-        
+        messages.append(
+            {'role': 'system', 'content': self._input.system_message})
+
         for msg in chat_history:
             messages.append(msg)
-            
-        messages.append({'role': 'user', 'content': [msg.dict() for msg in self._input.messages]})
-
-        openai.api_key = self._env['openai_api_key']
-        result = openai.chat.completions.create(
-                model=self._config.model,
-                messages=messages,
-                temperature=self._config.temperature,
-                stream=True,
-            ) 
-        
-        
+
+        messages.append({'role': 'user', 'content': [
+                        msg.dict() for msg in self._input.messages]})
+
+        openai_client = OpenAI(api_key=self._env['openai_api_key'])
+        result = openai_client.chat.completions.create(
+            model=self._config.model,
+            messages=messages,
+            temperature=self._config.temperature,
+            stream=True,
+        )
+
         for data in result:
-            if data.get('object') and data.get('object') == 'chat.completion.chunk' and data.get('choices') and len(data.get('choices')) > 0 and data['choices'][0].get('delta') and data['choices'][0]['delta'].get('content'):
+            if data.object == 'chat.completion.chunk' and len(data.choices) > 0 and data.choices[0].delta and data.choices[0].delta.content:
                 async_to_sync(output_stream.write)(
                     ChatCompletionsVisionOutput(
-                        result=data['choices'][0]['delta']['content']
+                        result=data.choices[0].delta.content
                     ))
-        
 
         output = self._output_stream.finalize()
 

llmstack/processors/providers/promptly/text_chat.py

@@ -3,17 +3,19 @@
 import logging
 import uuid
 from enum import Enum
-from typing import Optional
-from typing import List
+from typing import List, Optional
 
-import openai
 from asgiref.sync import async_to_sync
 from django import db
+from openai import AzureOpenAI, OpenAI
 from pydantic import Field
 
 from llmstack.datasources.models import DataSource
 from llmstack.datasources.types import DataSourceTypeFactory
-from llmstack.processors.providers.api_processor_interface import ApiProcessorInterface, ApiProcessorSchema
+from llmstack.processors.providers.api_processor_interface import (
+    ApiProcessorInterface,
+    ApiProcessorSchema,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -188,7 +190,6 @@ def fetch_datasource_docs(datasource_uuid):
 
     def process(self) -> dict:
         input = self._input.dict()
-        importlib.reload(openai)
         output_stream = self._output_stream
         docs = self._search_datasources(input)
 
@@ -227,10 +228,12 @@ def process(self) -> dict:
         )
 
         if self._env['azure_openai_api_key'] and self._config.use_azure_if_available:
-            openai.api_type = 'azure'
-            openai.api_key = self._env['azure_openai_api_key']
-            openai.api_base = f"https://{self._env['azure_openai_endpoint']}.openai.azure.com"
-            openai.api_version = '2023-03-15-preview'
+            openai_client = AzureOpenAI(
+                api_key=self._env['azure_openai_api_key'],
+                api_version='2023-03-15-preview',
+                azure_endpoint=self._env['azure_openai_endpoint'],
+            )
+
             model = self._config.dict().get('model', 'gpt-3.5-turbo')
             engine = 'gpt-4'
             if model == 'gpt-3.5-turbo':
@@ -240,30 +243,35 @@ def process(self) -> dict:
             elif model == 'gpt-4-32k':
                 engine = 'gpt-4-32k'
 
-            result = openai.ChatCompletion.create(
+            result = openai_client.chat.completions.create(
                 engine=engine,
                 messages=[system_message] +
                 [context_message] + self._chat_history,
                 temperature=self._config.temperature,
                 stream=True,
             )
         elif self._env['localai_base_url'] and self._config.use_localai_if_available:
-            if self._env['localai_api_key']:
-                openai.api_key = self._env['localai_api_key']
-            openai.api_base = self._env['localai_base_url']
+            openai_client = OpenAI(
+                api_key=self._env['localai_api_key'],
+                base_url=self._env['localai_base_url'],
+            ) if self._env['localai_api_key'] else OpenAI(
+                base_url=self._env['localai_base_url'],
+            )
             model = self._config.dict().get('model', 'gpt-3.5-turbo')
 
-            result = openai.ChatCompletion.create(
+            result = openai_client.chat.completions.create(
                 model=model,
                 messages=[system_message] +
                 [context_message] + self._chat_history,
                 temperature=self._config.temperature,
                 stream=True,
             )
         elif self._env['openai_api_key'] is not None:
-            openai.api_key = self._env['openai_api_key']
+            openai_client = OpenAI(
+                api_key=self._env['openai_api_key'],
+            )
             model = self._config.dict().get('model', 'gpt-3.5-turbo')
-            result = openai.ChatCompletion.create(
+            result = openai_client.chat.completions.create(
                 model=model,
                 messages=[system_message] +
                 [context_message] + self._chat_history,
@@ -274,10 +282,10 @@ def process(self) -> dict:
             raise Exception('No OpenAI API key provided')
 
         for data in result:
-            if data.get('object') and data.get('object') == 'chat.completion.chunk' and data.get('choices') and len(data.get('choices')) > 0 and data['choices'][0].get('delta') and data['choices'][0]['delta'].get('content'):
+            if data.object == 'chat.completion.chunk' and len(data.choices) > 0 and data.choices[0].delta and data.choices[0].delta.content:
                 async_to_sync(output_stream.write)(
                     TextChatOutput(
-                        answer=data['choices'][0]['delta']['content']
+                        answer=data.choices[0].delta.content
                     ))
 
         if len(docs) > 0: