Forge - LLM functions

forge/autogpt/sdk/llm.py

@@ -0,0 +1,69 @@
+import typing
+
+import openai
+from tenacity import retry, stop_after_attempt, wait_random_exponential
+
+from .forge_log import ForgeLogger
+
+LOG = ForgeLogger(__name__)
+
+
+@retry(wait=wait_random_exponential(min=1, max=40), stop=stop_after_attempt(3))
+def chat_completion_request(
+    messages, functions=None, function_call=None, model=str, custom_labels=None
+) -> typing.Union[typing.Dict[str, typing.Any], Exception]:
+    """Generate a response to a list of messages using OpenAI's API"""
+    try:
+        kwargs = {
+            "model": model,
+            "messages": messages,
+        }
+
+        if functions:
+            kwargs["functions"] = functions
+
+        if function_call:
+            kwargs["function_call"] = function_call
+
+        if custom_labels:
+            kwargs["headers"] = {}
+            for label in custom_labels.keys():
+                # This is an example showing adding in the labels as helicone properties
+                kwargs["headers"][f"Helicone-Property-{label}"] = custom_labels[label]
+
+        resp = openai.ChatCompletion.create(**kwargs)
+
+        return resp
+    except Exception as e:
+        LOG.error("Unable to generate ChatCompletion response")
+        LOG.error(f"Exception: {e}")
+        raise
+
+
+@retry(wait=wait_random_exponential(min=1, max=40), stop=stop_after_attempt(3))
+async def create_embedding_request(
+    messages, model="text-embedding-ada-002"
+) -> typing.Union[typing.Dict[str, typing.Any], Exception]:
+    """Generate an embedding for a list of messages using OpenAI's API"""
+    try:
+        return await openai.Embedding.acreate(
+            input=[f"{m['role']}: {m['content']}" for m in messages],
+            engine=model,
+        )
+    except Exception as e:
+        LOG.error("Unable to generate ChatCompletion response")
+        LOG.error(f"Exception: {e}")
+        raise
+
+
+@retry(wait=wait_random_exponential(min=1, max=40), stop=stop_after_attempt(3))
+async def transcribe_audio(
+    audio_file: str,
+) -> typing.Union[typing.Dict[str, typing.Any], Exception]:
+    """Transcribe an audio file using OpenAI's API"""
+    try:
+        return await openai.Audio.transcribe(model="whisper-1", file=audio_file)
+    except Exception as e:
+        LOG.error("Unable to generate ChatCompletion response")
+        LOG.error(f"Exception: {e}")
+        raise