File size: 4,121 Bytes
0703e2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """OpenAI-compatible and Anthropic chat models as FGAO LLMs (sync and native async).
``OpenAILLM`` also serves any OpenAI-compatible endpoint (vLLM, TGI, Ollama,
LM Studio, ...) through ``base_url``. Clients are created lazily, so a process
that only uses the async API never opens a sync client and vice versa.
"""
from __future__ import annotations
from typing import Any, Dict, List, Optional
from ..rag.prompts import SYSTEM_PROMPT
def _messages(system: Optional[str], prompt: str) -> List[Dict[str, str]]:
return ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": prompt}]
class OpenAILLM:
def __init__(self, model: str, *, client: Any = None, aclient: Any = None, base_url: Optional[str] = None,
api_key: Optional[str] = None, max_tokens: int = 128, temperature: float = 0.0,
system: Optional[str] = SYSTEM_PROMPT, timeout: Optional[float] = None,
max_retries: Optional[int] = None) -> None:
self.model, self.max_tokens, self.temperature, self.system = model, max_tokens, temperature, system
self._kw = {k: v for k, v in (("base_url", base_url), ("api_key", api_key), ("timeout", timeout),
("max_retries", max_retries)) if v is not None}
self._client, self._aclient = client, aclient
@property
def client(self):
if self._client is None:
from openai import OpenAI
self._client = OpenAI(**self._kw)
return self._client
@property
def aclient(self):
if self._aclient is None:
from openai import AsyncOpenAI
self._aclient = AsyncOpenAI(**self._kw)
return self._aclient
def _args(self, prompt: str) -> Dict[str, Any]:
return {"model": self.model, "messages": _messages(self.system, prompt), "max_tokens": self.max_tokens,
"temperature": self.temperature}
def complete(self, prompt: str) -> str:
out = self.client.chat.completions.create(**self._args(prompt))
return (out.choices[0].message.content or "").strip()
async def acomplete(self, prompt: str) -> str:
out = await self.aclient.chat.completions.create(**self._args(prompt))
return (out.choices[0].message.content or "").strip()
class AnthropicLLM:
def __init__(self, model: str, *, client: Any = None, aclient: Any = None, api_key: Optional[str] = None,
max_tokens: int = 128, temperature: float = 0.0, system: Optional[str] = SYSTEM_PROMPT,
timeout: Optional[float] = None, max_retries: Optional[int] = None) -> None:
self.model, self.max_tokens, self.temperature, self.system = model, max_tokens, temperature, system
self._kw = {k: v for k, v in (("api_key", api_key), ("timeout", timeout), ("max_retries", max_retries))
if v is not None}
self._client, self._aclient = client, aclient
@property
def client(self):
if self._client is None:
from anthropic import Anthropic
self._client = Anthropic(**self._kw)
return self._client
@property
def aclient(self):
if self._aclient is None:
from anthropic import AsyncAnthropic
self._aclient = AsyncAnthropic(**self._kw)
return self._aclient
def _args(self, prompt: str) -> Dict[str, Any]:
kw: Dict[str, Any] = {"model": self.model, "max_tokens": self.max_tokens, "temperature": self.temperature,
"messages": [{"role": "user", "content": prompt}]}
if self.system:
kw["system"] = self.system
return kw
def complete(self, prompt: str) -> str:
out = self.client.messages.create(**self._args(prompt))
return "".join(getattr(b, "text", "") for b in out.content).strip()
async def acomplete(self, prompt: str) -> str:
out = await self.aclient.messages.create(**self._args(prompt))
return "".join(getattr(b, "text", "") for b in out.content).strip()
__all__ = ["OpenAILLM", "AnthropicLLM"]
|