add handling on error (langchain-ai#541)

apify · Jan 5, 2023 · 5aefc2b · 5aefc2b
1 parent 1631981
commit 5aefc2b
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 10 deletions.
diff --git a/langchain/agents/agent.py b/langchain/agents/agent.py
@@ -272,9 +272,14 @@ def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]:
                     self.callback_manager.on_tool_start(
                         {"name": str(chain)[:60] + "..."}, output, color="green"
                     )
-                # We then call the tool on the tool input to get an observation
-                observation = chain(output.tool_input)
-                color = color_mapping[output.tool]
+                try:
+                    # We then call the tool on the tool input to get an observation
+                    observation = chain(output.tool_input)
+                    color = color_mapping[output.tool]
+                except Exception as e:
+                    if self.verbose:
+                        self.callback_manager.on_tool_error(e)
+                    raise e
             else:
                 if self.verbose:
                     self.callback_manager.on_tool_start(

diff --git a/langchain/chains/base.py b/langchain/chains/base.py
@@ -138,7 +138,12 @@ def __call__(
             self.callback_manager.on_chain_start(
                 {"name": self.__class__.__name__}, inputs
             )
-        outputs = self._call(inputs)
+        try:
+            outputs = self._call(inputs)
+        except Exception as e:
+            if self.verbose:
+                self.callback_manager.on_chain_error(e)
+            raise e
         if self.verbose:
             self.callback_manager.on_chain_end(outputs)
         self._validate_outputs(outputs)

diff --git a/langchain/llms/base.py b/langchain/llms/base.py
@@ -73,7 +73,12 @@ def generate(
                 self.callback_manager.on_llm_start(
                     {"name": self.__class__.__name__}, prompts
                 )
-            output = self._generate(prompts, stop=stop)
+            try:
+                output = self._generate(prompts, stop=stop)
+            except Exception as e:
+                if self.verbose:
+                    self.callback_manager.on_llm_error(e)
+                raise e
             if self.verbose:
                 self.callback_manager.on_llm_end(output)
             return output
@@ -90,11 +95,18 @@ def generate(
             else:
                 missing_prompts.append(prompt)
                 missing_prompt_idxs.append(i)
-        self.callback_manager.on_llm_start(
-            {"name": self.__class__.__name__}, missing_prompts
-        )
-        new_results = self._generate(missing_prompts, stop=stop)
-        self.callback_manager.on_llm_end(new_results)
+        if self.verbose:
+            self.callback_manager.on_llm_start(
+                {"name": self.__class__.__name__}, missing_prompts
+            )
+        try:
+            new_results = self._generate(missing_prompts, stop=stop)
+        except Exception as e:
+            if self.verbose:
+                self.callback_manager.on_llm_error(e)
+            raise e
+        if self.verbose:
+            self.callback_manager.on_llm_end(new_results)
         for i, result in enumerate(new_results.generations):
             existing_prompts[missing_prompt_idxs[i]] = result
             prompt = prompts[i]