From 0b86412fc0c5c62c25403074e6785ab532316ecc Mon Sep 17 00:00:00 2001 From: Cheng Zhang Date: Mon, 23 Feb 2026 13:02:26 +0000 Subject: [PATCH] new type hints --- .../control/resources/langgraph/__init__.py | 2 +- .../control/resources/langgraph/_executor.py | 8 +++++--- .../control/resources/langgraph/_runner.py | 2 +- src/sequrity/control/types/headers.py | 16 +++++++++------- 4 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/sequrity/control/resources/langgraph/__init__.py b/src/sequrity/control/resources/langgraph/__init__.py index 725c9ee..4da889c 100644 --- a/src/sequrity/control/resources/langgraph/__init__.py +++ b/src/sequrity/control/resources/langgraph/__init__.py @@ -22,7 +22,7 @@ def __init__(self, transport: ControlSyncTransport) -> None: def run( self, - graph: "StateGraph", + graph: StateGraph, initial_state: dict, model: str, *, diff --git a/src/sequrity/control/resources/langgraph/_executor.py b/src/sequrity/control/resources/langgraph/_executor.py index 5491234..6101835 100644 --- a/src/sequrity/control/resources/langgraph/_executor.py +++ b/src/sequrity/control/resources/langgraph/_executor.py @@ -4,6 +4,8 @@ Python code and maps nodes to tool definitions. """ +from __future__ import annotations + import json from typing import Callable @@ -29,7 +31,7 @@ class LangGraphExecutor: def __init__( self, - graph: "StateGraph", + graph: StateGraph, node_functions: dict[str, Callable] | None = None, internal_node_mapping: dict[str, str] | None = None, ): @@ -65,7 +67,7 @@ def __init__( # Generate code once at initialization self.generated_code = self._graph_to_code(graph, self.node_functions) - def _extract_function_map(self, graph: "StateGraph") -> dict[str, Callable]: + def _extract_function_map(self, graph: StateGraph) -> dict[str, Callable]: """Extract node functions from LangGraph StateGraph.""" function_map = {} @@ -160,7 +162,7 @@ def execute_tool_call(self, tool_call: dict, rest_api_type: RestApiType = RestAp return result - def _graph_to_code(self, graph: "StateGraph", function_map: dict[str, Callable]) -> str: + def _graph_to_code(self, graph: StateGraph, function_map: dict[str, Callable]) -> str: """Convert a LangGraph StateGraph into executable Python code.""" nodes = graph.nodes edges = list(graph.edges) diff --git a/src/sequrity/control/resources/langgraph/_runner.py b/src/sequrity/control/resources/langgraph/_runner.py index 272dc2b..9d48d5a 100644 --- a/src/sequrity/control/resources/langgraph/_runner.py +++ b/src/sequrity/control/resources/langgraph/_runner.py @@ -220,7 +220,7 @@ def _run_messages_loop( def run_graph_sync( transport: ControlSyncTransport, model: str, - graph: "StateGraph", + graph: StateGraph, initial_state: dict, *, initial_state_meta: MetaData | None = None, diff --git a/src/sequrity/control/types/headers.py b/src/sequrity/control/types/headers.py index a3c2b2a..f385757 100644 --- a/src/sequrity/control/types/headers.py +++ b/src/sequrity/control/types/headers.py @@ -4,6 +4,8 @@ HTTP headers: ``X-Features``, ``X-Policy``, and ``X-Config``. """ +from __future__ import annotations + from typing import Literal, TypeAlias, overload from pydantic import BaseModel, ConfigDict, Field @@ -99,7 +101,7 @@ def _build( finance_guardrail: bool, url_blocker: bool, file_blocker: bool, - ) -> "FeaturesHeader": + ) -> FeaturesHeader: classifiers: list[TaggerConfig] = [] if toxicity_filter: classifiers.append(TaggerConfig(name="toxicity_filter")) @@ -131,7 +133,7 @@ def single_llm( finance_guardrail: bool = False, url_blocker: bool = False, file_blocker: bool = False, - ) -> "FeaturesHeader": + ) -> FeaturesHeader: """Create a Single LLM features configuration.""" return cls._build( "single-llm", @@ -152,7 +154,7 @@ def dual_llm( finance_guardrail: bool = False, url_blocker: bool = False, file_blocker: bool = False, - ) -> "FeaturesHeader": + ) -> FeaturesHeader: """Create a Dual LLM features configuration.""" return cls._build( "dual-llm", @@ -283,7 +285,7 @@ def dual_llm( branching_meta_policy_producers: set[str] | None = None, branching_meta_policy_tags: set[str] | None = None, branching_meta_policy_consumers: set[str] | None = None, - ) -> "SecurityPolicyHeader": + ) -> SecurityPolicyHeader: """Create a Dual LLM security policy configuration.""" return cls.model_validate( { @@ -317,7 +319,7 @@ def single_llm( default_allow: bool = True, enable_llm_blocked_tag: bool = True, llm_blocked_tag_enforcement_level: Literal["hard", "soft"] = "hard", - ) -> "SecurityPolicyHeader": + ) -> SecurityPolicyHeader: """Create a Single LLM security policy configuration.""" return cls.model_validate( { @@ -607,7 +609,7 @@ def single_llm( min_num_tools_for_filtering: int | None = 10, clear_session_meta: Literal["never", "every_attempt", "every_turn"] = "never", max_n_turns: int | None = 50, - ) -> "FineGrainedConfigHeader": + ) -> FineGrainedConfigHeader: """Create a Single LLM fine-grained configuration.""" return cls( fsm=FsmOverrides( @@ -643,7 +645,7 @@ def dual_llm( include_program: bool | None = None, include_policy_check_history: bool | None = None, include_namespace_snapshot: bool | None = None, - ) -> "FineGrainedConfigHeader": + ) -> FineGrainedConfigHeader: """Create a Dual LLM fine-grained configuration.""" fsm = FsmOverrides( min_num_tools_for_filtering=min_num_tools_for_filtering,