Add logprobs in gen output

langchain-ai · Dec 18, 2023 · bbca143 · bbca143
1 parent 1acc7ff
commit bbca143
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/libs/community/langchain_community/chat_models/openai.py b/libs/community/langchain_community/chat_models/openai.py
@@ -454,9 +454,12 @@ def _create_chat_result(self, response: Union[dict, BaseModel]) -> ChatResult:
             response = response.dict()
         for res in response["choices"]:
             message = convert_dict_to_message(res["message"])
+            generation_info = dict(finish_reason=res.get("finish_reason"))
+            if "logprobs" in res:
+                generation_info["logprobs"] = res["logprobs"]
             gen = ChatGeneration(
                 message=message,
-                generation_info=dict(finish_reason=res.get("finish_reason")),
+                generation_info=generation_info,
             )
             generations.append(gen)
         token_usage = response.get("usage", {})