llama-index
67 строк · 2.5 Кб
1from typing import Any, Callable, Dict, Optional, Sequence2
3from llama_index.legacy.callbacks import CallbackManager4from llama_index.legacy.constants import DEFAULT_NUM_OUTPUTS, DEFAULT_TEMPERATURE5from llama_index.legacy.core.llms.types import ChatMessage, LLMMetadata6from llama_index.legacy.llms.everlyai_utils import everlyai_modelname_to_contextsize7from llama_index.legacy.llms.generic_utils import get_from_param_or_env8from llama_index.legacy.llms.openai import OpenAI9from llama_index.legacy.types import BaseOutputParser, PydanticProgramMode10
11EVERLYAI_API_BASE = "https://everlyai.xyz/hosted"12DEFAULT_MODEL = "meta-llama/Llama-2-7b-chat-hf"13
14
15class EverlyAI(OpenAI):16def __init__(17self,18model: str = DEFAULT_MODEL,19temperature: float = DEFAULT_TEMPERATURE,20max_tokens: int = DEFAULT_NUM_OUTPUTS,21additional_kwargs: Optional[Dict[str, Any]] = None,22max_retries: int = 10,23api_key: Optional[str] = None,24callback_manager: Optional[CallbackManager] = None,25system_prompt: Optional[str] = None,26messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,27completion_to_prompt: Optional[Callable[[str], str]] = None,28pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,29output_parser: Optional[BaseOutputParser] = None,30) -> None:31additional_kwargs = additional_kwargs or {}32callback_manager = callback_manager or CallbackManager([])33
34api_key = get_from_param_or_env("api_key", api_key, "EverlyAI_API_KEY")35
36super().__init__(37model=model,38temperature=temperature,39max_tokens=max_tokens,40api_base=EVERLYAI_API_BASE,41api_key=api_key,42additional_kwargs=additional_kwargs,43max_retries=max_retries,44callback_manager=callback_manager,45system_prompt=system_prompt,46messages_to_prompt=messages_to_prompt,47completion_to_prompt=completion_to_prompt,48pydantic_program_mode=pydantic_program_mode,49output_parser=output_parser,50)51
52@classmethod53def class_name(cls) -> str:54return "EverlyAI_LLM"55
56@property57def metadata(self) -> LLMMetadata:58return LLMMetadata(59context_window=everlyai_modelname_to_contextsize(self.model),60num_output=self.max_tokens,61is_chat_model=True,62model_name=self.model,63)64
65@property66def _is_chat_model(self) -> bool:67return True68