Lint and mypy fixes

Author: ganeshyb

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/client/tools/_client.py

@@ -13,58 +13,6 @@
 from .operations._operations import MCPToolsOperations, RemoteToolsOperations
 from ._utils._model_base import InvocationPayloadBuilder
 from ._model_base import FoundryTool, ToolSource
-	
-class AzureAITool:
-	"""Azure AI tool wrapper for invocation.
-	
-	Represents a single tool that can be invoked either via MCP protocol or
-	Azure AI Tools API. This class provides a convenient interface for tool
-	invocation and exposes tool metadata.
-
-	:ivar str name: The name of the tool.
-	:ivar str description: Human-readable description of what the tool does.
-	:ivar dict metadata: Additional metadata about the tool from the API.
-	:ivar ~Tool_Client.models.ToolSource source:
-		The source of the tool (MCP_TOOLS or REMOTE_TOOLS).
-
-	.. admonition:: Example:
-
-		.. literalinclude:: ../samples/simple_example.py
-			:start-after: [START use_tool]
-			:end-before: [END use_tool]
-			:language: python
-			:dedent: 4
-			:caption: Using an AzureAITool instance.
-	"""
-
-	def __init__(self, client: "AzureAIToolClient", descriptor: FoundryTool) -> None:
-		"""Initialize an Azure AI Tool.
-		
-		:param client: Parent client instance for making API calls.
-		:type client: AzureAIToolClient
-		:param descriptor: Tool descriptor containing metadata and configuration.
-		:type descriptor: ~Tool_Client.models.FoundryTool
-		"""
-		self._client = client
-		self._descriptor = descriptor
-		self.name = descriptor.name
-		self.description = descriptor.description
-		self.metadata = dict(descriptor.metadata)
-		self.source = descriptor.source
-
-	def invoke(self, *args: Any, **kwargs: Any) -> Any:
-		"""Invoke the tool synchronously.
-		
-		:param args: Positional arguments to pass to the tool.
-		:param kwargs: Keyword arguments to pass to the tool.
-		:return: The result from the tool invocation.
-		:rtype: Any
-		"""
-		payload = InvocationPayloadBuilder.build_payload(args, kwargs, {})
-		return self._client._invoke_tool(self._descriptor, payload)
-
-	def __call__(self, *args: Any, **kwargs: Any) -> Any:
-		return self.invoke(*args, **kwargs)
 
 class AzureAIToolClient:
 	"""Synchronous client for aggregating tools from Azure AI MCP and Tools APIs.
@@ -189,11 +137,9 @@ def invoke_tool(
 		return self._invoke_tool(descriptor, payload, **kwargs)
 
 	def _resolve_tool_descriptor(
-		self, tool: Union[AzureAITool, str, FoundryTool]
+		self, tool: Union[str, FoundryTool]
 	) -> FoundryTool:
 		"""Resolve a tool reference to a descriptor."""
-		if isinstance(tool, AzureAITool):
-			return tool._descriptor
 		if isinstance(tool, FoundryTool):
 			return tool
 		if isinstance(tool, str):

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/client/tools/aio/_client.py

@@ -16,58 +16,6 @@
 if TYPE_CHECKING:
     from azure.core.credentials_async import AsyncTokenCredential
 
-class AzureAITool:
-	"""Azure AI tool wrapper for invocation.
-	
-	Represents a single tool that can be invoked either via MCP protocol or
-	Azure AI Tools API. This class provides a convenient interface for tool
-	invocation and exposes tool metadata.
-
-	:ivar str name: The name of the tool.
-	:ivar str description: Human-readable description of what the tool does.
-	:ivar dict metadata: Additional metadata about the tool from the API.
-	:ivar ~Tool_Client.models.ToolSource source:
-		The source of the tool (MCP_TOOLS or REMOTE_TOOLS).
-
-	.. admonition:: Example:
-
-		.. literalinclude:: ../samples/simple_example.py
-			:start-after: [START use_tool]
-			:end-before: [END use_tool]
-			:language: python
-			:dedent: 4
-			:caption: Using an AzureAITool instance.
-	"""
-
-	def __init__(self, client: "AzureAIToolClient", descriptor: FoundryTool) -> None:
-		"""Initialize an Azure AI Tool.
-		
-		:param client: Parent client instance for making API calls.
-		:type client: AzureAIToolClient
-		:param descriptor: Tool descriptor containing metadata and configuration.
-		:type descriptor: ~Tool_Client.models.FoundryTool
-		"""
-		self._client = client
-		self._descriptor = descriptor
-		self.name = descriptor.name
-		self.description = descriptor.description
-		self.metadata = dict(descriptor.metadata)
-		self.source = descriptor.source
-
-	async def invoke(self, *args: Any, **kwargs: Any) -> Any:
-		"""Invoke the tool asynchronously.
-
-		:param args: Positional arguments to pass to the tool.
-		:param kwargs: Keyword arguments to pass to the tool.
-		:return: The result from the tool invocation.
-		:rtype: Any
-		"""
-		payload = InvocationPayloadBuilder.build_payload(args, kwargs, {})
-		return await self._client._invoke_tool(self._descriptor, payload)
-
-	async def __call__(self, *args: Any, **kwargs: Any) -> Any:
-		return await self.invoke(*args, **kwargs)
-
 class AzureAIToolClient:
 	"""Asynchronous client for aggregating tools from Azure AI MCP and Tools APIs.
 
@@ -177,7 +125,7 @@ async def _invoker(*args, **kwargs):
 
 	async def invoke_tool(
 		self,
-		tool: Union[AzureAITool, str, FoundryTool],
+		tool: Union[str, FoundryTool],
 		*args: Any,
 		**kwargs: Any,
 	) -> Any:
@@ -193,11 +141,9 @@ async def invoke_tool(
 		return await self._invoke_tool(descriptor, payload, **kwargs)
 
 	async def _resolve_tool_descriptor(
-		self, tool: Union[AzureAITool, str, FoundryTool]
+		self, tool: Union[str, FoundryTool]
 	) -> FoundryTool:
 		"""Resolve a tool reference to a descriptor."""
-		if isinstance(tool, AzureAITool):
-			return tool._descriptor
 		if isinstance(tool, FoundryTool):
 			return tool
 		if isinstance(tool, str):

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/constants.py

@@ -11,3 +11,4 @@ class Constants:
     AGENT_LOG_LEVEL = "AGENT_LOG_LEVEL"
     AGENT_DEBUG_ERRORS = "AGENT_DEBUG_ERRORS"
     ENABLE_APPLICATION_INSIGHTS_LOGGER = "AGENT_APP_INSIGHTS_ENABLED"
+    AZURE_AI_WORKSPACE_ENDPOINT = "AZURE_AI_WORKSPACE_ENDPOINT"

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py

@@ -88,22 +88,23 @@ def set_run_context_to_context_var(self, run_context):
         ctx.update(res)
         request_context.set(ctx)
 
-    def set_user_info_to_context_var(self, request):
-        user_info: UserInfo = {}
+    def set_user_info_to_context_var(self, request) -> UserInfo:
+        user_info: UserInfo = None
         try:
             object_id_header = request.headers.get("x-aml-oid", None)
-            tenant_id_header = request.headers.get("x-aml-tenant-id", None)
-
-            if object_id_header:
-                user_info["object_id"] = object_id_header
-            if tenant_id_header:
-                user_info["tenant_id"] = tenant_id_header
+            tenant_id_header = request.headers.get("x-aml-tid", None)
+            if not object_id_header and not tenant_id_header:
+                return None
+            user_info = UserInfo(
+                objectId=object_id_header,
+                tenantId=tenant_id_header
+            )
 
         except Exception as e:
             logger.error(f"Failed to parse X-User-Info header: {e}", exc_info=True)
         if user_info:
             ctx = request_context.get() or {}
-            for key, value in user_info.items():
+            for key, value in user_info.to_dict().items():
                 ctx[f"azure.ai.agentserver.user.{key}"] = str(value)
             request_context.set(ctx)
         return user_info
@@ -340,12 +341,26 @@ def get_tool_client(
         logger.debug("Creating AzureAIToolClient with tools: %s", tools)
         if not self.credentials:
             raise ValueError("Credentials are required to create Tool Client.")
-        return AzureAIToolClient(
+        
+        workspace_endpoint = os.getenv(Constants.AZURE_AI_WORKSPACE_ENDPOINT)
+        if workspace_endpoint:
+            agent_name = os.getenv(Constants.AGENT_NAME)
+            if not agent_name:
+                raise ValueError("AGENT_NAME environment variable is required when using workspace endpoint.")
+            return AzureAIToolClient(
+            endpoint=workspace_endpoint,
+            credential=self.credentials,
+            tools=tools,
+            user=user_info,
+            agent_name=agent_name,
+            )
+        else:
+            return AzureAIToolClient(
             endpoint=os.getenv(Constants.AZURE_AI_PROJECT_ENDPOINT),
             credential=self.credentials,
-            tools = tools,
-            user = user_info,
-        )
+            tools=tools,
+            user=user_info,
+            )
 
 
 def _event_to_sse_chunk(event: ResponseStreamEvent) -> str:

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/common/agent_run_context.py

@@ -7,7 +7,7 @@
 from ...models.projects import AgentId, AgentReference, ResponseConversation1
 from .id_generator.foundry_id_generator import FoundryIdGenerator
 from .id_generator.id_generator import IdGenerator
-
+from ...client.tools._model_base import UserInfo
 logger = get_logger()
 
 
@@ -19,7 +19,7 @@ def __init__(self, payload: dict, **kwargs: Any) -> None:
         self._response_id = self._id_generator.response_id
         self._conversation_id = self._id_generator.conversation_id
         self._stream = self.request.get("stream", False)
-        self._user_info = kwargs.get("user_info", {})
+        self._user_info = kwargs.get("user_info", None)
         self._agent_tools = kwargs.get("agent_tools", [])
 
     @property
@@ -70,7 +70,7 @@ def get_tools(self) -> list:
             return self._agent_tools
         return request_tools
 
-    def get_user_info(self) -> dict:
+    def get_user_info(self) -> UserInfo:
         return self._user_info
 
 

sdk/agentserver/azure-ai-agentserver-langgraph/azure/ai/agentserver/langgraph/__init__.py

@@ -7,6 +7,7 @@
 
 from ._version import VERSION
 from .tool_client import ToolClient
+from .langgraph import LangGraphAdapter
 
 if TYPE_CHECKING:  # pragma: no cover
     from . import models
@@ -19,7 +20,6 @@ def from_langgraph(
     state_converter: Optional["models.LanggraphStateConverter"] = None,
     **kwargs: Any
 ) -> "LangGraphAdapter":
-    from .langgraph import LangGraphAdapter
 
     return LangGraphAdapter(agent, credentials=credentials, state_converter=state_converter, **kwargs)
 

sdk/agentserver/azure-ai-agentserver-langgraph/azure/ai/agentserver/langgraph/langgraph.py

@@ -70,7 +70,7 @@ def __init__(
         :param state_converter: custom state converter. Required if graph state is not MessagesState.
         :type state_converter: Optional[LanggraphStateConverter]
         """
-        super().__init__(credentials=credentials, **kwargs)
+        super().__init__(credentials=credentials, **kwargs) # pylint: disable=unexpected-keyword-arg
         self._graph_or_factory: Union[CompiledStateGraph, GraphFactory] = graph
         self._resolved_graph: "Optional[CompiledStateGraph]" = None
         self.azure_ai_tracer = None
@@ -140,7 +140,7 @@ async def _resolve_graph(self, context: AgentRunContext):
 
 
             # Create ToolClient with credentials
-            tool_client = self.get_tool_client(tools = context.get_tools(), user_info = context.get_user_info())
+            tool_client = self.get_tool_client(tools = context.get_tools(), user_info = context.get_user_info()) # pylint: disable=no-member
             tool_client_wrapper = ToolClient(tool_client)
             tools = await tool_client_wrapper.list_tools()
             # Call the factory function with ToolClient
@@ -153,7 +153,7 @@ async def _resolve_graph(self, context: AgentRunContext):
                 self._resolved_graph = result
 
             # Validate and set up state converter if not already set from initialization
-            if not self.state_converter:
+            if not self.state_converter and self._resolved_graph is not None:
                 if is_state_schema_valid(self._resolved_graph.builder.state_schema):
                     self.state_converter = LanggraphMessageStateConverter()
                 else:
@@ -178,13 +178,13 @@ async def _resolve_graph_for_request(self, context: AgentRunContext):
         logger.debug("Resolving fresh graph from factory function for request")
 
         # Create ToolClient with credentials
-        tool_client = self.get_tool_client(tools = context.get_tools(), user_info = context.get_user_info())
+        tool_client = self.get_tool_client(tools = context.get_tools(), user_info = context.get_user_info()) # pylint: disable=no-member
         tool_client_wrapper = ToolClient(tool_client)
         tools = await tool_client_wrapper.list_tools()
         # Call the factory function with ToolClient
         # Support both sync and async factories
         import inspect
-        result = self._graph_or_factory(tools)
+        result = self._graph_or_factory(tools)  # type: ignore[operator]
         if inspect.iscoroutine(result):
             graph = await result
         else:

sdk/agentserver/azure-ai-agentserver-langgraph/azure/ai/agentserver/langgraph/tool_client.py

@@ -64,7 +64,7 @@ def __init__(self, tool_client: "AzureAIToolClient") -> None:
         :type tool_client: ~azure.ai.agentserver.core.client.tools.aio.AzureAIToolClient
         """
         self._tool_client = tool_client
-        self._langchain_tools_cache: List[StructuredTool] = None
+        self._langchain_tools_cache: Optional[List[StructuredTool]] = None
 
     async def list_tools(self) -> List[StructuredTool]:
         """List all available tools as LangChain BaseTool instances.
@@ -115,7 +115,7 @@ def _convert_to_langchain_tool(self, azure_tool: "FoundryTool") -> StructuredToo
         # Create a Pydantic model for the tool's input schema
         args_schema = self._create_pydantic_model(
             tool_name=azure_tool.name,
-            schema=input_schema
+            schema=dict(input_schema)
         )
 
         # Create an async function that invokes the tool
@@ -176,13 +176,13 @@ def _create_pydantic_model(
                 )
             else:
                 field_definitions[prop_name] = (
-                    Optional[prop_type],
-                    Field(None, description=prop_description)
+                    prop_type,
+                    Field(default=None, description=prop_description)
                 )
 
         # Create the model dynamically
-        model_name = f"{tool_name.replace('-', '_').replace(' ', '_').title()}Input"
-        return create_model(model_name, **field_definitions)
+        model_name = f"{tool_name.replace('-', '_').replace(' ', '_').title()}-Input"
+        return create_model(model_name, **field_definitions)  # type: ignore[call-overload]
 
     def _json_type_to_python_type(self, json_type: str) -> type:
         """Convert JSON schema type to Python type.