LangChain Integration

pip install agentsecretstore langchain langchain-google-genai langchain-community

import asyncio
from langchain.tools import BaseTool
from langchain_core.callbacks import CallbackManagerForToolRun
from agentsecretstore import AgentVault
from pydantic import BaseModel, Field

class GetSecretInput(BaseModel):
    path: str = Field(description="Secret path, e.g. 'production/gemini/GEMINI_API_KEY'")

class VaultSecretTool(BaseTool):
    """LangChain tool that retrieves a secret from Agent Secret Store."""

    name: str = "get_secret"
    description: str = (
        "Retrieve a credential or API key from the secure vault. "
        "Use this before making any external API calls that require authentication. "
        "Input: the secret path (e.g. 'production/gemini/GEMINI_API_KEY')."
    )
    args_schema: type[BaseModel] = GetSecretInput

    # The agent key is injected from the runtime environment.
    _agent_key: str

    def __init__(self, agent_key: str):
        super().__init__()
        self._agent_key = agent_key

    def _run(
        self,
        path: str,
        run_manager: CallbackManagerForToolRun | None = None,
    ) -> str:
        """Sync wrapper for the async vault call."""
        loop = asyncio.new_event_loop()
        try:
            return loop.run_until_complete(self._async_run(path))
        finally:
            loop.close()

    async def _arun(
        self,
        path: str,
        run_manager: CallbackManagerForToolRun | None = None,
    ) -> str:
        return await self._async_run(path)

    async def _async_run(self, path: str) -> str:
        from agentsecretstore import AgentVault

        async with AgentVault(agent_key=self._agent_key) as vault:
            return await vault.get_secret(path)

import asyncio
import os
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from agentsecretstore import AgentVault

async def build_agent(agent_key: str) -> AgentExecutor:
    """Build a LangChain agent with vault-backed credential access."""

    vault_tool = VaultSecretTool(agent_key=agent_key)

    # Agent only has access to secrets allowed by its agent-key scope
    tools = [vault_tool]

    prompt = ChatPromptTemplate.from_messages([
        ("system",
         "You are an AI assistant with access to a secure credential vault. "
         "When you need to call an external API, first retrieve the required "
         "credentials using the get_secret tool. Never ask the user for API keys."),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ])

    # The LLM key is fetched from the vault, not from process.env
    async with AgentVault(agent_key=agent_key) as v:
        gemini_secret = await v.get_secret("production/gemini/GEMINI_API_KEY")

    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash",
        api_key=gemini_secret,  # fetched from vault
        temperature=0,
    )

    agent = create_tool_calling_agent(llm, tools, prompt)

    return AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        max_iterations=10,
    )

async def main():
    agent_executor = await build_agent(agent_key=os.environ["ASS_AGENT_KEY"])

    result = await agent_executor.ainvoke({
        "input": "Search for recent AI research papers using Tavily and summarize them.",
        "chat_history": [],
    })

    print(result["output"])

asyncio.run(main())

from langchain.tools import StructuredTool
from agentsecretstore import AgentVault
import os

async def create_research_agent():
    """Agent that uses multiple services, each credential fetched from vault."""

    async def search_web(query: str) -> str:
        """Search using Tavily. Fetches API key from vault."""
        async with AgentVault(agent_key=os.environ["ASS_AGENT_KEY"]) as vault:
            key = await vault.get_secret("production/tavily/TAVILY_API_KEY")
        from langchain_community.tools.tavily_search import TavilySearchResults
        search = TavilySearchResults(api_key=key, max_results=5)
        return str(await search.ainvoke(query))

    async def send_slack_notification(message: str) -> str:
        """Send Slack notification. Fetches token from vault."""
        import httpx
        async with AgentVault(agent_key=os.environ["ASS_AGENT_KEY"]) as vault:
            token = await vault.get_secret("production/slack/SLACK_BOT_TOKEN")
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                "https://slack.com/api/chat.postMessage",
                headers={"Authorization": f"Bearer {token}"},
                json={"channel": "#ai-agents", "text": message},
            )
        return f"Slack notification sent: {resp.status_code}"

    tools = [
        StructuredTool.from_function(
            coroutine=search_web,
            name="web_search",
            description="Search the web for information",
        ),
        StructuredTool.from_function(
            coroutine=send_slack_notification,
            name="slack_notify",
            description="Send a Slack notification to the team",
        ),
    ]

    return tools