add support for AWS Bedrock as LLM provider

[pslusarz]

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

[codiumai-pr-agent-pro]

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.

[codiumai-pr-agent-pro]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	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