prompt-optimizer / backend.py
SharvNey's picture
Create backend.py
1dd3a75 verified
import os
from huggingface_hub import InferenceClient
# ---- TOKEN ----
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
raise RuntimeError("HF_TOKEN not found. Add it in Space → Settings → Repository secrets")
# ---- MODEL ----
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
client = InferenceClient(model=MODEL_NAME, token=HF_TOKEN)
def run_llm(prompt: str) -> str:
"""Run prompt using conversational API (supported by this model)."""
try:
messages = [{"role": "user", "content": prompt}]
resp = client.chat_completion(
messages=messages,
max_tokens=300,
temperature=0.7,
)
return resp.choices[0].message["content"].strip()
except Exception as e:
# raise clean error (no stacktrace in UI)
raise RuntimeError(f"LLM call failed: {e}")
def generate_variants(base_prompt: str):
meta_prompt = f"""
Rewrite the following prompt in 5 improved ways:
1. More specific
2. More technical
3. More concise
4. More creative
5. Step-by-step style
Return ONLY the rewritten prompts, one per line.
Original prompt:
{base_prompt}
"""
text = run_llm(meta_prompt)
variants = []
for line in text.split("\n"):
line = line.strip("-• ").strip()
if line:
variants.append(line)
# Fallback if parsing fails
if len(variants) < 3:
variants = [
f"You are an expert. {base_prompt}",
f"Explain step by step: {base_prompt}",
f"Give a concise answer: {base_prompt}",
f"Provide a technical explanation: {base_prompt}",
f"Explain creatively: {base_prompt}",
]
return variants