afm-langchain 0.1.0.dev1__tar.gz

afm_langchain-0.1.0.dev1/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: afm-langchain
+Version: 0.1.0.dev1
+Summary: AFM LangChain execution backend
+License-Expression: Apache-2.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: afm-core
+Requires-Dist: langchain>=1.2.8
+Requires-Dist: langchain-openai>=1.1.7
+Requires-Dist: langchain-anthropic>=1.3.1
+Requires-Dist: mcp>=1.26.0
+Requires-Dist: langchain-mcp-adapters>=0.2.1
+Requires-Python: >=3.11
+Project-URL: Repository, https://github.com/wso2/reference-implementations-afm
+Description-Content-Type: text/markdown
+

afm_langchain-0.1.0.dev1/pyproject.toml

@@ -0,0 +1,39 @@
+[project]
+name = "afm-langchain"
+version = "0.1.0.dev1"
+description = "AFM LangChain execution backend"
+readme = "README.md"
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Programming Language :: Python :: 3.11",
+]
+license = "Apache-2.0"
+requires-python = ">=3.11"
+urls = { Repository = "https://github.com/wso2/reference-implementations-afm" }
+dependencies = [
+    "afm-core",
+    "langchain>=1.2.8",
+    "langchain-openai>=1.1.7",
+    "langchain-anthropic>=1.3.1",
+    "mcp>=1.26.0",
+    "langchain-mcp-adapters>=0.2.1",
+]
+
+[project.entry-points."afm.runner"]
+langchain = "afm_langchain.backend:LangChainRunner"
+
+[build-system]
+requires = ["uv_build>=0.9.28,<0.10.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "afm_langchain"
+
+[tool.uv.sources]
+afm-core = { workspace = true }
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["src"]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"

afm_langchain-0.1.0.dev1/src/afm_langchain/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2025
+# Licensed under the Apache License, Version 2.0
+
+"""AFM LangChain execution backend."""
+
+from .backend import LangChainRunner
+
+__all__ = ["LangChainRunner"]

afm_langchain-0.1.0.dev1/src/afm_langchain/backend.py

@@ -0,0 +1,295 @@
+# Copyright (c) 2025
+# Licensed under the Apache License, Version 2.0
+
+from __future__ import annotations
+
+import json
+import logging
+from types import TracebackType
+from typing import Any
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.tools import BaseTool
+
+from afm.exceptions import AgentError, InputValidationError, OutputValidationError
+from afm.models import (
+    AFMRecord,
+    Interface,
+    Signature,
+)
+from afm.schema_validator import (
+    build_output_schema_instruction,
+    coerce_output_to_schema,
+    validate_input,
+)
+
+from .providers import create_model_provider
+from .tools.mcp import MCPManager
+
+logger = logging.getLogger(__name__)
+
+
+class LangChainRunner:
+    def __init__(
+        self,
+        afm: AFMRecord,
+        *,
+        model: BaseChatModel | None = None,
+        tools: list[BaseTool] | None = None,
+    ) -> None:
+        self._afm = afm
+        self._base_model = model or create_model_provider(afm.metadata.model)
+        self._model = self._base_model  # Will be updated with tools when connected
+        self._sessions: dict[str, list[HumanMessage | AIMessage]] = {}
+
+        # MCP management
+        self._mcp_manager = MCPManager.from_afm(afm)
+        self._external_tools = tools or []
+        self._mcp_tools: list[BaseTool] = []
+        self._connected = False
+
+        # Cache the active interface for signature validation
+        self._interface = self._get_primary_interface()
+        self._signature = self._get_signature()
+
+    async def __aenter__(self) -> "LangChainRunner":
+        await self.connect()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.disconnect()
+
+    async def connect(self) -> None:
+        if self._connected:
+            return
+
+        if self._mcp_manager is not None:
+            logger.info(f"Connecting to MCP servers: {self._mcp_manager.server_names}")
+            self._mcp_tools = await self._mcp_manager.get_tools()
+            logger.info(f"Loaded {len(self._mcp_tools)} MCP tools")
+
+        # Bind tools to model if any are available
+        all_tools = self._get_all_tools()
+        if all_tools:
+            self._model = self._base_model.bind_tools(all_tools)
+            logger.info(f"Bound {len(all_tools)} tools to model")
+
+        self._connected = True
+
+    async def disconnect(self) -> None:
+        if not self._connected:
+            return
+
+        # Clear MCP tools and reset model
+        self._mcp_tools = []
+        self._model = self._base_model
+        self._connected = False
+
+        if self._mcp_manager is not None:
+            self._mcp_manager.clear_cache()
+            logger.info("Disconnected from MCP servers")
+
+    def _get_all_tools(self) -> list[BaseTool]:
+
+        return self._external_tools + self._mcp_tools
+
+    @property
+    def afm(self) -> AFMRecord:
+        return self._afm
+
+    @property
+    def name(self) -> str:
+        return self._afm.metadata.name or "AFM Agent"
+
+    @property
+    def description(self) -> str | None:
+        return self._afm.metadata.description
+
+    @property
+    def system_prompt(self) -> str:
+        return f"""# Role
+
+{self._afm.role}
+
+# Instructions
+
+{self._afm.instructions}"""
+
+    @property
+    def max_iterations(self) -> int | None:
+        return self._afm.metadata.max_iterations
+
+    @property
+    def tools(self) -> list[BaseTool]:
+        return self._get_all_tools()
+
+    @property
+    def signature(self) -> Signature:
+        return self._signature
+
+    def _get_primary_interface(self) -> Interface | None:
+        interfaces = self._afm.metadata.interfaces
+        if interfaces and len(interfaces) > 0:
+            return interfaces[0]
+        return None
+
+    def _get_signature(self) -> Signature:
+        if self._interface is not None:
+            return self._interface.signature
+        # Default: string input/output
+        return Signature()
+
+    def _get_session_history(self, session_id: str) -> list[HumanMessage | AIMessage]:
+        if session_id not in self._sessions:
+            self._sessions[session_id] = []
+        return self._sessions[session_id]
+
+    def _prepare_input(self, input_data: str | dict[str, Any]) -> str:
+        input_schema = self._signature.input
+
+        # Validate input
+        validate_input(input_data, input_schema)
+
+        # Convert to string for the LLM
+        if isinstance(input_data, str):
+            return input_data
+        return json.dumps(input_data)
+
+    def _build_messages(
+        self,
+        user_input: str,
+        session_history: list[HumanMessage | AIMessage],
+    ) -> list[SystemMessage | HumanMessage | AIMessage | ToolMessage]:
+        messages: list[SystemMessage | HumanMessage | AIMessage | ToolMessage] = [
+            SystemMessage(content=self.system_prompt)
+        ]
+
+        # Add conversation history
+        messages.extend(session_history)
+
+        # Add the current user message
+        # If output schema is not string, append schema instructions
+        output_schema = self._signature.output
+        if output_schema.type != "string":
+            schema_instruction = build_output_schema_instruction(output_schema)
+            user_input = user_input + schema_instruction
+
+        messages.append(HumanMessage(content=user_input))
+
+        return messages
+
+    def _extract_response_content(self, response: Any) -> str:
+        if isinstance(response, AIMessage):
+            content = response.content
+        else:
+            content = str(response)
+
+        if not isinstance(content, str):
+            content = str(content)
+
+        return content
+
+    async def arun(
+        self,
+        input_data: str | dict[str, Any],
+        *,
+        session_id: str = "default",
+    ) -> str | dict[str, Any]:
+        try:
+            # Prepare and validate input
+            user_input = self._prepare_input(input_data)
+
+            # Get session history
+            session_history = self._get_session_history(session_id)
+
+            # Save the original input for history before schema augmentation
+            original_input = user_input
+
+            # Build messages
+            messages: list[Any] = self._build_messages(user_input, session_history)
+
+            # Max iterations for tool use
+            max_iterations = (
+                self.max_iterations if self.max_iterations is not None else 10
+            )
+            iterations = 0
+            response = None
+
+            # Main agent loop to handle tool calls
+            while iterations < max_iterations:
+                # Invoke the LLM asynchronously
+                response = await self._model.ainvoke(messages)
+
+                # If no tool calls, we're done
+                if not response.tool_calls:
+                    break
+
+                # Add the assistant message (containing tool calls) to the conversation
+                messages.append(response)
+
+                # Execute tool calls
+                for tool_call in response.tool_calls:
+                    # Find the tool
+                    tool_name = tool_call["name"]
+                    tool = next((t for t in self.tools if t.name == tool_name), None)
+
+                    if tool is None:
+                        tool_output = f"Error: Tool '{tool_name}' not found."
+                    else:
+                        try:
+                            # Run the tool
+                            tool_output = await tool.ainvoke(tool_call["args"])
+                        except Exception as e:
+                            tool_output = f"Error executing tool '{tool_name}': {e}"
+
+                    # Add tool response to messages
+                    messages.append(
+                        ToolMessage(
+                            content=str(tool_output),
+                            tool_call_id=tool_call["id"],
+                        )
+                    )
+
+                iterations += 1
+
+            if iterations >= max_iterations and response and response.tool_calls:
+                logger.warning(
+                    f"Max iterations ({max_iterations}) reached with "
+                    f"{len(response.tool_calls)} pending tool calls: "
+                    f"{[tc['name'] for tc in response.tool_calls]}"
+                )
+
+            if response is None:
+                raise AgentError("No response from LLM")
+
+            # Extract content from response
+            response_content = self._extract_response_content(response)
+
+            # Validate and coerce output
+            output_schema = self._signature.output
+            result = coerce_output_to_schema(response_content, output_schema)
+
+            # Update session history
+            session_history.append(HumanMessage(content=original_input))
+            session_history.append(AIMessage(content=response_content))
+
+            return result
+
+        except InputValidationError:
+            raise
+        except OutputValidationError:
+            raise
+        except Exception as e:
+            if isinstance(e, AgentError):
+                raise
+            raise AgentError(f"Agent execution failed: {e}") from e
+
+    def clear_history(self, session_id: str = "default") -> None:
+        if session_id in self._sessions:
+            del self._sessions[session_id]

afm_langchain-0.1.0.dev1/src/afm_langchain/providers.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2025
+# Licensed under the Apache License, Version 2.0
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+from langchain_core.language_models import BaseChatModel
+
+from afm.exceptions import ProviderError
+from afm.models import ClientAuthentication, Model
+
+if TYPE_CHECKING:
+    from langchain_anthropic import ChatAnthropic
+    from langchain_openai import ChatOpenAI
+
+
+DEFAULT_OPENAI_MODEL = "gpt-4o"
+DEFAULT_ANTHROPIC_MODEL = "claude-sonnet-4-5"
+
+DEFAULT_OPENAI_URL = "https://api.openai.com/v1"
+DEFAULT_ANTHROPIC_URL = "https://api.anthropic.com"
+
+# Environment variable names for API keys
+OPENAI_API_KEY_ENV = "OPENAI_API_KEY"
+ANTHROPIC_API_KEY_ENV = "ANTHROPIC_API_KEY"
+
+
+def create_model_provider(afm_model: Model | None = None) -> BaseChatModel:
+    if afm_model is None:
+        return _create_openai_model(None)
+
+    provider = (afm_model.provider or "openai").lower()
+
+    match provider:
+        case "openai":
+            return _create_openai_model(afm_model)
+        case "anthropic":
+            return _create_anthropic_model(afm_model)
+        case _:
+            raise ProviderError(f"Unsupported provider: {provider}", provider=provider)
+
+
+def _create_openai_model(afm_model: Model | None) -> ChatOpenAI:
+    try:
+        from langchain_openai import ChatOpenAI
+    except ImportError as e:
+        raise ProviderError(
+            "langchain-openai package is required for OpenAI models. "
+            "Install it with: pip install langchain-openai",
+            provider="openai",
+        ) from e
+
+    api_key = _get_api_key(
+        afm_model.authentication if afm_model else None,
+        OPENAI_API_KEY_ENV,
+        "openai",
+    )
+    model_name = (
+        afm_model.name if afm_model and afm_model.name else DEFAULT_OPENAI_MODEL
+    )
+    base_url = afm_model.url if afm_model and afm_model.url else None
+
+    kwargs: dict = {
+        "api_key": api_key,
+        "model": model_name,
+    }
+    if base_url:
+        kwargs["base_url"] = base_url
+
+    return ChatOpenAI(**kwargs)
+
+
+def _create_anthropic_model(afm_model: Model) -> ChatAnthropic:
+    try:
+        from langchain_anthropic import ChatAnthropic
+    except ImportError as e:
+        raise ProviderError(
+            "langchain-anthropic package is required for Anthropic models. "
+            "Install it with: pip install langchain-anthropic",
+            provider="anthropic",
+        ) from e
+
+    api_key = _get_api_key(
+        afm_model.authentication,
+        ANTHROPIC_API_KEY_ENV,
+        "anthropic",
+    )
+    model_name = afm_model.name if afm_model.name else DEFAULT_ANTHROPIC_MODEL
+    base_url = afm_model.url if afm_model.url else None
+
+    kwargs: dict = {
+        "api_key": api_key,
+        "model": model_name,
+    }
+    if base_url:
+        kwargs["base_url"] = base_url
+
+    return ChatAnthropic(**kwargs)
+
+
+def _get_api_key(
+    auth: ClientAuthentication | None,
+    env_var: str,
+    provider: str,
+) -> str:
+    if auth is not None:
+        auth_type = auth.type.lower()
+
+        if auth_type == "bearer" and auth.token:
+            return auth.token
+        elif auth_type == "api-key" and auth.api_key:
+            return auth.api_key
+        elif auth_type == "basic":
+            raise ProviderError(
+                "Basic authentication is not supported for LLM providers. "
+                "Use 'bearer' or 'api-key' authentication type.",
+                provider=provider,
+            )
+        else:
+            auth_dict = auth.model_dump(exclude_none=True)
+            for key in ["token", "api_key", "key", "apiKey"]:
+                if key in auth_dict and auth_dict[key]:
+                    return auth_dict[key]
+
+    # Fall back to environment variable
+    api_key = os.environ.get(env_var)
+    if api_key:
+        return api_key
+
+    raise ProviderError(
+        f"No API key found. Provide authentication in the model config "
+        f"or set the {env_var} environment variable.",
+        provider=provider,
+    )

afm_langchain-0.1.0.dev1/src/afm_langchain/tools/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2025
+# Licensed under the Apache License, Version 2.0

afm_langchain-0.1.0.dev1/src/afm_langchain/tools/mcp.py

@@ -0,0 +1,244 @@
+# Copyright (c) 2025
+# Licensed under the Apache License, Version 2.0
+
+from __future__ import annotations
+
+import logging
+
+import httpx
+from langchain_core.tools import BaseTool
+from langchain_mcp_adapters.client import MultiServerMCPClient
+from langchain_mcp_adapters.sessions import StreamableHttpConnection
+
+from afm.exceptions import (
+    MCPAuthenticationError,
+    MCPConnectionError,
+    MCPError,
+)
+from afm.models import (
+    AFMRecord,
+    ClientAuthentication,
+    MCPServer,
+    ToolFilter,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class BearerAuth(httpx.Auth):
+    def __init__(self, token: str) -> None:
+        self.token = token
+
+    def auth_flow(self, request: httpx.Request):
+        request.headers["Authorization"] = f"Bearer {self.token}"
+        yield request
+
+
+class ApiKeyAuth(httpx.Auth):
+    def __init__(self, api_key: str, header_name: str = "Authorization") -> None:
+        self.api_key = api_key
+        self.header_name = header_name
+
+    def auth_flow(self, request: httpx.Request):
+        request.headers[self.header_name] = self.api_key
+        yield request
+
+
+def build_httpx_auth(auth: ClientAuthentication | None) -> httpx.Auth | None:
+    if auth is None:
+        return None
+
+    auth_type = auth.type.lower()
+
+    if auth_type == "bearer":
+        if auth.token is None:
+            raise MCPAuthenticationError("Bearer auth requires 'token' field")
+        return BearerAuth(auth.token)
+
+    elif auth_type == "basic":
+        if auth.username is None or auth.password is None:
+            raise MCPAuthenticationError(
+                "Basic auth requires 'username' and 'password' fields"
+            )
+        return httpx.BasicAuth(auth.username, auth.password)
+
+    elif auth_type == "api-key":
+        if auth.api_key is None:
+            raise MCPAuthenticationError("API key auth requires 'api_key' field")
+        return ApiKeyAuth(auth.api_key)
+
+    elif auth_type in ("oauth2", "jwt"):
+        raise MCPAuthenticationError(
+            f"Authentication type '{auth_type}' not yet supported"
+        )
+
+    else:
+        raise MCPAuthenticationError(f"Unsupported authentication type: {auth_type}")
+
+
+def filter_tools(
+    tools: list[BaseTool],
+    tool_filter: ToolFilter | None,
+) -> list[BaseTool]:
+    if tool_filter is None:
+        return tools
+
+    allow = tool_filter.allow
+    deny = tool_filter.deny
+
+    # No filters specified
+    if allow is None and deny is None:
+        return tools
+
+    # Build a set of tool names for efficient lookup
+    tool_names = {tool.name for tool in tools}
+
+    if allow is not None:
+        # Start with allowed tools only
+        allowed_set = set(allow) & tool_names
+    else:
+        # Start with all tools
+        allowed_set = tool_names
+
+    if deny is not None:
+        # Remove denied tools
+        allowed_set -= set(deny)
+
+    # Filter the tools list maintaining order
+    return [tool for tool in tools if tool.name in allowed_set]
+
+
+class MCPClient:
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        authentication: ClientAuthentication | None = None,
+        tool_filter: ToolFilter | None = None,
+    ) -> None:
+        """Initialize an MCP client."""
+        self.name = name
+        self.url = url
+        self.authentication = authentication
+        self.tool_filter = tool_filter
+        self._tools: list[BaseTool] | None = None
+
+    @classmethod
+    def from_mcp_server(cls, server: MCPServer) -> "MCPClient":
+        transport = server.transport
+
+        if transport.type != "http":
+            raise MCPError(
+                f"Unsupported transport type: {transport.type}. Only 'http' is supported for now.",
+                server_name=server.name,
+            )
+
+        return cls(
+            name=server.name,
+            url=transport.url,
+            authentication=transport.authentication,
+            tool_filter=server.tool_filter,
+        )
+
+    def _build_connection_config(self) -> StreamableHttpConnection:
+        config: StreamableHttpConnection = {
+            "transport": "streamable_http",
+            "url": self.url,
+        }
+
+        # Add authentication if configured
+        auth = build_httpx_auth(self.authentication)
+        if auth is not None:
+            config["auth"] = auth
+
+        return config
+
+    async def get_tools(self) -> list[BaseTool]:
+        try:
+            # Create a client for just this server
+            client = MultiServerMCPClient({self.name: self._build_connection_config()})
+
+            # Get tools from the server
+            tools = await client.get_tools(server_name=self.name)
+
+            # Apply filtering
+            filtered_tools = filter_tools(tools, self.tool_filter)
+
+            logger.info(
+                f"MCP server '{self.name}': loaded {len(filtered_tools)} tools "
+                f"(filtered from {len(tools)})"
+            )
+
+            return filtered_tools
+
+        except MCPError:
+            # Re-raise MCPError subclasses (like MCPAuthenticationError) to preserve diagnostics
+            raise
+        except Exception as e:
+            raise MCPConnectionError(
+                f"Failed to connect: {e}",
+                server_name=self.name,
+            ) from e
+
+
+class MCPManager:
+    def __init__(self, servers: list[MCPServer]) -> None:
+        """Initialize the MCP manager."""
+        self._servers = servers
+        self._clients: list[MCPClient] = []
+        self._tools: list[BaseTool] | None = None
+
+        # Create clients for each server
+        for server in servers:
+            try:
+                client = MCPClient.from_mcp_server(server)
+                self._clients.append(client)
+            except MCPError as e:
+                logger.warning(f"Skipping MCP server: {e}")
+
+    @classmethod
+    def from_afm(cls, afm: AFMRecord) -> "MCPManager | None":
+        tools_config = afm.metadata.tools
+        if tools_config is None:
+            return None
+
+        mcp_servers = tools_config.mcp
+        if mcp_servers is None or len(mcp_servers) == 0:
+            return None
+
+        return cls(mcp_servers)
+
+    @property
+    def server_names(self) -> list[str]:
+        return [client.name for client in self._clients]
+
+    async def get_tools(self) -> list[BaseTool]:
+        if self._tools is not None:
+            return self._tools
+
+        all_tools: list[BaseTool] = []
+        errors: list[str] = []
+
+        # Get tools from each client individually to handle per-server filtering
+        for client in self._clients:
+            try:
+                tools = await client.get_tools()
+                all_tools.extend(tools)
+            except MCPConnectionError as e:
+                errors.append(str(e))
+                logger.error(f"Failed to get tools from server '{client.name}': {e}")
+
+        if errors and not all_tools:
+            raise MCPConnectionError(
+                f"Failed to connect to any MCP server: {'; '.join(errors)}"
+            )
+
+        # Only cache if all servers succeeded; partial results are not
+        # cached so that failed servers can be retried on the next call.
+        if not errors:
+            self._tools = all_tools
+
+        return all_tools
+
+    def clear_cache(self) -> None:
+        self._tools = None