from recap.agents.base import BaseAgent
from pydantic import BaseModel, Field
from dataclasses import dataclass, asdict
from typing import Optional, Literal
from pathlib import Path
import pandas as pd
import json
import yaml
[docs]
@dataclass
class AnalyticsPlan:
operation: Literal["group_agg", "column_agg", "sort", "value_counts"]
# --- group_agg: group by columns, aggregate another column ---
# --- column_agg: empty groupby, aggregate single column ---
group_by: Optional[list[str]] = None
agg_column: Optional[str] = None
aggregation: Optional[Literal["sum", "mean", "count", "min", "max"]] = None
# --- sort + limit (applies to any operation) ---
sort_column: Optional[str] = None
sort_desc: bool = True
top_k: Optional[int] = None
# --- value_counts: counts values or percentages of single column
# distribution of a single column ---
column: Optional[str] = None # for value_counts only
normalize: bool = False # return percentages
[docs]
class AnalyticsAgent(BaseAgent):
"""
Agent responsible for performing analytical operations on retrieved review
data.
This agent uses an LLM to generate an analytics plan from a natural
language query, then executes that plan against records loaded from
retrieval artifacts. Supported operations include grouping, aggregation,
sorting, value counting, and selecting the top results.
"""
description = (
"The analytics agent is a tool used to derive analytics on a set of retrieval artifact_refs. It answers questions"
"requiring grouping, aggregration (mean,max,min,sum), top_k or sorting of columns of the provided results. Note retrievals"
"should be broad enough to allow analytic tool to analyze"
)
args_schema = AnalyticsAgentInput
name = "analytics"
def __init__(
self,
model,
artifact_dir="./artifacts",
) -> None:
super().__init__(model=model)
self.artifact_dir = Path(artifact_dir)
def _get_schema(self) -> dict:
"""
A function to output schema format for the retrieval agent. In future
iterations this should be pulled programmatically to make agent more
generalizable
Returns:
Dictionary with the schema of the ChromaDB metadata
"""
schema = {
"fields": [
{
"name": "content",
"type": "string",
"description": "The text content of the review",
"example": "...",
},
{
"name": "paper_id",
"type": "string",
"description": "ID for the paper usually in format [year]-[number]",
"example": "2018-12",
},
{
"name": "year",
"type": "number",
"description": "Year the article was published",
"example": 2018,
},
{
"name": "openreview_id",
"type": "string",
"description": "ID for the paper usually in format Paper[number]",
"example": "Paper45",
},
{
"name": "reviewer_ratings",
"type": "string",
"description": "List of individual rating",
"example": "R1: 8, R2: 7, R3: 7. ",
},
{
"name": "average_rating",
"type": "number",
"description": "The average rating given by reviews as a float",
"example": 7.333333,
},
{
"name": "decision",
"type": "string",
"description": "The determination of whether paper was accepted or not it has these values: Accept (Oral), Accept (Presentation), Accept (Workshop), Accept (Spotlight), Reject, Reject (Invite to Workshop Track)",
"example": "Accept (Poster)",
},
{
"name": "title",
"type": "string",
"description": "Title of the paper being reviewed",
"example": "Spectral Normalization for Generative Adversarial Networks",
},
{
"name": "link",
"type": "string",
"description": "Link to the open review page for reviews in format https://openreview.net/forum?id=[reviewid]",
"example": "https://openreview.net/forum?id=B1QRgziT-",
},
{
"name": "word_count",
"type": "number",
"description": "The number of words within a review",
"example": 93,
},
{
"name": "sentence_classification",
"type": "list",
"description": "List of classifications of sentences making up the review",
"example": "['abstract', 'abstract', 'strength', 'rebuttal_process', 'strength', 'decision']",
},
{
"name": "segment_classification",
"type": "list",
"description": "List of classifications of sentences making up the review",
"example": "['abstract', 'abstract', 'strength', 'rebuttal_process', 'strength', 'decision']",
},
]
}
return schema
[docs]
def generate_plan(
self,
query: str,
df: pd.DataFrame, # TODO: Derive schema from dataframe to make it more generic
) -> AnalyticsPlan:
msg = f"Schema: {yaml.dump(self._get_schema())}\nQuestion: {query}"
agent_state = self.agent.invoke({"messages": [("user", msg)]}) # type: ignore[override]
print(agent_state["messages"][-1].content)
json_response, _ = self.parse_json_response(agent_state["messages"][-1].content)
# Parse LLM Response into Analytics Plan
if not isinstance(json_response, dict):
raise Exception("No Valid JSON response produced")
return AnalyticsPlan(**json_response)
[docs]
def execute_plan(self, plan: AnalyticsPlan, df: pd.DataFrame) -> pd.DataFrame:
"""Given a pandas DataFrame and AnaltyicsPlan execute plan
and return the resulting dataframe
Args:
df: DataFrame to run analytics on
plan: AnalyticsPlan
Returns:
Pandas DataFrame with executed analytics
"""
# 2. Main operation
if plan.operation == "group_agg":
if plan.agg_column:
result = (
df.groupby(plan.group_by)[plan.agg_column]
.agg(plan.aggregation)
.reset_index()
)
elif plan.operation == "single_agg":
result = df[agg_column].agg(plan.aggregation)
elif plan.operation == "value_counts":
result = (
df[plan.column].value_counts(normalize=plan.normalize).reset_index()
)
result.columns = [plan.column, "percentage" if plan.normalize else "count"]
elif plan.operation == "sort":
result = df.sort_values(
plan.sort_column or df.columns[-1], ascending=not plan.sort_desc
)
else:
result = df
# 3. Optional sort + limit on the result
if plan.sort_column and plan.operation != "sort":
result = result.sort_values(plan.sort_column, ascending=not plan.sort_desc)
if plan.top_k:
result = result.head(plan.top_k)
return result
[docs]
def invoke(
self,
query: str,
artifact_refs: list[str],
) -> str:
"""
Given a query and a set of artifact_refs run analysis on the
provided artifacts (groupby, average, top_k,e tc.)
Args:
query - List of string contents to summarize
artifact_refs - a list of retrieval references to use
Returns:
Str of the JSON with analytics plan and results
"""
contents = []
# Load artifact contents
if artifact_refs:
for artifact_ref in artifact_refs:
with open(artifact_ref, "r") as fp:
try:
ref_contents = json.load(fp=fp)
except:
ref_content = []
for ref_content in ref_contents:
contents.append(ref_content)
df = pd.DataFrame(contents)
analytics_plan = self.generate_plan(query, df)
results = self.execute_plan(analytics_plan, df)
return json.dumps(
{"plan": asdict(analytics_plan), "analytics_tsv": results.to_csv(sep="\t")}
)