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")