feat: Added support for moderations in generation_info

langchain-ai · Nov 18, 2024 · 1810726 · 1810726
1 parent 3b34b94
commit 1810726
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 91 deletions.
diff --git a/libs/ibm/langchain_ibm/llms.py b/libs/ibm/langchain_ibm/llms.py
@@ -323,9 +323,13 @@ def _create_llm_result(self, response: List[dict]) -> LLMResult:
             results = res.get("results")
             if results:
                 finish_reason = results[0].get("stop_reason")
+                moderations = results[0].get("moderations")
                 gen = Generation(
                     text=results[0].get("generated_text"),
-                    generation_info={"finish_reason": finish_reason},
+                    generation_info={
+                        "finish_reason": finish_reason,
+                        "moderations": moderations,
+                    },
                 )
                 generations.append([gen])
         final_token_usage = self._extract_token_usage(response)
@@ -345,13 +349,15 @@ def _stream_response_to_generation_chunk(
             return GenerationChunk(text="")
 
         finish_reason = stream_response["results"][0].get("stop_reason", None)
+        moderations = stream_response["results"][0].get("moderations", None)
 
         return GenerationChunk(
             text=stream_response["results"][0]["generated_text"],
             generation_info=dict(
                 finish_reason=(
                     None if finish_reason == "not_finished" else finish_reason
                 ),
+                moderations=moderations,
                 llm_output={
                     "model_id": self.model_id,
                     "deployment_id": self.deployment_id,