ConversationalRetrievalChain not working with stream callback handler for ChatGoogleGenerativeAI
Labels: name: auto:bug, name: area: models

### System Info

aiohttp==3.9.1
aiosignal==1.3.1
annotated-types==0.6.0
anyio==4.2.0
attrs==23.1.0
cachetools==5.3.2
certifi==2023.11.17
charset-normalizer==3.3.2
colorama==0.4.6
dataclasses-json==0.6.3
distro==1.8.0
frozenlist==1.4.1
google-ai-generativelanguage==0.4.0
google-api-core==2.15.0
google-auth==2.25.2
google-generativeai==0.3.2
googleapis-common-protos==1.62.0
greenlet==3.0.2
grpcio==1.60.0
grpcio-status==1.60.0
h11==0.14.0
httpcore==1.0.2
httpx==0.26.0
idna==3.6
jsonpatch==1.33
jsonpointer==2.4
langchain==0.0.352
langchain-community==0.0.5
langchain-core==0.1.2
langchain-google-genai==0.0.5
langsmith==0.0.72
marshmallow==3.20.1
multidict==6.0.4
mypy-extensions==1.0.0
numpy==1.26.2
openai==1.6.0
packaging==23.2
Pillow==10.1.0
proto-plus==1.23.0
protobuf==4.25.1
pyasn1==0.5.1
pyasn1-modules==0.3.0
pydantic==2.5.2
pydantic_core==2.14.5
PyYAML==6.0.1
requests==2.31.0
rsa==4.9
sniffio==1.3.0
SQLAlchemy==2.0.23
tenacity==8.2.3
tqdm==4.66.1
typing-inspect==0.9.0
typing_extensions==4.9.0
urllib3==2.1.0
yarl==1.9.4


### Who can help?

_No response_

### Information

- [ ] The official example notebooks/scripts
- [ ] My own modified scripts

### Related Components

- [ ] LLMs/Chat Models
- [ ] Embedding Models
- [ ] Prompts / Prompt Templates / Prompt Selectors
- [ ] Output Parsers
- [ ] Document Loaders
- [ ] Vector Stores / Retrievers
- [ ] Memory
- [ ] Agents / Agent Executors
- [ ] Tools / Toolkits
- [ ] Chains
- [ ] Callbacks/Tracing
- [ ] Async

### Reproduction

```
class StreamingLLMCallbackHandler(BaseCallbackHandler):
    """Callback handler for streaming LLM responses to a queue."""

    def __init__(self):
        """
        Initialize the StreamingLLMCallbackHandler.
        """
        self._is_done = False
        self._queue = Queue()

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """
        Run on new LLM token. Only available when streaming is enabled.

        Args:
            token (str): The new LLM token.
            **kwargs (Any): Additional keyword arguments.
        """
        if is_dev_mode():
            print(token, end='')
        self._queue.put(EventData(content=token))

    def on_llm_end(self, response: LLMResult, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any):
        """
        Run when LLM ends running.

        Args:
            response (LLMResult): The LLM processing result.
            run_id (UUID): The unique identifier for the current run.
            parent_run_id (Optional[UUID]): The unique identifier for the parent run.
            **kwargs (Any): Additional keyword arguments.
        """
        self._queue.put(EventData(content=None, finish_reason='done'))
        self._is_done = True

  
  class LangChainChatService(ChatBaseService):
    """
    Service class for handling chat operations using LangChain models.

    This class extends the ChatBaseService and provides methods for streaming chat and one-time HTTP chat using
    LangChain models.

    Methods:
        stream_chat_async: Stream chat using similarity search and AI.
        http_chat_async: Perform one-time HTTP chat using similarity search and AI.
        _qa_task: Helper method to execute the QA task in a separate thread.
        _get_qa_chain: Get the conversational retrieval chain for handling chat.
        _publish_chat_history: Publish chat history to Kafka.

    Attributes:
        model (LangchainChatModel): The LangChain chat model.
        request_manager: The API request manager.
        tool_query: The tool query for chat operations.
    """

    def __init__(self, model: LangchainChatModel, tool_query=None):
        """
        Initializes the LangChainChatService.

        Args:
            model (LangchainChatModel): The LangChain chat model.
            tool_query (Optional[str]): The tool query for chat operations.
        """
        super().__init__()
        self.model = model
        self.request_manager = api_request_manager_var.get()
        self.tool_query = tool_query
     async def http_chat_async(self) -> dict:
        """
        Perform one-time HTTP chat using similarity search and AI.

        Returns:
            dict: The response from the chat operation.
        """
        formatted_chat_history = [] if self.tool_query else self.get_formatted_chat_history(self.model.chat_history)
        qa = self._get_qa_chain(callbacks=[streaming_handler])
        qa.return_source_documents = self.model.return_source_documents
        qa.return_generated_question = True

        query_start = time.time()
        question = self.tool_query or self.model.query
        qa_response = await qa.ainvoke({"question": question, "chat_history": formatted_chat_history})
        query_end = time.time()

        result = {
            'query_result': qa_response.get("answer"),
            'query_time': int((query_end - query_start) * 1000),
            'generated_question': qa_response.get('generated_question'),
            'source_documents': [document.__dict__ for document in qa_response.get("source_documents", [])],
        }

        self._publish_chat_history(result)

        return result
 def _get_qa_chain(self, callbacks: Callbacks = None) -> BaseConversationalRetrievalChain:
        """
        Get the conversational retrieval chain for handling chat.

        Args:
            callbacks (Callbacks): The callbacks to be used.

        Returns:
            BaseConversationalRetrievalChain: The conversational retrieval chain.
        """
        collection_name = get_langchain_collection_name(self.model.client_id)
        connection_args = {"host": AppConfig.vector_db_host, "port": AppConfig.vector_db_port}
        embeddings = LLMSelector(self.model).get_embeddings()
        vector_store = Milvus(embeddings, collection_name=collection_name, connection_args=connection_args)

        expression = get_expression_to_fetch_db_text_from_ids(**self.model.model_dump())
       # this LLMSelector class return ChatGoogleGenerativeAI instance with streaming
        qa_llm = LLMSelector(self.model).get_language_model(streaming=self.model.stream_response, callbacks=callbacks)
        condense_question_llm = LLMSelector(self.model).get_language_model()
        prompt_selector = get_prompt_selector(human_context=self.model.human_context, system_context=self.model.system_context)

        qa = ConversationalRetrievalChain.from_llm(
            llm=qa_llm,
            retriever=vector_store.as_retriever(search_type="similarity", search_kwargs={"k": self.model.similarity_top_k, 'expr': expression}),
            condense_question_llm=condense_question_llm,
            combine_docs_chain_kwargs={"prompt": prompt_selector.get_prompt(qa_llm)}
        )
        return qa
```

### Expected behavior

Chain should trigger `on_llm_new_token` method of callback handler when streaming true for ChatGoogleGenerativeAI.