modelos-secure 0.0.0-stage → 0.1.0

package/python/modelos_secure/hook.py

@@ -0,0 +1,542 @@
+"""The hooks for Python apps: the same behaviour and control protocol as hooks/node/src/hook.mjs.
+
+Standard library only, Python 3.10+. Customers run this, not us, so when it meets something it
+does not understand it leaves the app's own behaviour alone. docs/spec/05-scan-and-cli.md#5-start
+"""
+
+import copy
+import functools
+import http.server
+import inspect
+import json
+import os
+import random
+import socket
+import sys
+import threading
+import types
+import urllib.parse
+
+from modelos_secure import clients
+
+_MARK = "__modelos_secure_wrapped__"
+
+_tools = {}
+_ways_in = {}
+_files = set()
+_wrapped = []
+_noted = set()
+_state = None
+
+
+def _fresh_state():
+    return {
+        "attempt": None,
+        "held": [],
+        "modelCalls": [],
+        "ran": [],
+        "plantDelivered": False,
+        "retrieved": [],
+    }
+
+
+def _note(error):
+    text = f"modelos-secure hook: {error}"
+    if text not in _noted:
+        _noted.add(text)
+        print(text, file=sys.stderr)
+
+
+def _safely(fn, *args, fallback=None):
+    """Runs a piece of the hook. If it fails, the app carries on as if the hook were not there."""
+    try:
+        return fn(*args)
+    except Exception as error:
+        _note(repr(error))
+        return fallback
+
+
+# Planted content and canaries, into whatever shape the way in returns.
+
+
+def _longest_string_key(fields):
+    best = None
+    for key, value in fields.items():
+        if isinstance(value, str) and (best is None or len(value) > len(fields[best])):
+            best = key
+    return best
+
+
+def _plant_into(result, text):
+    """Returns result itself, so nothing is delivered, when it is a shape with no place for text."""
+    if result is None:
+        return text
+    if isinstance(result, str):
+        return f"{result}\n\n{text}"
+    if isinstance(result, (list, tuple)):
+        if len(result) == 0 or isinstance(result[0], str):
+            return [*result, text]
+        planted = _plant_into(result[0], text)
+        return result if planted is result[0] else [*result, planted]
+    if isinstance(result, dict):
+        for key, value in result.items():
+            if isinstance(value, list):
+                return {**result, key: _plant_into(value, text)}
+        key = _longest_string_key(result)
+        return {**result, key: f"{result[key]}\n\n{text}"} if key else {**result, "note": text}
+    fields = getattr(result, "__dict__", None)
+    if fields:
+        # Objects with text in them, such as a LangChain Document.
+        key = _longest_string_key(fields)
+        if key:
+            planted = copy.copy(result)
+            setattr(planted, key, f"{fields[key]}\n\n{text}")
+            return planted
+    return result
+
+
+def _add_canary(result, record):
+    if isinstance(result, str):
+        return f"{result}\nCustomer reference: {record}"
+    if isinstance(result, dict):
+        return {**result, "customer_reference": record}
+    return result
+
+
+def _passage_text(result):
+    """Q2 grades answers against what was really read, so this is the text and not the way in's name."""
+    try:
+        if isinstance(result, str):
+            text = result
+        elif isinstance(result, (list, tuple)):
+            text = "\n".join(str(getattr(item, "page_content", item)) for item in result)
+        else:
+            text = str(getattr(result, "page_content", result))
+    except Exception:  # A reading we cannot understand must never break the customer's app.
+        return ""
+    return text[:8000]
+
+
+def _after_run(name, result):
+    attempt = _state["attempt"]
+    way_in = _ways_in.get(name)
+    out = result
+    if way_in is not None and attempt:
+        _state["retrieved"].append(_passage_text(result))
+        plant = attempt.get("plant")
+        if plant and plant.get("wayIn") == way_in["name"]:
+            out = _plant_into(out, plant["text"])
+            _state["plantDelivered"] = out is not result
+    if attempt and _tools.get(name, {}).get("can") == "read private data":
+        out = _add_canary(out, attempt["canaries"]["record"])
+    return out
+
+
+def _finish(name, result):
+    # A generator hands its content out later, where nothing here can change it.
+    if inspect.isgenerator(result) or inspect.isasyncgen(result):
+        return result
+    return _safely(_after_run, name, result, fallback=result)
+
+
+# The one place a tool call passes through. Held tools never run.
+
+
+def _args_object(entry, args, kwargs):
+    if len(args) == 1 and isinstance(args[0], dict) and not kwargs:
+        return dict(args[0])
+    if entry["signature"] is not None:
+        bound = entry["signature"].bind_partial(*args, **kwargs)
+        out = {}
+        for key, value in bound.arguments.items():
+            kind = entry["signature"].parameters[key].kind
+            if kind is inspect.Parameter.VAR_KEYWORD:
+                out.update(value)
+            elif key not in ("self", "cls"):
+                out[key] = value
+        return out
+    params = _tools.get(entry["name"], {}).get("params") or []
+    named = {params[i] if i < len(params) else f"arg{i}": arg for i, arg in enumerate(args)}
+    return {**named, **kwargs}
+
+
+def _hold(entry, args, kwargs):
+    """A realistic answer, so the app carries on as if the action happened."""
+    name = entry["name"]
+    recorded = _safely(_args_object, entry, args, kwargs, fallback={"args": list(args), **kwargs})
+    _state["held"].append({"tool": name, "args": recorded})
+    return {"status": "ok", "id": f"{name}-{random.randint(10000, 99999)}"}
+
+
+def _is_held(name):
+    return _tools.get(name, {}).get("mode") == "hold"
+
+
+def _call(key, args, kwargs):
+    entry = _wrapped[key]
+    name = entry["name"]
+    if _is_held(name):
+        held = _hold(entry, args, kwargs)
+        if inspect.isasyncgenfunction(entry["original"]):
+            return _async_once(held)
+        return iter([held]) if inspect.isgeneratorfunction(entry["original"]) else held
+    if name in _tools:
+        _state["ran"].append(name)
+    result = entry["original"](*args, **kwargs)
+    if inspect.isawaitable(result):
+        return _finish_later(name, result)
+    return _finish(name, result)
+
+
+async def _acall(key, args, kwargs):
+    entry = _wrapped[key]
+    name = entry["name"]
+    if _is_held(name):
+        return _hold(entry, args, kwargs)
+    if name in _tools:
+        _state["ran"].append(name)
+    return _finish(name, await entry["original"](*args, **kwargs))
+
+
+async def _finish_later(name, awaitable):
+    return _finish(name, await awaitable)
+
+
+async def _async_once(value):
+    yield value
+
+
+# Wrapping. Replacing the function's code in place, rather than the name in its module, means that
+# every reference to the function now runs the wrapper: other modules that imported it, a
+# {"send_email": send_email} table, the list given to a model client, a LangChain tool.
+
+
+def _signature_of(fn):
+    try:
+        return inspect.signature(fn)
+    except (TypeError, ValueError, NameError):
+        return None
+
+
+def _register(name, original, signature):
+    _wrapped.append({"name": name, "original": original, "signature": signature})
+    return len(_wrapped) - 1
+
+
+_TRAMPOLINE = """def _(*args, **kwargs):
+    return __import__("sys").modules["modelos_secure.hook"]._call({key}, args, kwargs)
+"""
+_ASYNC_TRAMPOLINE = """async def _(*args, **kwargs):
+    return await __import__("sys").modules["modelos_secure.hook"]._acall({key}, args, kwargs)
+"""
+
+
+def _swap_code(name, fn):
+    signature = _signature_of(fn)
+    original = types.FunctionType(
+        fn.__code__, fn.__globals__, fn.__name__, fn.__defaults__, fn.__closure__
+    )
+    original.__kwdefaults__ = fn.__kwdefaults__
+    key = _register(name, original, signature)
+    source = _ASYNC_TRAMPOLINE if inspect.iscoroutinefunction(fn) else _TRAMPOLINE
+    module_code = compile(source.format(key=key), fn.__code__.co_filename, "exec")
+    code = next(c for c in module_code.co_consts if isinstance(c, types.CodeType))
+    fn.__code__ = code.replace(co_firstlineno=fn.__code__.co_firstlineno, co_name=fn.__name__)
+    fn.__defaults__ = None
+    fn.__kwdefaults__ = None
+    if signature is not None:
+        # Frameworks read a tool's arguments from its signature, so it has to stay the same.
+        fn.__signature__ = signature
+    setattr(fn, _MARK, True)
+
+
+def _outer_wrapper(name, fn):
+    """For what cannot be changed in place: decorated functions with closures, partials, objects."""
+    key = _register(name, fn, _signature_of(fn))
+    if inspect.iscoroutinefunction(fn):
+
+        @functools.wraps(fn)
+        async def wrapper(*args, **kwargs):
+            return await _acall(key, args, kwargs)
+
+    else:
+
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            return _call(key, args, kwargs)
+
+    setattr(wrapper, _MARK, True)
+    return wrapper
+
+
+def _wrap_callable(name, value):
+    """Returns what should now stand where value stood."""
+    if not callable(value) or getattr(value, _MARK, False):
+        return value
+    if isinstance(value, types.FunctionType) and not value.__code__.co_freevars:
+        _swap_code(name, value)
+        return value
+    return _outer_wrapper(name, value)
+
+
+def _listed(name):
+    return name in _tools or name in _ways_in
+
+
+def _is_langchain_tool(value):
+    return any(c.__module__.startswith("langchain_core.tools") for c in type(value).__mro__)
+
+
+def _wrap_langchain_tool(tool):
+    name = getattr(tool, "name", None)
+    if not isinstance(name, str) or not _listed(name):
+        return
+    for field in ("func", "coroutine"):
+        fn = getattr(tool, field, None)
+        if fn is not None:
+            wrapped = _wrap_callable(name, fn)
+            if wrapped is not fn:
+                setattr(tool, field, wrapped)
+
+
+def _scan_class(cls):
+    for attr, value in list(vars(cls).items()):
+        if not _listed(attr):
+            continue
+        inner = value.__func__ if isinstance(value, (staticmethod, classmethod)) else value
+        wrapped = _wrap_callable(attr, inner)
+        if wrapped is not inner:
+            setattr(cls, attr, type(value)(wrapped) if inner is not value else wrapped)
+
+
+def _scan_module(module):
+    namespace = vars(module)
+    for attr, value in list(namespace.items()):
+        if attr.startswith("__"):
+            continue
+        if isinstance(value, type):
+            if value.__module__ == module.__name__:
+                _safely(_scan_class, value)
+        elif _is_langchain_tool(value):
+            _safely(_wrap_langchain_tool, value)
+        elif _listed(attr):
+            wrapped = _safely(_wrap_callable, attr, value, fallback=value)
+            if wrapped is not value:
+                namespace[attr] = wrapped
+
+
+def _is_app_file(path):
+    if not path or not path.endswith(".py"):
+        return False
+    if _files:
+        return os.path.realpath(path) in _files
+    return "site-packages" not in path and not path.startswith(sys.base_prefix)
+
+
+def _scan_main():
+    """The script Python was started with never finishes loading, so it is scanned when a review begins."""
+    main = sys.modules.get("__main__")
+    if main is not None and _is_app_file(getattr(main, "__file__", None)):
+        _scan_module(main)
+
+
+# Model calls: what the client patches need from here.
+
+
+def _canary_line():
+    attempt = _state["attempt"]
+    canary = attempt["canaries"].get("prompt") if attempt else None
+    return f"Internal reference: {canary}" if canary else None
+
+
+def _see_model_call(model, tool_calls, input_tokens, output_tokens):
+    _state["modelCalls"].append(
+        {
+            "model": model,
+            "toolCalls": tool_calls,
+            "inputTokens": input_tokens or 0,
+            "outputTokens": output_tokens or 0,
+        }
+    )
+
+
+# Modules as they load.
+
+
+def _after_load(module):
+    patch = clients.LIBRARIES.get(module.__name__)
+    if patch:
+        _safely(patch, module)
+    if _is_app_file(getattr(module, "__file__", None)):
+        _safely(_scan_module, module)
+
+
+class _Finder:
+    """Finds nothing itself. It only adds a step after the modules we care about have loaded."""
+
+    @staticmethod
+    def find_spec(name, path=None, target=None):
+        for finder in sys.meta_path:
+            if finder is _Finder or not hasattr(finder, "find_spec"):
+                continue
+            spec = finder.find_spec(name, path, target)
+            if spec is not None:
+                _safely(_watch_spec, spec)
+                return spec
+        return None
+
+
+def _watch_spec(spec):
+    if spec.name not in clients.LIBRARIES and not _is_app_file(spec.origin):
+        return
+    original = spec.loader.exec_module
+
+    def exec_module(module):
+        original(module)
+        _after_load(module)
+
+    spec.loader.exec_module = exec_module
+
+
+# The control port the runner talks to. Loopback only, and it dies with the app.
+
+
+class _Control(http.server.BaseHTTPRequestHandler):
+    def _send(self, body):
+        data = json.dumps(body, default=str).encode("utf-8")
+        self.send_response(200)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", str(len(data)))
+        self.end_headers()
+        self.wfile.write(data)
+
+    def _handle(self):
+        global _state
+        length = int(self.headers.get("content-length") or 0)
+        body = self.rfile.read(length) if length else b""
+        if self.path == "/begin":
+            _safely(_scan_main)
+            _state = _fresh_state()
+            _state["attempt"] = json.loads(body.decode("utf-8"))
+            return self._send({"ok": True})
+        if self.path == "/end":
+            seen = _state
+            _state = _fresh_state()
+            return self._send(
+                {
+                    "heldCalls": seen["held"],
+                    "modelCalls": seen["modelCalls"],
+                    "ranTools": seen["ran"],
+                    "plantDelivered": seen["plantDelivered"],
+                    "retrieved": seen["retrieved"],
+                }
+            )
+        if self.path == "/hello":
+            return self._send({"ok": True, "hook": "python"})
+        self.send_response(404)
+        self.send_header("content-length", "0")
+        self.end_headers()
+        return None
+
+    do_GET = _handle
+    do_POST = _handle
+
+    def log_message(self, *details):
+        pass
+
+
+def _announce(port):
+    """A plain socket, because a server may fork while this runs in its other thread. A forked child
+    keeps any lock urllib held at that moment, with nothing left to release it, and hangs."""
+    runner = urllib.parse.urlsplit(os.environ.get("MODELOS_SECURE_RUNNER_URL", ""))
+    if not runner.hostname:
+        return
+    body = json.dumps({"instance": os.environ.get("MODELOS_SECURE_INSTANCE"), "port": port})
+    request = (
+        f"POST /ready HTTP/1.1\r\nhost: {runner.netloc}\r\ncontent-type: application/json\r\n"
+        f"content-length: {len(body)}\r\nconnection: close\r\n\r\n{body}"
+    )
+    with socket.create_connection((runner.hostname, runner.port), timeout=5) as connection:
+        connection.sendall(request.encode("utf-8"))
+        connection.recv(1024)
+
+
+def _serve():
+    server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _Control)
+    server.daemon_threads = True
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+    threading.Thread(
+        target=_safely, args=(_announce, server.server_address[1]), daemon=True
+    ).start()
+
+
+def _serve_once_app_port_is_taken():
+    """Only the process that takes the app's port is the app.
+
+    PYTHONPATH reaches every Python process the app starts, and those are not the app: gcloud, run
+    by google-auth while it finds credentials, would otherwise announce a control port and exit,
+    leaving the runner talking to a dead port. docs/journal/adr/0020-announce-on-the-app-port.md
+    """
+    wanted = os.environ.get("PORT", "")
+    if not wanted.isdigit():
+        _serve()
+        return
+    wanted = int(wanted)
+    claimed = False
+
+    def claim(address):
+        nonlocal claimed
+        if not claimed and isinstance(address, tuple) and address[1] == wanted:
+            claimed = True
+            _serve()
+
+    init = socket.socket.__init__
+    bind = socket.socket.bind
+    listen = socket.socket.listen
+
+    @functools.wraps(init)
+    def watched_init(self, *args, **kwargs):
+        init(self, *args, **kwargs)
+        # Flask's reloader hands its child a socket that is already bound and listening.
+        fileno = kwargs.get("fileno", args[3] if len(args) > 3 else None)
+        if fileno is not None and not claimed:
+            _safely(lambda: claim(self.getsockname()))
+
+    @functools.wraps(bind)
+    def watched_bind(self, address):
+        bind(self, address)
+        _safely(claim, address)
+
+    @functools.wraps(listen)
+    def watched_listen(self, *args):
+        listen(self, *args)
+        # A reloading server binds in one process and listens in another.
+        _safely(lambda: claim(self.getsockname()))
+
+    socket.socket.__init__ = watched_init
+    socket.socket.bind = watched_bind
+    socket.socket.listen = watched_listen
+
+    def after_fork():
+        # A forked server worker cannot use the parent's thread, so it serves its own control port.
+        # Any other fork, such as a pool of workers, is not the app.
+        nonlocal claimed
+        if claimed:
+            _serve()
+
+    os.register_at_fork(after_in_child=lambda: _safely(after_fork))
+
+
+def install():
+    global _state
+    config = json.loads(os.environ["MODELOS_SECURE_HOOK"])
+    _state = _fresh_state()
+    _tools.update({tool["name"]: tool for tool in config.get("tools", [])})
+    _ways_in.update({way["function"]: way for way in config.get("waysIn", []) if way.get("function")})
+    _files.update(os.path.realpath(path) for path in config.get("files", []))
+    clients.connect(_canary_line, _see_model_call, _safely)
+    sys.meta_path.insert(0, _Finder)
+    for name in list(clients.LIBRARIES):
+        if name in sys.modules:
+            _after_load(sys.modules[name])
+    _serve_once_app_port_is_taken()

package/python/sitecustomize.py

@@ -0,0 +1,13 @@
+# Python imports this module before anything else when its folder is on PYTHONPATH.
+# docs/spec/05-scan-and-cli.md#5-start. Outside a review it does nothing, and nothing here may ever
+# stop the customer's app from starting.
+import os
+import sys
+
+if os.environ.get("MODELOS_SECURE_HOOK") and sys.version_info >= (3, 10):
+    try:
+        from modelos_secure import hook
+
+        hook.install()
+    except Exception as error:
+        print(f"modelos-secure hook: not attached: {error}", file=sys.stderr)