Source code for langchain.llms.fireworks
"""Wrapper around Fireworks APIs"""
import json
import logging
from typing import (
    Any,
    Dict,
    List,
    Optional,
    Set,
    Tuple,
    Union,
)
import requests
from pydantic import Field, root_validator
from langchain.callbacks.manager import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain.llms.base import BaseLLM
from langchain.schema import Generation, LLMResult
from langchain.utils import get_from_dict_or_env
logger = logging.getLogger(__name__)
[docs]class BaseFireworks(BaseLLM):
    """Wrapper around Fireworks large language models."""
    model_id: str = Field(
        "accounts/fireworks/models/fireworks-llama-v2-7b-chat", alias="model"
    )
    """Model name to use."""
    temperature: float = 0.7
    """What sampling temperature to use."""
    max_tokens: int = 512
    """The maximum number of tokens to generate in the completion.
    -1 returns as many tokens as possible given the prompt and
    the models maximal context size."""
    top_p: float = 1
    """Total probability mass of tokens to consider at each step."""
    fireworks_api_key: Optional[str] = None
    """Api key to use fireworks API"""
    batch_size: int = 20
    """Batch size to use when passing multiple documents to generate."""
    request_timeout: Optional[Union[float, Tuple[float, float]]] = None
    """Timeout for requests to Fireworks completion API. Default is 600 seconds."""
    max_retries: int = 6
    """Maximum number of retries to make when generating."""
    @property