From 228669caa1b5938a2f94ccfd220cf58c0bcc997b Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 23 Aug 2023 10:02:13 -0700
Subject: [PATCH] adding support for finetuned completion models

---
 litellm/__init__.py | 12 +++++++++++-
 litellm/main.py     |  4 +++-
 pyproject.toml      |  2 +-
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index c88d6e698a..39a3f89b90 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -24,6 +24,16 @@ caching = False
 caching_with_models = False  # if you want the caching key to be model + prompt
 debugger = False
 model_cost = {
+    "babbage-002": {
+        "max_tokens": 16384,
+        "input_cost_per_token": 0.0000004,
+        "output_cost_per_token": 0.0000004,
+    },
+    "davinci-002": {
+        "max_tokens": 16384,
+        "input_cost_per_token": 0.000002,
+        "output_cost_per_token": 0.000002,
+    },
     "gpt-3.5-turbo": {
         "max_tokens": 4000,
         "input_cost_per_token": 0.0000015,
@@ -137,7 +147,7 @@ open_ai_chat_completion_models = [
     "gpt-3.5-turbo-0613",
     "gpt-3.5-turbo-16k-0613",
 ]
-open_ai_text_completion_models = ["text-davinci-003"]
+open_ai_text_completion_models = ["text-davinci-003", "babbage-002", "davinci-002"]
 
 cohere_models = [
     "command-nightly",
diff --git a/litellm/main.py b/litellm/main.py
index 87a41b8af4..0e54af94e3 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -248,7 +248,9 @@ def completion(
                 original_response=response,
                 additional_args={"headers": litellm.headers},
             )
-        elif model in litellm.open_ai_text_completion_models:
+        elif (model in litellm.open_ai_text_completion_models or 
+              "ft:babbage-002" in model or # support for finetuned completion models
+              "ft:davinci-002" in model):
             openai.api_type = "openai"
             openai.api_base = (
                 litellm.api_base
diff --git a/pyproject.toml b/pyproject.toml
index e21c91218a..94a7262f0d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.457"
+version = "0.1.458"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"