From 3bed01f5c17fe50d7ad6495285d6a83e5f665a06 Mon Sep 17 00:00:00 2001
From: Krish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 27 Jul 2023 20:44:52 -0700
Subject: [PATCH] Update main.py

---
 litellm/main.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/litellm/main.py b/litellm/main.py
index cbefea5a9c..5f325f2590 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -34,6 +34,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)