phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat - use n in mock completion

Browse source
This commit is contained in:
Ishaan Jaff 2024-06-25 10:50:47 -07:00
parent 165fb72e31
commit ae01cdfc3b
2 changed files with 23 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/llms/prompt_templates/factory.py
View file

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

17
litellm/main.py
View file

@ -428,6 +428,7 @@ def mock_completion(
model: str, model: str,
messages: List, messages: List,
stream: Optional[bool] = False, stream: Optional[bool] = False,
n: Optional[int] = None,
mock_response: Union[str, Exception, dict] = "This is a mock request", mock_response: Union[str, Exception, dict] = "This is a mock request",
mock_tool_calls: Optional[List] = None, mock_tool_calls: Optional[List] = None,
logging=None, logging=None,
@ -496,8 +497,19 @@ def mock_completion(
model_response, mock_response=mock_response, model=model model_response, mock_response=mock_response, model=model
) )
return response return response
if n is None:
model_response["choices"][0]["message"]["content"] = mock_response 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["created"] = int(time.time())
model_response["model"] = model model_response["model"] = model
@ -944,6 +956,7 @@ def completion(
model, model,
messages, messages,
stream=stream, stream=stream,
n=n,
mock_response=mock_response, mock_response=mock_response,
mock_tool_calls=mock_tool_calls, mock_tool_calls=mock_tool_calls,
logging=logging, logging=logging,