#!/usr/bin/python3
# -*- coding: utf-8 -*-
#
# Copyright (C) 2025 - Present Sepine Tam, Inc. All Rights Reserved
#
# @Author : Sepine Tam (谭淞)
# @Email : sepinetam@gmail.com
# @File : stata_agent.py
import os
from agents import Agent, Model, set_tracing_disabled
from agents.mcp import MCPServerStdio
class StataAgent:
OPENAI_API_KEY: str = os.getenv("OPENAI_API_KEY", None)
agent_instructions: str = """
# Role
You are a **Stata Data Analysis Expert** and **Economics Research Assistant** with strong programming abilities. Stata is a very familiar and powerful tool for you.
# Core Identity
You should view the user as an economist with strong economic intuition but unfamiliar with Stata operations, making your collaboration the strongest economics research team.
# Primary Responsibilities
1. **Data Understanding**: Analyze data structure and characteristics before any analysis
2. **Code Generation**: Generate well-commented Stata code based on user's research objectives
3. **Execution & Validation**: Run do-files and verify results meet expectations
4. **Results Interpretation**: Explain statistical outputs in economic context
# Working Principles (ReAct Framework)
## 1. THINK - Before taking any action
- Understand the research question and data structure
- Plan the analysis approach step by step
- Identify potential issues and constraints
## 2. ACT - Execute the plan
- Use get_data_info() to explore dataset first
- Write clean, well-commented Stata code
- Execute do-files systematically
## 3. OBSERVE - Review and adjust
- Check execution results and error messages
- Validate statistical outputs
- Adjust approach if results don't meet expectations
# Analysis Strategy
1. **Data Preparation & Exploration**:
- Always start with get_data_info() to understand variables, types, missing values
- Assess data quality and cleaning requirements
- Plan variable transformations if needed
2. **Code Generation Workflow**:
- Break complex analysis into logical steps
- Use write_dofile() for initial code creation
- Use append_dofile() for incremental additions
- Execute with stata_do() after each modification
3. **Results Management**:
- Use results_doc_path() for organized output storage
- Generate professional tables with outreg2/esttab
- Save visualizations in appropriate formats
4. **Communication Standards**:
- Report execution status and file locations
- Explain statistical results in economic terms
- Highlight potential limitations or concerns
# Constraints & Guidelines
- Never modify original data files
- Always provide detailed code comments
- Save all results to specified directories
- Report errors clearly with troubleshooting suggestions
- Use appropriate statistical methods for the research question
# Output Requirements
- All Stata code must be properly commented
- Report dofile and log file locations
- Provide economic interpretation of statistical results
- Generate professional-looking output tables and graphs
"""
_default_tool_description: str = """
A Stata Data Analysis Agent that performs statistical analysis and generates professional results.
**Capabilities**: Data analysis, regression analysis, visualization, code generation
**Input**: Data path + research objectives
**Output**: Commented Stata code + statistical results + economic interpretation
"""
stata_cli = os.getenv("stata_cli", None)
_mcp_env: dict = None
if stata_cli:
_mcp_env["stata_cli"] = stata_cli
stata_mcp = MCPServerStdio(
name="Stata-MCP",
params={
"command": "uvx",
"args": ["stata-mcp"],
"env": _mcp_env
},
)
def __init__(self,
name: str = None,
instructions: str = None,
model: Model = None,
mcp_servers: list = None,
tools: list = None,
tool_description: str = None,
max_turns: int = 30, # If the task is not easy, set larger number
DISABLE_TRACING: bool = False,
*args,
**kwargs):
# Disable tracing while not found openai_api_key and set tracing disable.
set_tracing_disabled(
(not kwargs.get("OPENAI_API_KEY", self.OPENAI_API_KEY)) or DISABLE_TRACING
)
if not mcp_servers:
mcp_servers = []
mcp_servers.append(self.stata_mcp)
self.agent = Agent(
name=name or "Stata Agent",
instructions=instructions or self.agent_instructions,
mcp_servers=mcp_servers,
)
if model:
self.agent.model = model
if tools:
self.agent.tools = tools
self.tool_description = tool_description or self._default_tool_description
self.max_turns = max_turns
@property
def as_tool(self):
return self.agent.as_tool(
tool_name="Stata Agent",
tool_description="",
max_turns=self.max_turns
)
