|
|
|
@ -1,8 +1,13 @@
|
|
|
|
|
from typing import Any, Callable, Dict, Iterator, List, Mapping, Optional, Tuple, Union
|
|
|
|
|
|
|
|
|
|
import fireworks
|
|
|
|
|
import fireworks.client
|
|
|
|
|
from pydantic import root_validator
|
|
|
|
|
from typing import (
|
|
|
|
|
Any,
|
|
|
|
|
AsyncIterator,
|
|
|
|
|
Callable,
|
|
|
|
|
Dict,
|
|
|
|
|
Iterator,
|
|
|
|
|
List,
|
|
|
|
|
Optional,
|
|
|
|
|
Union,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
from langchain.adapters.openai import convert_message_to_dict
|
|
|
|
|
from langchain.callbacks.manager import (
|
|
|
|
@ -11,6 +16,7 @@ from langchain.callbacks.manager import (
|
|
|
|
|
)
|
|
|
|
|
from langchain.chat_models.base import BaseChatModel
|
|
|
|
|
from langchain.llms.base import create_base_retry_decorator
|
|
|
|
|
from langchain.pydantic_v1 import Field, root_validator
|
|
|
|
|
from langchain.schema.messages import (
|
|
|
|
|
AIMessage,
|
|
|
|
|
AIMessageChunk,
|
|
|
|
@ -30,12 +36,12 @@ from langchain.utils.env import get_from_dict_or_env
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _convert_delta_to_message_chunk(
|
|
|
|
|
_dict: Mapping[str, Any], default_class: type[BaseMessageChunk]
|
|
|
|
|
_dict: Any, default_class: type[BaseMessageChunk]
|
|
|
|
|
) -> BaseMessageChunk:
|
|
|
|
|
"""Convert a delta response to a message chunk."""
|
|
|
|
|
role = _dict.role
|
|
|
|
|
content = _dict.content or ""
|
|
|
|
|
additional_kwargs = {}
|
|
|
|
|
additional_kwargs: Dict = {}
|
|
|
|
|
|
|
|
|
|
if role == "user" or default_class == HumanMessageChunk:
|
|
|
|
|
return HumanMessageChunk(content=content)
|
|
|
|
@ -51,7 +57,7 @@ def _convert_delta_to_message_chunk(
|
|
|
|
|
return default_class(content=content)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:
|
|
|
|
|
def convert_dict_to_message(_dict: Any) -> BaseMessage:
|
|
|
|
|
"""Convert a dict response to a message."""
|
|
|
|
|
role = _dict.role
|
|
|
|
|
content = _dict.content or ""
|
|
|
|
@ -59,7 +65,7 @@ def convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:
|
|
|
|
|
return HumanMessage(content=content)
|
|
|
|
|
elif role == "assistant":
|
|
|
|
|
content = _dict.content
|
|
|
|
|
additional_kwargs = {}
|
|
|
|
|
additional_kwargs: Dict = {}
|
|
|
|
|
return AIMessage(content=content, additional_kwargs=additional_kwargs)
|
|
|
|
|
elif role == "system":
|
|
|
|
|
return SystemMessage(content=content)
|
|
|
|
@ -73,13 +79,23 @@ class ChatFireworks(BaseChatModel):
|
|
|
|
|
"""Fireworks Chat models."""
|
|
|
|
|
|
|
|
|
|
model: str = "accounts/fireworks/models/llama-v2-7b-chat"
|
|
|
|
|
model_kwargs: Optional[dict] = {"temperature": 0.7, "max_tokens": 512, "top_p": 1}
|
|
|
|
|
model_kwargs: dict = Field(
|
|
|
|
|
default_factory=lambda: {
|
|
|
|
|
"temperature": 0.7,
|
|
|
|
|
"max_tokens": 512,
|
|
|
|
|
"top_p": 1,
|
|
|
|
|
}.copy()
|
|
|
|
|
)
|
|
|
|
|
fireworks_api_key: Optional[str] = None
|
|
|
|
|
max_retries: int = 20
|
|
|
|
|
|
|
|
|
|
@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("") from e
|
|
|
|
|
fireworks_api_key = get_from_dict_or_env(
|
|
|
|
|
values, "fireworks_api_key", "FIREWORKS_API_KEY"
|
|
|
|
|
)
|
|
|
|
@ -105,14 +121,14 @@ class ChatFireworks(BaseChatModel):
|
|
|
|
|
"messages": message_dicts,
|
|
|
|
|
**self.model_kwargs,
|
|
|
|
|
}
|
|
|
|
|
response = completion_with_retry(self, **params)
|
|
|
|
|
response = completion_with_retry(self, run_manager=run_manager, **params)
|
|
|
|
|
return self._create_chat_result(response)
|
|
|
|
|
|
|
|
|
|
async def _agenerate(
|
|
|
|
|
self,
|
|
|
|
|
messages: List[BaseMessage],
|
|
|
|
|
stop: Optional[List[str]] = None,
|
|
|
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
|
|
|
|
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
|
|
|
|
|
**kwargs: Any,
|
|
|
|
|
) -> ChatResult:
|
|
|
|
|
message_dicts = self._create_message_dicts(messages, stop)
|
|
|
|
@ -121,13 +137,15 @@ class ChatFireworks(BaseChatModel):
|
|
|
|
|
"messages": message_dicts,
|
|
|
|
|
**self.model_kwargs,
|
|
|
|
|
}
|
|
|
|
|
response = await acompletion_with_retry(self, **params)
|
|
|
|
|
response = await acompletion_with_retry(self, run_manager=run_manager, **params)
|
|
|
|
|
return self._create_chat_result(response)
|
|
|
|
|
|
|
|
|
|
def _combine_llm_outputs(self, llm_outputs: List[Optional[dict]]) -> dict:
|
|
|
|
|
if llm_outputs[0] is None:
|
|
|
|
|
return {}
|
|
|
|
|
return llm_outputs[0]
|
|
|
|
|
|
|
|
|
|
def _create_chat_result(self, response: Mapping[str, Any]) -> ChatResult:
|
|
|
|
|
def _create_chat_result(self, response: Any) -> ChatResult:
|
|
|
|
|
generations = []
|
|
|
|
|
for res in response.choices:
|
|
|
|
|
message = convert_dict_to_message(res.message)
|
|
|
|
@ -141,7 +159,7 @@ class ChatFireworks(BaseChatModel):
|
|
|
|
|
|
|
|
|
|
def _create_message_dicts(
|
|
|
|
|
self, messages: List[BaseMessage], stop: Optional[List[str]]
|
|
|
|
|
) -> Tuple[List[Dict[str, Any]]]:
|
|
|
|
|
) -> List[Dict[str, Any]]:
|
|
|
|
|
message_dicts = [convert_message_to_dict(m) for m in messages]
|
|
|
|
|
return message_dicts
|
|
|
|
|
|
|
|
|
@ -160,7 +178,7 @@ class ChatFireworks(BaseChatModel):
|
|
|
|
|
"stream": True,
|
|
|
|
|
**self.model_kwargs,
|
|
|
|
|
}
|
|
|
|
|
for chunk in completion_with_retry(self, **params):
|
|
|
|
|
for chunk in completion_with_retry(self, run_manager=run_manager, **params):
|
|
|
|
|
choice = chunk.choices[0]
|
|
|
|
|
chunk = _convert_delta_to_message_chunk(choice.delta, default_chunk_class)
|
|
|
|
|
finish_reason = choice.finish_reason
|
|
|
|
@ -174,9 +192,9 @@ class ChatFireworks(BaseChatModel):
|
|
|
|
|
self,
|
|
|
|
|
messages: List[BaseMessage],
|
|
|
|
|
stop: Optional[List[str]] = None,
|
|
|
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
|
|
|
|
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
|
|
|
|
|
**kwargs: Any,
|
|
|
|
|
) -> Iterator[ChatGenerationChunk]:
|
|
|
|
|
) -> AsyncIterator[ChatGenerationChunk]:
|
|
|
|
|
message_dicts = self._create_message_dicts(messages, stop)
|
|
|
|
|
default_chunk_class = AIMessageChunk
|
|
|
|
|
params = {
|
|
|
|
@ -185,7 +203,9 @@ class ChatFireworks(BaseChatModel):
|
|
|
|
|
"stream": True,
|
|
|
|
|
**self.model_kwargs,
|
|
|
|
|
}
|
|
|
|
|
async for chunk in await acompletion_with_retry_streaming(self, **params):
|
|
|
|
|
async for chunk in await acompletion_with_retry_streaming(
|
|
|
|
|
self, run_manager=run_manager, **params
|
|
|
|
|
):
|
|
|
|
|
choice = chunk.choices[0]
|
|
|
|
|
chunk = _convert_delta_to_message_chunk(choice.delta, default_chunk_class)
|
|
|
|
|
finish_reason = choice.finish_reason
|
|
|
|
@ -202,6 +222,8 @@ def completion_with_retry(
|
|
|
|
|
**kwargs: Any,
|
|
|
|
|
) -> Any:
|
|
|
|
|
"""Use tenacity to retry the completion call."""
|
|
|
|
|
import fireworks.client
|
|
|
|
|
|
|
|
|
|
retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)
|
|
|
|
|
|
|
|
|
|
@retry_decorator
|
|
|
|
@ -219,6 +241,8 @@ async def acompletion_with_retry(
|
|
|
|
|
**kwargs: Any,
|
|
|
|
|
) -> Any:
|
|
|
|
|
"""Use tenacity to retry the async completion call."""
|
|
|
|
|
import fireworks.client
|
|
|
|
|
|
|
|
|
|
retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)
|
|
|
|
|
|
|
|
|
|
@retry_decorator
|
|
|
|
@ -236,6 +260,8 @@ async def acompletion_with_retry_streaming(
|
|
|
|
|
**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)
|
|
|
|
|
|
|
|
|
|
@retry_decorator
|
|
|
|
@ -254,6 +280,8 @@ def _create_retry_decorator(
|
|
|
|
|
] = None,
|
|
|
|
|
) -> Callable[[Any], Any]:
|
|
|
|
|
"""Define retry mechanism."""
|
|
|
|
|
import fireworks.client
|
|
|
|
|
|
|
|
|
|
errors = [
|
|
|
|
|
fireworks.client.error.RateLimitError,
|
|
|
|
|
fireworks.client.error.ServiceUnavailableError,
|
|
|
|
|