Code2AMR++ (#5)

api/Dockerfile

@@ -4,7 +4,6 @@ WORKDIR /
 
 COPY pyproject.toml pyproject.toml
 COPY poetry.lock poetry.lock
-COPY api api
 COPY README.md README.md
 
 FROM python:3.10
@@ -16,6 +15,7 @@ RUN pip install --no-cache-dir poetry==1.5.1
 RUN poetry config virtualenvs.create false && \
  poetry install --no-root --no-cache --extras api
 
+COPY api api
 WORKDIR /api
 
 EXPOSE 8000

api/server.py

@@ -5,10 +5,14 @@
 import pypdf
 
 from typing import List, Optional
+from enum import Enum
 
 from fastapi import FastAPI, Response, status, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 
+class EquationType(Enum):
+ LATEX = "latex"
+ MATHML = "mathml"
 
 def build_api(*args) -> FastAPI:
  api = FastAPI(
@@ -48,30 +52,57 @@ def get_status(extraction_job_id: str):
  return {"status": status, "result": result}
 
 
-@app.post("/mathml_to_amr")
-def mathml_to_amr(payload: List[str], model: str = "petrinet"):
- """Post MathML to skema service to get AMR return
+@app.post("/equations_to_amr")
+def equations_to_amr(payload: List[str], 
+ equation_type: EquationType, 
+ model: str = "petrinet", 
+ name: Optional[str] = None, 
+ description: Optional[str] = None):
+ """Post equations and store an AMR to TDS
 
  Args:
- payload (List[str]): A list of MathML strings representing the functions that are used to convert to AMR
+ ```
+ payload (List[str]): A list of Latex or MathML strings representing the functions that are used to convert to AMR
+ equation_type (str): [latex, mathml]
  model (str, optional): AMR model return type. Defaults to "petrinet". Options: "regnet", "petrinet".
+ name (str, optional): the name to set on the newly created model
+ description (str, optional): the description to set on the newly created model 
+ ```
  """
  from utils import create_job
-
- operation_name = "operations.put_mathml_to_skema"
- options = {"mathml": payload, "model": model}
+
+ operation_name = "operations.equations_to_amr"
+ options = {"equations": payload, 
+ "equation_type": equation_type.value, 
+ "model": model,
+ "name": name,
+ "description": description}
 
  resp = create_job(operation_name=operation_name, options=options)
 
  return resp
 
-
 @app.post("/code_to_amr")
-def code_to_amr(artifact_id: str):
+def code_to_amr(artifact_id: str, 
+ name: Optional[str] = None, 
+ description: Optional[str] = None):
+ """
+ Converts a code artifact to an AMR. Assumes that the code file is the first
+ file (and only) attached to the artifact.
+
+ Args:
+ ```
+ artifact_id (str): the id of the code artifact
+ name (str, optional): the name to set on the newly created model
+ description (str, optional): the description to set on the newly created model
+ ```
+ """
  from utils import create_job
 
  operation_name = "operations.code_to_amr"
- options = {"artifact_id": artifact_id}
+ options = {"artifact_id": artifact_id,
+ "name": name,
+ "description": description}
 
  resp = create_job(operation_name=operation_name, options=options)
 

workers/operations.py

@@ -29,29 +29,41 @@
 
 
 # Worker jobs for TA1 services
-def put_mathml_to_skema(*args, **kwargs):
- # Get vars
- mathml = kwargs.get("mathml")
+def equations_to_amr(*args, **kwargs):
+ equation_type = kwargs.get("equation_type")
+ equations = kwargs.get("equations")
  model = kwargs.get("model")
+ name = kwargs.get("name")
+ description = kwargs.get("description")
 
- # PUT the mathml to the skema endpoint.
- skema_mathml_url = SKEMA_API + "/mathml/amr"
-
- headers = {"Content-Type": "application/json"}
-
- put_payload = {"mathml": mathml, "model": model}
-
- amr_response = requests.put(
- skema_mathml_url, data=json.dumps(put_payload, default=str), headers=headers
- )
-
+ if equation_type == "mathml":
+ # PUT the mathml to the skema endpoint.
+ logger.info("Processing mathml")
+ url = f"{SKEMA_API}/mathml/amr"
+ put_payload = {"mathml": equations, "model": model}
+ elif equation_type == "latex":
+ logger.info("Processing latex")
+ url = f"{UNIFIED_API}/workflows/latex/equations-to-amr"
+ put_payload = {"equations": equations, "model": model}
+
+ headers = {"Content-Type": "application/json"} 
+
+ logger.info(f"Sending equations of type {equation_type} to TA1")
+ if equation_type == "mathml":
+ amr_response = requests.put(
+ url, data=json.dumps(put_payload, default=str), headers=headers
+ )
+ elif equation_type == "latex": 
+ amr_response = requests.post(
+ url, data=json.dumps(put_payload, default=str), headers=headers
+ )
  try:
  amr_json = amr_response.json()
  except:
- logger.error("Failed to parse response from TA1 Service")
+ logger.error(f"Failed to parse response from TA1 Service: {amr_response.text}")
 
  if amr_response.status_code == 200 and amr_json:
- tds_responses = put_amr_to_tds(amr_json)
+ tds_responses = put_amr_to_tds(amr_json, name, description)
 
  response = {
  "status_code": amr_response.status_code,
@@ -355,31 +367,34 @@ def link_amr(*args, **kwargs):
 # 60e539e4-6969-4369-a358-c601a3a583da
 def code_to_amr(*args, **kwargs):
  artifact_id = kwargs.get("artifact_id")
+ name = kwargs.get("name")
+ description = kwargs.get("description")
 
  artifact_json, downloaded_artifact = get_artifact_from_tds(artifact_id=artifact_id)
 
  code_blob = downloaded_artifact.decode("utf-8")
-
  code_amr_workflow_url = f"{UNIFIED_API}/workflows/code/snippets-to-pn-amr"
 
  request_payload = {
  "files": [artifact_json.get("file_names")[0]],
- "blobs": [code_blob],
+ "blobs": [code_blob]
  }
 
+ logger.info(f"Sending code to TA1 service with artifact id: {artifact_id}")
  amr_response = requests.post(
  code_amr_workflow_url, json=json.loads(json.dumps(request_payload))
  )
+ logger.info(f"Response received from TA1 with status code: {amr_response.status_code}")
 
  amr_json = amr_response
 
  try:
  amr_json = amr_response.json()
  except:
- logger.error("Failed to parse response from TA1 Service")
+ logger.error(f"Failed to parse response from TA1 Service:\n{amr_response.text}")
 
  if amr_response.status_code == 200 and amr_json:
- tds_responses = put_amr_to_tds(amr_json)
+ tds_responses = put_amr_to_tds(amr_json, name, description)
 
  response = {
  "status_code": amr_response.status_code,
@@ -391,6 +406,7 @@ def code_to_amr(*args, **kwargs):
 
  return response
  else:
+ logger.error(f"Code extraction failure: {amr_response.text}")
  response = {
  "status_code": amr_response.status_code,
  "amr": None,

workers/utils.py

@@ -18,12 +18,17 @@
 TDS_API = os.getenv("TDS_URL")
 
 
-def put_amr_to_tds(amr_payload):
+def put_amr_to_tds(amr_payload, name=None, description=None):
  # Expects json amr payload and puts it to TDS models and model-configurations, returning an ID.
 
  headers = {"Content-Type": "application/json"}
 
- logger.info(amr_payload)
+ if name:
+ amr_payload['name'] = name
+ if description:
+ amr_payload['description'] = description
+
+ logger.debug(amr_payload)
 
  # Create TDS model
  tds_models = f"{TDS_API}/models"
@@ -49,6 +54,8 @@ def put_amr_to_tds(amr_payload):
 
  config_id = config_response.json().get("id")
 
+ logger.info(f"Created model in TDS with id {model_id}")
+ logger.info(f"Created model config in TDS with id {config_id}")
  return {"model_id": model_id, "configuration_id": config_id}