diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index e359d36f4c..a97d6812c8 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/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" diff --git a/litellm/main.py b/litellm/main.py index 307659c8a2..07d7be2bad 100644 --- a/litellm/main.py +++ b/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,