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

veris_ai/jaeger_interface/client.py

@@ -0,0 +1,233 @@
+"""Synchronous Jaeger Query Service client built on **requests**.
+
+This implementation keeps dependencies minimal while providing fully-typed
+*pydantic* models for both **request** and **response** bodies.
+"""
+
+import json
+from typing import Any, Dict, List, Optional, Self
+
+import requests
+
+from .models import GetTraceResponse, SearchResponse, Span, Tag, Trace
+
+
+__all__ = ["JaegerClient"]
+
+
+class JaegerClient:  # noqa: D101
+    def __init__(
+        self,
+        base_url: str,
+        *,
+        timeout: float | None = 10.0,
+        session: requests.Session | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        """Create a new *JaegerClient* instance.
+
+        Args:
+            base_url: Base URL of the Jaeger Query Service (e.g. ``http://localhost:16686``).
+            timeout: Request timeout in **seconds** (applied to every call).
+            session: Optional pre-configured :class:`requests.Session` to reuse.
+            headers: Optional default headers to send with every request.
+        """
+        # Normalise to avoid trailing slash duplicates
+        self._base_url = base_url.rstrip("/")
+        self._timeout = timeout
+        self._external_session = session  # If provided we won't close it
+        self._headers = headers or {}
+
+    # ---------------------------------------------------------------------
+    # Internal helpers
+    # ---------------------------------------------------------------------
+
+    def _make_session(self) -> tuple[requests.Session, bool]:  # noqa: D401
+        """Return a *(session, should_close)* tuple.
+
+        If an external session was supplied we **must not** close it after the
+        request, hence the boolean flag letting callers know whether they are
+        responsible for closing the session.
+        """
+        if self._external_session is not None:
+            return self._external_session, False
+
+        # Reuse the session opened via the context manager if available
+        if hasattr(self, "_session_ctx"):
+            return self._session_ctx, False
+
+        session = requests.Session()
+        session.headers.update(self._headers)
+        return session, True
+
+    def _span_matches_tags(self, span: Span, span_tags: Dict[str, Any]) -> bool:
+        """Check if a span matches any of the provided tags (OR logic)."""
+        if not span.tags or not span_tags:
+            return False
+        
+        # Convert span tags to a dict for easier lookup
+        span_tag_dict = {tag.key: tag.value for tag in span.tags}
+        
+        # OR logic: return True if ANY tag matches
+        for key, value in span_tags.items():
+            if span_tag_dict.get(key) == value:
+                return True
+        
+        return False
+
+    def _filter_spans(
+        self,
+        traces: List[Trace],
+        span_tags: Optional[Dict[str, Any]],
+        span_operations: Optional[List[str]] = None
+    ) -> List[Trace]:
+        """
+        Filter spans within traces based on span_tags (OR logic) and/or span_operations (OR logic).
+        If both are provided, a span must match at least one tag AND at least one operation.
+        """
+        if not span_tags and not span_operations:
+            return traces
+
+        filtered_traces = []
+        for trace in traces:
+            filtered_spans = []
+            for span in trace.spans:
+                tag_match = True
+                op_match = True
+
+                if span_tags:
+                    tag_match = self._span_matches_tags(span, span_tags)
+                if span_operations:
+                    op_match = span.operationName in span_operations
+
+                # If both filters are provided, require both to match (AND logic)
+                if tag_match and op_match:
+                    filtered_spans.append(span)
+
+            if filtered_spans:
+                filtered_trace = Trace(
+                    traceID=trace.traceID,
+                    spans=filtered_spans,
+                    process=trace.process,
+                    warnings=trace.warnings
+                )
+                filtered_traces.append(filtered_trace)
+
+        return filtered_traces
+
+    # ---------------------------------------------------------------------
+    # Public API
+    # ---------------------------------------------------------------------
+
+    def search(
+        self,
+        service: Optional[str] = None,
+        *,
+        limit: Optional[int] = None,
+        tags: Optional[Dict[str, Any]] = None,
+        operation: Optional[str] = None,
+        span_tags: Optional[Dict[str, Any]] = None,
+        span_operations: Optional[List[str]] = None,
+        **kwargs: Any
+    ) -> SearchResponse:  # noqa: D401
+        """Search traces using the *v1* ``/api/traces`` endpoint with optional span filtering.
+
+        Args:
+            service: Service name to search for. If not provided, searches across all services.
+            limit: Maximum number of traces to return.
+            tags: Dictionary of tag filters for trace-level filtering (AND-combined).
+            operation: Operation name to search for.
+            span_tags: Dictionary of tag filters for span-level filtering (OR-combined AND-combined with span_operations).
+                      Applied client-side after retrieving traces.
+            span_operations: List of operation names to search for (OR-combined AND-combined with span_tags).
+            **kwargs: Additional parameters to pass to the Jaeger API.
+
+        Returns:
+            Parsed :class:`~veris_ai.jaeger_interface.models.SearchResponse` model
+            with spans filtered according to span_tags if provided.
+        """
+        # Build params for the Jaeger API (excluding span_tags)
+        params: Dict[str, Any] = {}
+        
+        if service is not None:
+            params["service"] = service
+        
+        if limit is not None:
+            params["limit"] = limit
+        
+        if operation is not None:
+            params["operation"] = operation
+        
+        if tags:
+            # Convert tags to JSON string as expected by Jaeger API
+            params["tags"] = json.dumps(tags)
+        
+        # Add any additional parameters
+        params.update(kwargs)
+        
+        session, should_close = self._make_session()
+        try:
+            url = f"{self._base_url}/api/traces"
+            response = session.get(url, params=params, timeout=self._timeout)
+            response.raise_for_status()
+            data = response.json()
+        finally:
+            if should_close:
+                session.close()
+        
+        # Parse the response
+        search_response = SearchResponse.model_validate(data)  # type: ignore[arg-type]
+        
+        # Apply span-level filtering if span_tags is provided
+        if span_tags or span_operations and search_response.data and isinstance(search_response.data, list):
+            filtered_traces = self._filter_spans(search_response.data, span_tags, span_operations)
+            search_response.data = filtered_traces
+            # Update the total to reflect filtered results
+            if search_response.total is not None:
+                search_response.total = len(filtered_traces)
+        
+        return search_response
+
+    def get_trace(self, trace_id: str) -> GetTraceResponse:  # noqa: D401
+        """Retrieve a single trace by *trace_id*.
+
+        Args:
+            trace_id: The Jaeger trace identifier.
+
+        Returns:
+            Parsed :class:`~veris_ai.jaeger_interface.models.GetTraceResponse` model.
+        """
+        if not trace_id:
+            error_msg = "trace_id must be non-empty"
+            raise ValueError(error_msg)
+
+        session, should_close = self._make_session()
+        try:
+            url = f"{self._base_url}/api/traces/{trace_id}"
+            response = session.get(url, timeout=self._timeout)
+            response.raise_for_status()
+            data = response.json()
+        finally:
+            if should_close:
+                session.close()
+        return GetTraceResponse.model_validate(data)  # type: ignore[arg-type]
+
+    # ------------------------------------------------------------------
+    # Context-manager helpers (optional but convenient)
+    # ------------------------------------------------------------------
+
+    def __enter__(self) -> Self:
+        """Enter the context manager."""
+        self._session_ctx, self._should_close_ctx = self._make_session()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: Any | None,
+    ) -> None:
+        """Exit the context manager."""
+        # Only close if we created the session
+        if getattr(self, "_should_close_ctx", False):
+            self._session_ctx.close()

veris_ai/jaeger_interface/models.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from pydantic import BaseModel, Field, ConfigDict
+
+__all__ = [
+    "Tag",
+    "Process",
+    "Span",
+    "Trace",
+    "SearchResponse",
+    "GetTraceResponse",
+]
+
+
+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

veris_ai/models.py

@@ -0,0 +1,11 @@
+"""Models for the VERIS SDK."""
+
+from enum import Enum
+
+
+class ResponseExpectation(str, Enum):
+    """Expected response behavior for tool mocking."""
+    
+    AUTO = "auto"
+    REQUIRED = "required"
+    NONE = "none"

veris_ai/tool_mock.py

@@ -8,13 +8,15 @@ 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.models import ResponseExpectation
+from veris_ai.utils import convert_to_type, extract_json_schema
 
 logger = logging.getLogger(__name__)
 
@@ -90,72 +92,168 @@ 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"))
+            
+            # Check if the original function is async
+            is_async = inspect.iscoroutinefunction(func)
+
+            def create_mock_payload(
+                *args: tuple[object, ...],
+                **kwargs: dict[str, object],
+            ) -> tuple[dict[str, Any], Any]:
+                """Create the mock payload - shared logic for both sync and async."""
+                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
+                # Convert expects_response to response_expectation enum
+                if expects_response is False:
+                    response_expectation = ResponseExpectation.NONE
+                elif expects_response is True:
+                    response_expectation = ResponseExpectation.REQUIRED
+                else:
+                    response_expectation = ResponseExpectation.AUTO
+                
+                payload = {
+                    "session_id": self.session_id,
+                    "response_expectation": response_expectation.value,
+                    "cache_response": bool(cache_response) if cache_response is not None else False,
+                    "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
+                return payload, return_type_obj
+
+            @wraps(func)
+            async def async_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__}")
+                payload, return_type_obj = create_mock_payload(*args, **kwargs)
+
+                # 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}")
+
+                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)
+
+            @wraps(func)
+            def sync_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 func(*args, **kwargs)
+                
+                logger.info(f"Simulating function: {func.__name__}")
+                payload, return_type_obj = create_mock_payload(*args, **kwargs)
+
+                # Send request to endpoint with timeout (synchronous)
+                with httpx.Client(timeout=timeout) as client:
+                    response = client.post(endpoint, json=payload)
+                    response.raise_for_status()
+                    mock_result = response.json()
+                    logger.info(f"Mock response: {mock_result}")
+
+                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 the appropriate wrapper based on whether the function is async
+            return async_wrapper if is_async else sync_wrapper
+
+        return decorator
+
+    def stub(self, return_value: Any) -> Callable:  # noqa: ANN401
+        """Decorator for stubbing toolw calls."""
+
+        def decorator(func: Callable) -> Callable:
+            # Check if the original function is async
+            is_async = inspect.iscoroutinefunction(func)
+            
+            @wraps(func)
+            async def async_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 func(*args, **kwargs)
+                logger.info(f"Simulating function: {func.__name__}")
+                return return_value
+            
+            @wraps(func)
+            def sync_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 func(*args, **kwargs)
+                logger.info(f"Simulating function: {func.__name__}")
+                return return_value
+
+            # Return the appropriate wrapper based on whether the function is async
+            return async_wrapper if is_async else sync_wrapper
+
+        return decorator
 
 
 veris = VerisSDK()