The desk's guardrails

import re

from langchain.chat_models import BaseChatModel
from langchain.messages import AIMessage, ToolMessage
from langchain_core.outputs import ChatGeneration, ChatResult


class ShopModel(BaseChatModel):
    tools: list = []

    @property
    def _llm_type(self):
        return "shop"

    def bind_tools(self, tools, **kwargs):
        return self.model_copy(update={"tools": tools})   # a copy holding the tools

    def _generate(self, messages, stop=None, run_manager=None, **kwargs):
        message = self.decide(messages)                   # the reply comes from decide
        return ChatResult(generations=[ChatGeneration(message=message)])

    def decide(self, messages):
        results = []                                # the tool results at the end
        for m in reversed(messages):
            if not isinstance(m, ToolMessage):
                break
            results.insert(0, m.text)
        if results:                                 # results are back: answer with them
            return AIMessage(" ".join(results))
        text = messages[-1].text
        orders = re.findall(r"\b[A-Z]\d+\b", text)
        tool = "refund_order" if "refund" in text.lower() else "lookup_order"
        if orders and tool in [t.name for t in self.tools]:   # one call per order id
            calls = [{"name": tool, "args": {"order_id": o}, "id": f"call_{o}"}
                     for o in orders]
            return AIMessage("", tool_calls=calls)
        if orders:                                  # that tool is not bound
            return AIMessage(f"I have no way to look up {orders[0]} yet.")
        return AIMessage("Hello. Which order is this about?")

from dataclasses import dataclass

from langchain.tools import ToolRuntime, tool

ORDERS = {"A17": ("ravi", "shipped on 3 March"), "C40": ("mei", "waiting for stock")}


@dataclass
class Customer:
    name: str


@tool
def lookup_order(order_id: str, runtime: ToolRuntime[Customer]) -> str:
    """Look up one of the customer's orders by its id, such as A17."""
    owner, status = ORDERS.get(order_id, (None, None))
    if owner != runtime.context.name:
        return f"{order_id} is not one of your orders."
    return f"{order_id} {status}."


@tool
def refund_order(order_id: str, runtime: ToolRuntime[Customer]) -> str:
    """Refund one of the customer's orders in full. This cannot be undone."""
    owner, _ = ORDERS.get(order_id, (None, None))
    if owner != runtime.context.name:
        return f"{order_id} is not one of your orders, so it cannot be refunded."
    return f"Refunded {order_id}."

from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

POLICIES = {
    "refunds.md": "Refunds go back to the card you paid with. They take up to 5 working days to arrive."
                  "\n\nYou can ask for a refund within 30 days of delivery. Opened items can be refunded if they are faulty.",
    "shipping.md": "Standard shipping takes 3 to 5 working days. Shipping is free on orders over 50 euros."
                   "\n\nExpress shipping arrives the next working day and costs 9 euros.",
    "accounts.md": "To reset your password, use the reset link on the sign-in page. Support staff never ask for your password.",
}

docs = [Document(page_content=text, metadata={"source": name}) for name, text in POLICIES.items()]
splitter = RecursiveCharacterTextSplitter(chunk_size=120, chunk_overlap=0, add_start_index=True)
chunks = splitter.split_documents(docs)

import re
import zlib

from langchain_core.embeddings import Embeddings

COMMON = {"a", "an", "and", "are", "can", "do", "does", "for", "how", "i",
          "if", "is", "it", "my", "of", "on", "the", "to", "what", "with", "you", "your"}


class WordEmbeddings(Embeddings):
    def embed_query(self, text):
        vector = [0.0] * 256
        for word in re.findall(r"[a-z]+", text.lower()):
            if word not in COMMON:
                vector[zlib.crc32(word.rstrip("s").encode()) % 256] += 1.0
        return vector

    def embed_documents(self, texts):
        return [self.embed_query(text) for text in texts]

from langchain.tools import tool
from langchain_core.vectorstores import InMemoryVectorStore
from policies import chunks
from word_embeddings import WordEmbeddings

store = InMemoryVectorStore(WordEmbeddings())
store.add_documents(chunks)


@tool
def search_policies(query: str) -> str:
    """Search the shop's policies on refunds, shipping and accounts.
    Pass the customer's question, word for word, as the query."""
    found = [doc for doc, score in store.similarity_search_with_score(query, k=2) if score >= 0.3]
    if not found:
        return "No policy covers this."
    return "\n".join(f"[{doc.metadata['source']}] {doc.page_content}" for doc in found)

import re

from langchain.messages import AIMessage
from shop_model import ShopModel


class DeskModel(ShopModel):
    def decide(self, messages):
        last = messages[-1]
        if last.type == "tool" and last.text == "No policy covers this.":
            return AIMessage("Our policies do not cover that. A person will reply.")
        if last.type == "tool" or re.findall(r"\b[A-Z]\d+\b", last.text):
            return super().decide(messages)
        query = {"name": "search_policies", "args": {"query": last.text}, "id": "call_p"}
        return AIMessage("", tool_calls=[query])