代码初始化

inference/inference.py

@@ -80,6 +80,33 @@ async def chat_completion(token: str = Body(...), messages: list = Body(...)):
         raise HTTPException(status_code=500, detail=f"Error generating chat completion: {str(e)}")
 
 
+@router.post("/format-prompt/")
+async def chat_completion(token: str = Body(...), messages: list = Body(...)):
+    try:
+        # 创建 InferenceClient
+        client = InferenceClient(api_key=token)
+        print("问题：", messages)
+        messages.append({
+            "role": "system",
+            "content": "You are a highly intelligent image generation text optimizer. Your role is to take the user's "
+                       "input prompt and enhance it, ensuring it is precise, vivid, and detailed, while maintaining "
+                       "its intended meaning. Always provide the enhanced version in English."
+        })
+
+        # 使用 chat API 请求生成模型的回答  备选模型  google/gemma-2-2b-it
+        completion = client.chat.completions.create(
+            model="mistralai/Mistral-Nemo-Instruct-2407",
+            messages=messages,
+            max_tokens=500
+        )
+
+        # 返回对话信息
+        return {"message": completion.choices[0].message}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error generating chat completion: {str(e)}")
+
+
 @router.post("/api-inference/")
 async def api_inference(
         authorization: str = Header(...),