Skip to main content
API7.ai Logo

ai-rag

The ai-rag plugin implements the retrieval step of a Retrieval-Augmented Generation (RAG) request flow. It generates an embedding from the request and performs a vector search. It then adds the retrieved content to the protocol-specific LLM input and removes the ai_rag object before proxying the request.

The current implementation supports Azure OpenAI for embeddings and Azure AI Search for vector search. Use the ai-proxy plugin in the same request flow to proxy the augmented request to the LLM provider. The plugin does not create or populate a search index; prepare the index and its content before sending requests through APISIX.


Behavior by Request Format​

The plugin enriches Chat Completions, Responses API, Anthropic Messages, and Bedrock Converse requests using each protocol's native prompt structure.

The gateway identifies each request by checking URI-specific rules before body-only rules:

  • Bedrock Converse requires a URI ending in /converse and a messages array.
  • Anthropic Messages requires a URI ending in /v1/messages.
  • Responses API requires a URI ending in /v1/responses and an input field.
  • Chat Completions uses a messages array.
  • Embeddings uses input after the earlier rules do not match.
  • Other non-empty JSON objects use passthrough after none of the earlier rules match.
Request formatContext enrichment
Bedrock ConverseAppends the retrieved context as a user message in messages.
Anthropic MessagesAppends the retrieved context as a user message in messages.
Responses APIAppends the retrieved context to input.
Chat CompletionsAppends the retrieved context as a user message in messages.
EmbeddingsDoes not enrich the request. The nested ai_rag.embeddings object configures the embedding input used internally for retrieval.
Other JSON (passthrough)Does not enrich the request.

Verify Upstream TLS​

In API7 Enterprise from version 3.9.10 and APISIX from version 3.17.0, ssl_verify defaults to true for calls to the embedding and vector-search services. Configure a trusted certificate chain for both endpoints before enabling the plugin in production. Setting ssl_verify to false preserves connectivity to an endpoint with an untrusted certificate, but should be limited to temporary migration or development use.

Example​

To follow along the example, create an Azure account and complete the following steps:

Save the API keys and endpoints to environment variables:

# replace with your values

export AZ_OPENAI_DOMAIN=https://your-openai-resource.openai.azure.com
export AZ_OPENAI_API_KEY=your-azure-openai-api-key
export AZ_CHAT_ENDPOINT=${AZ_OPENAI_DOMAIN}/openai/deployments/gpt-4o/chat/completions?api-version=2024-02-15-preview
export AZ_EMBEDDING_MODEL=text-embedding-3-large
export AZ_EMBEDDINGS_ENDPOINT=${AZ_OPENAI_DOMAIN}/openai/deployments/${AZ_EMBEDDING_MODEL}/embeddings?api-version=2023-05-15

export AZ_AI_SEARCH_SVC_DOMAIN=https://your-search-service.search.windows.net
export AZ_AI_SEARCH_KEY=your-azure-ai-search-api-key
export AZ_AI_SEARCH_INDEX=vectest
export AZ_AI_SEARCH_ENDPOINT=${AZ_AI_SEARCH_SVC_DOMAIN}/indexes/${AZ_AI_SEARCH_INDEX}/docs/search?api-version=2024-07-01

Integrate with Azure for RAG-Enhanced Responses​

The following example demonstrates how you can use the ai-proxy plugin to proxy requests to Azure OpenAI LLM and use the ai-rag plugin to generate embeddings and perform vector search to enhance LLM responses.

Create a route as such:

curl "http://127.0.0.1:9180/apisix/admin/routes" -X PUT \
  -H "X-API-KEY: ${ADMIN_API_KEY}" \
  --data-binary @- <<EOF
  {
  "id": "ai-rag-route",
  "uri": "/rag",
  "plugins": {
    "ai-rag": {
      "embeddings_provider": {
        "azure_openai": {
          "endpoint": "$AZ_EMBEDDINGS_ENDPOINT",
          "api_key": "$AZ_OPENAI_API_KEY"
        }
      },
      "vector_search_provider": {
        "azure_ai_search": {
          "endpoint": "$AZ_AI_SEARCH_ENDPOINT",
          "api_key": "$AZ_AI_SEARCH_KEY"
        }
      }
    },
    "ai-proxy": {
      "provider": "openai",
      "auth": {
        "header": {
          "api-key": "$AZ_OPENAI_API_KEY"
        }
      },
      "model": "gpt-4o",
      "override": {
        "endpoint": "$AZ_CHAT_ENDPOINT"
      }
    }
  }
}
EOF

Send a POST request to the route with the vector fields name, embedding model dimensions, and an input prompt in the request body:

curl "http://127.0.0.1:9080/rag" -X POST \
  -H "Content-Type: application/json" \
  -d '{
    "ai_rag":{
      "vector_search":{
        "fields":"contentVector"
      },
      "embeddings":{
        "input":"Which Azure services are good for DevOps?",
        "dimensions":1024
      }
    }
  }'

You should receive an HTTP/1.1 200 OK response similar to the following:

{
  "choices": [
    {
      "content_filter_results": {
        ...
      },
      "finish_reason": "length",
      "index": 0,
      "logprobs": null,
      "message": {
        "content": "Here is a list of Azure services categorized along with a brief description of each based on the provided JSON data:\n\n### Developer Tools\n- **Azure DevOps**: A suite of services that help you plan, build, and deploy applications, including Azure Boards, Azure Repos, Azure Pipelines, Azure Test Plans, and Azure Artifacts.\n- **Azure DevTest Labs**: A fully managed service to create, manage, and share development and test environments in Azure, supporting custom templates, cost management, and integration with Azure DevOps.\n\n### Containers\n- **Azure Kubernetes Service (AKS)**: A managed container orchestration service based on Kubernetes, simplifying deployment and management of containerized applications with features like automatic upgrades and scaling.\n- **Azure Container Instances**: A serverless container runtime to run and scale containerized applications without managing the underlying infrastructure.\n- **Azure Container Registry**: A fully managed Docker registry service to store and manage container images and artifacts.\n\n### Web\n- **Azure App Service**: A fully managed platform for building, deploying, and scaling web apps, mobile app backends, and RESTful APIs with support for multiple programming languages.\n- **Azure SignalR Service**: A fully managed real-time messaging service to build and scale real-time web applications.\n- **Azure Static Web Apps**: A serverless hosting service for modern web applications using static front-end technologies and serverless APIs.\n\n### Compute\n- **Azure Virtual Machines**: Infrastructure-as-a-Service (IaaS) offering for deploying and managing virtual machines in the cloud.\n- **Azure Functions**: A serverless compute service to run event-driven code without managing infrastructure.\n- **Azure Batch**: A job scheduling service to run large-scale parallel and high-performance computing (HPC) applications.\n- **Azure Service Fabric**: A platform to build, deploy, and manage scalable and reliable microservices and container-based applications.\n- **Azure Quantum**: A quantum computing service to build and run quantum applications.\n- **Azure Stack Edge**: A managed edge computing appliance to run Azure services and AI workloads on-premises or at the edge.\n\n### Security\n- **Azure Bastion**: A fully managed service providing secure and scalable remote access to virtual machines.\n- **Azure Security Center**: A unified security management service to protect workloads across Azure and on-premises infrastructure.\n- **Azure DDoS Protection**: A cloud-based service to protect applications and resources from distributed denial-of-service (DDoS) attacks.\n\n### Databases\n",
        "role": "assistant"
      }
    }
  ],
  "created": 1740625850,
  "id": "chatcmpl-B54gQdumpfioMPIybFnirr6rq9ZZS",
  "model": "gpt-4o-2024-05-13",
  "object": "chat.completion",
  "prompt_filter_results": [
    {
      "prompt_index": 0,
      "content_filter_results": {
        ...
      }
    }
  ],
  "system_fingerprint": "fp_65792305e4",
  "usage": {
    ...
  }
}