Spaces:
Running
Running
| """Interactive demo for the lethe memory server (keyword + neighbour ranking, no embeddings).""" | |
| import os | |
| import time | |
| import uuid | |
| import gradio as gr | |
| import httpx | |
| from fastembed import TextEmbedding | |
| from lethe import Lethe | |
| EMBED = TextEmbedding("BAAI/bge-small-en-v1.5") # 384-dim, runs on CPU via ONNX | |
| BASE = "http://127.0.0.1:8787" | |
| TOKEN = "demo" | |
| SEED = [ | |
| "Priya is allergic to peanuts and avoids anything cooked in groundnut oil.", | |
| "Priya's flight to Denver leaves on Friday at 7:40 in the morning.", | |
| "The team agreed to ship the billing fix before the quarterly review.", | |
| "Marcus prefers meetings after 2 pm because he runs in the morning.", | |
| "The staging database password rotates every 30 days.", | |
| "Priya booked a vegetarian meal for the Denver flight.", | |
| "Marcus will present the fraud dashboard at the quarterly review.", | |
| ] | |
| BENCH = """| Ranking (LoCoMo, evidence in top 5) | Score | | |
| |---|---| | |
| | Keyword only | 0.438 | | |
| | Embedding similarity only | 0.479 | | |
| | Equal mix of both | 0.493 | | |
| | Mix + neighbouring turns, hand-set weights | 0.545 | | |
| | Learned from one score per recall | 0.568 ± 0.006 | | |
| | **Learned from "which memories helped"** | **0.618 ± 0.002** | | |
| Source: the repository README. Feedback in that benchmark is simulated from evidence labels, and on | |
| LongMemEval the gain is much smaller (0.729 vs 0.727 for the equal mix). This live demo has no embedding | |
| model and does not reproduce those numbers.""" | |
| def _wait_ready() -> None: | |
| for _ in range(50): | |
| try: | |
| httpx.get(BASE + "/healthz", timeout=1) | |
| return | |
| except httpx.HTTPError: | |
| time.sleep(0.2) | |
| def embed(text: str) -> list[float]: | |
| return next(iter(EMBED.embed([text]))).tolist() | |
| def _client(ns: str, with_embeddings: bool = True) -> Lethe: | |
| # Memories are always stored with an embedding; recall can leave it out to rank by keywords and neighbours only. | |
| return Lethe(BASE, namespace=ns, token=TOKEN, embed=embed if with_embeddings else None) | |
| _SEEDED: set[str] = set() | |
| def session(request: gr.Request) -> str: | |
| """Namespace for this visitor, seeded with the sample memories on first use (works for browser and API callers).""" | |
| ns = "s-" + (request.session_hash or uuid.uuid4().hex)[:10] | |
| if ns not in _SEEDED: | |
| with _client(ns) as c: | |
| for t in SEED: | |
| c.remember(t) | |
| _SEEDED.add(ns) | |
| return ns | |
| def add_memory(text, request: gr.Request): | |
| if not text.strip(): | |
| return "Nothing to store.", gr.update() | |
| with _client(session(request)) as c: | |
| c.remember(text.strip()[:1000]) | |
| return "Stored.", "" | |
| def do_recall(query, explore, use_embeddings, request: gr.Request): | |
| if not query.strip(): | |
| return gr.update(choices=[]), "", None | |
| with _client(session(request), use_embeddings) as c: | |
| r = c.recall(query.strip()[:500], k=5, explore=explore) | |
| labels = [f"{m.score:.3f} {m.text}" for m in r.memories] | |
| state = {"id": r.id, "ids": [m.id for m in r.memories], "labels": labels} | |
| msg = "Tick the memories that actually answer the question, then send feedback." | |
| return gr.update(choices=labels, value=[]), msg, state | |
| def send_feedback(state, picked, request: gr.Request): | |
| if not state: | |
| return "Run a recall first." | |
| useful = [i for i, lab in zip(state["ids"], state["labels"]) if lab in (picked or [])] | |
| reward = len(useful) / max(len(state["ids"]), 1) | |
| with _client(session(request)) as c: | |
| c.feedback(state["id"], reward, useful=useful) | |
| return f"Feedback sent: {len(useful)} useful of {len(state['ids'])}, reward {reward:.2f}." | |
| with gr.Blocks(title="lethe: agent memory that learns") as demo: | |
| gr.Markdown( | |
| "# lethe: agent memory that learns what to recall\n" | |
| "Real Go server running in this Space. [Code](https://github.com/shaikn6/lethe), MIT. " | |
| "Embeddings come from a small CPU model; switch them off to see keyword-only ranking." | |
| ) | |
| state = gr.State(None) | |
| with gr.Tab("Recall and teach"): | |
| q = gr.Textbox(label="Question", placeholder="What can Priya not eat?") | |
| explore = gr.Checkbox(False, label="Explore (tries slightly different rankings so it can learn, so results vary from run to run)") | |
| emb = gr.Checkbox(True, label="Use embeddings (BAAI/bge-small-en-v1.5). Off = keyword and neighbour ranking only") | |
| go = gr.Button("Recall", variant="primary") | |
| note = gr.Markdown() | |
| picks = gr.CheckboxGroup(label="Top 5 (score, memory)", choices=[]) | |
| fb = gr.Button("Send feedback") | |
| out = gr.Markdown() | |
| go.click(do_recall, [q, explore, emb], [picks, note, state]) | |
| fb.click(send_feedback, [state, picks], out) | |
| with gr.Tab("Add memory"): | |
| t = gr.Textbox(label="Memory text", lines=2) | |
| add = gr.Button("Store") | |
| msg = gr.Markdown() | |
| add.click(add_memory, [t], [msg, t]) | |
| with gr.Tab("Measured results"): | |
| gr.Markdown(BENCH) | |
| if __name__ == "__main__": | |
| _wait_ready() | |
| demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860"))) | |