The desk that remembers with a session

from agents import Agent, Runner, function_tool, set_tracing_disabled, SQLiteSession
from agents.models.interface import Model
from agents.items import ModelResponse
from agents.usage import Usage
from openai.types.responses import (
    ResponseOutputMessage, ResponseOutputText, ResponseFunctionToolCall,
)
import re
set_tracing_disabled(True)


def _message(text):
    return ResponseOutputMessage(
        id="msg", role="assistant", type="message", status="completed",
        content=[ResponseOutputText(text=text, type="output_text", annotations=[])],
    )


def _tool_call(name, arguments):
    return ResponseFunctionToolCall(
        id="fc", call_id="call_1", name=name, arguments=arguments, type="function_call",
    )


def all_user_text(input):
    if isinstance(input, str):
        return input
    parts = []
    for it in input:
        d = it if isinstance(it, dict) else it.__dict__
        if d.get("role") == "user":
            c = d.get("content")
            if isinstance(c, str):
                parts.append(c)
            elif isinstance(c, list):
                parts.append(" ".join(
                    (p if isinstance(p, dict) else p.__dict__).get("text", "") for p in c))
    return " ".join(parts)


def latest_user_text(input):
    if isinstance(input, str):
        return input
    for it in reversed(input):
        d = it if isinstance(it, dict) else it.__dict__
        if d.get("role") == "user":
            c = d.get("content")
            if isinstance(c, str):
                return c
            if isinstance(c, list):
                return " ".join(
                    (p if isinstance(p, dict) else p.__dict__).get("text", "") for p in c)
    return ""


def newest_item(input):
    if isinstance(input, str) or not input:
        return {}
    last = input[-1]
    return last if isinstance(last, dict) else last.__dict__


class DeskMemoModel(Model):
    async def get_response(self, system_instructions, input, model_settings, tools,
                           output_schema, handoffs, tracing, **k):
        # A lookup ran this turn: answer from its result.
        if newest_item(input).get("type") == "function_call_output":
            return ModelResponse(output=[_message(newest_item(input).get("output"))],
                                 usage=Usage(), response_id=None)
        latest = latest_user_text(input).lower()
        # Read the order id from the WHOLE history, so a later turn still has it.
        found = re.search(r"\b[A-Z]\d+\b", all_user_text(input))
        order_id = found.group(0) if found else None
        if tools and order_id and "order" in latest:      # a named order: look it up
            return ModelResponse(output=[_tool_call("lookup_order",
                '{"order_id": "%s"}' % order_id)], usage=Usage(), response_id=None)
        if order_id and any(w in latest for w in           # a follow-up about that order
                            ("ship", "arriv", "deliver", "status", "yet", "where")):
            return ModelResponse(output=[_message(
                "Your order %s shipped on 3 March and arrives 7 March." % order_id)],
                usage=Usage(), response_id=None)
        return ModelResponse(output=[_message("How can I help with your order?")],
                             usage=Usage(), response_id=None)

    async def stream_response(self, *a, **k):
        raise NotImplementedError


@function_tool
def lookup_order(order_id: str) -> str:
    "Look up an order by its id."
    return f"Order {order_id}: shipped on 3 March, arriving 7 March."


desk = Agent(name="Shop desk", instructions="Help shoppers with their orders.",
             tools=[lookup_order], model=DeskMemoModel())

# Two turns share one session: turn 2 names no order id and still answers about A17.
session = SQLiteSession("customer-1")
print("Turn 1:", Runner.run_sync(desk, "Where is my order A17?", session=session).final_output)
print("Turn 2:", Runner.run_sync(desk, "Has it shipped yet?", session=session).final_output)

# No session: the same turn 2 has no earlier order in view, so it cannot answer.
print("No-session turn 2:", Runner.run_sync(desk, "Has it shipped yet?").final_output)

# An unrelated follow-up on a remembering session does not print the order line.
other = SQLiteSession("customer-2")
Runner.run_sync(desk, "Where is my order A17?", session=other)
print("Unrelated turn 2:", Runner.run_sync(desk, "Do you sell gift cards?", session=other).final_output)