Testing the desk agent

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])

from langchain.agents.middleware import before_agent
from langchain.messages import AIMessage


@before_agent(can_jump_to=["end"])
def no_passwords(state, runtime):
    if "password" in state["messages"][-1].text.lower():
        answer = AIMessage("I cannot help with passwords. Please use the reset link.")
        return {"messages": [answer], "jump_to": "end"}

from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware, ModelCallLimitMiddleware, PIIMiddleware
from langgraph.checkpoint.memory import InMemorySaver
from desk_model import DeskModel
from password_check import no_passwords
from search import search_policies
from desk_tools import Customer, lookup_order, refund_order

agent = create_agent(
    DeskModel(),
    system_prompt="You are the support assistant for a small online shop. Answer in one or two short sentences, using only what the tools returned. If a tool says an order is not the customer's, say exactly that. Add nothing the tools did not say.",
    tools=[lookup_order, refund_order, search_policies],
    context_schema=Customer,
    middleware=[
        no_passwords,
        PIIMiddleware("credit_card", strategy="mask"),
        ModelCallLimitMiddleware(run_limit=6),
        HumanInTheLoopMiddleware(interrupt_on={"refund_order": True}),
    ],
    checkpointer=InMemorySaver(),
)

from langgraph.types import Command
from desk import Customer, agent


def say(who, text, thread):
    config = {"configurable": {"thread_id": thread}}
    result = agent.invoke({"messages": [{"role": "user", "content": text}]}, config,
                          context=Customer(who), version="v2")
    if result.interrupts:
        print(f"{who}: {text}\n  paused for approval: {result.interrupts[0].value['action_requests'][0]['args']}")
        result = agent.invoke(Command(resume={"decisions": [{"type": "approve"}]}), config,
                              context=Customer(who), version="v2")
        text = "(approved)"
    print(f"{who}: {text}\n  desk: {result.value['messages'][-1].text}")