Create tracing haystack

conftest.py

@@ -182,6 +182,7 @@ def pytest_ignore_collect(collection_path, config):
             "tests/gemini",
             "tests/groq",
             "tests/h2o",
+            "tests/haystack",
             "tests/johnsnowlabs",
             "tests/keras",
             "tests/keras_core",

docs/api_reference/source/python_api/mlflow.haystack.rst

@@ -0,0 +1,7 @@
+mlflow.haystack
+===============
+
+.. automodule:: mlflow.haystack
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/docs/genai/tracing/integrations/listing/haystack.mdx

@@ -0,0 +1,48 @@
+---
+sidebarTitle: Haystack
+title: Haystack
+---
+
+import { APILink } from "@site/src/components/APILink";
+
+# Tracing Haystack
+
+![Haystack tracing via autolog](/images/llms/haystack/haystack-tracing.png)
+
+MLflow Tracing provides automatic tracing capability when using Haystack pipelines and components.
+When Haystack auto-tracing is enabled by calling the <APILink fn="mlflow.haystack.autolog" /> function,
+usage of Haystack pipelines and components will automatically record generated traces during interactive development.
+
+## Example Usage
+
+```python
+import mlflow
+from haystack import Pipeline
+from haystack.components.generators import OpenAIGenerator
+from haystack.components.builders.prompt_builder import PromptBuilder
+
+# Turn on auto tracing for Haystack by calling mlflow.haystack.autolog()
+mlflow.haystack.autolog()
+
+# Initialize the pipeline
+pipeline = Pipeline()
+
+# Configure the LLM component
+llm = OpenAIGenerator(model="gpt-4o-mini")
+prompt_template = """Answer the question. {{question}}"""
+prompt_builder = PromptBuilder(template=prompt_template)
+
+# Build the pipeline
+pipeline = Pipeline()
+pipeline.add_component("prompt_builder", prompt_builder)
+pipeline.add_component("llm", llm)
+pipeline.connect("prompt_builder", "llm")
+
+# Run the pipeline
+question = "Who lives in Paris?"
+results = pipeline.run({"prompt_builder": {"question": question}})
+```
+
+## Disable auto-tracing
+
+Auto tracing for Haystack can be disabled globally by calling `mlflow.haystack.autolog(disable=True)` or `mlflow.autolog(disable=True)`.

docs/src/components/TracingIntegrations/index.tsx

@@ -149,6 +149,12 @@ const TRACING_INTEGRATIONS: TracingIntegration[] = [
     logoPath: '/images/logos/txtai-logo.png',
     link: '/genai/tracing/integrations/listing/txtai',
   },
+  {
+    id: 'haystack',
+    name: 'Haystack',
+    logoPath: '/images/logos/haystack-logo.png',
+    link: '/genai/tracing/integrations/listing/haystack',
+  },
 ];
 
 interface TracingIntegrationsProps {

examples/haystack/agent_tracing.py

@@ -0,0 +1,161 @@
+"""
+Example demonstrating MLflow tracing for Haystack Agents.
+
+This example shows how MLflow captures the full execution flow of an Agent,
+including its internal calls to chat generators and tool invokers.
+"""
+
+import asyncio
+import datetime
+
+from haystack.components.agents import Agent
+from haystack.components.generators import OpenAIGenerator
+from haystack.dataclasses import ChatMessage
+from haystack.tools import Tool
+
+import mlflow
+
+# Turn on auto tracing for Haystack
+mlflow.haystack.autolog()
+
+
+def add_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+    return a + b
+
+
+def multiply_numbers(a: int, b: int) -> int:
+    """Multiply two numbers together."""
+    return a * b
+
+
+# Example 1: Simple Agent without tools (behaves like ChatGenerator)
+def simple_agent_example():
+    print("=== Simple Agent Example (No Tools) ===")
+
+    # Create a chat generator
+    llm = OpenAIGenerator(model="gpt-4o-mini")
+
+    # Create an agent without tools
+    agent = Agent(chat_generator=llm, system_prompt="You are a helpful assistant.")
+
+    # Warm up the agent
+    agent.warm_up()
+
+    # Run the agent
+    messages = [ChatMessage.from_user("What is the capital of France?")]
+    result = agent.run(messages=messages)
+
+    print("User:", messages[0].content)
+    print("Agent:", result["last_message"].content)
+    print()
+
+
+# Example 2: Agent with tools
+def agent_with_tools_example():
+    print("=== Agent with Tools Example ===")
+
+    # Create tools
+    add_tool = Tool(
+        name="add_numbers",
+        description="Add two numbers together",
+        function=add_numbers,
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "integer", "description": "First number"},
+                "b": {"type": "integer", "description": "Second number"},
+            },
+            "required": ["a", "b"],
+        },
+    )
+
+    multiply_tool = Tool(
+        name="multiply_numbers",
+        description="Multiply two numbers together",
+        function=multiply_numbers,
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "integer", "description": "First number"},
+                "b": {"type": "integer", "description": "Second number"},
+            },
+            "required": ["a", "b"],
+        },
+    )
+
+    # Create a chat generator
+    llm = OpenAIGenerator(model="gpt-4o-mini")
+
+    # Create an agent with tools
+    agent = Agent(
+        chat_generator=llm,
+        tools=[add_tool, multiply_tool],
+        system_prompt="You are a helpful math assistant. Use the tools provided to help with calculations.",
+    )
+
+    # Warm up the agent
+    agent.warm_up()
+
+    # Run the agent with a calculation request
+    messages = [ChatMessage.from_user("What is 15 + 27, and what is the result multiplied by 3?")]
+    result = agent.run(messages=messages)
+
+    print("User:", messages[0].content)
+    print("Agent:", result["last_message"].content)
+    print("\nNumber of messages exchanged:", len(result["messages"]))
+    print()
+
+
+# Example 3: Async Agent
+async def async_agent_example():
+    print("=== Async Agent Example ===")
+
+    # Create a simple tool
+    def get_time() -> str:
+        """Get the current time."""
+        return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+    time_tool = Tool(
+        name="get_time",
+        description="Get the current time",
+        function=get_time,
+        parameters={"type": "object", "properties": {}},
+    )
+
+    # Create components
+    llm = OpenAIGenerator(model="gpt-4o-mini")
+
+    agent = Agent(
+        chat_generator=llm,
+        tools=[time_tool],
+        system_prompt="You are a helpful assistant that can tell the time.",
+    )
+
+    # Warm up the agent
+    agent.warm_up()
+
+    # Run the agent asynchronously
+    messages = [ChatMessage.from_user("What time is it now?")]
+    result = await agent.run_async(messages=messages)
+
+    print("User:", messages[0].content)
+    print("Agent:", result["last_message"].content)
+    print()
+
+
+# Run examples
+if __name__ == "__main__":
+    # Run synchronous examples
+    simple_agent_example()
+    agent_with_tools_example()
+
+    # Run asynchronous example
+    asyncio.run(async_agent_example())
+
+    print("=== Tracing Complete ===")
+    print("The Agent execution traces have been logged to MLflow.")
+    print("You can view the hierarchical trace structure showing:")
+    print("- Agent.run as the parent span")
+    print("- chat_generator and tool_invoker as child spans")
+    print("- All tool invocations and their results")

examples/haystack/tracing.py

@@ -0,0 +1,65 @@
+"""
+This is an example for leveraging MLflow's auto tracing capabilities for Haystack.
+
+For more information about MLflow Tracing, see: https://mlflow.org/docs/latest/llms/tracing/index.html
+"""
+
+from haystack import Pipeline
+from haystack.components.builders.prompt_builder import PromptBuilder
+from haystack.components.generators import OpenAIGenerator
+from haystack.core.pipeline.async_pipeline import AsyncPipeline
+
+import mlflow
+
+# Turn on auto tracing for Haystack by calling mlflow.haystack.autolog()
+# This will automatically trace:
+# - Pipeline executions (both sync and async)
+# - Individual component executions
+# - Token usage for LLM components
+# - Component metadata and parameters
+mlflow.haystack.autolog()
+
+
+# Example 1: Synchronous Pipeline
+def sync_pipeline_example():
+    print("=== Synchronous Pipeline Example ===")
+
+    pipeline = Pipeline()
+    llm = OpenAIGenerator(model="gpt-4o-mini")
+
+    prompt_template = """Answer the question. {{question}}"""
+    prompt_builder = PromptBuilder(template=prompt_template)
+
+    pipeline = Pipeline()
+    pipeline.add_component("prompt_builder", prompt_builder)
+    pipeline.add_component("llm", llm)
+    pipeline.connect("prompt_builder", "llm")
+
+    question = "Who lives in Paris?"
+    results = pipeline.run({"prompt_builder": {"question": question}})
+
+    print("Question:", question)
+    print("Answer:", results["llm"]["replies"][0])
+    print()
+
+
+# Example 2: Asynchronous Pipeline
+async def async_pipeline_example():
+    print("=== Asynchronous Pipeline Example ===")
+
+    pipeline = AsyncPipeline()
+
+    llm = OpenAIGenerator(model="gpt-4o-mini")
+    prompt_template = """Tell me about: {{topic}}"""
+    prompt_builder = PromptBuilder(template=prompt_template)
+
+    pipeline.add_component("prompt_builder", prompt_builder)
+    pipeline.add_component("llm", llm)
+    pipeline.connect("prompt_builder", "llm")
+
+    topic = "artificial intelligence"
+    results = await pipeline.run_async({"prompt_builder": {"topic": topic}})
+
+    print("Topic:", topic)
+    print("Response:", results["llm"]["replies"][0])
+    print()

mlflow/__init__.py

@@ -71,6 +71,7 @@
 gemini = LazyLoader("mlflow.gemini", globals(), "mlflow.gemini")
 groq = LazyLoader("mlflow.groq", globals(), "mlflow.groq")
 h2o = LazyLoader("mlflow.h2o", globals(), "mlflow.h2o")
+haystack = LazyLoader("mlflow.haystack", globals(), "mlflow.haystack")
 johnsnowlabs = LazyLoader("mlflow.johnsnowlabs", globals(), "mlflow.johnsnowlabs")
 keras = LazyLoader("mlflow.keras", globals(), "mlflow.keras")
 langchain = LazyLoader("mlflow.langchain", globals(), "mlflow.langchain")
@@ -124,6 +125,7 @@
         gemini,
         groq,
         h2o,
+        haystack,
         johnsnowlabs,
         keras,
         langchain,