Source code for langchain.callbacks.arize_callback
from datetime import datetime
from typing import Any, Dict, List, Optional
from langchain.callbacks.base import BaseCallbackHandler
from langchain.callbacks.utils import import_pandas
from langchain.schema import AgentAction, AgentFinish, LLMResult
[docs]class ArizeCallbackHandler(BaseCallbackHandler):
    """Callback Handler that logs to Arize."""
[docs]    def __init__(
        self,
        model_id: Optional[str] = None,
        model_version: Optional[str] = None,
        SPACE_KEY: Optional[str] = None,
        API_KEY: Optional[str] = None,
    ) -> None:
        """Initialize callback handler."""
        super().__init__()
        self.model_id = model_id
        self.model_version = model_version
        self.space_key = SPACE_KEY
        self.api_key = API_KEY
        self.prompt_records: List[str] = []
        self.response_records: List[str] = []
        self.prediction_ids: List[str] = []
        self.pred_timestamps: List[int] = []
        self.response_embeddings: List[float] = []
        self.prompt_embeddings: List[float] = []
        self.prompt_tokens = 0
        self.completion_tokens = 0
        self.total_tokens = 0
        self.step = 0
        from arize.pandas.embeddings import EmbeddingGenerator, UseCases
        from arize.pandas.logger import Client
        self.generator = EmbeddingGenerator.from_use_case(
            use_case=UseCases.NLP.SEQUENCE_CLASSIFICATION,
            model_name="distilbert-base-uncased",
            tokenizer_max_length=512,
            batch_size=256,
        )
        self.arize_client = Client(space_key=SPACE_KEY, api_key=API_KEY)