Skip to content

Commit

Permalink
chore(llmobs): extract bedrock LLMObs tests to separate file (#9408)
Browse files Browse the repository at this point in the history
This PR moves LLMObs tests from the test_bedrock.py file into
test_bedrock_llmobs.py to ease maintainability for LLMObs tests in the
future. No functionality is added/removed.

## Checklist

- [x] Change(s) are motivated and described in the PR description
- [x] Testing strategy is described if automated tests are not included
in the PR
- [x] Risks are described (performance impact, potential for breakage,
maintainability)
- [x] Change is maintainable (easy to change, telemetry, documentation)
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed or label `changelog/no-changelog` is set
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/))
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))
- [x] If this PR changes the public interface, I've notified
`@DataDog/apm-tees`.

## Reviewer Checklist

- [x] Title is accurate
- [x] All changes are related to the pull request's stated goal
- [x] Description motivates each change
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- [x] Testing strategy adequately addresses listed risks
- [x] Change is maintainable (easy to change, telemetry, documentation)
- [x] Release note makes sense to a user of the library
- [x] Author has acknowledged and discussed the performance implications
of this PR as reported in the benchmarks PR comment
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
  • Loading branch information
Yun-Kim authored May 29, 2024
1 parent a864142 commit d89da44
Show file tree
Hide file tree
Showing 3 changed files with 379 additions and 308 deletions.
88 changes: 88 additions & 0 deletions tests/contrib/botocore/bedrock_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
import os


try:
import vcr
except ImportError:
vcr = None
get_request_vcr = None


_MODELS = {
"ai21": "ai21.j2-mid-v1",
"amazon": "amazon.titan-tg1-large",
"anthropic": "anthropic.claude-instant-v1",
"anthropic_message": "anthropic.claude-3-sonnet-20240229-v1:0",
"cohere": "cohere.command-light-text-v14",
"meta": "meta.llama2-13b-chat-v1",
}

_REQUEST_BODIES = {
"ai21": {
"prompt": "Explain like I'm a five-year old: what is a neural network?",
"temperature": 0.9,
"topP": 1.0,
"maxTokens": 10,
"stopSequences": [],
},
"amazon": {
"inputText": "Command: can you explain what Datadog is to someone not in the tech industry?",
"textGenerationConfig": {"maxTokenCount": 50, "stopSequences": [], "temperature": 0, "topP": 0.9},
},
"anthropic": {
"prompt": "\n\nHuman: %s\n\nAssistant: What makes you better than Chat-GPT or LLAMA?",
"temperature": 0.9,
"top_p": 1,
"top_k": 250,
"max_tokens_to_sample": 50,
"stop_sequences": ["\n\nHuman:"],
},
"anthropic_message": {
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "summarize the plot to the lord of the rings in a dozen words"}],
}
],
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 50,
"temperature": 0,
},
"cohere": {
"prompt": "\n\nHuman: %s\n\nAssistant: Can you explain what a LLM chain is?",
"temperature": 0.9,
"p": 1.0,
"k": 0,
"max_tokens": 10,
"stop_sequences": [],
"stream": False,
"num_generations": 1,
},
"meta": {
"prompt": "What does 'lorem ipsum' mean?",
"temperature": 0.9,
"top_p": 1.0,
"max_gen_len": 60,
},
}


# VCR is used to capture and store network requests made to OpenAI and other APIs.
# This is done to avoid making real calls to the API which could introduce
# flakiness and cost.
# To (re)-generate the cassettes: pass a real API key with
# {PROVIDER}_API_KEY, delete the old cassettes and re-run the tests.
# NOTE: be sure to check that the generated cassettes don't contain your
# API key. Keys should be redacted by the filter_headers option below.
# NOTE: that different cassettes have to be used between sync and async
# due to this issue: https://github.com/kevin1024/vcrpy/issues/463
# between cassettes generated for requests and aiohttp.
def get_request_vcr():
return vcr.VCR(
cassette_library_dir=os.path.join(os.path.dirname(__file__), "bedrock_cassettes/"),
record_mode="once",
match_on=["path"],
filter_headers=["authorization", "X-Amz-Security-Token"],
# Ignore requests to the agent
ignore_localhost=True,
)
Loading

0 comments on commit d89da44

Please sign in to comment.