Toggle Menu
Prev Up Next

# Source code for langchain_community.llms.fireworks

    import asyncio
    from concurrent.futures import ThreadPoolExecutor
    from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Union

    from langchain_core._api.deprecation import deprecated
    from langchain_core.callbacks import (
        AsyncCallbackManagerForLLMRun,
        CallbackManagerForLLMRun,
    )
    from langchain_core.language_models.llms import BaseLLM, create_base_retry_decorator
    from langchain_core.outputs import Generation, GenerationChunk, LLMResult
    from langchain_core.pydantic_v1 import Field, SecretStr, root_validator
    from langchain_core.utils import convert_to_secret_str
    from langchain_core.utils.env import get_from_dict_or_env

    def _stream_response_to_generation_chunk(
        stream_response: Any,
    ) -> GenerationChunk:
        """Convert a stream response to a generation chunk."""
        return GenerationChunk(
            text=stream_response.choices[0].text,
            generation_info=dict(
                finish_reason=stream_response.choices[0].finish_reason,
                logprobs=stream_response.choices[0].logprobs,
            ),
        )

    @deprecated(
        since="0.0.26",
        removal="0.3",
        alternative_import="langchain_fireworks.Fireworks",
    )
    class Fireworks(BaseLLM):
        """Fireworks models."""

        model: str = "accounts/fireworks/models/llama-v2-7b-chat"
        model_kwargs: dict = Field(
            default_factory=lambda: {
                "temperature": 0.7,
                "max_tokens": 512,
                "top_p": 1,
            }.copy()
        )
        fireworks_api_key: Optional[SecretStr] = None
        max_retries: int = 20
        batch_size: int = 20
        use_retry: bool = True

        @property
        def lc_secrets(self) -> Dict[str, str]:
            return {"fireworks_api_key": "FIREWORKS_API_KEY"}

        @classmethod
        def is_lc_serializable(cls) -> bool:
            return True

        @classmethod
        def get_lc_namespace(cls) -> List[str]:
            """Get the namespace of the langchain object."""
            return ["langchain", "llms", "fireworks"]

        @root_validator()
        def validate_environment(cls, values: Dict) -> Dict:
            """Validate that api key in environment."""
            try:
                import fireworks.client
            except ImportError as e:
                raise ImportError(
                    "Could not import fireworks-ai python package. "
                    "Please install it with `pip install fireworks-ai`."
                ) from e
            fireworks_api_key = convert_to_secret_str(
                get_from_dict_or_env(values, "fireworks_api_key", "FIREWORKS_API_KEY")
            )
            fireworks.client.api_key = fireworks_api_key.get_secret_value()
            return values

        @property
        def _llm_type(self) -> str:
            """Return type of llm."""
            return "fireworks"

        def _generate(
            self,
            prompts: List[str],
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> LLMResult:
            """Call out to Fireworks endpoint with k unique prompts.
            Args:
                prompts: The prompts to pass into the model.
                stop: Optional list of stop words to use when generating.
            Returns:
                The full LLM output.
            """
            params = {
                "model": self.model,
                **self.model_kwargs,
            }
            sub_prompts = self.get_batch_prompts(prompts)
            choices = []
            for _prompts in sub_prompts:
                response = completion_with_retry_batching(
                    self,
                    self.use_retry,
                    prompt=_prompts,
                    run_manager=run_manager,
                    stop=stop,
                    **params,
                )
                choices.extend(response)

            return self.create_llm_result(choices, prompts)

        async def _agenerate(
            self,
            prompts: List[str],
            stop: Optional[List[str]] = None,
            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> LLMResult:
            """Call out to Fireworks endpoint async with k unique prompts."""
            params = {
                "model": self.model,
                **self.model_kwargs,
            }
            sub_prompts = self.get_batch_prompts(prompts)
            choices = []
            for _prompts in sub_prompts:
                response = await acompletion_with_retry_batching(
                    self,
                    self.use_retry,
                    prompt=_prompts,
                    run_manager=run_manager,
                    stop=stop,
                    **params,
                )
                choices.extend(response)

            return self.create_llm_result(choices, prompts)

        def get_batch_prompts(
            self,
            prompts: List[str],
        ) -> List[List[str]]:
            """Get the sub prompts for llm call."""
            sub_prompts = [
                prompts[i : i + self.batch_size]
                for i in range(0, len(prompts), self.batch_size)
            ]
            return sub_prompts

        def create_llm_result(self, choices: Any, prompts: List[str]) -> LLMResult:
            """Create the LLMResult from the choices and prompts."""
            generations = []
            for i, _ in enumerate(prompts):
                sub_choices = choices[i : (i + 1)]
                generations.append(
                    [
                        Generation(
                            text=choice.__dict__["choices"][0].text,
                        )
                        for choice in sub_choices
                    ]
                )
            llm_output = {"model": self.model}
            return LLMResult(generations=generations, llm_output=llm_output)

        def _stream(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> Iterator[GenerationChunk]:
            params = {
                "model": self.model,
                "prompt": prompt,
                "stream": True,
                **self.model_kwargs,
            }
            for stream_resp in completion_with_retry(
                self, self.use_retry, run_manager=run_manager, stop=stop, **params
            ):
                chunk = _stream_response_to_generation_chunk(stream_resp)
                if run_manager:
                    run_manager.on_llm_new_token(chunk.text, chunk=chunk)
                yield chunk

        async def _astream(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> AsyncIterator[GenerationChunk]:
            params = {
                "model": self.model,
                "prompt": prompt,
                "stream": True,
                **self.model_kwargs,
            }
            async for stream_resp in await acompletion_with_retry_streaming(
                self, self.use_retry, run_manager=run_manager, stop=stop, **params
            ):
                chunk = _stream_response_to_generation_chunk(stream_resp)
                if run_manager:
                    await run_manager.on_llm_new_token(chunk.text, chunk=chunk)
                yield chunk

    def conditional_decorator(
        condition: bool, decorator: Callable[[Any], Any]
    ) -> Callable[[Any], Any]:
        """Conditionally apply a decorator.

        Args:
            condition: A boolean indicating whether to apply the decorator.
            decorator: A decorator function.

        Returns:
            A decorator function.
        """

        def actual_decorator(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
            if condition:
                return decorator(func)
            return func

        return actual_decorator

    def completion_with_retry(
        llm: Fireworks,
        use_retry: bool,
        *,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Any:
        """Use tenacity to retry the completion call."""
        import fireworks.client

        retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)

        @conditional_decorator(use_retry, retry_decorator)
        def _completion_with_retry(**kwargs: Any) -> Any:
            return fireworks.client.Completion.create(
                **kwargs,
            )

        return _completion_with_retry(**kwargs)

    async def acompletion_with_retry(
        llm: Fireworks,
        use_retry: bool,
        *,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Any:
        """Use tenacity to retry the completion call."""
        import fireworks.client

        retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)

        @conditional_decorator(use_retry, retry_decorator)
        async def _completion_with_retry(**kwargs: Any) -> Any:
            return await fireworks.client.Completion.acreate(
                **kwargs,
            )

        return await _completion_with_retry(**kwargs)

    def completion_with_retry_batching(
        llm: Fireworks,
        use_retry: bool,
        *,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Any:
        """Use tenacity to retry the completion call."""
        import fireworks.client

        prompt = kwargs["prompt"]
        del kwargs["prompt"]

        retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)

        @conditional_decorator(use_retry, retry_decorator)
        def _completion_with_retry(prompt: str) -> Any:
            return fireworks.client.Completion.create(**kwargs, prompt=prompt)

        def batch_sync_run() -> List:
            with ThreadPoolExecutor() as executor:
                results = list(executor.map(_completion_with_retry, prompt))
            return results

        return batch_sync_run()

    async def acompletion_with_retry_batching(
        llm: Fireworks,
        use_retry: bool,
        *,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Any:
        """Use tenacity to retry the completion call."""
        import fireworks.client

        prompt = kwargs["prompt"]
        del kwargs["prompt"]

        retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)

        @conditional_decorator(use_retry, retry_decorator)
        async def _completion_with_retry(prompt: str) -> Any:
            return await fireworks.client.Completion.acreate(**kwargs, prompt=prompt)

        def run_coroutine_in_new_loop(
            coroutine_func: Any, *args: Dict, **kwargs: Dict
        ) -> Any:
            new_loop = asyncio.new_event_loop()
            try:
                asyncio.set_event_loop(new_loop)
                return new_loop.run_until_complete(coroutine_func(*args, **kwargs))
            finally:
                new_loop.close()

        async def batch_sync_run() -> List:
            with ThreadPoolExecutor() as executor:
                results = list(
                    executor.map(
                        run_coroutine_in_new_loop,
                        [_completion_with_retry] * len(prompt),
                        prompt,
                    )
                )
            return results

        return await batch_sync_run()

    async def acompletion_with_retry_streaming(
        llm: Fireworks,
        use_retry: bool,
        *,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Any:
        """Use tenacity to retry the completion call for streaming."""
        import fireworks.client

        retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)

        @conditional_decorator(use_retry, retry_decorator)
        async def _completion_with_retry(**kwargs: Any) -> Any:
            return fireworks.client.Completion.acreate(
                **kwargs,
            )

        return await _completion_with_retry(**kwargs)

    def _create_retry_decorator(
        llm: Fireworks,
        *,
        run_manager: Optional[
            Union[AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun]
        ] = None,
    ) -> Callable[[Any], Any]:
        """Define retry mechanism."""
        import fireworks.client

        errors = [
            fireworks.client.error.RateLimitError,
            fireworks.client.error.InternalServerError,
            fireworks.client.error.BadGatewayError,
            fireworks.client.error.ServiceUnavailableError,
        ]
        return create_base_retry_decorator(
            error_types=errors, max_retries=llm.max_retries, run_manager=run_manager
        )