veris-ai 1.11.1__py3-none-any.whl → 1.12.0__py3-none-any.whl

veris_ai/tool_mock.py

@@ -1,3 +1,4 @@
+import base64
 import inspect
 import json
 import logging
@@ -22,8 +23,9 @@ logger = logging.getLogger(__name__)
 
 T = TypeVar("T")
 
-# Context variable to store session_id for each call
+# Context variables to store session_id and thread_id for each call
 _session_id_context: ContextVar[str | None] = ContextVar("veris_session_id", default=None)
+_thread_id_context: ContextVar[str | None] = ContextVar("veris_thread_id", default=None)
 
 
 class VerisSDK:
@@ -38,15 +40,111 @@ class VerisSDK:
         """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."""
+    @property
+    def thread_id(self) -> str | None:
+        """Get the thread_id from context variable."""
+        return _thread_id_context.get()
+
+    def _set_session_id(self, session_id: str) -> None:
+        """Set the session_id in context variable (private method)."""
         _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."""
+    def _set_thread_id(self, thread_id: str) -> None:
+        """Set the thread_id in context variable (private method)."""
+        _thread_id_context.set(thread_id)
+
+    def parse_token(self, token: str) -> None:
+        """Parse and set session_id and thread_id from a base64-encoded token.
+
+        The token must be a base64-encoded JSON object containing:
+        {"session_id": "...", "thread_id": "..."}
+
+        Both session_id and thread_id are required fields.
+
+        Args:
+            token: Base64-encoded JSON token
+
+        Raises:
+            ValueError: If token is not valid or missing required fields
+        """
+        try:
+            # Decode base64 JSON token
+            decoded = base64.b64decode(token.encode("utf-8")).decode("utf-8")
+            token_data = json.loads(decoded)
+
+            if not isinstance(token_data, dict):
+                raise ValueError("Token must decode to a JSON object")
+
+            if "session_id" not in token_data:
+                raise ValueError("Token must contain 'session_id' field")
+
+            if "thread_id" not in token_data:
+                raise ValueError("Token must contain 'thread_id' field")
+
+            self._set_session_id(token_data["session_id"])
+            self._set_thread_id(token_data["thread_id"])
+            logger.info(
+                f"Session ID set to {token_data['session_id']}, "
+                f"Thread ID set to {token_data['thread_id']} - mocking enabled"
+            )
+
+        except (ValueError, json.JSONDecodeError) as e:
+            error_msg = f"Invalid token format: {e}"
+            raise ValueError(error_msg) from e
+
+    def set_session_id(self, token: str) -> None:
+        """DEPRECATED: Use parse_token() instead.
+
+        Set the session_id and thread_id from a base64-encoded token.
+
+        Args:
+            token: Base64-encoded JSON token
+        """
+        logger.warning(
+            "set_session_id() is deprecated. Use parse_token() instead. "
+            "This method will be removed in a future version."
+        )
+        # For backwards compatibility, allow tokens without thread_id
+        try:
+            decoded = base64.b64decode(token.encode("utf-8")).decode("utf-8")
+            token_data = json.loads(decoded)
+
+            if not isinstance(token_data, dict):
+                raise ValueError("Token must decode to a JSON object")
+
+            if "session_id" not in token_data:
+                raise ValueError("Token must contain 'session_id' field")
+
+            self._set_session_id(token_data["session_id"])
+            if "thread_id" in token_data:
+                self._set_thread_id(token_data["thread_id"])
+            logger.info(f"Session ID set to {token_data['session_id']}")
+
+        except (ValueError, json.JSONDecodeError) as e:
+            error_msg = f"Invalid token format: {e}"
+            raise ValueError(error_msg) from e
+
+    def _clear_session_id(self) -> None:
+        """Clear the session_id from context variable (private method)."""
         _session_id_context.set(None)
-        logger.info("Session ID cleared")
+
+    def _clear_thread_id(self) -> None:
+        """Clear the thread_id from context variable (private method)."""
+        _thread_id_context.set(None)
+
+    def clear_context(self) -> None:
+        """Clear the session_id and thread_id from context variables."""
+        self._clear_session_id()
+        self._clear_thread_id()
+        logger.info("Session ID and Thread ID cleared - mocking disabled")
+
+    def clear_session_id(self) -> None:
+        """DEPRECATED: Use clear_context() instead."""
+        logger.warning(
+            "clear_session_id() is deprecated. Use clear_context() instead. "
+            "This method will be removed in a future version."
+        )
+        self.clear_context()
 
     @property
     def fastapi_mcp(self) -> Any | None:  # noqa: ANN401
@@ -69,7 +167,7 @@ class VerisSDK:
             token: str | None = Depends(oauth2_scheme),
         ) -> None:
             if token:
-                self.set_session_id(token)
+                self.parse_token(token)
 
         # Create auth config with dependencies
         auth_config = AuthConfig(
@@ -176,11 +274,13 @@ class VerisSDK:
                     )
                     return await func(*args, **kwargs)
                 parameters = get_function_parameters(func, args, kwargs)
+                thread_id = _thread_id_context.get()
                 return mock_tool_call(
                     func,
                     session_id,
                     parameters,
                     options,
+                    thread_id,
                 )
 
             @wraps(func)
@@ -196,11 +296,13 @@ class VerisSDK:
                     )
                     return func(*args, **kwargs)
                 parameters = get_function_parameters(func, args, kwargs)
+                thread_id = _thread_id_context.get()
                 return mock_tool_call(
                     func,
                     session_id,
                     parameters,
                     options,
+                    thread_id,
                 )
 
             # Return the appropriate wrapper based on whether the function is async
@@ -251,11 +353,29 @@ class VerisSDK:
 )
 def mock_tool_call(
     func: Callable,
-    session_id: str,
+    session_id: str,  # noqa: ARG001
     parameters: dict[str, dict[str, str]],
     options: ToolCallOptions | None = None,
+    thread_id: str | None = None,
 ) -> object:
-    """Mock tool call."""
+    """Mock tool call.
+
+    Args:
+        func: Function being mocked
+        session_id: Session ID (kept for backwards compatibility, not used)
+        parameters: Function parameters
+        options: Tool call options
+        thread_id: Thread ID to use as session_id in API request (required)
+
+    Raises:
+        ValueError: If thread_id is not provided
+    """
+    if thread_id is None:
+        raise ValueError(
+            "thread_id is required for mocking. "
+            "Use parse_token() to set both session_id and thread_id."
+        )
+
     options = options or ToolCallOptions()
     api_client = get_api_client()
     endpoint = api_client.tool_mock_endpoint
@@ -269,6 +389,8 @@ def mock_tool_call(
     # Get function docstring
     docstring = inspect.getdoc(func) or ""
 
+    # Use thread_id as session_id in the payload
+    payload_session_id = thread_id
     # Clean up parameters for V3 - just send values, not the nested dict
     clean_params: dict[str, Any] = {}
     for key, value in parameters.items():
@@ -281,7 +403,7 @@ def mock_tool_call(
 
     # Determine response expectation
     payload = {
-        "session_id": session_id,
+        "session_id": payload_session_id,
         "response_expectation": options.response_expectation.value,
         "cache_response": bool(options.cache_response),
         "tool_call": {

{veris_ai-1.11.1.dist-info → veris_ai-1.12.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veris-ai
-Version: 1.11.1
+Version: 1.12.0
 Summary: A Python package for Veris AI tools
 Project-URL: Homepage, https://github.com/veris-ai/veris-python-sdk
 Project-URL: Bug Tracker, https://github.com/veris-ai/veris-python-sdk/issues

{veris_ai-1.11.1.dist-info → veris_ai-1.12.0.dist-info}/RECORD

@@ -4,13 +4,13 @@ veris_ai/agents_wrapper.py,sha256=N-YZ0qvIEJ8rlFZAzVKg6eWKa7I_qIt-khlg_zL-qY8,11
 veris_ai/api_client.py,sha256=BQ2Fn5pO7uaWmWNcLOoZO8o-D39OPRmclPROwJvXBUc,3199
 veris_ai/models.py,sha256=xKeheSJQle2tBeJG1DsGJzMDwv24p5jECjX6RAa39n4,495
 veris_ai/observability.py,sha256=eSIXmk6fpOAoWM-sDbsvzyUASh1ZwU6tRIPduy09RxY,4206
-veris_ai/tool_mock.py,sha256=--t0ZhjplJXzBt1F--XPE57ZyMBjB9SOjO2OzuoWkbk,12861
+veris_ai/tool_mock.py,sha256=h_nS0tTOl8eI_Hb-HNuurjbBPp0ThdMnJwsMXfRyT7I,17489
 veris_ai/utils.py,sha256=hJetCiN8Bubhy0nqSoS1C2awN9cdkKuHM1v7YhtwtTs,10066
 veris_ai/jaeger_interface/README.md,sha256=kd9rKcE5xf3EyNaiHu0tjn-0oES9sfaK6Ih-OhhTyCM,2821
 veris_ai/jaeger_interface/__init__.py,sha256=KD7NSiMYRG_2uF6dOLKkGG5lNQe4K9ptEwucwMT4_aw,1128
 veris_ai/jaeger_interface/client.py,sha256=yJrh86wRR0Dk3Gq12DId99WogcMIVbL0QQFqVSevvlE,8772
 veris_ai/jaeger_interface/models.py,sha256=e64VV6IvOEFuzRUgvDAMQFyOZMRb56I-PUPZLBZ3rX0,1864
-veris_ai-1.11.1.dist-info/METADATA,sha256=H5_LRx0-4gOLrEj7iVMyhC1XJDHzbVwed-Dubugv8bE,15259
-veris_ai-1.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-veris_ai-1.11.1.dist-info/licenses/LICENSE,sha256=2g4i20atAgtD5einaKzhQrIB-JrPhyQgD3bC0wkHcCI,1065
-veris_ai-1.11.1.dist-info/RECORD,,
+veris_ai-1.12.0.dist-info/METADATA,sha256=NvHOxiR8I1gNUBmytOWjB5JvMy_bY9JYhObeoiGmPk8,15259
+veris_ai-1.12.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+veris_ai-1.12.0.dist-info/licenses/LICENSE,sha256=2g4i20atAgtD5einaKzhQrIB-JrPhyQgD3bC0wkHcCI,1065
+veris_ai-1.12.0.dist-info/RECORD,,