Refactor OpenAI API call logic

- Remove litellm import and usage
- Remove conditional logic choosing between OpenAI API and litellm
- Call OpenAI API directly in all cases
- Update token counting to handle Claude models separately from GPT models

devchat/openai/openai_chat.py

@@ -2,7 +2,6 @@
 import os
 from pydantic import BaseModel, Field
 import openai
-from litellm import completion
 from devchat.chat import Chat
 from devchat.utils import get_user_info, user_id
 from .openai_message import OpenAIMessage
@@ -68,22 +67,10 @@ def complete_response(self, prompt: OpenAIPrompt) -> str:
             config_params['function_call'] = 'auto'
         config_params['stream'] = False
 
-        api_key = os.environ.get("OPENAI_API_KEY")
-
-        if api_key.startswith("DC."):
-            response = openai.ChatCompletion.create(
-                messages=prompt.messages,
-                **config_params
-            )
-        else:
-            if config_params["model"].startswith("gpt-"):
-                # call gpt- model by openai api and openai api key
-                response = openai.ChatCompletion.create(
-                    messages=prompt.messages,
-                    **config_params
-                )
-            else:
-                response = completion(messages=prompt.messages, **config_params, api_key=api_key)
+        response = openai.ChatCompletion.create(
+            messages=prompt.messages,
+            **config_params
+        )
         return str(response)
 
     def stream_response(self, prompt: OpenAIPrompt) -> Iterator:
@@ -94,21 +81,8 @@ def stream_response(self, prompt: OpenAIPrompt) -> Iterator:
             config_params['function_call'] = 'auto'
         config_params['stream'] = True
 
-        # read environment variable
-        api_key = os.environ.get("OPENAI_API_KEY")
-
-        if api_key.startswith("DC."):
-            response = openai.ChatCompletion.create(
-                messages=prompt.messages,
-                **config_params
-            )
-        else:
-            if config_params["model"].startswith("gpt-"):
-                # call gpt- model by openai api and openai api key
-                response = openai.ChatCompletion.create(
-                    messages=prompt.messages,
-                    **config_params
-                )
-            else:
-                response = completion(**config_params, messages=prompt.messages, api_key=api_key)
+        response = openai.ChatCompletion.create(
+            messages=prompt.messages,
+            **config_params
+        )
         return response

devchat/utils.py

@@ -11,6 +11,7 @@
 from litellm import token_counter
 
 
+encoding = tiktoken.get_encoding("cl100k_base")
 log_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 
 
@@ -208,7 +209,10 @@ def _count_tokens(encoding: tiktoken.Encoding, string: str) -> int:
 
 def openai_message_tokens(message: dict, model: str) -> int:
     """Returns the number of tokens used by a message."""
-    return token_counter(model=model, text=str(message))
+    if "claude" in model:
+        return token_counter(model=model, text=str(message))
+    else:
+        return len(encoding.encode(str(message)))
 
 
 def openai_response_tokens(message: dict, model: str) -> int: