Skip to content

Add support for custom LLM parameters - #1961

Closed
bobbywlindsey wants to merge 16 commits into
crewAIInc:mainfrom
bobbywlindsey:main
Closed

Add support for custom LLM parameters#1961
bobbywlindsey wants to merge 16 commits into
crewAIInc:mainfrom
bobbywlindsey:main

Conversation

@bobbywlindsey

@bobbywlindsey bobbywlindsey commented Jan 23, 2025

Copy link
Copy Markdown
Contributor

Custom deployment of LLMs sometimes require parameters that either don't exist in the crewai.LLM class parameters or require different naming conventions.

This change adds the capability to pass in custom parameters to the crewai.LLM class which then forwards them on for LiteLLM completion. Usage would look like the following:

Agent(
    config=self.agents_config['researcher'],
    verbose=True,
    llm=LLM(
        model="sagemaker/<my-endpoint>",
        temperature=0.6,
        top_p=0.9,
        custom_params={"details": True, "max_new_tokens": 1000}
    )
)

@joaomdmoura

Copy link
Copy Markdown
Collaborator

Disclaimer: This review was made by a crew of AI Agents.

Code Review Comment for SageMaker LLM Provider Implementation

Overview

This pull request introduces Amazon SageMaker as a new LLM (Large Language Model) provider within CrewAI, encompassing both documentation and code modifications across multiple files.

Documentation Changes (docs/concepts/llms.mdx)

Strengths:

  • The documentation structure follows existing conventions and clearly outlines the setup process for SageMaker.
  • Includes comprehensive examples for configuring SageMaker, which is beneficial for new users.

Suggestions for Improvement:

  • Add explicit details regarding SageMaker endpoint requirements.
  • Outline region-specific considerations to help users avoid common pitfalls.

Recommended Documentation Enhancements:

<Accordion title="Amazon SageMaker">
    ```python
    # Required environment variables
    AWS_ACCESS_KEY_ID=<your-access-key>
    AWS_SECRET_ACCESS_KEY=<your-secret-key>
    AWS_DEFAULT_REGION=<your-region>
    
    # Optional SageMaker-specific configurations
    SAGEMAKER_ENDPOINT_CONFIG={
        "InferenceComponentName": "<component-name>",
        "InitialInstanceCount": 1
    }
    ```

    Example usage:
    ```python
    llm = LLM(
        model="sagemaker/<my-endpoint>",
        temperature=0.7,
        max_tokens=2000
    )
    ```
    
    *Note: Ensure your SageMaker endpoint is deployed and accessible in the specified region. For details on endpoint configuration, see AWS SageMaker documentation.*
</Accordion>

CLI Constants Changes (src/crewai/cli/constants.py)

Issues Identified:

  • Successfully removed duplicated AWS credentials handling between Bedrock and SageMaker, improving clarity and reducing redundancy.
  • Cleaned up unnecessary constants, enhancing maintainability.

Improvement Suggestions:

Consider implementing a shared AWS credentials handler across providers to streamline user input processes:

# src/crewai/cli/constants.py

AWS_COMMON_CREDENTIALS = [
    {
        "prompt": "Enter your AWS Access Key ID (press Enter to skip)",
        "key_name": "AWS_ACCESS_KEY_ID",
    },
    {
        "prompt": "Enter your AWS Secret Access Key (press Enter to skip)",
        "key_name": "AWS_SECRET_ACCESS_KEY",
    },
    {
        "prompt": "Enter your AWS Region Name (press Enter to skip)",
        "key_name": "AWS_REGION_NAME",
    },
]

PROVIDER_CONFIGS = {
    "bedrock": AWS_COMMON_CREDENTIALS,
    "sagemaker": AWS_COMMON_CREDENTIALS,
}

LLM Class Changes (src/crewai/llm.py)

Strengths:

  • Clean implementation of custom parameters provides flexibility while preserving backwards compatibility.

Key Issues Identified:

  • Potential for KeyError when accessing custom parameters needs addressing.
  • Missing type hints for custom_params can cause confusion about expected data types.
  • Lack of validation for custom parameters could lead to unexpected behavior.

Recommended Improvements:

To enhance flexibility and reliability, consider implementing error handling and validation for custom parameters:

from typing import Dict, Any, Optional, List

class LLM:
    def __init__(self, model: str, timeout: int = 600, temperature: float = 0.7, max_tokens: Optional[int] = None,
                 api_version: Optional[str] = None, api_key: Optional[str] = None,
                 callbacks: List[Any] = [], custom_params: Optional[Dict[str, Any]] = None):
        self.model = model
        self.timeout = timeout
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.api_version = api_version
        self.api_key = api_key
        self.callbacks = callbacks
        self.custom_params = custom_params or {}

    def call(self, prompt: str, tools: Optional[List[Dict]] = None) -> str:
        try:
            params = {
                "model": self.model,
                "messages": self._build_messages(prompt),
                "temperature": self.temperature,
                "max_tokens": self.max_tokens,
                "api_version": self.api_version,
                "api_key": self.api_key,
                "stream": False,
                "tools": tools,
                **self.custom_params
            }
            return self._execute_call(params)
        except Exception as e:
            raise LLMError(f"Error calling LLM: {str(e)}")

Security Considerations

  • Ensure AWS credentials are handled securely.
  • Implement validation for custom parameters to prevent injection attacks.
  • Consider logging for custom parameter usage to facilitate debugging.

Performance Considerations

  • Implement caching for AWS credential resolution.
  • Consider connection pooling for SageMaker endpoints.
  • Improve timeout management for SageMaker calls.

Testing Recommendations

  • Develop unit tests for custom parameter handling.
  • Include integration tests with SageMaker endpoints.
  • Add error handling tests for invalid configurations.

Overall Assessment

The implementation is solid but could benefit from additional error handling and type safety. The removal of duplicate AWS credentials handling is a significant positive change that enhances maintainability. Implementing the suggested improvements would further boost reliability and usability while maintaining the integrity of the existing code. These changes lay a strong foundation for SageMaker integration in CrewAI.

@bhancockio

Copy link
Copy Markdown
Contributor

We updated LLM.py so that you can now pass in custom params:

class LLM:
    def __init__(
        self,
        model: str,
        timeout: Optional[Union[float, int]] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        n: Optional[int] = None,
        stop: Optional[Union[str, List[str]]] = None,
        max_completion_tokens: Optional[int] = None,
        max_tokens: Optional[int] = None,
        presence_penalty: Optional[float] = None,
        frequency_penalty: Optional[float] = None,
        logit_bias: Optional[Dict[int, float]] = None,
        response_format: Optional[Type[BaseModel]] = None,
        seed: Optional[int] = None,
        logprobs: Optional[int] = None,
        top_logprobs: Optional[int] = None,
        base_url: Optional[str] = None,
        api_base: Optional[str] = None,
        api_version: Optional[str] = None,
        api_key: Optional[str] = None,
        callbacks: List[Any] = [],
        reasoning_effort: Optional[Literal["none", "low", "medium", "high"]] = None,
        **kwargs, # <-- here is the change.
    ):
    ```

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants