try to improve upon the `training-data` command.

Stream: general

[Zihao Zhang]

@Kim Morrison Hi, Kim! I am trying to write a function in Basic.lean of your project https://github.com/kim-em/lean-training-data.
I want achieve this goal: when there is only one goal of info.gsBefore and only one goal of info.gsAfter, I want make the latter's conclusion a new hyposis of the former's tactic state.Here is an example:
given the two tactic state before and after rw [pow_two]:
goal1:

a b : ℝ
h1 : a = 1
h2 : b = 2
⊢ (a + b) ^ 2 = a ^ 2 + 2 * a * b + b ^ 2

goal2:

a b : ℝ
h1 : a = 1
h2 : b = 2
⊢ (a + b) * (a + b) = a ^ 2 + 2 * a * b + b ^ 2

my goal is to make goal1:

h1 : a = 1
h2 : b = 2
h_1: (a + b) * (a + b) = a ^ 2 + 2 * a * b + b ^ 2
⊢ (a + b) ^ 2 = a ^ 2 + 2 * a * b + b ^ 2

But it is quite hard, I don't know how to continue, here is my updated code.

def addSimpleHypothesisAndPrint (info : TacticInfo) (ctx : ContextInfo) : IO Unit := do
  info.runMetaMGoalsBefore ctx fun gsBefore =>
    match gsBefore with
    | [] => do IO.println "gsBefore has no goals!"
    | _ :: _ :: _ => do IO.println "gsBefore has more than one goal!"
    | [goal1] => do
          let t←ctx.runMetaM {} <| Meta.withMCtx info.mctxAfter <| (
            match info.goalsAfter with
            | [] => do
              IO.println "gsAfter has no goals!"
              pure (mkConst ``Unit)
            | _ :: _ :: _ => do
              IO.println "gsAfter has more than one goal!"
              pure (mkConst ``Unit)
            | [goal2] => do
              let t ← goal2.getType
              pure t
          )


          if (← Meta.isDefEq t (mkConst ``Unit)) then
            return
          else
            let p ← mkFreshExprMVar t MetavarKind.syntheticOpaque `h_1
            goal1.withContext do
              let (_, goal1) ← MVarId.intro1P $ ← goal1.assert `h_1 t p
              let (_, g) ← goal1.revert (clearAuxDeclsInsteadOfRevert := true) (← goal1.getDecl).lctx.getFVarIds
              let ty ← instantiateMVars (← g.getType)
              if ty.hasExprMVar then
                throwError "Extracted goal has metavariables: {ty}"
              let (ty,_) ← Lean.Elab.Term.TermElabM.run (Term.levelMVarToParam ty)
              let seenLevels := collectLevelParams {} ty
              let (levelNames, _) ← Lean.Elab.Term.TermElabM.run Term.getLevelNames
              let levels := levelNames.filter fun u => seenLevels.visitedLevel.contains (.param u)
              addAndCompile <| Declaration.axiomDecl
                { name := `extracted_1
                  levelParams := levels
                  isUnsafe := false
                  type := ty }
              let sig ← addMessageContext <| MessageData.signature `extracted_1
              let cmd := if ← Meta.isProp ty then "theorem" else "def"
              let sig←sig.toString
              IO.println (s!"{cmd} {sig} := sorry")

Could you help me?