use liteLLM in the generator

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 ]
 dependencies = [
   "jinja2",
-  "openai",
+  "litellm",
 ]
 
 [project.urls]

src/banks/extensions/generate.py

@@ -1,15 +1,12 @@
 # SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <[email protected]>
 #
 # SPDX-License-Identifier: MIT
-import openai
+from typing import cast
+
+from litellm import completion, ModelResponse
 from jinja2 import nodes
 from jinja2.ext import Extension
 
-CHAT_MODELS = [
-    "gpt-4",
-    "gpt-4-32k",
-    "gpt-3.5-turbo",
-]
 DEFAULT_MODEL = "gpt-3.5-turbo"
 
 
@@ -45,13 +42,10 @@ def _generate(self, text, model_name=DEFAULT_MODEL):
 
         To tweak the prompt used to generate content, change the variable `messages` .
         """
-        content = openai.ChatCompletion.create(
-            model=model_name,
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": text},
-            ],
-            temperature=0.5,
-        )["choices"][0]["message"]["content"]
-
-        return content
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": text},
+        ]
+
+        response: ModelResponse = cast(ModelResponse, completion(model=model_name, messages=messages))
+        return response["choices"][0]["message"]["content"]

tests/test_run_prompt.py

@@ -17,5 +17,4 @@ def test_run_prompt_process():
     p = Prompt.from_template("run_prompt_process.jinja")
     env.extensions["banks.extensions.generate.GenerateExtension"]._generate = mock.MagicMock(return_value="foo bar baz")
 
-    # print(p.text({"topic": "climate change"}))
     assert p.text({"topic": "climate change"}) == "FOO BAR BAZ"