3
from typing import Any, Dict, Iterator, List, Optional, Union
7
from rayllm.common.models import ChatCompletion, Completion, Model
8
from rayllm.common.utils import (
28
"""Stores information about the Aviary backend configuration."""
30
def __init__(self, backend_url: str, token: str):
31
assert "::param" not in backend_url, "backend_url not set correctly"
32
assert "::param" not in token, "token not set correctly"
34
self.backend_url = backend_url
36
self.bearer = f"Bearer {token}" if token else ""
39
class URLNotSetException(Exception):
43
def get_aviary_backend(verbose: Optional[bool] = None):
45
Establishes a connection to the Aviary backed after establishing
46
the information using environmental variables.
48
For direct connection to the aviary backend (e.g. running on the same cluster),
49
no AVIARY_TOKEN is required. Otherwise, the AVIARY_URL and AVIARY_TOKEN environment
50
variables are required.
53
verbose: Whether to print the connecting message.
56
An instance of the AviaryResource class.
58
aviary_url = os.getenv("AVIARY_URL", os.getenv("OPENAI_API_BASE"))
60
raise URLNotSetException("AVIARY_URL or OPENAI_API_BASE must be set")
62
aviary_token = os.getenv("AVIARY_TOKEN", os.getenv("OPENAI_API_KEY")) or ""
64
aviary_url += "/v1" if not aviary_url.endswith("/v1") else ""
67
verbose = os.environ.get("AVIARY_SILENT", "0") == "0"
69
print(f"Connecting to Aviary backend at: {aviary_url}")
70
return AviaryResource(aviary_url, aviary_token)
73
def get_openai_client() -> openai.Client:
74
"""Get an OpenAI Client connected to the ray-llm backend."""
75
backend = get_aviary_backend()
76
openai_client = openai.Client(base_url=backend.backend_url, api_key=backend.token)
80
def models() -> List[str]:
81
"""List available models"""
82
models = get_openai_client().models.list()
83
return [model.id for model in models.data]
86
def metadata(model_id: str) -> Dict[str, Dict[str, Any]]:
87
"""Get model metadata"""
88
metadata = get_openai_client().models.retrieve(model_id).model_dump()
95
use_prompt_format: bool = True,
97
) -> Dict[str, Union[str, float, int]]:
98
"""Get completions from Aviary models."""
99
kwargs.setdefault("max_tokens", None)
100
if _is_aviary_model(model):
101
if use_prompt_format:
102
result = get_openai_client().chat.completions.create(
104
messages=[{"role": "user", "content": prompt}],
109
result = get_openai_client().completions.create(
115
return result.model_dump()
116
llm = _get_langchain_model(model)
117
return llm.predict(prompt)
123
use_prompt_format: bool = True,
125
) -> Dict[str, Union[str, float, int]]:
127
"'query' is deprecated, please use 'completions' instead",
131
return completions(model, prompt, use_prompt_format, **kwargs)
142
use_prompt_format: bool = True,
144
) -> Iterator[Dict[str, Union[str, float, int]]]:
145
"""Query Aviary and stream response"""
146
kwargs.setdefault("max_tokens", None)
147
if _is_aviary_model(model):
148
if use_prompt_format:
149
result = get_openai_client().chat.completions.create(
151
messages=[{"role": "user", "content": prompt}],
156
result = get_openai_client().completions.create(
162
return _iterator(result)
165
raise RuntimeError("Streaming is currently only supported for aviary models")
168
def run(*models: List[str], blocking: bool = True) -> None:
169
"""Run Aviary on the local ray cluster
172
*models: Models to run.
173
blocking: Whether to block the CLI until the application is ready.
175
NOTE: This only works if you are running this command
176
on the Ray or Anyscale cluster directly. It does not
177
work from a general machine which only has the url and token
181
from rayllm.backend.server.run import run
183
run(*models, blocking=blocking)
186
def shutdown() -> None:
187
"""Shutdown the Aviary backend server"""
189
from ray import serve