import time
from typing import Any, Dict, List, Optional
from uuid import UUID
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_core.utils import guard_import
from langchain_community.callbacks.utils import import_pandas
# Define constants
# LLMResult keys
TOKEN_USAGE = "token_usage"
TOTAL_TOKENS = "total_tokens"
PROMPT_TOKENS = "prompt_tokens"
COMPLETION_TOKENS = "completion_tokens"
RUN_ID = "run_id"
MODEL_NAME = "model_name"
GOOD = "good"
BAD = "bad"
NEUTRAL = "neutral"
SUCCESS = "success"
FAILURE = "failure"
# Default values
DEFAULT_MAX_TOKEN = 65536
DEFAULT_MAX_DURATION = 120000
# Fiddler specific constants
PROMPT = "prompt"
RESPONSE = "response"
CONTEXT = "context"
DURATION = "duration"
FEEDBACK = "feedback"
LLM_STATUS = "llm_status"
FEEDBACK_POSSIBLE_VALUES = [GOOD, BAD, NEUTRAL]
# Define a dataset dictionary
_dataset_dict = {
PROMPT: ["fiddler"] * 10,
RESPONSE: ["fiddler"] * 10,
CONTEXT: ["fiddler"] * 10,
FEEDBACK: ["good"] * 10,
LLM_STATUS: ["success"] * 10,
MODEL_NAME: ["fiddler"] * 10,
RUN_ID: ["123e4567-e89b-12d3-a456-426614174000"] * 10,
TOTAL_TOKENS: [0, DEFAULT_MAX_TOKEN] * 5,
PROMPT_TOKENS: [0, DEFAULT_MAX_TOKEN] * 5,
COMPLETION_TOKENS: [0, DEFAULT_MAX_TOKEN] * 5,
DURATION: [1, DEFAULT_MAX_DURATION] * 5,
}
[docs]def import_fiddler() -> Any:
"""Import the fiddler python package and raise an error if it is not installed."""
return guard_import("fiddler", pip_name="fiddler-client")
# First, define custom callback handler implementations
[docs]class FiddlerCallbackHandler(BaseCallbackHandler):
[docs] def __init__(
self,
url: str,
org: str,
project: str,
model: str,
api_key: str,
) -> None:
"""
Initialize Fiddler callback handler.
Args:
url: Fiddler URL (e.g. https://demo.fiddler.ai).
Make sure to include the protocol (http/https).
org: Fiddler organization id
project: Fiddler project name to publish events to
model: Fiddler model name to publish events to
api_key: Fiddler authentication token
"""
super().__init__()
# Initialize Fiddler client and other necessary properties
self.fdl = import_fiddler()
self.pd = import_pandas()
self.url = url
self.org = org
self.project = project
self.model = model
self.api_key = api_key
self._df = self.pd.DataFrame(_dataset_dict)
self.run_id_prompts: Dict[UUID, List[str]] = {}
self.run_id_response: Dict[UUID, List[str]] = {}
self.run_id_starttime: Dict[UUID, int] = {}
# Initialize Fiddler client here
self.fiddler_client = self.fdl.FiddlerApi(url, org_id=org, auth_token=api_key)
if self.project not in self.fiddler_client.get_project_names():
print( # noqa: T201
f"adding project {self.project}." "This only has to be done once."
)
try:
self.fiddler_client.add_project(self.project)
except Exception as e:
print( # noqa: T201
f"Error adding project {self.project}:"
"{e}. Fiddler integration will not work."
)
raise e
dataset_info = self.fdl.DatasetInfo.from_dataframe(
self._df, max_inferred_cardinality=0
)
# Set feedback column to categorical
for i in range(len(dataset_info.columns)):
if dataset_info.columns[i].name == FEEDBACK:
dataset_info.columns[i].data_type = self.fdl.DataType.CATEGORY
dataset_info.columns[i].possible_values = FEEDBACK_POSSIBLE_VALUES
elif dataset_info.columns[i].name == LLM_STATUS:
dataset_info.columns[i].data_type = self.fdl.DataType.CATEGORY
dataset_info.columns[i].possible_values = [SUCCESS, FAILURE]
if self.model not in self.fiddler_client.get_model_names(self.project):
if self.model not in self.fiddler_client.get_dataset_names(self.project):
print( # noqa: T201
f"adding dataset {self.model} to project {self.project}."
"This only has to be done once."
)
try:
self.fiddler_client.upload_dataset(
project_id=self.project,
dataset_id=self.model,
dataset={"train": self._df},
info=dataset_info,
)
except Exception as e:
print( # noqa: T201
f"Error adding dataset {self.model}: {e}."
"Fiddler integration will not work."
)
raise e
model_info = self.fdl.ModelInfo.from_dataset_info(
dataset_info=dataset_info,
dataset_id="train",
model_task=self.fdl.ModelTask.LLM,
features=[PROMPT, CONTEXT, RESPONSE],
target=FEEDBACK,
metadata_cols=[
RUN_ID,
TOTAL_TOKENS,
PROMPT_TOKENS,
COMPLETION_TOKENS,
MODEL_NAME,
DURATION,
],
custom_features=self.custom_features,
)
print( # noqa: T201
f"adding model {self.model} to project {self.project}."
"This only has to be done once."
)
try:
self.fiddler_client.add_model(
project_id=self.project,
dataset_id=self.model,
model_id=self.model,
model_info=model_info,
)
except Exception as e:
print( # noqa: T201
f"Error adding model {self.model}: {e}."
"Fiddler integration will not work."
)
raise e
@property
def custom_features(self) -> list:
"""
Define custom features for the model to automatically enrich the data with.
Here, we enable the following enrichments:
- Automatic Embedding generation for prompt and response
- Text Statistics such as:
- Automated Readability Index
- Coleman Liau Index
- Dale Chall Readability Score
- Difficult Words
- Flesch Reading Ease
- Flesch Kincaid Grade
- Gunning Fog
- Linsear Write Formula
- PII - Personal Identifiable Information
- Sentiment Analysis
"""
return [
self.fdl.Enrichment(
name="Prompt Embedding",
enrichment="embedding",
columns=[PROMPT],
),
self.fdl.TextEmbedding(
name="Prompt CF",
source_column=PROMPT,
column="Prompt Embedding",
),
self.fdl.Enrichment(
name="Response Embedding",
enrichment="embedding",
columns=[RESPONSE],
),
self.fdl.TextEmbedding(
name="Response CF",
source_column=RESPONSE,
column="Response Embedding",
),
self.fdl.Enrichment(
name="Text Statistics",
enrichment="textstat",
columns=[PROMPT, RESPONSE],
config={
"statistics": [
"automated_readability_index",
"coleman_liau_index",
"dale_chall_readability_score",
"difficult_words",
"flesch_reading_ease",
"flesch_kincaid_grade",
"gunning_fog",
"linsear_write_formula",
]
},
),
self.fdl.Enrichment(
name="PII",
enrichment="pii",
columns=[PROMPT, RESPONSE],
),
self.fdl.Enrichment(
name="Sentiment",
enrichment="sentiment",
columns=[PROMPT, RESPONSE],
),
]
def _publish_events(
self,
run_id: UUID,
prompt_responses: List[str],
duration: int,
llm_status: str,
model_name: Optional[str] = "",
token_usage_dict: Optional[Dict[str, Any]] = None,
) -> None:
"""
Publish events to fiddler
"""
prompt_count = len(self.run_id_prompts[run_id])
df = self.pd.DataFrame(
{
PROMPT: self.run_id_prompts[run_id],
RESPONSE: prompt_responses,
RUN_ID: [str(run_id)] * prompt_count,
DURATION: [duration] * prompt_count,
LLM_STATUS: [llm_status] * prompt_count,
MODEL_NAME: [model_name] * prompt_count,
}
)
if token_usage_dict:
for key, value in token_usage_dict.items():
df[key] = [value] * prompt_count if isinstance(value, int) else value
try:
if df.shape[0] > 1:
self.fiddler_client.publish_events_batch(self.project, self.model, df)
else:
df_dict = df.to_dict(orient="records")
self.fiddler_client.publish_event(
self.project, self.model, event=df_dict[0]
)
except Exception as e:
print( # noqa: T201
f"Error publishing events to fiddler: {e}. continuing..."
)
[docs] def on_llm_start(
self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
) -> Any:
run_id = kwargs[RUN_ID]
self.run_id_prompts[run_id] = prompts
self.run_id_starttime[run_id] = int(time.time() * 1000)
[docs] def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
flattened_llmresult = response.flatten()
run_id = kwargs[RUN_ID]
run_duration = int(time.time() * 1000) - self.run_id_starttime[run_id]
model_name = ""
token_usage_dict = {}
if isinstance(response.llm_output, dict):
token_usage_dict = {
k: v
for k, v in response.llm_output.items()
if k in [TOTAL_TOKENS, PROMPT_TOKENS, COMPLETION_TOKENS]
}
model_name = response.llm_output.get(MODEL_NAME, "")
prompt_responses = [
llmresult.generations[0][0].text for llmresult in flattened_llmresult
]
self._publish_events(
run_id,
prompt_responses,
run_duration,
SUCCESS,
model_name,
token_usage_dict,
)
[docs] def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
run_id = kwargs[RUN_ID]
duration = int(time.time() * 1000) - self.run_id_starttime[run_id]
self._publish_events(
run_id, [""] * len(self.run_id_prompts[run_id]), duration, FAILURE
)
