phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix using mock completion

Browse source
This commit is contained in:
Ishaan Jaff 2024-06-25 11:14:40 -07:00
parent 343e3f3e90
commit ccf1bbc5d7
3 changed files with 57 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

39
litellm/utils.py
View file

@ -9708,18 +9708,45 @@ class TextCompletionStreamWrapper:
raise StopAsyncIteration
def mock_completion_streaming_obj(model_response, mock_response, model):
def mock_completion_streaming_obj(
model_response, mock_response, model, n: Optional[int] = None
):
for i in range(0, len(mock_response), 3):
completion_obj = {"role": "assistant", "content": mock_response[i : i + 3]}
model_response.choices[0].delta = completion_obj
completion_obj = Delta(role="assistant", content=mock_response[i : i + 3])
if n is None:
model_response.choices[0].delta = completion_obj
else:
_all_choices = []
for j in range(n):
_streaming_choice = litellm.utils.StreamingChoices(
index=j,
delta=litellm.utils.Delta(
role="assistant", content=mock_response[i : i + 3]
),
)
_all_choices.append(_streaming_choice)
model_response.choices = _all_choices
yield model_response
async def async_mock_completion_streaming_obj(model_response, mock_response, model):
async def async_mock_completion_streaming_obj(
model_response, mock_response, model, n: Optional[int] = None
):
for i in range(0, len(mock_response), 3):
completion_obj = Delta(role="assistant", content=mock_response[i : i + 3])
model_response.choices[0].delta = completion_obj
model_response.choices[0].finish_reason = "stop"
if n is None:
model_response.choices[0].delta = completion_obj
else:
_all_choices = []
for j in range(n):
_streaming_choice = litellm.utils.StreamingChoices(
index=j,
delta=litellm.utils.Delta(
role="assistant", content=mock_response[i : i + 3]
),
)
_all_choices.append(_streaming_choice)
model_response.choices = _all_choices
yield model_response