feat(core): Model Fallback Support for BaseChatModel

libs/core/langchain_core/language_models/chat_models.py

@@ -63,6 +63,7 @@
 from langchain_core.rate_limiters import BaseRateLimiter
 from langchain_core.runnables import RunnableMap, RunnablePassthrough
 from langchain_core.runnables.config import ensure_config, run_in_executor
+from langchain_core.runnables.fallbacks import RunnableWithFallbacks
 from langchain_core.tracers._streaming import _StreamingCallbackHandler
 from langchain_core.utils.function_calling import (
     convert_to_json_schema,
@@ -318,6 +319,9 @@
     rate_limiter: BaseRateLimiter | None = Field(default=None, exclude=True)
     "An optional rate limiter to use for limiting the number of requests."
 
+    fallback: Optional[RunnableWithFallbacks] = Field(default=None, exclude=True)
+    """A fallback type to use if provided."""
+
     disable_streaming: bool | Literal["tool_calling"] = False
     """Whether to disable streaming for this model.
 
@@ -397,6 +401,10 @@
         stop: list[str] | None = None,
         **kwargs: Any,
     ) -> AIMessage:
+        fallback_object = self.get_fallback()
+        if fallback_object:
+            self._reset_fallback()
+            return fallback_object.invoke(input, config, stop=stop, **kwargs)
         config = ensure_config(config)
         return cast(
             "AIMessage",
@@ -424,6 +432,10 @@
         stop: list[str] | None = None,
         **kwargs: Any,
     ) -> AIMessage:
+        fallback_object = self.get_fallback()
+        if fallback_object:
+            self._reset_fallback()
+            return await fallback_object.ainvoke(input, config, stop=stop, **kwargs)
         config = ensure_config(config)
         llm_result = await self.agenerate_prompt(
             [self._convert_input(input)],
@@ -483,6 +495,12 @@
         stop: list[str] | None = None,
         **kwargs: Any,
     ) -> Iterator[AIMessageChunk]:
+        fallback_object = self.get_fallback()
+        if fallback_object:
+            self._reset_fallback()
+            for chunk in fallback_object.stream(input, config, stop=stop, **kwargs):
+                yield chunk
+            return
         if not self._should_stream(async_api=False, **{**kwargs, "stream": True}):
             # Model doesn't implement streaming, so use default implementation
             yield cast(
@@ -599,6 +617,14 @@
         stop: list[str] | None = None,
         **kwargs: Any,
     ) -> AsyncIterator[AIMessageChunk]:
+        fallback_object = self.get_fallback()
+        if fallback_object:
+            self._reset_fallback()
+            async for chunk in fallback_object.astream(
+                input, config, stop=stop, **kwargs
+            ):
+                yield chunk
+            return
         if not self._should_stream(async_api=True, **{**kwargs, "stream": True}):
             # No async or sync stream is implemented, so fall back to ainvoke
             yield cast(
@@ -710,6 +736,22 @@
             LLMResult(generations=[[generation]]),
         )
 
+    def _reset_fallback(self) -> BaseChatModel:
+        """Reset the fallback."""
+        self.fallback = None
+        return self
+
+    def set_fallback(
+        self, fallback: Optional[RunnableWithFallbacks] = None
+    ) -> BaseChatModel:
+        """Add fallbacks to the model."""
+        self.fallback = fallback
+        return self
+
+    def get_fallback(self) -> Optional[RunnableWithFallbacks]:
+        """Get the fallback."""
+        return self.fallback
+
     # --- Custom methods ---
 
     def _combine_llm_outputs(self, llm_outputs: list[dict | None]) -> dict:  # noqa: ARG002