import base64
import re
import logging
from typing import Optional, List, Union, Literal
from mira.openai_utility import OpenAIClient, ImageFmts
from mira.sources.sympy_ode import template_model_from_sympy_odes
from mira.sources.sympy_ode.constants import (
ODE_IMAGE_PROMPT,
ODE_CONCEPTS_PROMPT_TEMPLATE,
ODE_PDF_PROMPT,
ODE_MULTIPLE_IMAGE_PROMPT
)
ContentType = Literal["pdf", "image", "text"]
logger = logging.getLogger(__name__)
ode_pattern_raw = r"(odes\s*=\s*\[.*?\])\s*"
ode_pattern = re.compile(ode_pattern_raw, re.DOTALL)
[docs]class CodeExecutionError(Exception):
"""An error raised when there is an error executing the code"""
[docs]def pdf_file_to_odes_str(
pdf_path: str,
client: OpenAIClient
) -> str:
"""Get an ODE string from a PDF file depicting an ODE system
Parameters
----------
pdf_path :
The path to the PDF file
client :
A :class:`mira.openai_utility.OpenAIClient` instance
Returns
-------
:
The ODE string extracted from the PDF. The string should contain the code
necessary to define the ODEs using sympy.
"""
with open(pdf_path, "rb") as f:
pdf_bytes = f.read()
return pdf_to_odes_str(pdf_bytes, client)
[docs]def pdf_to_odes_str(
pdf_bytes: bytes,
client: OpenAIClient,
) -> str:
"""Get an ODE string from PDF bytes depicting an ODE system
Parameters
----------
pdf_bytes :
The bytes of the PDF file
client :
The OpenAI client
Returns
-------
:
The ODE string extracted from the PDF. The string should contain the code
necessary to define the ODEs using sympy.
"""
base64_pdf = base64.b64encode(pdf_bytes).decode('utf-8')
response = extract_ode_str_from_base64_pdf(
base64_pdf=base64_pdf,
client=client
)
return response
[docs]def image_file_to_odes_str(
image_path: Union[str, List[str]],
client: OpenAIClient,
) -> str:
"""
Get an ODE string from an image file or a list of image files depicting an
ODE system
Parameters
----------
image_path :
The path to the image file or a list of paths to each image file
client :
A :class:`mira.openai_utility.OpenAIClient` instance
Returns
-------
:
The ODE string extracted from the image(s). The string should contain the
code necessary to define the ODEs using sympy.
"""
if isinstance(image_path, str):
with open(image_path, "rb") as f:
image_bytes = f.read()
image_format = image_path.split(".")[-1]
logger.info("Single image passed")
return image_to_odes_str(image_bytes, client, image_format)
else:
image_bytes_list = []
image_format_list = []
for path in image_path:
with open(path, "rb") as f:
image_bytes = f.read()
image_bytes_list.append(image_bytes)
image_format = path.split(".")[-1]
image_format_list.append(image_format)
logger.info(f"{len(image_bytes_list)} images passed")
return image_to_odes_str(image_bytes_list, client, image_format_list)
[docs]def image_to_odes_str(
image_bytes: Union[bytes, List[bytes]],
client: OpenAIClient,
image_format: Union[ImageFmts, List[ImageFmts]]
) -> str:
"""Get an ODE string from an image or a list of images depicting an ODE system
Parameters
----------
image_bytes :
The bytes of the image or a list of bytes for each image
client :
The OpenAI client
image_format :
The format of the image or a list of formats for each image.
Returns
-------
:
The ODE string extracted from the image. The string should contain the code
necessary to define the ODEs using sympy.
"""
if not isinstance(image_bytes, List):
base64_image = base64.b64encode(image_bytes).decode('utf-8')
response = extract_ode_str_from_base64_image(base64_image=base64_image,
image_format=image_format,
client=client)
else:
base64_image_list = []
for image_byte in image_bytes:
base64_image = base64.b64encode(image_byte).decode('utf-8')
base64_image_list.append(base64_image)
response = extract_ode_str_from_base64_image(base64_image=base64_image_list,
image_format=image_format,
client=client)
return response
[docs]def clean_response(response: str) -> str:
"""Clean up the response from the OpenAI chat completion
Parameters
----------
response :
The response from the OpenAI chat completion.
Returns
-------
:
The cleaned up response. The response is stripped of the code block
markdown, i.e. the triple backticks and the language specifier. It is also
stripped of leading and trailing whitespaces.
"""
response = response.replace("```python", "")
response = response.replace("```", "")
return response.strip()
[docs]def get_concepts_from_odes(
ode_str: str,
client: OpenAIClient,
) -> Optional[dict]:
"""Get the concepts data from the ODEs defined in the code snippet
Parameters
----------
ode_str :
The string containing the code snippet defining the ODEs
client :
The openAI client
Returns
-------
:
The concepts data in the form of a dictionary, or None if no concepts data
could be generated.
"""
# Cut out the part of the code that defines the `odes` variable
# Regular expression to capture the `odes` variable assignment and definition
match = re.search(ode_pattern, ode_str)
if match:
odes_code = match.group(1)
else:
raise ValueError("No code snippet defining the variable `odes` found")
# Prompt the OpenAI chat completion to create the concepts data dictionary
prompt = ODE_CONCEPTS_PROMPT_TEMPLATE.substitute(ode_insert=odes_code)
response = client.run_chat_completion(prompt)
# Clean up the response
response_text = clean_response(response.message.content)
# Extract the concepts from the response using exec
locals_dict = locals()
exec(response_text, globals(), locals_dict)
concept_data = locals_dict.get("concept_data")
assert concept_data is not None, "The code should define a variable called `concept_data`"
return concept_data
[docs]def test_execution(code: str) -> tuple[bool, str]:
"""Test if code executes successfully
Parameters
----------
code :
The Python code
Returns
-------
:
Tuple of (success, error_message). success is True if code executed
and defined `odes`, False otherwise. error_message is empty on success.
"""
try:
namespace = {}
exec("import sympy", namespace)
exec(code, namespace)
if 'odes' in namespace:
return (True, "")
return (False, "Code executed but 'odes' was not defined")
except Exception as e:
return (False, f"{type(e).__name__}: {e}")
[docs]def test_ode_model(odes) -> tuple[bool, str]:
"""Test if a list of SymPy ODEs can be converted to a TemplateModel
Parameters
----------
odes :
The list of sympy.Eq ODEs to test
Returns
-------
:
Tuple of (success, error_message).
"""
try:
if not odes:
return (False, "Empty ODE list")
template_model_from_sympy_odes(odes)
return (True, "")
except Exception as e:
return (False, f"{type(e).__name__}: {e}")