Skip to content

Commit 945c791

Browse files
juandzambrano12Xpitfire
juandzambrano12
authored and
Xpitfire
committed
lean4+ engine
1 parent 1809309 commit 945c791

File tree

3 files changed

+253
-0
lines changed

3 files changed

+253
-0
lines changed

‎notebooks/examples/Lean engine.png

191 KB
Loading

‎notebooks/examples/lean.py

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
from symai.backend.engines.lean.engine_lean4 import LeanEngine
2+
# Example usage
3+
if __name__ == "__main__":
4+
# Initialize LeanEngine
5+
engine = LeanEngine()
6+
7+
# Sample Lean code
8+
code = '''
9+
theorem and_commutative (A B : Prop) : A ∧ B → B ∧ A :=
10+
fun h : A ∧ B => ⟨h.right, h.left⟩dsdsdds
11+
'''
12+
13+
# Execute the Lean code
14+
print("Running LeanEngine with a sample Lean theorem...")
15+
results, metadata = engine.forward(code)
16+
17+
# Check if results are valid and print the output
18+
if results:
19+
print("Results:", results[0].value['output'])
20+
print("Metadata:", metadata)

‎symai/backend/engines/lean/engine_lean4.py

+233
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
import os
2+
import subprocess
3+
import docker
4+
import paramiko
5+
import tempfile
6+
from typing import Any, List, Tuple, Optional, Dict
7+
from ...base import Engine
8+
from ....symbol import Result
9+
10+
class LeanResult(Result):
11+
"""
12+
Represents the result of executing a Lean code snippet.
13+
14+
Attributes:
15+
_value (Dict[str, str]): A dictionary containing the output of the Lean execution.
16+
"""
17+
def __init__(self, value: Dict[str, str]) -> None:
18+
"""
19+
Initializes a new LeanResult instance.
20+
21+
Args:
22+
value (Dict[str, str]): The result output of the Lean code execution.
23+
"""
24+
super().__init__(value)
25+
self._value = value
26+
27+
class LeanEngine(Engine):
28+
"""
29+
Engine for executing Lean code within a Docker container, providing SSH access for execution.
30+
31+
Attributes:
32+
ssh_host (str): The SSH host, defaulting to 'localhost'.
33+
ssh_port (int): The SSH port, defaulting to 2222.
34+
ssh_user (str): The SSH username, defaulting to 'root'.
35+
ssh_key_path (str): The path to the SSH private key, defaulting to '~/.ssh/id_rsa'.
36+
docker_client (docker.DockerClient): The Docker client used to manage containers.
37+
container (docker.models.containers.Container): The Docker container used for executing Lean code.
38+
"""
39+
40+
def __init__(
41+
self,
42+
ssh_host: str = 'localhost',
43+
ssh_port: int = 2222,
44+
ssh_user: str = 'root',
45+
ssh_key_path: str = '~/.ssh/id_rsa'
46+
) -> None:
47+
"""
48+
Initializes the LeanEngine with SSH and Docker configurations.
49+
50+
Args:
51+
ssh_host (str): The SSH host, defaulting to 'localhost'.
52+
ssh_port (int): The SSH port, defaulting to 2222.
53+
ssh_user (str): The SSH username, defaulting to 'root'.
54+
ssh_key_path (str): The path to the SSH private key, defaulting to '~/.ssh/id_rsa'.
55+
"""
56+
super().__init__()
57+
self.ssh_host: str = ssh_host
58+
self.ssh_port: int = ssh_port
59+
self.ssh_user: str = ssh_user
60+
self.ssh_key_path: str = os.path.expanduser(ssh_key_path)
61+
self.docker_client: docker.DockerClient = docker.from_env()
62+
self.container: docker.models.containers.Container = self._ensure_container()
63+
self._setup_ssh()
64+
65+
def id(self) -> str:
66+
"""
67+
Returns the identifier for the engine.
68+
69+
Returns:
70+
str: The identifier of the LeanEngine, 'lean4'.
71+
"""
72+
return 'lean4'
73+
74+
def _ensure_container(self) -> docker.models.containers.Container:
75+
"""
76+
Ensures the Docker container for Lean execution exists, creating it if necessary.
77+
78+
Returns:
79+
docker.models.containers.Container: The Docker container instance used for Lean code execution.
80+
"""
81+
container_name: str = "lean-container"
82+
83+
try:
84+
existing_container: docker.models.containers.Container = self.docker_client.containers.get(container_name)
85+
existing_container.remove(force=True)
86+
except docker.errors.NotFound:
87+
print(f"No existing container named '{container_name}' found. Proceeding to create a new one.")
88+
89+
dockerfile: str = """
90+
FROM buildpack-deps:buster
91+
92+
ENV ELAN_HOME=/usr/local/elan \
93+
PATH=/usr/local/elan/bin:$PATH \
94+
LEAN_VERSION=leanprover/lean4:nightly
95+
96+
RUN apt-get update && apt-get install -y openssh-server curl && rm -rf /var/lib/apt/lists/*
97+
98+
RUN curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | sh -s -- -y --default-toolchain $LEAN_VERSION; \
99+
elan default $LEAN_VERSION; \
100+
elan --version; \
101+
lean --version; \
102+
leanc --version; \
103+
lake --version;
104+
105+
RUN mkdir /var/run/sshd && echo 'root:root' | chpasswd && sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
106+
107+
EXPOSE 22
108+
CMD ["/usr/sbin/sshd", "-D"]
109+
"""
110+
with tempfile.NamedTemporaryFile("w", delete=False) as temp_dockerfile:
111+
temp_dockerfile.write(dockerfile)
112+
dockerfile_path: str = temp_dockerfile.name
113+
114+
image: docker.models.images.Image
115+
image, _ = self.docker_client.images.build(path=os.path.dirname(dockerfile_path), dockerfile=dockerfile_path, tag="lean4-container-image")
116+
os.remove(dockerfile_path)
117+
118+
container: docker.models.containers.Container = self.docker_client.containers.run(
119+
image.id,
120+
detach=True,
121+
name=container_name,
122+
ports={'22/tcp': self.ssh_port}
123+
)
124+
return container
125+
126+
def _setup_ssh(self) -> None:
127+
"""
128+
Sets up SSH access to the Docker container, including generating an SSH key pair if necessary,
129+
and configuring the container to accept SSH connections using the generated key.
130+
"""
131+
if not os.path.exists(self.ssh_key_path):
132+
subprocess.run(['ssh-keygen', '-t', 'rsa', '-b', '2048', '-f', self.ssh_key_path, '-N', ''], check=True)
133+
134+
subprocess.run(['docker', 'exec', self.container.id, 'mkdir', '-p', '/root/.ssh'], check=True)
135+
subprocess.run(['docker', 'cp', f'{self.ssh_key_path}.pub', f'{self.container.id}:/root/.ssh/authorized_keys'], check=True)
136+
subprocess.run(['docker', 'exec', self.container.id, 'chmod', '600', '/root/.ssh/authorized_keys'], check=True)
137+
subprocess.run(['docker', 'exec', self.container.id, 'chown', 'root:root', '/root/.ssh/authorized_keys'], check=True)
138+
139+
def forward(self, argument: Any) -> Tuple[List[LeanResult], dict]:
140+
"""
141+
Executes Lean code provided as a string or as an object property.
142+
143+
Args:
144+
argument (Any): The Lean code to be executed, either as a string or wrapped in an object.
145+
146+
Returns:
147+
Tuple[List[LeanResult], dict]: A tuple containing the result of the Lean execution and associated metadata.
148+
"""
149+
code: str = argument if isinstance(argument, str) else argument.prop.prepared_input
150+
151+
rsp: Optional[LeanResult] = None
152+
err: Optional[str] = None
153+
tmpfile_path: Optional[str] = None
154+
metadata: Dict[str, Any] = {}
155+
try:
156+
with tempfile.NamedTemporaryFile(delete=False, suffix=".lean") as tmpfile:
157+
tmpfile.write(code.encode())
158+
tmpfile_path = tmpfile.name
159+
160+
output, exec_metadata = self._execute_lean(tmpfile_path)
161+
metadata.update(exec_metadata)
162+
163+
if output:
164+
rsp = LeanResult({'output': output})
165+
else:
166+
metadata['status'] = 'no_output'
167+
except Exception as e:
168+
err = str(e)
169+
metadata.update({'status': 'error', 'message': err})
170+
print(f"Error during Lean execution: {err}")
171+
finally:
172+
if tmpfile_path and os.path.exists(tmpfile_path):
173+
os.remove(tmpfile_path)
174+
if self.container:
175+
print(f"Killing Docker container '{self.container.id}'...")
176+
self.container.remove(force=True)
177+
178+
return [rsp] if rsp else [], metadata
179+
180+
def _execute_lean(self, filepath: str) -> Tuple[str, dict]:
181+
"""
182+
Executes a Lean script within the Docker container via SSH.
183+
184+
Args:
185+
filepath (str): The path to the Lean file to be executed.
186+
187+
Returns:
188+
Tuple[str, dict]: The output from the Lean execution and associated status metadata.
189+
"""
190+
try:
191+
ssh = paramiko.SSHClient()
192+
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
193+
ssh.connect(self.ssh_host, port=self.ssh_port, username=self.ssh_user, key_filename=self.ssh_key_path)
194+
195+
elan_path: str = "/usr/local/elan/bin/elan"
196+
lean_path: str = "/usr/local/elan/bin/lean"
197+
198+
stdin, stdout, stderr = ssh.exec_command(f"{elan_path} default stable && {lean_path} --version")
199+
output: str = stdout.read().decode()
200+
error: str = stderr.read().decode()
201+
print("SSH Command Output:", output)
202+
print("SSH Command Error:", error)
203+
204+
sftp = ssh.open_sftp()
205+
remote_path: str = f"/root/{os.path.basename(filepath)}"
206+
sftp.put(filepath, remote_path)
207+
sftp.close()
208+
209+
stdin, stdout, stderr = ssh.exec_command(f"{lean_path} {remote_path}")
210+
output = stdout.read().decode()
211+
error = stderr.read().decode()
212+
213+
ssh.exec_command(f"rm {remote_path}")
214+
ssh.close()
215+
216+
if "error" in output.lower() or "error" in error.lower():
217+
return output, {'status': 'failure'}
218+
elif not output and not error:
219+
return "Lean program halted successfully with no output.", {'status': 'success'}
220+
else:
221+
return output, {'status': 'success'}
222+
223+
except Exception as e:
224+
raise RuntimeError(f"SSH command execution failed: {str(e)}")
225+
226+
def prepare(self, argument: Any) -> None:
227+
"""
228+
Prepares the input for Lean execution by processing and converting it into the appropriate format.
229+
230+
Args:
231+
argument (Any): The input to be processed and prepared.
232+
"""
233+
argument.prop.prepared_input = str(argument.prop.processed_input)

0 commit comments

Comments
 (0)