Merge pull request #8 from zakhar-kogan/main

Added OpenRouter support

.env.example

@@ -0,0 +1,5 @@
+OPENAI_API_KEY = ""
+COHERE_API_KEY = ""
+OPENROUTER_API_KEY = ""
+OR_SITE_URL = ""
+OR_APP_NAME = "LiteLLM Example app"

.gitignore

@@ -0,0 +1 @@
+.env

build/lib/litellm/main.py

@@ -1,4 +1,7 @@
-import os, openai, cohere
+import os, openai, cohere, dotenv
+
+# Loading env variables using dotenv
+dotenv.load_dotenv()
 
 ####### COMPLETION MODELS ###################
 open_ai_chat_completion_models = [
@@ -13,6 +16,18 @@ cohere_models = [
     'command-nightly',
 ]
 
+openrouter_models = [
+    'google/palm-2-codechat-bison',
+    'google/palm-2-chat-bison',
+    'openai/gpt-3.5-turbo',
+    'openai/gpt-3.5-turbo-16k',
+    'openai/gpt-4-32k',
+    'anthropic/claude-2',
+    'anthropic/claude-instant-v1',
+    'meta-llama/llama-2-13b-chat',
+    'meta-llama/llama-2-70b-chat'
+]
+
 ####### EMBEDDING MODELS ###################
 open_ai_embedding_models = [
     'text-embedding-ada-002'
@@ -34,6 +49,32 @@ def completion(model, messages, azure=False):
       engine=model,
       messages = messages
     )
+  elif "replicate" in model: 
+    prompt = " ".join([message["content"] for message in messages])
+    output = replicate.run(
+      model,
+      input={
+        "prompt": prompt,
+      })
+    print(f"output: {output}")
+    response = ""
+    for item in output: 
+      print(f"item: {item}")
+      response += item
+    new_response = {
+      "choices": [
+        {
+          "finish_reason": "stop",
+          "index": 0,
+          "message": {
+              "content": response,
+              "role": "assistant"
+          }
+        }
+      ]
+    }
+    print(f"new response: {new_response}")
+    response = new_response
   elif model in cohere_models:
     cohere_key = os.environ.get("COHERE_API_KEY")
     co = cohere.Client(cohere_key)
@@ -76,6 +117,22 @@ def completion(model, messages, azure=False):
         model=model,
         prompt = prompt
     )
+
+  elif model in openrouter_models:
+    openai.api_base = "https://openrouter.ai/api/v1"
+    openai.api_key = os.environ.get("OPENROUTER_API_KEY")
+
+    prompt = " ".join([message["content"] for message in messages])
+
+    response = openai.ChatCompletion.create(
+      model=model,
+      messages=messages,
+      headers={ 
+        "HTTP-Referer": os.environ.get("OR_SITE_URL"), # To identify your app
+        "X-Title": os.environ.get("OR_APP_NAME") 
+        },
+    )
+    reply = response.choices[0].message
   return response
 
 
@@ -100,4 +157,3 @@ def embedding(model, input=[], azure=False):
 
 #############################################
 #############################################
-

completion_test.py

@@ -1,4 +1,4 @@
-from main import completion
+from litellm.main import completion
 import os
 
 ## Configs for Models ##
@@ -28,3 +28,7 @@ response = completion("command-nightly", messages)
 print("\nCohere call")
 print(response)
 
+# openrouter call
+response = completion("google/palm-2-codechat-bison", messages)
+print("\OpenRouter call")
+print(response)

litellm/main.py

@@ -1,4 +1,7 @@
-import os, openai, cohere
+import os, openai, cohere, dotenv
+
+# Loading env variables using dotenv
+dotenv.load_dotenv()
 
 ####### COMPLETION MODELS ###################
 open_ai_chat_completion_models = [
@@ -13,6 +16,18 @@ cohere_models = [
     'command-nightly',
 ]
 
+openrouter_models = [
+    'google/palm-2-codechat-bison',
+    'google/palm-2-chat-bison',
+    'openai/gpt-3.5-turbo',
+    'openai/gpt-3.5-turbo-16k',
+    'openai/gpt-4-32k',
+    'anthropic/claude-2',
+    'anthropic/claude-instant-v1',
+    'meta-llama/llama-2-13b-chat',
+    'meta-llama/llama-2-70b-chat'
+]
+
 ####### EMBEDDING MODELS ###################
 open_ai_embedding_models = [
     'text-embedding-ada-002'
@@ -102,6 +117,22 @@ def completion(model, messages, azure=False):
         model=model,
         prompt = prompt
     )
+
+  elif model in openrouter_models:
+    openai.api_base = "https://openrouter.ai/api/v1"
+    openai.api_key = os.environ.get("OPENROUTER_API_KEY")
+
+    prompt = " ".join([message["content"] for message in messages])
+
+    response = openai.ChatCompletion.create(
+      model=model,
+      messages=messages,
+      headers={ 
+        "HTTP-Referer": os.environ.get("OR_SITE_URL"), # To identify your app
+        "X-Title": os.environ.get("OR_APP_NAME") 
+        },
+    )
+    reply = response.choices[0].message
   return response