github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-09-17 15:35:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[patch]: mlflow handle empty chunk(#24689)

Browse Source
This commit is contained in:
Sunish Sheth
2024-07-25 20:36:29 -07:00
committed by GitHub
parent 20690db482
commit 59880a9147
1 changed files with 25 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
libs/community/langchain_community/chat_models/mlflow.py
View File

@@ -176,29 +176,36 @@ class ChatMlflow(BaseChatModel):
chunk_iter = self._client.predict_stream(endpoint=self.endpoint, inputs=data)
first_chunk_role = None
for chunk in chunk_iter:
choice = chunk["choices"][0]
if chunk["choices"]:
choice = chunk["choices"][0]
chunk_delta = choice["delta"]
if first_chunk_role is None:
first_chunk_role = chunk_delta.get("role")
chunk = ChatMlflow._convert_delta_to_message_chunk(
chunk_delta, first_chunk_role
)
chunk_delta = choice["delta"]
if first_chunk_role is None:
first_chunk_role = chunk_delta.get("role")
generation_info = {}
if finish_reason := choice.get("finish_reason"):
generation_info["finish_reason"] = finish_reason
if logprobs := choice.get("logprobs"):
generation_info["logprobs"] = logprobs
chunk_message = ChatMlflow._convert_delta_to_message_chunk(
chunk_delta, first_chunk_role
)
chunk = ChatGenerationChunk(
message=chunk, generation_info=generation_info or None
)
generation_info = {}
if finish_reason := choice.get("finish_reason"):
generation_info["finish_reason"] = finish_reason
if logprobs := choice.get("logprobs"):
generation_info["logprobs"] = logprobs
if run_manager:
run_manager.on_llm_new_token(chunk.text, chunk=chunk, logprobs=logprobs)
chunk = ChatGenerationChunk(
message=chunk_message, generation_info=generation_info or None
)
yield chunk
if run_manager:
run_manager.on_llm_new_token(
chunk.text, chunk=chunk, logprobs=logprobs
)
yield chunk
else:
# Handle the case where choices are empty if needed
continue
@property
def _identifying_params(self) -> Dict[str, Any]: