Ruff >= 0.16.0 (#1557)

* ruff>=0.16.0
* Fixed all ruff issues with --fix and --unsafe-fixes
* Codex (GPT-5.6 Sol High) fixed remaining Ruff errors: https://gist.github.com/simonw/53404d27979d28f66ae59564d9fb3382
* Ruff target-version = "py310"
This commit is contained in:
Simon Willison
2026-07-25 14:16:47 -07:00
committed by GitHub
parent 1dc49af9a9
commit c12fd50b9f
45 changed files with 985 additions and 967 deletions
-3
View File
@@ -1,6 +1,3 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from subprocess import PIPE, Popen
# This file is execfile()d with the current directory set to its
+6 -5
View File
@@ -1,8 +1,9 @@
import llm
import random
import time
from typing import Optional
from pydantic import field_validator, Field
from pydantic import Field, field_validator
import llm
@llm.hookimpl
@@ -35,10 +36,10 @@ class Markov(llm.Model):
can_stream = True
class Options(llm.Options):
length: Optional[int] = Field(
length: int | None = Field(
description="Number of words to generate", default=None
)
delay: Optional[float] = Field(
delay: float | None = Field(
description="Seconds to delay between each token", default=None
)
+53 -52
View File
@@ -1,8 +1,19 @@
from .hookspecs import hookimpl
import inspect
import json
import os
import pathlib
import struct
from collections.abc import Callable
from typing import Any
import click
from .embeddings import Collection
from .errors import (
ModelError,
NeedsKeyException,
)
from .hookspecs import hookimpl
from .models import (
AsyncConversation,
AsyncKeyModel,
@@ -10,7 +21,6 @@ from .models import (
AsyncResponse,
Attachment,
CancelToolCall,
PauseChain,
Conversation,
EmbeddingModel,
EmbeddingModelWithAliases,
@@ -18,6 +28,7 @@ from .models import (
Model,
ModelWithAliases,
Options,
PauseChain,
Prompt,
Response,
Tool,
@@ -34,33 +45,20 @@ from .parts import (
tool_message,
user,
)
from .utils import schema_dsl, Fragment
from .embeddings import Collection
from .plugins import load_plugins, pm
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
from .utils import Fragment, schema_dsl
__all__ = [
"AsyncConversation",
"AsyncKeyModel",
"AsyncModel",
"AsyncResponse",
"assistant",
"Attachment",
"CancelToolCall",
"Collection",
"Conversation",
"Fragment",
"get_async_model",
"get_key",
"get_model",
"hookimpl",
"KeyModel",
"Message",
"Model",
@@ -70,16 +68,21 @@ __all__ = [
"PauseChain",
"Prompt",
"Response",
"schema_dsl",
"system",
"Template",
"Tool",
"Toolbox",
"ToolCall",
"tool_message",
"ToolOutput",
"ToolResult",
"Toolbox",
"Usage",
"assistant",
"get_async_model",
"get_key",
"get_model",
"hookimpl",
"schema_dsl",
"system",
"tool_message",
"user",
"user_dir",
]
@@ -106,12 +109,12 @@ def get_plugins(all=False):
return plugins
def get_models_with_aliases() -> List["ModelWithAliases"]:
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] = {}
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():
@@ -129,7 +132,7 @@ def get_models_with_aliases() -> List["ModelWithAliases"]:
return model_aliases
def _get_loaders(hook_method) -> Dict[str, Callable]:
def _get_loaders(hook_method) -> dict[str, Callable]:
load_plugins()
loaders = {}
@@ -145,32 +148,32 @@ def _get_loaders(hook_method) -> Dict[str, Callable]:
return loaders
def get_template_loaders() -> Dict[str, Callable[[str], Template]]:
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[
def get_fragment_loaders() -> dict[
str,
Callable[[str], Union[Fragment, Attachment, List[Union[Fragment, Attachment]]]],
Callable[[str], Fragment | Attachment | list[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]]]:
def get_tools() -> dict[str, Tool | type[Toolbox]]:
"""Return all tools (llm.Tool and llm.Toolbox) registered by plugins."""
load_plugins()
tools: Dict[str, Union[Tool, Type[Toolbox]]] = {}
tools: dict[str, 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,
tool_or_function: Tool | type[Toolbox] | Callable[..., Any],
name: str | None = None,
) -> None:
tool: Union[Tool, Type[Toolbox], None] = None
tool: Tool | type[Toolbox] | None = None
# If it's a Toolbox class, set the plugin field on it
if inspect.isclass(tool_or_function):
@@ -181,9 +184,7 @@ def get_tools() -> Dict[str, Union[Tool, Type[Toolbox]]]:
tool.name = name or tool.__name__
else:
raise TypeError(
"Toolbox classes must inherit from llm.Toolbox, {} does not.".format(
tool_or_function.__name__
)
f"Toolbox classes must inherit from llm.Toolbox, {tool_or_function.__name__} does not."
)
# If it's already a Tool instance, use it directly
@@ -231,12 +232,12 @@ def get_tools() -> Dict[str, Union[Tool, Type[Toolbox]]]:
return tools
def get_embedding_models_with_aliases() -> List["EmbeddingModelWithAliases"]:
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] = {}
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():
@@ -273,7 +274,7 @@ def get_embedding_model(name):
raise UnknownModelError("Unknown model: " + str(name))
def get_embedding_model_aliases() -> Dict[str, EmbeddingModel]:
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:
@@ -282,7 +283,7 @@ def get_embedding_model_aliases() -> Dict[str, EmbeddingModel]:
return model_aliases
def get_async_model_aliases() -> Dict[str, AsyncModel]:
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:
@@ -294,7 +295,7 @@ def get_async_model_aliases() -> Dict[str, AsyncModel]:
return async_model_aliases
def get_model_aliases() -> Dict[str, Model]:
def get_model_aliases() -> dict[str, Model]:
model_aliases = {}
for model_with_aliases in get_models_with_aliases():
if model_with_aliases.model:
@@ -308,19 +309,19 @@ class UnknownModelError(KeyError):
pass
def get_models() -> List[Model]:
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]:
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:
def get_async_model(name: str | None = None) -> AsyncModel:
"Get an async model by name or alias"
aliases = get_async_model_aliases()
name = name or get_default_model()
@@ -339,7 +340,7 @@ def get_async_model(name: Optional[str] = None) -> AsyncModel:
raise UnknownModelError("Unknown model: " + name)
def get_model(name: Optional[str] = None, _skip_async: bool = False) -> Model:
def get_model(name: str | None = None, _skip_async: bool = False) -> Model:
"Get a model by name or alias"
aliases = get_model_aliases()
name = name or get_default_model()
@@ -361,14 +362,14 @@ def get_model(name: Optional[str] = None, _skip_async: bool = False) -> Model:
def get_key(
explicit_key: Optional[str] = None,
key_alias: Optional[str] = None,
env_var: Optional[str] = None,
explicit_key: str | None = None,
key_alias: str | None = None,
env_var: str | None = None,
*,
alias: Optional[str] = None,
env: Optional[str] = None,
input: Optional[str] = None,
) -> Optional[str]:
alias: str | None = None,
env: str | None = None,
input: str | None = None,
) -> str | None:
"""
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.
@@ -459,7 +460,7 @@ def remove_alias(alias):
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))
raise KeyError(f"No such alias: {alias}")
del current[alias]
path.write_text(json.dumps(current, indent=4) + "\n")
+151 -175
View File
@@ -1,50 +1,69 @@
import asyncio
import click
from click_default_group import DefaultGroup
from dataclasses import asdict
from importlib.metadata import version
import base64
import inspect
import io
import json
import os
import pathlib
import re
import readline
import shutil
import sys
import textwrap
import warnings
from collections.abc import Iterable
from dataclasses import asdict
from importlib.metadata import version
from runpy import run_module
from typing import Any, cast
import click
import httpx
import pydantic
import sqlite_utils
import yaml
from click_default_group import DefaultGroup
from sqlite_utils.utils import Format, rows_from_file
from llm import (
Attachment,
AsyncConversation,
AsyncKeyModel,
AsyncResponse,
Attachment,
CancelToolCall,
Collection,
Conversation,
Fragment,
KeyModel,
Response,
Template,
Tool,
Toolbox,
UnknownModelError,
KeyModel,
encode,
get_async_model,
get_default_model,
get_default_embedding_model,
get_embedding_models_with_aliases,
get_embedding_model_aliases,
get_default_model,
get_embedding_model,
get_plugins,
get_tools,
get_embedding_model_aliases,
get_embedding_models_with_aliases,
get_fragment_loaders,
get_template_loaders,
get_model,
get_model_aliases,
get_models_with_aliases,
user_dir,
set_alias,
set_default_model,
set_default_embedding_model,
get_plugins,
get_template_loaders,
get_tools,
remove_alias,
set_alias,
set_default_embedding_model,
set_default_model,
user_dir,
)
from llm.models import _BaseConversation, ChainResponse
from llm.models import ChainResponse, _BaseConversation
from .migrations import migrate
from .plugins import pm, load_plugins
from .plugins import load_plugins, pm
from .utils import (
ensure_fragment,
extract_fenced_code_block,
@@ -63,22 +82,6 @@ from .utils import (
token_usage_string,
truncate_string,
)
import base64
import httpx
import inspect
import pathlib
import pydantic
import re
import readline
from runpy import run_module
import shutil
import sqlite_utils
from sqlite_utils.utils import rows_from_file, Format
import sys
import textwrap
from typing import cast, Dict, Optional, Iterable, List, Union, Tuple, Type, Any
import warnings
import yaml
warnings.simplefilter("ignore", ResourceWarning)
@@ -130,12 +133,12 @@ def validate_fragment_alias(ctx, param, value):
def resolve_fragments(
db: sqlite_utils.Database, fragments: Iterable[str], allow_attachments: bool = False
) -> List[Union[Fragment, Attachment]]:
) -> list[Fragment | Attachment]:
"""
Resolve fragment strings into a mixed of llm.Fragment() and llm.Attachment() objects.
"""
def _load_by_alias(fragment: str) -> Tuple[Optional[str], Optional[str]]:
def _load_by_alias(fragment: str) -> tuple[str | None, str | None]:
rows = list(
db.query(
"""
@@ -152,9 +155,9 @@ def resolve_fragments(
return None, None
# The fragment strings could be URLs or paths or plugin references
resolved: List[Union[Fragment, Attachment]] = []
resolved: list[Fragment | Attachment] = []
for fragment in fragments:
if fragment.startswith("http://") or fragment.startswith("https://"):
if fragment.startswith(("http://", "https://")):
llm_version = version("llm")
headers = {"User-Agent": f"llm/{llm_version} (https://llm.datasette.io/)"}
client = httpx.Client(
@@ -169,7 +172,7 @@ def resolve_fragments(
prefix, rest = fragment.split(":", 1)
loaders = get_fragment_loaders()
if prefix not in loaders:
raise FragmentNotFound("Unknown fragment prefix: {}".format(prefix))
raise FragmentNotFound(f"Unknown fragment prefix: {prefix}")
loader = loaders[prefix]
try:
result = loader(rest)
@@ -179,15 +182,11 @@ def resolve_fragments(
isinstance(r, Attachment) for r in result
):
raise FragmentNotFound(
"Fragment loader {} returned a disallowed attachment".format(
prefix
)
f"Fragment loader {prefix} returned a disallowed attachment"
)
resolved.extend(result)
except Exception as ex:
raise FragmentNotFound(
"Could not load fragment {}: {}".format(fragment, ex)
)
except Exception as ex: # noqa: BLE001
raise FragmentNotFound(f"Could not load fragment {fragment}: {ex}")
else:
# Try from the DB
content, source = _load_by_alias(fragment)
@@ -239,8 +238,6 @@ def process_fragments_in_chat(
class AttachmentError(Exception):
"""Exception raised for errors in attachment resolution."""
pass
def resolve_attachment(value):
"""
@@ -310,7 +307,7 @@ def resolve_attachment_with_type(value: str, mimetype: str) -> Attachment:
return attachment
def attachment_types_callback(ctx, param, values) -> List[Attachment]:
def attachment_types_callback(ctx, param, values) -> list[Attachment]:
collected = []
for value, mimetype in values:
collected.append(resolve_attachment_with_type(value, mimetype))
@@ -667,7 +664,7 @@ def prompt(
try:
to_save["model"] = model_aliases[model_id].model_id
except KeyError:
raise click.ClickException("'{}' is not a known model".format(model_id))
raise click.ClickException(f"'{model_id}' is not a known model")
prompt = read_prompt()
if prompt:
to_save["prompt"] = prompt
@@ -828,11 +825,11 @@ def prompt(
if options:
# Validate with pydantic
try:
validated_options = dict(
(key, value)
validated_options = {
key: value
for key, value in model.Options(**dict(options))
if value is not None
)
}
except pydantic.ValidationError as ex:
raise click.ClickException(render_errors(ex.errors()))
@@ -914,7 +911,7 @@ def prompt(
response.astream_events(),
show_reasoning=not hide_reasoning,
)
print("")
print()
else:
response = prompt_method(
prompt,
@@ -949,7 +946,7 @@ def prompt(
response.stream_events(),
show_reasoning=not hide_reasoning,
)
print("")
print()
else:
text = response.text()
if extract or extract_last:
@@ -975,7 +972,7 @@ def prompt(
# Show token usage to stderr in yellow
click.echo(
click.style(
"Token usage: {}".format(response_object.token_usage()),
f"Token usage: {response_object.token_usage()}",
fg="yellow",
bold=True,
),
@@ -1154,7 +1151,7 @@ def chat(
try:
model = get_model(model_id)
except KeyError:
raise click.ClickException("'{}' is not a known model".format(model_id))
raise click.ClickException(f"'{model_id}' is not a known model")
if conversation is None:
# Start a fresh conversation for this chat
@@ -1172,11 +1169,11 @@ def chat(
validated_options = get_model_options(model.model_id)
if options:
try:
validated_options = dict(
(key, value)
validated_options = {
key: value
for key, value in model.Options(**dict(options))
if value is not None
)
}
except pydantic.ValidationError as ex:
raise click.ClickException(render_errors(ex.errors()))
@@ -1217,7 +1214,7 @@ def chat(
except FragmentNotFound as ex:
raise click.ClickException(str(ex))
click.echo("Chatting with {}".format(model.model_id))
click.echo(f"Chatting with {model.model_id}")
click.echo("Type 'exit' or 'quit' to exit")
click.echo("Type '!multi' to enter multiple lines, then '!end' to finish")
click.echo("Type '!edit' to open your default editor and modify the prompt")
@@ -1306,14 +1303,14 @@ def chat(
show_reasoning=not hide_reasoning,
)
response.log_to_db(db)
print("")
print()
def load_conversation(
conversation_id: Optional[str],
conversation_id: str | None,
async_=False,
database=None,
) -> Optional[_BaseConversation]:
) -> _BaseConversation | None:
log_path = pathlib.Path(database) if database else logs_db_path()
db = sqlite_utils.Database(log_path)
migrate(db)
@@ -1327,9 +1324,7 @@ def load_conversation(
try:
row = cast(sqlite_utils.db.Table, db["conversations"]).get(conversation_id)
except sqlite_utils.db.NotFoundError:
raise click.ClickException(
"No conversation found with id={}".format(conversation_id)
)
raise click.ClickException(f"No conversation found with id={conversation_id}")
# Inflate that conversation
conversation_class = AsyncConversation if async_ else Conversation
response_class = AsyncResponse if async_ else Response
@@ -1399,7 +1394,7 @@ def keys_get(name):
try:
click.echo(keys[name])
except KeyError:
raise click.ClickException("No key found with name '{}'".format(name))
raise click.ClickException(f"No key found with name '{name}'")
@keys.command(name="set")
@@ -1449,7 +1444,7 @@ def logs_status():
"Show current status of database logging"
path = logs_db_path()
if not path.exists():
click.echo("No log database found at {}".format(path))
click.echo(f"No log database found at {path}")
return
if logs_on():
click.echo("Logging is ON for all prompts".format())
@@ -1457,12 +1452,10 @@ def logs_status():
click.echo("Logging is OFF".format())
db = sqlite_utils.Database(path)
migrate(db)
click.echo("Found log database at {}".format(path))
click.echo(f"Found log database at {path}")
click.echo("Number of conversations logged:\t{}".format(db["conversations"].count))
click.echo("Number of responses logged:\t{}".format(db["responses"].count))
click.echo(
"Database file size: \t\t{}".format(_human_readable_size(path.stat().st_size))
)
click.echo(f"Database file size: \t\t{_human_readable_size(path.stat().st_size)}")
@logs.command(name="backup")
@@ -1474,11 +1467,9 @@ def backup(path):
db = sqlite_utils.Database(logs_path)
try:
db.execute("vacuum into ?", [str(path)])
except Exception as ex:
except Exception as ex: # noqa: BLE001
raise click.ClickException(str(ex))
click.echo(
"Backed up {} to {}".format(_human_readable_size(path.stat().st_size), path)
)
click.echo(f"Backed up {_human_readable_size(path.stat().st_size)} to {path}")
@logs.command(name="on")
@@ -1688,7 +1679,7 @@ def logs_list(
path = database
path = pathlib.Path(path or logs_db_path())
if not path.exists():
raise click.ClickException("No log database found at {}".format(path))
raise click.ClickException(f"No log database found at {path}")
db = sqlite_utils.Database(path)
migrate(db)
@@ -1706,7 +1697,7 @@ def logs_list(
if flag[1]
]
)
raise click.ClickException("Cannot use --short and {} together".format(invalid))
raise click.ClickException(f"Cannot use --short and {invalid} together")
if response and not current_conversation and not conversation_id:
current_conversation = True
@@ -1747,7 +1738,7 @@ def logs_list(
limit = ""
if count is not None and count > 0:
limit = " limit {}".format(count)
limit = f" limit {count}"
sql_format = {
"limit": limit,
@@ -1797,7 +1788,7 @@ def logs_list(
)
"""
exists_clauses.append(exists_clause)
sql_params["f{}".format(i)] = fragment_hash
sql_params[f"f{i}"] = fragment_hash
where_bits.append(" and ".join(exists_clauses))
@@ -1900,16 +1891,14 @@ def logs_list(
response = row["response"] or ""
try:
decoded = json.loads(response)
new_items = []
if (
isinstance(decoded, dict)
and (data_key in decoded)
and all(isinstance(item, dict) for item in decoded[data_key])
):
for item in decoded[data_key]:
new_items.append(item)
new_items = list(decoded[data_key])
else:
new_items.append(decoded)
new_items = [decoded]
if data_ids:
for item in new_items:
item[find_unused_key(item, "response_id")] = row["id"]
@@ -2058,7 +2047,7 @@ def logs_list(
while "`" * num_backticks in value:
num_backticks += 1
fence = "`" * num_backticks
return textwrap.indent("{}\n{}\n{}".format(fence, value, fence), " ")
return textwrap.indent(f"{fence}\n{value}\n{fence}", " ")
def _inline_code(value):
num_backticks = 1
@@ -2066,21 +2055,19 @@ def logs_list(
num_backticks += 1
delimiter = "`" * num_backticks
if value.startswith("`") or value.endswith("`"):
return "{} {} {}".format(delimiter, value, delimiter)
return "{}{}{}".format(delimiter, value, delimiter)
return f"{delimiter} {value} {delimiter}"
return f"{delimiter}{value}{delimiter}"
def _format_tool_call_arguments(arguments):
if not isinstance(arguments, dict) or not arguments:
return " Arguments: {}".format(_inline_code(json.dumps(arguments)))
return f" Arguments: {_inline_code(json.dumps(arguments))}"
lines = []
for key, value in arguments.items():
if isinstance(value, str):
lines.append(" {}:".format(key))
lines.append(f" {key}:")
lines.append(_fenced_block(value))
else:
lines.append(
" {}: {}".format(key, _inline_code(json.dumps(value)))
)
lines.append(f" {key}: {_inline_code(json.dumps(value))}")
return "\n".join(lines)
def _token_usage_markdown(input_tokens, output_tokens, token_details):
@@ -2088,7 +2075,7 @@ def logs_list(
if token_details:
details = _inline_code(json.dumps(token_details))
if usage:
return "{}, {}".format(usage, details)
return f"{usage}, {details}"
return details
return usage
@@ -2207,9 +2194,9 @@ def logs_list(
options = json.loads(options)
if options:
options_text = "\n".join(
"- {}: {}".format(key, value) for key, value in options.items()
f"- {key}: {value}" for key, value in options.items()
)
click.echo("\n## Options\n\n{}".format(options_text))
click.echo(f"\n## Options\n\n{options_text}")
if row["system"] != current_system:
if row["system"] is not None:
click.echo("\n## System\n\n{}".format(row["system"]))
@@ -2255,7 +2242,7 @@ def logs_list(
desc += attachment["url"]
elif attachment.get("content"):
desc += f"<{attachment['content_length']:,} bytes>"
attachments += "\n - {}".format(desc)
attachments += f"\n - {desc}"
click.echo(
"- **{}**: `{}`<br>\n{}{}{}".format(
tool_result["name"],
@@ -2300,7 +2287,7 @@ def logs_list(
if row["schema_json"]:
try:
parsed = json.loads(response)
response = "```json\n{}\n```".format(json.dumps(parsed, indent=2))
response = f"```json\n{json.dumps(parsed, indent=2)}\n```"
except ValueError:
pass
if row.get("reasoning"):
@@ -2318,7 +2305,7 @@ def logs_list(
)
click.echo("")
if response:
click.echo("{}\n".format(response))
click.echo(f"{response}\n")
if usage:
token_usage = _token_usage_markdown(
row["input_tokens"],
@@ -2326,7 +2313,7 @@ def logs_list(
json.loads(row["token_details"]) if row["token_details"] else None,
)
if token_usage:
click.echo("## Token usage\n\n{}\n".format(token_usage))
click.echo(f"## Token usage\n\n{token_usage}\n")
@cli.group(
@@ -2400,7 +2387,7 @@ def render_model_with_aliases(
initial_indent=" ",
subsequent_indent=" ",
)
output += "\n Attachment types:\n{}".format(wrapper.fill(attachment_types))
output += f"\n Attachment types:\n{wrapper.fill(attachment_types)}"
features = (
[]
+ (["streaming"] if model.can_stream else [])
@@ -2410,14 +2397,14 @@ def render_model_with_aliases(
)
if options and features:
output += "\n Features:\n{}".format(
"\n".join(" - {}".format(feature) for feature in features)
"\n".join(f" - {feature}" for feature in features)
)
if options and hasattr(model, "needs_key") and model.needs_key:
output += "\n Keys:"
if hasattr(model, "needs_key") and model.needs_key:
output += "\n key: {}".format(model.needs_key)
output += f"\n key: {model.needs_key}"
if hasattr(model, "key_env_var") and model.key_env_var:
output += "\n env_var: {}".format(model.key_env_var)
output += f"\n env_var: {model.key_env_var}"
return output
@@ -2430,7 +2417,7 @@ def render_model_with_options(model_id, *, async_=False):
async_=async_,
models_that_have_shown_options=set(),
)
raise click.ClickException("'{}' is not a known model".format(model_id))
raise click.ClickException(f"'{model_id}' is not a known model")
@models.command(name="list")
@@ -2453,13 +2440,11 @@ def models_list(options, async_, schemas, tools, query, model_ids):
for model_with_aliases in get_models_with_aliases():
if async_ and not model_with_aliases.async_model:
continue
if query:
# Only show models where every provided query string matches
if not all(model_with_aliases.matches(q) for q in query):
continue
if model_ids:
if not model_matches_id_or_alias(model_with_aliases, model_ids):
continue
# Only show models where every provided query string matches
if query and not all(model_with_aliases.matches(q) for q in query):
continue
if model_ids and not model_matches_id_or_alias(model_with_aliases, model_ids):
continue
if schemas and not model_with_aliases.model.supports_schema:
continue
if tools and not model_with_aliases.model.supports_tools:
@@ -2488,7 +2473,7 @@ def models_default(model):
model = get_model(model)
set_default_model(model.model_id)
except KeyError:
raise click.ClickException("Unknown model: {}".format(model))
raise click.ClickException(f"Unknown model: {model}")
@cli.group(
@@ -2541,7 +2526,7 @@ def templates_show(name):
raise click.ClickException(f"Template '{name}' not found or invalid")
click.echo(
yaml.dump(
dict((k, v) for k, v in template.model_dump().items() if v is not None),
{k: v for k, v in template.model_dump().items() if v is not None},
indent=4,
default_flow_style=False,
)
@@ -2621,7 +2606,7 @@ def schemas_list(path, database, queries, full, json_, nl):
path = database
path = pathlib.Path(path or logs_db_path())
if not path.exists():
raise click.ClickException("No log database found at {}".format(path))
raise click.ClickException(f"No log database found at {path}")
db = sqlite_utils.Database(path)
migrate(db)
@@ -2630,9 +2615,9 @@ def schemas_list(path, database, queries, full, json_, nl):
if queries:
where_bits = ["schemas.content like ?" for _ in queries]
where_sql += " where {}".format(" and ".join(where_bits))
params.extend("%{}%".format(q) for q in queries)
params.extend(f"%{q}%" for q in queries)
sql = """
sql = f"""
select
schemas.id,
schemas.content,
@@ -2641,9 +2626,9 @@ def schemas_list(path, database, queries, full, json_, nl):
from schemas
join responses
on responses.schema_id = schemas.id
{} group by responses.schema_id
{where_sql} group by responses.schema_id
order by recently_used
""".format(where_sql)
"""
rows = db.query(sql, params)
if json_ or nl:
@@ -2697,7 +2682,7 @@ def schemas_show(schema_id, path, database):
path = database
path = pathlib.Path(path or logs_db_path())
if not path.exists():
raise click.ClickException("No log database found at {}".format(path))
raise click.ClickException(f"No log database found at {path}")
db = sqlite_utils.Database(path)
migrate(db)
@@ -2816,7 +2801,7 @@ def tools_list(tool_defs, json_, python_tools):
"{}{}{}\n".format(
tool.name,
sig,
" (plugin: {})".format(tool.plugin) if tool.plugin else "",
f" (plugin: {tool.plugin})" if tool.plugin else "",
)
)
if tool.description:
@@ -2829,12 +2814,7 @@ def tools_list(tool_defs, json_, python_tools):
.replace("(self, ", "(")
.replace("(self)", "()")
)
click.echo(
" {}{}\n".format(
tool.name,
sig,
)
)
click.echo(f" {tool.name}{sig}\n")
if tool.description:
click.echo(textwrap.indent(tool.description.strip(), " ") + "\n")
@@ -2974,12 +2954,10 @@ def fragments_list(queries, aliases, json_):
db = sqlite_utils.Database(logs_db_path())
migrate(db)
params = {}
param_count = 0
where_bits = []
if aliases:
where_bits.append("fragment_aliases.alias is not null")
for q in queries:
param_count += 1
for param_count, q in enumerate(queries, start=1):
p = f"p{param_count}"
params[p] = q
where_bits.append(f"""
@@ -2990,7 +2968,7 @@ def fragments_list(queries, aliases, json_):
where = "\n and\n ".join(where_bits)
if where:
where = " where " + where
sql = """
sql = f"""
select
fragments.hash,
json_group_array(fragment_aliases.alias) filter (
@@ -3008,7 +2986,7 @@ def fragments_list(queries, aliases, json_):
group by
fragments.id, fragments.hash, fragments.content, fragments.datetime_utc, fragments.source
order by fragments.datetime_utc
""".format(where=where)
"""
results = list(db.query(sql, params))
for result in results:
result["aliases"] = json.loads(result["aliases"])
@@ -3417,11 +3395,8 @@ def embed_multi(
if not input_path and not sql and not files:
raise click.UsageError("Either --sql or input path or --files is required")
if files:
if input_path or sql or format:
raise click.UsageError(
"Cannot use --files with --sql, input path or --format"
)
if files and (input_path or sql or format):
raise click.UsageError("Cannot use --files with --sql, input path or --format")
if database:
db = sqlite_utils.Database(database)
@@ -3473,7 +3448,7 @@ def embed_multi(
if content is None:
# Log to stderr
click.echo(
"Could not decode text in file {}".format(path),
f"Could not decode text in file {path}",
err=True,
)
else:
@@ -3483,7 +3458,7 @@ def embed_multi(
rows = iterate_files()
elif sql:
rows = db.query(sql)
count_sql = "select count(*) as c from ({})".format(sql)
count_sql = f"select count(*) as c from ({sql})"
expected_length = next(db.query(count_sql))["c"]
else:
@@ -3498,11 +3473,15 @@ def embed_multi(
for _ in load_rows(fp):
expected_length += 1
rows = load_rows(
open(input_path, "rb")
if input_path != "-"
else io.BufferedReader(sys.stdin.buffer)
)
if input_path != "-":
def rows_from_input():
with open(input_path, "rb") as fp:
yield from load_rows(fp)
rows = rows_from_input()
else:
rows = load_rows(io.BufferedReader(sys.stdin.buffer))
except json.JSONDecodeError as ex:
raise click.ClickException(str(ex))
@@ -3510,11 +3489,11 @@ def embed_multi(
rows, label="Embedding", show_percent=True, length=expected_length
) as rows:
def tuples() -> Iterable[Tuple[str, Union[bytes, str]]]:
def tuples() -> Iterable[tuple[str, bytes | str]]:
for row in rows:
values = list(row.values())
id: str = prefix + str(values[0])
content: Optional[Union[bytes, str]] = None
content: bytes | str | None = None
if binary:
content = cast(bytes, values[1])
else:
@@ -3633,9 +3612,8 @@ def embed_models_list(query):
"List available embedding models"
output = []
for model_with_aliases in get_embedding_models_with_aliases():
if query:
if not all(model_with_aliases.matches(q) for q in query):
continue
if query and not all(model_with_aliases.matches(q) for q in query):
continue
s = str(model_with_aliases.model)
if model_with_aliases.aliases:
s += " (aliases: {})".format(", ".join(model_with_aliases.aliases))
@@ -3665,7 +3643,7 @@ def embed_models_default(model, remove_default):
model = get_embedding_model(model)
set_default_embedding_model(model.model_id)
except KeyError:
raise click.ClickException("Unknown embedding model: {}".format(model))
raise click.ClickException(f"Unknown embedding model: {model}")
@cli.group(
@@ -3697,7 +3675,7 @@ def embed_db_collections(database, json_):
database = database or (user_dir() / "embeddings.db")
db = sqlite_utils.Database(str(database))
if not db["collections"].exists():
raise click.ClickException("No collections table found in {}".format(database))
raise click.ClickException(f"No collections table found in {database}")
rows = db.query("""
select
collections.name,
@@ -3938,7 +3916,7 @@ def _human_readable_size(size_bytes):
size_bytes /= 1024.0
i += 1
return "{:.2f}{}".format(size_bytes, size_name[i])
return f"{size_bytes:.2f}{size_name[i]}"
def logs_on():
@@ -4048,7 +4026,7 @@ def _parse_yaml_template(name, content):
try:
loaded = yaml.safe_load(content)
except yaml.YAMLError as ex:
raise LoadTemplateError("Invalid YAML: {}".format(str(ex)))
raise LoadTemplateError(f"Invalid YAML: {ex!s}")
if isinstance(loaded, str):
return Template(name=name, prompt=loaded)
loaded["name"] = name
@@ -4062,12 +4040,12 @@ def _parse_yaml_template(name, content):
def load_template(name: str) -> Template:
"Load template, or raise LoadTemplateError(msg)"
if name.startswith("https://") or name.startswith("http://"):
if name.startswith(("https://", "http://")):
response = httpx.get(name)
try:
response.raise_for_status()
except httpx.HTTPStatusError as ex:
raise LoadTemplateError("Could not load template {}: {}".format(name, ex))
raise LoadTemplateError(f"Could not load template {name}: {ex}")
return _parse_yaml_template(name, response.text)
potential_path = pathlib.Path(name)
@@ -4076,12 +4054,12 @@ def load_template(name: str) -> Template:
prefix, rest = name.split(":", 1)
loaders = get_template_loaders()
if prefix not in loaders:
raise LoadTemplateError("Unknown template prefix: {}".format(prefix))
raise LoadTemplateError(f"Unknown template prefix: {prefix}")
loader = loaders[prefix]
try:
return loader(rest)
except Exception as ex:
raise LoadTemplateError("Could not load template {}: {}".format(name, ex))
except Exception as ex: # noqa: BLE001
raise LoadTemplateError(f"Could not load template {name}: {ex}")
# Try local file
if potential_path.exists():
@@ -4098,7 +4076,7 @@ def load_template(name: str) -> Template:
return template_obj
def _tools_from_code(code_or_path: str) -> List[Tool]:
def _tools_from_code(code_or_path: str) -> list[Tool]:
"""
Treat all Python functions in the code as tools
"""
@@ -4106,13 +4084,13 @@ def _tools_from_code(code_or_path: str) -> List[Tool]:
try:
code_or_path = pathlib.Path(code_or_path).read_text()
except FileNotFoundError:
raise click.ClickException("File not found: {}".format(code_or_path))
namespace: Dict[str, Any] = {}
raise click.ClickException(f"File not found: {code_or_path}")
namespace: dict[str, Any] = {}
tools = []
try:
exec(code_or_path, namespace)
exec(code_or_path, namespace) # noqa: S102
except SyntaxError as ex:
raise click.ClickException("Error in --functions definition: {}".format(ex))
raise click.ClickException(f"Error in --functions definition: {ex}")
# Register all callables in the locals dict:
for name, value in namespace.items():
if callable(value) and not name.startswith("_"):
@@ -4123,7 +4101,7 @@ def _tools_from_code(code_or_path: str) -> List[Tool]:
def _debug_tool_call(_, tool_call, tool_result):
click.echo(
click.style(
"\nTool call: {}({})".format(tool_call.name, tool_call.arguments),
f"\nTool call: {tool_call.name}({tool_call.arguments})",
fg="yellow",
bold=True,
),
@@ -4134,7 +4112,7 @@ def _debug_tool_call(_, tool_call, tool_result):
if tool_result.attachments:
attachments += "\nAttachments:\n"
for attachment in tool_result.attachments:
attachments += f" {repr(attachment)}\n"
attachments += f" {attachment!r}\n"
try:
output = json.dumps(json.loads(tool_result.output), indent=2)
@@ -4152,7 +4130,7 @@ def _debug_tool_call(_, tool_call, tool_result):
if tool_result.exception:
click.echo(
click.style(
" Exception: {}".format(tool_result.exception),
f" Exception: {tool_result.exception}",
fg="red",
bold=True,
),
@@ -4163,7 +4141,7 @@ def _debug_tool_call(_, tool_call, tool_result):
def _approve_tool_call(_, tool_call):
click.echo(
click.style(
"Tool call: {}({})".format(tool_call.name, tool_call.arguments),
f"Tool call: {tool_call.name}({tool_call.arguments})",
fg="yellow",
bold=True,
),
@@ -4174,18 +4152,16 @@ def _approve_tool_call(_, tool_call):
def _gather_tools(
tool_specs: List[str], python_tools: List[str]
) -> List[Union[Tool, Type[Toolbox]]]:
tools: List[Union[Tool, Type[Toolbox]]] = []
tool_specs: list[str], python_tools: list[str]
) -> list[Tool | type[Toolbox]]:
tools: list[Tool | type[Toolbox]] = []
if python_tools:
for code_or_path in python_tools:
tools.extend(_tools_from_code(code_or_path))
registered_tools = get_tools()
registered_classes = dict(
(key, value)
for key, value in registered_tools.items()
if inspect.isclass(value)
)
registered_classes = {
key: value for key, value in registered_tools.items() if inspect.isclass(value)
}
bad_tools = [
tool for tool in tool_specs if tool.split("(")[0] not in registered_tools
]
+87 -96
View File
@@ -1,3 +1,17 @@
import datetime
import json
import os
from collections.abc import AsyncGenerator, Iterable, Iterator
from enum import Enum
from typing import Any, cast
import click
import httpx
import openai
import yaml
from pydantic import Field, create_model, field_validator
import llm
from llm import (
AsyncConversation,
AsyncKeyModel,
@@ -9,36 +23,13 @@ from llm import (
Response,
hookimpl,
)
import llm
from llm.parts import StreamEvent
from llm.utils import (
dicts_to_table_string,
remove_dict_none_values,
logging_client,
remove_dict_none_values,
simplify_usage_dict,
)
import click
import datetime
from enum import Enum
import httpx
import openai
import os
from pydantic import create_model, field_validator, Field
from typing import (
Any,
AsyncGenerator,
cast,
Dict,
List,
Iterable,
Iterator,
Optional,
Union,
)
import json
import yaml
@hookimpl
@@ -396,7 +387,7 @@ class OpenAIEmbeddingModel(EmbeddingModel):
self.openai_model_id = openai_model_id
self.dimensions = dimensions
def embed_batch(self, items: Iterable[Union[str, bytes]]) -> Iterator[List[float]]:
def embed_batch(self, items: Iterable[str | bytes]) -> Iterator[list[float]]:
kwargs = {
"input": items,
"model": self.openai_model_id,
@@ -447,12 +438,12 @@ def register_commands(cli):
"created": created_str,
}
)
done = dicts_to_table_string("id owned_by created".split(), to_print)
done = dicts_to_table_string(["id", "owned_by", "created"], to_print)
print("\n".join(done))
class SharedOptions(llm.Options):
temperature: Optional[float] = Field(
temperature: float | None = Field(
description=(
"What sampling temperature to use, between 0 and 2. Higher values like "
"0.8 will make the output more random, while lower values like 0.2 will "
@@ -462,10 +453,10 @@ class SharedOptions(llm.Options):
le=2,
default=None,
)
max_tokens: Optional[int] = Field(
max_tokens: int | None = Field(
description="Maximum number of tokens to generate.", default=None
)
top_p: Optional[float] = Field(
top_p: float | None = Field(
description=(
"An alternative to sampling with temperature, called nucleus sampling, "
"where the model considers the results of the tokens with top_p "
@@ -477,7 +468,7 @@ class SharedOptions(llm.Options):
le=1,
default=None,
)
frequency_penalty: Optional[float] = Field(
frequency_penalty: float | None = Field(
description=(
"Number between -2.0 and 2.0. Positive values penalize new tokens based "
"on their existing frequency in the text so far, decreasing the model's "
@@ -487,7 +478,7 @@ class SharedOptions(llm.Options):
le=2,
default=None,
)
presence_penalty: Optional[float] = Field(
presence_penalty: float | None = Field(
description=(
"Number between -2.0 and 2.0. Positive values penalize new tokens based "
"on whether they appear in the text so far, increasing the model's "
@@ -497,18 +488,18 @@ class SharedOptions(llm.Options):
le=2,
default=None,
)
stop: Optional[str] = Field(
stop: str | None = Field(
description=("A string where the API will stop generating further tokens."),
default=None,
)
logit_bias: Optional[Union[dict, str]] = Field(
logit_bias: dict | str | None = Field(
description=(
"Modify the likelihood of specified tokens appearing in the completion. "
'Pass a JSON string like \'{"1712":-100, "892":-100, "1489":-100}\''
),
default=None,
)
seed: Optional[int] = Field(
seed: int | None = Field(
description="Integer seed to attempt to sample deterministically",
default=None,
)
@@ -584,7 +575,7 @@ def build_options_class(
):
fields = {
"json_object": (
Optional[bool],
bool | None,
Field(
description="Output a valid JSON object {...}. Prompt must mention JSON.",
default=None,
@@ -593,7 +584,7 @@ def build_options_class(
}
if chat_completions:
fields["chat_completions"] = (
Optional[bool],
bool | None,
Field(
description=(
"Force the use of the older /v1/chat/completions endpoint "
@@ -609,7 +600,7 @@ def build_options_class(
)
image_detail_values = enum_values_sentence(image_detail_enum)
fields["image_detail"] = (
Optional[image_detail_enum],
image_detail_enum | None,
Field(
description=(
"Controls the detail level for image attachments. Supported values are "
@@ -620,7 +611,7 @@ def build_options_class(
)
if reasoning:
fields["reasoning_effort"] = (
Optional[ReasoningEffortEnum],
ReasoningEffortEnum | None,
Field(
description=(
"Constraints effort on reasoning for reasoning models. Currently "
@@ -633,7 +624,7 @@ def build_options_class(
)
if verbosity:
fields["verbosity"] = (
Optional[VerbosityEnum],
VerbosityEnum | None,
Field(
description=(
"Controls how verbose the model's response should be. Supported "
@@ -741,7 +732,7 @@ class _Shared:
)
def __str__(self) -> str:
return "OpenAI Chat: {}".format(self.model_id)
return f"OpenAI Chat: {self.model_id}"
def _append_llm_message(self, out, message, current_system, image_detail=None):
"""Translate one llm.Message into one (or more) OpenAI message
@@ -827,10 +818,10 @@ class _Shared:
def build_messages(self, prompt, conversation, image_detail=None):
"""Translate prompt.messages into OpenAI's wire format."""
messages: List[Dict[str, Any]] = []
messages: list[dict[str, Any]] = []
if image_detail is not None:
image_detail = image_detail.value
current_system: Optional[str] = None
current_system: str | None = None
for msg in prompt.messages:
current_system = self._append_llm_message(
messages, msg, current_system, image_detail=image_detail
@@ -915,9 +906,9 @@ class Chat(_Shared, KeyModel):
prompt: Prompt,
stream: bool,
response: Response,
conversation: Optional[Conversation] = None,
key: Optional[str] = None,
) -> Iterator[Union[str, StreamEvent]]:
conversation: Conversation | None = None,
key: str | None = None,
) -> Iterator[str | StreamEvent]:
if prompt.system and not self.allows_system_prompt:
raise NotImplementedError("Model does not support system prompts")
messages = self.build_messages(
@@ -1033,9 +1024,9 @@ class AsyncChat(_Shared, AsyncKeyModel):
prompt: Prompt,
stream: bool,
response: AsyncResponse,
conversation: Optional[AsyncConversation] = None,
key: Optional[str] = None,
) -> AsyncGenerator[Union[str, StreamEvent], None]:
conversation: AsyncConversation | None = None,
key: str | None = None,
) -> AsyncGenerator[str | StreamEvent, None]:
if prompt.system and not self.allows_system_prompt:
raise NotImplementedError("Model does not support system prompts")
messages = self.build_messages(
@@ -1166,30 +1157,30 @@ class _SharedResponses(_Shared):
"""Mixin that translates llm.Prompt into Responses API parameters."""
def __str__(self) -> str:
return "OpenAI Responses: {}".format(self.model_id)
return f"OpenAI Responses: {self.model_id}"
def _delegate_chat_kwargs(self):
"""Return constructor kwargs that mirror this Responses model so we
can build a sibling Chat / AsyncChat instance for the
``-o chat_completions 1`` opt-out path."""
return dict(
model_id=self.model_id,
key=self.key,
model_name=self.model_name,
api_base=self.api_base,
api_type=self.api_type,
api_version=self.api_version,
api_engine=self.api_engine,
headers=self.headers,
can_stream=self.can_stream,
vision=self.vision,
reasoning=self._reasoning,
verbosity=self._verbosity,
image_detail_original=self._image_detail_original,
supports_schema=self.supports_schema,
supports_tools=self.supports_tools,
allows_system_prompt=self.allows_system_prompt,
)
return {
"model_id": self.model_id,
"key": self.key,
"model_name": self.model_name,
"api_base": self.api_base,
"api_type": self.api_type,
"api_version": self.api_version,
"api_engine": self.api_engine,
"headers": self.headers,
"can_stream": self.can_stream,
"vision": self.vision,
"reasoning": self._reasoning,
"verbosity": self._verbosity,
"image_detail_original": self._image_detail_original,
"supports_schema": self.supports_schema,
"supports_tools": self.supports_tools,
"allows_system_prompt": self.allows_system_prompt,
}
def _build_responses_input(self, prompt, image_detail=None):
"""Translate prompt.messages into a (input_items, instructions) tuple
@@ -1207,8 +1198,8 @@ class _SharedResponses(_Shared):
ToolResultPart,
)
items: List[Dict[str, Any]] = []
instructions: Optional[str] = None
items: list[dict[str, Any]] = []
instructions: str | None = None
for msg in prompt.messages:
if msg.role == "system":
@@ -1217,11 +1208,11 @@ class _SharedResponses(_Shared):
instructions = text
continue
text_bits: List[str] = []
attachment_items: List[Dict[str, Any]] = []
tool_call_items: List[Dict[str, Any]] = []
tool_result_items: List[Dict[str, Any]] = []
reasoning_items: List[Dict[str, Any]] = []
text_bits: list[str] = []
attachment_items: list[dict[str, Any]] = []
tool_call_items: list[dict[str, Any]] = []
tool_result_items: list[dict[str, Any]] = []
reasoning_items: list[dict[str, Any]] = []
for part in msg.parts:
if isinstance(part, TextPart):
@@ -1256,7 +1247,7 @@ class _SharedResponses(_Shared):
if enc or rid:
# Round-trip a previous reasoning item so the model
# can pick up where it left off mid-tool-call.
item: Dict[str, Any] = {"type": "reasoning"}
item: dict[str, Any] = {"type": "reasoning"}
if rid:
item["id"] = rid
if enc:
@@ -1277,7 +1268,7 @@ class _SharedResponses(_Shared):
if msg.role == "user":
if attachment_items:
content: List[Dict[str, Any]] = []
content: list[dict[str, Any]] = []
if text_bits:
content.append(
{"type": "input_text", "text": "".join(text_bits)}
@@ -1308,7 +1299,7 @@ class _SharedResponses(_Shared):
top_p = opts.pop("top_p", None)
seed = opts.pop("seed", None)
kwargs: Dict[str, Any] = {}
kwargs: dict[str, Any] = {}
if max_tokens is None and self.default_max_tokens is not None:
max_tokens = self.default_max_tokens
if max_tokens is not None:
@@ -1328,7 +1319,7 @@ class _SharedResponses(_Shared):
if reasoning:
kwargs["reasoning"] = reasoning
text: Dict[str, Any] = {}
text: dict[str, Any] = {}
if verbosity:
text["verbosity"] = verbosity
if prompt.options.json_object:
@@ -1400,7 +1391,7 @@ class _SharedResponses(_Shared):
enc = getattr(item, "encrypted_content", None)
summary = getattr(item, "summary", None)
text = self._reasoning_text_from_item(item) if include_text else ""
meta: Dict[str, Any] = {}
meta: dict[str, Any] = {}
if rid:
meta["id"] = rid
if enc:
@@ -1413,7 +1404,7 @@ class _SharedResponses(_Shared):
s.model_dump() if hasattr(s, "model_dump") else dict(s)
for s in summary
]
except Exception:
except Exception: # noqa: BLE001
meta["summary"] = list(summary)
return StreamEvent(
type="reasoning",
@@ -1484,9 +1475,9 @@ class Responses(_SharedResponses, KeyModel):
prompt: Prompt,
stream: bool,
response: Response,
conversation: Optional[Conversation] = None,
key: Optional[str] = None,
) -> Iterator[Union[str, StreamEvent]]:
conversation: Conversation | None = None,
key: str | None = None,
) -> Iterator[str | StreamEvent]:
if getattr(prompt.options, "chat_completions", None):
chat = Chat(**self._delegate_chat_kwargs())
yield from chat.execute(prompt, stream, response, conversation, key)
@@ -1518,8 +1509,8 @@ class Responses(_SharedResponses, KeyModel):
stream=True,
**kwargs,
)
tool_call_meta: Dict[str, Dict[str, str]] = {}
final_response_dict: Optional[Dict[str, Any]] = None
tool_call_meta: dict[str, dict[str, str]] = {}
final_response_dict: dict[str, Any] | None = None
reasoning_items_with_streamed_text = set()
for event in stream_obj:
etype = getattr(event, "type", None)
@@ -1711,9 +1702,9 @@ class AsyncResponses(_SharedResponses, AsyncKeyModel):
prompt: Prompt,
stream: bool,
response: AsyncResponse,
conversation: Optional[AsyncConversation] = None,
key: Optional[str] = None,
) -> AsyncGenerator[Union[str, StreamEvent], None]:
conversation: AsyncConversation | None = None,
key: str | None = None,
) -> AsyncGenerator[str | StreamEvent, None]:
if getattr(prompt.options, "chat_completions", None):
chat = AsyncChat(**self._delegate_chat_kwargs())
async for event in chat.execute(
@@ -1748,8 +1739,8 @@ class AsyncResponses(_SharedResponses, AsyncKeyModel):
stream=True,
**kwargs,
)
tool_call_meta: Dict[str, Dict[str, str]] = {}
final_response_dict: Optional[Dict[str, Any]] = None
tool_call_meta: dict[str, dict[str, str]] = {}
final_response_dict: dict[str, Any] | None = None
reasoning_items_with_streamed_text = set()
async for event in stream_obj:
etype = getattr(event, "type", None)
@@ -1881,7 +1872,7 @@ class AsyncResponses(_SharedResponses, AsyncKeyModel):
class Completion(Chat):
class Options(SharedOptions):
logprobs: Optional[int] = Field(
logprobs: int | None = Field(
description="Include the log probabilities of most likely N per token",
default=None,
le=5,
@@ -1892,16 +1883,16 @@ class Completion(Chat):
self.default_max_tokens = default_max_tokens
def __str__(self) -> str:
return "OpenAI Completion: {}".format(self.model_id)
return f"OpenAI Completion: {self.model_id}"
def execute(
self,
prompt: Prompt,
stream: bool,
response: Response,
conversation: Optional[Conversation] = None,
key: Optional[str] = None,
) -> Iterator[Union[str, StreamEvent]]:
conversation: Conversation | None = None,
key: str | None = None,
) -> Iterator[str | StreamEvent]:
if prompt.system:
raise NotImplementedError(
"System prompts are not supported for OpenAI completion models"
@@ -1949,7 +1940,7 @@ def not_nulls(data) -> dict:
return {key: value for key, value in data if value is not None}
def combine_chunks(chunks: List) -> dict:
def combine_chunks(chunks: list) -> dict:
content = ""
role = None
finish_reason = None
+28 -25
View File
@@ -1,21 +1,24 @@
from .models import EmbeddingModel
from .embeddings_migrations import embeddings_migrations
from dataclasses import dataclass
import hashlib
from itertools import islice
import json
import time
from collections.abc import Iterable
from dataclasses import dataclass
from itertools import islice
from typing import Any, cast
from sqlite_utils import Database
from sqlite_utils.db import Table
import time
from typing import cast, Any, Dict, Iterable, List, Optional, Tuple, Union
from .embeddings_migrations import embeddings_migrations
from .models import EmbeddingModel
@dataclass
class Entry:
id: str
score: Optional[float]
content: Optional[str] = None
metadata: Optional[Dict[str, Any]] = None
score: float | None
content: str | None = None
metadata: dict[str, Any] | None = None
class Collection:
@@ -25,10 +28,10 @@ class Collection:
def __init__(
self,
name: str,
db: Optional[Database] = None,
db: Database | None = None,
*,
model: Optional[EmbeddingModel] = None,
model_id: Optional[str] = None,
model: EmbeddingModel | None = None,
model_id: str | None = None,
create: bool = True,
) -> None:
"""
@@ -115,8 +118,8 @@ class Collection:
def embed(
self,
id: str,
value: Union[str, bytes],
metadata: Optional[Dict[str, Any]] = None,
value: str | bytes,
metadata: dict[str, Any] | None = None,
store: bool = False,
) -> None:
"""
@@ -152,7 +155,7 @@ class Collection:
def embed_multi(
self,
entries: Iterable[Tuple[str, Union[str, bytes]]],
entries: Iterable[tuple[str, str | bytes]],
store: bool = False,
batch_size: int = 100,
) -> None:
@@ -172,7 +175,7 @@ class Collection:
def embed_multi_with_metadata(
self,
entries: Iterable[Tuple[str, Union[str, bytes], Optional[Dict[str, Any]]]],
entries: Iterable[tuple[str, str | bytes, dict[str, Any] | None]],
store: bool = False,
batch_size: int = 100,
) -> None:
@@ -237,11 +240,11 @@ class Collection:
def similar_by_vector(
self,
vector: List[float],
vector: list[float],
number: int = 10,
skip_id: Optional[str] = None,
prefix: Optional[str] = None,
) -> List[Entry]:
skip_id: str | None = None,
prefix: str | None = None,
) -> list[Entry]:
"""
Find similar items in the collection by a given vector.
@@ -295,8 +298,8 @@ class Collection:
]
def similar_by_id(
self, id: str, number: int = 10, prefix: Optional[str] = None
) -> List[Entry]:
self, id: str, number: int = 10, prefix: str | None = None
) -> list[Entry]:
"""
Find similar items in the collection by a given ID.
@@ -324,8 +327,8 @@ class Collection:
)
def similar(
self, value: Union[str, bytes], number: int = 10, prefix: Optional[str] = None
) -> List[Entry]:
self, value: str | bytes, number: int = 10, prefix: str | None = None
) -> list[Entry]:
"""
Find similar items in the collection by a given value.
@@ -360,7 +363,7 @@ class Collection:
self.db.execute("delete from collections where id = ?", [self.id])
@staticmethod
def content_hash(input: Union[str, bytes]) -> bytes:
def content_hash(input: str | bytes) -> bytes:
"Hash content for deduplication. Override to change hashing behavior."
if isinstance(input, str):
input = input.encode("utf8")
+2 -1
View File
@@ -1,7 +1,8 @@
from sqlite_migrate import Migrations
import hashlib
import time
from sqlite_migrate import Migrations
embeddings_migrations = Migrations("llm.embeddings")
+1 -2
View File
@@ -1,5 +1,4 @@
from pluggy import HookimplMarker
from pluggy import HookspecMarker
from pluggy import HookimplMarker, HookspecMarker
hookspec = HookspecMarker("llm")
hookimpl = HookimplMarker("llm")
+2 -2
View File
@@ -1,7 +1,7 @@
import datetime
from typing import Callable, List
from collections.abc import Callable
MIGRATIONS: List[Callable] = []
MIGRATIONS: list[Callable] = []
migration = MIGRATIONS.append
+321 -329
View File
File diff suppressed because it is too large Load Diff
+32 -34
View File
@@ -13,7 +13,7 @@ content are equal.
import base64
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from typing import Any
from .models import Attachment
from .serialization import (
@@ -29,7 +29,7 @@ from .serialization import (
def _attachment_to_dict(att: Attachment) -> AttachmentDict:
d: Dict[str, Any] = {}
d: dict[str, Any] = {}
if att.type:
d["type"] = att.type
if att.url:
@@ -43,7 +43,7 @@ def _attachment_to_dict(att: Attachment) -> AttachmentDict:
def _attachment_from_dict(d: AttachmentDict) -> Attachment:
raw_content = d.get("content")
content_bytes: Optional[bytes] = None
content_bytes: bytes | None = None
if isinstance(raw_content, str):
content_bytes = base64.b64decode(raw_content)
return Attachment(
@@ -107,10 +107,10 @@ class Part:
@dataclass
class TextPart(Part):
text: str = ""
provider_metadata: Optional[Dict[str, Any]] = None
provider_metadata: dict[str, Any] | None = None
def to_dict(self) -> TextPartDict:
d: Dict[str, Any] = {"type": "text", "text": self.text}
d: dict[str, Any] = {"type": "text", "text": self.text}
if self.provider_metadata:
d["provider_metadata"] = self.provider_metadata
return d # type: ignore[return-value]
@@ -130,10 +130,10 @@ class ReasoningPart(Part):
text: str = ""
redacted: bool = False
provider_metadata: Optional[Dict[str, Any]] = None
provider_metadata: dict[str, Any] | None = None
def to_dict(self) -> ReasoningPartDict:
d: Dict[str, Any] = {"type": "reasoning", "text": self.text}
d: dict[str, Any] = {"type": "reasoning", "text": self.text}
if self.redacted:
d["redacted"] = True
if self.provider_metadata:
@@ -151,13 +151,13 @@ class ToolCallPart(Part):
"""
name: str = ""
arguments: Dict[str, Any] = field(default_factory=dict)
tool_call_id: Optional[str] = None
arguments: dict[str, Any] = field(default_factory=dict)
tool_call_id: str | None = None
server_executed: bool = False
provider_metadata: Optional[Dict[str, Any]] = None
provider_metadata: dict[str, Any] | None = None
def to_dict(self) -> ToolCallPartDict:
d: Dict[str, Any] = {
d: dict[str, Any] = {
"type": "tool_call",
"name": self.name,
"arguments": self.arguments,
@@ -177,14 +177,14 @@ class ToolResultPart(Part):
name: str = ""
output: str = ""
tool_call_id: Optional[str] = None
tool_call_id: str | None = None
server_executed: bool = False
attachments: List[Any] = field(default_factory=list)
exception: Optional[str] = None
provider_metadata: Optional[Dict[str, Any]] = None
attachments: list[Any] = field(default_factory=list)
exception: str | None = None
provider_metadata: dict[str, Any] | None = None
def to_dict(self) -> ToolResultPartDict:
d: Dict[str, Any] = {
d: dict[str, Any] = {
"type": "tool_result",
"name": self.name,
"output": self.output,
@@ -206,11 +206,11 @@ class ToolResultPart(Part):
class AttachmentPart(Part):
"""An inline attachment (image, audio, file)."""
attachment: Optional[Attachment] = None
provider_metadata: Optional[Dict[str, Any]] = None
attachment: Attachment | None = None
provider_metadata: dict[str, Any] | None = None
def to_dict(self) -> AttachmentPartDict:
d: Dict[str, Any] = {"type": "attachment"}
d: dict[str, Any] = {"type": "attachment"}
if self.attachment:
d["attachment"] = _attachment_to_dict(self.attachment)
if self.provider_metadata:
@@ -229,11 +229,11 @@ class Message:
"""
role: str
parts: List[Part] = field(default_factory=list)
provider_metadata: Optional[Dict[str, Any]] = None
parts: list[Part] = field(default_factory=list)
provider_metadata: dict[str, Any] | None = None
def to_dict(self) -> MessageDict:
d: Dict[str, Any] = {
d: dict[str, Any] = {
"role": self.role,
"parts": [p.to_dict() for p in self.parts],
}
@@ -250,13 +250,13 @@ class Message:
)
def normalize_parts(items: Any) -> List[Part]:
def normalize_parts(items: Any) -> list[Part]:
"""Normalize helper inputs to a list of Part objects.
Accepts str (→ TextPart), Attachment (→ AttachmentPart), Part
(passed through), or a list/tuple of those (flattened one level).
"""
out: List[Part] = []
out: list[Part] = []
for item in items:
if isinstance(item, Part):
out.append(item)
@@ -271,7 +271,7 @@ def normalize_parts(items: Any) -> List[Part]:
return out
def system(*items: Any, provider_metadata: Optional[Dict[str, Any]] = None) -> Message:
def system(*items: Any, provider_metadata: dict[str, Any] | None = None) -> Message:
"Build a Message with role='system'."
return Message(
role="system",
@@ -280,7 +280,7 @@ def system(*items: Any, provider_metadata: Optional[Dict[str, Any]] = None) -> M
)
def user(*items: Any, provider_metadata: Optional[Dict[str, Any]] = None) -> Message:
def user(*items: Any, provider_metadata: dict[str, Any] | None = None) -> Message:
"Build a Message with role='user'."
return Message(
role="user",
@@ -289,9 +289,7 @@ def user(*items: Any, provider_metadata: Optional[Dict[str, Any]] = None) -> Mes
)
def assistant(
*items: Any, provider_metadata: Optional[Dict[str, Any]] = None
) -> Message:
def assistant(*items: Any, provider_metadata: dict[str, Any] | None = None) -> Message:
"Build a Message with role='assistant'."
return Message(
role="assistant",
@@ -301,7 +299,7 @@ def assistant(
def tool_message(
*items: Any, provider_metadata: Optional[Dict[str, Any]] = None
*items: Any, provider_metadata: dict[str, Any] | None = None
) -> Message:
"Build a Message with role='tool' (typically wrapping ToolResultParts)."
return Message(
@@ -343,10 +341,10 @@ class StreamEvent:
type: str # "text" / "reasoning" / "tool_call_name" /
# "tool_call_args" / "tool_result"
chunk: str
part_index: Optional[int] = None
tool_call_id: Optional[str] = None
part_index: int | None = None
tool_call_id: str | None = None
server_executed: bool = False
tool_name: Optional[str] = None
tool_name: str | None = None
redacted: bool = False
provider_metadata: Optional[Dict[str, Any]] = None
provider_metadata: dict[str, Any] | None = None
message_index: int = 0
+4 -2
View File
@@ -1,8 +1,10 @@
import importlib
from importlib import metadata
import os
import pluggy
import sys
from importlib import metadata
import pluggy
from . import hookspecs
DEFAULT_PLUGINS = (
+21 -21
View File
@@ -32,7 +32,7 @@ keys may be absent from a serialized payload; required keys must
always be present.
"""
from typing import Any, Dict, List, Literal, Union
from typing import Any, Literal
# NotRequired moved to typing in 3.11; use typing_extensions for 3.10
# support. typing_extensions is a transitive dep via pydantic.
@@ -75,7 +75,7 @@ class AttachmentDict(TypedDict, total=False):
class TextPartDict(TypedDict):
type: Literal["text"]
text: str
provider_metadata: NotRequired[Dict[str, Any]]
provider_metadata: NotRequired[dict[str, Any]]
class ReasoningPartDict(TypedDict):
@@ -85,19 +85,19 @@ class ReasoningPartDict(TypedDict):
# reasoning (OpenAI GPT-5, Gemini without thoughts). The token
# total lives on response usage, not on the Part.
redacted: NotRequired[bool]
provider_metadata: NotRequired[Dict[str, Any]]
provider_metadata: NotRequired[dict[str, Any]]
class ToolCallPartDict(TypedDict):
type: Literal["tool_call"]
name: str
arguments: Dict[str, Any]
arguments: dict[str, Any]
tool_call_id: NotRequired[str]
# True for provider-executed calls (Anthropic web search, Gemini code
# execution). Adapters use this to restore provider-side blocks on
# the next turn.
server_executed: NotRequired[bool]
provider_metadata: NotRequired[Dict[str, Any]]
provider_metadata: NotRequired[dict[str, Any]]
class ToolResultPartDict(TypedDict):
@@ -107,23 +107,23 @@ class ToolResultPartDict(TypedDict):
tool_call_id: NotRequired[str]
server_executed: NotRequired[bool]
exception: NotRequired[str]
attachments: NotRequired[List[AttachmentDict]]
provider_metadata: NotRequired[Dict[str, Any]]
attachments: NotRequired[list[AttachmentDict]]
provider_metadata: NotRequired[dict[str, Any]]
class AttachmentPartDict(TypedDict):
type: Literal["attachment"]
attachment: NotRequired[AttachmentDict]
provider_metadata: NotRequired[Dict[str, Any]]
provider_metadata: NotRequired[dict[str, Any]]
PartDict = Union[
TextPartDict,
ReasoningPartDict,
ToolCallPartDict,
ToolResultPartDict,
AttachmentPartDict,
]
PartDict = (
TextPartDict
| ReasoningPartDict
| ToolCallPartDict
| ToolResultPartDict
| AttachmentPartDict
)
"""Discriminated union of Part dict shapes. Use with
``pydantic.TypeAdapter(PartDict)`` to validate / dispatch by ``type``.
"""
@@ -140,8 +140,8 @@ class MessageDict(TypedDict):
"""
role: str
parts: List[PartDict]
provider_metadata: NotRequired[Dict[str, Any]]
parts: list[PartDict]
provider_metadata: NotRequired[dict[str, Any]]
# ---- Response + nested shapes -----------------------------------------------
@@ -152,8 +152,8 @@ class PromptDict(TypedDict):
full input chain that was sent for this turn plus any options that
apply."""
messages: List[MessageDict]
options: NotRequired[Dict[str, Any]]
messages: list[MessageDict]
options: NotRequired[dict[str, Any]]
system: NotRequired[str]
@@ -163,7 +163,7 @@ class UsageDict(TypedDict, total=False):
input: int
output: int
details: Dict[str, Any]
details: dict[str, Any]
class ResponseDict(TypedDict):
@@ -174,7 +174,7 @@ class ResponseDict(TypedDict):
model: str
prompt: PromptDict
messages: List[MessageDict]
messages: list[MessageDict]
# Audit fields — present on a freshly-serialized response, optional
# on hand-constructed ones.
id: NotRequired[str]
+23 -22
View File
@@ -1,6 +1,7 @@
from pydantic import BaseModel, ConfigDict
import string
from typing import Optional, Any, Dict, List, Tuple
from typing import Any
from pydantic import BaseModel, ConfigDict
class AttachmentType(BaseModel):
@@ -12,20 +13,20 @@ class Template(BaseModel):
"""A reusable prompt template."""
name: str
prompt: Optional[str] = None
system: Optional[str] = None
attachments: Optional[List[str]] = None
attachment_types: Optional[List[AttachmentType]] = None
model: Optional[str] = None
defaults: Optional[Dict[str, Any]] = None
options: Optional[Dict[str, Any]] = None
extract: Optional[bool] = None # For extracting fenced code blocks
extract_last: Optional[bool] = None
schema_object: Optional[dict] = None
fragments: Optional[List[str]] = None
system_fragments: Optional[List[str]] = None
tools: Optional[List[str]] = None
functions: Optional[str] = None
prompt: str | None = None
system: str | None = None
attachments: list[str] | None = None
attachment_types: list[AttachmentType] | None = None
model: str | None = None
defaults: dict[str, Any] | None = None
options: dict[str, Any] | None = None
extract: bool | None = None # For extracting fenced code blocks
extract_last: bool | None = None
schema_object: dict | None = None
fragments: list[str] | None = None
system_fragments: list[str] | None = None
tools: list[str] | None = None
functions: str | None = None
model_config = ConfigDict(extra="forbid")
@@ -39,8 +40,8 @@ class Template(BaseModel):
self._functions_is_trusted = False
def evaluate(
self, input: str, params: Optional[Dict[str, Any]] = None
) -> Tuple[Optional[str], Optional[str]]:
self, input: str, params: dict[str, Any] | None = None
) -> tuple[str | None, str | None]:
"""Evaluate the template with the given input and parameters, returning (prompt, system)."""
params = params or {}
params["input"] = input
@@ -48,8 +49,8 @@ class Template(BaseModel):
for k, v in self.defaults.items():
if k not in params:
params[k] = v
prompt: Optional[str] = None
system: Optional[str] = None
prompt: str | None = None
system: str | None = None
if not self.prompt:
system = self.interpolate(self.system, params)
prompt = input
@@ -68,7 +69,7 @@ class Template(BaseModel):
return all_vars
@classmethod
def interpolate(cls, text: Optional[str], params: Dict[str, Any]) -> Optional[str]:
def interpolate(cls, text: str | None, params: dict[str, Any]) -> str | None:
"""Substitute template variables in text with values from params, raising MissingVariables if any are absent."""
if not text:
return text
@@ -83,7 +84,7 @@ class Template(BaseModel):
return string_template.substitute(**params)
@staticmethod
def extract_vars(string_template: string.Template) -> List[str]:
def extract_vars(string_template: string.Template) -> list[str]:
"""Extract and return the list of named variable identifiers from a string.Template."""
return [
match.group("named")
+2 -2
View File
@@ -1,6 +1,6 @@
import time
from datetime import datetime, timezone
from importlib.metadata import version
import time
def llm_version() -> str:
@@ -12,7 +12,7 @@ def llm_time() -> dict:
"Returns the current time, as local time and UTC"
# Get current times
utc_time = datetime.now(timezone.utc)
local_time = datetime.now()
local_time = datetime.now(timezone.utc).astimezone()
# Get timezone information
local_tz_name = time.tzname[time.localtime().tm_isdst]
+37 -30
View File
@@ -1,19 +1,18 @@
import click
import hashlib
import httpx
import itertools
import json
import pathlib
import puremagic
import re
import sqlite_utils
import textwrap
from typing import Any, List, Dict, Optional, Tuple, Type
import os
import pathlib
import re
import textwrap
import threading
import time
from typing import Final
from typing import Any, Final
import click
import httpx
import puremagic
import sqlite_utils
from ulid import ULID
MIME_TYPE_FIXES = {
@@ -34,7 +33,7 @@ class Fragment(str):
return hashlib.sha256(self.encode("utf-8")).hexdigest()
def mimetype_from_string(content) -> Optional[str]:
def mimetype_from_string(content) -> str | None:
try:
type_ = puremagic.from_string(content, mime=True)
return MIME_TYPE_FIXES.get(type_, type_)
@@ -42,7 +41,7 @@ def mimetype_from_string(content) -> Optional[str]:
return None
def mimetype_from_path(path) -> Optional[str]:
def mimetype_from_path(path) -> str | None:
try:
type_ = puremagic.from_file(path, mime=True)
return MIME_TYPE_FIXES.get(type_, type_)
@@ -51,8 +50,8 @@ def mimetype_from_path(path) -> Optional[str]:
def dicts_to_table_string(
headings: List[str], dicts: List[Dict[str, str]]
) -> List[str]:
headings: list[str], dicts: list[dict[str, str]]
) -> list[str]:
max_lengths = [len(h) for h in headings]
# Compute maximum length for each column
@@ -179,7 +178,7 @@ def token_usage_string(input_tokens, output_tokens, token_details) -> str:
return ", ".join(bits)
def extract_fenced_code_block(text: str, last: bool = False) -> Optional[str]:
def extract_fenced_code_block(text: str, last: bool = False) -> str | None:
"""
Extracts and returns Markdown fenced code block found in the given text.
@@ -217,7 +216,7 @@ def extract_fenced_code_block(text: str, last: bool = False) -> Optional[str]:
return None
def make_schema_id(schema: dict) -> Tuple[str, str]:
def make_schema_id(schema: dict) -> tuple[str, str]:
schema_json = json.dumps(schema, separators=(",", ":"))
schema_id = hashlib.blake2b(schema_json.encode(), digest_size=16).hexdigest()
return schema_id, schema_json
@@ -282,9 +281,9 @@ def resolve_schema_input(db, schema_input, load_template):
template = load_template(name)
schema_object = template.schema_object
except ValueError:
raise click.ClickException("Invalid template: {}".format(name))
raise click.ClickException(f"Invalid template: {name}")
if not schema_object:
raise click.ClickException("Template '{}' has no schema".format(name))
raise click.ClickException(f"Template '{name}' has no schema")
return template.schema_object
if schema_input.strip().startswith("{"):
try:
@@ -351,7 +350,7 @@ def schema_summary(schema: dict) -> str:
return ""
def schema_dsl(schema_dsl: str, multi: bool = False) -> Dict[str, Any]:
def schema_dsl(schema_dsl: str, multi: bool = False) -> dict[str, Any]:
"""
Build a JSON schema from a concise schema string.
@@ -372,7 +371,7 @@ def schema_dsl(schema_dsl: str, multi: bool = False) -> Dict[str, Any]:
}
# Initialize the schema dictionary with required elements
json_schema: Dict[str, Any] = {"type": "object", "properties": {}, "required": []}
json_schema: dict[str, Any] = {"type": "object", "properties": {}, "required": []}
# Check if the schema is newline-separated or comma-separated
if "\n" in schema_dsl:
@@ -487,9 +486,13 @@ def ensure_fragment(db, content):
source = content.source
with db.conn:
db.execute(sql, {"hash": hash_id, "content": content, "source": source})
return list(
db.query("select id from fragments where hash = :hash", {"hash": hash_id})
)[0]["id"]
return next(
iter(
db.query(
"select id from fragments where hash = :hash", {"hash": hash_id}
)
)
)["id"]
def ensure_tool(db, tool):
@@ -509,9 +512,13 @@ def ensure_tool(db, tool):
"plugin": tool.plugin,
},
)
return list(
db.query("select id from tools where hash = :hash", {"hash": tool.hash()})
)[0]["id"]
return next(
iter(
db.query(
"select id from tools where hash = :hash", {"hash": tool.hash()}
)
)
)["id"]
def maybe_fenced_code(content: str) -> str:
@@ -551,7 +558,7 @@ def has_plugin_prefix(value: str) -> bool:
return bool(_plugin_prefix_re.match(value))
def _parse_kwargs(arg_str: str) -> Dict[str, Any]:
def _parse_kwargs(arg_str: str) -> dict[str, Any]:
"""Parse key=value pairs where each value is valid JSON."""
tokens = []
buf = []
@@ -588,7 +595,7 @@ def _parse_kwargs(arg_str: str) -> Dict[str, Any]:
if buf:
tokens.append("".join(buf).strip())
kwargs: Dict[str, Any] = {}
kwargs: dict[str, Any] = {}
for token in tokens:
if not token:
continue
@@ -605,7 +612,7 @@ def _parse_kwargs(arg_str: str) -> Dict[str, Any]:
return kwargs
def instantiate_from_spec(class_map: Dict[str, Type], spec: str):
def instantiate_from_spec(class_map: dict[str, type], spec: str):
"""
Instantiate a class from a specification string with flexible argument formats.
@@ -665,7 +672,7 @@ def instantiate_from_spec(class_map: Dict[str, Type], spec: str):
return cls(**kw)
# Starts with quote / number / [ / t f n for single positional JSON value
if re.match(r'\s*(["\[\d\-]|true|false|null)', arg_body, re.I):
if re.match(r'\s*(["\[\d\-]|true|false|null)', arg_body, re.IGNORECASE):
try:
positional_value = json.loads(arg_body)
except json.JSONDecodeError as e:
@@ -682,7 +689,7 @@ TIMESTAMP_LEN = 6
RANDOMNESS_LEN = 10
_lock: Final = threading.Lock()
_last: Optional[bytes] = None # 16-byte last produced ULID
_last: bytes | None = None # 16-byte last produced ULID
def monotonic_ulid() -> ULID:
+1 -1
View File
@@ -53,7 +53,7 @@ dev = [
"mypy>=1.10.0",
"black>=26.3.1",
"pytest-recording",
"ruff",
"ruff>=0.16.0",
"syrupy",
"types-click",
"types-PyYAML",
+1 -3
View File
@@ -1,4 +1,2 @@
line-length = 160
[lint]
select = ["E4", "E7", "E9", "F"]
target-version = "py310"
+12 -11
View File
@@ -1,14 +1,15 @@
import importlib.metadata
import pytest
import sqlite_utils
import json
import sqlite3
import llm
import llm_echo
from llm.plugins import pm
import pytest
import sqlite_utils
from pydantic import Field
from pytest_httpx import IteratorStream
from typing import Optional
import llm
from llm.plugins import pm
def pytest_configure(config):
@@ -23,8 +24,8 @@ def pytest_report_header(config):
conn.close()
sqlite_utils_version = importlib.metadata.version("sqlite-utils")
return [
"SQLite: {}".format(version),
"sqlite-utils: {}".format(sqlite_utils_version),
f"SQLite: {version}",
f"sqlite-utils: {sqlite_utils_version}",
]
@@ -63,13 +64,13 @@ def env_setup(monkeypatch, user_path):
class MockModel(llm.Model):
model_id = "mock"
attachment_types = {"image/png", "audio/wav"}
attachment_types = frozenset({"image/png", "audio/wav"})
can_stream = True
supports_schema = True
supports_tools = True
class Options(llm.Options):
max_tokens: Optional[int] = Field(
max_tokens: int | None = Field(
description="Maximum number of tokens to generate.", default=None
)
@@ -302,7 +303,7 @@ def stream_events():
}
)
).encode("utf-8")
yield "data: [DONE]\n\n".encode("utf-8")
yield b"data: [DONE]\n\n"
@pytest.fixture
@@ -408,7 +409,7 @@ def stream_completion_events():
}
)
).encode("utf-8")
yield "data: [DONE]\n\n".encode("utf-8")
yield b"data: [DONE]\n\n"
@pytest.fixture
+17 -14
View File
@@ -1,10 +1,12 @@
from click.testing import CliRunner
from llm.cli import cli
import llm
import json
import pytest
import re
import pytest
from click.testing import CliRunner
import llm
from llm.cli import cli
@pytest.mark.parametrize("model_id_or_alias", ("gpt-3.5-turbo", "chatgpt"))
def test_set_alias(model_id_or_alias):
@@ -30,16 +32,17 @@ def test_cli_aliases_list(args):
runner = CliRunner()
result = runner.invoke(cli, args)
assert result.exit_code == 0
for line in (
"3.5 : gpt-3.5-turbo\n"
"chatgpt : gpt-3.5-turbo\n"
"chatgpt-16k : gpt-3.5-turbo-16k\n"
"3.5-16k : gpt-3.5-turbo-16k\n"
"4 : gpt-4\n"
"gpt4 : gpt-4\n"
"e-demo : embed-demo (embedding)\n"
"ada : text-embedding-ada-002 (embedding)\n"
).split("\n"):
for line in [
"3.5 : gpt-3.5-turbo",
"chatgpt : gpt-3.5-turbo",
"chatgpt-16k : gpt-3.5-turbo-16k",
"3.5-16k : gpt-3.5-turbo-16k",
"4 : gpt-4",
"gpt4 : gpt-4",
"e-demo : embed-demo (embedding)",
"ada : text-embedding-ada-002 (embedding)",
"",
]:
line = line.strip()
if not line:
continue
+2 -1
View File
@@ -1,6 +1,7 @@
import llm
import pytest
import llm
@pytest.mark.asyncio
async def test_async_model(async_mock_model):
+4 -1
View File
@@ -7,9 +7,10 @@ paths exercise real registered models with identical behaviour.
import json
import llm
import pytest
import llm
# ---- basic sanity: both variants are registered --------------------
@@ -91,6 +92,7 @@ async def test_async_from_row_response_messages_synthesized(tmp_path):
response.messages from _chunks+_tool_calls so follow-up chains
don't silently drop the assistant turn."""
import sqlite_utils
from llm.migrations import migrate
model = llm.get_async_model("echo")
@@ -123,6 +125,7 @@ async def test_async_load_conversation_follow_up_preserves_chain(tmp_path):
load_conversation, a follow-up turn's prompt.messages is the full
[user, assistant, user] chain — not missing the assistant."""
import sqlite_utils
from llm.cli import load_conversation
from llm.migrations import migrate
+7 -5
View File
@@ -1,10 +1,12 @@
from click.testing import CliRunner
import os
import sys
from unittest.mock import ANY
import pytest
from click.testing import CliRunner
import llm
from llm import cli
import pytest
TINY_PNG = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\xa6\x00\x00\x01\x1a"
@@ -47,8 +49,8 @@ def test_prompt_attachment(mock_model, logs_db, attachment_type, attachment_cont
conversation = conversations[0]
assert conversation["model"] == "mock"
assert conversation["name"] == "describe file"
response = list(logs_db["responses"].rows)[0]
attachment = list(logs_db["attachments"].rows)[0]
response = next(iter(logs_db["responses"].rows))
attachment = next(iter(logs_db["attachments"].rows))
assert attachment == {
"id": ANY,
"type": attachment_type,
@@ -56,7 +58,7 @@ def test_prompt_attachment(mock_model, logs_db, attachment_type, attachment_cont
"url": None,
"content": attachment_content,
}
prompt_attachment = list(logs_db["prompt_attachments"].rows)[0]
prompt_attachment = next(iter(logs_db["prompt_attachments"].rows))
assert prompt_attachment["attachment_id"] == attachment["id"]
assert prompt_attachment["response_id"] == response["id"]
+9 -7
View File
@@ -1,12 +1,14 @@
from click.testing import CliRunner
import re
from unittest.mock import ANY
import json
import llm.cli
import pytest
import sqlite_utils
import re
import sys
import textwrap
from unittest.mock import ANY
import pytest
import sqlite_utils
from click.testing import CliRunner
import llm.cli
@pytest.mark.xfail(sys.platform == "win32", reason="Expected to fail on Windows")
@@ -390,7 +392,7 @@ def test_chat_fragments(tmpdir):
output = runner.invoke(
llm.cli.cli,
["chat", "-m", "echo", "-f", path1],
input=("hi\n!fragment {}\nquit\n".format(path2)),
input=(f"hi\n!fragment {path2}\nquit\n"),
).output
assert '"prompt": "one' in output
assert '"prompt": "two"' in output
+4 -2
View File
@@ -1,7 +1,9 @@
from click.testing import CliRunner
import sys
import llm.cli
import pytest
from click.testing import CliRunner
import llm.cli
@pytest.mark.xfail(sys.platform == "win32", reason="Expected to fail on Windows")
+5 -3
View File
@@ -1,9 +1,11 @@
from click.testing import CliRunner
import json
import llm
from llm.cli import cli
import pytest
import sqlite_utils
from click.testing import CliRunner
import llm
from llm.cli import cli
@pytest.fixture
+5 -3
View File
@@ -1,8 +1,10 @@
from click.testing import CliRunner
from llm.cli import cli
import pytest
import json
import pytest
from click.testing import CliRunner
from llm.cli import cli
@pytest.mark.parametrize(
"args,expected_options,expected_error",
+7 -5
View File
@@ -1,9 +1,11 @@
import json
import llm
from llm.embeddings import Entry
from unittest.mock import ANY
import pytest
import sqlite_utils
from unittest.mock import ANY
import llm
from llm.embeddings import Entry
def test_demo_plugin():
@@ -20,7 +22,7 @@ def test_demo_plugin():
)
def test_embed_huge_list(batch_size, expected_batches):
model = llm.get_embedding_model("embed-demo")
huge_list = ("hello {}".format(i) for i in range(1000))
huge_list = (f"hello {i}" for i in range(1000))
kwargs = {}
if batch_size:
kwargs["batch_size"] = batch_size
@@ -120,7 +122,7 @@ def test_embed_multi(with_metadata, batch_size, expected_batches):
collection = llm.Collection("test", db, model_id="embed-demo")
model = collection.model()
assert getattr(model, "batch_count", 0) == 0
ids_and_texts = ((str(i), "hello {}".format(i)) for i in range(1000))
ids_and_texts = ((str(i), f"hello {i}") for i in range(1000))
kwargs = {}
if batch_size is not None:
kwargs["batch_size"] = batch_size
+8 -6
View File
@@ -1,13 +1,15 @@
from click.testing import CliRunner
from llm.cli import cli
from llm import Collection
import json
import pathlib
import pytest
import sqlite_utils
import sys
from unittest.mock import ANY
import pytest
import sqlite_utils
from click.testing import CliRunner
from llm import Collection
from llm.cli import cli
@pytest.mark.parametrize(
"format_,expected",
@@ -356,7 +358,7 @@ def test_embed_multi_files_binary_store(tmpdir):
assert result.exit_code == 0
db = sqlite_utils.Database(str(db_path))
assert db["embeddings"].count == 1
row = list(db["embeddings"].rows)[0]
row = next(iter(db["embeddings"].rows))
assert row == {
"collection_id": 1,
"id": "file.bin",
+3 -2
View File
@@ -1,6 +1,7 @@
import llm
import pytest
import numpy as np
import pytest
import llm
@pytest.mark.parametrize(
+9 -7
View File
@@ -1,12 +1,14 @@
from click.testing import CliRunner
import os
import textwrap
from importlib.metadata import version
from unittest import mock
import sqlite_utils
import yaml
from click.testing import CliRunner
from llm.cli import cli
from llm.migrations import migrate
from unittest import mock
import os
import yaml
import sqlite_utils
import textwrap
def test_fragments_set_show_remove(user_path):
@@ -138,7 +140,7 @@ def test_fragment_url_user_agent(mocked_openai_chat, user_path):
# Verify the User-Agent header was sent for the fragment URL request
requests = mocked_openai_chat.get_requests()
fragment_request = [r for r in requests if "example.com" in str(r.url)][0]
fragment_request = next(r for r in requests if "example.com" in str(r.url))
llm_version = version("llm")
expected_user_agent = f"llm/{llm_version} (https://llm.datasette.io/)"
assert fragment_request.headers["User-Agent"] == expected_user_agent
+6 -4
View File
@@ -1,10 +1,12 @@
from click.testing import CliRunner
import json
from llm.cli import cli
import pathlib
import pytest
import sys
import pytest
from click.testing import CliRunner
from llm.cli import cli
@pytest.mark.xfail(sys.platform == "win32", reason="Expected to fail on Windows")
@pytest.mark.parametrize("env", ({}, {"LLM_USER_PATH": "/tmp/llm-keys-test"}))
@@ -81,7 +83,7 @@ def test_uses_correct_key(mocked_openai_chat, monkeypatch, tmpdir):
def assert_key(key):
request = mocked_openai_chat.get_requests()[-1]
assert request.headers["Authorization"] == "Bearer {}".format(key)
assert request.headers["Authorization"] == f"Bearer {key}"
runner = CliRunner()
+8 -6
View File
@@ -1,14 +1,16 @@
from click.testing import CliRunner
import llm
from llm.cli import cli
from llm.models import Usage
import json
import os
import pathlib
from pydantic import BaseModel
from unittest import mock
import pytest
import sqlite_utils
from unittest import mock
from click.testing import CliRunner
from pydantic import BaseModel
import llm
from llm.cli import cli
from llm.models import Usage
def test_version():
+11 -10
View File
@@ -1,19 +1,21 @@
from click.testing import CliRunner
from llm.cli import cli
from llm.migrations import migrate
from llm.utils import monotonic_ulid
from llm import Fragment
import datetime
import json
import pathlib
import pytest
import re
import sqlite_utils
import sys
import textwrap
import time
from ulid import ULID
import pytest
import sqlite_utils
import yaml
from click.testing import CliRunner
from ulid import ULID
from llm import Fragment
from llm.cli import cli
from llm.migrations import migrate
from llm.utils import monotonic_ulid
SINGLE_ID = "5843577700ba729bb14c327b30441885"
MULTI_ID = "4860edd987df587d042a9eb2b299ce5c"
@@ -371,7 +373,6 @@ def test_logs_filtered(user_path, model, path_option):
("llama", ["-m", "davinci"], ["doc1", "doc3"]),
("llama", ["-m", "davinci2"], []),
# Adding -l/--latest should return latest first (order by id desc)
("llama", [], ["doc1", "doc3"]),
("llama", ["-l"], ["doc3", "doc1"]),
("llama", ["--latest"], ["doc3", "doc1"]),
),
@@ -1098,7 +1099,7 @@ def test_logs_resolved_model(logs_db, mock_model, async_mock_model, async_):
assert result.exit_code == 0
# Should have logged the resolved model name
assert logs_db["responses"].count
response = list(logs_db["responses"].rows)[0]
response = next(iter(logs_db["responses"].rows))
assert response["model"] == "mock"
assert response["resolved_model"] == "resolved-mock"
+4 -3
View File
@@ -1,9 +1,10 @@
import llm
from llm.migrations import migrate
from llm.embeddings_migrations import embeddings_migrations
import pytest
import sqlite_utils
import llm
from llm.embeddings_migrations import embeddings_migrations
from llm.migrations import migrate
EXPECTED = {
"id": str,
"model": str,
+1 -1
View File
@@ -22,7 +22,7 @@ def _sse(delta, finish_reason=None, usage=None, tool_calls=None):
chunk["choices"][0]["delta"]["tool_calls"] = tool_calls
if usage is not None:
chunk["usage"] = usage
return f"data: {json.dumps(chunk)}\n\n".encode("utf-8")
return f"data: {json.dumps(chunk)}\n\n".encode()
def _text_stream():
+8 -7
View File
@@ -3,16 +3,17 @@
import json
import os
import llm
import pytest
from pytest_httpx import IteratorStream
import llm
API_KEY = os.environ.get("PYTEST_OPENAI_API_KEY", None) or "badkey"
def _responses_sse(event_type, data):
data = {"type": event_type, **data}
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode("utf-8")
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
def _responses_reasoning_summary_stream():
@@ -239,7 +240,7 @@ def test_responses_input_translation():
model = llm.get_model("gpt-5.5")
class FakePrompt:
messages = [
messages = (
Message(role="system", parts=[TextPart(text="be brief")]),
Message(role="user", parts=[TextPart(text="2 + 2?")]),
Message(
@@ -256,7 +257,7 @@ def test_responses_input_translation():
role="tool",
parts=[ToolResultPart(name="add", output="4", tool_call_id="call_abc")],
),
]
)
items, instructions = model._build_responses_input(FakePrompt())
assert instructions == "be brief"
@@ -282,11 +283,11 @@ def test_responses_input_translation_assistant_text_uses_easy_input_message():
model = llm.get_model("gpt-5.5")
class FakePrompt:
messages = [
messages = (
Message(role="user", parts=[TextPart(text="hello")]),
Message(role="assistant", parts=[TextPart(text="first-ok")]),
Message(role="user", parts=[TextPart(text="what next?")]),
]
)
items, instructions = model._build_responses_input(FakePrompt())
@@ -526,7 +527,7 @@ def test_responses_tool_use_streaming(vcr):
)
output = "".join(chain)
assert "2869461" in output.replace(",", "")
first, second = chain._responses
first, _second = chain._responses
assert first.tool_calls()[0].arguments == {"a": 1231, "b": 2331}
+2 -1
View File
@@ -87,9 +87,10 @@ async def test_async_prompt_options_and_kwargs_conflict_raises(async_mock_model)
class Options(llm.Options):
from typing import Optional as _Opt
from pydantic import Field as _Field
max_tokens: _Opt[int] = _Field(default=None)
max_tokens: int | None = _Field(default=None)
async def execute(self, prompt, stream, response, conversation):
yield "ok"
+10 -7
View File
@@ -1,5 +1,7 @@
import json
import pytest
import llm
@@ -446,7 +448,7 @@ class TestStreamEventsFromStreamEventPlugin:
response = mock_model.prompt("hi")
response.text()
with pytest.raises(ValueError, match="part_index"):
response.messages() # noqa: B018
response.messages()
def test_provider_metadata_merges_last_wins(self, mock_model):
events = [
@@ -909,8 +911,8 @@ class TestPromptMessagesSynthesis:
]
def test_tool_results_become_tool_role_message(self, mock_model):
from llm.models import Prompt
from llm import ToolResult
from llm.models import Prompt
tr = ToolResult(name="t", output="ok", tool_call_id="c1")
p = Prompt(None, model=mock_model, tool_results=[tr])
@@ -1120,6 +1122,7 @@ class TestSqliteRehydrateMessages:
self, mock_model, tmp_path
):
import sqlite_utils
from llm.migrations import migrate
mock_model.enqueue(["answer text"])
@@ -1149,8 +1152,9 @@ class TestSqliteRehydrateMessages:
"""End-to-end: a follow-up turn via load_conversation must send
[user(q1), assistant(a1), user(q2)] — not drop the assistant."""
import sqlite_utils
from llm.migrations import migrate
from llm.cli import load_conversation
from llm.migrations import migrate
mock_model.enqueue(["first answer"])
mock_model.enqueue(["second answer"])
@@ -1179,6 +1183,7 @@ class TestSqliteRehydrateMessages:
assistant tool_use. Otherwise Anthropic sees an orphan
tool_result at the start of the continued request."""
import sqlite_utils
from llm.cli import load_conversation
from llm.migrations import migrate
@@ -1704,8 +1709,7 @@ class TestChainPropagatesSystem:
yield "done"
return
msgs = self._queue.pop(0)
for m in msgs:
yield m
yield from msgs
if not response._tool_calls:
response.add_tool_call(tool_call)
@@ -1731,8 +1735,7 @@ class TestChainPropagatesSystem:
yield "done"
return
msgs = self._queue.pop(0)
for m in msgs:
yield m
yield from msgs
if not response._tool_calls:
response.add_tool_call(tool_call)
+24 -19
View File
@@ -1,15 +1,16 @@
from click.testing import CliRunner
import click
import importlib
import json
import llm
from llm.tools import llm_version, llm_time
from llm import cli, hookimpl, plugins, get_template_loaders, get_fragment_loaders
import pathlib
import re
from unittest.mock import ANY
import click
import pytest
import textwrap
from click.testing import CliRunner
import llm
from llm import cli, get_fragment_loaders, get_template_loaders, hookimpl, plugins
from llm.tools import llm_time, llm_version
def test_register_commands():
@@ -176,16 +177,20 @@ def test_register_fragment_loaders(logs_db, httpx_mock):
cli.cli, ["-m", "echo", "-f", "mixed:x"], catch_exceptions=False
)
assert result3.exit_code == 0
result3.output.strip == textwrap.dedent("""\
system:
prompt:
one:x
attachments:
- https://example.com/attachment.png
""").strip()
assert json.loads(result3.output) == {
"prompt": "one:x",
"system": "",
"attachments": [
{
"type": None,
"path": None,
"url": "https://example.com/attachment.png",
"id": ANY,
}
],
"stream": True,
"previous": [],
}
finally:
plugins.pm.unregister(name="FragmentLoadersPlugin")
@@ -383,7 +388,7 @@ def test_register_tools(tmpdir, logs_db):
assert '"output": "HI"' in result4.output
# Now check in the database
tool_row = [row for row in logs_db["tools"].rows][0]
tool_row = next(iter(logs_db["tools"].rows))
assert tool_row["name"] == "upper"
assert tool_row["plugin"] == "ToolsPlugin"
@@ -760,7 +765,7 @@ def test_register_toolbox(tmpdir, logs_db):
[
"prompt",
"-T",
"Filesystem({})".format(json.dumps(str(my_dir2))),
f"Filesystem({json.dumps(str(my_dir2))})",
json.dumps({"tool_calls": [{"name": "Filesystem_list_files"}]}),
"-m",
"echo",
@@ -894,7 +899,7 @@ def test_toolbox_logging_async(logs_db, tmpdir):
"-T",
"Memory",
"--tool",
"Filesystem({})".format(json.dumps(str(path))),
f"Filesystem({json.dumps(str(path))})",
json.dumps(
{
"tool_calls": [
+4 -3
View File
@@ -7,16 +7,17 @@ dependency.
"""
import json
import pytest
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
import llm
from llm.serialization import (
AttachmentPartDict,
MessageDict,
PartDict,
ResponseDict,
ReasoningPartDict,
ResponseDict,
TextPartDict,
ToolCallPartDict,
ToolResultPartDict,
@@ -147,7 +148,7 @@ class TestPartDiscriminatedUnion:
TypeAdapter(PartDict).validate_python(d)
def test_unknown_type_rejected(self):
with pytest.raises(Exception):
with pytest.raises(ValidationError):
TypeAdapter(PartDict).validate_python({"type": "nonsense", "text": "x"})
+10 -8
View File
@@ -1,15 +1,17 @@
from click.testing import CliRunner
from importlib.metadata import version
import json
import os
import pathlib
import textwrap
from importlib.metadata import version
from unittest import mock
import pytest
import yaml
from click.testing import CliRunner
from llm import Template, Toolbox, hookimpl, user_dir
from llm.cli import cli
from llm.plugins import pm
import os
from unittest import mock
import pathlib
import pytest
import textwrap
import yaml
@pytest.mark.parametrize(
+10 -10
View File
@@ -1,16 +1,18 @@
import asyncio
import re
from click.testing import CliRunner
from importlib.metadata import version
import json
import llm
from llm import cli, CancelToolCall
from llm.migrations import migrate
from llm.tools import llm_time
import os
import re
import time
from importlib.metadata import version
import pytest
import sqlite_utils
import time
from click.testing import CliRunner
import llm
from llm import CancelToolCall, cli
from llm.migrations import migrate
from llm.tools import llm_time
API_KEY = os.environ.get("PYTEST_OPENAI_API_KEY", None) or "badkey"
@@ -544,7 +546,6 @@ def test_chain_sync_cancel_only_first_of_two():
if tool.name == "t1":
raise CancelToolCall("skip1")
# allow t2
return None
calls = [
{"name": "t1"},
@@ -583,7 +584,6 @@ async def test_chain_async_cancel_only_first_of_two():
async def before(tool, tool_call):
if tool.name == "t1":
raise CancelToolCall("skip1")
return None
calls = [
{"name": "t1"},
+4 -2
View File
@@ -1,7 +1,9 @@
import os
import pytest
import llm
from llm.tools import llm_version
import os
import pytest
API_KEY = os.environ.get("PYTEST_OPENAI_API_KEY", None) or "badkey"
+19 -14
View File
@@ -1,15 +1,17 @@
import json
import pytest
from llm import Toolbox, get_key
from llm.utils import (
extract_fenced_code_block,
instantiate_from_spec,
maybe_fenced_code,
monotonic_ulid,
schema_dsl,
simplify_usage_dict,
truncate_string,
monotonic_ulid,
)
from llm import get_key, Toolbox
@pytest.mark.parametrize(
@@ -83,22 +85,28 @@ def test_simplify_usage_dict(input_data, expected_output):
None,
],
[
"First code block:\n\n```python\ndef foo():\n return 'bar'\n```\n\n"
"Second code block:\n\n```javascript\nfunction foo() {\n return 'bar';\n}\n```",
(
"First code block:\n\n```python\ndef foo():\n return 'bar'\n```\n\n"
"Second code block:\n\n```javascript\nfunction foo() {\n return 'bar';\n}\n```"
),
False,
"def foo():\n return 'bar'\n",
],
[
"First code block:\n\n```python\ndef foo():\n return 'bar'\n```\n\n"
"Second code block:\n\n```javascript\nfunction foo() {\n return 'bar';\n}\n```",
(
"First code block:\n\n```python\ndef foo():\n return 'bar'\n```\n\n"
"Second code block:\n\n```javascript\nfunction foo() {\n return 'bar';\n}\n```"
),
True,
"function foo() {\n return 'bar';\n}\n",
],
[
"First code block:\n\n```python\ndef foo():\n return 'bar'\n```\n\n"
# This one has trailing whitespace after the second code block:
# https://github.com/simonw/llm/pull/718#issuecomment-2613177036
"Second code block:\n\n```javascript\nfunction foo() {\n return 'bar';\n}\n``` ",
(
"First code block:\n\n```python\ndef foo():\n return 'bar'\n```\n\n"
# This one has trailing whitespace after the second code block:
# https://github.com/simonw/llm/pull/718#issuecomment-2613177036
"Second code block:\n\n```javascript\nfunction foo() {\n return 'bar';\n}\n``` "
),
True,
"function foo() {\n return 'bar';\n}\n",
],
@@ -268,8 +276,6 @@ def test_schema_dsl_multi():
("Hello \n\t world!", 12, True, True, "Hello world!"),
# Edge cases
("12345", 5, False, False, "12345"),
("123456", 5, False, False, "12..."),
("12345", 5, False, True, "12345"), # Unchanged for exact fit
("123456", 5, False, False, "12..."), # Regular truncation for small max_length
# Very long string
("A" * 200, 10, False, False, "AAAAAAA..."),
@@ -287,7 +293,6 @@ def test_schema_dsl_multi():
True,
"12345...",
), # Too small for keep_end, use regular
("1234567890", 9, False, True, "12... 90"), # Just enough for keep_end
],
)
def test_truncate_string(text, max_length, normalize_whitespace, keep_end, expected):
@@ -328,7 +333,7 @@ def test_test_truncate_string_keep_end(
assert result == expected_full
# Only check prefix/suffix when we expect truncation with keep_end
if prefix_len is not None and len(text) > max_length and max_length >= 9:
if prefix_len is not None and len(text) > max_length >= 9:
assert result[:prefix_len] == text[:prefix_len]
assert result[-prefix_len:] == text[-prefix_len:]
assert "... " in result