fireworks

FireworksGen class for text generation using Fireworks API.

Classes:

• FireworksGen –

  Fireworks AI generation model for text generation.

FireworksGen

FireworksGen(
    model_name: Optional[str] = None,
    temperature: Optional[float] = None,
    prompt_template: str = '',
    output_max_length: int = 500,
    device: str = 'auto',
    structured_output: Optional[Type[BaseModel]] = None,
    system_message: str = '',
    api_params: dict[str, Any] = DEFAULT_API_PARAMS,
    api_key: str = '',
    cache: Cache | None = None,
    logs: dict[str, Any] | None = None,
)

Bases: GenerationBase

Fireworks AI generation model for text generation.

Methods:

• apply –

  Apply attached configuration to the step.

• generate –

  Generate text using Fireworks AI's API.

• process –

  Generate a result from the current pipeline content.

Source code in src/rago/generation/base.py

def __init__(
    self,
    model_name: Optional[str] = None,
    temperature: Optional[float] = None,
    prompt_template: str = '',
    output_max_length: int = 500,
    device: str = 'auto',
    structured_output: Optional[Type[BaseModel]] = None,
    system_message: str = '',
    api_params: dict[str, Any] = DEFAULT_API_PARAMS,
    api_key: str = '',
    cache: Cache | None = None,
    logs: dict[str, Any] | None = None,
) -> None:
    super().__init__()
    self.api_key = api_key
    self.cache = cache
    self.logs = logs if logs is not None else {}

    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.output_max_length = (
        output_max_length or self.default_output_max_length
    )
    self.temperature = (
        temperature
        if temperature is not None
        else self.default_temperature
    )
    self.prompt_template = prompt_template or self.default_prompt_template
    self.structured_output = structured_output
    if api_params is DEFAULT_API_PARAMS:
        api_params = deepcopy(self.default_api_params or {})

    self.system_message = system_message
    self.api_params = api_params

    if device not in ['cpu', 'cuda', 'auto']:
        raise Exception(
            f'Device {device} not supported. Options: cpu, cuda, auto.'
        )

    cuda_available = torch.cuda.is_available()
    self.device_name = (
        'cpu' if device == 'cpu' or not cuda_available else 'cuda'
    )
    self.device = torch.device(self.device_name)

    self._validate()
    self._load_optional_modules()
    self._setup()

apply

apply(parameters: Any) -> None

Apply attached configuration to the step.

Source code in src/rago/base.py

def apply(self, parameters: Any) -> None:
    """Apply attached configuration to the step."""
    if parameters is None:
        return

    if _is_cache_backend(parameters):
        self.cache = parameters
        return

    if _is_vector_db(parameters):
        setattr(self, 'db', parameters)
        return

    if _is_text_splitter(parameters):
        setattr(self, 'splitter', parameters)
        return

    for key, value in config_to_dict(parameters).items():
        if key == 'cache':
            self.cache = value
        elif key == 'logs':
            self.logs = value if value is not None else {}
        else:
            setattr(self, key, value)

generate

generate(query: str, data: list[str]) -> str | BaseModel

Generate text using Fireworks AI's API.

Source code in src/rago/generation/fireworks.py

def generate(self, query: str, data: list[str]) -> str | BaseModel:
    """Generate text using Fireworks AI's API."""
    input_text = self._format_prompt(query, data)

    api_params = self.api_params or self.default_api_params

    messages = []
    if self.system_message:
        messages.append({'role': 'system', 'content': self.system_message})
    messages.append({'role': 'user', 'content': input_text})

    model_params = {
        'model': self.model_name,
        'messages': messages,
        'max_tokens': self.output_max_length,
        'temperature': self.temperature,
        **api_params,
    }

    if self.structured_output:
        model_params['response_model'] = self.structured_output
        response = self.model.chat.completions.create(**model_params)
        # self.logs['model_params'] = model_params
        return cast(BaseModel, response)

    response = self.model.chat.completions.create(**model_params)
    # self.logs['model_params'] = model_params
    return cast(str, response.choices[0].message.content.strip())

process

process(inp: Input) -> Output

Generate a result from the current pipeline content.

Source code in src/rago/generation/base.py

def process(self, inp: Input) -> Output:
    """Generate a result from the current pipeline content."""
    query = str(inp.query)
    data = [
        str(item)
        for item in ensure_list(
            inp.get('content', inp.get('data', inp.get('source')))
        )
    ]
    result = self.generate(query, data)
    output = Output.from_input(inp)
    output.result = result
    output.content = _serialize_generation_result(result)
    output.data = output.content
    return output