Refactoring the Report Generation Agent demo UI and evaluation script

.env.example

@@ -9,7 +9,12 @@ DEFAULT_PLANNER_MODEL="gemini-2.5-pro"
 DEFAULT_WORKER_MODEL="gemini-2.5-flash"
 DEFAULT_EVALUATOR_MODEL="gemini-2.5-pro"
 
-# LangFuse (optional, for tracing)
+# LangFuse for agent execution tracing and evaluations
 LANGFUSE_SECRET_KEY="sk-lf-..."
 LANGFUSE_PUBLIC_KEY="pk-lf-..."
 LANGFUSE_HOST="https://us.cloud.langfuse.com"
+
+# Report Generation (all optional, defaults are in implementations/report_generation/env_vars.py)
+REPORT_GENERATION_OUTPUT_PATH="..."
+REPORT_GENERATION_DB_PATH="..."
+REPORT_GENERATION_LANGFUSE_PROJECT_NAME="..."

aieng-eval-agents/aieng/agent_evals/langfuse.py

@@ -1,11 +1,13 @@
 """Functions and objects pertaining to Langfuse."""
 
 import base64
+import json
 import logging
 import os
 
 import logfire
 import nest_asyncio
+from aieng.agent_evals.async_client_manager import AsyncClientManager
 from aieng.agent_evals.configs import Configs
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
@@ -75,3 +77,48 @@ def setup_langfuse_tracer(service_name: str = "aieng-eval-agents") -> "trace.Tra
     # Set the global default tracer provider
     trace.set_tracer_provider(trace_provider)
     return trace.get_tracer(__name__)
+
+
+async def upload_dataset_to_langfuse(dataset_path: str, dataset_name: str):
+    """Upload a dataset to Langfuse.
+
+    Parameters
+    ----------
+    dataset_path : str
+        Path to the dataset to upload. The dataset must be a json file
+        containing a list of dictionaries. Each dictionary must contain a
+        `input` and `expected_output` keys. Additionally, it can include
+        an `id` key that will be added to the metadata of the dataset item.
+    dataset_name : str
+        Name of the dataset to upload.
+    """
+    # Get the client manager singleton instance and langfuse client
+    client_manager = AsyncClientManager.get_instance()
+    langfuse_client = client_manager.langfuse_client
+
+    # Load the ground truth dataset from the file path
+    logger.info(f"Loading dataset from '{dataset_path}'")
+    with open(dataset_path, "r") as file:
+        dataset = json.load(file)
+
+    # Create the dataset in Langfuse
+    langfuse_client.create_dataset(name=dataset_name)
+
+    # Upload each item to the dataset
+    for item in dataset:
+        assert "input" in item, "`input` is required for all items in the dataset"
+        assert "expected_output" in item, "`expected_output` is required for all items in the dataset"
+
+        langfuse_client.create_dataset_item(
+            dataset_name=dataset_name,
+            input=item["input"],
+            expected_output=item["expected_output"],
+            metadata={
+                "id": item.get("id", None),
+            },
+        )
+
+    logger.info(f"Uploaded {len(dataset)} items to dataset '{dataset_name}'")
+
+    # Gracefully close the services
+    await client_manager.close()

aieng-eval-agents/aieng/agent_evals/report_generation/agent.py

@@ -0,0 +1,82 @@
+"""
+Definitions for the the report generation agent.
+
+Example
+-------
+>>> from aieng.agent_evals.report_generation.agent import get_report_generation_agent
+>>> from aieng.agent_evals.report_generation.prompts import MAIN_AGENT_INSTRUCTIONS
+>>> agent = get_report_generation_agent(
+>>>     instructions=MAIN_AGENT_INSTRUCTIONS,
+>>>     sqlite_db_path=Path("data/OnlineRetail.db"),
+>>>     reports_output_path=Path("reports/"),
+>>>     langfuse_project_name="Report Generation",
+>>> )
+"""
+
+from pathlib import Path
+
+import agents
+from aieng.agent_evals.async_client_manager import AsyncClientManager
+from aieng.agent_evals.langfuse import setup_langfuse_tracer
+from aieng.agent_evals.report_generation.file_writer import ReportFileWriter
+
+
+def get_report_generation_agent(
+    instructions: str,
+    sqlite_db_path: Path,
+    reports_output_path: Path,
+    langfuse_project_name: str | None,
+) -> agents.Agent:
+    """
+    Define the report generation agent.
+
+    Parameters
+    ----------
+    instructions : str
+        The instructions for the agent.
+    sqlite_db_path : Path
+        The path to the SQLite database.
+    reports_output_path : Path
+        The path to the reports output directory.
+    langfuse_project_name : str | None
+        The name of the Langfuse project to use for tracing.
+
+    Returns
+    -------
+    agents.Agent
+        The report generation agent.
+    """
+    # Setup langfuse tracing if project name is provided
+    if langfuse_project_name:
+        setup_langfuse_tracer(langfuse_project_name)
+
+    # Get the client manager singleton instance
+    client_manager = AsyncClientManager.get_instance()
+    report_file_writer = ReportFileWriter(reports_output_path)
+
+    # Define an agent using the OpenAI Agent SDK
+    return agents.Agent(
+        name="Report Generation Agent",  # Agent name for logging and debugging purposes
+        instructions=instructions,  # System instructions for the agent
+        # Tools available to the agent
+        # We wrap the `execute_sql_query` and `write_report_to_file` methods
+        # with `function_tool`, which will construct the tool definition JSON
+        # schema by extracting the necessary information from the method
+        # signature and docstring.
+        tools=[
+            agents.function_tool(
+                client_manager.sqlite_connection(sqlite_db_path).execute,
+                name_override="execute_sql_query",
+                description_override="Execute a SQL query against the SQLite database.",
+            ),
+            agents.function_tool(
+                report_file_writer.write,
+                name_override="write_report_to_file",
+                description_override="Write the report data to a downloadable XLSX file.",
+            ),
+        ],
+        model=agents.OpenAIChatCompletionsModel(
+            model=client_manager.configs.default_worker_model,
+            openai_client=client_manager.openai_client,
+        ),
+    )