llama-index
60 строк · 2.0 Кб
1from typing import Any, Dict, Optional
2
3from llama_index.legacy.bridge.pydantic import Field
4from llama_index.legacy.constants import (
5DEFAULT_CONTEXT_WINDOW,
6DEFAULT_NUM_OUTPUTS,
7DEFAULT_TEMPERATURE,
8)
9from llama_index.legacy.core.llms.types import LLMMetadata
10from llama_index.legacy.llms.generic_utils import get_from_param_or_env
11from llama_index.legacy.llms.openai_like import OpenAILike
12
13DEFAULT_API_BASE = "https://openrouter.ai/api/v1"
14DEFAULT_MODEL = "gryphe/mythomax-l2-13b"
15
16
17class OpenRouter(OpenAILike):
18model: str = Field(
19description="The OpenRouter model to use. See https://openrouter.ai/models for options."
20)
21context_window: int = Field(
22default=DEFAULT_CONTEXT_WINDOW,
23description="The maximum number of context tokens for the model. See https://openrouter.ai/models for options.",
24gt=0,
25)
26is_chat_model: bool = Field(
27default=True,
28description=LLMMetadata.__fields__["is_chat_model"].field_info.description,
29)
30
31def __init__(
32self,
33model: str = DEFAULT_MODEL,
34temperature: float = DEFAULT_TEMPERATURE,
35max_tokens: int = DEFAULT_NUM_OUTPUTS,
36additional_kwargs: Optional[Dict[str, Any]] = None,
37max_retries: int = 5,
38api_base: Optional[str] = DEFAULT_API_BASE,
39api_key: Optional[str] = None,
40**kwargs: Any,
41) -> None:
42additional_kwargs = additional_kwargs or {}
43
44api_base = get_from_param_or_env("api_base", api_base, "OPENROUTER_API_BASE")
45api_key = get_from_param_or_env("api_key", api_key, "OPENROUTER_API_KEY")
46
47super().__init__(
48model=model,
49temperature=temperature,
50max_tokens=max_tokens,
51api_base=api_base,
52api_key=api_key,
53additional_kwargs=additional_kwargs,
54max_retries=max_retries,
55**kwargs,
56)
57
58@classmethod
59def class_name(cls) -> str:
60return "OpenRouter_LLM"
61