Introduction

Building production-grade agentic systems requires more than just connecting an LLM to data sources. It demands a robust architecture that can scale, observe, and evaluate agent behavior in real-time. This article explores how to combine three powerful technologies—LlamaIndex, Apache Spark, and MLflow—to create a sophisticated agentic system capable of processing large-scale data while maintaining complete observability and evaluation capabilities.

We'll build a system that leverages Spark's distributed computing for data processing, LlamaIndex for intelligent agent orchestration, and MLflow for comprehensive experiment tracking and evaluation metrics.

Architecture Overview

The architecture consists of four main layers:

1. Data Layer: Apache Spark for distributed data processing and feature engineering
2. Indexing Layer: LlamaIndex for creating embeddings and vector indexes
3. Agent Layer: ReAct agents with tool definitions for reasoning and decision-making
4. Observability Layer: MLflow for tracking experiments, metrics, and agent evaluations

This layered approach allows each component to scale independently while maintaining clear boundaries and communication protocols.

Data Processing with Spark

The foundation of our system is Apache Spark, which handles large-scale data processing and feature engineering. Instead of loading all data into memory, we process data in batches and create distributed indexes that LlamaIndex can query efficiently.

from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler
import mlflow

# Initialize Spark Session with MLflow tracking
spark = SparkSession.builder \
    .appName("LlamaIndex-Spark-Agent") \
    .config("spark.sql.adaptive.enabled", "true") \
    .getOrCreate()

# Log Spark configuration to MLflow
mlflow.log_param("spark_version", spark.version)
mlflow.log_param("spark_cores", spark.sparkContext.defaultParallelism)

# Load and process data
df = spark.read.parquet("s3://data-bucket/raw-data")

# Feature engineering
assembler = VectorAssembler(
    inputCols=["feature1", "feature2", "feature3"],
    outputCol="features"
)

processed_df = assembler.transform(df)
processed_df.write.mode("overwrite").parquet("s3://data-bucket/processed-data")

Integrating LlamaIndex with Spark Data

LlamaIndex provides flexibility in how you structure your data sources. Rather than loading all Spark data into memory, we create a custom data loader that fetches chunks on-demand from the Spark cluster.

from llama_index.core import Document
from llama_index.core.data_structs import IndexDict
from llama_index.vector_stores import MilvusVectorStore
from llama_index.core import VectorStoreIndex
from pyspark.sql import DataFrame

class SparkDataLoader:
    def __init__(self, spark_df: DataFrame, chunk_size: int = 100):
        self.spark_df = spark_df
        self.chunk_size = chunk_size
    
    def load_documents(self) -> list[Document]:
        """Convert Spark DataFrame to LlamaIndex Documents"""
        documents = []
        
        # Convert to Pandas for processing
        pdf = self.spark_df.toPandas()
        
        for idx, row in pdf.iterrows():
            doc = Document(
                text=f"Data Record {idx}: {row.to_dict()}",
                metadata={
                    "row_id": int(idx),
                    "source": "spark",
                    "timestamp": str(row.get('timestamp', '')),
                    **{k: str(v) for k, v in row.items() if k != 'text'}
                }
            )
            documents.append(doc)
        
        return documents

# Set up Milvus vector store for efficient similarity search
vector_store = MilvusVectorStore(
    url="http://milvus:19530",
    collection_name="llamaindex_spark_data",
    dim=384  # Dimension for embeddings
)

# Create index from Spark data
spark_df = spark.read.parquet("s3://data-bucket/processed-data")
loader = SparkDataLoader(spark_df, chunk_size=100)
documents = loader.load_documents()

index = VectorStoreIndex.from_documents(
    documents, 
    vector_store=vector_store
)

Building Agents with MLflow Observability

The key to production-grade agentic systems is comprehensive observability. MLflow allows us to track every agent interaction, decision, and evaluation metric.

from llama_index.core.agent import ReActAgent
from llama_index.llm.openai import OpenAI
from llama_index.core.tools import QueryEngineTool, FunctionTool
import mlflow
from functools import wraps
import json

# MLflow wrapper for agent interactions
def log_agent_interaction(agent_name: str):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            with mlflow.start_run(nested=True):
                mlflow.set_tag("agent", agent_name)
                mlflow.set_tag("interaction_type", func.__name__)
                
                # Log input parameters
                mlflow.log_params({
                    k: str(v) for k, v in kwargs.items()
                })
                
                # Execute function
                result = func(*args, **kwargs)
                
                # Log output
                mlflow.log_text(json.dumps(result), "output.json")
                
                return result
        return wrapper
    return decorator

class ObservableAgent:
    def __init__(self, llm_model: str = "gpt-4"):
        self.llm = OpenAI(model=llm_model)
        self.query_engine = None
        
    @log_agent_interaction("query_executor")
    def execute_query(self, user_query: str) -> str:
        """Execute a query with full observability"""
        response = self.query_engine.query(user_query)
        return str(response)
    
    def create_agent(self, tools: list):
        """Create a ReAct agent with MLflow hooks"""
        with mlflow.start_run():
            mlflow.log_param("num_tools", len(tools))
            mlflow.log_text(
                json.dumps([{"name": t.metadata.name} for t in tools]), 
                "tools.json"
            )
            
            agent = ReActAgent.from_tools(
                tools, 
                llm=self.llm,
                verbose=True
            )
            
            return agent

# Usage
observable_agent = ObservableAgent(llm_model="gpt-4")
observable_agent.query_engine = index.as_query_engine()

tools = [
    QueryEngineTool(
        query_engine=observable_agent.query_engine,
        metadata=type('obj', (object,), {
            'name': 'spark_data_search',
            'description': 'Search through processed Spark data'
        })()
    )
]

agent = observable_agent.create_agent(tools)

Real-Time Evaluation with MLflow

Beyond observability, we need to evaluate agent performance in real-time. MLflow provides a framework for logging custom metrics and comparing runs across multiple experiments.

import mlflow
from typing import Any
from datetime import datetime
import numpy as np

class AgentEvaluator:
    def __init__(self, experiment_name: str):
        mlflow.set_experiment(experiment_name)
        self.start_time = datetime.now()
    
    def evaluate_response(self, 
                         user_query: str, 
                         agent_response: str,
                         ground_truth: str = None) -> dict:
        """Evaluate agent response quality"""
        
        metrics = {
            "response_length": len(agent_response.split()),
            "query_length": len(user_query.split()),
            "latency_ms": 0  # Calculate actual latency
        }
        
        # Semantic similarity with ground truth
        if ground_truth:
            from sklearn.feature_extraction.text import TfidfVectorizer
            from sklearn.metrics.pairwise import cosine_similarity
            
            vectorizer = TfidfVectorizer()
            tfidf_matrix = vectorizer.fit_transform([agent_response, ground_truth])
            similarity = cosine_similarity(tfidf_matrix[0], tfidf_matrix[1])[0][0]
            metrics["semantic_similarity"] = float(similarity)
        
        # Log metrics to MLflow
        with mlflow.start_run(nested=True):
            for metric_name, value in metrics.items():
                mlflow.log_metric(metric_name, value)
            
            mlflow.log_param("query", user_query)
            mlflow.log_text(agent_response, "response.txt")
        
        return metrics
    
    def compare_experiments(self, experiment_ids: list[int]):
        """Compare metrics across multiple experiment runs"""
        client = mlflow.tracking.MlflowClient()
        
        results = []
        for exp_id in experiment_ids:
            runs = client.search_runs(experiment_ids=[exp_id])
            for run in runs:
                results.append({
                    "run_id": run.info.run_id,
                    "metrics": run.data.metrics,
                    "params": run.data.params
                })
        
        return results

# Usage
evaluator = AgentEvaluator("llamaindex-spark-mlflow-agents")

# Simulate agent interactions
for i in range(10):
    query = f"Analyze trend for period {i}"
    response = agent.chat(query)
    
    metrics = evaluator.evaluate_response(
        user_query=query,
        agent_response=str(response),
        ground_truth=f"Expected analysis for period {i}"
    )

Best Practices

• Batch Processing: Use Spark's distributed capabilities to process data in batches rather than streaming to avoid bottlenecks
• Vector Store Selection: Choose vector stores (Milvus, Pinecone) that scale with your data size and query volume
• Experiment Tracking: Track every agent configuration and evaluation run in MLflow for reproducibility
• Custom Metrics: Define domain-specific metrics that matter for your use case and track them consistently
• Error Handling: Implement comprehensive error handling with MLflow logging to track agent failures
• Cost Optimization: Monitor token usage and API costs within MLflow to optimize LLM spending
• Data Privacy: Implement proper data masking when logging sensitive information to MLflow

Production Deployment Considerations

When deploying this system to production, consider:

• Scaling: Use Spark on Kubernetes to automatically scale compute based on data volume
• Latency: Cache common queries and embeddings to reduce response times
• Monitoring: Set up alerts in MLflow for performance degradation or high error rates
• Version Control: Track model versions, prompts, and tool definitions in MLflow
• A/B Testing: Use MLflow to run parallel agent configurations and compare effectiveness

Conclusion

Combining LlamaIndex, Spark, and MLflow creates a powerful foundation for production agentic systems. This architecture provides:

• Scalability: Process large datasets with Spark's distributed computing
• Intelligence: Leverage sophisticated agent reasoning with LlamaIndex
• Observability: Track every aspect of agent behavior with MLflow
• Reliability: Evaluate and improve agent performance through systematic experimentation

By adopting this approach, you can build agentic systems that are not just intelligent but also transparent, measurable, and continuously improving.