Added generic event handler for both tockens and functions calls

[andrewBatutin]

Description

Main motivation for this PR is to sync with JS langchain langchain-ai/langchainjs#2025

Added on_event callback that works for both token and openai function calls in streaming mode

Twitter: @ShelfDev

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Ignored Deployment

Name	Status	Preview	Comments	Updated (UTC)
langchain	⬜️ Ignored (Inspect)	Visit Preview		Aug 24, 2023 4:16pm

[hinthornw]

How about changing

def on_llm_new_token(
        self,
        token: str,
        **kwargs: Any,
    ) -> None:

To

def on_llm_new_token(
        self,
        token: str,
        chunk: Optional[Union[GenerationChunk,ChatGenerationChunk]] = None,
        **kwargs: Any,
    ) -> None:
        """Run when LLM generates a new token.

        Args:
            token (str): The new token.
            chunk (GenerationChunk | ChatGenerationChunk): The new generated chunk, containing
                content and other information.
        """

There's room for a generic on_event() callback but not scoped to the llm run manager and the input arg would have to be "Any" typed

[andrewBatutin]

How about changing

def on_llm_new_token(
        self,
        token: str,
        **kwargs: Any,
    ) -> None:

To

def on_llm_new_token(
        self,
        token: str,
        chunk: Optional[Union[GenerationChunk,ChatGenerationChunk]] = None,
        **kwargs: Any,
    ) -> None:
        """Run when LLM generates a new token.

        Args:
            token (str): The new token.
            chunk (GenerationChunk | ChatGenerationChunk): The new generated chunk, containing
                content and other information.
        """

There's room for a generic on_event() callback but not scoped to the llm run manager and the input arg would have to be "Any" typed

@hinthornw implemented, please check🙏

[hinthornw]

I think we also ened to update the BaseTracer

def on_llm_new_token(
        self,
        token: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Run on new LLM token. Only available when streaming is enabled."""
        if not run_id:
            raise TracerException("No run_id provided for on_llm_new_token callback.")

        run_id_ = str(run_id)
        llm_run = self.run_map.get(run_id_)
        if llm_run is None or llm_run.run_type != "llm":
            raise TracerException(f"No LLM Run found to be traced for {run_id}")
        llm_run.events.append(
            {
                "name": "new_token",
                "time": datetime.utcnow(),
                "kwargs": {"token": token},
            },
        )
        ```
        
        and add the chunk to the events dict

[andrewBatutin]

I think we also ened to update the BaseTracer

def on_llm_new_token(
        self,
        token: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Run on new LLM token. Only available when streaming is enabled."""
        if not run_id:
            raise TracerException("No run_id provided for on_llm_new_token callback.")

        run_id_ = str(run_id)
        llm_run = self.run_map.get(run_id_)
        if llm_run is None or llm_run.run_type != "llm":
            raise TracerException(f"No LLM Run found to be traced for {run_id}")
        llm_run.events.append(
            {
                "name": "new_token",
                "time": datetime.utcnow(),
                "kwargs": {"token": token},
            },
        )
        ```
        
        and add the chunk to the events dict

added , please review @hinthornw

[hinthornw]

LGTM pending lint passing. Thank you! 🎉