import base64
import re
from typing import Optional, List
from mira.metamodel import TemplateModel
from mira.openai 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_pattern = r"(odes\s*=\s*\[.*?\])\s*"
pattern = re.compile(ode_pattern, re.DOTALL)
[docs]def image_file_to_odes_str(
image_path: str,
client: OpenAIClient,
) -> str:
"""Get an ODE string from an image file depicting an ODE system
Parameters
----------
image_path :
The path to the image file
client :
A :class:`mira.openai.OpenAIClient` instance
Returns
-------
:
The ODE string extracted from the image. The string should contain the code
necessary to define the ODEs using sympy.
"""
with open(image_path, "rb") as f:
image_bytes = f.read()
image_format = image_path.split(".")[-1]
return image_to_odes_str(image_bytes, client, image_format)
[docs]def image_to_odes_str(
image_bytes: bytes,
client: OpenAIClient,
image_format: ImageFmts = "png"
) -> str:
"""Get an ODE string from an image depicting an ODE system
Parameters
----------
image_bytes :
The bytes of the image
client :
The OpenAI client
image_format :
The format of the image. The default is "png".
Returns
-------
:
The ODE string extracted from the image. The string should contain the code
necessary to define the ODEs using sympy.
"""
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)
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(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 execute_template_model_from_sympy_odes(
ode_str,
attempt_grounding: bool,
client: OpenAIClient,
) -> TemplateModel:
"""Create a TemplateModel from the sympy ODEs defined in the code snippet string
Parameters
----------
ode_str :
The code snippet defining the ODEs
attempt_grounding :
Whether to attempt grounding the concepts in the ODEs. This will prompt the
OpenAI chat completion to create concepts data to provide grounding for the
concepts in the ODEs. The concepts data is then used to create the TemplateModel.
client :
The OpenAI client
Returns
-------
:
The TemplateModel created from the sympy ODEs.
"""
# FixMe, for now use `exec` on the code, but need to find a safer way to execute
# the code
# Import sympy just in case the code snippet does not import it
import sympy
odes: List[sympy.Eq] = None
# Execute the code and expose the `odes` variable to the local scope
local_dict = locals()
exec(ode_str, globals(), local_dict)
# `odes` should now be defined in the local scope
odes = local_dict.get("odes")
assert odes is not None, "The code should define a variable called `odes`"
if attempt_grounding:
concept_data = get_concepts_from_odes(ode_str, client)
else:
concept_data = None
return template_model_from_sympy_odes(odes, concept_data=concept_data)