Integration: AISIX AI Gateway

pip install haystack-ai==3.1.1

_format_version: "1"

provider_keys:
  - display_name: haystack-provider
    provider: openai
    adapter: openai
    api_key: ${UPSTREAM_API_KEY}
    api_base: ${UPSTREAM_API_BASE}

models:
  - display_name: rag-primary
    provider: openai
    model_name: ${UPSTREAM_PRIMARY_MODEL}
    provider_key: haystack-provider
  - display_name: rag-backup
    provider: openai
    model_name: ${UPSTREAM_BACKUP_MODEL}
    provider_key: haystack-provider
  - display_name: rag-model
    routing:
      strategy: failover
      retries: 0
      max_fallbacks: 1
      targets:
        - model: rag-primary
        - model: rag-backup

api_keys:
  - display_name: haystack-caller
    key_env: HAYSTACK_CALLER_KEY
    allowed_models:
      - rag-model

aisix validate --resources /etc/aisix/resources.yaml

export AISIX_BASE_URL="http://localhost:3000"
export HAYSTACK_CALLER_KEY="replace-with-a-gateway-caller-key"

import os

from haystack import Document, Pipeline
from haystack.components.builders import ChatPromptBuilder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.dataclasses import ChatMessage
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.utils import Secret

document_store = InMemoryDocumentStore()
document_store.write_documents(
    [
        Document(
            content=(
                "If the checkout service returns HTTP 503, fail over to the "
                "secondary region and page the on-call engineer."
            )
        ),
        Document(
            content=(
                "If checkout latency exceeds two seconds, inspect the payment "
                "provider dashboard before scaling the application."
            )
        ),
    ]
)

template = [
    ChatMessage.from_system(
        "Answer only from the supplied runbook documents."
    ),
    ChatMessage.from_user(
        "Runbook documents:\n"
        "{% for document in documents %}"
        "- {{ document.content }}\n"
        "{% endfor %}"
        "Question: {{ question }}"
    ),
]

pipeline = Pipeline()
pipeline.add_component(
    "retriever",
    InMemoryBM25Retriever(document_store=document_store, top_k=1),
)
pipeline.add_component(
    "prompt_builder",
    ChatPromptBuilder(
        template=template,
        required_variables={"documents", "question"},
    ),
)
pipeline.add_component(
    "llm",
    OpenAIChatGenerator(
        api_key=Secret.from_env_var("HAYSTACK_CALLER_KEY"),
        model="rag-model",
        api_base_url=f"{os.environ['AISIX_BASE_URL'].rstrip('/')}/v1",
        max_retries=0,
        generation_kwargs={"temperature": 0},
    ),
)

pipeline.connect("retriever.documents", "prompt_builder.documents")
pipeline.connect("prompt_builder.prompt", "llm.messages")

question = "What should we do when checkout returns HTTP 503?"
result = pipeline.run(
    {
        "retriever": {"query": question},
        "prompt_builder": {"question": question},
    }
)

print(result["llm"]["replies"][0].text)