pal-code 0.1.0__py3-none-any.whl

pal/local_llm/components/huggingface_local.py

@@ -0,0 +1,217 @@
+"""
+Fully in-process local core: the model is loaded directly into this
+Python process via transformers, no server (Ollama/llama.cpp/vLLM)
+involved at all. This is the deployment option AIOS 3.2 calls out
+explicitly -- "whether the LLM is hosted on cloud or on local device,
+what hardware conditions the LLM requires" -- that a server-based local
+core doesn't cover, since a server is still a separate process/host
+with its own port.
+
+Heavy deps (torch, transformers) are imported lazily inside __init__,
+not at module import time, so importing pal.llm_wrapper doesn't
+require them unless this specific core is actually built.
+
+This is also the only core with direct logits access, which is what
+lets it implement the beam-search step/snapshot/restore methods the
+logits-based context interrupt mechanism needs (see
+orchestrator/components/context_interrupt.py) -- no cloud or
+local-server backend exposes raw per-step logits over its API, so
+those cores are permanently limited to text-based snapshotting.
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any
+
+from pal.llm_wrapper.components.base import (
+    DeploymentType,
+    LLMResponse,
+    LLMSyscall,
+    LLMWrapper,
+    WrapperCapabilities,
+    WrapperConfig,
+)
+
+
+@dataclass
+class BeamSearchState:
+    """The "intermediate search tree structure" a logits-based snapshot
+    preserves: every beam's full token sequence so far and its
+    cumulative log-probability score. Restoring from this resumes
+    decoding exactly where generation left off, one token at a time --
+    unlike a text snapshot, nothing has to be re-tokenized or
+    re-interpreted by the model to continue.
+
+    Deliberately does NOT snapshot the model's KV cache (past_key_values)
+    -- that would make resumption cheaper (no need to reprocess the
+    prefix) but ties the snapshot to a single model instance's live
+    memory, and isn't picklable/inspectable the way this plain
+    tensor-of-token-ids state is. The tradeoff: resuming here recomputes
+    the forward pass over the full sequence-so-far rather than just the
+    new token, which is the honest cost of keeping snapshots simple and
+    portable.
+    """
+
+    input_ids: Any             # torch.Tensor, shape (beam_width, seq_len)
+    beam_scores: Any           # torch.Tensor, shape (beam_width,)
+    prompt_len: int
+    step: int = 0
+    finished: bool = False
+
+
+class HuggingFaceLocalLLM(LLMWrapper):
+    name = "huggingface_local"
+
+    def __init__(self, config: WrapperConfig):
+        super().__init__(config)
+        try:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError as e:
+            raise ImportError(
+                "transformers and torch are required for the in-process local "
+                "core. Run: pip install torch transformers"
+            ) from e
+
+        self._torch = torch
+        self.model_name = config.model
+        # "auto" lets accelerate/torch pick the best available device;
+        # explicit device (e.g. "cuda", "cpu", "mps") pins it, matching
+        # the "hardware conditions the LLM requires" concern from AIOS.
+        device = config.device or "auto"
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            device_map=device,
+            torch_dtype=config.extra.get("torch_dtype", "auto"),
+        )
+
+    def _render_prompt(self, messages: list[dict]) -> str:
+        if hasattr(self.tokenizer, "apply_chat_template") and self.tokenizer.chat_template:
+            return self.tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
+        # No chat template on this tokenizer/model -- fall back to a
+        # plain role-prefixed transcript rather than failing outright.
+        return "\n".join(f"{m.get('role', 'user')}: {m.get('content', '')}" for m in messages) + "\nassistant:"
+
+    def chat(self, syscall: LLMSyscall) -> LLMResponse:
+        prompt = self._render_prompt(syscall.messages)
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+
+        t0 = time.time()
+        with self._torch.no_grad():
+            output_ids = self.model.generate(
+                **inputs,
+                max_new_tokens=syscall.max_tokens or 512,
+                temperature=syscall.temperature if syscall.temperature is not None else 1.0,
+                do_sample=syscall.temperature not in (None, 0),
+                **syscall.params,
+            )
+        latency = time.time() - t0
+
+        prompt_tokens = inputs["input_ids"].shape[-1]
+        completion_tokens = output_ids.shape[-1] - prompt_tokens
+        text = self.tokenizer.decode(output_ids[0][prompt_tokens:], skip_special_tokens=True)
+
+        return LLMResponse(
+            text=text,
+            raw=output_ids,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            meta={"latency": latency, "device": str(self.model.device)},
+        )
+
+    def token_count(self, text: str) -> int:
+        return len(self.tokenizer.encode(text))
+
+    # -- logits-based beam search: step / snapshot / restore -------------------
+    #
+    # Not part of LLMWrapper's abstract contract (see base.py's LLMWrapper
+    # docstring) -- only reachable via capabilities().supports_logits_
+    # snapshot, and only ever called by ContextInterruptController.
+
+    def init_beam_search(self, syscall: LLMSyscall, beam_width: int = 1) -> BeamSearchState:
+        prompt = self._render_prompt(syscall.messages)
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        input_ids = inputs["input_ids"].repeat(beam_width, 1)
+        beam_scores = self._torch.zeros(beam_width, device=self.model.device)
+        return BeamSearchState(input_ids=input_ids, beam_scores=beam_scores, prompt_len=input_ids.shape[-1])
+
+    def beam_search_step(self, state: BeamSearchState, max_new_tokens: int = 512) -> BeamSearchState:
+        """Advances every beam by exactly one token. This single-token
+        granularity is the actual snapshot boundary: the context
+        interrupt mechanism can pause between ANY two calls to this
+        method, not just between whole streamed chunks the way the
+        text-based path is limited to."""
+        if state.finished:
+            return state
+
+        torch = self._torch
+        with torch.no_grad():
+            outputs = self.model(state.input_ids)
+            next_token_logits = outputs.logits[:, -1, :]  # (beam_width, vocab)
+            next_token_logprobs = torch.log_softmax(next_token_logits, dim=-1)
+            beam_width = state.input_ids.shape[0]
+
+            if beam_width == 1:
+                # Beam width 1 (this project's default, matching the
+                # AIOS example): greedy is equivalent to beam search
+                # with one beam, and avoids the topk/reshape machinery
+                # below for the common case.
+                next_token_id = next_token_logprobs.argmax(dim=-1, keepdim=True)
+                new_scores = state.beam_scores + next_token_logprobs.gather(-1, next_token_id).squeeze(-1)
+                new_input_ids = torch.cat([state.input_ids, next_token_id], dim=-1)
+            else:
+                # True beam search: expand every beam by its full
+                # vocabulary of candidates, then keep only the top
+                # `beam_width` sequences overall (not top-1-per-beam --
+                # that's what lets a single strong beam claim more than
+                # one slot next step, same as standard beam search).
+                vocab_size = next_token_logprobs.shape[-1]
+                candidate_scores = (state.beam_scores.unsqueeze(-1) + next_token_logprobs).view(-1)
+                top_scores, top_indices = candidate_scores.topk(beam_width)
+                beam_indices = top_indices // vocab_size
+                token_indices = (top_indices % vocab_size).unsqueeze(-1)
+                new_input_ids = torch.cat([state.input_ids[beam_indices], token_indices], dim=-1)
+                new_scores = top_scores
+
+        generated_len = new_input_ids.shape[-1] - state.prompt_len
+        eos_id = self.tokenizer.eos_token_id
+        finished = generated_len >= max_new_tokens or (
+            eos_id is not None and bool((new_input_ids[:, -1] == eos_id).all())
+        )
+
+        return BeamSearchState(
+            input_ids=new_input_ids,
+            beam_scores=new_scores,
+            prompt_len=state.prompt_len,
+            step=state.step + 1,
+            finished=finished,
+        )
+
+    def decode_beam_search(self, state: BeamSearchState) -> str:
+        """Decodes the current best beam's generated tokens. Called both
+        mid-generation (to snapshot a human-readable partial answer
+        alongside the raw state) and at completion."""
+        best_idx = int(state.beam_scores.argmax())
+        generated = state.input_ids[best_idx, state.prompt_len:]
+        return self.tokenizer.decode(generated, skip_special_tokens=True)
+
+    def capabilities(self) -> WrapperCapabilities:
+        return WrapperCapabilities(
+            deployment_type=DeploymentType.LOCAL_IN_PROCESS,
+            # The KV cache lives in this same process's memory for the
+            # lifetime of self.model -- there's no cross-request reuse
+            # concept to expose since nothing is reprocessed for a
+            # repeated turn within one generate() call.
+            supports_kv_cache=False,
+            supports_embeddings=False,
+            supports_streaming=False,
+            billed_per_token=False,
+            # Direct logits access via self.model(...) -- the only core
+            # in this project that can offer this.
+            supports_logits_snapshot=True,
+        )

pal/local_llm/components/local_llm.py

@@ -0,0 +1,60 @@
+"""The Local LLM component: models that run on this machine, behind a local server (Ollama, vLLM, LM Studio, llama.cpp).
+
+It is the same LiteLLM wrapper underneath, so a local model has the same interface, token and error format as a cloud one, but it differs
+where it matters: there is no key to protect (the OpenAI-style local servers get a placeholder, never a real key from the environment),
+a model's first call may load it (a longer time-out), it costs nothing per token, and it can tell you which models it has installed.
+
+The in-process Hugging Face model is the other local kind (huggingface_local.py): LiteLLM calls model servers and cannot load a model into
+this process.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from pal.llm_wrapper.components.litellm_wrapper import LOCAL_PREFIX, LiteLLMWrapper
+
+DEFAULT_URLS = {"ollama": "http://localhost:11434", "vllm": "http://localhost:8000/v1", "lmstudio": "http://localhost:1234/v1", "llamacpp": "http://localhost:8080/v1"}
+LOCAL_BACKENDS = tuple(LOCAL_PREFIX)
+
+
+class LocalLLM(LiteLLMWrapper):
+    name = "local_llm"
+    IS_LOCAL = True
+
+    def __init__(self, config):
+        if config.backend not in LOCAL_PREFIX:
+            raise ValueError(f"{config.backend!r} is not a local runtime (local runtimes: {', '.join(LOCAL_BACKENDS)})")
+        super().__init__(config)
+
+    @property
+    def url(self) -> str:
+        return (self.config.base_url or self.config.host or DEFAULT_URLS[self.config.backend]).rstrip("/")
+
+    def installed_models(self, timeout: float = 3.0, get: Any = None) -> list[str]:
+        """The models this server has, as it names them: Ollama's /api/tags, or the OpenAI-style /models of the others. An empty list when the
+        server cannot be reached (see `is_running`). `get` is for tests; it defaults to `requests.get`."""
+        if get is None:
+            import requests
+            get = requests.get
+        try:
+            if self.config.backend == "ollama":
+                r = get(f"{self.url}/api/tags", timeout=timeout)
+                r.raise_for_status()
+                return sorted(str(m.get("name") or m.get("model")) for m in r.json().get("models", []) if m.get("name") or m.get("model"))
+            r = get(f"{self.url}/models", timeout=timeout)
+            r.raise_for_status()
+            return sorted(str(m["id"]) for m in r.json().get("data", []) if m.get("id"))
+        except Exception:  # noqa: BLE001  (not running, not reachable, or not the API we expected)
+            return []
+
+    def is_running(self, timeout: float = 3.0, get: Any = None) -> bool:
+        """Is the server up? (It answers a models request, even with no models yet.)"""
+        if get is None:
+            import requests
+            get = requests.get
+        try:
+            path = "/api/tags" if self.config.backend == "ollama" else "/models"
+            get(f"{self.url}{path}", timeout=timeout).raise_for_status()
+            return True
+        except Exception:  # noqa: BLE001
+            return False

pal/logging_setup.py

@@ -0,0 +1,18 @@
+"""Logging setup for entry points. Library code only does `logging.getLogger(__name__)`; only the
+CLI/server calls configure_logging(). Secrets are never logged: API keys never enter a log call."""
+from __future__ import annotations
+
+import logging
+import sys
+
+
+def configure_logging(level: str = "INFO") -> None:
+    root = logging.getLogger("pal")
+    if root.handlers:
+        root.setLevel(level.upper())
+        return
+    handler = logging.StreamHandler(sys.stderr)
+    handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s"))
+    root.addHandler(handler)
+    root.setLevel(level.upper())
+    root.propagate = False

pal/memory/__init__.py

@@ -0,0 +1,7 @@
+from pal.memory.components.lru_k import LRUKTracker
+from pal.memory.components.manager import MemoryManager
+from pal.memory.components.note import MemoryNote
+from pal.memory.components.storage_manager import StorageManager
+from pal.memory.pipeline import MemoryPipeline
+
+__all__ = ["MemoryPipeline", "MemoryManager", "MemoryNote", "StorageManager", "LRUKTracker"]

pal/memory/components/__init__.py

@@ -0,0 +1,6 @@
+from pal.memory.components.lru_k import LRUKTracker
+from pal.memory.components.manager import MemoryManager
+from pal.memory.components.note import MemoryNote
+from pal.memory.components.storage_manager import StorageManager
+
+__all__ = ["LRUKTracker", "MemoryManager", "MemoryNote", "StorageManager"]

pal/memory/components/lru_k.py

@@ -0,0 +1,56 @@
+"""
+LRU-K tracking (O'Neil, O'Neil & Weikum, 1993) -- the eviction policy
+AIOS section 3.5 specifies for its memory manager, used here to decide
+which RAM-cached MemoryNotes get swapped to disk once the cache nears
+capacity.
+
+Plain LRU only looks at the single most recent access, so a note that
+was accessed once, long ago, and a note accessed constantly until just
+a moment ago can look identically "recently used" the instant both get
+touched again -- LRU can't tell a one-off touch from a genuinely hot
+item. LRU-K's backward K-distance (how long ago an item's K-th most
+recent access happened, not just its most recent one) is what "at
+least K times recently" in the spec actually cashes out to: an item
+needs a real access HISTORY, not just a single recent touch, to count
+as hot.
+"""
+import math
+import time
+from collections import deque
+
+
+class LRUKTracker:
+    def __init__(self, k: int = 2):
+        self.k = k
+        self._history: dict[str, deque] = {}
+
+    def record_access(self, item_id: str, when: float = None) -> None:
+        when = when if when is not None else time.time()
+        hist = self._history.setdefault(item_id, deque(maxlen=self.k))
+        hist.append(when)
+
+    def backward_k_distance(self, item_id: str, now: float = None) -> float:
+        """How long ago this item's K-th most recent access happened.
+        Items with fewer than K recorded accesses ever get infinite
+        distance -- they haven't earned "accessed at least K times
+        recently" status yet, so they're always evicted ahead of
+        anything that has."""
+        now = now if now is not None else time.time()
+        hist = self._history.get(item_id)
+        if not hist or len(hist) < self.k:
+            return math.inf
+        return now - hist[0]  # oldest of the last K accesses
+
+    def forget(self, item_id: str) -> None:
+        self._history.pop(item_id, None)
+
+    def eviction_order(self, candidate_ids: list[str], now: float = None) -> list[str]:
+        """Returns candidate_ids ordered worst-to-best for retention --
+        the FIRST id in the result is the one to evict first (largest
+        backward K-distance = coldest). Python's sort is stable, so
+        ties (most commonly: several items all still below K accesses,
+        all at infinite distance) keep their relative order from
+        candidate_ids -- in practice, whichever of them was inserted
+        into the RAM cache earliest."""
+        now = now if now is not None else time.time()
+        return sorted(candidate_ids, key=lambda i: self.backward_k_distance(i, now), reverse=True)

pal/memory/components/manager.py

@@ -0,0 +1,98 @@
+"""
+Runtime memory manager for agent interaction history (AIOS section
+3.5). Unlike a traditional OS memory manager, which manages physical
+RAM pages, this one manages MemoryNote objects -- conversation turns,
+tool-call results, and similar agent history -- between an in-process
+RAM cache and on-disk persistence via StorageManager.
+
+Every write is persisted to disk immediately: this project has no
+concept of "unsaved" agent memory, durability isn't optional here, only
+RAM residency is. The RAM cache exists purely as a bounded-size fast
+path so repeated access to hot notes doesn't round-trip through disk
+every time. When the cache's size crosses `high_water_mark` (default
+80%) of `block_limit`, the manager evicts the coldest RAM-resident
+notes by LRU-K backward-distance (components/lru_k.py) until it's back
+under the mark. "Evicted" only means dropped from the RAM cache --
+nothing is lost, since it was already durably written to disk on
+admission; the next read() for that id reloads it from disk and
+re-admits it to RAM (a swap-in).
+"""
+import time
+from typing import Optional
+
+from pal.memory.components.lru_k import LRUKTracker
+from pal.memory.components.note import MemoryNote
+from pal.memory.components.storage_manager import StorageManager
+
+
+class MemoryManager:
+    def __init__(
+        self,
+        storage: StorageManager,
+        block_limit: int = 200,
+        high_water_mark: float = 0.8,
+        k: int = 2,
+    ):
+        self.storage = storage
+        self.block_limit = block_limit
+        self.high_water_mark = high_water_mark
+        self.ram: dict[str, MemoryNote] = {}
+        self.lru_k = LRUKTracker(k=k)
+
+        # Warm the RAM cache from whatever's already durably persisted,
+        # most-recently-touched first, up to the block limit -- so a
+        # fresh process (a new session against an existing memory
+        # store) doesn't start with a cold cache and immediately thrash
+        # on its first few reads.
+        for note in sorted(self.storage.read_all(), key=lambda n: n.last_accessed, reverse=True)[: self.block_limit]:
+            self.ram[note.id] = note
+            self.lru_k.record_access(note.id, note.last_accessed)
+
+    def usage_ratio(self) -> float:
+        return len(self.ram) / self.block_limit if self.block_limit else 0.0
+
+    def write(self, note: MemoryNote) -> None:
+        """Allocates (or updates) a note: persisted to disk immediately,
+        cached in RAM, and recorded as an access. May trigger eviction
+        if this pushes the RAM cache over its high-water mark."""
+        self.storage.write(note)
+        self.ram[note.id] = note
+        self.lru_k.record_access(note.id)
+        self._maybe_evict()
+
+    def read(self, note_id: str) -> Optional[MemoryNote]:
+        """Reads a note by id -- from RAM if resident, else swaps it in
+        from disk. Returns None if the id doesn't exist anywhere."""
+        note = self.ram.get(note_id)
+        if note is None:
+            note = self.storage.read(note_id)
+            if note is None:
+                return None
+            self.ram[note_id] = note  # swap-in
+
+        note.retrieval_count += 1
+        note.last_accessed = time.time()
+        self.lru_k.record_access(note_id, note.last_accessed)
+        self.storage.write(note)  # persist the updated access metadata
+        self._maybe_evict()
+        return note
+
+    def delete(self, note_id: str) -> None:
+        self.ram.pop(note_id, None)
+        self.storage.delete(note_id)
+        self.lru_k.forget(note_id)
+
+    def all_ids(self) -> list[str]:
+        """Every note id known to the system, RAM-resident or swapped
+        to disk."""
+        return self.storage.list_ids()
+
+    def _maybe_evict(self) -> None:
+        while self.usage_ratio() >= self.high_water_mark and self.ram:
+            victim_id = self.lru_k.eviction_order(list(self.ram.keys()))[0]
+            # Already durable on disk from write() -- eviction here
+            # only drops the RAM cache entry, nothing more needs to be
+            # written. LRU-K access history is deliberately kept (not
+            # forgotten) for the evicted item, so if it's swapped back
+            # in later its recency history isn't reset to zero.
+            del self.ram[victim_id]

pal/memory/components/note.py

@@ -0,0 +1,54 @@
+"""
+A single persistent memory unit. Design adapted from AIOS
+(github.com/agiresearch/AIOS aios/memory/note.py) but trimmed to what
+Pal actually needs -- no evolution_history/links graph, since that
+requires the heavier retrieval stack (chromadb/qdrant) AIOS uses, which
+this lightweight module deliberately avoids.
+"""
+import time
+import uuid
+from typing import Optional
+
+
+class MemoryNote:
+    def __init__(
+        self,
+        content: str,
+        id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        timestamp: Optional[float] = None,
+        last_accessed: Optional[float] = None,
+        retrieval_count: int = 0,
+        tags: Optional[list[str]] = None,
+    ):
+        self.id = id or str(uuid.uuid4())
+        self.content = content
+        self.session_id = session_id
+        now = time.time()
+        self.timestamp = timestamp if timestamp is not None else now
+        self.last_accessed = last_accessed if last_accessed is not None else now
+        self.retrieval_count = retrieval_count
+        self.tags = tags or []
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "content": self.content,
+            "session_id": self.session_id,
+            "timestamp": self.timestamp,
+            "last_accessed": self.last_accessed,
+            "retrieval_count": self.retrieval_count,
+            "tags": self.tags,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "MemoryNote":
+        return cls(
+            content=d["content"],
+            id=d.get("id"),
+            session_id=d.get("session_id"),
+            timestamp=d.get("timestamp"),
+            last_accessed=d.get("last_accessed"),
+            retrieval_count=d.get("retrieval_count", 0),
+            tags=d.get("tags", []),
+        )

pal/memory/components/storage_manager.py

@@ -0,0 +1,60 @@
+"""
+Minimal disk-backed persistence for MemoryNotes: one JSON file per note
+under `root`, keyed by note id.
+
+Stands in for the full Storage Manager AIOS section 3.6 describes (a
+general read/write file-access layer for agents) -- this project
+doesn't implement that broader module yet, only the narrow slice the
+memory manager's RAM<->disk swapping needs: durably persist a note,
+load it back by id, delete it, and enumerate every persisted id.
+"""
+import json
+from pathlib import Path
+from typing import Optional
+
+from pal.memory.components.note import MemoryNote
+
+
+class StorageManager:
+    def __init__(self, root: Path):
+        self.root = Path(root)
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def _path(self, note_id: str) -> Path:
+        return self.root / f"{note_id}.json"
+
+    def write(self, note: MemoryNote) -> None:
+        with open(self._path(note.id), "w", encoding="utf-8") as f:
+            json.dump(note.to_dict(), f)
+
+    def read(self, note_id: str) -> Optional[MemoryNote]:
+        path = self._path(note_id)
+        if not path.exists():
+            return None
+        try:
+            with open(path, encoding="utf-8") as f:
+                return MemoryNote.from_dict(json.load(f))
+        except (json.JSONDecodeError, KeyError):
+            # Corrupt/partial file (e.g. killed mid-write) -- treat as
+            # missing rather than raising, same tolerance the old
+            # single-file store had for a corrupt memory.json.
+            return None
+
+    def delete(self, note_id: str) -> None:
+        path = self._path(note_id)
+        if path.exists():
+            path.unlink()
+
+    def exists(self, note_id: str) -> bool:
+        return self._path(note_id).exists()
+
+    def list_ids(self) -> list[str]:
+        return [p.stem for p in self.root.glob("*.json")]
+
+    def read_all(self) -> list[MemoryNote]:
+        notes = []
+        for note_id in self.list_ids():
+            note = self.read(note_id)
+            if note is not None:
+                notes.append(note)
+        return notes