Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add support for AWS Bedrock as LLM provider #61

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

pslusarz
Copy link

@pslusarz pslusarz commented Oct 2, 2024

User description

  • resolved dependency conflicts that are currently present in main that prevent this from being built from a clean checkout
  • pinned litellm at 1.48.2 until ([Bug]: cross-region-inferencing max_tokens_to_sample being sent to Amazon Bedrock models BerriAI/litellm#6003) is solved
  • introduced code for bedrock models. Not happy with this implementation, but I would need to rewrite the entire ai_handler module to make it right. Ideally, this should just pass everything to litellm and not get in the way. This was tested manually by running solve_my_problem example against both openai and aws bedrock
  • updated readme with bedrock instructions
  • left config in the same state as in main branch, so defaulting to openai
  • added a couple of useful litellm switches to the config, although did not implement a full set of config options at this point

PR Type

enhancement, documentation, dependencies


Description

  • Added support for AWS Bedrock as a new LLM provider in the AiHandler class, including a new Provider enum and refactored chat completion logic.
  • Updated the README with detailed instructions for configuring AWS Bedrock and clarified setup steps for different providers.
  • Modified the configuration file to include examples for AWS Bedrock models and added new options for litellm.
  • Updated and pinned several dependencies to resolve conflicts and ensure compatibility, including fastapi, litellm, tiktoken, and boto3.

Changes walkthrough 📝

Relevant files
Enhancement
ai_handler.py
Add support for AWS Bedrock and refactor provider handling

alpha_codium/llm/ai_handler.py

  • Introduced a new Provider enum to handle different LLM providers.
  • Added support for AWS Bedrock as a provider in the AiHandler class.
  • Refactored chat completion logic into separate methods for each
    provider.
  • Improved logging to include provider and model information.
  • +100/-61
    Documentation
    README.md
    Update README with AWS Bedrock configuration instructions

    README.md

  • Added instructions for configuring AWS Bedrock as a model provider.
  • Clarified setup steps for OpenAI and AWS Bedrock.
  • Fixed a minor typo in the documentation.
  • +10/-2   
    Configuration changes
    configuration.toml
    Update configuration for AWS Bedrock and litellm options 

    alpha_codium/settings/configuration.toml

  • Added commented-out configuration examples for AWS Bedrock models.
  • Updated default model to gpt-3.5-turbo-16k.
  • Added new configuration options for litellm.
  • +7/-2     
    Dependencies
    requirements.txt
    Update and pin dependencies for compatibility                       

    requirements.txt

  • Updated fastapi to version 0.115.0.
  • Pinned litellm to version 1.48.2.
  • Updated tiktoken and boto3 to more flexible version specifications.
  • +5/-5     

    💡 PR-Agent usage: Comment /help "your question" on any pull request to receive relevant information

    …project to build (it was not building on the main branch)
    
    - pinned litellm at 1.48.2 until (BerriAI/litellm#6003) is solved
    - introduced code for bedrock models. Not happy with this implementation, but I would need to rewrite the entire ai_handler module to make it right. Ideally, this should just pass everything to litellm and not get in the way. This was tested manually by running solve_my_problem example against both openai and aws bedrock
    - updated readme with bedrock instructions
    - left config in the same state as in main branch, so defaulting to openai
    - added a couple of useful litellm switches to the config, although did not implement a full set of config options at this point
    @qodo-merge-pro qodo-merge-pro bot added documentation Improvements or additions to documentation enhancement New feature or request dependencies labels Oct 2, 2024
    Copy link

    qodo-merge-pro bot commented Oct 2, 2024

    PR Reviewer Guide 🔍

    Here are some key observations to aid the review process:

    ⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪
    🧪 No relevant tests
    🔒 Security concerns

    Sensitive information exposure:
    The README.md file instructs users to store their OpenAI API key in a .secrets.toml file. While this is better than hardcoding, it's still not the most secure method. Consider recommending the use of environment variables or a secure secret management system for API keys and other sensitive information.

    ⚡ Recommended focus areas for review

    Code Complexity
    The AiHandler class has become more complex with the addition of multiple providers. Consider refactoring to use a strategy pattern or factory method to handle different providers more cleanly.

    Error Handling
    The error handling in the chat_completion method is not provider-specific. Consider implementing more granular error handling for each provider.

    Configuration Management
    The configuration file now includes commented-out options for different models. Consider using a more structured approach to manage multiple configurations, such as separate files for different environments.

    Copy link

    qodo-merge-pro bot commented Oct 2, 2024

    PR Code Suggestions ✨

    Explore these optional code suggestions:

    CategorySuggestion                                                                                                                                    Score
    Maintainability
    Refactor the initialization logic into separate methods for each provider to improve readability and maintainability

    The init method is becoming complex with multiple provider checks. Consider
    refactoring this into a separate method for better readability and maintainability.

    alpha_codium/llm/ai_handler.py [26-63]

     def __init__(self):
         self.provider = Provider.UNKNOWN
    -    """
    -    Initializes the OpenAI API key and other settings from a configuration file.
    -    Raises a ValueError if the OpenAI key is missing.
    -    """
         self.limiter = AsyncLimiter(get_settings().config.max_requests_per_minute)
    -    if get_settings().get("config.model").lower().startswith("bedrock"):
    +    self._setup_provider()
    +    self.azure = False
    +    litellm.set_verbose = get_settings().get("litellm.set_verbose", False)
    +    litellm.drop_params = get_settings().get("litellm.drop_params", False)
    +
    +def _setup_provider(self):
    +    model = get_settings().get("config.model").lower()
    +    if model.startswith("bedrock"):
             self.provider = Provider.AWSBEDROCK
    -    elif "gpt" in get_settings().get("config.model").lower():
    -        self.provider = Provider.OPENAI
    -        try:
    -            openai.api_key = get_settings().openai.key
    -            litellm.openai_key = get_settings().openai.key
    -        except AttributeError as e:
    -            raise ValueError("OpenAI key is required") from e
    -    elif "deepseek" in get_settings().get("config.model"):
    -        self.provider = Provider.DEEPSEEK_LEGACY
    -        litellm.register_prompt_template(
    -            model="huggingface/deepseek-ai/deepseek-coder-33b-instruct",
    -            roles={
    -                "system": {
    -                    "pre_message": "",
    -                    "post_message": "\n"
    -                },
    -                "user": {
    -                    "pre_message": "### Instruction:\n",
    -                    "post_message": "\n### Response:\n"
    -                },
    -            },
    -        )
    +    elif "gpt" in model:
    +        self._setup_openai()
    +    elif "deepseek" in model:
    +        self._setup_deepseek()
     
    +def _setup_openai(self):
    +    self.provider = Provider.OPENAI
    +    try:
    +        openai.api_key = get_settings().openai.key
    +        litellm.openai_key = get_settings().openai.key
    +    except AttributeError as e:
    +        raise ValueError("OpenAI key is required") from e
    +
    +def _setup_deepseek(self):
    +    self.provider = Provider.DEEPSEEK_LEGACY
    +    litellm.register_prompt_template(
    +        model="huggingface/deepseek-ai/deepseek-coder-33b-instruct",
    +        roles={
    +            "system": {"pre_message": "", "post_message": "\n"},
    +            "user": {"pre_message": "### Instruction:\n", "post_message": "\n### Response:\n"},
    +        },
    +    )
    +
    • Apply this suggestion
    Suggestion importance[1-10]: 9

    Why: This refactoring significantly enhances the readability and maintainability of the code by separating provider-specific logic into distinct methods, making the codebase easier to understand and modify.

    9
    Enhancement
    Use a dictionary to map providers to their respective chat completion methods for improved maintainability and extensibility

    Consider using a dictionary to map providers to their respective chat completion
    methods. This would simplify the chat_completion method and make it easier to add
    new providers in the future.

    alpha_codium/llm/ai_handler.py [84-106]

    -if self.provider == Provider.DEEPSEEK_LEGACY:
    -    response = await self._deepseek_chat_completion(
    -        model=model,
    -        system=system,
    -        user=user,
    -        temperature=temperature,
    -        frequency_penalty=frequency_penalty,
    -    )
    -elif self.provider == Provider.OPENAI:
    -    response = await self._openai_chat_completion(
    -        model=model,
    -        system=system,
    -        user=user,
    -        temperature=temperature,
    -        frequency_penalty=frequency_penalty,
    -    )
    -else:
    -    response = await self._awsbedrock_chat_completion(
    -        model=model,
    -        system=system,
    -        user=user,
    -        temperature=temperature
    -    )
    +provider_methods = {
    +    Provider.DEEPSEEK_LEGACY: self._deepseek_chat_completion,
    +    Provider.OPENAI: self._openai_chat_completion,
    +    Provider.AWSBEDROCK: self._awsbedrock_chat_completion,
    +}
    +chat_method = provider_methods.get(self.provider, self._awsbedrock_chat_completion)
    +response = await chat_method(
    +    model=model,
    +    system=system,
    +    user=user,
    +    temperature=temperature,
    +    frequency_penalty=frequency_penalty,
    +)
    • Apply this suggestion
    Suggestion importance[1-10]: 8

    Why: This suggestion improves the maintainability and extensibility of the code by using a dictionary to map providers to their respective methods, simplifying the logic and making it easier to add new providers in the future.

    8
    Consistency
    Add the frequency_penalty parameter to the AWS Bedrock chat completion method for consistency with other provider methods

    The _awsbedrock_chat_completion method is missing the frequency_penalty parameter,
    which is used in other provider methods. Consider adding this parameter for
    consistency across all provider methods.

    alpha_codium/llm/ai_handler.py [164-174]

    -async def _awsbedrock_chat_completion(self, model, system, user, temperature):
    +async def _awsbedrock_chat_completion(self, model, system, user, temperature, frequency_penalty):
         response = await acompletion(
             model=model,
             user=user,
             messages=[
                 {"role": "system", "content": system},
                 {"role": "user", "content": user},
             ],
    -        temperature=temperature
    +        temperature=temperature,
    +        frequency_penalty=frequency_penalty
         )
         return response
    • Apply this suggestion
    Suggestion importance[1-10]: 7

    Why: Adding the frequency_penalty parameter to the AWS Bedrock method ensures consistency across all provider methods, which is important for maintaining uniform behavior and reducing potential errors.

    7

    💡 Need additional feedback ? start a PR chat

    Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
    Labels
    dependencies documentation Improvements or additions to documentation enhancement New feature or request Review effort [1-5]: 4
    Projects
    None yet
    Development

    Successfully merging this pull request may close these issues.

    1 participant