Skip to main content

Custom RAG Agent with LangChain

note

This article is for custom agent development: writing a Python agent class because the built-in configurable agent is not enough. For standard search and retrieval-augmented generation, configure the built-in agent with the CLI instead - see Configuring Tools.

Use this pattern when you want to run a LangChain retrieval pipeline inside an Agents SDK agent. The example below uses ZAVLangchainStore as the retriever source and ChatOpenAI as the LangChain chat model.

Prerequisites​

Install the LangChain extra:

pip install zetaalpha.agents[langchain,langchain-openai]

Step 1: Create a Custom Agent​

Change to your agents project directory and scaffold a custom agent:

za agents add --custom "my_langchain_rag_agent"

This creates my_langchain_rag_agent.py in your agents project:

<agents project>/
├── .gitignore
├── __init__.py
├── my_langchain_rag_agent.py
├── agent_setups.json
└── env/
    └── agent_setups.json

Step 2: Implement the LangChain RAG Pipeline​

Replace the contents of my_langchain_rag_agent.py with:

from typing import List, Optional

from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

from zav.agents_sdk import ChatAgent, ChatAgentClassRegistry, ChatMessage
from zav.agents_sdk.adapters import ZAVLangchainStore


@ChatAgentClassRegistry.register()
class LangChainRAG(ChatAgent):
    agent_name = "my_langchain_rag_agent"

    def __init__(
        self,
        zav_langchain_store: ZAVLangchainStore,
        client: ChatOpenAI,
    ):
        self.retriever = zav_langchain_store.as_retriever()
        self.llm = client

    async def execute(self, conversation: List[ChatMessage]) -> Optional[ChatMessage]:
        prompt = hub.pull("rlm/rag-prompt")

        def format_docs(docs):
            return "\n\n".join(doc.page_content for doc in docs)

        rag_chain = (
            {"context": self.retriever | format_docs, "question": RunnablePassthrough()}
            | prompt
            | self.llm
            | StrOutputParser()
        )

        answer_msg = await rag_chain.ainvoke(conversation[-1].content)
        return ChatMessage(
            sender="bot",
            content=answer_msg,
        )

The agent converts the Zeta Alpha store into a LangChain retriever, uses the LangChain Hub RAG prompt, and returns the chain output as a chat message.

Step 3: Run the Agent Locally​

Run the local development UI:

za agents dev --reload

You can also serve the agent as an API:

za agents serve --reload

Then call the chat endpoint:

curl -X POST "http://localhost:8000/chats/responses?tenant=zetaalpha" \
-H "Content-Type: application/json" \
-d '{
  "agent_identifier": "my_langchain_rag_agent",
  "conversation": [
    {"sender": "user", "content": "What is a transformer?"}
  ]
}'

The response should contain the LangChain-generated answer.