Groups of order $p^3$

As @Arturo Magidin suggests in a comment, you only need to tweak your generators.

For suppose $a$ and $b$ have order $p^2$; then by taking suitable powers we can assume $a^p=c$, $b^p=c^{-1}$ where $c$ is a generator of the centre.

Now use the magic fact: as the derived group is central we have $$a^n b^n=(ab)^n [b,a]^{n\choose 2}$$.

(See https://en.wikipedia.org/wiki/Commutator).

Now take $n=p$, an odd prime, so that $p$ divides $p\choose 2$: we will get $$c c^{-1}=(ab)^p$$ and so can replace the generators $a,b$ by new generators $a,ab$ of orders $p^2, p$.

Case 2

Suppose $G$ has an element $a$ of order $p^2$. This case seems to require a certain minimum of cabalistic manipulation, I have not encountered a way of deriving the required result without a bit of slog. The following account still leaves a little arithmetic to be verified by the reader.

let $A = \langle a \rangle$. then $Z \subset A$ since $x \not \in \langle a\rangle \implies \langle a,x\rangle = G$ and G is non-abelian. so $Z = \langle a^p\rangle$.

continue with $x$ representing an element in $G \setminus A$, and use an overbar to indicate images in the factor group $G/Z$.

since $\bar G$ is abelian the inner automorphism of $\bar A$ corresponding to $\bar x$ is trivial. i.e. $\bar x^{-1} \bar a \bar x = \bar a$.

the pre-image of this relationship in $G$ is $(xa^{mp})^{-1}a^{np+1}xa^{mp} = a^{n'p+1}$ for some $m,n,n' \in [1,p) \cap \mathbb{N}$. As $\langle a^p\rangle = Z$ this simplifies to: $$ x^{-1}ax = a^{jp + 1} \tag{1} $$

where $j=n'-n$.

now (1), together with $a^{p^2}=1$ implies that: $$ x^{-k}ax^k= a^{(jp+1)^k} = a^{kjp+1} $$

for some $k$ we have $kj \equiv 1 \pmod{p}$, so set $b_1 = x^k$

now we have $b_1^{-1}ab_1 = a^{p+1}$

it follows that: $$ b_1^{-p}ab_1^p =a \tag{2} $$

construction of the other generator, forcing |b|= $p$

from (2) $b_1^p$ commutes with $a$ hence $b_1^p = a^k$ for some $k$. but now raising this to the power $p$ gives $1=a^{kp}$ so that $k$ is a multiple of $p$, say $k=hp$.

suppose $b_1$ has order $p^2$ with $b_1^p = a^{hp}$ define $b = b_1a^{-h}$. Since $a^jb_1 = b_1a^{j(p+1)}$ we may write $a^{-h}b_1 = b_1a^{-h(p+1)}$ hence: $$ b^p = (b_1a^{-h})^p = b_1^pa^{-h\sum_{s=0}^{p-1} (p+1)^s}=b_1^pa^{-hp} = 1 \tag{2} $$

so $G = \langle a,b\rangle$, $|a|=p^2$, $|b|=p$, $b^{-1}ab = a^{p+1}$.

the last equality in (2) follows from: $$ \sum_{s=0}^{p-1} (p+1)^s = \frac{(p+1)^p - 1}{p} \equiv p \pmod{p^2} $$