human cleanup

vibe_bot/llama_wrapper.py

@@ -5,9 +5,12 @@
 import openai
 from typing import Iterable
 from openai.types.chat import ChatCompletionMessageParam
+from openai._types import FileTypes, SequenceNotStr
+from typing import Union
+from io import BufferedReader, BytesIO
 
 
-def chat_completion_think(
+def chat_completion(
     system_prompt: str,
     user_prompt: str,
     openai_url: str,
@@ -80,35 +83,56 @@ def chat_completion_instruct(
         return ""
 
 
-def image_generation(prompt: str, n=1) -> str:
-    client = openai.OpenAI(base_url=OPENAI_API_IMAGE_ENDPOINT, api_key="placeholder")
+def image_generation(prompt: str, openai_url: str, openai_api_key: str, n=1) -> str:
+    """Generates an image using the given prompt and returns the base64 encoded image data
+
+    Returns:
+        str: The base64 encoded image data. Decode and write to a file.
+    """
+    client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key)
     response = client.images.generate(
         prompt=prompt,
         n=n,
         size="1024x1024",
     )
     if response.data:
-        return response.data[0].url
+        return response.data[0].b64_json or ""
     else:
         return ""
 
 
-def image_edit(image, mask, prompt, n=1, size="1024x1024"):
-    client = openai.OpenAI(base_url=OPENAI_API_EDIT_ENDPOINT, api_key="placeholder")
+def image_edit(
+    image: BufferedReader | BytesIO,
+    prompt: str,
+    openai_url: str,
+    openai_api_key: str,
+    n=1,
+) -> str:
+    client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key)
     response = client.images.edit(
         image=image,
-        mask=mask,
         prompt=prompt,
         n=n,
-        size=size,
+        size="1024x1024",
     )
-    return response.data[0].url
+    if response.data:
+        return response.data[0].b64_json or ""
+    else:
+        return ""
 
 
-def embeddings(text, model="text-embedding-3-small"):
-    client = openai.OpenAI(base_url=OPENAI_API_EMBED_ENDPOINT, api_key="placeholder")
+def embedding(
+    text: str, openai_url: str, openai_api_key: str, model: str
+) -> list[float]:
+    client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key)
     response = client.embeddings.create(
-        input=text,
-        model=model,
+        input=[text], model=model, encoding_format="float"
     )
-    return response.data[0].embedding
+    if response:
+        raw_data = response[0].embedding  # type: ignore
+        # The result could be an array of floats or an array of an array of floats.
+        try:
+            return raw_data[0]
+        except Exception:
+            return raw_data
+    return []