122 lines
4.8 KiB
Python
122 lines
4.8 KiB
Python
"""Apply deployment instructions at the model boundary without editing chat state."""
|
|
|
|
from functools import lru_cache
|
|
|
|
from langchain_core.messages import BaseMessage, SystemMessage
|
|
from langchain_core.prompt_values import ChatPromptValue
|
|
from pydantic import Field
|
|
|
|
from deerflow.config.system_settings import load_system_settings
|
|
from deerflow.models.message_ordering import normalize_langchain_messages
|
|
|
|
|
|
def append_model_prompt_suffix(
|
|
messages: list[BaseMessage],
|
|
model_name: str,
|
|
*,
|
|
enabled: bool = True,
|
|
) -> list[BaseMessage]:
|
|
if not enabled or not model_name:
|
|
return messages
|
|
setting = load_system_settings().model_prompt_suffixes.get(model_name)
|
|
if setting is None or not setting.enabled or not setting.prompt.strip():
|
|
return messages
|
|
suffix = setting.prompt.strip()
|
|
normalized = normalize_langchain_messages(messages)
|
|
if not normalized or not isinstance(normalized[0], SystemMessage):
|
|
return [SystemMessage(content=suffix), *normalized]
|
|
system = normalized[0]
|
|
content = system.content
|
|
# invoke converts input before generate; fallback streaming can also pass
|
|
# through both entry points. Do not append the same instructions twice.
|
|
if isinstance(content, str):
|
|
if content == suffix or content.endswith("\n\n" + suffix):
|
|
return normalized
|
|
content = content + "\n\n" + suffix
|
|
else:
|
|
block = {"type": "text", "text": "\n\n" + suffix}
|
|
if content and content[-1] == block:
|
|
return normalized
|
|
content = [*content, block]
|
|
return [system.model_copy(update={"content": content}), *normalized[1:]]
|
|
|
|
|
|
def append_payload_prompt_suffix(
|
|
payload: dict,
|
|
model_name: str,
|
|
*,
|
|
enabled: bool = True,
|
|
) -> None:
|
|
"""Append after OpenAI transport adapters have added their own instructions."""
|
|
key = "messages" if "messages" in payload else "input"
|
|
if not model_name or not isinstance(payload.get(key), list):
|
|
return
|
|
messages = payload[key]
|
|
system_index = next((i for i in range(len(messages) - 1, -1, -1) if messages[i].get("role") in {"system", "developer"}), None)
|
|
system = messages[system_index] if system_index is not None else None
|
|
original = [SystemMessage(content=system["content"])] if system else []
|
|
updated = append_model_prompt_suffix(original, model_name, enabled=enabled)
|
|
if updated is original:
|
|
return
|
|
content = updated[0].content
|
|
if key == "input" and isinstance(content, list):
|
|
content = [{**block, "type": "input_text"} if isinstance(block, dict) and block.get("type") == "text" else block for block in content]
|
|
if system is not None:
|
|
payload[key] = [*messages]
|
|
payload[key][system_index] = {**system, "content": content}
|
|
else:
|
|
payload[key] = [{"role": "system", "content": content}, *messages]
|
|
|
|
|
|
@lru_cache(maxsize=64)
|
|
def with_model_prompt_suffix(model_class: type) -> type:
|
|
"""Keep the provider's tool binding/structured output and native streaming.
|
|
|
|
Read settings per request, including for already-created model instances.
|
|
The local routing name is excluded from serialization/provider parameters.
|
|
"""
|
|
|
|
from deerflow.models.patched_openai import PatchedChatOpenAI
|
|
|
|
uses_payload_suffix = issubclass(model_class, PatchedChatOpenAI)
|
|
|
|
class ModelWithPromptSuffix(model_class):
|
|
deerflow_prompt_model_name: str = Field(default="", exclude=True)
|
|
deerflow_prompt_suffix_enabled: bool = Field(default=True, exclude=True)
|
|
|
|
def _convert_input(self, model_input):
|
|
prompt = super()._convert_input(model_input)
|
|
if uses_payload_suffix:
|
|
return prompt
|
|
return ChatPromptValue(messages=append_model_prompt_suffix(
|
|
prompt.to_messages(),
|
|
self.deerflow_prompt_model_name,
|
|
enabled=self.deerflow_prompt_suffix_enabled,
|
|
))
|
|
|
|
def generate(self, messages, *args, **kwargs):
|
|
if uses_payload_suffix:
|
|
return super().generate(messages, *args, **kwargs)
|
|
return super().generate([
|
|
append_model_prompt_suffix(
|
|
batch,
|
|
self.deerflow_prompt_model_name,
|
|
enabled=self.deerflow_prompt_suffix_enabled,
|
|
)
|
|
for batch in messages
|
|
], *args, **kwargs)
|
|
|
|
async def agenerate(self, messages, *args, **kwargs):
|
|
if uses_payload_suffix:
|
|
return await super().agenerate(messages, *args, **kwargs)
|
|
return await super().agenerate([
|
|
append_model_prompt_suffix(
|
|
batch,
|
|
self.deerflow_prompt_model_name,
|
|
enabled=self.deerflow_prompt_suffix_enabled,
|
|
)
|
|
for batch in messages
|
|
], *args, **kwargs)
|
|
|
|
return ModelWithPromptSuffix
|