modelos-secure 0.0.0-stage → 0.1.0

package/package.json

@@ -1,6 +1,20 @@
 {
   "name": "modelos-secure",
-  "version": "0.0.0-stage",
-  "stub": true,
-  "description": "Temporary package placeholder for staged publishing"
+  "version": "0.1.0",
+  "description": "Tests AI systems for Security, Safety and Quality. One command connects a system from its code.",
+  "license": "UNLICENSED",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/mn-chetan/modelos-secure.git"
+  },
+  "type": "module",
+  "bin": {
+    "modelos-secure": "cli.js"
+  },
+  "exports": {
+    "./hook": "./hook/hook.mjs"
+  },
+  "engines": {
+    "node": ">=20"
+  }
 }

package/python/modelos_secure/clients.py

@@ -0,0 +1,444 @@
+"""The model clients: the prompt canary going in, and what the model asked for coming out.
+
+Whatever happens here, the app's own call goes ahead as if the hook were not there.
+"""
+
+import contextvars
+import functools
+import inspect
+import json
+
+_MARK = "__modelos_secure_wrapped__"
+
+# A framework's model call is reported once, by the framework's own patch, not again by the client
+# underneath it. A context variable follows the call through threads and tasks.
+_inside_framework = contextvars.ContextVar("modelos_secure_inside_framework", default=False)
+
+_canary_line = None
+_see_model_call = None
+_safely = None
+
+
+def connect(canary_line, see_model_call, safely):
+    """What this module needs from the hook, which holds the state of the current attempt."""
+    global _canary_line, _see_model_call, _safely
+    _canary_line = canary_line
+    _see_model_call = see_model_call
+    _safely = safely
+
+
+def _with_line(content, line):
+    """Adds a line to a message's content, whether it is text or a list of parts."""
+    if isinstance(content, str):
+        return f"{content}\n{line}"
+    if isinstance(content, list):
+        return [*content, {"type": "text", "text": line}]
+    return content
+
+
+def _parse_args(text):
+    if isinstance(text, dict):
+        return text
+    try:
+        parsed = json.loads(text)
+        return parsed if isinstance(parsed, dict) else {"value": parsed}
+    except (TypeError, ValueError):
+        return {"raw": str(text)}
+
+
+class _Watched:
+    """A stream the app reads as usual. What the model asked for is known once it has been read."""
+
+    def __init__(self, stream, reader, report):
+        self._stream = stream
+        self._reader = reader
+        self._report = report
+        self._iterator = None
+        self._done = False
+
+    def __getattr__(self, name):
+        return getattr(self._stream, name)
+
+    @property
+    def __class__(self):
+        return type(self._stream)
+
+    def _see(self, chunk):
+        _safely(self._reader.add, chunk)
+        return chunk
+
+    def _finish(self):
+        if not self._done:
+            self._done = True
+            _safely(self._report, self._reader.finish())
+
+    def __iter__(self):
+        self._iterator = iter(self._stream)
+        return self
+
+    def __next__(self):
+        if self._iterator is None:
+            self._iterator = iter(self._stream)
+        try:
+            return self._see(next(self._iterator))
+        except StopIteration:
+            self._finish()
+            raise
+
+    def __aiter__(self):
+        self._iterator = self._stream.__aiter__()
+        return self
+
+    async def __anext__(self):
+        if self._iterator is None:
+            self._iterator = self._stream.__aiter__()
+        try:
+            return self._see(await self._iterator.__anext__())
+        except StopAsyncIteration:
+            self._finish()
+            raise
+
+    def __enter__(self):
+        self._stream.__enter__()
+        return self
+
+    def __exit__(self, *details):
+        self._finish()
+        return self._stream.__exit__(*details)
+
+    async def __aenter__(self):
+        await self._stream.__aenter__()
+        return self
+
+    async def __aexit__(self, *details):
+        self._finish()
+        return await self._stream.__aexit__(*details)
+
+
+def _watch_call(original, prepare, read, stream_reader=None, always_streams=False):
+    """Wraps a model client's method. prepare adds the canary; read says what came back."""
+
+    def before(args, kwargs):
+        if _inside_framework.get():
+            return args, kwargs, False
+        prepared = _safely(prepare, args, kwargs, fallback=None)
+        args, kwargs = prepared if prepared else (args, kwargs)
+        return args, kwargs, True
+
+    def after(args, kwargs, result):
+        def report(found):
+            calls, input_tokens, output_tokens = found
+            _see_model_call(_model_argument(args, kwargs), calls, input_tokens, output_tokens)
+
+        if stream_reader and (always_streams or kwargs.get("stream")):
+            return _Watched(result, stream_reader(), report)
+        found = _safely(read, result)
+        if found:
+            _safely(report, found)
+        return result
+
+    if inspect.iscoroutinefunction(original):
+
+        @functools.wraps(original)
+        async def wrapper(*args, **kwargs):
+            args, kwargs, watching = before(args, kwargs)
+            result = await original(*args, **kwargs)
+            return after(args, kwargs, result) if watching else result
+
+    else:
+
+        @functools.wraps(original)
+        def wrapper(*args, **kwargs):
+            args, kwargs, watching = before(args, kwargs)
+            result = original(*args, **kwargs)
+            if not watching:
+                return result
+            if inspect.isawaitable(result):
+                return _after_awaited(args, kwargs, result, after)
+            return after(args, kwargs, result)
+
+    setattr(wrapper, _MARK, True)
+    return wrapper
+
+
+async def _after_awaited(args, kwargs, awaitable, after):
+    return after(args, kwargs, await awaitable)
+
+
+def _model_argument(args, kwargs):
+    model = kwargs.get("model")
+    return model if isinstance(model, str) else None
+
+
+def _replace(cls, attr, make):
+    original = vars(cls).get(attr)
+    if original is not None and not getattr(original, _MARK, False):
+        setattr(cls, attr, make(original))
+
+
+# openai: chat completions and responses.
+
+
+def _prepare_chat(args, kwargs):
+    line = _canary_line()
+    messages = kwargs.get("messages")
+    if not line or not isinstance(messages, list):
+        return args, kwargs
+    for index, message in enumerate(messages):
+        if isinstance(message, dict) and message.get("role") in ("system", "developer"):
+            changed = {**message, "content": _with_line(message.get("content"), line)}
+            return args, {**kwargs, "messages": [*messages[:index], changed, *messages[index + 1 :]]}
+    return args, kwargs
+
+
+def _read_chat(response):
+    if not hasattr(response, "choices") and hasattr(response, "parse"):
+        response = response.parse()
+    calls = []
+    for choice in response.choices or []:
+        for call in getattr(choice.message, "tool_calls", None) or []:
+            function = getattr(call, "function", None)
+            if function:
+                calls.append({"tool": function.name, "args": _parse_args(function.arguments)})
+    usage = getattr(response, "usage", None)
+    return calls, getattr(usage, "prompt_tokens", 0), getattr(usage, "completion_tokens", 0)
+
+
+class _ChatStreamReader:
+    def __init__(self):
+        self.calls = {}
+        self.usage = None
+
+    def add(self, chunk):
+        self.usage = getattr(chunk, "usage", None) or self.usage
+        for choice in getattr(chunk, "choices", None) or []:
+            for delta in getattr(choice.delta, "tool_calls", None) or []:
+                call = self.calls.setdefault(delta.index, {"tool": "", "text": ""})
+                function = delta.function
+                if function is not None:
+                    call["tool"] += function.name or ""
+                    call["text"] += function.arguments or ""
+
+    def finish(self):
+        calls = [{"tool": c["tool"], "args": _parse_args(c["text"])} for c in self.calls.values()]
+        usage = self.usage
+        return calls, getattr(usage, "prompt_tokens", 0), getattr(usage, "completion_tokens", 0)
+
+
+def _patch_openai_chat(module):
+    for name in ("Completions", "AsyncCompletions"):
+        cls = getattr(module, name, None)
+        if isinstance(cls, type):
+            _replace(
+                cls,
+                "create",
+                lambda original: _watch_call(original, _prepare_chat, _read_chat, _ChatStreamReader),
+            )
+
+
+def _prepare_responses(args, kwargs):
+    line = _canary_line()
+    if not line:
+        return args, kwargs
+    if isinstance(kwargs.get("instructions"), str):
+        return args, {**kwargs, "instructions": f"{kwargs['instructions']}\n{line}"}
+    items = kwargs.get("input")
+    if isinstance(items, list):
+        prepared = _prepare_chat(args, {**kwargs, "messages": items})
+        changed = dict(prepared[1])
+        changed["input"] = changed.pop("messages")
+        return args, changed
+    return args, kwargs
+
+
+def _read_responses(response):
+    if not hasattr(response, "output") and hasattr(response, "parse"):
+        response = response.parse()
+    calls = [
+        {"tool": item.name, "args": _parse_args(item.arguments)}
+        for item in response.output or []
+        if getattr(item, "type", None) == "function_call"
+    ]
+    usage = getattr(response, "usage", None)
+    return calls, getattr(usage, "input_tokens", 0), getattr(usage, "output_tokens", 0)
+
+
+class _ResponsesStreamReader:
+    def __init__(self):
+        self.found = ([], 0, 0)
+
+    def add(self, event):
+        if getattr(event, "type", None) == "response.completed":
+            self.found = _read_responses(event.response)
+
+    def finish(self):
+        return self.found
+
+
+def _patch_openai_responses(module):
+    for name in ("Responses", "AsyncResponses"):
+        cls = getattr(module, name, None)
+        if isinstance(cls, type):
+            _replace(
+                cls,
+                "create",
+                lambda original: _watch_call(
+                    original, _prepare_responses, _read_responses, _ResponsesStreamReader
+                ),
+            )
+
+
+# google.genai: every real model call, which the client's automatic function calling loops over.
+
+
+def _with_instruction(instruction, line):
+    if isinstance(instruction, str):
+        return f"{instruction}\n{line}"
+    if isinstance(instruction, list):
+        return [*instruction, line]
+    parts = getattr(instruction, "parts", None)
+    if parts:
+        return instruction.model_copy(update={"parts": [*parts, type(parts[0])(text=line)]})
+    if isinstance(instruction, dict) and isinstance(instruction.get("parts"), list):
+        return {**instruction, "parts": [*instruction["parts"], {"text": line}]}
+    return instruction
+
+
+def _prepare_genai(args, kwargs):
+    line = _canary_line()
+    config = kwargs.get("config")
+    if not line or config is None:
+        return args, kwargs
+    if isinstance(config, dict):
+        for key in ("system_instruction", "systemInstruction"):
+            if config.get(key):
+                changed = {**config, key: _with_instruction(config[key], line)}
+                return args, {**kwargs, "config": changed}
+        return args, kwargs
+    instruction = getattr(config, "system_instruction", None)
+    if instruction:
+        changed = config.model_copy(
+            update={"system_instruction": _with_instruction(instruction, line)}
+        )
+        return args, {**kwargs, "config": changed}
+    return args, kwargs
+
+
+def _read_genai(response):
+    calls = [
+        {"tool": call.name, "args": dict(call.args or {})}
+        for call in getattr(response, "function_calls", None) or []
+    ]
+    usage = getattr(response, "usage_metadata", None)
+    return (
+        calls,
+        getattr(usage, "prompt_token_count", 0),
+        getattr(usage, "candidates_token_count", 0),
+    )
+
+
+class _GenAiStreamReader:
+    def __init__(self):
+        self.calls = []
+        self.usage = (0, 0)
+
+    def add(self, chunk):
+        calls, input_tokens, output_tokens = _read_genai(chunk)
+        self.calls += calls
+        self.usage = (input_tokens or self.usage[0], output_tokens or self.usage[1])
+
+    def finish(self):
+        return self.calls, self.usage[0], self.usage[1]
+
+
+def _patch_genai(module):
+    for name in ("Models", "AsyncModels"):
+        cls = getattr(module, name, None)
+        if not isinstance(cls, type):
+            continue
+        _replace(
+            cls, "_generate_content", lambda original: _watch_call(original, _prepare_genai, _read_genai)
+        )
+        _replace(
+            cls,
+            "_generate_content_stream",
+            lambda original: _watch_call(
+                original, _prepare_genai, _read_genai, _GenAiStreamReader, always_streams=True
+            ),
+        )
+
+
+# LangChain chat models: every provider's model passes through these two.
+
+
+def _prepare_langchain_messages(messages):
+    line = _canary_line()
+    if not line or not isinstance(messages, list):
+        return messages
+    for index, message in enumerate(messages):
+        if getattr(message, "type", None) == "system":
+            changed = message.model_copy(update={"content": _with_line(message.content, line)})
+            return [*messages[:index], changed, *messages[index + 1 :]]
+    return messages
+
+
+def _read_langchain(model, result):
+    message = result.generations[0].message
+    calls = [
+        {"tool": call["name"], "args": call["args"]}
+        for call in getattr(message, "tool_calls", None) or []
+    ]
+    usage = getattr(message, "usage_metadata", None) or {}
+    _see_model_call(model, calls, usage.get("input_tokens", 0), usage.get("output_tokens", 0))
+
+
+def _langchain_model_name(chat_model):
+    return getattr(chat_model, "model_name", None) or getattr(chat_model, "model", None)
+
+
+def _patch_langchain(module):
+    cls = getattr(module, "BaseChatModel", None)
+    if not isinstance(cls, type):
+        return
+
+    def sync_call(original):
+        @functools.wraps(original)
+        def wrapper(self, messages, *args, **kwargs):
+            messages = _safely(_prepare_langchain_messages, messages, fallback=messages)
+            token = _inside_framework.set(True)
+            try:
+                result = original(self, messages, *args, **kwargs)
+            finally:
+                _inside_framework.reset(token)
+            _safely(_read_langchain, _langchain_model_name(self), result)
+            return result
+
+        setattr(wrapper, _MARK, True)
+        return wrapper
+
+    def async_call(original):
+        @functools.wraps(original)
+        async def wrapper(self, messages, *args, **kwargs):
+            messages = _safely(_prepare_langchain_messages, messages, fallback=messages)
+            token = _inside_framework.set(True)
+            try:
+                result = await original(self, messages, *args, **kwargs)
+            finally:
+                _inside_framework.reset(token)
+            _safely(_read_langchain, _langchain_model_name(self), result)
+            return result
+
+        setattr(wrapper, _MARK, True)
+        return wrapper
+
+    _replace(cls, "_generate_with_cache", sync_call)
+    _replace(cls, "_agenerate_with_cache", async_call)
+
+
+LIBRARIES = {
+    "openai.resources.chat.completions.completions": _patch_openai_chat,
+    "openai.resources.chat.completions": _patch_openai_chat,
+    "openai.resources.responses.responses": _patch_openai_responses,
+    "google.genai.models": _patch_genai,
+    "langchain_core.language_models.chat_models": _patch_langchain,
+}