diff --git a/ai_agent.py b/ai_agent.py
index 217f188..561f212 100644
--- a/ai_agent.py
+++ b/ai_agent.py
@@ -11,7 +11,7 @@ from typing import List, Tuple, Optional, Union, Dict, Awaitable
 from openrouter import OpenRouter, RetryConfig
 from openrouter.components import AssistantMessage, AssistantMessageTypedDict, ChatMessageContentItemTypedDict, \
     ChatMessageToolCall, MessageTypedDict, SystemMessageTypedDict
-from openrouter.errors import ResponseValidationError, ChatError
+from openrouter.errors import ResponseValidationError, OpenRouterError
 from openrouter.utils import BackoffStrategy
 
 from fal_client import AsyncClient as FalClient
@@ -308,7 +308,7 @@ class AiAgent:
                 except Exception:
                     pass
             raise e
-        except ChatError as e:
+        except OpenRouterError as e:
             if e.message == "Provider returned error":
                 body = json.loads(e.body)
                 try: