From d0bf79758d0ed5b52dd5ab4242d258d1d521ff96 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 5 Oct 2023 11:50:09 -0700
Subject: [PATCH] v0 add litellm

---
 main.py        | 5 +++--
 pyproject.toml | 1 +
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/main.py b/main.py
index b772289..570f09d 100644
--- a/main.py
+++ b/main.py
@@ -2,6 +2,7 @@
 import json
 import re
 import openai
+import litellm
 import requests
 from bs4 import BeautifulSoup
 from graphviz import Digraph
@@ -20,7 +21,7 @@
 app = Flask(__name__)
 
 # Set your OpenAI API key
-openai.api_key = os.getenv("OPENAI_API_KEY")
+litellm.api_key = os.getenv("OPENAI_API_KEY")
 response_data = ""
 
 # If Neo4j credentials are set, then Neo4j is used to store information
@@ -130,7 +131,7 @@ def get_response_data():
         return jsonify({"error": "No input provided"}), 400
     print("starting openai call")
     try:
-        completion: KnowledgeGraph = openai.ChatCompletion.create(
+        completion: KnowledgeGraph = litellm.completion(
             model="gpt-3.5-turbo-16k",
             messages=[
                 {
diff --git a/pyproject.toml b/pyproject.toml
index 93e94cc..9fe255b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,6 +7,7 @@ authors = ["Your Name <you@example.com>"]
 [tool.poetry.dependencies]
 python = ">=3.10.0,<3.11"
 openai = "^0.28.0"
+litellm = ">=0.1.820"
 bs4 = "^0.0.1"
 matplotlib = "^3.7.2"
 networkx = "^3.1"