How to wrap AzureOpenAIEmbeddings to emit custom events with number of input_tokens?

[ai-learner-00]

Since there are no callbacks for embeddings Total token count of openai callback does not count embedding usage, I am trying to emit a custom event that will include input_tokens: int calculated by tiktoken.

from typing import Any, Dict, List, Optional
from uuid import UUID

from langchain_core.callbacks import AsyncCallbackHandler
from langchain_core.callbacks.manager import (
    adispatch_custom_event,
)
from langchain_core.runnables import RunnableLambda
from langchain_core.runnables.config import RunnableConfig


class AsyncCustomCallbackHandler(AsyncCallbackHandler):
    async def on_custom_event(
        self,
        name: str,
        data: Any,
        *,
        run_id: UUID,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        print(
            f"Received event {name} with data: {data}, with tags: {tags}, with metadata: {metadata} and run_id: {run_id}"
        )


@RunnableLambda
async def bar(x: str, config: RunnableConfig) -> str:
    """An example that shows how to manually propagate config.

    You must do this if you're running python<=3.10.
    """
    await adispatch_custom_event("event1", {"x": x}, config=config)
    await adispatch_custom_event("event2", 5, config=config)
    return x


async_handler = AsyncCustomCallbackHandler()
await bar.ainvoke("a", {"callbacks": [async_handler], "tags": ["foo", "bar"]})

How should I approach wrapping the AzureOpenAIEmbeddings class to emit events similar to llm_start and llm_end? I realized it doesn't even have a callbacks parameter.

embeddings = AzureOpenAIEmbeddings(
    azure_deployment="text-embedding-ada-002", callbacks=[] # No parameter named "callbacks"
)

[Uvi-12]

To achieve this, you can create a wrapper class around AzureOpenAIEmbeddings. The wrapper will calculate token usage using tiktoken, emit custom events like llm_start and llm_end before and after calling the embedding method, and delegate the actual embedding process to the original class. Since AzureOpenAIEmbeddings doesn’t have a callbacks parameter, you need to handle the event dispatching manually by emitting events using adispatch_custom_event within the wrapper.

This approach allows you to track metrics (like token counts) and integrate with LangChain’s callback system without modifying the base class.