Improve error handling (#10)

* chore: update the exponential backoff error handling

* chore: update versions of langgraph and langchain

* chore: update poetry lock file

* fix: remove the unused functions from util

Author: akshseh

bedrock_deep_research.py

@@ -13,6 +13,7 @@
 from bedrock_deep_research import BedrockDeepResearch
 from bedrock_deep_research.config import DEFAULT_TOPIC, SUPPORTED_MODELS, Configuration
 from bedrock_deep_research.model import Section
+from bedrock_deep_research.utils import CustomError
 
 logger = logging.getLogger(__name__)
 LOGLEVEL = os.environ.get("LOGLEVEL", "INFO").upper()
@@ -146,20 +147,28 @@ def render_initial_form():
                     with st.spinner(
                         "Please wait while the article outline is being generated..."
                     ):
-                        response = st.session_state.bedrock_deep_research.start(
-                            topic)
-
-                        logger.debug(f"Outline response: {response}")
-
-                        state = st.session_state.bedrock_deep_research.get_state()
-
-                        article = Article(
-                            title=state.values["title"],
-                            sections=state.values["sections"],
-                        )
-                        st.session_state.article = article.render_outline()
-                        st.session_state.stage = "outline_feedback"
-                        st.rerun()
+                        try:
+                            response = st.session_state.bedrock_deep_research.start(
+                                topic)
+                        except CustomError as e:
+                            logger.error(f"Bedrock ClientError: {e}")
+                            raise e
+
+                        except Exception as e:
+                            logger.error(
+                                f"An error occurred while creating the outline: {e}")
+                            raise e
+                        else:
+                            logger.debug(f"Outline response: {response}")
+                            state = st.session_state.bedrock_deep_research.get_state()
+
+                            article = Article(
+                                title=state.values["title"],
+                                sections=state.values["sections"],
+                            )
+                            st.session_state.article = article.render_outline()
+                            st.session_state.stage = "outline_feedback"
+                            st.rerun()
     except Exception as e:
         logger.error(f"An error occurred: {e}")
         raise
@@ -288,7 +297,7 @@ def on_accept_outline_button_click():
                     st.session_state.text_error = ""
                     st.rerun()
                 except Exception as e:
-                    st.error(f"An error occurred: {e}")
+                    st.error(f"An error occurred in article creation: {e}")
 
     except Exception as e:
         logger.error(f"An error occurred: {e}")

bedrock_deep_research/nodes/article_head_image_generator.py

@@ -28,7 +28,7 @@ def __init__(self, message):
         self.message = message
 
 
-@exponential_backoff_retry(Exception, max_retries=10)
+@exponential_backoff_retry(ClientError, max_retries=10)
 def generate_image(model_id, body):
     """
     Generate an image using Amazon Nova Canvas model on demand.
@@ -44,7 +44,6 @@ def generate_image(model_id, body):
 
     accept = "application/json"
     content_type = "application/json"
-
     response = bedrock.invoke_model(
         body=body, modelId=model_id, accept=accept, contentType=content_type
     )
@@ -143,7 +142,6 @@ def __call__(self, state: ArticleState, config: RunnableConfig):
             logger.error("A bedrock client error occurred:", message)
         except Exception as e:
             logger.error(
-
                 "An error occurred during ArticleHeadImageGenerator:", e)
 
         logger.info("Generated head image: %s", image_path)

bedrock_deep_research/nodes/article_outline_generator.py

@@ -67,7 +67,6 @@ def __call__(self, state: ArticleState, config: RunnableConfig):
         outline = self.generate_outline(
             configurable.planner_model, configurable.max_tokens, system_prompt, user_prompt)
 
-
         logger.info(f"Generated sections: {outline.sections}")
         sections = [
             Section(section_number=i, name=section.name,

bedrock_deep_research/nodes/final_sections_writer.py

@@ -1,3 +1,4 @@
+from botocore.exceptions import ClientError
 from langchain_aws import ChatBedrock
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.runnables import RunnableConfig
@@ -87,7 +88,7 @@ def __call__(self, state: SectionState, config: RunnableConfig):
 
         return {"completed_sections": [section]}
 
-    @exponential_backoff_retry(Exception, max_retries=10)
+    @exponential_backoff_retry(ClientError, max_retries=10)
     def _generate_final_sections(
         self,
         model: ChatBedrock,

bedrock_deep_research/nodes/initial_researcher.py

@@ -2,6 +2,7 @@
 import logging
 from typing import List
 
+from botocore.exceptions import ClientError
 from langchain_aws import ChatBedrock
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.runnables import RunnableConfig
@@ -61,7 +62,7 @@ def __call__(self, state: ArticleInputState, config: RunnableConfig):
 
         return {"source_str": source_str}
 
-    @exponential_backoff_retry(Exception, max_retries=10)
+    @exponential_backoff_retry(ClientError, max_retries=10)
     def generate_search_queries(self, model_id: str, max_tokens: int, system_prompt: str, user_prompt: str) -> List[str]:
         planner_model = ChatBedrock(
             model_id=model_id, max_tokens=max_tokens)

bedrock_deep_research/nodes/section_search_query_generator.py

@@ -1,5 +1,6 @@
 import logging
 
+from botocore.client import ClientError
 from langchain_aws import ChatBedrock
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.runnables import RunnableConfig
@@ -54,7 +55,7 @@ def __call__(self, state: SectionState, config: RunnableConfig):
         return {"search_queries": queries.queries}
 
 
-@exponential_backoff_retry(Exception, max_retries=10)
+@exponential_backoff_retry(ClientError, max_retries=10)
 def generate_section_queries(configurable: Configuration, section: Section) -> Queries:
     planner_model = ChatBedrock(
         model_id=configurable.planner_model, max_tokens=configurable.max_tokens

bedrock_deep_research/nodes/section_writer.py

@@ -1,6 +1,7 @@
 import logging
 from typing import List, Literal
 
+from botocore.exceptions import ClientError
 from langchain_aws import ChatBedrock
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.runnables import RunnableConfig
@@ -155,7 +156,7 @@ def __call__(self, state: SectionState, config: RunnableConfig) -> Command[Liter
                 goto=SectionWebResearcher.N,
             )
 
-    @exponential_backoff_retry(Exception, max_retries=10)
+    @exponential_backoff_retry(ClientError, max_retries=10)
     def _generate_section_content(
         self,
         model: ChatBedrock,
@@ -183,7 +184,7 @@ def _generate_section_content(
 
         return section_content.content
 
-    @exponential_backoff_retry(Exception, max_retries=10)
+    @exponential_backoff_retry(ClientError, max_retries=10)
     def _grade_section_content(
         self, model: ChatBedrock, system_prompt: str, section: Section
     ) -> Feedback:

bedrock_deep_research/utils.py

@@ -1,11 +1,22 @@
 import logging
 import random
+import re
 import time
 from functools import wraps
 
+from botocore.exceptions import ClientError
+
 logger = logging.getLogger(__name__)
 
 
+class CustomError(Exception):
+    def __init__(self, message):
+        self.message = message
+
+    def __str__(self):
+        return self.message
+
+
 def exponential_backoff_retry(
     ExceptionToCheck, max_retries: int = 5, initial_delay: float = 1.0
 ):
@@ -14,6 +25,7 @@ def exponential_backoff_retry(
 
     Args:
         func: Function to retry
+        ExceptionToCheck: Exception class to check for retrying
         max_retries: Maximum number of retry attempts
         initial_delay: Initial delay in seconds before first retry
     """
@@ -26,26 +38,37 @@ def wrapper(*args, **kwargs):
             for attempt in range(max_retries + 1):
                 try:
                     return func(*args, **kwargs)
-
-                except ExceptionToCheck as e:
-
-                    if attempt == max_retries:
+                except ClientError as e:
+                    if e.response['Error']['Code'] == 'ExpiredTokenException':
                         logger.error(
-                            f"Execution failed after {attempt} attempts")
+                            "Expired token error. Please check/update your Security Token included in the request")
+                        # Do not try max_retry times
+                        attempt = max_retries
+                        raise CustomError(
+                            message="Expired Token. Please update the AWS credentials, to connect to the boto Client.")
+                    elif e.response['Error']['Code'] == 'ThrottlingException':
+                        if attempt == max_retries:
+                            logger.error(
+                                f"Error code: {e.response['Error']['Code']}"
+                                f"Execution failed after {max_retries} attempts due to throttling. Try again later.")
+                            raise CustomError(
+                                message=f"Throttling Exception raised.. Retry limit of {max_retries} retries reached.")
+                        logger.info(
+                            f"Attempt {attempt+1} failed due to throttling. Retrying...")
+                        # Add jitter to avoid thundering herd problem
+                        jitter = random.uniform(0, 0.1 * delay)
+                        sleep_time = delay + jitter
+                        logger.debug(
+                            f"Retrying in {sleep_time:.2f} seconds..."
+                        )
+                        time.sleep(sleep_time)
+                        delay *= 2  # Exponential backoff
+                    else:
+                        logger.error(f"Client Error Raised: {e}")
                         raise e
-
-                    # Add jitter to avoid thundering herd problem
-                    jitter = random.uniform(0, 0.1 * delay)
-                    sleep_time = delay + jitter
-
-                    logger.debug(
-                        f"Attempt {attempt + 1}/{max_retries} failed. {str(e)}"
-                        f"Retrying in {sleep_time:.2f} seconds..."
-                    )
-
-                    time.sleep(sleep_time)
-                    delay *= 2  # Exponential backoff
-
+                except ExceptionToCheck as e:
+                    logger.error(f"Error raised by {func.__name__}: {e}")
+                    raise e
         return wrapper
 
     return decorator
@@ -74,3 +97,12 @@ def format_web_search(search_response, max_tokens_per_source, include_raw_conten
             formatted_text += f"Full source content limited to {max_tokens_per_source} tokens: {raw_content}\n\n"
 
     return formatted_text.strip()
+
+
+def extract_xml_content(text: str, tag_name: str) -> str | None:
+    pattern = f"<{tag_name}>(.*?)</{tag_name}>"
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    else:
+        return None