From c14b757e8cbe1c512af94f9853a8cb96300a3e0d Mon Sep 17 00:00:00 2001
From: abhijitjavelin <abhijit.hema@getjavelin.io>
Date: Thu, 8 May 2025 16:18:25 +0530
Subject: [PATCH] feat: img gen support for unified endpoints

---
 examples/openai/img_generations_example.py | 58 ++++++++++++++++++++++
 javelin_sdk/client.py                      |  6 +++
 2 files changed, 64 insertions(+)
 create mode 100644 examples/openai/img_generations_example.py

diff --git a/examples/openai/img_generations_example.py b/examples/openai/img_generations_example.py
new file mode 100644
index 0000000..d9ab8a3
--- /dev/null
+++ b/examples/openai/img_generations_example.py
@@ -0,0 +1,58 @@
+import base64
+from openai import OpenAI
+from javelin_sdk import JavelinClient, JavelinConfig
+import os
+import dotenv
+
+dotenv.load_dotenv()
+
+# Load API keys from environment variables
+JAVELIN_API_KEY = os.getenv("JAVELIN_API_KEY")
+LLM_API_KEY = os.getenv("LLM_API_KEY")
+BASE_URL = os.getenv("BASE_URL")
+
+# Configure Javelin
+config = JavelinConfig(
+    base_url=BASE_URL,
+    javelin_api_key=JAVELIN_API_KEY,
+    llm_api_key=LLM_API_KEY,
+)
+javelin_client = JavelinClient(config)
+
+client = OpenAI(api_key=LLM_API_KEY)
+route_name = "openai_univ"  # define your universal route name here
+javelin_client.register_openai(client, route_name=route_name)
+
+# --- Example 1: Edit an image ---
+# result = client.images.edit(
+#     model="gpt-image-1",
+#     image=open("examples/dog.png", "rb"),
+#     prompt="an angry dog"
+# )
+# image_base64 = result.data[0].b64_json
+# image_bytes = base64.b64decode(image_base64)
+# with open("angry_dog_2.png", "wb") as f:
+#     f.write(image_bytes)
+
+# --- Example 2: Create image variations ---
+# response = client.images.create_variation(
+#     image=open("examples/dog.png", "rb"),
+#     n=2,
+#     size="1024x1024"
+# )
+# for idx, img_data in enumerate(response.data):
+#     image_bytes = base64.b64decode(img_data.b64_json)
+#     with open(f"dog_variation_{idx+1}.png", "wb") as f:
+#         f.write(image_bytes)
+
+# --- Example 3: Generate an image ---
+img = client.images.generate(
+    model="gpt-image-1",
+    prompt="A friendly dog playing in a park.",
+    n=1,
+    size="1024x1024"
+)
+
+image_bytes = base64.b64decode(img.data[0].b64_json)
+with open("generated_image.png", "wb") as f:
+    f.write(image_bytes)
diff --git a/javelin_sdk/client.py b/javelin_sdk/client.py
index 2a73581..f8843d8 100644
--- a/javelin_sdk/client.py
+++ b/javelin_sdk/client.py
@@ -71,6 +71,9 @@ class JavelinClient:
         "chat.completions.create": "chat",
         "completions.create": "text_completion",
         "embeddings.create": "embeddings",
+        "images.generate": "image_generation",
+        "images.edit": "image_editing",
+        "images.create_variation": "image_variation",
     }
 
     def __init__(self, config: JavelinConfig) -> None:
@@ -205,6 +208,9 @@ def register_provider(
                 "chat_completions_create": openai_client.chat.completions.create,
                 "completions_create": openai_client.completions.create,
                 "embeddings_create": openai_client.embeddings.create,
+                "images_generate": openai_client.images.generate,
+                "images_edit": openai_client.images.edit,
+                "images_create_variation": openai_client.images.create_variation,
             }
 
         # Patch methods with tracing and header updates