assathe/data-analyst-agent2

1# main.py2import os3import tempfile4import json5import traceback6from fastapi import FastAPI, UploadFile, File, Form, Body, Request7from fastapi.responses import JSONResponse8from fastapi.middleware.cors import CORSMiddleware9import shutil10from typing import List, Optional, Union, Dict11import asyncio12from claude_client import query_claude13from code_executor import execute_user_code14from dotenv import load_dotenv15load_dotenv()16 17app = FastAPI(title="Data Analyst Agent", version="1.0.0")18 19# Add CORS middleware for Vercel deployment20app.add_middleware(21    CORSMiddleware,22    allow_origins=["*"],23    allow_credentials=True,24    allow_methods=["*"],25    allow_headers=["*"],26)27 28SYSTEM_PROMPT = """29You are an expert data analyst agent capable of web scraping, data analysis, and visualization. You must write executable Python code to answer questions.30Available libraries: pandas, numpy, matplotlib, seaborn, plotly, requests, BeautifulSoup, duckdb, json, base64, io, re, datetime, time31CRITICAL REQUIREMENTS:321. Generate ONLY valid Python code inside `````` markdown blocks332. NO explanations or text outside code blocks343. Handle errors gracefully with try-except blocks354. For plots: save as PNG, convert to base64, format as "data:image/png;base64,..."365. Keep base64 images under 100KB376. Assign final result to variable `final_answer`387. For JSON responses, ensure proper formatting398. For web scraping, use requests + BeautifulSoup409. For data analysis, use pandas/duckdb as appropriate4110. Always include proper imports at the top4211. IMPORTANT: Clean data thoroughly - remove currency symbols, commas, handle non-numeric strings4312. Use regex and pandas methods to clean text data before converting to numbers44DATA CLEANING GUIDELINES:45- Remove $ signs, commas, and other currency formatting46- Extract only numeric parts from strings47- Handle footnote references and citations in scraped data48- Use pd.to_numeric() with errors='coerce' for safe conversion49- Always validate data types before mathematical operations50RESPONSE FORMATS:51- Single values: final_answer = value52- JSON arrays: final_answer = [item1, item2, item3, item4]53- JSON objects: final_answer = {"key1": "value1", "key2": "value2"}54- Base64 images: final_answer = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAA..."55For plots:56import matplotlib.pyplot as plt57import base6458import io59 60Your plotting code here61plt.figure(figsize=(8, 6))62 63... plot creation ...64Save plot to base6465buffer = io.BytesIO()66plt.savefig(buffer, format='png', dpi=100, bbox_inches='tight')67buffer.seek(0)68image_base64 = base64.b64encode(buffer.getvalue()).decode()69plt.close()70final_answer = f"data:image/png;base64,{image_base64}"71 72Remember: NO text outside code blocks. Only executable Python code.73"""74 75def get_python_code_from_claude(question_text: str, file_contexts: dict, error_context: str = "") -> str:76    """Get Python code from Claude for data analysis tasks."""77 78    file_info = ""79    if file_contexts:80        file_info = "\n\nAvailable files:\n"81        for filename, filepath in file_contexts.items():82            if filename.endswith('.csv'):83                file_info += f"- {filename}: CSV file at '{filepath}'\n"84            elif filename.endswith('.json'):85                file_info += f"- {filename}: JSON file at '{filepath}'\n"86            elif filename.endswith('.txt'):87                file_info += f"- {filename}: Text file at '{filepath}'\n"88            else:89                file_info += f"- {filename}: File at '{filepath}'\n"90 91    error_info = ""92    if error_context:93        error_info = f"\n\nPREVIOUS ERROR TO FIX:\n{error_context}\n"94 95    user_content = f"""Question: {question_text}{file_info}{error_info}96Generate Python code to answer this question. Use file_contexts dictionary to access file paths."""97 98    messages = [99        {"role": "user", "content": user_content}100    ]101 102    return query_claude(messages, system_message=SYSTEM_PROMPT)103 104def extract_python_code(claude_response: str) -> str:105    """Extract Python code from Claude's response."""106    import re107 108    # Look for code blocks109    code_blocks = re.findall(r'``````', claude_response, re.DOTALL)110 111    if code_blocks:112        return code_blocks[0].strip()113 114    # If no code blocks found, look for any Python-like content115    lines = claude_response.split('\n')116    python_lines = []117    in_code = False118 119    for line in lines:120        if any(keyword in line for keyword in ['import ', 'def ', 'class ', 'if ', 'for ', 'while ', 'try:', 'final_answer']):121            in_code = True122        if in_code:123            python_lines.append(line)124 125    if python_lines:126        return '\n'.join(python_lines)127 128    # Last resort: return the whole response129    return claude_response130 131@app.post("/api/")132async def handle_request(request: Request):133    temp_dir = tempfile.mkdtemp()134    try:135        content_type = request.headers.get("content-type", "")136        question_text = None137        file_contexts = {}138 139        if "multipart/form-data" in content_type:140            form = await request.form()141            for key, value in form.items():142                if isinstance(value, UploadFile):143                    path = os.path.join(temp_dir, value.filename or key)144                    content = await value.read()145                    with open(path, "wb") as f:146                        f.write(content)147                    # Accept either key or filename containing 'questions.txt'148                    if key.lower().endswith("questions.txt") or (value.filename and value.filename.lower().endswith("questions.txt")):149                        question_text = content.decode("utf-8").strip()150                    else:151                        file_contexts[value.filename or key] = path152                elif isinstance(value, str):153                    # If field sent as plain text154                    if key.lower().endswith("questions.txt"):155                        question_text = value.strip()156        elif "application/json" in content_type:157            body = await request.json()158            # ----------- MODIFIED FIELD PICKING LOGIC HERE ----------159            question_text = (160                body.get("questions", "") or161                body.get("question", "") or162                body.get("prompt", "")163            )164            if question_text is not None:165                question_text = str(question_text).strip()166            else:167                question_text = ""168            incoming_files = body.get("files", [])169            for i, item in enumerate(incoming_files):170                path = os.path.join(temp_dir, f"file_{i}")171                with open(path, "w", encoding="utf-8") as f:172                    f.write(item)173                file_contexts[f"file_{i}"] = path174        else:175            return JSONResponse(176                status_code=400, content={177                    "error": "Unsupported content-type",178                    "received_content_type": content_type179                }180            )181        if not question_text:182            # Improved error message -- show all fields received183            try:184                body_keys = []185                if "application/json" in content_type:186                    body = await request.json()187                    body_keys = list(body.keys())188                error_fields = {189                    "error": "No prompt text found in any uploaded file/field",190                    "hint": "For application/json, provide any of ['questions','question','prompt']",191                    "received_fields": body_keys192                }193            except Exception:194                error_fields = {195                    "error": "No prompt text found in any uploaded file/field",196                    "hint": "For application/json, provide any of ['questions','question','prompt']"197                }198            return JSONResponse(status_code=400, content=error_fields)199 200        # --- Now invoke your agent logic ---201        globals_dict = {"file_contexts": file_contexts, "temp_dir": temp_dir}202        max_attempts, attempt, success = 3, 0, False203        last_error = ""204        while attempt < max_attempts and not success:205            attempt += 1206            try:207                claude_response = get_python_code_from_claude(208                    question_text, file_contexts, last_error if attempt > 1 else ""209                )210                python_code = extract_python_code(claude_response)211                if not python_code:212                    raise ValueError("No valid Python code found from model")213                result = execute_user_code(python_code, globals_dict.copy())214                if result["success"]:215                    success = True216                    final_answer = result["globals"].get("final_answer", None)217                    if isinstance(final_answer, (dict, list)):218                        return JSONResponse(content=final_answer)219                    elif isinstance(final_answer, str):220                        try:221                            parsed = json.loads(final_answer)222                            return JSONResponse(content=parsed)223                        except Exception:224                            return JSONResponse(content={"result": final_answer})225                    else:226                        return JSONResponse(content={"result": str(final_answer)})227                else:228                    last_error = result["error"]229            except Exception as e:230                last_error = f"Attempt {attempt} failed: {e}\n{traceback.format_exc()}"231        # If we got here => failed after attempts232        return JSONResponse(233            status_code=500,234            content={"error": "Code execution failed", "details": last_error}235        )236    except Exception as e:237        return JSONResponse(238            status_code=500,239            content={"error": "Request processing failed", "details": traceback.format_exc()}240        )241    finally:242        shutil.rmtree(temp_dir, ignore_errors=True)243 244@app.get("/")245async def root():246    """Health check endpoint."""247    return {"status": "Data Analyst Agent is running", "version": "1.0.0"}248 249@app.get("/health")250async def health():251    """Health check for deployment."""252    return {"status": "healthy"}253 254if __name__ == "__main__":255    import uvicorn256    uvicorn.run(app, host="0.0.0.0", port=8000)