Add Pro prompt redaction integration

flocks/session/lifecycle/compaction/compaction.py

@@ -1190,6 +1190,7 @@ async def process(
                         focus_instruction=focus_instruction,
                         previous_summary=previous_summary,
                         chat_messages=chat_messages,
+                        session_id=session_id,
                     )
                 except RuntimeError as _e:
                     # No provider configured — long cooldown (hermes: 600s)

flocks/session/lifecycle/compaction/summary.py

@@ -13,8 +13,16 @@
 # here in the future).
 ProgressCallback = Callable[[str, Dict[str, Any]], Awaitable[None]]
 
+from flocks.provider.provider import ChatMessage
 from flocks.utils.log import Log
 from flocks.session.prompt import SessionPrompt
+from flocks.session.llm_hook_utils import (
+    apply_hook_request_output,
+    restore_text_with_replacements,
+    serialize_chat_message,
+    stream_text_replacements_from_hook_output,
+    strip_think_blocks,
+)
 from .models import DEFAULT_COMPACTION_PROMPT_WITH_PREVIOUS
 
 log = Log.create(service="session.compaction.summarization")
@@ -159,16 +167,72 @@ async def _llm_chat_with_timeout(
     messages: list,
     max_tokens: int,
     timeout: int = COMPACTION_TIMEOUT_SECONDS,
+    session_id: Optional[str] = None,
+    purpose: str = "compaction_summary",
 ) -> Any:
     """Call provider_client.chat with a timeout guard."""
-    return await asyncio.wait_for(
+    provider_options: Dict[str, Any] = {"max_tokens": max_tokens}
+    replacements: list[tuple[str, str]] = []
+    if session_id:
+        try:
+            from flocks.hooks.pipeline import HookPipeline, HookStage
+
+            provider_id = (
+                getattr(provider_client, "provider_id", None)
+                or getattr(provider_client, "id", None)
+                or provider_client.__class__.__name__
+            )
+            llm_hook_metadata = {
+                "sessionID": session_id,
+                "agent": "session.compaction",
+                "step": None,
+                "model": {
+                    "providerID": provider_id,
+                    "modelID": model_id,
+                },
+                "purpose": purpose,
+            }
+            if await HookPipeline.has_stage_handlers(HookStage.LLM_BEFORE, llm_hook_metadata):
+                llm_before_ctx = await HookPipeline.run_llm_before({
+                    **llm_hook_metadata,
+                    "request": {
+                        "messageCount": len(messages),
+                        "messages": [serialize_chat_message(message) for message in messages],
+                        "toolCount": 0,
+                        "tools": [],
+                        "providerOptions": dict(provider_options),
+                        "providerToolsEnabled": False,
+                    },
+                })
+                hook_output = getattr(llm_before_ctx, "output", None) or {}
+                messages, provider_options = apply_hook_request_output(
+                    messages,
+                    provider_options,
+                    hook_output,
+                )
+                replacements = stream_text_replacements_from_hook_output(hook_output)
+        except Exception as hook_err:
+            log.error("compaction.llm_before_hook.error", {
+                "session_id": session_id,
+                "purpose": purpose,
+                "error": str(hook_err),
+            })
+            raise RuntimeError("compaction llm_before hook failed; request was not sent") from hook_err
+
+    provider_options.setdefault("max_tokens", max_tokens)
+    response = await asyncio.wait_for(
         provider_client.chat(
             model_id=model_id,
             messages=messages,
-            max_tokens=max_tokens,
+            **provider_options,
         ),
         timeout=timeout,
     )
+    if replacements and response is not None and isinstance(getattr(response, "content", None), str):
+        response.content = restore_text_with_replacements(response.content, replacements)
+    if response is not None and isinstance(getattr(response, "content", None), str):
+        response.content = strip_think_blocks(response.content)
+    return response
 
 
 async def summarize_single_pass(
@@ -181,6 +245,7 @@ async def summarize_single_pass(
     focus_instruction: Optional[str] = None,
     previous_summary: Optional[str] = None,
     chat_messages: Optional[list] = None,
+    session_id: Optional[str] = None,
 ) -> Optional[str]:
     """Generate summary in a single LLM call.
 
@@ -201,8 +266,6 @@ async def summarize_single_pass(
     as "merge new turns into the prior summary" rather than compressing from
     scratch.
     """
-    from flocks.provider.provider import ChatMessage
-
     if chat_messages:
         # Per-message truncation path (hermes-style): every turn contributes
         # a capped fragment (head + tail per message), so early decisions
@@ -236,6 +299,8 @@ async def summarize_single_pass(
             model_id=model_id,
             messages=[ChatMessage(role="user", content=request)],
             max_tokens=max_tokens,
+            session_id=session_id,
+            purpose="compaction_summary",
         )
     except asyncio.TimeoutError:
         log.error("compaction.single_pass.timeout", {
@@ -452,6 +517,8 @@ async def summarize_chunked_iterative(
                 messages=[ChatMessage(role="user", content=chunk_prompt)],
                 max_tokens=chunk_max_tokens,
                 timeout=COMPACTION_TIMEOUT_SECONDS,
+                session_id=session_id,
+                purpose="compaction_summary_chunk",
             )
             duration_ms = (time.perf_counter() - started) * 1000
             if resp and resp.content:

flocks/session/lifecycle/title.py

@@ -12,6 +12,12 @@
 
 from flocks.utils.log import Log
 from flocks.provider.provider import ChatMessage
+from flocks.session.llm_hook_utils import (
+    apply_hook_request_output,
+    restore_text_with_replacements,
+    serialize_chat_message,
+    stream_text_replacements_from_hook_output,
+)
 
 EventPublishCallback = Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]]
 
@@ -132,16 +138,61 @@ async def generate_title_after_first_message(
             # Send PROMPT_TITLE as system instruction, user question as user message
             title = ""
             try:
+                title_messages = [
+                    ChatMessage(role="system", content=_CANONICAL_TITLE_PROMPT),
+                    ChatMessage(role="user", content=question),
+                ]
+                provider_options: Dict[str, Any] = {"max_tokens": 50}
+                replacements: list[tuple[str, str]] = []
+                try:
+                    from flocks.hooks.pipeline import HookPipeline, HookStage
+
+                    llm_hook_metadata = {
+                        "sessionID": session_id,
+                        "agent": "session.title",
+                        "step": None,
+                        "model": {
+                            "providerID": provider_id,
+                            "modelID": model_id,
+                        },
+                        "purpose": "title_generation",
+                    }
+                    if await HookPipeline.has_stage_handlers(HookStage.LLM_BEFORE, llm_hook_metadata):
+                        llm_before_ctx = await HookPipeline.run_llm_before({
+                            **llm_hook_metadata,
+                            "request": {
+                                "messageCount": len(title_messages),
+                                "messages": [serialize_chat_message(message) for message in title_messages],
+                                "toolCount": 0,
+                                "tools": [],
+                                "providerOptions": dict(provider_options),
+                                "providerToolsEnabled": False,
+                            },
+                        })
+                        hook_output = getattr(llm_before_ctx, "output", None) or {}
+                        title_messages, provider_options = apply_hook_request_output(
+                            title_messages,
+                            provider_options,
+                            hook_output,
+                        )
+                        replacements = stream_text_replacements_from_hook_output(hook_output)
+                except Exception as hook_err:
+                    log.error("title.llm_before_hook.error", {
+                        "session_id": session_id,
+                        "error": str(hook_err),
+                    })
+                    raise RuntimeError("title llm_before hook failed; request was not sent") from hook_err
+
+                provider_options.setdefault("max_tokens", 50)
                 async for chunk in provider.chat_stream(
                     model_id,
-                    [
-                        ChatMessage(role="system", content=_CANONICAL_TITLE_PROMPT),
-                        ChatMessage(role="user", content=question),
-                    ],
-                    max_tokens=50,
+                    title_messages,
+                    **provider_options,
                 ):
                     if hasattr(chunk, 'delta') and chunk.delta:
                         title += chunk.delta
+                if replacements:
+                    title = restore_text_with_replacements(title, replacements)
             except Exception as llm_err:
                 log.warn("title.llm_failed", {
                     "session_id": session_id,

flocks/session/llm_hook_utils.py

@@ -0,0 +1,140 @@
+"""Shared helpers for LLM hook request/response payload handling."""
+
+from __future__ import annotations
+
+import re
+from typing import Any, Dict, List, Tuple
+
+from flocks.provider.provider import ChatMessage
+
+
+class StreamingTextReplacementBuffer:
+    """Incrementally replace streamed placeholders without leaking partial tokens."""
+
+    def __init__(self, replacements: List[Tuple[str, str]]):
+        self._replacements = [
+            (pattern, value)
+            for pattern, value in sorted(replacements, key=lambda item: len(item[0]), reverse=True)
+            if pattern
+        ]
+        self._buffer = ""
+        self._prefixes: set[str] = set()
+        self._max_pattern_len = 0
+        for pattern, _value in self._replacements:
+            self._max_pattern_len = max(self._max_pattern_len, len(pattern))
+            for index in range(1, len(pattern)):
+                self._prefixes.add(pattern[:index])
+
+    @property
+    def enabled(self) -> bool:
+        return bool(self._replacements)
+
+    def feed(self, text: str) -> str:
+        if not self.enabled or not text:
+            return text
+        self._buffer += text
+        keep = self._pending_suffix_length(self._buffer)
+        if keep:
+            flush_text = self._buffer[:-keep]
+            self._buffer = self._buffer[-keep:]
+        else:
+            flush_text = self._buffer
+            self._buffer = ""
+        return restore_text_with_replacements(flush_text, self._replacements)
+
+    def flush(self) -> str:
+        if not self.enabled or not self._buffer:
+            return ""
+        flush_text = self._buffer
+        self._buffer = ""
+        return restore_text_with_replacements(flush_text, self._replacements)
+
+    def _pending_suffix_length(self, text: str) -> int:
+        max_keep = min(len(text), max(self._max_pattern_len - 1, 0))
+        for length in range(max_keep, 0, -1):
+            if text[-length:] in self._prefixes:
+                return length
+        return 0
+
+
+def serialize_chat_message(message: ChatMessage) -> Dict[str, Any]:
+    payload = message.model_dump(exclude_none=True)
+    if not payload.get("custom_settings"):
+        payload.pop("custom_settings", None)
+    return payload
+
+
+def stream_text_replacements_from_hook_output(output: Dict[str, Any]) -> List[Tuple[str, str]]:
+    redaction = output.get("redaction") if isinstance(output, dict) else None
+    raw_items = redaction.get("streamTextReplacements") if isinstance(redaction, dict) else None
+    if not isinstance(raw_items, list):
+        return []
+
+    replacements: List[Tuple[str, str]] = []
+    for item in raw_items:
+        if not isinstance(item, dict):
+            continue
+        placeholder = item.get("placeholder")
+        value = item.get("value")
+        if isinstance(placeholder, str) and isinstance(value, str):
+            replacements.append((placeholder, value))
+    return replacements
+
+
+def restore_text_with_replacements(text: str, replacements: List[Tuple[str, str]]) -> str:
+    restored = text
+    for pattern, value in sorted(replacements, key=lambda item: len(item[0]), reverse=True):
+        if pattern:
+            restored = restored.replace(pattern, value)
+    return restored
+
+
+def restore_value_with_replacements(value: Any, replacements: List[Tuple[str, str]]) -> Any:
+    if isinstance(value, str):
+        return restore_text_with_replacements(value, replacements)
+    if isinstance(value, list):
+        return [restore_value_with_replacements(item, replacements) for item in value]
+    if isinstance(value, dict):
+        return {
+            key: restore_value_with_replacements(item, replacements)
+            for key, item in value.items()
+        }
+    return value
+
+
+_THINK_BLOCK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
+_THINK_TAG_RE = re.compile(r"</?think>", re.IGNORECASE)
+
+
+def strip_think_blocks(text: str) -> str:
+    if not isinstance(text, str) or "<think" not in text.lower():
+        return text
+    without_blocks = _THINK_BLOCK_RE.sub("", text)
+    return _THINK_TAG_RE.sub("", without_blocks).strip()
+
+
+def apply_hook_request_output(
+    messages: List[ChatMessage],
+    provider_options: Dict[str, Any],
+    output: Dict[str, Any],
+) -> tuple[List[ChatMessage], Dict[str, Any]]:
+    updated_request = output.get("request") if isinstance(output, dict) else None
+    if not isinstance(updated_request, dict):
+        return messages, provider_options
+
+    updated_messages = messages
+    raw_messages = updated_request.get("messages")
+    if isinstance(raw_messages, list):
+        updated_messages = [
+            message
+            if isinstance(message, ChatMessage)
+            else ChatMessage.model_validate(message)
+            for message in raw_messages
+        ]
+
+    updated_provider_options = provider_options
+    raw_provider_options = updated_request.get("providerOptions")
+    if isinstance(raw_provider_options, dict):
+        updated_provider_options = dict(raw_provider_options)
+
+    return updated_messages, updated_provider_options