feat - use n in mock completion

litellm/llms/prompt_templates/factory.py

@@ -135,7 +135,7 @@ def convert_to_ollama_image(openai_image_url: str):
 
 
 def ollama_pt(
-        model, messages
+    model, messages
 ):  # https://github.com/ollama/ollama/blob/af4cf55884ac54b9e637cd71dadfe9b7a5685877/docs/modelfile.md#template
     if "instruct" in model:
         prompt = custom_prompt(
@@ -185,19 +185,18 @@ def ollama_pt(
                     function_name: str = call["function"]["name"]
                     arguments = json.loads(call["function"]["arguments"])
 
-                    tool_calls.append({
-                        "id": call_id,
-                        "type": "function",
-                        "function": {
-                            "name": function_name,
-                            "arguments": arguments
+                    tool_calls.append(
+                        {
+                            "id": call_id,
+                            "type": "function",
+                            "function": {"name": function_name, "arguments": arguments},
                         }
-                    })
+                    )
 
                 prompt += f"### Assistant:\nTool Calls: {json.dumps(tool_calls, indent=2)}\n\n"
 
             elif "tool_call_id" in message:
-                prompt += f"### User:\n{message["content"]}\n\n"
+                prompt += f"### User:\n{message['content']}\n\n"
 
             elif content:
                 prompt += f"### {role.capitalize()}:\n{content}\n\n"

litellm/main.py

@@ -428,6 +428,7 @@ def mock_completion(
     model: str,
     messages: List,
     stream: Optional[bool] = False,
+    n: Optional[int] = None,
     mock_response: Union[str, Exception, dict] = "This is a mock request",
     mock_tool_calls: Optional[List] = None,
     logging=None,
@@ -496,8 +497,19 @@ def mock_completion(
                 model_response, mock_response=mock_response, model=model
             )
             return response
-
-        model_response["choices"][0]["message"]["content"] = mock_response
+        if n is None:
+            model_response["choices"][0]["message"]["content"] = mock_response
+        else:
+            _all_choices = []
+            for i in range(n):
+                _choice = litellm.utils.Choices(
+                    index=i,
+                    message=litellm.utils.Message(
+                        content=mock_response, role="assistant"
+                    ),
+                )
+                _all_choices.append(_choice)
+            model_response["choices"] = _all_choices
         model_response["created"] = int(time.time())
         model_response["model"] = model
 
@@ -944,6 +956,7 @@ def completion(
                 model,
                 messages,
                 stream=stream,
+                n=n,
                 mock_response=mock_response,
                 mock_tool_calls=mock_tool_calls,
                 logging=logging,