From b2686b81f84432f2293ccd092ff04672d7a0bbc7 Mon Sep 17 00:00:00 2001 From: veguAI <152010387+vegu-ai-tools@users.noreply.github.com> Date: Mon, 6 Jul 2026 01:18:02 +0300 Subject: [PATCH] refactor: extract shared api_handles_prompt_template mixin and chat message assembly (#70) * refactor: extract shared api_handles_prompt_template mixin and chat message assembly (#67) * review: direct attribute access in ApiHandlesPromptTemplateMixin (fail loudly on mis-wiring) --- src/talemate/client/api_handles.py | 52 ++++++ src/talemate/client/base.py | 28 ++++ src/talemate/client/llamacpp.py | 50 +++--- src/talemate/client/ollama.py | 36 ++--- src/talemate/client/openai_compat.py | 52 ++---- src/talemate/client/tabbyapi.py | 59 ++----- src/talemate/client/textgenwebui.py | 53 +++--- ...test_client_api_handles_prompt_template.py | 152 ++++++++++++++++-- 8 files changed, 297 insertions(+), 185 deletions(-) create mode 100644 src/talemate/client/api_handles.py diff --git a/src/talemate/client/api_handles.py b/src/talemate/client/api_handles.py new file mode 100644 index 00000000..fc717a33 --- /dev/null +++ b/src/talemate/client/api_handles.py @@ -0,0 +1,52 @@ +""" +Shared pieces for clients that expose the `api_handles_prompt_template` +config flag, which routes generation through the remote API's own +prompt-template rendering (chat-style) instead of Talemate's local +prompt template. +""" + +import pydantic + +from .base import ExtraField + +__all__ = [ + "ApiHandlesPromptTemplateConfig", + "ApiHandlesPromptTemplateMixin", + "api_handles_prompt_template_extra_fields", +] + + +class ApiHandlesPromptTemplateConfig(pydantic.BaseModel): + api_handles_prompt_template: bool = False + + +def api_handles_prompt_template_extra_fields( + description: str, + label: str = "API handles prompt template (chat/completions)", +) -> dict[str, ExtraField]: + return { + "api_handles_prompt_template": ExtraField( + name="api_handles_prompt_template", + type="bool", + label=label, + required=False, + description=description, + ), + } + + +class ApiHandlesPromptTemplateMixin: + """ + Provides the `api_handles_prompt_template` flag property and bypasses + local prompt-template rendering when the API applies the model's own + template. Mix in before ClientBase. + """ + + @property + def api_handles_prompt_template(self) -> bool: + return self.client_config.api_handles_prompt_template + + def prompt_template(self, system_message: str, prompt: str): + if self.api_handles_prompt_template: + return prompt + return super().prompt_template(system_message, prompt) diff --git a/src/talemate/client/base.py b/src/talemate/client/base.py index e77f5479..122eeb5d 100644 --- a/src/talemate/client/base.py +++ b/src/talemate/client/base.py @@ -830,6 +830,34 @@ class ClientBase: return prompt, coercion return prompt, None + def chat_messages_for_coercion( + self, prompt: str, kind: str + ) -> tuple[list[dict], str | None]: + """ + Assembles chat messages for clients that let the API handle the + prompt template: [system, user] plus an assistant pre-fill message + when the prompt carries a coercion marker. + + Returns the messages and the stripped coercion prompt (None when + the prompt has no coercion). Transport-specific handling of the + pre-fill (e.g. TabbyAPI's `"prefix": True`, text-generation-webui's + `continue_`) stays at the call site. + """ + prompt, coercion_prompt = self.split_prompt_for_coercion(prompt) + if coercion_prompt: + coercion_prompt = coercion_prompt.strip() + + messages = [ + {"role": "system", "content": self.get_system_message(kind)}, + {"role": "user", "content": prompt.strip()}, + ] + + if coercion_prompt: + self.log.debug("Adding coercion pre-fill", coercion_prompt=coercion_prompt) + messages.append({"role": "assistant", "content": coercion_prompt}) + + return messages, coercion_prompt + def rate_limit_update(self): """ Updates the rate limit counter for the client. diff --git a/src/talemate/client/llamacpp.py b/src/talemate/client/llamacpp.py index ad20a5f0..203c3724 100644 --- a/src/talemate/client/llamacpp.py +++ b/src/talemate/client/llamacpp.py @@ -1,15 +1,18 @@ import json import pydantic -import structlog import httpx from openai import AsyncOpenAI +from talemate.client.api_handles import ( + ApiHandlesPromptTemplateConfig, + ApiHandlesPromptTemplateMixin, + api_handles_prompt_template_extra_fields, +) from talemate.client.base import ( STOPPING_STRINGS, ClientBase, CommonDefaults, - ExtraField, ParameterReroute, ) from talemate.client.registry import register @@ -22,24 +25,28 @@ from talemate.client.vision import VisionConfig, vision_extra_fields, OpenAIVisi from talemate.config.schema import Client as BaseClientConfig from talemate.exceptions import GenerationProcessingError -log = structlog.get_logger("talemate.client.llamacpp") - APPLY_TEMPLATE_TIMEOUT = 30 -class Defaults(CommonDefaults, pydantic.BaseModel): +class Defaults(CommonDefaults, ApiHandlesPromptTemplateConfig): # llama.cpp `llama-server` defaults to port 8080 (see ggml-org/llama.cpp README) api_url: str = "http://localhost:8080" max_token_length: int = 8192 - api_handles_prompt_template: bool = False -class ClientConfig(ConcurrentInference, VisionConfig, BaseClientConfig): - api_handles_prompt_template: bool = False +class ClientConfig( + ConcurrentInference, ApiHandlesPromptTemplateConfig, VisionConfig, BaseClientConfig +): + pass @register() -class LlamaCppClient(ConcurrentInferenceMixin, OpenAIVisionMixin, ClientBase): +class LlamaCppClient( + ApiHandlesPromptTemplateMixin, + ConcurrentInferenceMixin, + OpenAIVisionMixin, + ClientBase, +): """ Client for ggml-org/llama.cpp `llama-server`. @@ -63,11 +70,8 @@ class LlamaCppClient(ConcurrentInferenceMixin, OpenAIVisionMixin, ClientBase): self_hosted: bool = True extra_fields: dict = pydantic.Field( default_factory=lambda: { - "api_handles_prompt_template": ExtraField( - name="api_handles_prompt_template", - type="bool", + **api_handles_prompt_template_extra_fields( label="API handles prompt template", - required=False, description="The prompt template is rendered by llama.cpp using the model's built-in chat template, and the prompt template selection below is ignored. Response pre-filling keeps working. Keep this disabled for full control of the prompt template in Talemate; enable it to trust that the template on the remote end is correct.", ), **vision_extra_fields(), @@ -75,10 +79,6 @@ class LlamaCppClient(ConcurrentInferenceMixin, OpenAIVisionMixin, ClientBase): } ) - @property - def api_handles_prompt_template(self) -> bool: - return self.client_config.api_handles_prompt_template - @property def supported_parameters(self): # Talemate inference params (see config.schema.InferenceParameters) that @@ -112,11 +112,6 @@ class LlamaCppClient(ConcurrentInferenceMixin, OpenAIVisionMixin, ClientBase): ), ] - def prompt_template(self, system_message: str, prompt: str): - if not self.api_handles_prompt_template: - return super().prompt_template(system_message, prompt) - return prompt - def tune_prompt_parameters(self, parameters: dict, kind: str): super().tune_prompt_parameters(parameters, kind) @@ -179,16 +174,7 @@ class LlamaCppClient(ConcurrentInferenceMixin, OpenAIVisionMixin, ClientBase): message is included. """ - prompt, coercion_prompt = self.split_prompt_for_coercion(prompt) - - messages = [ - {"role": "system", "content": self.get_system_message(kind)}, - {"role": "user", "content": prompt.strip()}, - ] - - if coercion_prompt: - self.log.debug("Adding coercion pre-fill", coercion_prompt=coercion_prompt) - messages.append({"role": "assistant", "content": coercion_prompt.strip()}) + messages, _ = self.chat_messages_for_coercion(prompt, kind) payload = {"messages": messages} if not self.reason_enabled: diff --git a/src/talemate/client/ollama.py b/src/talemate/client/ollama.py index f629a7a8..05aa9830 100644 --- a/src/talemate/client/ollama.py +++ b/src/talemate/client/ollama.py @@ -3,6 +3,11 @@ import httpx import ollama import time +from talemate.client.api_handles import ( + ApiHandlesPromptTemplateConfig, + ApiHandlesPromptTemplateMixin, + api_handles_prompt_template_extra_fields, +) from talemate.client.base import ( STOPPING_STRINGS, ClientBase, @@ -19,18 +24,17 @@ log = structlog.get_logger("talemate.client.ollama") FETCH_MODELS_INTERVAL = 15 -class OllamaClientDefaults(CommonDefaults): +class OllamaClientDefaults(CommonDefaults, ApiHandlesPromptTemplateConfig): api_url: str = "http://localhost:11434" # Default Ollama URL model: str = "" # Allow empty default, will fetch from Ollama - api_handles_prompt_template: bool = False -class ClientConfig(BaseClientConfig): - api_handles_prompt_template: bool = False +class ClientConfig(ApiHandlesPromptTemplateConfig, BaseClientConfig): + pass @register() -class OllamaClient(ClientBase): +class OllamaClient(ApiHandlesPromptTemplateMixin, ClientBase): """ Ollama client for generating text using locally hosted models. """ @@ -48,15 +52,10 @@ class OllamaClient(ClientBase): manual_model_choices: list[str] = [] # Will be overridden by finalize_status defaults: OllamaClientDefaults = OllamaClientDefaults() self_hosted: bool = True - extra_fields: dict[str, ExtraField] = { - "api_handles_prompt_template": ExtraField( - name="api_handles_prompt_template", - type="bool", - label="API handles prompt template", - required=False, - description="Let Ollama handle the prompt template. Only do this if you don't know which prompt template to use. Letting talemate handle the prompt template will generally lead to improved responses.", - ), - } + extra_fields: dict[str, ExtraField] = api_handles_prompt_template_extra_fields( + label="API handles prompt template", + description="Let Ollama handle the prompt template. Only do this if you don't know which prompt template to use. Letting talemate handle the prompt template will generally lead to improved responses.", + ) @property def supported_parameters(self): @@ -97,10 +96,6 @@ class OllamaClient(ClientBase): """ return not self.api_handles_prompt_template and not self.reason_enabled - @property - def api_handles_prompt_template(self) -> bool: - return self.client_config.api_handles_prompt_template - async def status(self): """ Send a request to the API to retrieve the loaded AI model name. @@ -161,11 +156,6 @@ class OllamaClient(ClientBase): async def get_model_name(self): return self.model - def prompt_template(self, system_message: str, prompt: str): - if not self.api_handles_prompt_template: - return super().prompt_template(system_message, prompt) - return prompt - def tune_prompt_parameters(self, parameters: dict, kind: str): """ Tune parameters for Ollama's generate endpoint. diff --git a/src/talemate/client/openai_compat.py b/src/talemate/client/openai_compat.py index 717bc3bf..8f668255 100644 --- a/src/talemate/client/openai_compat.py +++ b/src/talemate/client/openai_compat.py @@ -1,15 +1,16 @@ import random -import pydantic -import structlog from openai import AsyncOpenAI +from talemate.client.api_handles import ( + ApiHandlesPromptTemplateConfig, + ApiHandlesPromptTemplateMixin, + api_handles_prompt_template_extra_fields, +) from talemate.client.base import ClientBase, ExtraField, FieldGroup from talemate.client.registry import register from talemate.config.schema import Client as BaseClientConfig -log = structlog.get_logger("talemate.client.openai_compat") - EXPERIMENTAL_DESCRIPTION = """Use this client if you want to connect to a service implementing an OpenAI-compatible API. Success is going to depend on the level of compatibility. Use the actual OpenAI client if you want to connect to OpenAI's API.""" # Sampler parameters whose inclusion in the request payload is user-toggleable. @@ -44,12 +45,11 @@ def _send_parameter_field(param: str) -> ExtraField: ) -class Defaults(pydantic.BaseModel): +class Defaults(ApiHandlesPromptTemplateConfig): api_url: str = "http://localhost:5000" api_key: str = "" max_token_length: int = 8192 model: str = "" - api_handles_prompt_template: bool = False double_coercion: str = None rate_limit: int | None = None send_temperature: bool = True @@ -57,15 +57,14 @@ class Defaults(pydantic.BaseModel): send_presence_penalty: bool = True -class ClientConfig(BaseClientConfig): - api_handles_prompt_template: bool = False +class ClientConfig(ApiHandlesPromptTemplateConfig, BaseClientConfig): send_temperature: bool = True send_top_p: bool = True send_presence_penalty: bool = True @register() -class OpenAICompatibleClient(ClientBase): +class OpenAICompatibleClient(ApiHandlesPromptTemplateMixin, ClientBase): client_type = "openai_compat" conversation_retries = 0 config_cls = ClientConfig @@ -79,20 +78,12 @@ class OpenAICompatibleClient(ClientBase): defaults: Defaults = Defaults() self_hosted: bool | None = None extra_fields: dict[str, ExtraField] = { - "api_handles_prompt_template": ExtraField( - name="api_handles_prompt_template", - type="bool", - label="API handles prompt template (chat/completions)", - required=False, + **api_handles_prompt_template_extra_fields( description="The API handles the prompt template, meaning your choice in the UI for the prompt template below will be ignored. This is not recommended and should only be used if the API does not support the `completions` andpoint or you don't know which prompt template to use.", ), **{f"send_{p}": _send_parameter_field(p) for p in TOGGLEABLE_PARAMETERS}, } - @property - def api_handles_prompt_template(self) -> bool: - return self.client_config.api_handles_prompt_template - @property def send_temperature(self) -> bool: return self.client_config.send_temperature @@ -125,11 +116,6 @@ class OpenAICompatibleClient(ClientBase): params.append(param) return params - def prompt_template(self, system_message: str, prompt: str): - if not self.api_handles_prompt_template: - return super().prompt_template(system_message, prompt) - return prompt - async def get_model_name(self): return self.model @@ -149,25 +135,11 @@ class OpenAICompatibleClient(ClientBase): parameters=parameters, ) - if self.can_be_coerced: - prompt, coercion_prompt = self.split_prompt_for_coercion(prompt) - else: - coercion_prompt = None - - messages = [ - {"role": "system", "content": self.get_system_message(kind)}, - {"role": "user", "content": prompt.strip()}, - ] + messages, coercion_prompt = self.chat_messages_for_coercion(prompt, kind) if coercion_prompt: - log.debug("Adding coercion pre-fill", coercion_prompt=coercion_prompt) - messages.append( - { - "role": "assistant", - "content": coercion_prompt.strip(), - "prefix": True, - } - ) + # continue the pre-fill via the (non-standard) prefix flag + messages[-1]["prefix"] = True response = await client.chat.completions.create( model=self.model_name, diff --git a/src/talemate/client/tabbyapi.py b/src/talemate/client/tabbyapi.py index 20df0d85..506db6c0 100644 --- a/src/talemate/client/tabbyapi.py +++ b/src/talemate/client/tabbyapi.py @@ -1,33 +1,33 @@ import random import json import httpx -import pydantic -import structlog +from talemate.client.api_handles import ( + ApiHandlesPromptTemplateConfig, + ApiHandlesPromptTemplateMixin, + api_handles_prompt_template_extra_fields, +) from talemate.client.base import ClientBase, ExtraField, CommonDefaults from talemate.client.registry import register from talemate.client.utils import urljoin from talemate.config.schema import Client as BaseClientConfig -log = structlog.get_logger("talemate.client.tabbyapi") - EXPERIMENTAL_DESCRIPTION = """Use this client to use all of TabbyAPI's features. Note on EXL3 models: They seem to be very sensitive to `presence_penalty`, `frequency_penalty` and `repetition_penalty_range`. If you're getting gibberish output, try creating a new inference parameter group and turn those off or way down.""" -class Defaults(CommonDefaults, pydantic.BaseModel): +class Defaults(CommonDefaults, ApiHandlesPromptTemplateConfig): api_url: str = "http://localhost:5000/v1" api_key: str = "" max_token_length: int = 8192 model: str = "" - api_handles_prompt_template: bool = False double_coercion: str = None -class ClientConfig(BaseClientConfig): - api_handles_prompt_template: bool = False +class ClientConfig(ApiHandlesPromptTemplateConfig, BaseClientConfig): + pass @register() -class TabbyAPIClient(ClientBase): +class TabbyAPIClient(ApiHandlesPromptTemplateMixin, ClientBase): client_type = "tabbyapi" conversation_retries = 0 config_cls = ClientConfig @@ -41,19 +41,9 @@ class TabbyAPIClient(ClientBase): manual_model: bool = False defaults: Defaults = Defaults() self_hosted: bool = True - extra_fields: dict[str, ExtraField] = { - "api_handles_prompt_template": ExtraField( - name="api_handles_prompt_template", - type="bool", - label="API handles prompt template (chat/completions)", - required=False, - description="The API handles the prompt template, meaning your choice in the UI for the prompt template below will be ignored. This is not recommended and should only be used if the API does not support the `completions` endpoint or you don't know which prompt template to use.", - ) - } - - @property - def api_handles_prompt_template(self) -> bool: - return self.client_config.api_handles_prompt_template + extra_fields: dict[str, ExtraField] = api_handles_prompt_template_extra_fields( + description="The API handles the prompt template, meaning your choice in the UI for the prompt template below will be ignored. This is not recommended and should only be used if the API does not support the `completions` endpoint or you don't know which prompt template to use.", + ) @property def experimental(self): @@ -87,11 +77,6 @@ class TabbyAPIClient(ClientBase): "temperature", ] - def prompt_template(self, system_message: str, prompt: str): - if not self.api_handles_prompt_template: - return super().prompt_template(system_message, prompt) - return prompt - async def get_model_name(self): url = urljoin(self.api_url, "model") headers = { @@ -124,25 +109,11 @@ class TabbyAPIClient(ClientBase): parameters=parameters, ) - if self.can_be_coerced: - prompt, coercion_prompt = self.split_prompt_for_coercion(prompt) - else: - coercion_prompt = None - - messages = [ - {"role": "system", "content": self.get_system_message(kind)}, - {"role": "user", "content": prompt.strip()}, - ] + messages, coercion_prompt = self.chat_messages_for_coercion(prompt, kind) if coercion_prompt: - log.debug("Adding coercion pre-fill", coercion_prompt=coercion_prompt) - messages.append( - { - "role": "assistant", - "content": coercion_prompt.strip(), - "prefix": True, - } - ) + # TabbyAPI continues the pre-fill via its prefix flag + messages[-1]["prefix"] = True payload = { "model": self.model_name, diff --git a/src/talemate/client/textgenwebui.py b/src/talemate/client/textgenwebui.py index 1f41ccd8..b8cc7051 100644 --- a/src/talemate/client/textgenwebui.py +++ b/src/talemate/client/textgenwebui.py @@ -9,7 +9,12 @@ import pydantic import structlog from openai import AsyncOpenAI -from talemate.client.base import STOPPING_STRINGS, ClientBase, Defaults, ExtraField +from talemate.client.api_handles import ( + ApiHandlesPromptTemplateConfig, + ApiHandlesPromptTemplateMixin, + api_handles_prompt_template_extra_fields, +) +from talemate.client.base import STOPPING_STRINGS, ClientBase, Defaults from talemate.client.registry import register from talemate.client.vision import VisionConfig, vision_extra_fields, OpenAIVisionMixin from talemate.config.schema import Client as BaseClientConfig @@ -17,17 +22,18 @@ from talemate.config.schema import Client as BaseClientConfig log = structlog.get_logger("talemate.client.textgenwebui") -class TextGeneratorWebuiClientDefaults(Defaults): +class TextGeneratorWebuiClientDefaults(Defaults, ApiHandlesPromptTemplateConfig): api_key: str = "" - api_handles_prompt_template: bool = False -class ClientConfig(VisionConfig, BaseClientConfig): - api_handles_prompt_template: bool = False +class ClientConfig(ApiHandlesPromptTemplateConfig, VisionConfig, BaseClientConfig): + pass @register() -class TextGeneratorWebuiClient(OpenAIVisionMixin, ClientBase): +class TextGeneratorWebuiClient( + ApiHandlesPromptTemplateMixin, OpenAIVisionMixin, ClientBase +): auto_determine_prompt_template: bool = True remote_model_locked: bool = True finalizers: list[str] = [ @@ -46,32 +52,19 @@ class TextGeneratorWebuiClient(OpenAIVisionMixin, ClientBase): self_hosted: bool = True extra_fields: dict = pydantic.Field( default_factory=lambda: { - "api_handles_prompt_template": ExtraField( - name="api_handles_prompt_template", - type="bool", - label="API handles prompt template (chat/completions)", - required=False, + **api_handles_prompt_template_extra_fields( description="Requests go to the chat/completions API and text-generation-webui applies the model's prompt template, and the prompt template selection below is ignored. Response pre-filling keeps working. Keep this disabled for full control of the prompt template in Talemate; enable it to trust that the template on the remote end is correct.", ), **vision_extra_fields(), } ) - @property - def api_handles_prompt_template(self) -> bool: - return self.client_config.api_handles_prompt_template - @property def requires_reasoning_pattern(self) -> bool: # in chat mode the API separates reasoning into reasoning_content # deltas, which are captured during streaming return not self.api_handles_prompt_template - def prompt_template(self, system_message: str, prompt: str): - if not self.api_handles_prompt_template: - return super().prompt_template(system_message, prompt) - return prompt - def make_client(self) -> AsyncOpenAI: api_key = self.api_key or "sk-1234" base = self.api_url.rstrip("/") @@ -207,34 +200,26 @@ class TextGeneratorWebuiClient(OpenAIVisionMixin, ClientBase): async def generate(self, prompt: str, parameters: dict, kind: str): loop = asyncio.get_event_loop() if self.api_handles_prompt_template: - # resolve on the event loop - the executor thread cannot see the + # assemble on the event loop - the executor thread cannot see the # active_scene contextvar, which would drop persona instructions - system_message = self.get_system_message(kind) + messages, coercion_prompt = self.chat_messages_for_coercion(prompt, kind) return await loop.run_in_executor( - None, self._generate_chat, prompt, parameters, system_message + None, self._generate_chat, messages, coercion_prompt, parameters ) return await loop.run_in_executor( None, self._generate, prompt, parameters, kind ) - def _generate_chat(self, prompt: str, parameters: dict, system_message: str): + def _generate_chat( + self, messages: list[dict], coercion_prompt: str | None, parameters: dict + ): """ Generates text via the chat/completions endpoint, letting text-generation-webui apply the model's prompt template. Coercion is passed as a partial assistant message that the API continues via its `continue_` parameter. """ - prompt, coercion_prompt = self.split_prompt_for_coercion(prompt) - - messages = [ - {"role": "system", "content": system_message}, - {"role": "user", "content": prompt.strip()}, - ] - if coercion_prompt: - coercion_prompt = coercion_prompt.strip() - log.debug("Adding coercion pre-fill", coercion_prompt=coercion_prompt) - messages.append({"role": "assistant", "content": coercion_prompt}) parameters["continue_"] = True parameters["mode"] = "instruct" diff --git a/tests/test_client_api_handles_prompt_template.py b/tests/test_client_api_handles_prompt_template.py index 556a28b3..138d6605 100644 --- a/tests/test_client_api_handles_prompt_template.py +++ b/tests/test_client_api_handles_prompt_template.py @@ -1,7 +1,9 @@ -"""Unit tests for the `api_handles_prompt_template` option on the -Text-Generation-WebUI and llama.cpp clients (issue #65). +"""Unit tests for the `api_handles_prompt_template` option on the clients +that expose it (issues #65 and #67): TabbyAPI, OpenAI Compatible, Ollama, +Text-Generation-WebUI and llama.cpp. -Covers the pure-Python logic: the `prompt_template()` bypass, chat message +Covers the pure-Python logic: the shared `ApiHandlesPromptTemplateMixin` +`prompt_template()` bypass, the shared `chat_messages_for_coercion()` message assembly including coercion splitting, the transport payloads (mocked HTTP), and the coercion pre-fill echo stripping. Live generation against real backends is exercised separately. @@ -16,8 +18,14 @@ import pytest import talemate.config.state as config_state from talemate.client import llamacpp as llamacpp_module +from talemate.client import ollama as ollama_module +from talemate.client import openai_compat as openai_compat_module +from talemate.client import tabbyapi as tabbyapi_module from talemate.client import textgenwebui as textgenwebui_module from talemate.client.llamacpp import LlamaCppClient +from talemate.client.ollama import OllamaClient +from talemate.client.openai_compat import OpenAICompatibleClient +from talemate.client.tabbyapi import TabbyAPIClient from talemate.client.textgenwebui import TextGeneratorWebuiClient from talemate.context import ActiveScene from talemate.exceptions import GenerationProcessingError @@ -48,6 +56,30 @@ def _register_textgenwebui(name: str, **kwargs): return cfg +def _register_tabbyapi(name: str, **kwargs): + cfg = tabbyapi_module.ClientConfig( + type="tabbyapi", name=name, api_url="http://fake:5000/v1", **kwargs + ) + config_state.CONFIG.clients[name] = cfg + return cfg + + +def _register_openai_compat(name: str, **kwargs): + cfg = openai_compat_module.ClientConfig( + type="openai_compat", name=name, api_url="http://fake:5000", **kwargs + ) + config_state.CONFIG.clients[name] = cfg + return cfg + + +def _register_ollama(name: str, **kwargs): + cfg = ollama_module.ClientConfig( + type="ollama", name=name, api_url="http://fake:11434", **kwargs + ) + config_state.CONFIG.clients[name] = cfg + return cfg + + # --------------------------------------------------------------------------- # prompt_template bypass # --------------------------------------------------------------------------- @@ -80,12 +112,101 @@ class TestPromptTemplateBypass: "PROMPT<|BOT|>Certainly:" ) + def test_tabbyapi_flag_off_applies_local_template(self, cfg_isolation): + _register_tabbyapi("tabby_off") + client = TabbyAPIClient(name="tabby_off") + assert client.prompt_template("SYS", "PROMPT") == "SYS\nPROMPT" + + def test_tabbyapi_flag_on_returns_raw_prompt(self, cfg_isolation): + _register_tabbyapi("tabby_on", api_handles_prompt_template=True) + client = TabbyAPIClient(name="tabby_on") + assert client.prompt_template("SYS", "PROMPT<|BOT|>Certainly:") == ( + "PROMPT<|BOT|>Certainly:" + ) + + def test_openai_compat_flag_off_applies_local_template(self, cfg_isolation): + _register_openai_compat("oaic_off") + client = OpenAICompatibleClient(name="oaic_off") + assert client.prompt_template("SYS", "PROMPT") == "SYS\nPROMPT" + + def test_openai_compat_flag_on_returns_raw_prompt(self, cfg_isolation): + _register_openai_compat("oaic_on", api_handles_prompt_template=True) + client = OpenAICompatibleClient(name="oaic_on") + assert client.prompt_template("SYS", "PROMPT<|BOT|>Certainly:") == ( + "PROMPT<|BOT|>Certainly:" + ) + + def test_ollama_flag_off_applies_local_template(self, cfg_isolation): + _register_ollama("oll_off") + client = OllamaClient(name="oll_off") + assert client.prompt_template("SYS", "PROMPT") == "SYS\nPROMPT" + + def test_ollama_flag_on_returns_raw_prompt(self, cfg_isolation): + _register_ollama("oll_on", api_handles_prompt_template=True) + client = OllamaClient(name="oll_on") + assert client.prompt_template("SYS", "PROMPT<|BOT|>Certainly:") == ( + "PROMPT<|BOT|>Certainly:" + ) + def test_coercion_stays_enabled_with_flag_on(self, cfg_isolation): _register_llamacpp("lcpp_coerce", api_handles_prompt_template=True) _register_textgenwebui("tgw_coerce", api_handles_prompt_template=True) assert LlamaCppClient(name="lcpp_coerce").can_be_coerced assert TextGeneratorWebuiClient(name="tgw_coerce").can_be_coerced + def test_ollama_coercion_disabled_with_flag_on(self, cfg_isolation): + # ollama sends a raw prompt (no chat pre-fill), so coercion is only + # possible when talemate renders the prompt template itself + _register_ollama("oll_coerce", api_handles_prompt_template=True) + assert not OllamaClient(name="oll_coerce").can_be_coerced + + +# --------------------------------------------------------------------------- +# ClientBase.chat_messages_for_coercion +# --------------------------------------------------------------------------- + + +class TestChatMessagesForCoercion: + def test_coercion_appends_assistant_prefill(self, cfg_isolation): + _register_tabbyapi("tabby_msgs", api_handles_prompt_template=True) + client = TabbyAPIClient(name="tabby_msgs") + + messages, coercion = client.chat_messages_for_coercion( + "Describe the room.<|BOT|>The room is ", "narrate" + ) + + assert [m["role"] for m in messages] == ["system", "user", "assistant"] + assert messages[1]["content"] == "Describe the room." + assert messages[2]["content"] == "The room is" + assert coercion == "The room is" + + def test_no_coercion_returns_system_and_user_only(self, cfg_isolation): + _register_openai_compat("oaic_msgs", api_handles_prompt_template=True) + client = OpenAICompatibleClient(name="oaic_msgs") + + messages, coercion = client.chat_messages_for_coercion( + "Describe the room.", "narrate" + ) + + assert [m["role"] for m in messages] == ["system", "user"] + assert messages[1]["content"] == "Describe the room." + assert coercion is None + + def test_double_coercion_prepended_to_prefill(self, cfg_isolation): + _register_tabbyapi( + "tabby_dc", + api_handles_prompt_template=True, + double_coercion="Sure thing!", + ) + client = TabbyAPIClient(name="tabby_dc") + + messages, coercion = client.chat_messages_for_coercion( + "Describe the room.<|BOT|>The room is", "narrate" + ) + + assert coercion == "Sure thing!\n\nThe room is" + assert messages[2]["content"] == coercion + # --------------------------------------------------------------------------- # llama.cpp /apply-template @@ -254,9 +375,12 @@ class TestTextGenWebuiChatGenerate: "[DONE]", ] - response = client._generate_chat( - "Describe the room.<|BOT|>The room is", {"temperature": 0.7}, "SYS" - ) + messages = [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "Describe the room."}, + {"role": "assistant", "content": "The room is"}, + ] + response = client._generate_chat(messages, "The room is", {"temperature": 0.7}) assert response == " dark and quiet." @@ -287,9 +411,11 @@ class TestTextGenWebuiChatGenerate: "[DONE]", ] - response = client._generate_chat( - "Describe the room.", {"temperature": 0.7}, "SYS" - ) + messages = [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "Describe the room."}, + ] + response = client._generate_chat(messages, None, {"temperature": 0.7}) assert response == "The room is dark." assert client._reasoning_response == "Thinking about the room." @@ -311,9 +437,11 @@ class TestTextGenWebuiChatGenerate: "[DONE]", ] - response = client._generate_chat( - "Describe the room.", {"temperature": 0.7}, "SYS" - ) + messages = [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "Describe the room."}, + ] + response = client._generate_chat(messages, None, {"temperature": 0.7}) assert response == "The room is dark."