openai[patch]: Update openai chat model to new base class interface (#…

…19729)
langchain-ai · Mar 29, 2024 · d4673a3 · d4673a3
1 parent 23fcc14
commit d4673a3
Showing 1 changed file with 3 additions and 28 deletions.
diff --git a/libs/partners/openai/langchain_openai/chat_models/base.py b/libs/partners/openai/langchain_openai/chat_models/base.py
@@ -34,11 +34,7 @@
     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,28 +474,19 @@ def _stream(
             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(
         self,
         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,31 +569,19 @@ async def _astream(
             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(
         self,
         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)