참고소스 수정본

참고/guardrails-main/guardrails/formatters/__init__.py

@@ -0,0 +1,26 @@
+from guardrails.formatters.base_formatter import BaseFormatter, PassthroughFormatter
+
+try:
+    from guardrails.formatters.json_formatter import JsonFormatter
+except ImportError:
+    JsonFormatter = None
+
+
+def get_formatter(name: str, *args, **kwargs) -> BaseFormatter:
+    """Returns a class."""
+    name = name.lower()
+    if name == "jsonformer":
+        if JsonFormatter is None:
+            raise ValueError("jsonformatter requires transformers to be installed.")
+        return JsonFormatter(*args, **kwargs)
+    elif name == "none":
+        return PassthroughFormatter(*args, **kwargs)
+    raise ValueError(f"Unrecognized formatter '{name}'")
+
+
+__all__ = [
+    "get_formatter",
+    "BaseFormatter",
+    "PassthroughFormatter",
+    "JsonFormatter",
+]

참고/guardrails-main/guardrails/formatters/base_formatter.py

@@ -0,0 +1,32 @@
+from abc import ABC, abstractmethod
+
+from guardrails.llm_providers import (
+    ArbitraryCallable,
+    AsyncPromptCallableBase,
+    PromptCallableBase,
+)
+
+
+class BaseFormatter(ABC):
+    """A Formatter takes an LLM Callable and wraps the method into an abstract
+    callable.
+
+    Used to perform manipulations of the input or the output, like JSON
+    constrained- decoding.
+    """
+
+    @abstractmethod
+    def wrap_callable(self, llm_callable: PromptCallableBase) -> ArbitraryCallable: ...
+
+    @abstractmethod
+    def wrap_async_callable(
+        self, llm_callable: PromptCallableBase
+    ) -> AsyncPromptCallableBase: ...
+
+
+class PassthroughFormatter(BaseFormatter):
+    def wrap_callable(self, llm_callable: PromptCallableBase):  # type: ignore
+        return llm_callable  # Noop
+
+    def wrap_async_callable(self, llm_callable: PromptCallableBase):  # type: ignore
+        return llm_callable  # Noop

참고/guardrails-main/guardrails/formatters/json_formatter.py

@@ -0,0 +1,154 @@
+import json
+from typing import Dict, List, Optional, Union
+
+from guardrails.formatters.base_formatter import BaseFormatter
+from guardrails.llm_providers import (
+    ArbitraryCallable,
+    HuggingFacePipelineCallable,
+    HuggingFaceModelCallable,
+)
+
+
+def _deref_schema_path(schema: dict, path: Union[list, str]):
+    """Given a path like #/$defs/foo/bar/bez, nagivates into a JSONSchema dict
+    and pulls the respective sub-object."""
+    if isinstance(path, str):
+        path = path.split("/")
+    if path[0] == "#":
+        # The '#' indicates the root of the chain, so this is a first call.
+        # If we're at the root we want to make sure we have our '$defs'.
+        assert "$defs" in schema
+        return _deref_schema_path(schema, path[1:])
+    if len(path) == 1:
+        return schema[path[0]]
+    else:
+        return _deref_schema_path(schema[path[0]], path[1:])
+
+
+def _jsonschema_to_jsonformer(
+    schema: dict, path: Optional[list] = None, objdefs: Optional[dict] = None
+) -> dict:
+    """Converts the large-ish JSONSchema standard into the JSONFormer schema
+    format.
+
+    These are mostly identical, but the jsonschema supports '$defs' and
+    '$ref'. There's an additional inconsistency in the use 'integer'
+    versus 'number'.
+    """
+    # TODO: Can we use jsonref to replace references, since it's a dependency anyway?
+    if path is None:
+        path = []
+    if objdefs is None:
+        objdefs = {"$defs": dict()}
+
+    # We may get something we don't expect...
+    if not isinstance(schema, dict):
+        raise Exception(
+            f"Error: could not convert/parse base schema. Encountered `{schema}`"
+        )
+
+    if "$defs" in schema:
+        # We have some sub-schemas defined here.  We need to convert them.
+        # We may also need to handle sub-schema defs.
+        # For now, build a quick tree in the defs.
+        current = objdefs["$defs"]
+        for step in path:
+            if step not in current:
+                current[step] = dict()
+            current = current[step]
+        current.update(schema["$defs"])
+
+    result = dict()
+    for k, v in schema.items():
+        # Convert {"type": "integer"} to {"type": "number"} float is already 'number'.
+        if k == "type" and v == "integer":
+            result["type"] = "number"
+        elif k == "type" and v == "object":
+            result["type"] = "object"
+            result["properties"] = dict()
+            for subkey, subvalue in schema["properties"].items():  # Must be present.
+                path.append(subkey)
+                result["properties"][subkey] = _jsonschema_to_jsonformer(
+                    subvalue,
+                    path,
+                    objdefs,
+                )
+                assert path.pop() == subkey
+        elif k == "type" and v == "array":
+            result["type"] = "array"
+            result["items"] = _jsonschema_to_jsonformer(schema["items"], path, objdefs)
+        elif k == "$ref":
+            result = _jsonschema_to_jsonformer(
+                _deref_schema_path(objdefs, v), path, objdefs
+            )
+        else:
+            result[k] = v
+    return result
+
+
+class JsonFormatter(BaseFormatter):
+    """A formatter that uses Jsonformer to ensure the shape of structured data
+    for Hugging Face models."""
+
+    def __init__(self, schema: dict):
+        self.output_schema = _jsonschema_to_jsonformer(schema)
+
+    def wrap_callable(self, llm_callable) -> ArbitraryCallable:
+        # JSON Schema enforcement experiment.
+        from jsonformer import Jsonformer
+
+        if isinstance(llm_callable, HuggingFacePipelineCallable):
+            model = llm_callable.init_kwargs["pipeline"]
+
+            def fn(
+                *args,
+                messages: Optional[List[Dict[str, str]]] = None,
+                **kwargs,
+            ) -> str:
+                prompt = ""
+                for msg in messages:  # type: ignore
+                    prompt += msg["content"]
+
+                return json.dumps(
+                    Jsonformer(
+                        model=model.model,
+                        tokenizer=model.tokenizer,
+                        json_schema=self.output_schema,
+                        prompt=prompt,
+                    )()
+                )
+
+            return ArbitraryCallable(fn)
+        elif isinstance(llm_callable, HuggingFaceModelCallable):
+            # This will not work because 'model_generate' is the .gen method.
+            # model = self.api.init_kwargs["model_generate"]
+            # Use the __self__ to grab the base mode for passing into JF.
+            model = llm_callable.init_kwargs["model_generate"].__self__
+            tokenizer = llm_callable.init_kwargs["tokenizer"]
+
+            def fn(
+                *args,
+                messages: Optional[List[Dict[str, str]]] = None,
+                **kwargs,
+            ) -> str:
+                prompt = ""
+                for msg in messages:  # type: ignore
+                    prompt += msg["content"]
+
+                return json.dumps(
+                    Jsonformer(
+                        model=model,
+                        tokenizer=tokenizer,
+                        json_schema=self.output_schema,
+                        prompt=prompt,
+                    )()
+                )
+
+            return ArbitraryCallable(fn)
+        else:
+            raise ValueError(
+                "JsonFormatter can only be used with HuggingFace*Callable."
+            )
+
+    def wrap_async_callable(self, llm_callable):
+        raise NotImplementedError()