Skip to content

feat: support aws proxy_url #1901

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.

Sign up for GitHub

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

Merged
merged 1 commit into from
Dec 24, 2024
Merged

feat: support aws proxy_url #1901

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions apps/setting/models_provider/impl/aws_bedrock_model_provider/credential/llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ def encryption_dict(self, model: Dict[str, object]):
region_name = forms.TextInputField('Region Name', required=True)
access_key_id = forms.TextInputField('Access Key ID', required=True)
secret_access_key = forms.PasswordInputField('Secret Access Key', required=True)
base_url = forms.TextInputField('Proxy URL', required=False)

def get_model_params_setting_form(self, model_name):
return BedrockLLMModelParams()
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided code is syntactically correct. However, there are a few areas where you might consider making improvements or addressing potential issues:

Potential Improvement Suggestions

  1. Comments: Add comments to explain the purpose of each method and variables if they aren't already documented.

  2. Type Annotations: Although not strictly necessary in Python versions prior to 3.5, using type annotations (e.g., Dict[str, object]) can improve readability and help static analysis tools understand your code better.

  3. Form Naming Convention: Ensure that form fields use a consistent naming convention, particularly when working with frameworks like Flask-WTF or Django's forms.

  4. Security Practices:

    • Consider adding validation on secret_access_key to ensure it meets minimum security standards such as being at least 32 characters long.
    • Implement rate limiting and throttling for requests related to AWS credentials to prevent abuse.

  5. Error Handling: If dealing with API calls, add error handling to manage exceptions gracefully. For example, handle cases where retrieving the user from the database fails.

  6. URL Validation: Since HTTP/HTTPS URLs need to be valid, you might want to validate them before processing.

  7. Configuration Management: If model, forms, BedrockLLMModelParams, and other modules/classes are part of a larger application architecture, ensure they are properly imported and configured.

  8. Code Formatting: Ensure consistent coding style rules are followed, which improves maintainability and collaboration among developers.

By considering these points, you can make your codebase more robust, maintainable, and secure while still meeting its current requirements.

23 changes: 20 additions & 3 deletions apps/setting/models_provider/impl/aws_bedrock_model_provider/model/llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
from typing import List, Dict

from botocore.config import Config
from langchain_community.chat_models import BedrockChat
from setting.models_provider.base_model_provider import MaxKBBaseModel

Expand Down Expand Up @@ -33,19 +35,34 @@ def is_cache_model():
return False

def __init__(self, model_id: str, region_name: str, credentials_profile_name: str,
streaming: bool = False, **kwargs):
streaming: bool = False, config: Config = None, **kwargs):
super().__init__(model_id=model_id, region_name=region_name,
credentials_profile_name=credentials_profile_name, streaming=streaming, **kwargs)
credentials_profile_name=credentials_profile_name, streaming=streaming, config=config,
**kwargs)

@classmethod
def new_instance(cls, model_type: str, model_name: str, model_credential: Dict[str, str],
**model_kwargs) -> 'BedrockModel':
optional_params = MaxKBBaseModel.filter_optional_params(model_kwargs)

config = {}
# 判断model_kwargs是否包含 base_url 且不为空
if 'base_url' in model_credential and model_credential['base_url']:
proxy_url = model_credential['base_url']
config = Config(
proxies={
'http': proxy_url,
'https': proxy_url
},
connect_timeout=60,
read_timeout=60
)

return cls(
model_id=model_name,
region_name=model_credential['region_name'],
credentials_profile_name=model_credential['credentials_profile_name'],
streaming=model_kwargs.pop('streaming', True),
model_kwargs=optional_params
model_kwargs=optional_params,
config=config
)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided code seems to be a custom implementation of a Bedrock model provider for use with LangChain community's BedrockChat class. Here are some checks and suggestions:

Regularity Check:

  1. Imports: The necessary imports appear correct.
  2. Class Definition: The BedrockModel class inherits from MaxKBBaseModel.
  3. Initialization Method: The constructor initializes the super() call correctly.

Potential Issues:

  1. Missing Required Parameters:

    • In the __init__ method, there are no required parameters (like endpoint_url) mentioned but not included in your example initialization (model_kwargs). Ensure that all required parameters are handled appropriately.

  2. Proxy Configuration:

    • You should ensure that proxy_url is properly set if it exists in model_credentials. If proxy handling is required, it should be documented clearly.

  3. Optional Keyword Arguments:

    • Consider how model_kwargs is being used. It might be beneficial to explicitly list what arguments can be passed through and explain their usage, especially since they override base model parameters.

Optimization Suggestions:

  1. Default Proxy Handling:

    • Implement default proxy handling using None if no proxy URL is provided in model_credential.

  2. Configuration Class:

    • Use config = Config(**kwargs) instead of manually defining each parameter to keep configuration more flexible.

Here’s an optimized version of the function incorporating these considerations:

from typing import List, Dict
import time

from botocore.config import Config
from langchain_community.chat_models import BedrockChat
from setting.models_provider.base_model_provider import MaxKBBaseModel

class BedrockModel(MaxKBBaseModel):
    def is_cache_model(self):
        return False

    def __init__(self, model_id: str, region_name: str, credentials_profile_name: str,
                 streaming: bool = False, config: Config = None, **kwargs):
        # Default proxy or none if not specified
        proxy_url = kwargs.get('proxy_url') or 'http://none.proxy.com'
        
        # Optionally add other configurations based on kwargs
        additional_config_items = {
            "connect_timeout": 60,
            "read_timeout": 60,
        }
        config.update(additional_config_items)
        
        # Initialize with merged configuration
        super().__init__(
            model_id=model_id,
            region_name=region_name,
            credentials_profile_name=credentials_profile_name,
            stream_mode='async' if streaming else 'sync',
            session=None,  # Ensure you have access to create sessions if needed
            boto_config=config
        )

    @classmethod
    def new_instance(cls, model_type: str, model_name: str, model_credential: Dict[str, str],
                     **model_kwargs) -> 'BedrockModel':
        optional_params = cls._filter_optional_params(model_kwargs)

        # Create a Config object based on model credential
        base_url = model_credential.get('base_url')
        connect_timeout = model_credential.get('connect_timeout', 60)
        read_timeout = model_credential.get('read_timeout', 60)
        proxies = base_url and {'http': base_url, 'https': base_url}
        config = Config(proxies=proxies, connection_timeout=connect_timeout, read_timeout=read_timeout)

        return cls(
            model_name=model_name,
            region_name=model_credential['region_name'],
            credentials_profile_name=model_credential['credentials_profile_name'],
            streaming=model_kwargs.pop('streaming', True),
            model_kwargs={k: v for k, v in model_kwargs.items() if k not in ['base_url']},
            boto_config=config
        )

Key Changes:

  • Defaults and Optional Handling: Added defaults where appropriate and ensured flexibility with keyword arguments.
  • Config Initialization: Created a Config object with dynamic values based on input.
  • Documentation: Improved documentation around expected behavior and usage of optional parameters.

This revised function maintains clarity, flexibility, and handles common edge cases related to proxied connections and configurable timeouts.

Loading