diff --git a/libs/partners/openai/langchain_openai/chat_models/base.py b/libs/partners/openai/langchain_openai/chat_models/base.py index fcabe5c99a..4de8b2798a 100644 --- a/libs/partners/openai/langchain_openai/chat_models/base.py +++ b/libs/partners/openai/langchain_openai/chat_models/base.py @@ -34,11 +34,7 @@ from langchain_core.callbacks import ( CallbackManagerForLLMRun, ) from langchain_core.language_models import LanguageModelInput -from langchain_core.language_models.chat_models import ( - BaseChatModel, - agenerate_from_stream, - generate_from_stream, -) +from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import ( AIMessage, AIMessageChunk, @@ -478,8 +474,6 @@ class ChatOpenAI(BaseChatModel): chunk = ChatGenerationChunk( message=chunk, generation_info=generation_info or None ) - if run_manager: - run_manager.on_llm_new_token(chunk.text, chunk=chunk, logprobs=logprobs) yield chunk def _generate( @@ -487,19 +481,12 @@ class ChatOpenAI(BaseChatModel): messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, - stream: Optional[bool] = None, **kwargs: Any, ) -> ChatResult: - should_stream = stream if stream is not None else self.streaming - if should_stream: - stream_iter = self._stream( - messages, stop=stop, run_manager=run_manager, **kwargs - ) - return generate_from_stream(stream_iter) message_dicts, params = self._create_message_dicts(messages, stop) params = { **params, - **({"stream": stream} if stream is not None else {}), + **({"stream": self.streaming} if self.streaming else {}), **kwargs, } response = self.client.create(messages=message_dicts, **params) @@ -582,10 +569,6 @@ class ChatOpenAI(BaseChatModel): chunk = ChatGenerationChunk( message=chunk, generation_info=generation_info or None ) - if run_manager: - await run_manager.on_llm_new_token( - token=chunk.text, chunk=chunk, logprobs=logprobs - ) yield chunk async def _agenerate( @@ -593,20 +576,12 @@ class ChatOpenAI(BaseChatModel): messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, - stream: Optional[bool] = None, **kwargs: Any, ) -> ChatResult: - should_stream = stream if stream is not None else self.streaming - if should_stream: - stream_iter = self._astream( - messages, stop=stop, run_manager=run_manager, **kwargs - ) - return await agenerate_from_stream(stream_iter) - message_dicts, params = self._create_message_dicts(messages, stop) params = { **params, - **({"stream": stream} if stream is not None else {}), + **({"stream": self.streaming} if self.streaming else {}), **kwargs, } response = await self.async_client.create(messages=message_dicts, **params)