v0 litellm (#5076)

facebookresearch · Nov 3, 2023 · 6c80241 · 6c80241
1 parent c437a91
commit 6c80241
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/parlai/agents/openai_chat_completions/openai_chat_completions.py b/parlai/agents/openai_chat_completions/openai_chat_completions.py
@@ -11,9 +11,9 @@
 from parlai.core.message import Message
 
 try:
-    import openai
+    import litellm
 except ImportError:
-    raise ImportError('Please run `pip install openai`.')
+    raise ImportError('Please run `pip install litellm`.')
 
 
 class OpenaiChatCompletionsAgent(Agent):
@@ -127,7 +127,7 @@ def __init__(self, opt, shared=None):
         if self.init_prompt:
             self.turns.append({'role': 'system', 'content': self.init_prompt})
 
-        openai.api_key = opt.get('openai_api_key')
+        litellm.api_key = opt.get('openai_api_key')
 
     def reset(self):
         """
@@ -169,7 +169,7 @@ def act(self):
         return Message({'id': self.getID(), 'text': resp_txt, 'episode_done': False})
 
     def query_chat_completion_api(self):
-        response = openai.ChatCompletion.create(
+        response = litellm.completion(
             model=self.model_name,
             messages=self.turns,
             temperature=self.temperature,

diff --git a/requirements.txt b/requirements.txt
@@ -60,3 +60,4 @@ contractions~=0.1.72
 fsspec~=2022.2.0
 google-api-core<=2.11.0
 openai<=0.27.7
+litellm>=0.1.400