LiteLLM Model

rankify.generator.models.litellm_model

BaseRAGModel

Bases: ABC

Base RAG Model for Retrieval-Augmented Generation (RAG).

This is an abstract base class for implementing LLM endpoints in rankify. It defines the interface for generating responses and optional embedding generation.

Methods:

Name	Description
generate	str, **kwargs) -> str: Abstract method to generate a response based on the given prompt.
embed	str, **kwargs) -> List[float]: Optional method to generate embeddings for the given text.

Notes

• This class serves as a blueprint for RAG models like OpenAIModel and HuggingFaceModel.
• The embed method is optional and can be implemented if needed.
• This class needs to be extended to include new LLM endpoints in Rankify.

Source code in rankify/generator/models/base_rag_model.py

class BaseRAGModel(ABC):
    """
    **Base RAG Model** for Retrieval-Augmented Generation (RAG).

    This is an abstract base class for implementing LLM endpoints in rankify. 
    It defines the interface for generating responses and optional embedding generation.

    Methods:
        generate(prompt: str, **kwargs) -> str:
            Abstract method to generate a response based on the given prompt.
        embed(text: str, **kwargs) -> List[float]:
            Optional method to generate embeddings for the given text.

    Notes:
        - This class serves as a blueprint for RAG models like `OpenAIModel` and `HuggingFaceModel`.
        - The `embed` method is optional and can be implemented if needed.
        - This class needs to be extended to include new LLM endpoints in Rankify.
    """

    @abstractmethod
    def generate(self, prompt: str, **kwargs) -> str:
        """Generate a response based on the given prompt."""
        pass

    def embed(self, text: str, **kwargs) -> List[float]:
        """Optional: Generate embeddings for the given text."""
        raise NotImplementedError("Embedding is not required for this implementation.")

generate(prompt, **kwargs) abstractmethod

Generate a response based on the given prompt.

Source code in rankify/generator/models/base_rag_model.py

@abstractmethod
def generate(self, prompt: str, **kwargs) -> str:
    """Generate a response based on the given prompt."""
    pass

embed(text, **kwargs)

Optional: Generate embeddings for the given text.

Source code in rankify/generator/models/base_rag_model.py

def embed(self, text: str, **kwargs) -> List[float]:
    """Optional: Generate embeddings for the given text."""
    raise NotImplementedError("Embedding is not required for this implementation.")

PromptGenerator

PromptGenerator for Retrieval-Augmented Generation (RAG).

This class manages prompt construction for different RAG methods and model types in Rankify. It selects and formats prompt templates based on the specified method and model, enabling flexible and consistent prompt generation.

Attributes:

Name	Type	Description
method	str	The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").
prompt_template	PromptTemplate	The prompt template used for formatting prompts.
model_type	str	The type of model (e.g., "huggingface", "openai").

Methods:

Name	Description
generate_user_prompt	Generates a user prompt by formatting the question and contexts with the selected template.
_select_template	Selects the appropriate prompt template for the given method.

Notes

• Supports custom prompt templates via the custom_prompt argument.
• If no contexts are provided, generates prompts with only the question.
• Ensures consistent prompt formatting across different RAG methods and models.
• Automatically selects a default template if none is specified.

Source code in rankify/generator/prompt_generator.py

class PromptGenerator:
    """
    **PromptGenerator** for Retrieval-Augmented Generation (RAG).

    This class manages prompt construction for different RAG methods and model types in Rankify.
    It selects and formats prompt templates based on the specified method and model, enabling flexible and consistent prompt generation.

    Attributes:
        method (str): The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").
        prompt_template (PromptTemplate): The prompt template used for formatting prompts.
        model_type (str): The type of model (e.g., "huggingface", "openai").

    Methods:
        generate_user_prompt(question, contexts, custom_prompt=None) -> str:
            Generates a user prompt by formatting the question and contexts with the selected template.
        _select_template(method) -> PromptTemplate:
            Selects the appropriate prompt template for the given method.

    Notes:
        - Supports custom prompt templates via the `custom_prompt` argument.
        - If no contexts are provided, generates prompts with only the question.
        - Ensures consistent prompt formatting across different RAG methods and models.
        - Automatically selects a default template if none is specified.
    """
    def __init__(self, method: str, model_type: str, prompt_template: Optional[PromptTemplate] = None):
        """
        Initialize the PromptGenerator.

        Args:
            method (str): The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").
            model_type (str): The type of model (e.g., "huggingface", "openai").
            prompt_template (PromptTemplate, optional): Custom prompt template to use. If None, selects based on method.

        Notes:
            - If no custom template is provided, selects a default template for the specified method.
            - Stores the method, model type, and selected prompt template for prompt generation.
        """
        self.method = method
        if prompt_template is not None:
            self.prompt_template = prompt_template
        else:
            self.prompt_template = self._select_template(method)

    def _select_template(self, method: str) -> PromptTemplate:
        """
        Selects the appropriate prompt template for the given RAG method.

        Args:
            method (str): The RAG method or strategy.

        Returns:
            PromptTemplate: The selected prompt template.

        Notes:
            - If the method is not recognized, defaults to BASIC_RAG template.
        """
        method = method.lower()
        if method in PromptTemplate._value2member_map_:
            return PromptTemplate(method)
        return PromptTemplate.BASIC_RAG

    def generate_user_prompt(self, question: str, contexts: List[str], custom_prompt: Optional[str] = None) -> str:
        """
        Generates a user prompt by formatting the question and contexts with the selected template.

        Args:
            question (str): The question to be answered.
            contexts (List[str]): List of context passages to include in the prompt.
            custom_prompt (str, optional): Custom prompt template string. If provided, overrides the default template.

        Returns:
            str: The formatted prompt string.

        Notes:
            - If no contexts are provided, generates a prompt with only the question.
            - Custom prompt templates can use `{question}` and `{contexts}` placeholders.
            - Ensures consistent prompt formatting for all RAG methods and models.
        """
        if contexts is None:
            contexts = []
        context_str = "\n".join(contexts)
        if custom_prompt:
            return custom_prompt.format(question=question, contexts=context_str)
        template = DEFAULT_PROMPTS[self.prompt_template]
        return template.format(question=question, contexts=context_str)

__init__(method, model_type, prompt_template=None)

Initialize the PromptGenerator.

Parameters:

Name	Type	Description	Default
method	str	The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").	required
model_type	str	The type of model (e.g., "huggingface", "openai").	required
prompt_template	PromptTemplate	Custom prompt template to use. If None, selects based on method.	None

Notes

• If no custom template is provided, selects a default template for the specified method.
• Stores the method, model type, and selected prompt template for prompt generation.

Source code in rankify/generator/prompt_generator.py

def __init__(self, method: str, model_type: str, prompt_template: Optional[PromptTemplate] = None):
    """
    Initialize the PromptGenerator.

    Args:
        method (str): The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").
        model_type (str): The type of model (e.g., "huggingface", "openai").
        prompt_template (PromptTemplate, optional): Custom prompt template to use. If None, selects based on method.

    Notes:
        - If no custom template is provided, selects a default template for the specified method.
        - Stores the method, model type, and selected prompt template for prompt generation.
    """
    self.method = method
    if prompt_template is not None:
        self.prompt_template = prompt_template
    else:
        self.prompt_template = self._select_template(method)

generate_user_prompt(question, contexts, custom_prompt=None)

Generates a user prompt by formatting the question and contexts with the selected template.

Parameters:

Name	Type	Description	Default
question	str	The question to be answered.	required
contexts	List[str]	List of context passages to include in the prompt.	required
custom_prompt	str	Custom prompt template string. If provided, overrides the default template.	None

Returns:

Name	Type	Description
str	str	The formatted prompt string.

Notes

• If no contexts are provided, generates a prompt with only the question.
• Custom prompt templates can use {question} and {contexts} placeholders.
• Ensures consistent prompt formatting for all RAG methods and models.

Source code in rankify/generator/prompt_generator.py

def generate_user_prompt(self, question: str, contexts: List[str], custom_prompt: Optional[str] = None) -> str:
    """
    Generates a user prompt by formatting the question and contexts with the selected template.

    Args:
        question (str): The question to be answered.
        contexts (List[str]): List of context passages to include in the prompt.
        custom_prompt (str, optional): Custom prompt template string. If provided, overrides the default template.

    Returns:
        str: The formatted prompt string.

    Notes:
        - If no contexts are provided, generates a prompt with only the question.
        - Custom prompt templates can use `{question}` and `{contexts}` placeholders.
        - Ensures consistent prompt formatting for all RAG methods and models.
    """
    if contexts is None:
        contexts = []
    context_str = "\n".join(contexts)
    if custom_prompt:
        return custom_prompt.format(question=question, contexts=context_str)
    template = DEFAULT_PROMPTS[self.prompt_template]
    return template.format(question=question, contexts=context_str)

LitellmClient

Source code in rankify/utils/api/litellmclient.py

class LitellmClient:
    #  https://github.com/BerriAI/litellm
    def __init__(self, keys=None):
        self.api_key = keys

    def chat(self, return_text=True, *args, **kwargs):
        response = completion(api_key=self.api_key, *args, **kwargs)
        if return_text:
            response = response.choices[0].message.content
        return response

LitellmModel

Bases: BaseRAGModel

LiteLLM Model for Retrieval-Augmented Generation (RAG).

This class integrates LiteLLM's API for text generation in a RAG pipeline. It uses the LiteLLM API to generate responses based on input prompts.

Attributes:

Name	Type	Description
model_name	str	Name of the LiteLLM model.
prompt_generator	PromptGenerator	Instance for generating prompts.
client	LitellmClient	Client for interacting with the LiteLLM API.

Notes

• This model uses LiteLLM's API for text generation.
• It supports additional parameters like max_tokens and temperature.

Source code in rankify/generator/models/litellm_model.py

class LitellmModel(BaseRAGModel):
    """
    **LiteLLM Model** for Retrieval-Augmented Generation (RAG).

    This class integrates LiteLLM's API for text generation in a RAG pipeline. 
    It uses the LiteLLM API to generate responses based on input prompts.

    Attributes:
        model_name (str): Name of the LiteLLM model.
        prompt_generator (PromptGenerator): Instance for generating prompts.
        client (LitellmClient): Client for interacting with the LiteLLM API.

    Notes:
        - This model uses LiteLLM's API for text generation.
        - It supports additional parameters like `max_tokens` and `temperature`.
    """
    def __init__(self, model_name: str, api_key: str, prompt_generator: PromptGenerator):
        """
        Initialize the LitellmModel with the LitellmClient.

        :param model_name: Name of the LiteLLM model.
        :param api_key: API key for LiteLLM.
        :param prompt_generator: Instance of PromptGenerator for generating prompts.
        """
        self.model_name = model_name
        self.prompt_generator = prompt_generator

        self.client = LitellmClient(keys=api_key)

    def generate(self, prompt: str, **kwargs) -> str:
        """
        Generate a response using LiteLLM's API.

        Args:
            prompt (str): The input prompt for the model.
            **kwargs: Additional parameters for the LiteLLM API call, such as:
                - model (str): Model name to use (default: self.model_name).
                - max_tokens (int): Maximum number of tokens to generate (default: 128).
                - temperature (float): Sampling temperature (default: 0.7).

        Returns:
            str: The generated response as a string.

        Notes:
            - Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
            - All generation parameters can be overridden via `kwargs`.

        Example:
            ```python
            answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)
            ```
        """

        # Set default parameters for the LiteLLM API call
        kwargs.setdefault("model", self.model_name)
        kwargs.setdefault("max_tokens", 128)
        kwargs.setdefault("temperature", 0.7)

        # Call the LiteLLM API using the LitellmClient
        response = self.client.chat(messages=[{"role": "user", "content": prompt}], return_text=True, **kwargs)
        return response

__init__(model_name, api_key, prompt_generator)

Initialize the LitellmModel with the LitellmClient.

:param model_name: Name of the LiteLLM model. :param api_key: API key for LiteLLM. :param prompt_generator: Instance of PromptGenerator for generating prompts.

Source code in rankify/generator/models/litellm_model.py

def __init__(self, model_name: str, api_key: str, prompt_generator: PromptGenerator):
    """
    Initialize the LitellmModel with the LitellmClient.

    :param model_name: Name of the LiteLLM model.
    :param api_key: API key for LiteLLM.
    :param prompt_generator: Instance of PromptGenerator for generating prompts.
    """
    self.model_name = model_name
    self.prompt_generator = prompt_generator

    self.client = LitellmClient(keys=api_key)

generate(prompt, **kwargs)

Generate a response using LiteLLM's API.

Parameters:

Name	Type	Description	Default
prompt	str	The input prompt for the model.	required
**kwargs		Additional parameters for the LiteLLM API call, such as: - model (str): Model name to use (default: self.model_name). - max_tokens (int): Maximum number of tokens to generate (default: 128). - temperature (float): Sampling temperature (default: 0.7).	{}

Returns:

Name	Type	Description
str	str	The generated response as a string.

Notes

• Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
• All generation parameters can be overridden via kwargs.

Example

answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)

Source code in rankify/generator/models/litellm_model.py

def generate(self, prompt: str, **kwargs) -> str:
    """
    Generate a response using LiteLLM's API.

    Args:
        prompt (str): The input prompt for the model.
        **kwargs: Additional parameters for the LiteLLM API call, such as:
            - model (str): Model name to use (default: self.model_name).
            - max_tokens (int): Maximum number of tokens to generate (default: 128).
            - temperature (float): Sampling temperature (default: 0.7).

    Returns:
        str: The generated response as a string.

    Notes:
        - Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
        - All generation parameters can be overridden via `kwargs`.

    Example:
        ```python
        answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)
        ```
    """

    # Set default parameters for the LiteLLM API call
    kwargs.setdefault("model", self.model_name)
    kwargs.setdefault("max_tokens", 128)
    kwargs.setdefault("temperature", 0.7)

    # Call the LiteLLM API using the LitellmClient
    response = self.client.chat(messages=[{"role": "user", "content": prompt}], return_text=True, **kwargs)
    return response