[GSK-1590] native support for claude 3 and titan embeddings on Bedrock

[celmore25]

Description

As the Claude models become more popular, it is helpful to provide native support for these models in an AWS environment. This feature adds native support for using an LLM client with Amazon Bedrock. As a first step, it supports the Claude 3 family of text generation models as well as the Titan text embedding models. Having implemented this myself, I have seen the QA sets produced by Haiku to be a great blend of speed and useful for a very low cost.

Here is how I envision this feature being used in the current library's implementation. This follows the test set generation example here.

import os
import boto3
import pandas as pd

import giskard
from giskard.llm.client.bedrock import ClaudeBedrockClient
from giskard.llm.embeddings.bedrock import BedrockEmbedding
from giskard.rag import generate_testset, KnowledgeBase
from giskard.rag import QATestset

# setup the bedrock client and embedding model
bedrock_runtime = boto3.client("bedrock-runtime", region_name=os.environ["AWS_DEFAULT_REGION"])
oc = ClaudeBedrockClient(bedrock_runtime, model="anthropic.claude-3-haiku-20240307-v1:0")
embedding_model = BedrockEmbedding(bedrock_runtime, model="amazon.titan-embed-text-v1")
giskard.llm.set_default_client(oc)

# Load your data and initialize the KnowledgeBase
df = pd.read_csv("test_faqs.csv")
knowledge_base = KnowledgeBase(
    data=df,
    embedding_model=embedding_model,
)

# Generate a testset with 10 questions & answers for each question types (this will take a while)
testset = generate_testset(
    knowledge_base, 
    num_questions=60,
    language='en',  # optional, we'll auto detect if not provided
    agent_description="A customer support chatbot for Amazon SageMaker", # helps generating better questions
)

# Save the generated testset
testset.save("my_testset.jsonl")

# You can easily load it back
loaded_testset = QATestset.load("my_testset.jsonl")

# Convert it to a pandas dataframe
testset_df = loaded_testset.to_pandas()

Related Issue

GSK-1590

Type of Change

• 📚 Examples / docs / tutorials / dependencies update
• 🔧 Bug fix (non-breaking change which fixes an issue)
• 🥂 Improvement (non-breaking change which improves an existing feature)
• 🚀 New feature (non-breaking change which adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to change)
• 🔐 Security fix

Checklist

• I've read the CODE_OF_CONDUCT.md document.
• I've read the CONTRIBUTING.md guide.
• I've written tests for all new methods and classes that I created.
• I've written the docstring in Google format for all the methods and classes that I used.
• I've updated the pdm.lock running pdm update-lock (only applicable when pyproject.toml has been
  modified)

[luca-martial]

Thanks a lot for this contribution @celmore25 ! We'll review it this week

[kevinmessiaen]

Thanks for your contribution!

I've done a small modification to ensure that the Client will still work when instructed to use format or seed. Those params aren't necessary for the model to work properly for the scan and tests since the format are also instructed in our promps and seeds are nice to have.

[celmore25]

@kevinmessiaen - I not aware why those parameters were there; thanks for fixing that and for the clarification :)

Looks like the merge is blocked based on this failed test.

FAILED tests/integrations/test_mlflow.py::test_fast[drug_classification_data-drug_classification_model]

I am not sure what the issue is here. I did not see any change with our commits would have impacted that integration test. Any thoughts here?

[kevinmessiaen]

@kevinmessiaen - I not aware why those parameters were there; thanks for fixing that and for the clarification :)

Looks like the merge is blocked based on this failed test.

FAILED tests/integrations/test_mlflow.py::test_fast[drug_classification_data-drug_classification_model]

I am not sure what the issue is here. I did not see any change with our commits would have impacted that integration test. Any thoughts here?

It seems that the issue is global as we have it in other branches. I'll investigate and fix it on main so we can propagate the fix in other branches.

PS: It seems that the error was caused by the release to mlflow 2.12.1 which cause an error with imbalance-learn

[celmore25]

Thanks for making the fix @kevinmessiaen