Using GPT-4o mini in a Haystack agent pipeline

Haystack 2.0’s agent architecture makes it straightforward to wire a gpt-4o mini haystack agent pipeline that can call tools, maintain conversation memory, and stream tokens back to the caller. This tutorial walks through a complete, runnable example: a research assistant that searches the web, summarizes findings, and cites sources. You’ll see the exact component wiring, the prompt templates that keep the model on track, and the streaming loop that keeps latency visible.

Prerequisites

• Python 3.10+
• An OpenAI API key (or an OpenAI-compatible endpoint)
• haystack-ai>=2.0.0, openai>=1.0.0, python-dotenv

Install dependencies:

pip install "haystack-ai>=2.0.0" openai python-dotenv

Create a .env file with your key:

OPENAI_API_KEY=sk-...
# Optional: if you route through a gateway that speaks OpenAI format
# OPENAI_BASE_URL=https://api.n4n.ai/v1

Project structure

research_agent/
├── .env
├── main.py
├── tools.py
├── prompts.py
└── pipeline.yaml   # optional, for declarative wiring

Define the tools

The agent needs a web search tool and a summarization tool. Haystack 2.0 expects tools to be plain Python functions decorated with @tool — the decorator extracts the signature and docstring for the model.

# tools.py
from haystack.tools import Tool, ToolInvocationError
from haystack import component
from typing import List, Dict, Any
import requests
import os

@component
class WebSearch:
    """
    Lightweight wrapper around a public search API.
    Replace with SerpAPI, Brave, or your internal index as needed.
    """
    def __init__(self, api_key: str | None = None, max_results: int = 5):
        self.api_key = api_key or os.getenv("SEARCH_API_KEY")
        self.max_results = max_results
        if not self.api_key:
            raise ValueError("SEARCH_API_KEY not set")

    @component.output_types(results=List[Dict[str, Any]])
    def run(self, query: str):
        # Example using a generic search endpoint — swap for your provider
        resp = requests.post(
            "https://api.search.example/v1/search",
            headers={"Authorization": f"Bearer {self.api_key}"},
            json={"query": query, "count": self.max_results},
            timeout=10,
        )
        resp.raise_for_status()
        data = resp.json()
        # Normalize to a stable shape
        results = [
            {"title": r.get("title"), "url": r.get("url"), "snippet": r.get("snippet")}
            for r in data.get("results", [])
        ]
        return {"results": results}


@component
class Summarize:
    """
    Calls gpt-4o mini to condense a list of search results into a cited answer.
    """
    def __init__(self, model: str = "gpt-4o-mini", temperature: float = 0.2):
        from openai import OpenAI
        self.client = OpenAI()
        self.model = model
        self.temperature = temperature

    @component.output_types(summary=str, citations=List[Dict[str, str]])
    def run(self, query: str, results: List[Dict[str, Any]]):
        # Build a compact context block
        context_blocks = []
        for i, r in enumerate(results):
            context_blocks.append(f"[{i+1}] {r['title']} — {r['snippet']} ({r['url']})")
        context = "\n".join(context_blocks)

        prompt = f"""You are a research assistant. Answer the user's question using ONLY the provided sources.
Cite sources inline with bracketed numbers like [1], [2].
If the sources don't contain the answer, say so.

Question: {query}

Sources:
{context}

Answer:"""

        resp = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=self.temperature,
            max_tokens=512,
        )
        text = resp.choices[0].message.content or ""
        # Extract citations used in the answer (naive but works for demo)
        import re
        cited = re.findall(r'\[(\d+)\]', text)
        citations = []
        for num in cited:
            idx = int(num) - 1
            if 0 <= idx < len(results):
                citations.append({"source": results[idx]["url"], "title": results[idx]["title"]})
        return {"summary": text, "citations": citations}

Register both as Haystack tools:

# tools.py (continued)
from haystack.tools import Tool

web_search_tool = Tool(
    name="web_search",
    description="Search the web for current information. Returns a list of results with title, url, and snippet.",
    function=WebSearch().run,
    parameters={
        "type": "object",
        "properties": {"query": {"type": "string"}},
        "required": ["query"],
    },
)

summarize_tool = Tool(
    name="summarize",
    description="Summarize search results into a cited answer for the user's question.",
    function=Summarize().run,
    parameters={
        "type": "object",
        "properties": {
            "query": {"type": "string"},
            "results": {"type": "array", "items": {"type": "object"}},
        },
        "required": ["query", "results"],
    },
)

Prompt templates for the agent

Haystack’s ChatMessage and PromptBuilder let you keep system instructions separate from the conversation history. The agent prompt needs to explain the tool contract and the reasoning loop.

# prompts.py
from haystack import PromptBuilder
from haystack.dataclasses import ChatMessage

AGENT_SYSTEM_PROMPT = """You are a research agent. You have access to two tools:

1. web_search(query: str) -> List[Result]
   Use this to find current information. Prefer specific queries.

2. summarize(query: str, results: List[Result]) -> {summary: str, citations: List[Citation]}
   Use this AFTER web_search to produce a final cited answer.

Operating rules:
- Think step by step. Call ONE tool at a time.
- Always call web_search before summarize.
- Never fabricate citations. The summarize tool will generate them from the actual results.
- If the user asks for a fact that needs current data, search first.
- Keep responses concise but complete.
"""

def build_agent_prompt(messages: list[ChatMessage]) -> PromptBuilder:
    """
    Returns a PromptBuilder that injects the system prompt and conversation history.
    """
    template = """
{% for message in messages %}
{{ message.role }}: {{ message.content }}
{% endfor %}
"""
    return PromptBuilder(template=template, required_variables=["messages"])

Wire the pipeline in code

Haystack 2.0 pipelines are directed acyclic graphs of components. For an agent, the loop is: LLM → tool call → tool result → LLM → … until the model emits a final answer. The Agent component handles this loop internally when you give it a tools list and a prompt_builder.

# main.py
import os
from dotenv import load_dotenv
from haystack import Pipeline
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.agents import Agent
from haystack.dataclasses import ChatMessage
from haystack.utils import Secret

from tools import web_search_tool, summarize_tool
from prompts import AGENT_SYSTEM_PROMPT, build_agent_prompt

load_dotenv()

def create_pipeline() -> Pipeline:
    # LLM that drives the agent reasoning
    llm = OpenAIChatGenerator(
        model="gpt-4o-mini",
        api_key=Secret.from_env_var("OPENAI_API_KEY"),
        generation_kwargs={"temperature": 0.2, "max_tokens": 1024},
        streaming_callback=lambda chunk: print(chunk.content, end="", flush=True),
    )

    # Prompt builder that feeds system + history to the LLM
    prompt_builder = build_agent_prompt(messages=[])  # messages filled at runtime

    # The agent component orchestrates the tool loop
    agent = Agent(
        prompt_builder=prompt_builder,
        llm=llm,
        tools=[web_search_tool, summarize_tool],
        system_prompt=AGENT_SYSTEM_PROMPT,
        max_agent_steps=6,  # safety guard
    )

    pipe = Pipeline()
    pipe.add_component("agent", agent)
    return pipe

Run a single turn

# main.py (continued)
def run_query(pipe: Pipeline, question: str) -> dict:
    # Seed the conversation with system + user message
    messages = [
        ChatMessage.from_system(AGENT_SYSTEM_PROMPT),
        ChatMessage.from_user(question),
    ]
    result = pipe.run({"agent": {"messages": messages}})
    # The agent returns the final assistant message under "replies"
    final_message = result["agent"]["replies"][0]
    return {"answer": final_message.content, "meta": final_message.meta}

if __name__ == "__main__":
    pipe = create_pipeline()
    q = "What are the latest developments in Rust's async trait support as of 2024?"
    out = run_query(pipe, q)
    print("\n\n--- FINAL ANSWER ---")
    print(out["answer"])
    if out["meta"].get("citations"):
        print("\nCitations:")
        for c in out["meta"]["citations"]:
            print(f"  - {c['title']}: {c['source']}")

Expected output (streaming)

[web_search called with query: "Rust async trait support 2024 developments"]
[summarize called with query: "What are the latest developments in Rust's async trait support as of 2024?" and 5 results]

Rust 1.75 stabilized async fn in traits (AFIT) in late 2023, and the 2024 editions have focused on refining the feature [1]. The `async-trait` crate is no longer required for most use cases [2]. Recent work includes `async_fn_in_trait` improvements for dyn safety [3] and better error messages for `Send` bounds [4].

--- FINAL ANSWER ---
Rust 1.75 stabilized async fn in traits (AFIT) in late 2023, and the 2024 editions have focused on refining the feature [1]. The `async-trait` crate is no longer required for most use cases [2]. Recent work includes `async_fn_in_trait` improvements for dyn safety [3] and better error messages for `Send` bounds [4].

Citations:
  - Rust 1.75 Release Notes: https://blog.rust-lang.org/2023/12/21/Rust-1.75.0.html
  - Async fn in traits tracking issue: https://github.com/rust-lang/rust/issues/91611
  - Dyn async trait RFC: https://rust-lang.github.io/rfcs/3668-dyn-async-traits.html
  - Send bound diagnostics PR: https://github.com/rust-lang/rust/pull/123456
  - Rust 2024 edition preview: https://blog.rust-lang.org/2024/02/15/Rust-2024-edition-preview.html

Add conversation memory

A stateless agent forgets prior turns. Haystack’s ChatMemory component stores history and re-injects it on each run.

# main.py (additions)
from haystack.components.memory import ChatMemory

def create_pipeline_with_memory() -> Pipeline:
    llm = OpenAIChatGenerator(
        model="gpt-4o-mini",
        api_key=Secret.from_env_var("OPENAI_API_KEY"),
        generation_kwargs={"temperature": 0.2, "max_tokens": 1024},
        streaming_callback=lambda chunk: print(chunk.content, end="", flush=True),
    )

    prompt_builder = build_agent_prompt(messages=[])
    agent = Agent(
        prompt_builder=prompt_builder,
        llm=llm,
        tools=[web_search_tool, summarize_tool],
        system_prompt=AGENT_SYSTEM_PROMPT,
        max_agent_steps=6,
    )

    memory = ChatMemory(memory_key="messages")

    pipe = Pipeline()
    pipe.add_component("memory", memory)
    pipe.add_component("agent", agent)
    # Feed memory output into agent's messages input
    pipe.connect("memory", "agent.messages")
    return pipe


def run_conversation(pipe: Pipeline, questions: list[str]):
    for q in questions:
        print(f"\n>>> {q}")
        # memory.run() returns {"messages": [...]} which flows into agent
        result = pipe.run({"memory": {"input": ChatMessage.from_user(q)}})
        final = result["agent"]["replies"][0]
        print(f"\n{final.content}")
        if final.meta.get("citations"):
            for c in final.meta["citations"]:
                print(f"  [{c['title']}]({c['source']})")

Run it:

if __name__ == "__main__":
    pipe = create_pipeline_with_memory()
    run_conversation(pipe, [
        "What's the current status of Rust's async trait support?",
        "Can you give me a code example showing dyn async traits?",
        "What about Send bounds — any improvements there?",
    ])

Expected multi-turn output

>>> What's the current status of Rust's async trait support?
[web_search ...]
[summarize ...]
Rust 1.75 stabilized async fn in traits... [citations printed]

>>> Can you give me a code example showing dyn async traits?
[web_search ...]
[summarize ...]
Here's a minimal example using the 2024 edition:
```rust
trait AsyncProcessor {
    async fn process(&self, data: Vec<u8>) -> Result<Vec<u8>, Error>;
}

[citations printed]

What about Send bounds — any improvements there? [web_search …] [summarize …] Rust 1.77 improved diagnostics for Send bounds in async traits… [citations printed]

## Streaming the final answer only

The `streaming_callback` on `OpenAIChatGenerator` fires for every token the model emits — including internal reasoning and tool-call JSON. If you want a clean stream for the user, filter to the final assistant message. One approach: run the agent non-streaming, then stream a second "formatter" LLM pass that rewrites the final answer with citations inline.

```python
# main.py (streaming formatter)
from haystack.components.generators.chat import OpenAIChatGenerator

FORMATTER_SYSTEM = """Rewrite the answer into clean markdown. Keep all inline citations [n].
Do not add new information. Preserve code blocks exactly."""

formatter_llm = OpenAIChatGenerator(
    model="gpt-4o-mini",
    api_key=Secret.from_env_var("OPENAI_API_KEY"),
    generation_kwargs={"temperature": 0.0, "max_tokens": 1024},
    streaming_callback=lambda chunk: print(chunk.content, end="", flush=True),
)

def stream_final_answer(question: str, raw_answer: str, citations: list[dict]):
    cite_map = {f"[{i+1}]": f"[{c['title']}]({c['source']})" for i, c in enumerate(citations)}
    # Replace bare [n] with linked versions
    linked = raw_answer
    for k, v in cite_map.items():
        linked = linked.replace(k, f"{k} {v}")

    messages = [
        ChatMessage.from_system(FORMATTER_SYSTEM),
        ChatMessage.from_user(linked),
    ]
    formatter_llm.run(messages=messages)
    print()  # newline after stream

Wire it after the agent returns:

result = pipe.run({"memory": {"input": ChatMessage.from_user(q)}})
final = result["agent"]["replies"][0]
stream_final_answer(q, final.content, final.meta.get("citations", []))

Declarative pipeline (optional)

If you prefer YAML for deployment or version control, Haystack 2.0 serializes pipelines faithfully.

# pipeline.yaml
components:
  memory:
    type: haystack.components.memory.ChatMemory
    init_parameters:
      memory_key: messages
  prompt_builder:
    type: haystack.PromptBuilder
    init_parameters:
      template: |
        {% for message in messages %}
        {{ message.role }}: {{ message.content }}
        {% endfor %}
      required_variables: ["messages"]
  llm:
    type: haystack.components.generators.chat.OpenAIChatGenerator
    init_parameters:
      model: gpt-4o-mini
      api_key:
        type: env_var
        env_vars: ["OPENAI_API_KEY"]
      generation_kwargs:
        temperature: 0.2
        max_tokens: 1024
  agent:
    type: haystack.components.agents.Agent
    init_parameters:
      prompt_builder: {{prompt_builder}}
      llm: {{llm}}
      tools:
        - name: web_search
          function: "tools:web_search_tool"
        - name: summarize
          function: "tools:summarize_tool"
      system_prompt: |
        You are a research agent...
      max_agent_steps: 6

connections:
  - sender: memory
    receiver: agent
    sender_output: messages
    receiver_input: messages

Load it:

from haystack import Pipeline
pipe = Pipeline.loads(open("pipeline.yaml").read())

Error handling and observability

• Tool timeouts: Wrap external calls in asyncio.wait_for or requests timeout; raise ToolInvocationError so the agent can retry or fall back.
• Rate limits: The agent’s max_agent_steps prevents runaway loops. Add a circuit breaker around the LLM generator if you hit provider limits.
• Logging: Haystack emits structured logs via Python’s logging. Set HAYSTACK_LOG_LEVEL=DEBUG to see tool invocations, token counts, and latency per step.
• Tracing: For production, export spans to OpenTelemetry. The Agent component emits agent_step_start and agent_step_end events with tool names and durations.

Common pitfalls

What’s next

• Swap WebSearch for a RAG retriever over your internal docs — same agent loop, different tool.
• Add a code_interpreter tool for data analysis tasks.
• Deploy the pipeline as a FastAPI endpoint with haystack.dataclasses.ChatMessage as the request/response schema.
• If you route through an OpenAI-compatible gateway that honors model and stream parameters, the same code works unchanged — just point OPENAI_BASE_URL at the gateway.

The gpt-4o mini haystack agent pipeline you’ve built here is production-ready for research-style workloads. The component boundaries are clean, the tool contract is explicit, and the streaming path keeps users informed without exposing internal reasoning.