diff --git a/libs/partners/deepseek/langchain_deepseek/chat_models.py b/libs/partners/deepseek/langchain_deepseek/chat_models.py index f3bcc2d1583..2398b1f9881 100644 --- a/libs/partners/deepseek/langchain_deepseek/chat_models.py +++ b/libs/partners/deepseek/langchain_deepseek/chat_models.py @@ -1,9 +1,13 @@ """DeepSeek chat models.""" -from typing import Dict, Optional, Type, Union +from json import JSONDecodeError +from typing import Any, Dict, Iterator, List, Optional, Type, Union import openai -from langchain_core.messages import AIMessageChunk +from langchain_core.callbacks import ( + CallbackManagerForLLMRun, +) +from langchain_core.messages import AIMessageChunk, BaseMessage from langchain_core.outputs import ChatGenerationChunk, ChatResult from langchain_core.utils import from_env, secret_from_env from langchain_openai.chat_models.base import BaseChatOpenAI @@ -239,3 +243,41 @@ class ChatDeepSeek(BaseChatOpenAI): reasoning_content ) return generation_chunk + + def _stream( + self, + messages: List[BaseMessage], + stop: Optional[List[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> Iterator[ChatGenerationChunk]: + try: + yield from super()._stream( + messages, stop=stop, run_manager=run_manager, **kwargs + ) + except JSONDecodeError as e: + raise JSONDecodeError( + "DeepSeek API returned an invalid response. " + "Please check the API status and try again.", + e.doc, + e.pos, + ) from e + + def _generate( + self, + messages: List[BaseMessage], + stop: Optional[List[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> ChatResult: + try: + return super()._generate( + messages, stop=stop, run_manager=run_manager, **kwargs + ) + except JSONDecodeError as e: + raise JSONDecodeError( + "DeepSeek API returned an invalid response. " + "Please check the API status and try again.", + e.doc, + e.pos, + ) from e