google-genai[patch]: Invoke callback prior to yielding token (#17092)

- **Description:** Invoke callback prior to yielding token in stream and
astream methods for Google-genai,
  - **Issue:** the issue # 16913,
  - **Twitter handle:** Sparsh10649446

---------

Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
pull/17263/head
Sparsh Jain 8 months ago committed by GitHub
parent 7306600e2f
commit a2167614b7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -624,9 +624,9 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
generation_t=ChatGenerationChunk,
)
gen = cast(ChatGenerationChunk, _chat_result.generations[0])
yield gen
if run_manager:
run_manager.on_llm_new_token(gen.text)
yield gen
async def _astream(
self,
@ -654,9 +654,9 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
generation_t=ChatGenerationChunk,
)
gen = cast(ChatGenerationChunk, _chat_result.generations[0])
yield gen
if run_manager:
await run_manager.on_llm_new_token(gen.text)
yield gen
def _prepare_chat(
self,

Loading…
Cancel
Save