crystil 0.8.0__py3-none-any.whl

crystil/__init__.py

@@ -0,0 +1,102 @@
+r"""
+  ___             _   _ _
+ / __|_ _ _  _ __| |_(_) |
+| (__| '_| || (_-<  _| | |
+ \___|_|  \_, /__/\__|_|_|AI
+          |__/                             07312025 / optimus codex
+"""
+
+import os
+from uuid import uuid4
+
+from crystil._config import Config
+from crystil._errors import CrystilRequestInterceptedError
+from crystil._providers import Anthropic as LlmProviderAnthropic
+from crystil._providers import Google as LlmProviderGoogle
+from crystil._providers import Groq as LlmProviderGroq
+from crystil._providers import LangChain as LlmProviderLangChain
+from crystil._providers import OpenAi as LlmProviderOpenAi
+from crystil._providers import PydanticAi as LlmProviderPydanticAi
+from crystil._sentinel import Sentinel
+from crystil.api._workflow import Workflow, Workflows
+
+__all__ = ["Crystil", "CrystilRequestInterceptedError"]
+
+
+class Crystil:
+    def __init__(self, api_key=None):
+        if api_key is None:
+            api_key = os.environ.get("CRYSTIL_API_KEY", None)
+
+        if api_key is None:
+            raise RuntimeError(
+                "API key is missing. Either set the CRYSTIL_API_KEY environment "
+                + "variable or set the api_key parameter when instantiating Crystil."
+            )
+
+        self.config = Config()
+        self.config.api_key = api_key
+        self.config.tx_uuid = uuid4()
+        self.sentinel = Sentinel(self.config)
+
+        self.anthropic = LlmProviderAnthropic(self)
+        self.google = LlmProviderGoogle(self)
+        self.groq = LlmProviderGroq(self)
+        self.langchain = LlmProviderLangChain(self)
+        self.openai = LlmProviderOpenAi(self)
+        self.pydantic_ai = LlmProviderPydanticAi(self)
+
+        self.workflow = Workflow(self.config)
+        self.workflows = Workflows(self.config)
+
+    def attribution(
+        self,
+        parent_id=None,
+        parent_name=None,
+        subsidiary_id=None,
+        subsidiary_name=None,
+        # -- Deprecated parameters! They are here for backwards compatibility only.
+        parent_uuid=None,
+        subsidiary_uuid=None,
+    ):
+        if parent_id is None:
+            raise RuntimeError("a string parent_id is required")
+
+        parent_id = str(parent_id)
+
+        if len(parent_id) > 100:
+            raise RuntimeError("parent_id cannot be greater than 100 characters")
+
+        if parent_name is not None and len(parent_name) > 100:
+            raise RuntimeError("parent_name cannot be greater than 100 characters")
+
+        if subsidiary_name is not None and subsidiary_id is None:
+            raise RuntimeError(
+                "a string subsidiary_id is required if a subsidiary_name is provided"
+            )
+
+        if subsidiary_id is not None:
+            subsidiary_id = str(subsidiary_id)
+
+            if len(subsidiary_id) > 100:
+                raise RuntimeError(
+                    "subsidiary_id cannot be greater than 100 characters"
+                )
+
+        if subsidiary_name is not None and len(subsidiary_name) > 100:
+            raise RuntimeError("subsidiary_name cannot be greater than 100 characters")
+
+        subsidiary = None
+        if subsidiary_id is not None:
+            subsidiary = {"id": subsidiary_id, "name": subsidiary_name}
+
+        self.config.attribution = {
+            "parent": {"id": parent_id, "name": parent_name},
+            "subsidiary": subsidiary,
+        }
+
+        return self
+
+    def new_transaction(self):
+        self.config.tx_uuid = uuid4()
+        return self

crystil/_base.py

@@ -0,0 +1,443 @@
+r"""
+  ___             _   _ _
+ / __|_ _ _  _ __| |_(_) |
+| (__| '_| || (_-<  _| | |
+ \___|_|  \_, /__/\__|_|_|AI
+          |__/                             07312025 / optimus codex
+"""
+
+import base64
+import copy
+import json
+from typing import Any, Callable, Optional
+
+from google.protobuf import json_format
+
+from crystil._config import Config
+from crystil._constants import (
+    LANGCHAIN_CHATBEDROCK_CLIENT_TITLE,
+    LANGCHAIN_CLIENT_PROVIDER,
+)
+from crystil._errors import CrystilRequestInterceptedError
+from crystil._network import Api
+from crystil._utils import merge_chunk
+
+# These are attributes that are in ProtoBuf, but since they
+# are recursive they cause issues making into JSON. Also, they
+# are not part of the data we care about.
+PROTOBUF_SKIP_ATTRS = frozenset(
+    [
+        "_client",
+        "_http",
+        "_session",
+        "_transport",
+        "__objclass__",  # Descriptor metadata
+        "__doc__",  # Documentation
+        "_member_map_",  # Enum metadata
+        "_value2member_map_",  # Enum metadata
+        "_member_names_",  # Enum metadata
+    ]
+)
+
+
+class BaseClient:
+    def __init__(self, config: Config):
+        self.config = config
+        self.stream = False
+
+
+class BaseInvoke:
+    def __init__(self, config: Config, method):
+        self.config = config
+        self._method = method
+        self._client_provider = None
+        self._client_title = None
+        self._client_version = None
+        self._client_title_resolver: Optional[Callable[[dict[str, Any]], str]] = None
+        self._uses_protobuf = False
+        self._include_stream_options = False
+        self._stream_options_via_extra_body: bool = False
+        self._stream_response_extractor = None
+
+    def client_is_bedrock(self):
+        return (
+            self._client_provider == LANGCHAIN_CLIENT_PROVIDER
+            and self._client_title == LANGCHAIN_CHATBEDROCK_CLIENT_TITLE
+        )
+
+    def configure_for_streaming_usage(self, kwargs):
+        if self._include_stream_options:
+            if kwargs.get("stream", None) == True:
+                if self._stream_options_via_extra_body:
+                    # Groq's Python SDK doesn't expose `stream_options` as a
+                    # top-level kwarg (the function signature rejects it with
+                    # TypeError), but its underlying HTTP API supports it.
+                    # `extra_body` is the SDK's documented escape hatch for
+                    # arbitrary additional request-body fields.
+                    extra_body = kwargs.get("extra_body", None)
+                    if extra_body is None or not isinstance(extra_body, dict):
+                        kwargs["extra_body"] = {}
+                    stream_options = kwargs["extra_body"].get("stream_options", None)
+                    if stream_options is None or not isinstance(stream_options, dict):
+                        kwargs["extra_body"]["stream_options"] = {}
+                    kwargs["extra_body"]["stream_options"]["include_usage"] = True
+                else:
+                    stream_options = kwargs.get("stream_options", None)
+                    if stream_options is None or not isinstance(
+                        kwargs["stream_options"], dict
+                    ):
+                        kwargs["stream_options"] = {}
+
+                    kwargs["stream_options"]["include_usage"] = True
+
+        return kwargs
+
+    def dict_to_json(self, dict_):
+        result = {}
+        for key, value in dict_.items():
+            if key not in PROTOBUF_SKIP_ATTRS:
+                if isinstance(value, list):
+                    result[key] = self.list_to_json(value)
+                elif isinstance(value, dict):
+                    result[key] = self.dict_to_json(value)
+                else:
+                    if hasattr(value, "__dict__"):
+                        result[key] = self.dict_to_json(value.__dict__)
+                    elif isinstance(value, bytes):
+                        # If left as bytes, the payload sent to Collector won't work.
+                        # Gemini 3.0 passes a though_signature field is binary.
+                        result[key] = base64.b64encode(value).decode("utf-8")
+                    else:
+                        result[key] = value
+
+        return result
+
+    def _format_kwargs(self, kwargs):
+        formatted_kwargs = None
+
+        if self._uses_protobuf:
+            request = kwargs.get("request")
+            if request:
+                if kwargs["request"].__dict__ and kwargs["request"].__dict__.get("_pb"):
+                    # The "_pb" is for handling the protobuf structure.  It is used
+                    # in the LangChain Google code paths.
+                    formatted_kwargs = json.loads(
+                        json_format.MessageToJson(kwargs["request"].__dict__["_pb"])
+                    )
+            else:
+                formatted_kwargs = self.dict_to_json(copy.deepcopy(kwargs))
+        else:
+            formatted_kwargs = copy.deepcopy(kwargs)
+            if self.provider_is_langchain():
+                if "response_format" in formatted_kwargs and isinstance(
+                    formatted_kwargs["response_format"], object
+                ):
+                    """
+                    We are likely processing the result of LangChain's structured
+                    output runnable. The object defined in "response_format" is
+                    recursive (it refers to itself) so formatting it into a dictionary
+                    will result in an RecursionError. We also do not need the data in
+                    this object so we are going to discard it here.
+                    """
+
+                    del formatted_kwargs["response_format"]
+
+            formatted_kwargs = self.dict_to_json(formatted_kwargs)
+
+        return formatted_kwargs
+
+    def _format_payload(
+        self,
+        client_provider,
+        client_title,
+        client_version,
+        start_time,
+        end_time,
+        query,
+        response,
+    ):
+        response_json = self.response_to_json(response)
+
+        payload = {
+            "attribution": self.config.attribution,
+            "conversation": {
+                "client": {
+                    "provider": client_provider,
+                    "title": client_title,
+                    "version": client_version,
+                },
+                "query": query,
+                "response": response_json,
+            },
+            "meta": {
+                "api": {"key": self.config.api_key},
+                "fnfg": {
+                    "exc": None,
+                    "status": "succeeded",
+                },
+                "sdk": {"client": "python", "version": self.config.version},
+            },
+            "time": {"end": end_time, "start": start_time},
+            "tx": {"uuid": str(self.config.tx_uuid)},
+        }
+
+        return payload
+
+    def _format_response(self, raw_response):
+        formatted_response = copy.deepcopy(raw_response)
+        if self._uses_protobuf:
+            # If _uses_protobuf is true, it means the response needs special formatting.
+            # Below the _pb case is for protobuf, and the model_dump is for Pydantic which
+            # is used in the v4 version of google-genai.
+            if not isinstance(formatted_response, (list, dict)):
+                # The "_pb" is for handling the protobuf structure.  It is used
+                # in the LangChain Google code paths.
+                if formatted_response.__dict__ and formatted_response.__dict__.get(
+                    "_pb"
+                ):
+                    formatted_response = json.loads(
+                        json_format.MessageToJson(formatted_response.__dict__["_pb"])
+                    )
+                elif hasattr(formatted_response, "model_dump"):
+                    # New path: Pydantic-based response (google-genai >= v4)
+                    formatted_response = formatted_response.model_dump(mode="json")
+
+        return formatted_response
+
+    def _format_intercept_payload(
+        self,
+        client_provider: Optional[str],
+        client_title: str,
+        client_version: str,
+        kwargs: dict,
+    ) -> dict:
+        return {
+            "attribution": self.config.attribution,
+            "conversation": {
+                "client": {
+                    "provider": client_provider,
+                    "title": client_title,
+                    "version": client_version,
+                },
+                "request": self._format_kwargs(kwargs),
+            },
+            "meta": {
+                "sdk": {"client": "python", "version": self.config.version},
+            },
+        }
+
+    def get_response_content(self, raw_response):
+        if (
+            raw_response.__class__.__name__ == "LegacyAPIResponse"
+            and raw_response.__class__.__module__ == "openai._legacy_response"
+        ):
+            """
+            Library: langchain-openai
+            Version: > 0.3.31
+
+            Calling the chat / invoke method of the client no longer returns the JSON
+            response but instead an object that looks like an API response. This
+            object does not inherit from a base class we can reliably identify and
+            we do not want to force the OpenAI library as a dependency.
+            """
+
+            return json.loads(raw_response.text)
+
+        return raw_response
+
+    def list_to_json(self, list_):
+        result = []
+        for entry in list_:
+            if isinstance(entry, list):
+                result.append(self.list_to_json(entry))
+            elif isinstance(entry, dict):
+                result.append(self.dict_to_json(entry))
+            else:
+                if hasattr(entry, "__dict__"):
+                    # This is used in the uses_protobuf cases with LangChain Google.
+                    result.append(self.dict_to_json(entry.__dict__))
+                elif isinstance(entry, bytes):
+                    # If left as bytes, the payload sent to Collector won't work.
+                    result.append(base64.b64encode(entry).decode("utf-8"))
+                else:
+                    result.append(entry)
+
+        return result
+
+    def _make_relevance_intercept_request(self, **kwargs) -> tuple[bool, Optional[str]]:
+        response = Api(self.config).post(
+            "sentinel/relevance/intercept",
+            self._format_intercept_payload(
+                self._client_provider,
+                self._resolve_client_title(kwargs),
+                self._client_version,
+                kwargs,
+            ),
+            retryable=False,
+            timeout=self.config.secs_irrelevant_request_timeout,
+        )
+
+        # Default to relevant if the response is missing expected boolean.
+        relevant: bool = response.get("relevant", True)
+        reason: Optional[str] = None
+
+        if not relevant:
+            reason = response.get("reason", "Irrelevant request blocked.")
+
+        return relevant, reason
+
+    def provider_is_langchain(self):
+        return self._client_provider == LANGCHAIN_CLIENT_PROVIDER
+
+    def _raise_if_irrelevant(self, **kwargs):
+        try:
+            relevant, reason = self._make_relevance_intercept_request(**kwargs)
+        except Exception as e:
+            # The caller should be unaware of any unexpected
+            # errors that may occur.
+            return
+
+        if not relevant:
+            raise CrystilRequestInterceptedError(reason)
+
+    def response_to_json(self, response):
+        data = response
+        if isinstance(data, list):
+            result = self.list_to_json(data)
+        else:
+            if not isinstance(data, dict):
+                data = response.__dict__
+
+            result = {}
+
+            for key, value in data.items():
+                if key not in PROTOBUF_SKIP_ATTRS:
+                    if isinstance(value, list):
+                        result[key] = self.list_to_json(value)
+                    elif isinstance(value, dict):
+                        result[key] = self.dict_to_json(value)
+                    else:
+                        if hasattr(value, "__dict__"):
+                            result[key] = self.dict_to_json(value.__dict__)
+                        elif isinstance(value, bytes):
+                            # If left as bytes, the payload sent to Collector won't work
+                            result[key] = base64.b64encode(value).decode("utf-8")
+                        else:
+                            result[key] = value
+
+        return result
+
+    def set_client(self, provider, title, version):
+        self._client_provider = provider
+        self._client_title = title
+        self._client_version = version
+        return self
+
+    # Override the static client title with a per-call value derived from the
+    # request kwargs. Used by Groq, where the framework hosts many third-party
+    # models and the model-ID prefix carries the LLM family
+    # (`meta-llama/...` → `meta-llama`).
+    #
+    # Implementations are responsible for picking a sensible fallback
+    # (typically the framework name) when the kwargs don't yield a useful
+    # prefix. See `_clients.py::groq_title_from_model` for the canonical
+    # example.
+    def title_from_kwargs(self, fn: Callable[[dict[str, Any]], str]):
+        self._client_title_resolver = fn
+        return self
+
+    def _resolve_client_title(self, kwargs: dict[str, Any]):
+        if self._client_title_resolver is None:
+            return self._client_title
+        return self._client_title_resolver(kwargs)
+
+    # Opt in to injecting stream_options.include_usage = True for streaming calls.
+    # Used by chat-completions wrappers (OpenAI, LangChain-OpenAI, Groq).
+    #
+    # `via_extra_body=True` routes the injection through the upstream SDK's
+    # `extra_body` escape hatch instead of the top-level `stream_options` kwarg.
+    # Required for Groq's Python SDK, whose `create()` signature does not
+    # accept `stream_options` directly (TypeError) even though the underlying
+    # HTTP API supports it.
+    def include_stream_options(self, via_extra_body: bool = False):
+        self._include_stream_options = True
+        self._stream_options_via_extra_body = via_extra_body
+        return self
+
+    def stream_response_extractor(self, fn):
+        # Override how the final telemetry response is extracted from the merged
+        # stream accumulator.  The callable receives the merged chunk dict and
+        # returns the object that should be serialised into the payload.
+        #
+        # Default (None): use get_response_content(raw_response) as-is, which
+        # is correct for chat completions chunks that share a single schema.
+        #
+        # Responses API streaming: events have varying types and the actual
+        # Response is nested inside the terminal "response.completed" event:
+        #   {"type": "response.completed", "response": <Response>, ...}
+        # The extractor pulls out raw_response["response"] so the telemetry
+        # payload matches the non-streaming responses.create() shape.
+        self._stream_response_extractor = fn
+        return self
+
+    def _extract_stream_response(self, raw_response):
+        if self._stream_response_extractor is not None:
+            return self._stream_response_extractor(raw_response)
+        return self.get_response_content(raw_response)
+
+    def uses_protobuf(self):
+        self._uses_protobuf = True
+        return self
+
+
+class BaseIterator:
+    def __init__(self, config: Config, source_iterator):
+        self.config = config
+        self.source_iterator = source_iterator
+        self.iterator = None
+        self.raw_response = None
+
+    def configure_invoke(self, invoke: BaseInvoke):
+        self.invoke = invoke
+        return self
+
+    def configure_request(self, kwargs, time_start):
+        self._kwargs = kwargs
+        self._time_start = time_start
+        return self
+
+    def process_chunk(self, chunk):
+        if self.invoke._uses_protobuf is True:
+            formatted_chunk = copy.deepcopy(chunk)
+            if formatted_chunk.__dict__ and formatted_chunk.__dict__.get("_pb"):
+                # The "_pb" is for handling the protobuf structure.  It is used
+                # in the LangChain Google code paths.
+                self.raw_response.append(
+                    json.loads(
+                        json_format.MessageToJson(formatted_chunk.__dict__["_pb"])
+                    )
+                )
+            elif hasattr(formatted_chunk, "model_dump"):
+                # New path: Pydantic-based response (google-genai >= v4)
+                self.raw_response.append(formatted_chunk.model_dump(mode="json"))
+        else:
+            self.raw_response = merge_chunk(self.raw_response, chunk.__dict__)
+
+        return self
+
+    def set_raw_response(self):
+        if self.raw_response is not None:
+            return self
+
+        self.raw_response = {}
+        if self.invoke._uses_protobuf:
+            self.raw_response = []
+
+        return self
+
+
+class BaseProvider:
+    def __init__(self, parent):
+        self.client = None
+        self.parent = parent
+        self.config = parent.config