wehub-resource-sync
504 lines
15 KiB
Python
504 lines
15 KiB
Python
from .hookspecs import hookimpl
|
|
from .errors import (
|
|
ModelError,
|
|
NeedsKeyException,
|
|
)
|
|
from .models import (
|
|
AsyncConversation,
|
|
AsyncKeyModel,
|
|
AsyncModel,
|
|
AsyncResponse,
|
|
Attachment,
|
|
CancelToolCall,
|
|
PauseChain,
|
|
Conversation,
|
|
EmbeddingModel,
|
|
EmbeddingModelWithAliases,
|
|
KeyModel,
|
|
Model,
|
|
ModelWithAliases,
|
|
Options,
|
|
Prompt,
|
|
Response,
|
|
Tool,
|
|
Toolbox,
|
|
ToolCall,
|
|
ToolOutput,
|
|
ToolResult,
|
|
Usage,
|
|
)
|
|
from .parts import (
|
|
Message,
|
|
assistant,
|
|
system,
|
|
tool_message,
|
|
user,
|
|
)
|
|
from .utils import schema_dsl, Fragment
|
|
from .embeddings import Collection
|
|
from .templates import Template
|
|
from .plugins import pm, load_plugins
|
|
import click
|
|
from typing import Any, Dict, List, Optional, Callable, Type, Union
|
|
import inspect
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import struct
|
|
|
|
__all__ = [
|
|
"AsyncConversation",
|
|
"AsyncKeyModel",
|
|
"AsyncModel",
|
|
"AsyncResponse",
|
|
"assistant",
|
|
"Attachment",
|
|
"CancelToolCall",
|
|
"Collection",
|
|
"Conversation",
|
|
"Fragment",
|
|
"get_async_model",
|
|
"get_key",
|
|
"get_model",
|
|
"hookimpl",
|
|
"KeyModel",
|
|
"Message",
|
|
"Model",
|
|
"ModelError",
|
|
"NeedsKeyException",
|
|
"Options",
|
|
"PauseChain",
|
|
"Prompt",
|
|
"Response",
|
|
"schema_dsl",
|
|
"system",
|
|
"Template",
|
|
"Tool",
|
|
"Toolbox",
|
|
"ToolCall",
|
|
"tool_message",
|
|
"ToolOutput",
|
|
"ToolResult",
|
|
"Usage",
|
|
"user",
|
|
"user_dir",
|
|
]
|
|
DEFAULT_MODEL = "gpt-4o-mini"
|
|
|
|
|
|
def get_plugins(all=False):
|
|
plugins = []
|
|
plugin_to_distinfo = dict(pm.list_plugin_distinfo())
|
|
for plugin in pm.get_plugins():
|
|
if not all and plugin.__name__.startswith("llm.default_plugins."):
|
|
continue
|
|
plugin_info = {
|
|
"name": plugin.__name__,
|
|
"hooks": [h.name for h in pm.get_hookcallers(plugin)],
|
|
}
|
|
distinfo = plugin_to_distinfo.get(plugin)
|
|
if distinfo:
|
|
plugin_info["version"] = distinfo.version
|
|
plugin_info["name"] = (
|
|
getattr(distinfo, "name", None) or distinfo.project_name
|
|
)
|
|
plugins.append(plugin_info)
|
|
return plugins
|
|
|
|
|
|
def get_models_with_aliases() -> List["ModelWithAliases"]:
|
|
model_aliases = []
|
|
|
|
# Include aliases from aliases.json
|
|
aliases_path = user_dir() / "aliases.json"
|
|
extra_model_aliases: Dict[str, list] = {}
|
|
if aliases_path.exists():
|
|
configured_aliases = json.loads(aliases_path.read_text())
|
|
for alias, model_id in configured_aliases.items():
|
|
extra_model_aliases.setdefault(model_id, []).append(alias)
|
|
|
|
def register(model, async_model=None, aliases=None):
|
|
alias_list = list(aliases or [])
|
|
if model.model_id in extra_model_aliases:
|
|
alias_list.extend(extra_model_aliases[model.model_id])
|
|
model_aliases.append(ModelWithAliases(model, async_model, alias_list))
|
|
|
|
load_plugins()
|
|
pm.hook.register_models(register=register, model_aliases=model_aliases)
|
|
|
|
return model_aliases
|
|
|
|
|
|
def _get_loaders(hook_method) -> Dict[str, Callable]:
|
|
load_plugins()
|
|
loaders = {}
|
|
|
|
def register(prefix, loader):
|
|
suffix = 0
|
|
prefix_to_try = prefix
|
|
while prefix_to_try in loaders:
|
|
suffix += 1
|
|
prefix_to_try = f"{prefix}_{suffix}"
|
|
loaders[prefix_to_try] = loader
|
|
|
|
hook_method(register=register)
|
|
return loaders
|
|
|
|
|
|
def get_template_loaders() -> Dict[str, Callable[[str], Template]]:
|
|
"""Get template loaders registered by plugins."""
|
|
return _get_loaders(pm.hook.register_template_loaders)
|
|
|
|
|
|
def get_fragment_loaders() -> Dict[
|
|
str,
|
|
Callable[[str], Union[Fragment, Attachment, List[Union[Fragment, Attachment]]]],
|
|
]:
|
|
"""Get fragment loaders registered by plugins."""
|
|
return _get_loaders(pm.hook.register_fragment_loaders)
|
|
|
|
|
|
def get_tools() -> Dict[str, Union[Tool, Type[Toolbox]]]:
|
|
"""Return all tools (llm.Tool and llm.Toolbox) registered by plugins."""
|
|
load_plugins()
|
|
tools: Dict[str, Union[Tool, Type[Toolbox]]] = {}
|
|
|
|
# Variable to track current plugin name
|
|
current_plugin_name = None
|
|
|
|
def register(
|
|
tool_or_function: Union[Tool, Type[Toolbox], Callable[..., Any]],
|
|
name: Optional[str] = None,
|
|
) -> None:
|
|
tool: Union[Tool, Type[Toolbox], None] = None
|
|
|
|
# If it's a Toolbox class, set the plugin field on it
|
|
if inspect.isclass(tool_or_function):
|
|
if issubclass(tool_or_function, Toolbox):
|
|
tool = tool_or_function
|
|
if current_plugin_name:
|
|
tool.plugin = current_plugin_name
|
|
tool.name = name or tool.__name__
|
|
else:
|
|
raise TypeError(
|
|
"Toolbox classes must inherit from llm.Toolbox, {} does not.".format(
|
|
tool_or_function.__name__
|
|
)
|
|
)
|
|
|
|
# If it's already a Tool instance, use it directly
|
|
elif isinstance(tool_or_function, Tool):
|
|
tool = tool_or_function
|
|
if name:
|
|
tool.name = name
|
|
if current_plugin_name:
|
|
tool.plugin = current_plugin_name
|
|
|
|
# If it's a bare function, wrap it in a Tool
|
|
else:
|
|
tool = Tool.function(tool_or_function, name=name)
|
|
if current_plugin_name:
|
|
tool.plugin = current_plugin_name
|
|
|
|
# Get the name for the tool/toolbox
|
|
if tool:
|
|
# For Toolbox classes, use their name attribute or class name
|
|
if inspect.isclass(tool) and issubclass(tool, Toolbox):
|
|
prefix = name or getattr(tool, "name", tool.__name__) or ""
|
|
else:
|
|
prefix = name or tool.name or ""
|
|
|
|
suffix = 0
|
|
candidate = prefix
|
|
|
|
# Avoid name collisions
|
|
while candidate in tools:
|
|
suffix += 1
|
|
candidate = f"{prefix}_{suffix}"
|
|
|
|
tools[candidate] = tool
|
|
|
|
# Call each plugin's register_tools hook individually to track current_plugin_name
|
|
for plugin in pm.get_plugins():
|
|
current_plugin_name = pm.get_name(plugin)
|
|
hook_caller = pm.hook.register_tools
|
|
plugin_impls = [
|
|
impl for impl in hook_caller.get_hookimpls() if impl.plugin is plugin
|
|
]
|
|
for impl in plugin_impls:
|
|
impl.function(register=register)
|
|
|
|
return tools
|
|
|
|
|
|
def get_embedding_models_with_aliases() -> List["EmbeddingModelWithAliases"]:
|
|
model_aliases = []
|
|
|
|
# Include aliases from aliases.json
|
|
aliases_path = user_dir() / "aliases.json"
|
|
extra_model_aliases: Dict[str, list] = {}
|
|
if aliases_path.exists():
|
|
configured_aliases = json.loads(aliases_path.read_text())
|
|
for alias, model_id in configured_aliases.items():
|
|
extra_model_aliases.setdefault(model_id, []).append(alias)
|
|
|
|
def register(model, aliases=None):
|
|
alias_list = list(aliases or [])
|
|
if model.model_id in extra_model_aliases:
|
|
alias_list.extend(extra_model_aliases[model.model_id])
|
|
model_aliases.append(EmbeddingModelWithAliases(model, alias_list))
|
|
|
|
load_plugins()
|
|
pm.hook.register_embedding_models(register=register)
|
|
|
|
return model_aliases
|
|
|
|
|
|
def get_embedding_models():
|
|
models = []
|
|
|
|
def register(model, aliases=None):
|
|
models.append(model)
|
|
|
|
load_plugins()
|
|
pm.hook.register_embedding_models(register=register)
|
|
return models
|
|
|
|
|
|
def get_embedding_model(name):
|
|
aliases = get_embedding_model_aliases()
|
|
try:
|
|
return aliases[name]
|
|
except KeyError:
|
|
raise UnknownModelError("Unknown model: " + str(name))
|
|
|
|
|
|
def get_embedding_model_aliases() -> Dict[str, EmbeddingModel]:
|
|
model_aliases = {}
|
|
for model_with_aliases in get_embedding_models_with_aliases():
|
|
for alias in model_with_aliases.aliases:
|
|
model_aliases[alias] = model_with_aliases.model
|
|
model_aliases[model_with_aliases.model.model_id] = model_with_aliases.model
|
|
return model_aliases
|
|
|
|
|
|
def get_async_model_aliases() -> Dict[str, AsyncModel]:
|
|
async_model_aliases = {}
|
|
for model_with_aliases in get_models_with_aliases():
|
|
if model_with_aliases.async_model:
|
|
for alias in model_with_aliases.aliases:
|
|
async_model_aliases[alias] = model_with_aliases.async_model
|
|
async_model_aliases[model_with_aliases.model.model_id] = (
|
|
model_with_aliases.async_model
|
|
)
|
|
return async_model_aliases
|
|
|
|
|
|
def get_model_aliases() -> Dict[str, Model]:
|
|
model_aliases = {}
|
|
for model_with_aliases in get_models_with_aliases():
|
|
if model_with_aliases.model:
|
|
for alias in model_with_aliases.aliases:
|
|
model_aliases[alias] = model_with_aliases.model
|
|
model_aliases[model_with_aliases.model.model_id] = model_with_aliases.model
|
|
return model_aliases
|
|
|
|
|
|
class UnknownModelError(KeyError):
|
|
pass
|
|
|
|
|
|
def get_models() -> List[Model]:
|
|
"Get all registered models"
|
|
models_with_aliases = get_models_with_aliases()
|
|
return [mwa.model for mwa in models_with_aliases if mwa.model]
|
|
|
|
|
|
def get_async_models() -> List[AsyncModel]:
|
|
"Get all registered async models"
|
|
models_with_aliases = get_models_with_aliases()
|
|
return [mwa.async_model for mwa in models_with_aliases if mwa.async_model]
|
|
|
|
|
|
def get_async_model(name: Optional[str] = None) -> AsyncModel:
|
|
"Get an async model by name or alias"
|
|
aliases = get_async_model_aliases()
|
|
name = name or get_default_model()
|
|
try:
|
|
return aliases[name]
|
|
except KeyError:
|
|
# Does a sync model exist?
|
|
sync_model = None
|
|
try:
|
|
sync_model = get_model(name, _skip_async=True)
|
|
except UnknownModelError:
|
|
pass
|
|
if sync_model:
|
|
raise UnknownModelError("Unknown async model (sync model exists): " + name)
|
|
else:
|
|
raise UnknownModelError("Unknown model: " + name)
|
|
|
|
|
|
def get_model(name: Optional[str] = None, _skip_async: bool = False) -> Model:
|
|
"Get a model by name or alias"
|
|
aliases = get_model_aliases()
|
|
name = name or get_default_model()
|
|
try:
|
|
return aliases[name]
|
|
except KeyError:
|
|
# Does an async model exist?
|
|
if _skip_async:
|
|
raise UnknownModelError("Unknown model: " + name)
|
|
async_model = None
|
|
try:
|
|
async_model = get_async_model(name)
|
|
except UnknownModelError:
|
|
pass
|
|
if async_model:
|
|
raise UnknownModelError("Unknown model (async model exists): " + name)
|
|
else:
|
|
raise UnknownModelError("Unknown model: " + name)
|
|
|
|
|
|
def get_key(
|
|
explicit_key: Optional[str] = None,
|
|
key_alias: Optional[str] = None,
|
|
env_var: Optional[str] = None,
|
|
*,
|
|
alias: Optional[str] = None,
|
|
env: Optional[str] = None,
|
|
input: Optional[str] = None,
|
|
) -> Optional[str]:
|
|
"""
|
|
Return an API key based on a hierarchy of potential sources. You should use the keyword arguments,
|
|
the positional arguments are here purely for backwards-compatibility with older code.
|
|
|
|
:param input: Input provided by the user. This may be the key, or an alias of a key in keys.json.
|
|
:param alias: The alias used to retrieve the key from the keys.json file.
|
|
:param env: Name of the environment variable to check for the key as a final fallback.
|
|
"""
|
|
if alias:
|
|
key_alias = alias
|
|
if env:
|
|
env_var = env
|
|
if input:
|
|
explicit_key = input
|
|
stored_keys = load_keys()
|
|
# If user specified an alias, use the key stored for that alias
|
|
if explicit_key in stored_keys:
|
|
return stored_keys[explicit_key]
|
|
if explicit_key:
|
|
# User specified a key that's not an alias, use that
|
|
return explicit_key
|
|
# Stored key over-rides environment variables over-ride the default key
|
|
if key_alias in stored_keys:
|
|
return stored_keys[key_alias]
|
|
# Finally try environment variable
|
|
if env_var and os.environ.get(env_var):
|
|
return os.environ[env_var]
|
|
# Couldn't find it
|
|
return None
|
|
|
|
|
|
def load_keys():
|
|
path = user_dir() / "keys.json"
|
|
if path.exists():
|
|
return json.loads(path.read_text())
|
|
else:
|
|
return {}
|
|
|
|
|
|
def user_dir():
|
|
llm_user_path = os.environ.get("LLM_USER_PATH")
|
|
if llm_user_path:
|
|
path = pathlib.Path(llm_user_path)
|
|
else:
|
|
path = pathlib.Path(click.get_app_dir("io.datasette.llm"))
|
|
path.mkdir(exist_ok=True, parents=True)
|
|
return path
|
|
|
|
|
|
def set_alias(alias, model_id_or_alias):
|
|
"""
|
|
Set an alias to point to the specified model.
|
|
"""
|
|
path = user_dir() / "aliases.json"
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
if not path.exists():
|
|
path.write_text("{}\n")
|
|
try:
|
|
current = json.loads(path.read_text())
|
|
except json.decoder.JSONDecodeError:
|
|
# We're going to write a valid JSON file in a moment:
|
|
current = {}
|
|
# Resolve model_id_or_alias to a model_id
|
|
try:
|
|
model = get_model(model_id_or_alias)
|
|
model_id = model.model_id
|
|
except UnknownModelError:
|
|
# Try to resolve it to an embedding model
|
|
try:
|
|
model = get_embedding_model(model_id_or_alias)
|
|
model_id = model.model_id
|
|
except UnknownModelError:
|
|
# Set the alias to the exact string they provided instead
|
|
model_id = model_id_or_alias
|
|
current[alias] = model_id
|
|
path.write_text(json.dumps(current, indent=4) + "\n")
|
|
|
|
|
|
def remove_alias(alias):
|
|
"""
|
|
Remove an alias.
|
|
"""
|
|
path = user_dir() / "aliases.json"
|
|
if not path.exists():
|
|
raise KeyError("No aliases.json file exists")
|
|
try:
|
|
current = json.loads(path.read_text())
|
|
except json.decoder.JSONDecodeError:
|
|
raise KeyError("aliases.json file is not valid JSON")
|
|
if alias not in current:
|
|
raise KeyError("No such alias: {}".format(alias))
|
|
del current[alias]
|
|
path.write_text(json.dumps(current, indent=4) + "\n")
|
|
|
|
|
|
def encode(values):
|
|
return struct.pack("<" + "f" * len(values), *values)
|
|
|
|
|
|
def decode(binary):
|
|
return struct.unpack("<" + "f" * (len(binary) // 4), binary)
|
|
|
|
|
|
def cosine_similarity(a, b):
|
|
dot_product = sum(x * y for x, y in zip(a, b))
|
|
magnitude_a = sum(x * x for x in a) ** 0.5
|
|
magnitude_b = sum(x * x for x in b) ** 0.5
|
|
return dot_product / (magnitude_a * magnitude_b)
|
|
|
|
|
|
def get_default_model(filename="default_model.txt", default=DEFAULT_MODEL):
|
|
path = user_dir() / filename
|
|
if path.exists():
|
|
return path.read_text().strip()
|
|
else:
|
|
return default
|
|
|
|
|
|
def set_default_model(model, filename="default_model.txt"):
|
|
path = user_dir() / filename
|
|
if model is None and path.exists():
|
|
path.unlink()
|
|
else:
|
|
path.write_text(model)
|
|
|
|
|
|
def get_default_embedding_model():
|
|
return get_default_model("default_embedding_model.txt", None)
|
|
|
|
|
|
def set_default_embedding_model(model):
|
|
set_default_model(model, "default_embedding_model.txt")
|