dchatel-comfyui_davcha_llm/utils.py

import os
import re
import inspect
import contextlib
from rapidfuzz import process, fuzz
from llama_cpp import llama_chat_format
from gguf import GGUFReader

def get_chat_handler(model_path: str, mmproj_path: str, threshold: float = 70.0):
        """
        Dynamically extracts multimodal chat handlers and uses RapidFuzz to map 
        the model filename to the handler, respecting **kwargs and inheritance.
        """
        # Extract valid chat handlers based purely on naming convention
        handlers = {}
        for name, cls in inspect.getmembers(llama_chat_format, inspect.isclass):
            if name.endswith("ChatHandler"):
                # Strip "ChatHandler" and lowercase (e.g., "Gemma4ChatHandler" -> "gemma4")
                clean_name = name.replace("ChatHandler", "").lower()
                handlers[clean_name] = cls

        if not handlers:
            raise RuntimeError("No chat handlers found in llama_cpp.llama_chat_format.")

        filename = os.path.basename(model_path).lower()
        
        best_match = None
        
        # Try matching against projector type first
        gguf = GGUFReader(mmproj_path)
        projector_type = gguf.get_field('clip.projector_type')
        if projector_type is not None:
            projector_type = projector_type.contents()
            best_match = process.extractOne(
                query=projector_type,
                choices=list(handlers.keys()),
                scorer=fuzz.partial_ratio
            )
        
        if not best_match or best_match[1] < threshold:
            # Try matching against the filename
            # Normalize the model filename for fuzzy matching
            # E.g., "gemma-4-8b-multimodal-q4_k.gguf" -> "gemma48bmultimodalq4kgguf"
            normalized_filename = re.sub(r'[^a-z0-9]', '', filename)

            # Fuzzy match using RapidFuzz
            best_match = process.extractOne(
                query=normalized_filename,
                choices=list(handlers.keys()),
                scorer=fuzz.partial_ratio
            )

        if not best_match:
            raise ValueError("Could not find any suitable chat handlers.")

        matched_key, score, _ = best_match

        # Enforce confidence threshold
        if score < threshold:
            # Fallback on default Llava
            handler_cls = llama_chat_format.Llava15ChatHandler
        else:
            handler_cls = handlers[matched_key]
            print(f"[Auto-Detect] Filename '{filename}' mapped to '{handler_cls.__name__}' (Fuzzy Score: {score:.1f}%)")
        
        # Safe Instantiation
        # Since classes like Gemma4ChatHandler use **kwargs, we just pass mmproj_path
        # and let Python's native MRO (Method Resolution Order) handle the routing.
        kwargs = {}
        
        # Check if handler_cls supports 'enable_thinking'
        try:
            sig = inspect.signature(handler_cls)
            has_enable_thinking = "enable_thinking" in sig.parameters
        except ValueError:
            # Fallback if signature extraction fails
            has_enable_thinking = False

        if has_enable_thinking:
            kwargs["enable_thinking"] = False

        try:
            return handler_cls(mmproj_path=mmproj_path, **kwargs)
        except TypeError as e:
            raise TypeError(
                f"The matched handler '{handler_cls.__name__}' rejected the arguments. "
                f"Is this truly a multimodal model? Details: {e}"
            )

def strip_via_fuzzy_tags(text: str, llm, threshold: float = 85.0) -> str:
    """
    Finds all native tags and uses RapidFuzz on ALL possible pairs to 
    dynamically locate and strip the reasoning block, regardless of structural tags.
    """
    # 1. Extract all tags
    matches = list(re.finditer(r'<[^>]+>', text))
    verified =[]
    
    # 2. Verify they are single tokens in the model's vocabulary
    for m in matches:
        tag_str = m.group()
        tokens = llm.tokenize(tag_str.encode('utf-8'), special=True)
        if len(tokens) <= 2:
            verified.append(m)
            
    if not verified:
        return text
        
    # 3. Search ALL combinations of tags to find the reasoning pair
    best_score = 0
    best_pair = None
    
    # Compare every tag against every subsequent tag (n^2 complexity, but n is tiny)
    for i in range(len(verified)):
        for j in range(i + 1, len(verified)):
            score = fuzz.ratio(verified[i].group(), verified[j].group())
            if score > best_score:
                best_score = score
                best_pair = (verified[i], verified[j])
                
    # 4. Strip the thought block if a highly similar pair is found
    # 85.0+ safely captures <think>/</think> (93%) while rejecting <|im_start|>/<|im_end|> (77%)
    if best_pair and best_score > threshold:
        open_tag, close_tag = best_pair
        print(f"[Fuzzy Tags] Paired '{open_tag.group()}' and '{close_tag.group()}' (Score: {best_score:.1f}%).")
        
        # Reconstruct the string: Everything BEFORE the open tag + everything AFTER the close tag
        return text[:open_tag.start()] + text[close_tag.end():]

    # 5. Implicit Fallback
    # If the system prompt forced the model into a thought block, it might output
    # text immediately, followed ONLY by the closing tag (e.g. <channel|>).
    first_match = verified[0]
    if len(text[:first_match.start()].strip()) > 0:
        print(f"[Fuzzy Tags] Implicit thought block. Stripping up to '{first_match.group()}'.")
        return text[first_match.end():].strip()

    return text

@contextlib.contextmanager
def disable_thinking_context(llm, enable_thinking):
    """
    Dynamically intercepts prompt generation to bypass reasoning blocks safely.
    Works across Jinja templates and hardcoded Multimodal Python handlers without breaking the token ledger.
    """
    if enable_thinking:
        yield
        return
        
    handler = getattr(llm, "chat_handler", None)
    if not handler:
        yield
        return
        
    # 1. Patch Jinja-based handlers (Text models)
    handler_template = getattr(handler, "template", None)
    if handler_template and isinstance(handler_template, str):
        handler.template = "{% set enable_thinking = false %}\n" + handler_template

    # 2. Extract the bypass block for Multimodal handlers
    template = llm.metadata.get("tokenizer.chat_template", "")
    
    # Regex hunts for exactly what the model creator defined: 
    # e.g., {%- if not enable_thinking... -%}{{- '<|channel>thought\n<channel|>' -}}
    match = re.search(r'if not enable_thinking.*?\{\{-?\s*[\'"]([^\'"]+)[\'"]\s*-?\}\}', template, re.DOTALL | re.IGNORECASE)
    
    if match:
        empty_block = match.group(1).replace('\\n', '\n')
        print(f"[DavchaLLM] Extracted thought bypass block: {repr(empty_block)}")
    else:
        # Robust inference fallback if the regex fails to find the toggle
        if "<channel|>" in template or "<channel>" in template or "Gemma" in type(handler).__name__:
            # specific fix for Gemma 4 / HauhauCS
            empty_block = "<|channel|>thought\n<channel|>"
        elif "</thought>" in template:
            empty_block = "<thought>\n</thought>\n"
        else:
            empty_block = "<think>\n</think>\n"
        print(f"[DavchaLLM] Inferred thought bypass block: {repr(empty_block)}")

    # 3. Patch Multimodal handlers (Vision models) BEFORE the ledger is created!
    original_m2p = getattr(handler, "messages_to_prompt", None)
    
    if original_m2p:
        def patched_m2p(*args, **kwargs):
            # Call original to format the user's prompt and extract images
            result = original_m2p(*args, **kwargs)
            
            # Tokenize our dynamically found bypass block
            think_tokens = llm.tokenize(empty_block.encode('utf-8'), special=True)
            bos_id = llm.token_bos()
            if bos_id != -1 and len(think_tokens) > 0 and think_tokens[0] == bos_id:
                think_tokens = think_tokens[1:]
                
            # Safely inject the tokens/string at the end of the prompt sequence
            # The handler will now build its ledger based on this modified sequence natively!
            if isinstance(result, tuple):
                prompt = result[0]
                if isinstance(prompt, list):
                    prompt = list(prompt) + list(think_tokens)
                elif isinstance(prompt, str):
                    prompt = prompt + empty_block
                return (prompt, *result[1:])
            else:
                prompt = result
                if isinstance(prompt, list):
                    prompt = list(prompt) + list(think_tokens)
                elif isinstance(prompt, str):
                    prompt = prompt + empty_block
                return prompt
                
        # Assign the patch directly to the instance
        handler.messages_to_prompt = patched_m2p

    try:
        # Yield back to the execution block (which calls create_chat_completion)
        yield
    finally:
        # Safely restore everything back to its original state for the next generation
        if handler_template and isinstance(handler_template, str):
            handler.template = handler_template
            
        if original_m2p:
            # Deleting the instance attribute cleanly restores the original Class method
            try:
                del handler.messages_to_prompt
            except AttributeError:
                pass