veris-ai 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

veris_ai/jaeger_interface/models.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+__all__ = [
+    "Tag",
+    "Process",
+    "Span",
+    "Trace",
+    "SearchResponse",
+    "GetTraceResponse",
+    "SearchQuery",
+]
+
+
+class Tag(BaseModel):
+    """A Jaeger tag key/value pair."""
+
+    key: str
+    value: Any
+    type: str | None = None  # Jaeger uses an optional *type* field in v1
+
+
+class Process(BaseModel):
+    """Represents the *process* section of a Jaeger trace."""
+
+    serviceName: str = Field(alias="serviceName")  # noqa: N815
+    tags: list[Tag] | None = None
+
+
+class Span(BaseModel):
+    """Represents a single Jaeger span."""
+
+    traceID: str  # noqa: N815
+    spanID: str  # noqa: N815
+    operationName: str  # noqa: N815
+    startTime: int  # noqa: N815
+    duration: int
+    tags: list[Tag] | None = None
+    references: list[dict[str, Any]] | None = None
+    processID: str | None = None  # noqa: N815
+
+    model_config = ConfigDict(extra="allow")
+
+
+class Trace(BaseModel):
+    """A full Jaeger trace as returned by the Query API."""
+
+    traceID: str  # noqa: N815
+    spans: list[Span]
+    process: Process | dict[str, Process] | None = None
+    warnings: list[str] | None = None
+
+    model_config = ConfigDict(extra="allow")
+
+
+class _BaseResponse(BaseModel):
+    data: list[Trace] | Trace | None = None
+    errors: list[str] | None = None
+
+    # Allow any additional keys returned by Jaeger so that nothing gets
+    # silently dropped if the backend adds new fields we don’t know about.
+
+    model_config = ConfigDict(extra="allow")
+
+
+class SearchResponse(_BaseResponse):
+    """Response model for *search* or *find traces* requests."""
+
+    total: int | None = None
+    limit: int | None = None
+
+
+class GetTraceResponse(_BaseResponse):
+    """Response model for *get trace by id* requests."""
+
+    # Same as base but alias for clarity
+
+
+# ---------------------------------------------------------------------------
+# Query models
+# ---------------------------------------------------------------------------
+
+
+class SearchQuery(BaseModel):
+    """Minimal set of query parameters for the `/api/traces` endpoint.
+
+    Parameter interaction rules:
+
+    * **service** – global filter; *all* returned traces must belong to this
+      service.
+    * **operation** – optional secondary filter; returned traces must contain
+      *at least one span* whose ``operationName`` equals the provided value.
+    * **tags** – dictionary of key‒value pairs; each trace must include a span
+      that matches **all** of the pairs (logical AND).
+    * **limit** – applied *after* all other filters; truncates the final list
+      of traces to the requested maximum.
+
+    Any additional/unknown parameters are forwarded thanks to
+    ``extra = "allow"`` – this keeps the model future-proof.
+    """
+
+    # NOTE: Only the fields that are reliably supported by Jaeger’s REST API and
+    # work with the user’s deployment are kept. The model remains *open* to any
+    # extra parameters thanks to `extra = "allow"`.
+
+    service: str = Field(
+        ...,
+        description="Service name to search for. Example: 'veris-agent'",
+    )
+
+    limit: int | None = Field(
+        None,
+        description="Maximum number of traces to return. Example: 10",
+    )
+
+    tags: dict[str, Any] | None = Field(
+        None,
+        description=(
+            "Dictionary of tag filters (AND-combined). "
+            "Example: {'error': 'true', 'bt.metrics.time_to_first_token': '0.813544'}"
+        ),
+    )
+
+    operation: str | None = Field(
+        None,
+        description="Operation name to search for. Example: 'process_chat_message'",
+    )
+
+    model_config = ConfigDict(
+        extra="allow",  # allow additional query params implicitly
+        populate_by_name=True,
+        str_to_lower=False,
+    )
+
+    @field_validator("tags", mode="before")
+    @classmethod
+    def _empty_to_none(cls, v: dict[str, Any] | None) -> dict[str, Any] | None:  # noqa: D401, ANN102
+        return v or None
+
+    def to_params(self) -> dict[str, Any]:  # noqa: D401
+        """Translate the model into a *requests*/*httpx* compatible params dict."""
+        # Dump using aliases so ``span_kind`` becomes ``spanKind`` automatically.
+        params: dict[str, Any] = self.model_dump(exclude_none=True, by_alias=True)
+
+        # Convert tags to a JSON string if necessary – this matches what the UI sends.
+        if "tags" in params and isinstance(params["tags"], dict):
+            params["tags"] = json.dumps(params["tags"])
+
+        return params

veris_ai/tool_mock.py

@@ -8,13 +8,14 @@ from contextvars import ContextVar
 from functools import wraps
 from typing import (
     Any,
+    Literal,
     TypeVar,
     get_type_hints,
 )
 
 import httpx
 
-from veris_ai.utils import convert_to_type
+from veris_ai.utils import convert_to_type, extract_json_schema
 
 logger = logging.getLogger(__name__)
 
@@ -90,72 +91,109 @@ class VerisSDK:
             **params_dict,
         )
 
-    def mock(self, func: Callable) -> Callable:
+    def mock(
+        self,
+        mode: Literal["tool", "function"] = "tool",
+        expects_response: bool | None = None,
+        cache_response: bool | None = None,
+    ) -> Callable:
         """Decorator for mocking tool calls."""
-        endpoint = os.getenv("VERIS_MOCK_ENDPOINT_URL")
-        if not endpoint:
-            error_msg = "VERIS_MOCK_ENDPOINT_URL environment variable is not set"
-            raise ValueError(error_msg)
-        # Default timeout of 30 seconds
-        timeout = float(os.getenv("VERIS_MOCK_TIMEOUT", "30.0"))
-
-        @wraps(func)
-        async def wrapper(
-            *args: tuple[object, ...],
-            **kwargs: dict[str, object],
-        ) -> object:
-            # Check if we're in simulation mode
-            env_mode = os.getenv("ENV", "").lower()
-            if env_mode != "simulation":
-                # If not in simulation mode, execute the original function
-                return await func(*args, **kwargs)
-            logger.info(f"Simulating function: {func.__name__}")
-            sig = inspect.signature(func)
-            type_hints = get_type_hints(func)
-
-            # Extract return type object (not just the name)
-            return_type_obj = type_hints.pop("return", Any)
-
-            # Create parameter info
-            params_info = {}
-            bound_args = sig.bind(*args, **kwargs)
-            bound_args.apply_defaults()
-
-            for param_name, param_value in bound_args.arguments.items():
-                params_info[param_name] = {
-                    "value": param_value,
-                    "type": type_hints.get(param_name, Any).__name__,
+
+        def decorator(func: Callable) -> Callable:
+            """Decorator for mocking tool calls."""
+            endpoint = os.getenv("VERIS_MOCK_ENDPOINT_URL")
+            if not endpoint:
+                error_msg = "VERIS_MOCK_ENDPOINT_URL environment variable is not set"
+                raise ValueError(error_msg)
+            # Default timeout of 30 seconds
+            timeout = float(os.getenv("VERIS_MOCK_TIMEOUT", "90.0"))
+
+            @wraps(func)
+            async def wrapper(
+                *args: tuple[object, ...],
+                **kwargs: dict[str, object],
+            ) -> object:
+                # Check if we're in simulation mode
+                env_mode = os.getenv("ENV", "").lower()
+                if env_mode != "simulation":
+                    # If not in simulation mode, execute the original function
+                    return await func(*args, **kwargs)
+                logger.info(f"Simulating function: {func.__name__}")
+                sig = inspect.signature(func)
+                type_hints = get_type_hints(func)
+
+                # Extract return type object (not just the name)
+                return_type_obj = type_hints.pop("return", Any)
+                # Create parameter info
+                params_info = {}
+                bound_args = sig.bind(*args, **kwargs)
+                bound_args.apply_defaults()
+                _ = bound_args.arguments.pop("ctx", None)
+                _ = bound_args.arguments.pop("self", None)
+                _ = bound_args.arguments.pop("cls", None)
+
+                for param_name, param_value in bound_args.arguments.items():
+                    params_info[param_name] = {
+                        "value": str(param_value),
+                        "type": str(type_hints.get(param_name, Any)),
+                    }
+                # Get function docstring
+                docstring = inspect.getdoc(func) or ""
+                nonlocal expects_response
+                if expects_response is None and mode == "function":
+                    expects_response = False
+                # Prepare payload
+                payload = {
+                    "session_id": self.session_id,
+                    "expects_response": expects_response,
+                    "cache_response": cache_response,
+                    "tool_call": {
+                        "function_name": func.__name__,
+                        "parameters": params_info,
+                        "return_type": json.dumps(extract_json_schema(return_type_obj)),
+                        "docstring": docstring,
+                    },
                 }
 
-            # Get function docstring
-            docstring = inspect.getdoc(func) or ""
-            # Prepare payload
-            payload = {
-                "session_id": self.session_id,
-                "tool_call": {
-                    "function_name": func.__name__,
-                    "parameters": params_info,
-                    "return_type": return_type_obj.__name__,
-                    "docstring": docstring,
-                },
-            }
-
-            # Send request to endpoint with timeout
-            async with httpx.AsyncClient(timeout=timeout) as client:
-                response = await client.post(endpoint, json=payload)
-                response.raise_for_status()
-                mock_result = response.json()["result"]
-                logger.info(f"Mock response: {mock_result}")
-
-            # Parse the mock result if it's a string
-            if isinstance(mock_result, str):
-                with suppress(json.JSONDecodeError):
-                    mock_result = json.loads(mock_result)
-
-            # Convert the mock result to the expected return type
-            return convert_to_type(mock_result, return_type_obj)
-
-        return wrapper
+                # Send request to endpoint with timeout
+                async with httpx.AsyncClient(timeout=timeout) as client:
+                    response = await client.post(endpoint, json=payload)
+                    response.raise_for_status()
+                    mock_result = response.json()
+                    logger.info(f"Mock response: {mock_result}")
+
+                # Convert the mock result to the expected return type
+                if mode == "tool":
+                    return {"content": [{"type": "text", "text": mock_result}]}
+                # Parse the mock result if it's a string
+                # Extract result field for backwards compatibility
+                # Parse the mock result if it's a string
+                if isinstance(mock_result, str):
+                    with suppress(json.JSONDecodeError):
+                        mock_result = json.loads(mock_result)
+                        return convert_to_type(mock_result, return_type_obj)
+                return convert_to_type(mock_result, return_type_obj)
+
+            return wrapper
+
+        return decorator
+
+    def stub(self, return_value: Any) -> Callable:  # noqa: ANN401
+        """Decorator for stubbing tool calls."""
+
+        def decorator(func: Callable) -> Callable:
+            @wraps(func)
+            async def wrapper(*args: tuple[object, ...], **kwargs: dict[str, object]) -> object:
+                env_mode = os.getenv("ENV", "").lower()
+                if env_mode != "simulation":
+                    # If not in simulation mode, execute the original function
+                    return await func(*args, **kwargs)
+                logger.info(f"Simulating function: {func.__name__}")
+                return return_value
+
+            return wrapper
+
+        return decorator
 
 
 veris = VerisSDK()

veris_ai/utils.py

@@ -1,5 +1,10 @@
+import sys
+import types
+import typing
 from contextlib import suppress
-from typing import Any, Union, get_args, get_origin
+from typing import Any, ForwardRef, Literal, NotRequired, Required, Union, get_args, get_origin
+
+from pydantic import BaseModel
 
 
 def convert_to_type(value: object, target_type: type) -> object:
@@ -69,3 +74,197 @@ def _convert_simple_type(value: object, target_type: type) -> object:
             return target_type(**value)
 
     return target_type(value)
+
+
+def _resolve_forward_ref(ref: ForwardRef, module_context: types.ModuleType | None = None) -> Any:  # noqa: ANN401
+    """Resolve a ForwardRef to its actual type."""
+    if not isinstance(ref, ForwardRef):
+        return ref
+
+    # Try to evaluate the forward reference
+    try:
+        # Get the module's namespace for evaluation
+        namespace = dict(vars(module_context)) if module_context else {}
+
+        # Add common typing imports to namespace
+        namespace.update(
+            {
+                "Union": Union,
+                "Any": Any,
+                "Literal": Literal,
+                "Required": Required,
+                "NotRequired": NotRequired,
+                "List": list,
+                "Dict": dict,
+                "Optional": typing.Optional,
+                "Iterable": typing.Iterable,
+                "str": str,
+                "int": int,
+                "float": float,
+                "bool": bool,
+            },
+        )
+
+        # Try to import from the same module to resolve local references
+        if module_context and hasattr(module_context, "__name__"):
+            with suppress(Exception):
+                # Import all from the module to get access to local types
+                exec(f"from {module_context.__name__} import *", namespace)  # noqa: S102
+
+        # Get the forward reference string
+        ref_string = ref.__forward_arg__ if hasattr(ref, "__forward_arg__") else str(ref)
+
+        # Try to evaluate the forward reference string
+        return eval(ref_string, namespace, namespace)  # noqa: S307
+    except Exception:
+        # If we can't resolve it, return the ref itself
+        return ref
+
+
+def _unwrap_required(field_type: Any) -> tuple[Any, bool]:  # noqa: ANN401
+    """Unwrap Required/NotRequired and return the inner type and whether it's required."""
+    origin = get_origin(field_type)
+
+    # Check if it's Required or NotRequired
+    if origin is Required:
+        args = get_args(field_type)
+        return args[0] if args else field_type, True
+    if origin is NotRequired:
+        args = get_args(field_type)
+        return args[0] if args else field_type, False
+
+    # Default to required for TypedDict fields
+    return field_type, True
+
+
+def extract_json_schema(target_type: Any) -> dict:  # noqa: PLR0911, PLR0912, C901, ANN401
+    """Extract the JSON schema from a type or pydantic model.
+
+    Args:
+        target_type: The type or pydantic model to extract the JSON schema from.
+
+    Returns:
+        A dictionary representing the JSON schema.
+
+    Example:
+        >>> extract_json_schema(int)
+        {"type": "integer"}
+
+        >>> extract_json_schema(list[int])
+        {"type": "array", "items": {"type": "integer"}}
+
+        >>> extract_json_schema(list[User])
+        {"type": "array", "items": {"type": "object", "properties": {...}}}
+    """
+    # Handle Pydantic BaseModel instances or classes
+    if isinstance(target_type, type) and issubclass(target_type, BaseModel):
+        return target_type.model_json_schema()
+    if isinstance(target_type, BaseModel):
+        return target_type.model_json_schema()
+
+    # Handle TypedDict
+    if (
+        isinstance(target_type, type)
+        and hasattr(target_type, "__annotations__")
+        and hasattr(target_type, "__total__")
+    ):
+        # This is a TypedDict
+        properties = {}
+        required = []
+
+        # Get the module context for resolving forward references
+        module = sys.modules.get(target_type.__module__)
+
+        for field_name, field_type_annotation in target_type.__annotations__.items():
+            # Resolve forward references if present
+            resolved_type = field_type_annotation
+            if isinstance(resolved_type, ForwardRef):
+                resolved_type = _resolve_forward_ref(resolved_type, module)
+
+            # Unwrap Required/NotRequired
+            unwrapped_type, is_required = _unwrap_required(resolved_type)
+
+            # Extract schema for the unwrapped type
+            properties[field_name] = extract_json_schema(unwrapped_type)
+
+            # Add to required list if necessary
+            if is_required and getattr(target_type, "__total__", True):
+                required.append(field_name)
+
+        schema = {"type": "object", "properties": properties}
+        if required:
+            schema["required"] = required
+        return schema
+
+    # Handle built-in types
+    type_mapping = {
+        str: {"type": "string"},
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        bool: {"type": "boolean"},
+        type(None): {"type": "null"},
+        Any: {},  # Empty schema for Any type
+    }
+
+    if target_type in type_mapping:
+        return type_mapping[target_type]
+
+    # Handle generic types
+    origin = get_origin(target_type)
+
+    # Handle bare collection types
+    if target_type is list:
+        return {"type": "array"}
+    if target_type is dict:
+        return {"type": "object"}
+    if target_type is tuple:
+        return {"type": "array"}
+
+    # Handle Literal types
+    if origin is Literal:
+        values = get_args(target_type)
+        if len(values) == 1:
+            # Single literal value - use const
+            return {"const": values[0]}
+        # Multiple literal values - use enum
+        return {"enum": list(values)}
+
+    if origin is list:
+        args = get_args(target_type)
+        if args:
+            return {"type": "array", "items": extract_json_schema(args[0])}
+        return {"type": "array"}
+
+    if origin is dict:
+        args = get_args(target_type)
+        if len(args) == 2:  # noqa: PLR2004
+            # For typed dicts like dict[str, int]
+            return {
+                "type": "object",
+                "additionalProperties": extract_json_schema(args[1]),
+            }
+        return {"type": "object"}
+
+    if origin is Union:
+        args = get_args(target_type)
+        # Handle Optional types (Union[T, None])
+        if len(args) == 2 and type(None) in args:  # noqa: PLR2004
+            non_none_type = args[0] if args[1] is type(None) else args[1]
+            schema = extract_json_schema(non_none_type)
+            return {"anyOf": [schema, {"type": "null"}]}
+        # Handle general Union types
+        return {"anyOf": [extract_json_schema(arg) for arg in args]}
+
+    if origin is tuple:
+        args = get_args(target_type)
+        if args:
+            return {
+                "type": "array",
+                "prefixItems": [extract_json_schema(arg) for arg in args],
+                "minItems": len(args),
+                "maxItems": len(args),
+            }
+        return {"type": "array"}
+
+    # Default case for unknown types
+    return {"type": "object"}