Implement LLM router (#95)

skyvern/forge/sdk/api/llm/utils.py

@@ -0,0 +1,45 @@
+import base64
+from typing import Any
+
+import commentjson
+import litellm
+
+from skyvern.forge.sdk.api.llm.exceptions import EmptyLLMResponseError, InvalidLLMResponseFormat
+
+
+async def llm_messages_builder(
+    prompt: str,
+    screenshots: list[bytes] | None = None,
+) -> list[dict[str, Any]]:
+    messages: list[dict[str, Any]] = [
+        {
+            "type": "text",
+            "text": prompt,
+        }
+    ]
+
+    if screenshots:
+        for screenshot in screenshots:
+            encoded_image = base64.b64encode(screenshot).decode("utf-8")
+            messages.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{encoded_image}",
+                    },
+                }
+            )
+
+    return [{"role": "user", "content": messages}]
+
+
+def parse_api_response(response: litellm.ModelResponse) -> dict[str, str]:
+    try:
+        content = response.choices[0].message.content
+        content = content.replace("```json", "")
+        content = content.replace("```", "")
+        if not content:
+            raise EmptyLLMResponseError(str(response))
+        return commentjson.loads(content)
+    except Exception as e:
+        raise InvalidLLMResponseFormat(str(response)) from e