fix: unify generation outputs on newer openllm release (langchain-ai#…

…10523) update newer generation format from OpenLLm where it returns a dictionary for one shot generation cc @baskaryan Signed-off-by: Aaron <[email protected]> --------- Signed-off-by: Aaron <[email protected]>
apify · Sep 13, 2023 · ac9609f · ac9609f
1 parent 201b61d
commit ac9609f
Showing 1 changed file with 16 additions and 4 deletions.
diff --git a/libs/langchain/langchain/llms/openllm.py b/libs/langchain/langchain/llms/openllm.py
@@ -265,10 +265,16 @@ def _call(
             self._identifying_params["model_name"], **copied
         )
         if self._client:
-            return self._client.query(prompt, **config.model_dump(flatten=True))
+            o = self._client.query(prompt, **config.model_dump(flatten=True))
+            if isinstance(o, dict) and "text" in o:
+                return o["text"]
+            return o
         else:
             assert self._runner is not None
-            return self._runner(prompt, **config.model_dump(flatten=True))
+            o = self._runner(prompt, **config.model_dump(flatten=True))
+            if isinstance(o, dict) and "text" in o:
+                return o["text"]
+            return o
 
     async def _acall(
         self,
@@ -291,9 +297,12 @@ async def _acall(
             self._identifying_params["model_name"], **copied
         )
         if self._client:
-            return await self._client.acall(
+            o = await self._client.acall(
                 "generate", prompt, **config.model_dump(flatten=True)
             )
+            if isinstance(o, dict) and "text" in o:
+                return o["text"]
+            return o
         else:
             assert self._runner is not None
             (
@@ -304,6 +313,9 @@ async def _acall(
             generated_result = await self._runner.generate.async_run(
                 prompt, **generate_kwargs
             )
-            return self._runner.llm.postprocess_generate(
+            o = self._runner.llm.postprocess_generate(
                 prompt, generated_result, **postprocess_kwargs
             )
+            if isinstance(o, dict) and "text" in o:
+                return o["text"]
+            return o