veris-ai 0.2.1__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

@@ -4,137 +4,196 @@ import logging
 import os
 from collections.abc import Callable
 from contextlib import suppress
+from contextvars import ContextVar
 from functools import wraps
-from typing import Any, TypeVar, Union, get_args, get_origin, get_type_hints
+from typing import (
+    Any,
+    Literal,
+    TypeVar,
+    get_type_hints,
+)
 
 import httpx
 
+from veris_ai.utils import convert_to_type, extract_json_schema
+
 logger = logging.getLogger(__name__)
 
 T = TypeVar("T")
 
+# Context variable to store session_id for each call
+_session_id_context: ContextVar[str | None] = ContextVar("veris_session_id", default=None)
+
 
-class ToolMock:
+class VerisSDK:
     """Class for mocking tool calls."""
 
     def __init__(self) -> None:
         """Initialize the ToolMock class."""
+        self._mcp = None
+
+    @property
+    def session_id(self) -> str | None:
+        """Get the session_id from context variable."""
+        return _session_id_context.get()
+
+    def set_session_id(self, session_id: str) -> None:
+        """Set the session_id in context variable."""
+        _session_id_context.set(session_id)
+        logger.info(f"Session ID set to {session_id}")
+
+    def clear_session_id(self) -> None:
+        """Clear the session_id from context variable."""
+        _session_id_context.set(None)
+        logger.info("Session ID cleared")
+
+    @property
+    def fastapi_mcp(self) -> Any | None:  # noqa: ANN401
+        """Get the FastAPI MCP server."""
+        return self._mcp
+
+    def set_fastapi_mcp(self, **params_dict: Any) -> None:  # noqa: ANN401
+        """Set the FastAPI MCP server."""
+        from fastapi import Depends, Request  # noqa: PLC0415
+        from fastapi.security import OAuth2PasswordBearer  # noqa: PLC0415
+        from fastapi_mcp import (  # type: ignore[import-untyped] # noqa: PLC0415
+            AuthConfig,
+            FastApiMCP,
+        )
+
+        oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+        async def authenticate_request(
+            _: Request,
+            token: str = Depends(oauth2_scheme),  # noqa: ARG001
+        ) -> None:
+            self.set_session_id(token)
+
+        # Create auth config with dependencies
+        auth_config = AuthConfig(
+            dependencies=[Depends(authenticate_request)],
+        )
+
+        # Merge the provided params with our auth config
+        if "auth_config" in params_dict:
+            # Merge the provided auth config with our dependencies
+            provided_auth_config = params_dict.pop("auth_config")
+            if provided_auth_config.dependencies:
+                auth_config.dependencies.extend(provided_auth_config.dependencies)
+            # Copy other auth config properties if they exist
+            for field, value in provided_auth_config.model_dump(exclude_none=True).items():
+                if field != "dependencies" and hasattr(auth_config, field):
+                    setattr(auth_config, field, value)
+
+        # Create the FastApiMCP instance with merged parameters
+        self._mcp = FastApiMCP(
+            auth_config=auth_config,
+            **params_dict,
+        )
+
+    def mock(
+        self,
+        mode: Literal["tool", "function"] = "tool",
+        expects_response: bool | None = None,
+        cache_response: bool | None = None,
+    ) -> Callable:
+        """Decorator for mocking tool calls."""
 
-    def _convert_to_type(self, value: object, target_type: type) -> object:
-        """Convert a value to the specified type."""
-        if target_type is Any:
-            return value
-
-        if target_type in (str, int, float, bool):
-            return target_type(value)
-
-        origin = get_origin(target_type)
-        if origin is list:
-            if not isinstance(value, list):
-                error_msg = f"Expected list but got {type(value)}"
-                raise ValueError(error_msg)
-            item_type = get_args(target_type)[0]
-            return [self._convert_to_type(item, item_type) for item in value]
-
-        if origin is dict:
-            if not isinstance(value, dict):
-                error_msg = f"Expected dict but got {type(value)}"
+        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)
-            key_type, value_type = get_args(target_type)
-            return {
-                self._convert_to_type(k, key_type): self._convert_to_type(v, value_type)
-                for k, v in value.items()
-            }
-
-        if origin is Union:
-            error_msg = (
-                f"Could not convert {value} to any of the union types {get_args(target_type)}"
-            )
-            for possible_type in get_args(target_type):
-                try:
-                    return self._convert_to_type(value, possible_type)
-                except (ValueError, TypeError):
-                    continue
-            raise ValueError(error_msg)
-
-        # For other types, try direct conversion
-        return target_type(value)
-
-    def mock(self, 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", "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()
-
-            ctx = bound_args.arguments.pop("ctx", None)
-            session_id = None
-            if ctx:
-                try:
-                    session_id = ctx.request_context.lifespan_context.session_id
-                except AttributeError:
-                    logger.warning("Cannot get session_id from context.")
-
-            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__,
+            # 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": 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}")
+                # 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)
 
-                # Convert the mock result to the expected return type
-                return self._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 wrapper
+        return decorator
 
 
-veris = ToolMock()
+veris = VerisSDK()