MemoryModel: a rule-based chat model

"""A stand-in chat model for LangMem: rules instead of a neural network."""
import re

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.utils.function_calling import convert_to_openai_tool

# What the model can notice: a sentence shape, the kind of fact, and how to write it down.
RULES = [
    (r"\bmy name is ([A-Z]\w+)", "name", "The customer's name is {0}"),
    (r"\b(email|call|text) me\b", "contact", "Wants us to {0} them"),
    (r"\bi (?:live in|moved to) ([A-Z]\w+)", "city", "Lives in {0}"),
    (r"\border (A-\d{4}) arrived (\w+)", "issue", "Order {0} arrived {1}"),
]


def notice(text):
    """Every (kind, value, sentence) the rules find in the text."""
    found = []
    for pattern, kind, template in RULES:
        for hit in re.finditer(pattern, text, re.I):
            words = [w.lower() if kind == "contact" else w for w in hit.groups()]
            found.append((kind, words[0], template.format(*words)))
    return found


def holds(memory, kind):
    """Whether an existing memory holds this kind of fact: as a sentence, or as a profile field."""
    template = next(t for _, k, t in RULES if k == kind)
    return template.split("{")[0] in memory or f"'{kind}'" in memory


def call(name, args, number=0):
    return {"name": name, "args": args, "id": f"call_{number}"}


class MemoryModel(BaseChatModel):
    @property
    def _llm_type(self):
        return "memory-rules"

    def bind_tools(self, tools, **kwargs):
        return self.bind(tools=[convert_to_openai_tool(t) for t in tools], **kwargs)

    def _generate(self, messages, stop=None, run_manager=None, tools=None, **kwargs):
        tools = {t["function"]["name"]: t["function"]["parameters"] for t in tools or []}
        return ChatResult(generations=[ChatGeneration(message=self.decide(messages, tools))])

    def decide(self, messages, tools):
        text = "\n".join(str(m.content) for m in messages)
        if messages[-1].type == "tool":
            result = str(messages[-1].content)
            remembered = re.findall(r'"content":"(.*?)"', result)
            if remembered:
                return AIMessage("I remember: " + "; ".join(sorted(remembered)))
            return AIMessage("I have nothing saved about that yet." if result.startswith("[") else "Saved.")
        if "GeneralResponse" in tools:
            return self.improve_prompt(text)
        if tools:
            return self.use_tools(messages, text, tools)
        if "summary" in str(messages[-1].content).lower():
            orders = sorted(set(re.findall(r"A-\d{4}", text)))
            return AIMessage(f"The customer asked about orders {', '.join(orders)}.")
        return AIMessage("Thanks for your message.")

    def use_tools(self, messages, text, tools):
        said = "\n".join(str(m.content) for m in messages if m.type == "human")
        found = notice(said.split("</existing>")[-1])
        calls = []
        # Existing memories come as <instance id=...> blocks: update the ones a new fact replaces.
        for doc_id, body in re.findall(r"<instance id=(\S+) [^>]*>\n(.*?)\n</instance>", text, re.S):
            for kind, value, sentence in list(found):
                if holds(body, kind) and "PatchDoc" in tools:
                    path, new = ("/content", sentence) if "'content'" in body else (f"/{kind}", value)
                    patch = [{"op": "replace", "path": path, "value": new}]
                    calls.append(call("PatchDoc", {"json_doc_id": doc_id, "planned_edits": f"new {kind}", "patches": patch}, len(calls)))
                    found.remove((kind, value, sentence))
            forget = re.search(r"\bforget my (\w+)", said, re.I)
            if forget and "RemoveDoc" in tools and holds(body, forget.group(1).lower()):
                calls.append(call("RemoveDoc", {"json_doc_id": doc_id}, len(calls)))
        # New facts go to the tool that records memories.
        record = next((n for n in tools if n not in ("PatchDoc", "RemoveDoc") and "search" not in n), None)
        fields = tools.get(record, {}).get("properties", {})
        if record and "content" in fields:
            extra = {"action": "create"} if "action" in fields else {}
            calls += [call(record, {"content": s, **extra}, len(calls) + i) for i, (_, _, s) in enumerate(found)]
        elif record and found and not calls:
            calls.append(call(record, {k: v for k, v, _ in found if k in fields}, len(calls)))
        if calls:
            return AIMessage("", tool_calls=calls)
        # A question, and a tool to look the answer up with.
        search = next((n for n in tools if "search" in n), None)
        if search and str(messages[-1].content).endswith("?"):
            return AIMessage("", tool_calls=[call(search, {"query": str(messages[-1].content)})])
        return AIMessage("Noted.")

    def improve_prompt(self, text):
        prompt = re.search(r"<current_prompt>\n(.*?)\n</current_prompt>", text, re.S).group(1)
        notes = [n for n in re.findall(r"<feedback \d+>\n(.*?)\n</feedback", text, re.S) if n.strip()]
        new = prompt + "".join(f"\n- {n}" for n in notes)
        return AIMessage("", tool_calls=[call("GeneralResponse", {
            "logic": f"Turned {len(notes)} piece(s) of feedback into rules.", "update_prompt": bool(notes), "new_prompt": new})])


# A stand-in embedder for the LangGraph store: texts that share words get similar vectors.
COMMON = {"the", "and", "for", "you", "your", "was", "with", "that", "this", "how", "what", "does", "have"}


def embed(texts):
    vectors = []
    for text in texts:
        vector = [0.0] * 64
        for word in re.findall(r"[a-z]+", text.lower()):
            if len(word) > 2 and word not in COMMON:
                vector[sum(map(ord, word.rstrip("s"))) % 64] += 1.0
        length = sum(x * x for x in vector) ** 0.5 or 1.0
        vectors.append([x / length for x in vector])
    return vectors

from memory_model import notice

for kind, value, sentence in notice("My name is Asha. Please EMAIL me. Order A-2002 arrived late."):
    print(f"{kind:8} {value:6} {sentence}")