| """Model access for the graph-memory engines (HippoRAG 2 and ReMem). |
| |
| The engines call a language model through an object with `infer(messages, **kw)` returning |
| `(text, metadata, cache_hit)` and embed text through an object with `encode` / `batch_encode`. |
| `EngineChat` and `EngineEmbedder` provide those interfaces on top of the evaluation's chat client |
| and an OpenAI-compatible embeddings endpoint, so the engines use the configured backbone and the |
| token usage of every call is logged under the current phase. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import logging |
| import os |
| import sys |
| import time |
| import types |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from ..llm_client import ENDPOINTS, MAX_ATTEMPTS, SDK_MAX_RETRIES, USAGE, api_key_for, http_timeout |
|
|
| logger = logging.getLogger(__name__) |
|
|
| VENDOR = Path(__file__).resolve().parents[2] / "vendor" |
| _FENCE_OPEN, _FENCE_CLOSE = "```json\n", "```" |
|
|
|
|
| def add_engine_source(name: str) -> None: |
| """Put `vendor/<name>/src` on the import path.""" |
| src = VENDOR / name / "src" |
| if not src.exists(): |
| raise ImportError(f"{name} source folder not found at {src}") |
| if str(src) not in sys.path: |
| sys.path.insert(0, str(src)) |
|
|
|
|
| def stub_optional_imports() -> None: |
| """HippoRAG imports vllm at module level for its offline modes, which are not used here. |
| When vllm is absent, register empty modules carrying the names those imports expect.""" |
| if "vllm" in sys.modules: |
| return |
| try: |
| import vllm |
| return |
| except ImportError: |
| pass |
| names = {"vllm": ("SamplingParams", "LLM"), |
| "vllm.model_executor": (), |
| "vllm.model_executor.guided_decoding": (), |
| "vllm.model_executor.guided_decoding.guided_fields": ("GuidedDecodingRequest",)} |
| for module_name, classes in names.items(): |
| module = types.ModuleType(module_name) |
| for cls in classes: |
| setattr(module, cls, type(cls, (), {})) |
| sys.modules[module_name] = module |
| logger.info("vllm is not installed; empty placeholder modules registered") |
|
|
|
|
| def read_openie_docs(path: str) -> list: |
| """The `docs` list of an engine's OpenIE cache file; empty when the file is missing or unreadable.""" |
| if not os.path.exists(path): |
| return [] |
| try: |
| with open(path, encoding="utf-8") as fh: |
| return json.load(fh).get("docs", []) |
| except Exception as e: |
| logger.warning("could not read the OpenIE cache %s: %s", path, e) |
| return [] |
|
|
|
|
| class EngineChat: |
| """`infer` and `batch_infer` over the evaluation's chat client. |
| |
| With `max_tokens` set, each call sends an output limit: the engine's own `max_completion_tokens` |
| or `max_tokens` when it passes one, else `max_tokens`. With `max_tokens=None` no limit is sent |
| and the client's default applies. `echo=True` adds the messages and the reply to the metadata. |
| The metadata carries the reply's finish reason ("stop" when the client reports none), so an |
| engine can repair JSON that was cut off at the output limit ("length"). |
| """ |
|
|
| def __init__(self, client, *, temperature: float, max_tokens: int | None = None, echo: bool = False): |
| self._client = client |
| self._temperature = temperature |
| self._max_tokens = max_tokens |
| self._echo = echo |
|
|
| def infer(self, messages, **kw): |
| request = {"temperature": kw.get("temperature", self._temperature)} |
| if self._max_tokens is not None: |
| request["max_tokens"] = kw.get("max_completion_tokens", kw.get("max_tokens", self._max_tokens)) |
| if "response_format" in kw: |
| request["response_format"] = kw["response_format"] |
| try: |
| reply = self._client.chat(messages=messages, **request) |
| text = reply.content |
| wants_json = (kw.get("response_format") or {}).get("type") == "json_object" |
| if wants_json and text.startswith(_FENCE_OPEN) and text.endswith(_FENCE_CLOSE): |
| text = text[len(_FENCE_OPEN):-len(_FENCE_CLOSE)].strip() |
| except Exception as e: |
| logger.error("engine chat call failed: %s", e) |
| raise |
| meta = {"prompt": messages, "response": text} if self._echo else {} |
| meta.update(prompt_tokens=reply.input_tokens, completion_tokens=reply.output_tokens, |
| finish_reason=getattr(reply, "finish_reason", None) or "stop") |
| return text, meta, False |
|
|
| def batch_infer(self, messages_list, max_workers: int = 10, **kw): |
| """`infer` over several message lists in parallel; results keep the input order.""" |
| with ThreadPoolExecutor(max_workers=max_workers) as pool: |
| pending = [pool.submit(self.infer, messages, **kw) for messages in messages_list] |
| return [p.result() for p in pending] |
|
|
|
|
| class EngineEmbedder: |
| """`encode` / `batch_encode` over an OpenAI-compatible embeddings endpoint. |
| |
| Blank texts are sent as "empty"; an `instruction` keyword is prefixed to every text. Requests |
| hold at most `batch_size` texts and are retried up to MAX_ATTEMPTS times (a 400 response is |
| not retried). Rows are scaled to unit length unless `norm=False` is passed. |
| """ |
|
|
| def __init__(self, *, model: str | None, batch_size: int | None, dim: int | None = None, |
| api_key: str | None = None, base_url: str | None = None): |
| self.embedding_model_name = model |
| self.embedding_size = dim |
| self._batch_size = batch_size |
| self._api_key = api_key |
| self._base_url = base_url |
| self._endpoint = None |
|
|
| def _embeddings(self): |
| if self._endpoint is None: |
| from openai import OpenAI |
| base_url = self._base_url or ENDPOINTS["embedding"] |
| self._endpoint = OpenAI(api_key=self._api_key or api_key_for(base_url), base_url=base_url, |
| timeout=http_timeout(), max_retries=SDK_MAX_RETRIES).embeddings |
| logger.info("embedding endpoint ready, model=%s", self.embedding_model_name) |
| return self._endpoint |
|
|
| def _request(self, texts: list) -> list: |
| attempt = 0 |
| while True: |
| attempt += 1 |
| try: |
| response = self._embeddings().create(input=texts, model=self.embedding_model_name) |
| except Exception as e: |
| if attempt < MAX_ATTEMPTS and "Error code: 400" not in str(e): |
| time.sleep(1.5 * attempt) |
| continue |
| raise RuntimeError(f"embedding request failed after {attempt} attempt(s): {e}") from e |
| usage = getattr(response, "usage", None) |
| if usage is not None: |
| USAGE.log(self.embedding_model_name, "embed", int(getattr(usage, "prompt_tokens", 0) or 0), 0) |
| return [row.embedding for row in response.data] |
|
|
| def encode(self, texts, **kw) -> np.ndarray: |
| texts = [texts] if isinstance(texts, str) else texts |
| prefix = kw.get("instruction", "") |
| texts = [f"{prefix}{t if t and t.strip() else 'empty'}" for t in texts] |
| if not self._batch_size: |
| raise RuntimeError("set embedding.batch_size in backbone.yaml") |
| step = int(self._batch_size) |
| vectors = [] |
| for start in range(0, len(texts), step): |
| vectors.extend(self._request(texts[start:start + step])) |
| matrix = np.array(vectors) |
| if kw.get("norm", True): |
| lengths = np.linalg.norm(matrix, axis=1, keepdims=True) |
| matrix = matrix / np.where(lengths == 0, 1, lengths) |
| return matrix |
|
|
| batch_encode = encode |
|
|
