| """Generic LLM wrappers and coercion.""" |
| from __future__ import annotations |
|
|
| import asyncio |
| import hashlib |
| import inspect |
| import json |
| import os |
| import threading |
| from typing import Any, Callable, Dict, Optional |
|
|
|
|
| async def acomplete(llm: Any, prompt: str) -> str: |
| """Native ``acomplete`` when the LLM has one, otherwise ``complete`` in a worker thread.""" |
| fn = getattr(llm, "acomplete", None) |
| if fn is not None: |
| return await fn(prompt) |
| return await asyncio.to_thread(llm.complete, prompt) |
|
|
|
|
| class FunctionLLM: |
| """Wrap ``fn(prompt) -> str`` (or an ``async def`` function).""" |
|
|
| def __init__(self, fn: Callable[[str], Any]) -> None: |
| self.fn = fn |
| if inspect.iscoroutinefunction(fn): |
| self.acomplete = self._acomplete |
|
|
| def complete(self, prompt: str) -> str: |
| if inspect.iscoroutinefunction(self.fn): |
| raise TypeError("this LLM is async-only; use the async API (ainvoke / abatch / astream)") |
| return str(self.fn(prompt)) |
|
|
| async def _acomplete(self, prompt: str) -> str: |
| return str(await self.fn(prompt)) |
|
|
|
|
| class CachedLLM: |
| """Thread-safe prompt -> response cache in front of any LLM. |
| |
| With ``path`` the cache persists as JSON (written after every new |
| response). Useful for reproducible runs and to avoid paying twice for the |
| same deterministic call. |
| """ |
|
|
| def __init__(self, llm: Any, path: Optional[str] = None) -> None: |
| self.llm = as_llm(llm) |
| self.path = path |
| self._lock = threading.Lock() |
| self.cache: Dict[str, Dict[str, str]] = {} |
| if path and os.path.exists(path): |
| with open(path, encoding="utf-8") as fh: |
| self.cache = json.load(fh) |
| self.hits = 0 |
| self.misses = 0 |
|
|
| @staticmethod |
| def key(prompt: str) -> str: |
| return hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:16] |
|
|
| def complete(self, prompt: str) -> str: |
| k = self.key(prompt) |
| with self._lock: |
| hit = self.cache.get(k) |
| if hit is not None: |
| self.hits += 1 |
| return hit["response"] |
| response = self.llm.complete(prompt) |
| self._store(k, prompt, response) |
| return response |
|
|
| async def acomplete(self, prompt: str) -> str: |
| k = self.key(prompt) |
| with self._lock: |
| hit = self.cache.get(k) |
| if hit is not None: |
| self.hits += 1 |
| return hit["response"] |
| response = await acomplete(self.llm, prompt) |
| self._store(k, prompt, response) |
| return response |
|
|
| def _store(self, k: str, prompt: str, response: str) -> None: |
| with self._lock: |
| self.misses += 1 |
| self.cache[k] = {"prompt": prompt, "response": response} |
| if self.path: |
| tmp = self.path + ".tmp" |
| os.makedirs(os.path.dirname(os.path.abspath(self.path)), exist_ok=True) |
| with open(tmp, "w", encoding="utf-8") as fh: |
| json.dump(self.cache, fh, ensure_ascii=False) |
| os.replace(tmp, self.path) |
|
|
|
|
| def as_llm(x: Any): |
| """Accept an object with ``complete``, a LangChain model, or ``fn(prompt) -> str``.""" |
| if x is None: |
| raise ValueError("an LLM is required") |
| if hasattr(x, "complete") and callable(getattr(x, "complete")): |
| return x |
| if hasattr(x, "invoke") and callable(getattr(x, "invoke")): |
| from ..integrations.langchain import LangChainLLM |
|
|
| return LangChainLLM(x) |
| if callable(x): |
| return FunctionLLM(x) |
| raise TypeError(f"cannot use {type(x).__name__} as an LLM: give .complete(prompt), .invoke(prompt) or a function") |
|
|
