Pantograph/pantograph/gen_tactic.py

from pantograph.server import Server, ServerError, TacticFailure
from pantograph.expr import Variable, Goal, TacticCalc
import  unittest
import sglang as sgl

LEAN4_INTRO = '''/-- A sequence `u` of real numbers converges to `l` if `∀ ε > 0, ∃ N, ∀ n ≥ N, |u_n - l| ≤ ε`.
This condition will be spelled `seq_limit u l`. -/
def seq_limit (u : ℕ → ℝ) (l : ℝ) : Prop :=
∀ ε > 0, ∃ N, ∀ n ≥ N, |u n - l| ≤ ε

/- In the above definition, note that the `n`-th term of the sequence `u` is denoted
simply by `u n`.

Similarly, in the next definition, `f x` is what we would write `f(x)` on paper.
Also note that implication is denoted by a single arrow (we'll explain why later). -/

/-- A function`f : ℝ → ℝ` is continuous at `x₀` if
`∀ ε > 0, ∃ δ > 0, ∀ x, |x - x₀| ≤ δ ⇒ |f(x) - f(x₀)| ≤ ε`.
This condition will be spelled `continuous_at f x₀`.-/
def continuous_at (f : ℝ → ℝ) (x₀ : ℝ) : Prop :=
∀ ε > 0, ∃ δ > 0, ∀ x, |x - x₀| ≤ δ → |f x - f x₀| ≤ ε

/-- Now we claim that if `f` is continuous at `x₀` then it is sequentially continuous
at `x₀`: for any sequence `u` converging to `x₀`, the sequence `f ∘ u` converges
to `f x₀`.  -/
example (f : ℝ → ℝ) (u : ℕ → ℝ) (x₀ : ℝ) (hu : seq_limit u x₀) (hf : continuous_at f x₀) :
  seq_limit (f ∘ u) (f x₀) := by { -- This `by` keyword marks the beginning of the proof
  -- Put your text cursor here and watch the Lean InfoView panel to the right.
  -- Then move your cursor from line to line in the proof while monitoring the Infoview.

  -- Our goal is to prove that, for any positive `ε`, there exists a natural
  -- number `N` such that, for any natural number `n` at least `N`,
  --  `|f(u_n) - f(x₀)|` is at most `ε`.
  unfold seq_limit
  -- Fix a positive number `ε`.
  intros ε hε
  -- By assumption on `f` applied to this positive `ε`, we get a positive `δ`
  -- such that, for all real number `x`, if `|x - x₀| ≤ δ` then `|f(x) - f(x₀)| ≤ ε` (1).
  obtain ⟨δ, δ_pos, Hf⟩ : ∃ δ > 0, ∀ x, |x - x₀| ≤ δ → |f x - f x₀| ≤ ε := hf ε hε
  -- The assumption on `u` applied to this `δ` gives a natural number `N` such that
  -- for every natural number `n`, if `n ≥ N` then `|u_n - x₀| ≤ δ`   (2).
  obtain ⟨N, Hu⟩ : ∃ N, ∀ n ≥ N, |u n - x₀| ≤ δ := hu δ δ_pos
  -- Let's prove `N` is suitable.
  use N
  -- Fix `n` which is at least `N`. Let's prove `|f(u_n) - f(x₀)| ≤ ε`.
  intros n hn
  -- Thanks to (1) applied to `u_n`, it suffices to prove that `|u_n - x₀| ≤ δ`.
  apply Hf
  -- This follows from property (2) and our assumption on `n`.
  exact Hu n hn
  -- This finishes the proof!
  }

/-
Now that this proof is over, you can use the file explorer to the
left of this panel to open the file `Exercises > 01Rewriting.lean`.
-/'''

LEAN4_REWRITE = '''Rewrite tactic tutorial:
example (a b c : Nat) : a + b + c = a + c + b := by
  rw [Nat.add_assoc, Nat.add_comm b, ← Nat.add_assoc]

example (a b c : Nat) : a + b + c = a + c + b := by
  rw [Nat.add_assoc, Nat.add_assoc, Nat.add_comm b]

example (a b c : Nat) : a + b + c = a + c + b := by
  rw [Nat.add_assoc, Nat.add_assoc, Nat.add_comm _ b]

example (f : Nat → Nat) (a : Nat) (h : a + 0 = 0) : f a = f 0 := by
  rw [Nat.add_zero] at h
  rw [h]

def Tuple (α : Type) (n : Nat) :=
  { as : List α // as.length = n }

example (n : Nat) (h : n = 0) (t : Tuple α n) : Tuple α 0 := by
  rw [h] at t
  exact t
'''

@sgl.function
def multi_turn_question(s, question_1, question_2):
    s += sgl.system("You are a helpful assistant.")
    s += sgl.user(question_1)
    s += sgl.assistant(sgl.gen("answer_1", max_tokens=256))
    s += sgl.user(question_2)
    s += sgl.assistant(sgl.gen("answer_2", max_tokens=256))


@sgl.function
def select_tactic(s, server, state, goal_id, feedback_turns = 5):
    
    s += sgl.system("You are an expert in Lean. Choose the next one tactic to run given the current proof state and goals.")
    s += sgl.user(LEAN4_REWRITE)
    s += sgl.user("The current proof state: GoalState(state_id=0, goals=[Goal(variables=[], target='∀ (a b: Nat), (b = 2) -> 1 + a + 1 = a + b', name=None, is_conversion=False)])")
    s += sgl.assistant("```intros a b h```")
    s += sgl.user("The current proof state: GoalState(state_id=1, goals=[Goal(variables=[Variable(t='Nat', v=None, name='a'), Variable(t='Nat', v=None, name='b'), Variable(t='b = 2', v=None, name='h')], target='1 + a + 1 = a + b', name=None, is_conversion=False)])")
    s += sgl.assistant('TacticCalc("1 + a + 1 = a + 1 + 1")')
    s += sgl.user("The current proof state: " + str(state))
    for i in range(feedback_turns):
        with s.copy() as tmp:
            tmp += sgl.assistant(sgl.gen("tactic", max_tokens=64))
            print("==tmp===")
            print(tmp["tactic"])
            tactic = extract_code_from_llm_output(tmp["tactic"])
        s += sgl.assistant("```"+tactic+"```")
        success, new_state = apply_tactic(server, state, goal_id, tactic)
        print("===execute===")
        print(success, new_state )
        if not success:
            with s.user():
                s += "This answer got Lean compile error:\n" + str(new_state) + "\n"
                s += "Please try again by taking the Lean compiler feedback."
            
        else:
            return tactic, new_state
    return None, None

def apply_tactic(server, state, goal_id, tactic):
    try:
        new_state = server.goal_tactic(state, goal_id=goal_id, tactic=tactic)
    except ServerError as e:
        return False, e
    except TacticFailure as e:
        return False, e
    return True, new_state
    
def extract_code_from_llm_output(reply):
    i = reply.find("```lean")
    if i != -1:
        reply = reply[i + 7:]
        i = reply.find("```")
        reply = reply[:i]
        return reply
    i = reply.find("```")
    if i != -1:
        reply = reply[i + 3:]
        i = reply.find("```")
        reply = reply[:i]
        return reply
    return reply

class TestServerSGL(unittest.TestCase):

    def test_conv_calc_sgl(self):
        n_trails = 5
        sgl.set_default_backend(sgl.OpenAI("gpt-4"))

        server = Server()
        state0 = server.goal_start("∀ (a b: Nat), (b = 2) -> 1 + a + 1 = a + b")
        print("==========state0============")
        print(state0)
        variables = [
            Variable(name="a", t="Nat"),
            Variable(name="b", t="Nat"),
            Variable(name="h", t="b = 2"),
        ]

        state1 = server.goal_tactic(state0, goal_id=0, tactic="intro a b h")
        print("==========state1============")
        print(state1)
        state2 = server.goal_tactic(state1, goal_id=0, tactic=TacticCalc("1 + a + 1 = a + 1 + 1"))
        print("==========state2============")
        print(state2)
        self.assertEqual(state2.goals, [
            Goal(
                variables,
                target="1 + a + 1 = a + 1 + 1",
                name='calc',
            ),
            Goal(
                variables,
                target="a + 1 + 1 = a + b",
            ),
        ])
        state3 = None
        for i in range(n_trails):
            print(f"===============trail {str(i)}============")
            try:
                state = select_tactic.run(server, state2, goal_id = 1)
                state3 = state.ret_value
                for m in state.messages():
                    print(m["role"], ":", m["content"])

                print("\n-- new state --\n", state3)
                break
                
            except ServerError as e:
                print(f"server error: {e}")
                continue
        state3 = server.goal_tactic(state2, goal_id=1, tactic=TacticCalc("_ = a + 2"))


        print("==========state3============")
        print(state3)
        state4 = None
        for i in range(n_trails):
            print(f"===============trail {str(i)}============")
            try:
                state = select_tactic.run(server, state3, goal_id = 0)
                state4 = state.ret_value
                for m in state.messages():
                    print(m["role"], ":", m["content"])

                print("\n-- new state --\n", state4)
                break
                
            except ServerError as e:
                print(f"server error: {e}")
                continue

        state4 = server.goal_tactic(state3, goal_id=0, tactic="rw [Nat.add_assoc]")
        print("==========state4============")
        print(state4)        
        self.assertTrue(state4.is_solved)


    def test_sglang_openai(self):
        sgl.set_default_backend(sgl.OpenAI("gpt-4"))

        print('\n----- Test sglang ---')
        state = multi_turn_question.run(
            question_1="What is the capital of the United States?",
            question_2="List two local attractions.",
        )

        for m in state.messages():
            print(m["role"], ":", m["content"])

        print("\n-- answer_1 --\n", state["answer_1"])

   
if __name__ == '__main__':

    unittest.main()