Skip to content

OpenAI Model

rankify.generator.models.openai_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

OpenaiClient

Source code in rankify/utils/api/openaiclient.py
class OpenaiClient:
    def __init__(self, keys=None, base_url=None, start_id=None, proxy=None, api_version=None):

        if isinstance(keys, str):
            keys = [keys]
        if keys is None:
            raise "Please provide OpenAI Key."

        self.key = keys
        self.base_url = base_url
        self.key_id = start_id or 0
        self.key_id = self.key_id % len(self.key)
        self.api_key = self.key[self.key_id % len(self.key)]
        print(self.base_url)

        # Use AzureOpenAI when base_url points to an Azure endpoint
        if base_url and ".openai.azure.com" in base_url:
            try:
                from openai import AzureOpenAI
                import os
                # Extract api_version from env or passed param
                _api_version = (api_version
                                or os.getenv("RANKIFY_AZURE_API_VERSION", "2025-01-01-preview"))
                # Extract endpoint (everything before /openai/...)
                _endpoint = base_url.split("/openai/")[0] if "/openai/" in base_url else base_url
                self.client = AzureOpenAI(
                    azure_endpoint=_endpoint,
                    api_key=self.api_key,
                    api_version=_api_version,
                )
                # Store deployment name from the base_url path for chat calls
                self._azure_deployment = base_url.split("/deployments/")[-1].split("/")[0] if "/deployments/" in base_url else None
            except ImportError:
                self.client = OpenAI(api_key=self.api_key, base_url=base_url)
                self._azure_deployment = None
        else:
            self.client = OpenAI(api_key=self.api_key, base_url=base_url)
            self._azure_deployment = None


    def chat(self, *args, return_text=False, reduce_length=False, **kwargs):
        while True:
            try:
                completion = self.client.chat.completions.create(*args, **kwargs, timeout=30)
                break
            except Exception as e:
                print(str(e))
                if "This model's maximum context length is" in str(e):
                    print('reduce_length')
                    return 'ERROR::reduce_length'
                time.sleep(0.1)
        if return_text:
            completion = completion.choices[0].message.content
        return completion

    def text(self, *args, return_text=False, reduce_length=False, **kwargs):
        while True:
            try:
                completion = self.client.completions.create(
                    *args, **kwargs
                )
                break
            except Exception as e:
                print(e)
                if "This model's maximum context length is" in str(e):
                    print('reduce_length')
                    return 'ERROR::reduce_length'
                time.sleep(0.1)
        if return_text:
            completion = completion.choices[0].text
        return completion

OpenAIModel

Bases: BaseRAGModel

OpenAI Model for Retrieval-Augmented Generation (RAG).

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

Attributes:

Name Type Description
model_name str

Name of the OpenAI model (e.g., "gpt-3.5-turbo").

prompt_generator PromptGenerator

Instance for generating prompts.

client OpenaiClient

Client for interacting with the OpenAI API.

Notes
  • This model uses OpenAI's GPT models for text generation.
  • It supports additional parameters like max_tokens and temperature.
Source code in rankify/generator/models/openai_model.py
class OpenAIModel(BaseRAGModel):
    """
    **OpenAI Model** for Retrieval-Augmented Generation (RAG).

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

    Attributes:
        model_name (str): Name of the OpenAI model (e.g., "gpt-3.5-turbo").
        prompt_generator (PromptGenerator): Instance for generating prompts.
        client (OpenaiClient): Client for interacting with the OpenAI API.

    Notes:
        - This model uses OpenAI's GPT models for text generation.
        - It supports additional parameters like `max_tokens` and `temperature`.
    """
    def __init__(self, model_name: str, api_keys: list, prompt_generator: PromptGenerator, base_url: str = None):
        """
        Initialize the OpenAIModel with the OpenaiClient.

        :param model_name: Name of the OpenAI model (e.g., "gpt-3.5-turbo").
        :param api_keys: List of API keys for OpenAI.
        :param prompt_generator: Instance of PromptGenerator for generating prompts.
        :param base_url: Optional base URL for the OpenAI API.
        """
        self.model_name = model_name
        self.prompt_generator = prompt_generator

        self.client = OpenaiClient(keys=api_keys, base_url=base_url)

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

        Args:
            prompt (str): The input prompt for the model.
            **kwargs: Additional parameters for the OpenAI 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 OpenAI API call
        kwargs.setdefault("model", self.model_name)
        kwargs.setdefault("max_tokens", 128)
        kwargs.setdefault("temperature", 0.7)

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

__init__(model_name, api_keys, prompt_generator, base_url=None)

Initialize the OpenAIModel with the OpenaiClient.

:param model_name: Name of the OpenAI model (e.g., "gpt-3.5-turbo"). :param api_keys: List of API keys for OpenAI. :param prompt_generator: Instance of PromptGenerator for generating prompts. :param base_url: Optional base URL for the OpenAI API.

Source code in rankify/generator/models/openai_model.py
def __init__(self, model_name: str, api_keys: list, prompt_generator: PromptGenerator, base_url: str = None):
    """
    Initialize the OpenAIModel with the OpenaiClient.

    :param model_name: Name of the OpenAI model (e.g., "gpt-3.5-turbo").
    :param api_keys: List of API keys for OpenAI.
    :param prompt_generator: Instance of PromptGenerator for generating prompts.
    :param base_url: Optional base URL for the OpenAI API.
    """
    self.model_name = model_name
    self.prompt_generator = prompt_generator

    self.client = OpenaiClient(keys=api_keys, base_url=base_url)

generate(prompt, **kwargs)

Generate a response using OpenAI's API.

Parameters:

Name Type Description Default
prompt str

The input prompt for the model.

required
**kwargs

Additional parameters for the OpenAI 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/openai_model.py
def generate(self, prompt: str, **kwargs) -> str:
    """
    Generate a response using OpenAI's API.

    Args:
        prompt (str): The input prompt for the model.
        **kwargs: Additional parameters for the OpenAI 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 OpenAI API call
    kwargs.setdefault("model", self.model_name)
    kwargs.setdefault("max_tokens", 128)
    kwargs.setdefault("temperature", 0.7)

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