GLM 4.6 support

Author: m1rl0k

.env.example

@@ -108,12 +108,17 @@ REFRAG_ENCODER_MODEL=BAAI/bge-base-en-v1.5
 REFRAG_PHI_PATH=/work/models/refrag_phi_768_to_dmodel.json
 REFRAG_SENSE=heuristic
 
-# Llama.cpp sidecar (optional)
+# Llama.cpp sidecar (optional; REFRAG_RUNTIME=llamacpp)
 LLAMACPP_URL=http://llamacpp:8080
 REFRAG_DECODER_MODE=prompt  # prompt|soft
 
 REFRAG_SOFT_SCALE=1.0
 
+# GLM API provider (alternative to llamacpp; REFRAG_RUNTIME=glm)
+GLM_API_KEY=
+GLM_API_BASE=https://api.z.ai/api/paas/v4/
+GLM_MODEL=glm-4.6
+
 
 # Operational safeguards and timeouts
 # Limit explosion of micro-chunks on huge files (0 to disable)

requirements.txt

@@ -8,6 +8,7 @@ tree_sitter>=0.25.2
 tree_sitter_languages; python_version < "3.13"
 mcp==1.17.0
 fastmcp==2.12.4
+openai>=1.0
 
 # Test-only
 pytest

scripts/mcp_indexer_server.py

@@ -3312,21 +3312,29 @@ async def expand_query(query: Any = None, max_new: Any = None) -> Dict[str, Any]
                 cap = max(0, min(2, int(max_new)))
             except (ValueError, TypeError):
                 cap = 2
-        from scripts.refrag_llamacpp import LlamaCppRefragClient, is_decoder_enabled  # type: ignore
-        if not is_decoder_enabled():
-            return {"alternates": [], "hint": "decoder disabled: set REFRAG_DECODER=1 and start llamacpp (LLAMACPP_URL)"}
+        runtime = str(os.environ.get("REFRAG_RUNTIME", "llamacpp")).strip().lower()
+        dec_on = str(os.environ.get("REFRAG_DECODER", "0")).strip().lower() in {"1", "true", "yes", "on"}
+        if runtime == "glm":
+            if not (dec_on and os.environ.get("GLM_API_KEY")):
+                return {"alternates": [], "hint": "decoder disabled: set REFRAG_DECODER=1 and GLM_API_KEY; or use llamacpp"}
+            from scripts.refrag_glm import GLMRefragClient  # type: ignore
+            client = GLMRefragClient()
+        else:
+            from scripts.refrag_llamacpp import LlamaCppRefragClient, is_decoder_enabled  # type: ignore
+            if not is_decoder_enabled():
+                return {"alternates": [], "hint": "decoder disabled: set REFRAG_DECODER=1 and start llamacpp (LLAMACPP_URL)"}
+            client = LlamaCppRefragClient()
         if not qlist:
             return {"alternates": []}
         prompt = (
             "You expand code search queries. Given short queries, propose up to 2 compact alternates.\n"
             "Return JSON array of strings only. No explanations.\n"
             f"Queries: {qlist}\n"
         )
-        client = LlamaCppRefragClient()
         out = client.generate_with_soft_embeddings(
             prompt=prompt,
             max_tokens=int(os.environ.get("EXPAND_MAX_TOKENS", "64") or 64),
-            temperature=0.0,  # Always 0 for deterministic expansion
+            temperature=0.0,
             top_k=int(os.environ.get("EXPAND_TOP_K", "30") or 30),
             top_p=float(os.environ.get("EXPAND_TOP_P", "0.9") or 0.9),
             stop=["\n\n"],
@@ -4790,8 +4798,13 @@ def _ca_build_prompt(context_blocks: list[str], citations: list[Dict[str, Any]],
 
 
 def _ca_decode(prompt: str, *, mtok: int, temp: float, top_k: int, top_p: float, stops: list[str]) -> str:
-    from scripts.refrag_llamacpp import LlamaCppRefragClient  # type: ignore
-    client = LlamaCppRefragClient()
+    runtime = str(os.environ.get("REFRAG_RUNTIME", "llamacpp")).strip().lower()
+    if runtime == "glm":
+        from scripts.refrag_glm import GLMRefragClient  # type: ignore
+        client = GLMRefragClient()
+    else:
+        from scripts.refrag_llamacpp import LlamaCppRefragClient  # type: ignore
+        client = LlamaCppRefragClient()
     return client.generate_with_soft_embeddings(
         prompt=prompt,
         max_tokens=mtok,

scripts/refrag_glm.py

@@ -0,0 +1,57 @@
+"""
+Feature-flagged adapter for decoder-side ReFRAG using GLM (ZhipuAI).
+
+Safe defaults:
+- Only used when REFRAG_DECODER=1 and REFRAG_RUNTIME=glm
+- Requires GLM_API_KEY; optionally configure GLM_MODEL
+"""
+from __future__ import annotations
+import os
+from typing import Any, Optional
+
+
+class GLMRefragClient:
+    """GLM client exposing generate_with_soft_embeddings(prompt, ...).
+
+    Notes:
+    - soft_embeddings are ignored (GLM does not support KV/soft-embed injection)
+    - prompt-mode only; mirrors llama.cpp adapter surface
+    - Uses OpenAI SDK with custom base_url for GLM API
+    """
+
+    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None) -> None:
+        self.api_key = api_key or os.environ.get("GLM_API_KEY", "").strip()
+        if not self.api_key:
+            raise ValueError("GLM_API_KEY is required when using REFRAG_RUNTIME=glm")
+        self.base_url = base_url or os.environ.get("GLM_API_BASE", "https://api.z.ai/api/paas/v4/")
+        from openai import OpenAI
+        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
+
+    def generate_with_soft_embeddings(
+        self,
+        prompt: str,
+        soft_embeddings: Optional[list[list[float]]] = None,  # unused
+        max_tokens: int = 256,
+        **gen_kwargs: Any,
+    ) -> str:
+        model = os.environ.get("GLM_MODEL", "glm-4.6")
+        temperature = float(gen_kwargs.get("temperature", 0.2))
+        top_p = float(gen_kwargs.get("top_p", 0.95))
+        stop = gen_kwargs.get("stop")
+
+        try:
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=int(gen_kwargs.get("max_tokens", max_tokens)),
+                temperature=temperature,
+                top_p=top_p,
+                stop=stop if stop else None,
+            )
+            msg = response.choices[0].message
+            # GLM-4.6 uses reasoning_content for thinking models
+            content = getattr(msg, 'reasoning_content', None) or msg.content or ""
+            return content.strip()
+        except Exception as e:
+            raise RuntimeError(f"GLM completion failed: {e}")
+