Stata-MCP

#!/usr/bin/python3 # -*- coding: utf-8 -*- # # Copyright (C) 2025 - Present Sepine Tam, Inc. All Rights Reserved # # @Author : Sepine Tam (谭淞) # @Email : sepinetam@gmail.com # @File : repl_agents.py import asyncio import os import textwrap import uuid from pathlib import Path from agents import Agent, OpenAIChatCompletionsModel, Runner, set_tracing_disabled from agents.mcp import MCPServerStdio from agents.memory.sqlite_session import SQLiteSession from openai_api_polling.polling import APIPolling, ClientPolling set_tracing_disabled(disabled=True) def get_model_instructions(work_dir: str) -> str: """获取模型指令，包含当前工作目录信息""" return textwrap.dedent(f""" # 角色定义 # [定义助手的核心身份和专业领域] # 核心能力 # [列出主要技能和功能模块] # 工作原则 # [指导性行为准则和优先级] # - 所有操作都在当前工作目录进行: {work_dir} # - 执行 Stata 命令前先确保切换到正确的目录: cd "{work_dir}" # 分析流程 # [标准化的分析步骤和方法] # 输出标准 # [结果呈现的格式和质量要求] # 代码规范 # [生成代码的最佳实践] # 交互风格 # [与用户沟通的方式和特点] """) class REPLAgent: client_polling = ClientPolling( api_keys=[ os.getenv("STATA_MCP_API_KEY", os.getenv("OPENAI_API_KEY", None)) ], base_url=os.getenv("STATA_MCP_API_BASE_URL", os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")) ) llm = os.getenv( "STATA_MCP_MODEL", os.getenv( "OPENAI_MODEL", "gpt-3.5-turbo" ) ) def __init__(self, work_dir: str = "./", session_id: str = None): self.work_dir = Path(work_dir).expanduser().absolute() self.work_dir.mkdir(exist_ok=True) self.session_id = session_id or f"stata_session_{uuid.uuid4().hex[:8]}" session_db_path = self.work_dir / ".stata_sessions.db" self.session = SQLiteSession( session_id=self.session_id, db_path=session_db_path ) self.stata_mcp_server = MCPServerStdio( name="Stata-MCP", params={ "command": "uvx", "args": [ "stata-mcp" # or you can use the local beta version with the git clone repo. ], "env": { "STATA_MCP_CWD": self.work_dir.as_posix(), } }, cache_tools_list=True, client_session_timeout_seconds=60, max_retry_attempts=3 ) async_client = self.client_polling.async_client self.agent = Agent( name="Stata Assistant", instructions=get_model_instructions(self.work_dir.as_posix()), mcp_servers=[self.stata_mcp_server], model=OpenAIChatCompletionsModel( model=self.llm, openai_client=async_client ) ) print("Current Agent Config: ") print(f">>> API KEY\t: {APIPolling.mask_api_key(async_client.api_key, mask_len=20)}") print(f">>> BASE URL\t: {async_client.base_url}") print(f">>> MODEL\t: {self.llm}") def __del__(self): async def clean_up(): try: await self.stata_mcp_server.cleanup() except Exception: pass asyncio.run(clean_up()) async def invoke(self, query: str): await self.stata_mcp_server.connect() try: result = await Runner.run( self.agent, query, session=self.session, max_turns=30 ) except Exception as e: print("Something went wrong") return str(e) finally: await self.stata_mcp_server.cleanup() return result.final_output def run(self): print("Welcome to Stata-MCP built-in Agent!") print("You can type what you want to do with Stata-MCP after sign `> `") print("Type `/exit` or `bye` to exit") while True: query = input("> ") if query.lower() == "/exit" or query.lower() == "bye": break result = asyncio.run(self.invoke(query)) print(result) exit(0)