GraphQL for AI Agent Systems: Flexible Queries for Agent Data and Conversations

Why GraphQL Fits AI Agent Systems

AI agent systems are data-hungry. A single dashboard page might need the conversation history, agent configuration, tool call results, token usage stats, and active session count — all at once. With REST, that means five separate HTTP requests. With GraphQL, a single query fetches exactly what the client needs in one round trip.

GraphQL's typed schema also serves as a contract between your agent backend and any consumer — whether that is a monitoring dashboard, another agent, or a mobile app. The schema is self-documenting and introspectable, which eliminates the drift between API docs and actual behavior.

Designing the Agent Schema

Start by defining your core types using Strawberry, a modern Python GraphQL library that integrates cleanly with FastAPI:

import strawberry
from strawberry.fastapi import GraphQLRouter
from fastapi import FastAPI
from datetime import datetime
from typing import Optional

@strawberry.type
class Message:
    id: str
    role: str
    content: str
    tool_call_id: Optional[str]
    created_at: datetime

@strawberry.type
class TokenUsage:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int

@strawberry.type
class Conversation:
    id: str
    agent_id: str
    created_at: datetime
    status: str
    messages: list[Message]
    usage: TokenUsage

@strawberry.type
class Agent:
    id: str
    name: str
    model: str
    system_prompt: str
    conversations: list[Conversation]

This schema lets clients query at any depth. An agent monitoring tool can request just conversation IDs and statuses without pulling full message histories. A debugging view can request everything including tool call details.

Resolvers with Efficient Data Loading

The key to GraphQL performance is avoiding N+1 queries. Use dataloaders to batch database lookups:

from strawberry.dataloader import DataLoader

async def load_messages_batch(conversation_ids: list[str]) -> list[list[Message]]:
    # Single database query for all conversations
    rows = await db.fetch_all(
        "SELECT * FROM messages WHERE conversation_id = ANY($1) ORDER BY created_at",
        [conversation_ids],
    )
    grouped: dict[str, list[Message]] = {cid: [] for cid in conversation_ids}
    for row in rows:
        grouped[row["conversation_id"]].append(
            Message(
                id=row["id"],
                role=row["role"],
                content=row["content"],
                tool_call_id=row["tool_call_id"],
                created_at=row["created_at"],
            )
        )
    return [grouped[cid] for cid in conversation_ids]

message_loader = DataLoader(load_fn=load_messages_batch)

With this dataloader, querying 50 conversations with their messages results in exactly two database queries — one for conversations and one for all their messages — instead of 51.

Subscriptions for Real-Time Agent Streaming

GraphQL subscriptions are a natural fit for streaming agent responses token by token. Clients subscribe once and receive each chunk as it arrives:

import asyncio
from typing import AsyncGenerator

@strawberry.type
class StreamChunk:
    conversation_id: str
    content: str
    is_final: bool
    token_index: int

@strawberry.type
class Subscription:
    @strawberry.subscription
    async def agent_stream(
        self, conversation_id: str
    ) -> AsyncGenerator[StreamChunk, None]:
        queue = agent_streams.get(conversation_id)
        if queue is None:
            return
        index = 0
        while True:
            chunk = await queue.get()
            yield StreamChunk(
                conversation_id=conversation_id,
                content=chunk["text"],
                is_final=chunk.get("done", False),
                token_index=index,
            )
            index += 1
            if chunk.get("done"):
                break

The client opens a WebSocket connection, subscribes with a conversation ID, and receives each token as a StreamChunk. The is_final flag signals when the full response is complete.

Cursor-Based Pagination for Message History

Agent conversations can grow to thousands of messages. Use cursor-based pagination to let clients page through history efficiently:

@strawberry.type
class MessageEdge:
    cursor: str
    node: Message

@strawberry.type
class MessageConnection:
    edges: list[MessageEdge]
    has_next_page: bool
    end_cursor: Optional[str]

@strawberry.type
class Query:
    @strawberry.field
    async def conversation_messages(
        self,
        conversation_id: str,
        first: int = 20,
        after: Optional[str] = None,
    ) -> MessageConnection:
        query = "SELECT * FROM messages WHERE conversation_id = $1"
        params = [conversation_id]
        if after:
            query += " AND created_at > $2"
            params.append(decode_cursor(after))
        query += " ORDER BY created_at LIMIT $" + str(len(params) + 1)
        params.append(first + 1)
        rows = await db.fetch_all(query, params)
        has_next = len(rows) > first
        edges = [
            MessageEdge(cursor=encode_cursor(r["created_at"]), node=to_message(r))
            for r in rows[:first]
        ]
        return MessageConnection(
            edges=edges,
            has_next_page=has_next,
            end_cursor=edges[-1].cursor if edges else None,
        )

Wiring It Into FastAPI

app = FastAPI()
schema = strawberry.Schema(query=Query, subscription=Subscription)
graphql_app = GraphQLRouter(schema)
app.include_router(graphql_app, prefix="/graphql")

FAQ

When should I choose GraphQL over REST for an AI agent API?

Choose GraphQL when your consumers have varied data needs — dashboards, mobile apps, and other agents all querying the same backend but needing different fields. Stick with REST when your API is simple, has few consumers, or when you need strict HTTP caching.

How do I prevent expensive GraphQL queries from overwhelming my agent backend?

Implement query depth limiting and query cost analysis. Strawberry supports extensions that reject queries exceeding a maximum depth or estimated cost. You can also use persisted queries in production so that only pre-approved query shapes are allowed.

Can GraphQL subscriptions replace Server-Sent Events for agent streaming?

Yes, for most use cases. GraphQL subscriptions run over WebSocket and provide typed, schema-validated streaming. SSE is simpler if you just need a single text stream. Subscriptions are better when you need structured chunks with metadata like token counts and tool call signals.

---

#GraphQL #AIAgents #APIDesign #Strawberry #RealTime #AgenticAI #LearnAI #AIEngineering