mirror of
https://github.com/vegu-ai/talemate.git
synced 2026-08-29 10:08:58 +02:00
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)
This commit is contained in:
veguAI
52
src/talemate/client/api_handles.py
Normal file
52
src/talemate/client/api_handles.py
Normal file
@@ -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)
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
extra_fields: dict[str, ExtraField] = api_handles_prompt_template_extra_fields(
|
||||
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.",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
@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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
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 api_handles_prompt_template(self) -> bool:
|
||||
return self.client_config.api_handles_prompt_template
|
||||
|
||||
@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,
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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."
|
||||
|
||||
|
||||
Reference in New Issue
Block a user