From 4605924a447f23f65fc036ca9e5ce14683905b54 Mon Sep 17 00:00:00 2001
From: Liang Zhang <68933075+ualiangzhang@users.noreply.github.com>
Date: Tue, 23 Jan 2024 20:13:48 -0700
Subject: [PATCH] [Equations] Add the base64 support to equations_to_latex
 (#759)

## Summary of Changes
Add the base64 support to the `equations_to_latex` endpoint
---
 skema/rest/tests/test_eqn_to_latex.py | 34 +++++++++++++++++--
 skema/rest/workflows.py               | 48 +++++++++++++++++++++++++--
 2 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/skema/rest/tests/test_eqn_to_latex.py b/skema/rest/tests/test_eqn_to_latex.py
index 8ed6339b89d..89cfa81e7b9 100644
--- a/skema/rest/tests/test_eqn_to_latex.py
+++ b/skema/rest/tests/test_eqn_to_latex.py
@@ -1,3 +1,4 @@
+import base64
 from pathlib import Path
 from httpx import AsyncClient
 from skema.rest.workflows import app
@@ -5,8 +6,6 @@
 import json
 
 
-
-
 @pytest.mark.ci_only
 @pytest.mark.asyncio
 async def test_post_image_to_latex():
@@ -21,7 +20,7 @@ async def test_post_image_to_latex():
     endpoint = "/images/equations-to-latex"
     # see https://fastapi.tiangolo.com/advanced/async-tests/#async-tests
     async with AsyncClient(app=app, base_url="http://eqn-to-latex-test") as ac:
-      response = await ac.post(endpoint, files=files)
+        response = await ac.post(endpoint, files=files)
     expected = "\\frac{d H}{dt}=\\nabla \\cdot {(\\Gamma*H^{n+2}*\\left|\\nabla{H}\\right|^{n-1}*\\nabla{H})}"
     # check for route's existence
     assert (
@@ -35,3 +34,32 @@ async def test_post_image_to_latex():
     assert (
         json.loads(response.text) == expected
     ), f"Response should be {expected}, but instead received {response.text}"
+
+
+@pytest.mark.ci_only
+@pytest.mark.asyncio
+async def test_post_image_to_latex_base64():
+    """Test case for /images/base64/equations-to-latex endpoint."""
+    cwd = Path(__file__).parents[0]
+    image_path = cwd / "data" / "img2latex" / "halfar.png"
+    with Path(image_path).open("rb") as infile:
+        img_bytes = infile.read()
+    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
+
+    endpoint = "/images/base64/equations-to-latex"
+    # see https://fastapi.tiangolo.com/advanced/async-tests/#async-tests
+    async with AsyncClient(app=app, base_url="http://eqn-to-latex-base64-test") as ac:
+        response = await ac.post(endpoint, data=img_b64)
+    expected = "\\frac{d H}{dt}=\\nabla \\cdot {(\\Gamma*H^{n+2}*\\left|\\nabla{H}\\right|^{n-1}*\\nabla{H})}"
+    # check for route's existence
+    assert (
+        any(route.path == endpoint for route in app.routes) == True
+    ), "{endpoint} does not exist for app"
+    # check status code
+    assert (
+        response.status_code == 200
+    ), f"Request was unsuccessful (status code was {response.status_code} instead of 200)"
+    # check response
+    assert (
+        json.loads(response.text) == expected
+    ), f"Response should be {expected}, but instead received {response.text}"
\ No newline at end of file
diff --git a/skema/rest/workflows.py b/skema/rest/workflows.py
index c95c8d43cb3..5f0e53745db 100644
--- a/skema/rest/workflows.py
+++ b/skema/rest/workflows.py
@@ -12,11 +12,11 @@
 import json
 import requests
 
-from fastapi import APIRouter, Depends, File, UploadFile, FastAPI
+from fastapi import APIRouter, Depends, File, UploadFile, FastAPI, Request
 from starlette.responses import JSONResponse
 
 from skema.img2mml import eqn2mml
-from skema.img2mml.eqn2mml import image2mathml_db
+from skema.img2mml.eqn2mml import image2mathml_db, b64_image_to_mml
 from skema.img2mml.api import get_mathml_from_bytes
 from skema.rest import config, schema, utils, llm_proxy
 from skema.rest.proxies import SKEMA_RS_ADDESS
@@ -112,6 +112,48 @@ async def equations_to_latex(data: UploadFile, client: httpx.AsyncClient = Depen
         return f"Error: {response.status_code} {response.text}"
 
 
+# equation images -> base64 -> mml -> latex
+@router.post("/images/base64/equations-to-latex", summary="Equations (images) → MML → LaTeX")
+async def equations_to_latex(request: Request, client: httpx.AsyncClient = Depends(utils.get_client)):
+    """
+    Converts images of equations to LaTeX.
+
+    ### Python example
+
+    Endpoint for generating LaTeX from an input image.
+
+    ```
+    from pathlib import Path
+    import base64
+    import requests
+
+    url = "http://127.0.0.1:8000/workflows/images/base64/equations-to-latex"
+    with Path("test.png").open("rb") as infile:
+      img_bytes = infile.read()
+    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
+    r = requests.post(url, data=img_b64)
+    print(r.text)
+    ```
+    """
+    # Read image data
+    img_b64 = await request.body()
+    mml_res = b64_image_to_mml(img_b64)
+
+    # pass image bytes to get_mathml_from_bytes function
+    proxy_url = f"{SKEMA_RS_ADDESS}/mathml/latex"
+    print(f"MML:\t{mml_res}")
+    print(f"Proxying request to {proxy_url}")
+    response = await client.post(proxy_url, data=mml_res)
+    # Check the response
+    if response.status_code == 200:
+        # The request was successful
+        return response.text
+    else:
+        # The request failed
+        print(f"Error: {response.status_code}")
+        print(response.text)
+        return f"Error: {response.status_code} {response.text}"
+
 # tex equations -> pmml -> amr
 @router.post("/latex/equations-to-amr", summary="Equations (LaTeX) → pMML → AMR")
 async def equations_to_amr(data: schema.EquationLatexToAMR, client: httpx.AsyncClient = Depends(utils.get_client)):
@@ -162,7 +204,7 @@ async def equations_to_amr(data: schema.MmlToAMR, client: httpx.AsyncClient = De
     return res.json()
 
 
-# code snippets -> fn -> petrinet amr 
+# code snippets -> fn -> petrinet amr
 @router.post("/code/snippets-to-pn-amr", summary="Code snippets → PetriNet AMR")
 async def code_snippets_to_pn_amr(system: code2fn.System, client: httpx.AsyncClient = Depends(utils.get_client)):
     gromet = await code2fn.fn_given_filepaths(system)