RooAgent 0.1.0__tar.gz

rooagent-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Aman Desai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

rooagent-0.1.0/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.4
+Name: RooAgent
+Version: 0.1.0
+Summary: Agent For ROOT
+Author-email: Aman Desai <amanmukeshdesai@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/amanmdesai/RooAgent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: langchain
+Requires-Dist: langchain-core
+Requires-Dist: langgraph
+Requires-Dist: langchain-openai
+Requires-Dist: uproot
+Requires-Dist: pandas
+Requires-Dist: numpy
+Dynamic: license-file
+
+# RooAgent
+Agent for ROOT
+
+
+Download and Install: 
+
+```
+git clone https://github.com/amanmdesai/RooAgent.git
+cd RooAgent
+pip install .
+```
+
+Installation: 
+
+```
+pip install rooagent
+```
+
+External Dependencies
+
+Requires a pre-installed version of ROOT that is compatible to work with Python
+
+```
+CERN ROOT: - ROOT >= 6.34 (https://root.cern/)
+```
+
+Install ollama's gpt-oss model: 
+
+```
+curl -fsSL https://ollama.com/install.sh | sh
+ollama pull gpt-oss
+```
+

rooagent-0.1.0/README.md

@@ -0,0 +1,33 @@
+# RooAgent
+Agent for ROOT
+
+
+Download and Install: 
+
+```
+git clone https://github.com/amanmdesai/RooAgent.git
+cd RooAgent
+pip install .
+```
+
+Installation: 
+
+```
+pip install rooagent
+```
+
+External Dependencies
+
+Requires a pre-installed version of ROOT that is compatible to work with Python
+
+```
+CERN ROOT: - ROOT >= 6.34 (https://root.cern/)
+```
+
+Install ollama's gpt-oss model: 
+
+```
+curl -fsSL https://ollama.com/install.sh | sh
+ollama pull gpt-oss
+```
+

rooagent-0.1.0/RooAgent.egg-info/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.4
+Name: RooAgent
+Version: 0.1.0
+Summary: Agent For ROOT
+Author-email: Aman Desai <amanmukeshdesai@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/amanmdesai/RooAgent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: langchain
+Requires-Dist: langchain-core
+Requires-Dist: langgraph
+Requires-Dist: langchain-openai
+Requires-Dist: uproot
+Requires-Dist: pandas
+Requires-Dist: numpy
+Dynamic: license-file
+
+# RooAgent
+Agent for ROOT
+
+
+Download and Install: 
+
+```
+git clone https://github.com/amanmdesai/RooAgent.git
+cd RooAgent
+pip install .
+```
+
+Installation: 
+
+```
+pip install rooagent
+```
+
+External Dependencies
+
+Requires a pre-installed version of ROOT that is compatible to work with Python
+
+```
+CERN ROOT: - ROOT >= 6.34 (https://root.cern/)
+```
+
+Install ollama's gpt-oss model: 
+
+```
+curl -fsSL https://ollama.com/install.sh | sh
+ollama pull gpt-oss
+```
+

rooagent-0.1.0/RooAgent.egg-info/SOURCES.txt

@@ -0,0 +1,19 @@
+LICENSE
+README.md
+pyproject.toml
+RooAgent.egg-info/PKG-INFO
+RooAgent.egg-info/SOURCES.txt
+RooAgent.egg-info/dependency_links.txt
+RooAgent.egg-info/entry_points.txt
+RooAgent.egg-info/requires.txt
+RooAgent.egg-info/top_level.txt
+rooagent/__init__.py
+rooagent/agent.py
+rooagent/tools/__init__.py
+rooagent/tools/converter.py
+rooagent/tools/fit_tools.py
+rooagent/tools/histogram_tools.py
+rooagent/tools/plot_tools.py
+rooagent/tools/rdataframe_tools.py
+rooagent/tools/tfile.py
+rooagent/tools/utils.py

rooagent-0.1.0/RooAgent.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

rooagent-0.1.0/RooAgent.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+rooagent = rooagent.agent:main

rooagent-0.1.0/RooAgent.egg-info/requires.txt

@@ -0,0 +1,7 @@
+langchain
+langchain-core
+langgraph
+langchain-openai
+uproot
+pandas
+numpy

rooagent-0.1.0/RooAgent.egg-info/top_level.txt

@@ -0,0 +1 @@
+rooagent

rooagent-0.1.0/pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "RooAgent"
+version = "0.1.0"
+description = "Agent For ROOT"
+authors = [
+    { name="Aman Desai", email="amanmukeshdesai@gmail.com"}
+]
+readme = "README.md"
+license = {text = "MIT"}
+dependencies = [
+    "langchain", 
+    "langchain-core",
+    "langgraph",
+    "langchain-openai",
+    "uproot", 
+    "pandas", 
+    "numpy",
+]
+
+requires-python = ">=3.10"
+[project.scripts]
+rooagent = "rooagent.agent:main"
+
+[project.urls]
+Homepage = "https://github.com/amanmdesai/RooAgent"

rooagent-0.1.0/rooagent/__init__.py

@@ -0,0 +1,2 @@
+# Makes rooagent a package
+from .agent import model_with_tools, agent, tools

rooagent-0.1.0/rooagent/agent.py

@@ -0,0 +1,153 @@
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, START, END
+from langchain.messages import HumanMessage, SystemMessage, AnyMessage, ToolMessage
+from typing_extensions import TypedDict, Annotated
+from langgraph.prebuilt import ToolNode
+from langgraph.checkpoint.memory import InMemorySaver  
+import operator
+import os
+# Import all tools
+from .tools import *
+
+# -----------------------------
+# SYSTEM PROMPT (Clean Version)
+# -----------------------------
+SYSTEM_PROMPT = """
+ROOT-based HEP analysis assistant.
+
+RULES:
+1. Always use tools. Never guess.
+2. Verify everything (files, trees, branches) before use.
+3. Workflow: inspect file -> trees -> branches -> analyze.
+4. Auto-use single tree; list if multiple.
+5. Handle missing data and empty results safely.
+6. Compare signal vs background when relevant.
+7. Output: clear plots, labeled axes, report statistics.
+8. No assumptions. State uncertainty if unverified.
+
+Prioritize correctness, efficiency, reproducibility.
+"""
+
+# -----------------------------
+# Initialize Model (GitHub Copilot / GitHub Models API)
+# -----------------------------
+DEFAULT_MODEL_NAME = "gpt-4o-mini"
+MODEL_NAME = os.getenv("MODEL", DEFAULT_MODEL_NAME)
+
+model = ChatOpenAI(
+    model=MODEL_NAME,
+    api_key=os.getenv("GITHUB_TOKEN"),
+    base_url="https://models.inference.ai.azure.com"
+)
+
+# Bind tools
+tools = [
+    list_root_file_contents,
+    list_tree_branches,
+    get_histogram_stats,
+    apply_cut_and_count,
+    compute_significance,
+    plot_tree_variable,
+    compare_tree_variables,
+    root_tree_to_csv,
+    define_variable_and_plot,
+    fit_tree_variable,
+    fit_histogram,
+    draw_histograms_same_canvas,
+    draw_2d_histogram,
+    find_optimal_cut,
+    draw_1d_histogram,
+    define_variable,
+    draw_2d_histogram_from_tree,
+    generate_cutflow,
+    discover_root_data,
+    list_ttrees,
+]
+
+tools_by_name = {tool.name: tool for tool in tools}
+model_with_tools = model.bind_tools(tools)
+
+# -----------------------------
+# State Definition
+# -----------------------------
+class MessagesState(TypedDict):
+    messages: Annotated[list[AnyMessage], operator.add]
+    llm_calls: int
+
+# -----------------------------
+# Tool Node
+# -----------------------------
+def tool_node(state: MessagesState):
+    results = []
+
+    for tool_call in state["messages"][-1].tool_calls:
+
+        tool = tools_by_name[tool_call["name"]]
+        observation = tool.invoke(tool_call["args"])
+
+        results.append(
+            ToolMessage(
+                content=str(observation),
+                tool_call_id=tool_call["id"]
+            )
+        )
+
+    return {"messages": results}
+
+# -----------------------------
+# LLM Node
+# -----------------------------
+def llm_call(state: MessagesState):
+    response = model_with_tools.invoke(
+        [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
+    )
+
+    return {
+        "messages": [response],
+        "llm_calls": state.get("llm_calls", 0) + 1
+    }
+
+# -----------------------------
+# Routing
+# -----------------------------
+def should_continue(state: MessagesState):
+    last_message = state["messages"][-1]
+    if getattr(last_message, "tool_calls", None):
+        return "tool_node"
+    return END
+
+# -----------------------------
+# Build Graph
+# -----------------------------
+builder = StateGraph(MessagesState)
+
+builder.add_node("llm_call", llm_call)
+builder.add_node("tool_node",ToolNode(tools,handle_tool_errors=True)) 
+                                      # tool_node)
+
+builder.add_edge(START, "llm_call")
+builder.add_conditional_edges("llm_call", should_continue, ["tool_node", END])
+builder.add_edge("tool_node", "llm_call")
+
+agent = builder.compile()
+
+# -----------------------------
+# CLI
+# -----------------------------
+def main():
+    print(f"\nROOT Physics Analysis Agent using ({MODEL_NAME})")
+    print("Type 'exit' to quit\n")
+
+    state = {"messages": [], "llm_calls": 0}
+
+    while True:
+        user_input = input("User: ")
+
+        if user_input.lower() in ["exit", "quit"]:
+            print("Thanks for using RooAgent!")
+            break
+
+        # FIXED BUG HERE
+        state["messages"].append(HumanMessage(content=user_input))
+
+        state = agent.invoke(state)

rooagent-0.1.0/rooagent/tools/__init__.py

@@ -0,0 +1,29 @@
+from .fit_tools import *
+from .histogram_tools import *
+from .converter import *
+from .plot_tools import *
+from .rdataframe_tools import *
+from .tfile import *
+from .utils import *
+
+
+# ================= ROOT STYLE =================
+ROOT.gStyle.SetOptStat(0)
+
+ROOT.gStyle.SetTitleFont(42, "XYZ")
+ROOT.gStyle.SetLabelFont(42, "XYZ")
+
+ROOT.gStyle.SetTitleSize(0.05, "XYZ")
+ROOT.gStyle.SetLabelSize(0.04, "XYZ")
+
+ROOT.gStyle.SetPadLeftMargin(0.12)
+ROOT.gStyle.SetPadBottomMargin(0.12)
+
+ROOT.gStyle.SetLegendBorderSize(0)
+ROOT.gStyle.SetLegendFillColor(0)
+
+ROOT.gStyle.SetFrameLineWidth(2)
+ROOT.gStyle.SetLineWidth(2)
+
+ROOT.gStyle.SetPadGridX(False)
+ROOT.gStyle.SetPadGridY(False)

rooagent-0.1.0/rooagent/tools/converter.py

@@ -0,0 +1,37 @@
+from typing import List
+import ROOT
+import pandas as pd
+from langchain_core.tools import tool
+
+@tool
+def root_tree_to_csv(file_path: str, tree_name: str, branches: List[str], output_csv: str, max_vector_size: int = 5) -> str:
+    """
+    Convert ROOT TTree to CSV, flattening vector branches.
+
+    Parameters
+    ----------
+    file_path : str
+        Path to ROOT file.
+    tree_name : str
+        Name of TTree.
+    branches : List[str]
+        Branches to extract (scalar or vector).
+    output_csv : str
+        Output CSV file path.
+    max_vector_size : int
+        Max number of elements for vector branches (columns will be suffixed _0, _1, ...).
+    """
+    df = ROOT.RDataFrame(tree_name, file_path)
+    data = df.AsNumpy(branches)
+
+    flat_dict = {}
+    for branch, array in data.items():
+        if array.ndim == 1:  # scalar
+            flat_dict[branch] = array
+        else:  # vector branch
+            for i in range(max_vector_size):
+                flat_dict[f"{branch}_{i}"] = [a[i] if i < len(a) else None for a in array]
+
+    pandas_df = pd.DataFrame(flat_dict)
+    pandas_df.to_csv(output_csv, index=False)
+    return f"Flattened CSV saved to {output_csv}"

rooagent-0.1.0/rooagent/tools/fit_tools.py

@@ -0,0 +1,173 @@
+import ROOT
+from langchain_core.tools import tool
+
+@tool
+def fit_tree_variable(
+    file_path: str,
+    tree_name: str,
+    variable: str,
+    bins: int,
+    xmin: float,
+    xmax: float,
+    fit_function: str,
+    output_plot: str
+) -> str:
+    """
+    Create a histogram from a ROOT TTree variable, perform a statistical fit,
+    and save the resulting plot with legend.
+
+    Parameters
+    ----------
+    file_path : str
+        Path to the ROOT file containing the TTree.
+    tree_name : str
+        Name of the TTree to read.
+    variable : str
+        Name of the variable to histogram and fit.
+    bins : int
+        Number of bins in the histogram.
+    xmin : float
+        Minimum x-axis value.
+    xmax : float
+        Maximum x-axis value.
+    fit_function : str
+        ROOT fit function (examples: 'gaus', 'expo', 'pol1', 'pol2').
+    output_plot : str
+        File name where the canvas with histogram and fit will be saved.
+
+    Returns
+    -------
+    str
+        Summary of fit parameters and chi-square/NDF.
+    """
+    df = ROOT.RDataFrame(tree_name, file_path)
+    hist_ptr = df.Histo1D((variable, variable, bins, xmin, xmax), variable)
+    h = hist_ptr.GetValue()
+    ROOT.SetOwnership(h, False)
+
+    canvas = ROOT.TCanvas("c_fit", "", 900, 700)
+    #canvas.SetGrid()
+
+    # Perform fit
+    h.Fit(fit_function, "S")
+
+    # Style histogram
+    h.SetLineWidth(3)
+    h.SetLineColor(ROOT.kBlue + 1)
+    h.SetMarkerStyle(20)
+    h.SetMarkerSize(1.2)
+    h.GetXaxis().SetTitle(variable)
+    h.GetYaxis().SetTitle("Events")
+    h.Draw("E")
+
+    # Style fit function
+    func = h.GetFunction(fit_function)
+    if func:
+        func.SetLineColor(ROOT.kRed)
+        func.SetLineWidth(2)
+
+    # Add legend
+    legend = ROOT.TLegend(0.65, 0.7, 0.88, 0.88)
+    legend.SetBorderSize(0)
+    legend.SetFillStyle(0)
+    legend.AddEntry(h, "Histogram", "lep")
+    if func:
+        legend.AddEntry(func, f"Fit: {fit_function}", "l")
+    legend.Draw()
+
+    canvas.Update()
+    canvas.SaveAs(output_plot)
+
+    # Extract fit parameters
+    params = [f"p{i} = {func.GetParameter(i):.4f}" for i in range(func.GetNpar())]
+    chi2 = func.GetChisquare()
+    ndf = func.GetNDF()
+
+    return (
+        f"Fit function: {fit_function}\n"
+        f"Parameters: {', '.join(params)}\n"
+        f"Chi2/NDF: {chi2:.3f}/{ndf}"
+    )
+
+
+@tool
+def fit_histogram(
+    file_path: str,
+    hist_name: str,
+    fit_function: str,
+    output_plot: str
+) -> str:
+    """
+    Fit an existing histogram in a ROOT file, overlay the fit, and save with legend.
+
+    Parameters
+    ----------
+    file_path : str
+        Path to the ROOT file containing the histogram.
+    hist_name : str
+        Name of the histogram to fit.
+    fit_function : str
+        ROOT fit function (examples: 'gaus', 'expo', 'pol1', 'pol2').
+    output_plot : str
+        File name where the canvas with histogram and fit will be saved.
+
+    Returns
+    -------
+    str
+        Summary of fit parameters and chi-square/NDF.
+    """
+    f = ROOT.TFile.Open(file_path)
+    if not f or f.IsZombie():
+        return "Error: could not open ROOT file."
+
+    h = f.Get(hist_name)
+    if not h:
+        f.Close()
+        return f"Histogram '{hist_name}' not found."
+
+    ROOT.SetOwnership(h, False)
+
+    canvas = ROOT.TCanvas("c_fit_hist", "", 900, 700)
+    #canvas.SetGrid()
+
+    # Perform fit
+    h.Fit(fit_function, "S")
+
+    # Style histogram
+    h.SetLineWidth(3)
+    h.SetLineColor(ROOT.kBlue + 1)
+    h.SetMarkerStyle(20)
+    h.SetMarkerSize(1.2)
+    h.Draw("E")
+
+    # Style fit function
+    func = h.GetFunction(fit_function)
+    if func:
+        func.SetLineColor(ROOT.kRed)
+        func.SetLineWidth(2)
+
+    # Add legend
+    legend = ROOT.TLegend(0.65, 0.7, 0.88, 0.88)
+    legend.SetBorderSize(0)
+    legend.SetFillStyle(0)
+    legend.AddEntry(h, "Histogram", "lep")
+    if func:
+        legend.AddEntry(func, f"Fit: {fit_function}", "l")
+    legend.Draw()
+
+    canvas.Update()
+    canvas.SaveAs(output_plot)
+
+    # Extract fit parameters
+    params = [f"p{i}={func.GetParameter(i):.4f}" for i in range(func.GetNpar())]
+    chi2 = func.GetChisquare()
+    ndf = func.GetNDF()
+
+    f.Close()
+
+    return (
+        f"Histogram fit completed.\n"
+        f"Function: {fit_function}\n"
+        f"Parameters: {', '.join(params)}\n"
+        f"Chi2/NDF: {chi2:.3f}/{ndf}"
+    )

rooagent-0.1.0/rooagent/tools/histogram_tools.py

@@ -0,0 +1,38 @@
+from pydantic import BaseModel
+import ROOT
+from langchain_core.tools import tool
+
+
+@tool
+def get_histogram_stats(file_path: str, hist_name: str) -> str:
+    """
+    Retrieve basic statistical information (mean, RMS, entries) of a histogram
+    stored in a ROOT file.
+
+    Parameters
+    ----------
+    file_path : str
+        Path to the ROOT file containing the histogram.
+    hist_name : str
+        Name of the histogram within the ROOT file.
+
+    Returns
+    -------
+    str
+        A formatted string reporting the histogram's mean, RMS, and total entries.
+        Example:
+            "myHist -> Mean: 0.123, RMS: 1.234, Entries: 1000"
+
+    """
+    f = ROOT.TFile.Open(file_path)
+    if not f or f.IsZombie():
+        return f"Error: could not open file {file_path}."
+    h = f.Get(hist_name)
+    if not h:
+        f.Close()
+        return f"Error: histogram '{hist_name}' not found in file {file_path}."
+    mean = h.GetMean()
+    rms = h.GetRMS()
+    entries = int(h.GetEntries())
+    f.Close()
+    return f"{hist_name} -> Mean: {mean:.3f}, RMS: {rms:.3f}, Entries: {entries}"